DyT 颠覆归一化 | 无需归一化层的Transformer依然能够实现跨视觉语言任务性能SOTA

大模型向量数据库机器学习

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群

picture.image

picture.image

picture.image

picture.image

picture.image

导读

归一化层在现代神经网络中无处不在,并且长期以来一直被视为必不可少的。这项工作表明,无需归一化的Transformer可以通过一种非常简单的技术实现相同或更好的性能。作者引入了动态双曲正切函数(Dynamic Tanh,简称DyT),它是一个逐元素操作

,作为Transformer中归一化层的替代方案。DyT的灵感来源于观察到的Transformer中的层归一化通常会产生类似tanh的S型输入-输出映射。通过引入DyT,无需归一化的Transformer可以匹配或超越其归一化版本的性能,大多数情况下无需调整超参数。

作者验证了具有DyT的Transformer在多种设置下的有效性,从识别到生成,从监督学习到自监督学习,从计算机视觉到语言模型。这些发现挑战了传统观念,即归一化层在现代神经网络中是不可或缺的,并为它们在深度网络中的作用提供了新的见解。

  1. 引言

过去十年中,归一化层已经巩固了其在现代神经网络中最基本组成部分之一的地位。这一切都始于2015年批量归一化的发明,它使得视觉识别模型能够实现大幅度的加速和更好的收敛,并在随后的几年中迅速流行起来。自那时起,针对不同的网络架构或领域已经提出了许多归一化层的变体。如今,几乎所有现代网络都使用了归一化层,其中层归一化是最受欢迎的之一,尤其是在占主导地位的Transformer架构中。

规范化层的广泛应用主要是由其在优化方面的经验优势驱动的。除了实现更好的结果外,它们还有助于加速和稳定收敛。随着神经网络变得越来越宽和深,这种必要性变得更加关键。因此,规范化层被广泛认为对于深度网络的有效训练是至关重要的,如果不是必不可少的。这一信念微妙地体现在近年来,新型架构通常试图取代注意力或卷积层,但几乎总是保留规范化层。

本文通过引入Transformer中归一化层的简单替代方案来挑战这一信念。作者的探索从观察LN层将输入映射到具有tanh-like、S形曲线的输出开始,该曲线在缩放输入激活的同时压缩极端值。受这一洞察的启发,作者提出了一种称为动态tanh(DyT)的逐元素操作,定义为:

,其中

是一个可学习的参数。此操作旨在通过学习适当的缩放因子

并通过有界tanh函数压缩极端值来模拟LN的行为。值得注意的是,与归一化层不同,它无需计算激活统计信息即可实现这两种效果。

采用DyT非常简单,如图1所示:作者直接将现有的归一化层替换为DyT,应用于视觉和语言Transformer等架构中。作者通过实证证明,具有DyT的模型可以在各种设置下稳定训练并实现高最终性能。通常不需要调整原始架构上的训练超参数。作者的工作挑战了归一化层对训练现代神经网络不可或缺的观念,并提供了关于归一化层特性的实证见解。此外,初步测量表明,DyT可以提高训练和推理速度,使其成为面向效率的网络设计的候选方案。

2 背景:归一化层

作者首先回顾归一化层。大多数归一化层具有一个共同的公式。给定一个形状为

的输入

,其中

是批量大小,

是 Token 数量,

是每个 Token 的嵌入维度,输出通常计算如下:

其中

是一个很小的常数,

是形状为

的可学习向量参数。它们是“缩放”和“平移”仿射参数,允许输出处于任何范围内。

分别表示输入的均值和方差。不同的方法主要区别在于这两项统计量的计算方式。这导致

具有不同的维度,每个维度在计算过程中都应用了广播操作。

批归一化(BN)是第一个现代归一化层,它主要被用于卷积神经网络(ConvNet)模型中。其引入代表了深度学习架构设计的一个重大里程碑。BN计算批量和 Token 维度的均值和方差,具体如下:

在卷积神经网络中流行的其他归一化层,如组归一化和实例归一化,最初是为特定任务如目标检测和图像风格化而提出的。它们具有相同的整体公式,但在计算统计量的轴和范围内有所不同。

层归一化(LN)和均方根归一化(RMSNorm)是Transformer架构中使用的两种主要归一化层。LN为每个样本中的每个 Token 独立计算这些统计量,其中

。RMSNorm通过去除均值中心化步骤并使用

对输入进行归一化来简化LN。如今,大多数现代神经网络都使用LN,这得益于其简洁性和通用性。最近,RMSNorm在T5、LLaMA、Mistral、Qwen、InternLM和DeepSeek等语言模型中获得了流行。在本工作中,作者考察的Transformers都使用LN,除了LLaMA使用RMSNorm。

3 正则化层的作用是什么?

分析设置。作者首先对训练网络中归一化层的行为进行实证研究。为此分析,作者选取了在ImageNet-1K上训练的Vision Transformer模型(ViT-B),在LibriSpeech上训练的wav2vec 2.0大型Transformer模型,以及训练在ImageNet-1K上的Diffusion Transformer(DiT-XL)。在所有情况下,LN均应用于每个Transformer块以及最终的线性投影之前。

对于所有三个训练好的网络,作者采样一个迷你批次样本并通过网络进行前向传播。然后,作者测量归一化层的输入和输出,即归一化操作之前和之后的张量。由于LN保留了输入张量的维度,作者可以建立输入和输出张量元素之间的一一对应关系,从而直接可视化它们之间的关系。作者在图2中绘制了这些映射结果。

picture.image

具有层归一化的Tanh-like映射。对于所有三个模型,在早期的LN层(图2的第一列)中,作者发现这种输入输出关系主要是线性的,类似于c-y图中的直线。然而,更深层的LN层是作者进行更多有趣观察的地方。

从这些深层中一个显著的观察结果是,这些曲线的形状高度类似于由tanh函数表示的完整或部分S形曲线(见图3)。人们可能会预期LN层会以线性方式转换输入张量,因为减去均值并除以标准差是线性操作。LN以每个 Token 的方式进行归一化,仅线性转换每个 Token 的激活。由于 Token 具有不同的均值和标准差值,因此线性性在输入张量的所有激活上并不成立。尽管如此,实际的非线性转换与缩放后的tanh函数高度相似,这仍然让作者感到惊讶。

对于这种“S”形曲线,作者注意到其中心部分,即

值接近零的点,仍然主要呈线性形状。大多数点(约99%)都落在这个线性范围内。然而,仍有许多点明显超出了这个范围,这些点被认为是具有“极端”值,例如在ViT模型中

值大于50或小于-50的点。归一化层对这些值的主要影响是将它们压缩成更不极端的值,更接近大多数点的分布。这就是归一化层不能仅通过简单的仿射 Transformer Layer 来近似的原因。作者假设这种对极端值的非线性且不成比例的压缩效应是归一化层重要且不可或缺的原因。

近期Ni等人(2024年)的研究结果也突出了LN层引入的强烈非线性,展示了非线性如何增强模型的表示能力。此外,这种挤压行为与生物神经元对大输入的饱和特性相呼应,这一现象首次在大约一个世纪前被观察到。

通过token和channel进行归一化。一个LN层是如何对每个token执行线性变换,同时以这种非线性的方式压缩极端值的?为了理解这一点,作者分别按token和channel对点进行可视化。这通过图4中的第二和第三个子图(来自图2的ViT)进行展示,但使用了采样子集的点以增加清晰度。当作者选择要绘制的channel时,确保包括具有极端值的channel。

picture.image

在图4的左侧两个面板中,作者使用相同的颜色可视化每个 Token 的激活情况。作者观察到来自任何单个 Token 的所有点都形成了一条直线。然而,由于每个 Token 的方差不同,斜率也不同。输入范围较小的 Token 往往具有较小的方差,归一化层将使用较小的标准差来划分它们的激活,因此产生更大的斜率,形成一条直线。总体来看,它们形成了一个类似tanh函数的S形曲线。在右侧的两个面板中,作者使用相同的颜色为每个通道的激活着色。作者发现不同的通道往往具有截然不同的输入范围,只有少数几个通道(例如红色、绿色和粉色)表现出较大的极端值。这些是归一化层压缩得最多的通道。

4 动态Tanh(DyT)

受归一化层形状与缩放tanh函数相似性的启发,作者提出了动态tanh(DyT)作为归一化层的替代方案。给定一个输入张量

,一个DyT层的定义如下:

其中

是一个可学习的标量参数,允许根据输入的范围不同进行不同的缩放,以适应不同的

尺度(图2)。这也是作者将其整个操作命名为“动态”Tanh的原因。

是每个通道的可学习向量参数,与所有归一化层中使用的参数相同——它们允许输出缩放回任何尺度。这有时被视为一个独立的仿射层;就作者的目的而言,作者认为它们是DyT层的一部分,就像归一化层也包括它们一样。有关DyT在Pytorch-like伪代码中的实现,请参阅算法1。

picture.image

将DyT层集成到现有架构中非常简单:一个DyT层替换一个归一化层(见图1)。这适用于注意力块、FFN块以及最终的归一化层。尽管DyT看起来或被认为是一个激活函数,但本研究仅将其用于替换归一化层,而不改变原始架构中激活函数的任何部分,例如GELU或ReLU。网络的其它部分也保持不变。作者还观察到,对于DyT表现良好,几乎没有必要调整原始架构使用的超参数。

关于缩放参数。作者通常在归一化层之后,将

初始化为全1向量,将

初始化为全0向量。对于缩放参数

,默认初始化为0.5通常足够,除了在大语言模型(LLM)训练的情况下。关于

初始化的详细分析见第7节。除非特别说明,否则在后续实验中,

的初始化值均为0.5。

备注。DyT并非一种新的归一化层,因为它在正向传播过程中独立地对张量中的每个输入元素进行操作,而不计算统计或其他类型的聚合。然而,它确实保留了归一化层的效果,以非线性方式压缩极端值,同时几乎线性地变换输入的非常中心部分。

5 实验

为了展示DyT的有效性,作者在多种任务和领域中对Transformers以及其他一些现代架构进行了实验。在每个实验中,作者将原始架构中的LN或RMSNorm替换为DyT层,并遵循官方开源协议来训练和测试两种模型版本。关于重现作者结果的详细说明见附录A。值得注意的是,为了突出DyT的简单性,作者使用了与标准化版本相同的超参数。为了完整性,附录B提供了关于学习率调整和

初始值设置的额外实验结果。

监督学习在视觉领域。作者在ImageNet-1K分类任务上训练了“Base”和“Large”尺寸的Vision Transformer(ViT)和ConvNeXt。这些模型因其流行度和独特的操作而被选中:ViT中的注意力机制和ConvNeXt中的卷积操作。表1报告了Top-1分类准确率。DyT在两种架构和模型尺寸上均略优于LN。作者进一步在图5中绘制了ViT-B和ConvNeXt-B的训练损失。曲线显示,基于DyT和LN的模型的收敛行为高度一致。

picture.image

视觉中的自监督学习。作者与两种流行的视觉自监督学习方法进行了基准测试: Mask 自编码器(MAE)和DINO。两者默认使用视觉Transformer作为 Backbone 网络,但训练目标不同:MAE使用重建损失进行训练,而DINO使用联合嵌入损失。遵循标准的自监督学习协议,作者首先在ImageNet-1K上预训练模型,不使用任何标签,然后通过添加分类层并使用标签进行微调来测试预训练模型。微调结果如表2所示。DyT在自监督学习任务中始终与LN表现相当。

扩散模型。作者在ImageNet-1K上训练了三个不同尺寸的扩散Transformer(DiT)模型,分别为B、L和XL。各自的patch大小分别为4、4和2。请注意,在DiT中,LN层的仿射参数用于DiT中的类别条件,在作者的DyT实验中,作者保持这种方式,仅将归一化变换替换为tanh(ox)函数。训练完成后,作者使用标准的ImageNet“参考批次”评估Fréchet Inception Distance(FID)分数,结果如表3所示。DyT在FID上实现了与LN相当或更好的性能。

picture.image

大语言模型。作者预先训练了LLaMA 7B、13B、34B和70B模型,以评估DyT相对于RMSNorm在LLaMA中使用的默认归一化层的性能。这些模型在包含200B个 Token 的The Pile数据集上按照LLaMA中概述的原版配方进行训练。

picture.image

在具有DyT的LLaMA上,作者在初始嵌入层之后添加了一个可学习的标量参数,并调整了

的初始值,具体细节见第7节。作者报告了训练后的损失值,并遵循OpenLLaMA对模型在lm-eval的15个零样本任务上进行基准测试。如表4所示,DyT在所有四种模型规模上与RMSNorm的表现相当。图6展示了损失曲线,显示了所有模型规模之间存在相似的趋势,训练损失在整个训练过程中紧密对齐。

picture.image

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

picture.image

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

picture.image

6 分析

作者对DyT的重要性质进行了多项分析。首先,作者评估了它们的计算效率,随后进行了两项研究,分别探讨了tanh函数和可学习的尺度α的作用。最后,作者展示了与旨在去除归一化层的先前方法的比较。

6.1 动态阈值(DyT)效率

作者将LLaMA 7B模型与RMSNorm或DyT进行基准测试,通过测量使用单个4096个token的序列进行100次前向传递(推理)和100次前向-后向传递(训练)所需的总时间。表7报告了在Nvidia H100 GPU上以BF16精度运行时,所有RMSNorm或DyT层以及整个模型所需的时间。与RMSNorm层相比,DyT层显著减少了计算时间,在FP32精度下也观察到类似趋势。DyT可能是一个针对效率导向的网络设计的有希望的选项。

picture.image

6.2 tanh和

的消融

为了进一步研究tanh和

在DyT中的作用,作者进行了实验以评估当这些组件被更改或移除时模型的性能。

替换和移除tanh。作者将DyT层中的tanh替换为替代的压缩函数,具体为hardtanh和sigmoid(图7),同时保留可学习的缩放因子

。此外,作者还评估了完全移除tanh的影响,即用恒等函数替换它,同时仍保留

。如表8所示,压缩函数对于稳定的训练至关重要。使用恒等函数会导致训练不稳定和发散,而压缩函数则能实现稳定的训练。在压缩函数中,tanh表现最佳。这可能是由于其平滑性和零中心特性。

picture.image

移除

。接下来,作者评估在保留挤压函数(tanh、hardtanh 和 sigmoid)的同时移除可学习的

对模型的影响。如表9 所示,移除

导致所有挤压函数的性能均有所下降,突显了

在整体模型性能中的关键作用。

picture.image

图8 左侧:对于ViT-B模型中选定的两个DyT层,作者追踪每个epoch结束时激活的

和标准差的倒数

,观察到它们在训练过程中共同演变。右侧:作者将两个训练好的模型的最终

值(ViT-B和ConvNeXt-B)与输入激活的1/std值进行对比,展示了这两个值之间强烈的关联性。

6.3

的取值

在训练过程中,作者的分析表明,

与激活值的1/标准差紧密相关。如图8左面板所示,

在训练过程中先下降后上升,但始终与输入激活值的标准差保持一致性的波动。这支持了

在维持激活值在合适范围内的重要作用,这有助于实现稳定有效的训练。

在训练后,作者对训练网络中

的最终值进行了进一步分析,发现其与输入激活的1/标准差之间存在强烈的关联。如图8的右侧面板所示,较高的1/标准差值通常对应较大的

值,反之亦然。此外,作者观察到深层通常具有标准差较大的激活。这一趋势与深度残差网络的特征相符,如ConvNets中所示,以及在Transformer中的展示。

两种分析均表明,α函数部分充当了一种归一化机制,通过学习近似于输入激活值标准差的值。与LN不同,LN对每个 Token 的激活值进行归一化,α对整个输入激活值进行集体归一化。因此,α本身无法以非线性方式抑制极端值。

6.4 与其他方法的比较

为了进一步评估DyT的有效性,作者将其与其他一些方法进行了比较,这些方法同样允许在无需归一化层的情况下训练Transformer。这些方法可以大致分为基于初始化和基于权重归一化的方法。作者考虑了两种流行的基于初始化的方法,即Fixup和SkipInit。这两种方法的目标都是通过调整初始参数值来减轻训练过程中的不稳定性,防止在训练开始时出现大的梯度和激活,从而在没有归一化层的情况下实现稳定学习。相比之下,基于权重归一化的方法在整个训练过程中对网络权重施加约束,以在没有归一化层的情况下维持稳定的训练动态。作者包括了一种这样的方法,即

Reparam,它通过控制权重的谱范数来促进稳定学习。

7

的初始化

作者发现,调整

(记作

)的初始化很少能带来显著的性能提升。唯一的例外是LLM训练,其中对

的仔细调整可以带来明显的性能提升。在本节中,作者将详细阐述作者对

初始化影响的发现。

7.1 非LLM模型的

初始化

非LLM模型对

相对不敏感。图9展示了在不同任务中,改变

对验证性能的影响。所有实验都遵循各自食谱的原有设置和超参数。作者观察到,在广泛的

值范围内,性能保持稳定,通常在0.5到1.2之间的值能产生良好的结果。作者观察到,调整

通常只影响训练曲线的早期阶段。主要的例外是监督ViT-L实验,当

超过0.6时,训练变得不稳定并发散。在这种情况下,降低学习率可以恢复稳定性,具体细节如下。

picture.image

较小的

导致训练更加稳定。基于之前的观察,作者进一步分析了导致训练不稳定性的因素。作者的发现表明,增加模型大小或学习率都需要降低

以确保稳定训练。相反,较高的

需要较低的学习率以减轻训练不稳定性。图10显示了使用ImageNet-1K数据集的监督ViT训练稳定性的消融情况。作者改变了学习率、模型大小和

值。训练更大的模型更容易失败,需要更小的

值或学习率才能实现稳定训练。在类似条件下,基于LN的模型也观察到了类似的不稳定性模式,将

设置为稳定性模式与LN相似。

picture.image

图10 不同

值、学习率和模型大小下的稳定性。作者在 ImageNet-1K 数据集上训练监督式 ViT 模型,并观察到较大模型对于 LN 和 DyT 模型都更容易出现不稳定性。降低学习率或减少

可以提高稳定性。当

时,LN 显示出与 DyT 相似的稳定性。将

设置为默认值。基于作者的发现,作者将

设置为所有非 LLM 模型的默认值。此设置提供了与 LN 相当的训练稳定性,同时保持了强大的性能。

7.2 长语言模型(LLMs)的

初始化

调整

提升了大语言模型(LLM)的性能。如前所述,

的默认设置在大多数任务中通常表现良好。然而,作者发现调整

可以显著提高 LLM 的性能。作者通过在每个 LLaMA 模型上预训练 30B 个 Token 并比较它们的训练损失来调整

。表11 总结了每个模型的调整后的

值。出现了两个关键发现:

    1. 较大的模型需要更小的

值。一旦确定了较小模型的最佳

,可以相应地减少较大模型的搜索空间。 2. 2. 注意块中更高的

值可以提升性能。作者发现,将

初始化为更高的值用于注意力块中的 DyT 层,以及更低的值用于其他位置的 DyT 层(例如,在 FFN 块内或最终的线性投影之前)可以提升性能。

为了进一步说明

调节的 影响,图1l 展示了两个 LLaMA 模型的损失值 Heatmap 。这两个模型在注意力块中受益于更高的

,导致训练损失降低。

模型宽度主要决定了

的选择。作者还研究了模型宽度和深度对最优

的影响。作者发现,模型宽度在确定最优

方面至关重要,而模型深度的影响最小。表12展示了不同宽度和深度下的最优

值,表明较宽的网络从较小的

值中获益以实现最佳性能。另一方面,模型深度对

的选择几乎没有影响。

picture.image

表12 LLaMA训练中不同模型宽度和深度下的最优

(注意力/其他)。模型宽度对

的选择有显著影响,宽度更大的网络需要更小的值。相比之下,模型深度的影响可以忽略不计。如表12所示,网络越宽,对“注意力”和“其他”的初始化需要越不均匀。作者假设LLM的

初始化的敏感性与其与其他模型相比过大的宽度有关。

8 局限性

作者对使用LN或RMSNorm的网络进行实验,因为它们在Transformers和其他现代架构中非常流行。初步实验(见附录C)表明,DyT在经典网络如ResNets中直接替换BN存在困难。是否以及如何DyT能够适应具有其他类型归一化层的模型,还有待深入研究。

9 结论

在本工作中,作者证明了现代神经网络,特别是Transformer,可以在不使用归一化层的情况下进行训练。这是通过动态Tanh(DyT)实现的,它是传统归一化层的一个简单替代方案。它通过一个可学习的缩放因子

调整输入激活范围,然后通过S形tanh函数压缩极端值。尽管函数更简单,但它有效地捕捉了归一化层的行为。在多种设置下,使用DyT的模型与归一化模型的性能相当甚至更优。这些发现挑战了关于在训练现代神经网络中归一化层必要性的传统理解。作者的研究也有助于理解归一化层的机制,这是深度神经网络中最基本的构建模块之一。

参考

[1]. Transformers without Normalization

picture.image

扫码加入👉「集智书童」交流群

(备注:方向+学校/公司+昵称)

picture.image

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论