DyT：何凯明LeCun联手，删除归一化层，性能不减速度反升！

最新推荐文章于 2025-04-19 22:34:52 发布

Sherlock Ma

最新推荐文章于 2025-04-19 22:34:52 发布

阅读量1k

点赞数 19

分类专栏：经典网络架构 Transformer 文章标签： transformer 归一化人工智能神经网络 pytorch 论文阅读

本文链接：https://blog.csdn.net/sherlockMa/article/details/146315115

版权

经典网络架构同时被 2 个专栏收录

3 篇文章

订阅专栏

Transformer

1 篇文章

订阅专栏

1.简介

在过去的十年中，归一化层（Normalization Layer）已经巩固了其作为现代神经网络最基本组件之一的地位，它使得视觉识别模型的收敛速度大大加快，效果更好，并在随后的几年中迅速获得势头。如今，几乎所有现代网络都使用归一化图层。除了取得更好的结果外，它们还有助于加速和稳定收敛。随着神经网络变得更宽和更深，这种必要性变得越来越重要。近年来，新的架构经常试图取代注意或卷积层，但几乎总是保留规格化层。

Dynamic Tanh（DyT）是由刘壮团队联合何恺明、Yann LeCun等人提出的一种新型神经网络元素级运算，旨在替代传统的归一化层，如Layer Normalization（LN）和RMSNorm。

本文通过在Transformers中引入一个简单的规范化层替代方案来挑战这一信念。作者的探索从观察LN层将其输入映射到输出（具有双曲正切、S形曲线）开始，缩放输入激活，同时压缩极值。受此启发，作者提出了一种元素级的操作，称为动态双曲正切函数（DyT），定义为：DyT（x）= tanh（αx），其中α是一个可学习的参数。该操作旨在通过α学习适当的缩放因子并通过有界tanh函数挤压极值来模拟LN的行为。值得注意的是，与规格化层不同，它无需计算激活统计信息即可实现这两种效果。

前Salesforce首席科学家、搜索引擎You的CEO Richard Socher表示，这项成果加强了他之前的假设——原始的Transformer只是众多等效神经结构之中的一个。

github地址：GitHub - jiachenzhu/DyT: Code release for DynamicTanh (DyT)

论文地址：[2503.10622v1] Transformers without Normalization

2.论文详解

背景

大多数归一化层共享一个通用公式。给定形状为（B，T，C）的输入x，其中B是批量大小，T是令牌的数量，C是每个令牌的嵌入维数，输出通常计算为： $normalization(x)=\gamma \cdot (\frac{x-\mu }{\sqrt{\sigma ^2+\epsilon }})+\beta$ ，其中 $\epsilon$ 是一个小常数，γ和β是形状为（C）的可学习向量参数。它们是“缩放”和“移位”仿射参数，允许输出在任何范围内。µ和 $\sigma ^2$ 表示输入的平均值和方差。不同方法的主要区别在于如何计算这两个统计量。这导致µ和 $\sigma ^2$ 具有不同的维数，在计算过程中分别应用广播。

批归一化（Batch Normalization，简称BN）和层归一化（Layer Normalization，简称LN）都是深度学习中用于优化神经网络训练过程的技术，它们的主要目的是通过归一化操作来改善网络的训练性能。

批归一化是在训练过程中对每个小批量（batch）的数据进行归一化处理。具体来说，它对每个特征维度上的数据进行操作，计算该维度在当前小批量中的均值和方差，然后用这些统计量对数据进行标准化，使其均值为0，方差为1。这样做的好处是能够减少内部协变量偏移，即在训练过程中，网络中间层的输入分布会不断发生变化，这会影响网络的训练速度和稳定性。通过批归一化，可以使得网络中间层的输入分布更加稳定，从而加快训练速度，提高模型的收敛性能。此外，批归一化还可以在一定程度上起到正则化的作用，有助于缓解过拟合问题。不过，批归一化依赖于小批量的统计信息，当小批量大小较小时，计算得到的均值和方差可能会不够准确，从而影响归一化的效果。而且在推理阶段（模型训练完成后进行预测时），由于没有小批量的概念，需要使用训练阶段计算得到的全局均值和方差来进行归一化。
层归一化则是对单个样本的所有特征进行归一化操作，而不是像批归一化那样基于小批量的统计信息。它对每个样本的所有特征值计算均值和方差，然后对这些特征值进行标准化。层归一化的一个显著特点是它不依赖于小批量的大小，因此在处理小批量数据或者单样本数据时仍然能够有效工作。这使得它在一些特定的应用场景下，如自然语言处理中的序列模型（例如Transformer架构）中表现出色，因为这些模型在处理长序列数据时，小批量的划分可能会比较复杂，而层归一化可以更好地适应这种情况。不过，层归一化可能会在一定程度上减弱模型对不同特征的区分能力，因为它对每个样本的所有特征进行了统一的归一化处理，而没有像批归一化那样考虑不同特征维度之间的差异。

归一化层的功能？

分析设置：作者首先对训练网络中归一化层的行为进行了实证研究。作者采用ViT-B在ImageNet-1 K上进行训练、wav 2 vec 2.0在LibriSpeech上训练，以及在ImageNet-1K上训练的DiT-XL。在所有情况下，LN都应用于每个Transformer块中，并在最终线性层之前应用。

对于所有三个训练好的网络，作者对一小批样本进行采样，并在网络中进行前向传递。然后测量归一化层的输入和输出。由于LN保持了输入张量的维数，因此可以在输入和输出张量元素之间建立一对一的对应关系，从而可以直接可视化它们的关系。作者将结果映射绘制到下图中。

具有类似Tanh的映射：对于所有三个模型，在早期的LN层（第1列）中，作者发现这种输入-输出关系大部分是线性的，类似于x-y图中的直线。然而，更深的LN层是我们进行更有趣观察的地方。

从这些更深层的观察中可以明显看出，大多数曲线的形状与tanh函数所表示的完全或部分S形曲线非常相似。

对于这种S形曲线，作者注意到，由x值接近零的点表示的中心部分仍然主要是线性形状。大多数点（约99%）都在这个线性范围内。然而，仍有许多点明显超出该范围，被认为具有“极端”值，例如，在ViT模型中，x大于50或小于-50。归一化层对这些值的主要影响是将它们压缩为较小的极端值，与大多数点更加一致。这是无法通过简单的仿射变换层来近似归一化层的地方。作者假设，正是这种对极值的非线性和不成比例的挤压效应使得归一化层变得重要和不可或缺。

通过令牌和通道进行归一化：LN图层如何对每个标记执行线性变换，同时以这种非线性方式压缩极值？为了理解这一点，作者将分别按令牌和通道分组的点可视化。

在图4左边的两个面板中，作者使用相同的颜色来显示每个令牌的激活。可以观察到，来自任何单个记号的所有点都形成一条直线。然而，由于每个令牌具有不同的方差，因此斜率是不同的。具有较小输入x范围的令牌往往具有较小的方差，并且归一化层将使用较小的标准差划分其激活，从而在直线中产生较大的斜率。总的来说，它们形成了类似于双曲正切函数的S形曲线。在右边的两个面板中，作者使用相同的颜色为每个通道的激活着色。可以发现，不同的通道往往具有截然不同的输入范围，仅有少数通道（例如，红色、绿色和粉红色），表现出大的极值。这些是被归一化层挤压最多的通道。

图4：两个LN层的输出与输入，张量元素被着色以指示不同的通道和令牌维度。输入张量的形状为（样本、标记和通道），通过为相同的标记（左侧两个面板）和通道（右侧两个面板）分配一致的颜色来可视化元素。左两个面板：表示相同标记（相同颜色）的点形成跨不同通道的直线，因为LN针对每个标记跨通道线性地操作。有趣的是，当共同绘制时，这些线形成非线性双曲正切曲线。右侧两个面板：每个通道的输入跨越x轴上的不同范围，为整个双曲正切曲线提供不同的分段。某些频道（例如，红色、绿色和粉色）表现出更极端的x值，其被LN压扁。

动态Tanh（DyT）

受归一化层的形状和缩放的tanh函数之间的相似性的启发，作者提出了动态Tanh（DyT）作为归一化层的替代品。给定输入张量x，DyT层定义如下： $DyT(x)=\gamma \cdot tanh(\alpha x)+\beta$ ，其中，α是一个可学习的标量参数，允许根据输入的范围对输入进行不同的缩放，并考虑到x尺度的变化。这也是为什么作者将整个层命名为“动态”Tanh。γ和β是可学习的参数，与所有归一化层中使用的参数相同——它们允许输出回缩到任何尺度。这有时被认为是一个单独的仿射层;作者认为它们是DyT层的一部分，就像归一化层也包含它们一样。

虽然DyT看起来像或被认为是一个激活函数，但本研究仅使用它来替换归一化层，而不会改变原始架构中激活函数的任何部分，例如GELU或ReLU。网络的其他部分也保持完整。作者还观察到，几乎不需要调整DyT的原始架构所使用的超参数来表现良好。

关于缩放参数：作者总是简单地将γ初始化为一个全1的向量，将β初始化为一个全0的向量。对于缩放器参数α，默认初始化为0.5通常就足够了，LLM训练除外。

备注：DyT不是一种新类型的归一化层，因为它在前向传递期间独立地对来自张量的每个输入元素进行操作，而不计算统计信息或其他类型的聚合。然而，它确实保留了归一化层的效果，即以非线性方式压缩极值，同时几乎线性地变换输入的非常中心的部分。

实验

为了证明DyT的有效性，作者在不同的任务和领域范围内对Transformers和其他一些现代架构进行了实验。在每个实验中，作者将原始架构中的LN或RMSNorm替换为DyT层，并遵循官方开源协议来训练和测试两个版本的模型。

视觉领域的有监督学习。作者在ImageNet-1 K分类任务上训练ViT和ConvNeXt上的“Base”和“Large”尺寸。选择这些模型是因为它们的受欢迎程度和独特的操作：ViT中的注意力和ConvNeXt中的卷积。表1报告了排名前1的分类准确度。DyT在架构和模型大小方面的性能略优于LN。在图5中，作者进一步绘制了ViT-B和ConvNeXt-B的训练损失。这些曲线表明，DyT和基于LN的模型的收敛行为是高度一致的.

视觉自监督学习。作者使用两种流行的视觉自监督学习方法进行基准测试：掩蔽自动编码器（MAE）和DINO。默认情况下，两者都使用ViT作为骨干，但具有不同的训练目标：MAE使用重建损失进行训练，而DINO使用joint-embedding loss。根据标准的自监督学习协议，作者首先在ImageNet-1 K上不使用任何标签对模型进行预训练，然后通过附加分类层并使用标签对其进行微调来测试预训练的模型。微调结果见表2。在自监督学习任务中，DyT的表现始终与LN相当。

扩散模型。作者在ImageNet-1 K上训练了三个尺寸为B、L和XL的扩散Transformer（DiT）模型。patch大小分别为4、4和2。注意，在DiT中，LN层的仿射参数被用于DiT中的类条件化，并且我们在DyT实验中保持这种方式，只是用tanh（αx）函数替换了归一化变换。训练后，作者使用标准ImageNet的“参考批次”评价Fréchet起始距离（FID），如表3所示。DyT实现了与LN相当或更好的FID。

大型语言模型。作者对LLaMA 7B、13B、34B和70B模型进行了预训练，以评估DyT相对于LLaMA中使用的默认归一化层RMSNorm的性能。如表4所示，DyT在所有四种型号规格中的性能与RMSNorm相当。图6所示为损耗曲线，表明所有模型尺寸的趋势相似，整个训练过程中的训练损耗密切一致。

语音中的自监督学习。作者在LibriSpeech数据集上预训练了两个wav2vec2.0 Transformer模型。在表5中报告了最终验证损失。作者观察到DyT在两种模型尺寸中的性能与LN相当。

DNA序列建模。在长距离DNA序列建模任务中，作者预训练HyenaDNA模型和Caduceus模型。预训练使用来自GRCh38的人类参考基因组数据，并在GenomicBenchmarks上进行评估。结果如表6所示。DyT在此任务中的性能与LN相当。

分析

作者对DyT的重要性质进行了一些分析。

DyT的效率：作者使用RMSNorm或DyT对LLaMA 7 B模型进行基准测试，通过使用4096个令牌的单个序列测量100次向前传递（推理）和100次前向后向传递（训练）所花费的总时间。表7报告了在具有BF 16精度的Nvidia H100 GPU上运行时，所有RMS Norm或DyT层以及整个模型所需的时间。与RMS Norm层相比，DyT层显着减少了计算时间，在FP 32精度下观察到类似的趋势。DyT可能是面向效率的网络设计的一个有前途的选择。

tanh和α的消融：为了进一步研究tanh和α在DyT中的作用，作者进行了实验，以评估当这些成分被改变或去除时模型的性能。

作者用替代的压缩函数替换DyT层中的tanh，特别是hardtanh和sigmoid（图7），同时保持可学习的scaler α不变。此外，作者评估了完全删除tanh的影响，将其替换为恒等函数，同时仍然保留α。如表8所示，压缩功能对于稳定训练至关重要。使用恒等函数会导致不稳定的训练和发散，而压缩函数可以实现稳定的训练。在这些压缩函数中，tanh的性能最好。这可能是由于其平滑性和零中心属性。

三个挤压函数的曲线：tanh，hardtanh和sigmoid。所有三个函数都将输入压缩到一个有界范围内，但tanh（x）在DyT层中使用时实现了最佳性能。作者怀疑这是由于它的光滑性和零中心特性。

删除α：接下来，作者评估了删除可学习α而保留挤压函数（tanh、hardtanh和sigmoid）的影响。如表9所示，删除α会导致所有挤压函数的性能下降，从而突出了α在整体模型性能中的关键作用。

α的值：在训练期间，α在训练过程中先减小后增大，但总是与输入激活的标准差一致地波动。这支持了α在将激活维持在适当范围内的重要作用，从而导致稳定和有效的训练。在训练完成之后，作者进一步分析了训练网络中α的最终值，发现α与输入激活的1/std有很强的相关性。如图8的右图所示，1/std值越高，通常对应的α值越大，反之亦然。此外，作者还观察到，更深层的活化倾向于具有更大的标准偏差。这一趋势与深层残差网络的特征一致。

两种分析都表明，α通过学习接近输入激活的1/std的值而部分地作为归一化机制发挥作用。与LN不同，LN对每个令牌的激活进行归一化，而α对整个输入激活进行集体归一化。因此，α不能单独以非线性方式抑制极值。

与其他方法相比：为了进一步评估DyT的有效性，作者将其与其他方法进行了比较。

α的初始化

作者发现，调整α的初始化（表示为 $\alpha _0$ ）很少导致显著的性能改善。唯一的例外是LLM训练，在LLM训练时仔细调整 $\alpha _0$ 可产生显著的性能增益。

非LLM模型的α分布

作者观察到，性能在较宽的 $\alpha _0$ 值范围内保持稳定，在0.5和1.2之间的值通常会产生良好的结果。作者还观察到，调整 $\alpha _0$ 通常只影响训练曲线的早期阶段。主要的例外是ViT-L实验，当 $\alpha _0$ 超过0.6时，训练变得不稳定并发散。在这样的情况下，会降低学习速率。

较小的 $\alpha _0$ 导致更稳定的训练。
设置 $\alpha _0$ = 0.5作为默认值最佳。

LLM的α值

调整 $\alpha _0$ 可增强LLM性能。如前所述，默认设置 $\alpha _0$ = 0.5通常在大多数任务中表现良好。然而，作者发现调整 $\alpha _0$ 可以大大提高LLM的性能。作者通过在30B token上预训练每个模型并比较它们的训练损失来调整LLaMA模型的 $\alpha _0$ 。表11总结了每个模型的调整 $\alpha _0$ 值。出现了两个关键发现：

较大的模型需要较小的 $\alpha _0$ 值。一旦为较小的模型确定了最佳 $\alpha _0$ ，就可以相应地减少较大模型的搜索空间。
注意力块的 $\alpha _0$ 值越高，性能越好。作者发现，对于注意力块中的DyT层，初始化α具有较高的值，而对于其他位置中的DyT层，初始化α具有较低的值（即，在FFN块内或在最终线性投影之前）提高了性能。

为了进一步说明 $\alpha _0$ 调整的影响，图11显示了两个LLaMA模型的损失值热图。这两种模型都受益于注意力块中较高的 $\alpha _0$ ，从而减少了训练损失。

3.代码详解

环境安装

默认为Python3.12，pytorch环境为2.5.1

conda install pytorch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 pytorch-cuda=12.4 -c pytorch -c nvidia
pip install timm==1.0.15 tensorboard

dynamic_tanh.py

dynamic_tanh.py定义了DyT的基本形式，

import torch
import torch.nn as nn
from timm.layers import LayerNorm2d


class DynamicTanh(nn.Module):
    def __init__(self, normalized_shape, channels_last, alpha_init_value=0.5):
        # normalized_shape：定义了 weight 和 bias 的形状，类似于 nn.LayerNorm 的归一化维度。
        # channels_last：布尔值，指示数据的通道维度是否在最后。如果为 True，则假设输入张量的形状为 (N, H, W, C)；如果为 False，则假设形状为 (N, C, H, W)。
        # alpha_init_value：初始化参数 alpha 的值，默认为 0.5。
        super().__init__()
        self.normalized_shape = normalized_shape
        self.alpha_init_value = alpha_init_value
        self.channels_last = channels_last

        self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value)    # 一个标量参数，用于控制 torch.tanh 的缩放。
        self.weight = nn.Parameter(torch.ones(normalized_shape))       # 一个张量，形状为 normalized_shape，用于对输入张量进行逐元素乘法操作。
        self.bias = nn.Parameter(torch.zeros(normalized_shape))        # 一个张量，形状为 normalized_shape，用于对输入张量进行逐元素加法操作。

    def forward(self, x):
        x = torch.tanh(self.alpha * x)    
        if self.channels_last:
            x = x * self.weight + self.bias
        else:
            x = x * self.weight[:, None, None] + self.bias[:, None, None]    # 通过广播机制（[:, None, None]）将 weight 和 bias 应用于每个通道。
        return x

    def extra_repr(self):    # 在打印模块时显示额外信息
        return f"normalized_shape={self.normalized_shape}, alpha_init_value={self.alpha_init_value}, channels_last={self.channels_last}"


def convert_ln_to_dyt(module):    # 将模型中的 nn.LayerNorm 模块递归地替换为 DynamicTanh 模块。
    module_output = module
    if isinstance(module, nn.LayerNorm):
        module_output = DynamicTanh(module.normalized_shape, not isinstance(module, LayerNorm2d))
    for name, child in module.named_children():
        module_output.add_module(name, convert_ln_to_dyt(child))     # 对于模块的子模块，递归调用 convert_ln_to_dyt。
    del module
    return module_output

训练

ViT

要在 ImageNet-1K 上使用 ViT重现结果，请运行以下命令：(要获得 LN 的结果，请将 --dynamic_tanh 设为 false）。

torchrun --nnodes=8 --nproc_per_node=8 main.py \
    --model vit_base_patch16_224 \
    --drop_path 0.1 \
    --batch_size 64 \
    --lr 4e-3 \
    --update_freq 1 \
    --model_ema true \
    --model_ema_eval true \
    --data_path /path/to/imagenet \
    --output_dir /path/to/saving_dir \
    --dynamic_tanh true

对于ViT-L、ConvNext的结果，请参考项目主页

LLaMA

除了ViT，作者团队还提供了其他常用模型替换DyT后的结果：

我们以LLaMA为例：

1. 按照原始的 fms-fsdp 文档为项目设置所需的 Python 环境。除此之外，还需要一个额外的 Python 库来保存数据加载器检查点，以防需要在不重复训练数据的情况下恢复训练：

pip install torchdata

2. 由于原始 fms-fsdp 数据库没有提供标准数据集，因此作者使用了 Pile 数据集。为了简化过程，作者提供了一个脚本，用于从原始数据集生成标记化的 Arrow 文件。首先，确定要使用的world size，因为这将决定生成的 Arrow 文件总数。一般来说，创建的 Arrow 文件数量应与计划用于模型训练的 GPU 最大总数相同。确定世界大小后，运行以下命令：

python prepare_data.py \
  --rank $RANK \
  --world_size $WORLD_SIZE \
  --data_path /path/to/data \
  --output_path /path/to/output_dir \
  --max_num_tokens $MAX_NUM_TOKENS \
  --tokenizer $TOKENIZER

重要：您需要运行该命令 $WORLD_SIZE 次，每次使用不同的 $RANK 值，范围从 0 到 $WORLD_SIZE - 1。
对于只想tokenize一个子集的大型数据集，可适当设置 $MAX_NUM_TOKENS。tokenize数据中的标记总数将是 $MAX_NUM_TOKENS × $WORLD_SIZE。对于tokenizer，作者使用 “meta-llama/Llama-2-7b-chat-hf”。

3. 接着使用如下命令运行：

MODEL_ARGS="\
--model_variant=llama2_7b \
--ckpt_load_path=/checkpoint/path \
--ckpt_save_path=/checkpoint/path \
--data_path=/dataset/path \
--sharding_strategy=hsdp \
--fsdp_activation_checkpointing=False \
--selective_checkpointing=1 \
--mixed_precision=True \
--low_cpu_fsdp=True \
--batch_size=2 \
--learning_rate=0.0003 \
--checkpoint_interval=5000 \
--tracker=wandb \
--tracker_dir=/tracker/path \
--tracker_project_name=tracker_project_name \
--tracker_run_name=llama2_dyt_7b \
--attn_alpha_init_value=0.8 \
--ffn_alpha_init_value=0.2 \
--dec_alpha_init_value=0.2 
"
srun torchrun --nnodes=64 --nproc_per_node=8 main_training_llama.py ${MODEL_ARGS}

更多内容请参考代码项目主页：https://github.com/jiachenzhu/DyT

4.总结

DyT的提出具有重要的贡献和意义。首先，它提供了一种简单而有效的归一化层替代方案。传统的归一化层如Layer Normalization虽然在稳定训练和提升模型性能方面发挥了重要作用，但计算过程相对复杂，需要计算输入数据的均值和方差。DyT通过简单的α缩放和tanh函数操作，避免了这些复杂的计算，从而显著提高了训练和推理的效率。

其次，DyT在多种任务和模型架构中表现出色。实验结果表明，使用DyT的模型在视觉、语言、语音等多个领域的任务中，都能达到甚至超过传统归一化层的性能，且几乎不需要调整超参数。这表明DyT具有很强的普适性，可以广泛应用于不同的神经网络架构中。

此外，DyT的提出挑战了“归一化层对训练现代神经网络必不可少”这一传统观念。它为神经网络的设计和优化提供了新的思路，有助于推动深度学习领域的进一步发展。通过这种简单的替代方案，研究人员可以更深入地理解归一化层在神经网络中的作用，并探索更多可能的优化方法。

如果你觉得这篇文章对你有启发，或者解决了你的困惑，那就请动动手指，给我点个赞吧！你的点赞是我持续创作的动力，也是对我努力的最好鼓励。

同时，如果你不想错过更多干货内容，不妨关注我的博客，这样就能第一时间获取最新更新啦！如果这篇文章你特别喜欢，也可以收藏起来，方便以后随时回顾。

你的每一个支持，我都看在眼里，记在心里。感谢有你，让我们一起在知识的海洋里探索更多！