本文提出了一种名为**动态双曲正切(Dynamic Tanh, DyT)**的简单技术,用以替代Transformer中的归一化层(如层归一化LN和均方根归一化RMSNorm)。DyT通过引入可学习的缩放参数α和双曲正切函数tanh(αx),在无需计算激活统计量的情况下,实现了对输入激活值的缩放和极端值的非线性压缩。实验表明,DyT在多种任务(包括监督学习、自监督学习、扩散模型和大型语言模型)中均能匹配或超越传统归一化方法的性能,且无需调整超参数。此外,DyT在训练和推理速度上展现出显著优势,为高效网络设计提供了新思路。该研究挑战了归一化层在现代神经网络中不可或缺的传统认知,为理解归一化层的作用机制提供了新的视角。这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:
官方项目地址在这里,如下所示:
摘要
归一化层是现代神经网络中最基本的组件之一,长期以来被认为是必不可少的。本研究证明,无需归一化的Transformer可以通过一种非常简单的方法达到相同甚至更好的性能。我们引入了动态双曲正切(Dynamic Tanh,简称DyT),这是一种逐元素操作DyT(x)=tanh(αx),作为Transformer中归一化层的替代品。DyT的灵感来源于观察到Transformer中的层归一化(Layer Norm)通常会产生类似双曲正切(tanh)的S形输入-输出映射。通过引入DyT,无需归一化的Transformer能够在大多数情况下匹配甚至超过其归一化对应模型的性能,且通常无需超参数调整。我们在多种设置中验证了带有DyT的Transformer的有效性,涵盖从识别到生成、从监督学习到自监督学习,以及从计算机视觉到语言模型的广泛任务。这些发现挑战了归一化层在现代神经网络中不可或缺的传统认知,并为深度网络中归一化层的作用提供了新的见解。此外,初步测量表明,DyT可以提高训练和推理速度,使其成为面向效率的网络设计的候选方案。
图1 左侧:原始Transformer模块。右侧:带有我们提出的动态双曲正切(DyT)层的模块。DyT是常用层归一化(Layer Norm,Ba等人,2016)(在某些情况下为RMSNorm,Zhang和Sennrich,2019)层的直接替代品。使用DyT的Transformer在性能上能够匹配甚至超过其归一化对应模型。
1. 引言
在过去十年中,归一化层已成为现代神经网络中最基本的组件之一。这一趋势可以追溯到2015年批量归一化(Batch Normalization)的发明(Ioffe和Szegedy, 2015),它在视觉识别模型中实现了更快、更好的收敛,并在随后几年迅速普及。此后,针对不同网络架构或领域提出了许多归一化层的变体(Ba等人, 2016; Ulyanov等人, 2016; Wu和He, 2018; Zhang和Sennrich, 2019)。如今,几乎所有现代网络都使用归一化层,其中层归一化(Layer Norm,简称LN)(Ba等人, 2016)是最受欢迎的之一,尤其是在主导的Transformer架构中(Vaswani等人, 2017; Dosovitskiy等人, 2020)。
归一化层的广泛采用主要是由于它们在优化方面的经验性优势(Santurkar等人, 2018; Bjorck等人, 2018)。除了能够获得更好的结果外,它们还有助于加速和稳定收敛。随着神经网络变得越来越宽和深,这种必要性变得愈发关键(Brock等人, 2021a; Huang等人, 2023)。因此,归一化层被认为对于深度网络的有效训练至关重要,如果不是不可或缺的话。这种信念在近年来新架构的开发中得到了微妙的体现:尽管这些架构常常试图替换注意力或卷积层(Tolstikhin等人, 2021; Gu和Dao, 2023; Sun等人, 2024; Feng等人, 2024),但几乎总是保留归一化层。
本文通过引入一种简单的归一化层替代方案来挑战这一信念。我们的探索始于观察到LN层将输入映射到输出时呈现出类似双曲正切(tanh)的S形曲线,通过缩放输入激活值并压缩极端值来实现。受此启发,我们提出了一个逐元素操作,称为动态双曲正切(Dynamic Tanh,简称DyT),定义为:DyT(x)=tanh(αx),其中α是一个可学习的参数。这一操作旨在通过学习适当的缩放因子α并通过有界的tanh函数压缩极端值来模拟LN的行为。值得注意的是,与归一化层不同,它在无需计算激活统计量的情况下实现了这两种效果。
使用DyT非常简单,如图1所示:我们直接在诸如视觉和语言Transformer的架构中替换现有的归一化层。我们通过实验表明,带有DyT的模型可以在广泛的设置中稳定训练并达到高性能。它通常不需要对原始架构的训练超参数进行调整。我们的工作挑战了归一化层对于训练现代神经网络不可或缺的观念,并为归一化层的性质提供了经验性见解。此外,初步测量表明,DyT可以提高训练和推理速度,使其成为面向效率的网络设计的候选方案。
2. 背景:归一化层
我们首先回顾归一化层的相关内容。大多数归一化层共享一个通用公式。给定一个输入x,其形状为(B, T, C),其中B是批量大小,T是token数量,C是每个token的嵌入维度,输出通常计算为:
其中ϵ是一个小常数,γ和β是形状为(C,)的可学习向量参数。它们是“缩放”和“平移”仿射参数,允许输出处于任意范围。µ和σ²分别表示输入的均值和方差。不同方法主要区别在于这两个统计量的计算方式。这导致µ和σ²具有不同的维度,并在计算过程中应用广播。
批量归一化(BN)(Ioffe和Szegedy, 2015)是第一个现代归一化层,主要用于ConvNet模型(Szegedy等人, 2016; He等人, 2016; Xie等人, 2017)。它的引入是深度学习架构设计的一个重要里程碑。BN在批量和token维度上计算均值和方差,具体为:µₖ = (1/BT)∑ᵢⱼxᵢⱼₖ,σ²ₖ = (1/BT)∑ᵢⱼ(xᵢⱼₖ - µₖ)²。其他在ConvNets中流行的归一化层,如组归一化(Wu和He, 2018)和实例归一化(Ulyanov等人, 2016),最初是为特定任务(如目标检测和图像风格化)提出的。它们共享相同的总体公式,但在统计量计算的轴和范围上有所不同。
图2 展示了在Vision Transformer(ViT)(Dosovitskiy等人,2020)、wav2vec 2.0(一种用于语音的Transformer模型)(Baevski等人,2020)以及Diffusion Transformer(DiT)(Peebles和Xie,2023)中选定的层归一化(LN)层的输出与输入的关系。我们从每个模型中采样了一个小批量样本,并绘制了其中四个LN层的输入/输出值。这些输出值是在LN的仿射变换之前的。这些S形曲线与双曲正切(tanh)函数的形状高度相似(见图3)。早期层中更接近线性的形状也可以通过双曲正切曲线的中心部分来捕捉。这促使我们提出动态双曲正切(DyT)作为一种替代方案,并引入一个可学习的缩放因子α,以适应x轴上的不同尺度。
层归一化(LN)(Ba等人, 2016)和均方根归一化(RMSNorm)(Zhang和Sennrich, 2019)是Transformer架构中使用的两种主要归一化层。LN独立于每个样本中的每个token计算这些统计量,其中µᵢⱼ = (1/C)∑ₖxᵢⱼₖ,σ²ᵢⱼ = (1/C)∑ₖ(xᵢⱼₖ - µᵢⱼ)²。RMSNorm(Zhang和Sennrich, 2019)通过移除均值中心化步骤简化了LN,并使用µᵢⱼ = 0和σ²ᵢⱼ = (1/C)∑ₖx²ᵢⱼₖ对输入进行归一化。如今,由于其简单性和通用性,大多数现代神经网络使用LN。最近,RMSNorm在语言模型(如T5(Raffel等人, 2020)、LLaMA(Touvron等人, 2023a,b; Dubey等人, 2024)、Mistral(Jiang等人, 2023)、Qwen(Bai等人, 2023; Yang等人, 2024)、InternLM(Zhang等人, 2024; Cai等人, 2024)和DeepSeek(Liu等人, 2024; Guo等人, 2025))中变得流行。我们在这项工作中检查的Transformer都使用LN,除了LLaMA使用RMSNorm。
3. 归一化层的作用是什么?
分析设置。我们首先对训练网络中归一化层的行为进行实证研究。为此,我们采用了一个在ImageNet-1K(Deng等人, 2009)上训练的Vision Transformer模型(ViT-B)(Dosovitskiy等人, 2020)、一个在LibriSpeech(Panayotov等人, 2015)上训练的wav2vec 2.0 Large Transformer模型(Baevski等人, 2020)和一个在ImageNet-1K上训练的Diffusion Transformer(DiT-XL)(Peebles和Xie, 2023)。在所有情况下,LN都应用于每个Transformer块以及最终线性投影之前。
对于这三个训练好的网络,我们对一个迷你批次的样本进行前向传播,并测量归一化层的输入和输出,即归一化操作之前和之后的张量,以及在LN中的可学习仿射变换之前。由于LN保留了输入张量的维度,我们可以建立输入和输出张量元素之间的一一对应关系,从而可以直接可视化它们之间的关系。我们在图2中绘制了这些映射。
在所有三个模型中,早期LN层(图2的第一列)的输入-输出关系大多是线性的,在x-y图中类似于一条直线。然而,在更深的LN层中,我们观察到了更有趣的发现。
图4 展示了两个层归一化(LN)层的输出与输入关系,其中张量元素通过颜色表示不同的通道和token维度。输入张量的形状为(样本,token,通道),通过为相同的token(左两幅图)和通道(右两幅图)分配一致的颜色来可视化元素。左两幅图:表示相同token(颜色相同)的点在不同通道中形成直线,因为LN对每个token的通道进行线性操作。有趣的是,当这些直线集体绘制时,它们形成了一个非线性的双曲正切(tanh)形状曲线。右两幅图:每个通道的输入在x轴上跨越不同的范围,为整体的双曲正切形状曲线贡献了不同的段。某些通道(例如红色、绿色和粉色)表现出更极端的x值,这些值被LN压缩。
一个显著的观察结果是,这些曲线的形状高度类似于由双曲正切函数表示的完整或部分S形曲线(见图3)。人们可能期望LN层对输入张量进行线性变换,因为减去均值和除以标准差是线性操作。LN以每个token的方式进行归一化,只对每个token的激活值进行线性变换。由于token具有不同的均值和标准差值,线性关系在所有输入张量的激活值上并不成立。尽管如此,实际的非线性变换与缩放的tanh函数高度相似仍然令人惊讶。
对于这种S形曲线,我们注意到中心部分(x值接近零的点)仍然主要呈线性形状。大约99%的点落在这个线性范围内。然而,仍然有许多点明显超出了这个范围,这些点被认为是具有“极端”值的点,例如在ViT模型中x大于50或小于-50的点。归一化层对这些值的主要作用是将它们压缩成更不极端的值,使其更接近大多数点。这就是归一化层不能被简单仿射变换层近似的部分。我们假设这种对极端值的非线性和不成比例的压缩效应是使归一化层重要且不可或缺的原因。
Ni等人(2024)的最新研究也强调了LN层引入的强烈非线性,展示了这种非线性如何增强模型的表示能力。此外,这种压缩行为反映了生物神经元对大输入的饱和特性,这一现象在一个世纪前首次被观察到(Adrian, 1926; Adrian和Zotterman, 1926a,b)。
按token和通道进行归一化。LN层如何对每个token进行线性变换,同时又以这种非线性的方式压缩极端值呢?为了理解这一点,我们分别按token和通道对点进行分组并可视化。这在图4中绘制,取自图2中ViT的第二和第三个子图,但为了更清晰,只采样了一部分点。在选择通道进行绘制时,我们确保包括具有极端值的通道。
在图4的左两幅图中,我们使用相同的颜色对每个token的激活值进行可视化。我们观察到,来自任何单个token的所有点确实形成了一条直线。然而,由于每个token具有不同的方差,斜率也不同。输入x范围较小的token倾向于具有较小的方差,归一化层将使用较小的标准差对其激活值进行划分,因此会产生较大的斜率。集体来看,它们形成了一个类似于双曲正切函数的S形曲线。在右两幅图中,我们使用相同的颜色对每个通道的激活值进行可视化。我们发现不同通道的输入范围差异很大,只有少数通道(例如红色、绿色和粉色)表现出较大的极端值。这些是被归一化层压缩最多的通道。
4. 动态双曲正切(DyT)
受归一化层形状与缩放双曲正切函数相似性的启发,我们提出了动态双曲正切(DyT)作为归一化层的替代品。给定一个输入张量x,DyT层定义如下:
DyT(x) = γ * tanh(αx) + β (2)
其中α是一个可学习的标量参数,允许根据输入的范围以不同的方式缩放输入,以适应不同的x尺度(见图2)。这也是我们将其命名为“动态”双曲正切的原因。γ和β是与所有归一化层中使用的相同的每个通道的可学习向量参数——它们允许输出缩回到任何尺度。这有时被认为是一个单独的仿射层;为了我们的目的,我们将其视为DyT层的一部分,就像归一化层也包括它们一样。有关DyT层在Pytorch风格伪代码中的实现,请参见算法1。
算法1:DyT层的伪代码。
# 输入x的形状为[B, T, C] # B:批量大小,T:token,C:维度
class DyT(Module):
def __init__(self, C, init_α):
super().__init__()
self.α = Parameter(ones(1) * init_α)
self.γ = Parameter(ones(C))
self.β = Parameter(zeros(C))
def forward(self, x):
x = tanh(self.α * x)
return self.γ * x + self.β
将DyT层集成到现有架构中非常简单:一个DyT层替换一个归一化层(见图1)。这适用于注意力块、FFN块以及最终归一化层中的归一化层。尽管DyT看起来像或被认为是一个激活函数,但本研究仅用它替换归一化层,而不改变原始架构中的任何激活函数部分,如GELU或ReLU。网络的其他部分也保持不变。我们还观察到,在适应DyT时,几乎没有必要调整原始架构的超参数。
关于缩放参数。我们始终简单地将γ初始化为全1向量,将β初始化为全0向量,遵循归一化层的做法。对于缩放参数α,0.5的默认初始化通常就足够了,除了LLM训练。关于α初始化的详细分析在第7节中提供。除非另有明确说明,否则在随后的实验中,α初始化为0.5。
备注:DyT不是一种新的归一化层类型,因为它在前向传播期间独立于张量的每个输入元素进行操作,而不计算统计量或其他类型的聚合。然而,它确实保留了归一化层的效果,以非线性的方式压缩极端值,同时几乎线性地变换输入的中心部分。
5. 实验
为了证明DyT的有效性,我们在Transformer和一些其他现代架构上进行了实验,涵盖了各种任务和领域。在每个实验中,我们将原始架构中的LN或RMSNorm替换为DyT层,并按照官方开源协议训练和测试这两种模型版本。有关复现我们结果的详细说明,请参见附录A。值得注意的是,为了突出适应DyT的简单性,我们在实验中使用了与归一化对应模型相同的超参数。为了完整性,有关学习率调整和α初始值的额外实验结果在附录B中提供。
计算机视觉中的监督学习。我们在ImageNet-1K分类任务(Deng等人, 2009)上训练了“基础”和“大”尺寸的Vision Transformer(ViT)(Dosovitskiy等人, 2020)和ConvNeXt(Liu等人, 2022)。这些模型因其流行度和不同的操作(ViT中的注意力和ConvNeXt中的卷积)而被选中。表1报告了top-1分类精度。DyT在两种架构和模型尺寸上均略优于LN。我们进一步绘制了ViT-B和ConvNeXt-B的训练损失曲线,如图5所示。这些曲线表明,DyT和基于LN的模型的收敛行为高度一致。
# ViT-B和ConvNeXt-B模型的训练损失曲线。两种模型类型的损失曲线显示出相似的模式,表明LN和DyT可能具有相似的学习动态。
模型 | LN | DyT | 变化 |
---|---|---|---|
ViT-B | 82.3% | 82.5% | ↑ 0.2% |
ViT-L | 83.1% | 83.6% | ↑ 0.5% |
ConvNeXt-B | 83.7% | 83.7% | - |
ConvNeXt-L | 84.3% | 84.4% | ↑ 0.1% |
表1:ImageNet-1K上的监督分类精度。DyT在两种架构和模型尺寸上均实现了与LN相当或更好的性能。
计算机视觉中的自监督学习。我们使用两种流行的视觉自监督学习方法进行基准测试:掩码自编码器(MAE)(He等人, 2022)和DINO(Caron等人, 2021)。两者默认都使用Vision Transformers作为骨干网络,但训练目标不同:MAE使用重建损失进行训练,DINO使用联合嵌入损失(LeCun, 2022)。按照标准自监督学习协议,我们首先在ImageNet-1K上对模型进行预训练,不使用任何标签,然后通过附加分类层并使用标签进行微调来测试预训练模型。微调结果如表2所示。DyT在自监督学习任务中与LN表现相当。
模型 | LN | DyT | 变化 |
---|---|---|---|
MAE ViT-B | 83.2% | 83.2% | - |
MAE ViT-L | 85.5% | 85.4% | ↓ 0.1% |
DINO ViT-B (patch size 16) | 83.2% | 83.4% | ↑ 0.2% |
DINO ViT-B (patch size 8) | 84.1% | 84.5% | ↑ 0.4% |
表2:ImageNet-1K上的自监督学习精度。DyT在不同预训练方法和模型尺寸的自监督学习任务中表现相当。
扩散模型。我们在ImageNet-1K(Deng等人, 2009)上训练了三种不同尺寸的Diffusion Transformer(DiT)模型(Peebles和Xie, 2023),分别为B、L和XL。它们的patch大小分别为4、4和2。需要注意的是,在DiT中,LN层的仿射参数用于类别条件化,在我们的DyT实验中,我们保持这种方式,仅用tanh(αx)函数替换归一化变换。训练后,我们使用标准ImageNet“参考批次”评估Fréchet Inception Distance(FID)分数,如表3所示。DyT在FID方面实现了与LN相当或更好的性能。
模型 | LN | DyT | 变化 |
---|---|---|---|
DiT-B | 64.9 | 63.9 | ↓ 1.0 |
DiT-L | 45.9 | 45.7 | ↓ 0.2 |
DiT-XL | 19.9 | 20.8 | ↑ 0.9 |
表3:ImageNet上的图像生成质量(FID,越低越好)。DyT在各种DiT模型尺寸上实现了与LN相当或更好的FID分数。
大型语言模型。我们预训练了7B、13B、34B和70B尺寸的LLaMA模型(Touvron等人, 2023a,b; Dubey等人, 2024),以评估DyT相对于RMSNorm(Zhang和Sennrich, 2019)的性能,这是LLaMA中使用的默认归一化层。这些模型在The Pile数据集(Gao等人, 2020)上进行了训练,使用了200B tokens,遵循LLaMA中概述的原始配方(Touvron等人, 2023b)。对于使用DyT的LLaMA,我们在初始嵌入层之后添加了一个可学习的标量参数,并调整了α的初始值,详细信息在第7节中提供。我们在表4中报告训练后的损失值,并按照OpenLLaMA(Geng和Liu, 2023)的建议,在lm-eval(Gao等人)中的15个零样本任务上对模型进行基准测试。如表4所示,DyT在所有四种模型尺寸上均实现了与RMSNorm相当的零样本性能和训练损失。
# LLaMA预训练损失。DyT和RMSNorm模型的损失曲线在所有模型尺寸上紧密对齐。
模型 | RMSNorm | DyT | 变化 |
---|---|---|---|
LLaMA 7B | 0.513 / 1.59 | 0.513 / 1.60 | - / ↑ 0.01 |
LLaMA 13B | 0.529 / 1.53 | 0.529 / 1.54 | - / ↑ 0.01 |
LLaMA 34B | 0.536 / 1.50 | 0.536 / 1.50 | - / - |
LLaMA 70B | 0.549 / 1.45 | 0.549 / 1.45 | - / - |
表4:语言模型的训练损失和在15个零样本lm-eval任务上的平均性能。DyT在训练损失和零样本性能方面均实现了与RMSNorm相当的结果。
语音自监督学习。我们在LibriSpeech数据集(Panayotov等人, 2015)上预训练了两个wav2vec 2.0 Transformer模型(Baevski等人, 2020)。我们在表5中报告了最终验证损失。我们观察到DyT在两种模型尺寸上均与LN表现相当。
模型 | LN | DyT | 变化 |
---|---|---|---|
wav2vec 2.0 Base | 1.95 | 1.95 | - |
wav2vec 2.0 Large | 1.92 | 1.91 | ↓ 0.01 |
表5:LibriSpeech上的语音预训练验证损失。DyT在两种wav2vec 2.0模型上均实现了与LN相当的性能。
DNA序列建模。在长程DNA序列建模任务中,我们预训练了HyenaDNA模型(Nguyen等人, 2024)和Caduceus模型(Schiff等人, 2024)。预训练使用了人类参考基因组数据(GRCh38, 2013),并在GenomicBenchmarks(Grešová等人, 2023)上进行评估。结果如表6所示。DyT在该任务中保持了与LN相当的性能。
模型 | LN | DyT | 变化 |
---|---|---|---|
HyenaDNA(Nguyen等人, 2024) | 85.2% | 85.2% | - |
Caduceus(Schiff等人, 2024) | 86.9% | 86.9% | - |
表6:GenomicBenchmarks上的DNA分类精度,平均值取自GenomicBenchmarks中的每个数据集。DyT实现了与LN相当的性能。
6. 分析
我们对DyT的重要属性进行了几项分析。我们首先评估其计算效率,随后进行了两项研究,分别考察tanh函数和可学习尺度α的作用。最后,我们与旨在移除归一化层的先前方法进行了比较。
6.1 DyT的效率
我们通过测量100次前向传播(推理)和100次前向-后向传播(训练)所需的总时间来对LLaMA 7B模型中的RMSNorm或DyT进行基准测试,使用单个4096 tokens的序列。表7报告了在Nvidia H100 GPU上以BF16精度运行时,所有RMSNorm或DyT层以及整个模型所需的时间。与RMSNorm层相比,DyT层显著减少了计算时间,在FP32精度下也观察到了类似的趋势。DyT可能是面向效率的网络设计的一个有前途的选择。
模型 | 推理 | 训练 |
---|---|---|
RMSNorm | 2.1s | 8.3s |
DyT | 1.0s | 4.8s |
减少 | ↓ 52.4% | ↓ 42.2% |
模型 | 层 | 模型 |
---|---|---|
RMSNorm | 14.1s | 42.6s |
DyT | 13.0s | 39.1s |
减少 | ↓ 7.8% | ↓ 8.2% |
表7:LLaMA 7B在BF16精度下的推理和训练延迟。DyT在推理和训练时间上均实现了显著减少。
6.2 tanh和α的消融实验
为了进一步研究DyT中tanh和α的作用,我们进行了实验,以评估在改变或移除这些组件时模型的性能。
替换和移除tanh。我们用替代压缩函数hardtanh和sigmoid替换DyT层中的tanh(见图7),同时保持可学习缩放器α完整。此外,我们还通过用恒等函数替换tanh来完全移除它,同时仍然保留α。如表8所示,压缩函数对于稳定训练至关重要。使用恒等函数会导致训练不稳定和发散,而压缩函数则可以实现稳定训练。在压缩函数中,tanh表现最佳。这可能是因为它的平滑性和零中心属性。
模型 | 恒等函数 | tanh | hardtanh | sigmoid |
---|---|---|---|---|
ViT-S | 58.5% → 失败 | 80.3% | 79.9% | 79.6% |
ViT-B | 61.0% → 失败 | 82.5% | 82.2% | 81.6% |
表8:使用不同压缩函数的ImageNet-1K分类精度。所有实验均遵循原始LN模型的相同训练配方。压缩函数在防止发散中发挥关键作用,其中tanh在三种函数中表现最佳。“→ 失败”表示训练在取得一些进展后发散,前面的数字表示发散前达到的最高精度。
移除α。接下来,我们在保留压缩函数(tanh、hardtanh和sigmoid)的同时,评估移除可学习α的影响。如表9所示,移除α会导致所有压缩函数的性能下降,这突显了α在整体模型性能中的关键作用。
模型 | tanh | hardtanh | sigmoid |
---|---|---|---|
无α | 81.1% | 80.7% | 80.7% |
有α | 82.5% | 82.2% | 81.6% |
表9:使用ViT-B在ImageNet-1K上的分类精度。所有实验均遵循原始LN模型的相同训练配方。可学习的α对于增强模型性能至关重要。
6.3 α的值
在训练期间。我们的分析表明,α在训练期间始终与激活值的标准差一起变化。如图8左图所示,α在训练初期下降,随后上升,但始终与输入激活值的标准差一致波动。这支持了α在维持激活值在合适范围内的重要作用,从而实现稳定有效的训练。
在训练之后。我们对训练完成的网络中α的最终值进一步分析,发现其与输入激活值的1/std之间存在强相关性。如图8右图所示,较高的1/std值通常对应较大的α值,反之亦然。此外,我们观察到较深的层倾向于具有标准差较大的激活值。这一趋势与深度残差网络的特性一致,如Brock等人(2021a)对ConvNets的研究以及Sun等人(2025)对Transformers的研究所示。
这两项分析均表明,α部分地作为一种归一化机制,通过学习近似于输入激活值的1/std的值。与LN不同的是,LN对每个token的激活值进行归一化,而α则对整个输入激活值进行集体归一化。因此,单独的α无法以非线性方式抑制极端值。
6.4 与其他方法的比较
为了进一步评估DyT的有效性,我们将它与其他能够实现无需归一化层训练Transformer的方法进行了比较。这些方法大致可以分为基于初始化的方法和基于权重归一化的方法。我们考虑了两种流行的基于初始化的方法:Fixup(Zhang等人, 2019; Huang等人, 2020)和SkipInit(De和Smith, 2020; Bachlechner等人, 2021)。这两种方法均旨在通过调整初始参数值来减轻训练不稳定性,从而在训练开始时防止出现大的梯度和激活值,从而实现无需归一化层的稳定学习。相比之下,基于权重归一化的方法在整个训练过程中对网络权重施加约束,以维持在没有归一化层的情况下的稳定学习动态。我们包括了其中一种方法,σReparam(Zhai等人, 2023),它通过控制权重的谱范数来促进稳定学习。
我们在表10中总结了两种基于ViT的任务的结果。我们严格遵循这些方法各自论文中概述的原始协议。然而,我们发现这两种基于初始化的方法,Fixup和SkipInit,都需要显著降低学习率以防止训练发散。为了确保公平比较,我们对所有方法(包括DyT)进行了简单学习率搜索。这产生了与第5节中报告的结果不同的结果,其中未调整任何超参数。总体而言,结果表明DyT在不同配置中均一致优于所有其他测试方法。
模型 | LN | Fixup | SkipInit | σReparam | DyT |
---|---|---|---|---|---|
ViT-B | 82.3% | 77.2% | 74.1% | 82.5% | 82.8% |
ViT-L | 83.1% | 78.1% | 75.6% | 83.0% | 83.6% |
MAE ViT-B | 83.2% | 73.7% | 73.1% | 83.2% | 83.7% |
MAE ViT-L | 85.5% | 74.1% | 74.0% | 85.4% | 85.8% |
表10:ImageNet-1K上的分类精度。DyT在所有配置中均一致优于其他测试方法。
7. α的初始化
我们发现调整α的初始化(记作α₀)很少能带来显著的性能提升。唯一的例外是LLM训练,其中仔细调整α₀可以获得显著的性能提升。在本节中,我们详细阐述了我们关于α初始化的影响的发现。
7.1 非LLM模型的α初始化
非LLM模型对α₀相对不敏感。图9展示了在不同任务中改变α₀对验证性能的影响。所有实验均遵循各自配方的原始设置和超参数。我们观察到,在较宽的α₀值范围内,性能保持稳定,0.5到1.2之间的值通常可以获得良好的结果。我们观察到,调整α₀通常只影响训练曲线的早期阶段。唯一的例外是监督ViT-L实验,当α₀超过0.6时,训练变得不稳定并发散。在这种情况下,降低学习率可以恢复稳定性,详细信息如下所述。
较小的α₀可以实现更稳定的训练。基于之前的观察,我们进一步分析了导致训练不稳定的因素。我们的发现表明,增加模型大小或学习率需要降低α₀以确保稳定训练。相反,较高的α₀需要较低的学习率以减轻训练不稳定性。图10展示了在ImageNet-1K数据集上训练监督ViT时的训练稳定性消融实验。我们改变了学习率、模型大小和α₀值。训练更大的模型更容易失败,需要较小的α₀值或学习率以实现稳定训练。在类似条件下,LN基础模型也观察到了类似的不稳定性模式,将α₀设置为0.5可以实现与LN相似的稳定性模式。
将α₀设置为0.5作为默认值。基于我们的发现,我们将α₀ = 0.5设置为所有非LLM模型的默认值。这一设置提供了与LN相当的训练稳定性,同时保持了强大的性能。
7.2 LLM的α初始化
调整α₀可以提升LLM性能。如前所述,α₀ = 0.5的默认设置在大多数任务中表现良好。然而,我们发现调整α₀可以显著提升LLM性能。我们在30B tokens上预训练了每种LLaMA模型,并比较了它们的训练损失。表11总结了每种模型的调整后的α₀值。得出了两个关键发现:
-
较大的模型需要较小的α₀值。一旦确定了较小模型的最优α₀,就可以相应地缩小较大模型的搜索空间。
-
在注意力块中为α₀设置更高的值可以提升性能。我们发现,在注意力块中的DyT层初始化α时使用更高的值,而在其他位置(即在FFN块或最终线性投影之前)的DyT层初始化α时使用较低的值,可以提升性能。
模型 | 宽度 | 深度 | 最优α₀(注意力/其他) |
---|---|---|---|
LLaMA 7B | 4096 | 32 | 0.8/0.2 |
LLaMA 13B | 5120 | 40 | 0.6/0.15 |
LLaMA 34B | 8196 | 48 | 0.2/0.05 |
LLaMA 70B | 8196 | 80 | 0.2/0.05 |
表11:不同LLaMA模型的最优α₀。较大的模型需要较小的α₀值。我们发现,为注意力块中的DyT层初始化α时使用更高的值很重要。
为了进一步说明α₀调整的影响,图11展示了两种LLaMA模型的损失值热图。这两种模型都从注意力块中较高的α₀值中受益,从而降低了训练损失。
# 两种LLaMA模型在不同α₀设置下的损失值热图。两种模型都从注意力块中较高的α₀值中受益。
模型宽度主要决定α₀的选择。我们还研究了模型宽度和深度对最优α₀的影响。我们发现模型宽度是决定最优α₀的关键因素,而模型深度的影响可以忽略不计。表12展示了不同宽度和深度下的最优α₀值,表明较宽的网络需要较小的α₀值以实现最佳性能。另一方面,模型深度对α₀的选择几乎没有影响。
宽度/深度 | 8 | 16 | 32 | 64 |
---|---|---|---|---|
1024 | 1.0/1.0 | 1.0/1.0 | 1.0/1.0 | 1.0/1.0 |
2048 | 1.0/0.5 | 1.0/0.5 | 1.0/0.5 | 1.0/0.5 |
4096 | 0.8/0.2 | 0.8/0.2 | 0.8/0.2 | 0.8/0.2 |
8192 | 0.2/0.05 | 0.2/0.05 | 0.2/0.05 | 0.2/0.05 |
表12:不同宽度和深度的LLaMA训练中的最优α₀(注意力/其他)。模型宽度显著影响α₀的选择,较宽的网络需要较小的值。相比之下,模型深度对α₀的选择几乎没有影响。
正如表12所示,网络越宽,注意力和其他部分的初始化就越不均匀。我们假设LLM对α初始化的敏感性与其相对于其他模型的过度宽度有关。
8. 相关工作
归一化层的机制。有大量的工作研究了归一化层通过各种机制增强模型性能的作用。这些机制包括在训练期间稳定梯度流(Balduzzi等人, 2017; Daneshmand等人, 2020; Lubana等人, 2021),减少对权重初始化的敏感性(Zhang等人, 2019; De和Smith, 2020; Shao等人, 2020),调节异常特征值(Bjorck等人, 2018; Karakida等人, 2019),自动调整学习率(Arora等人, 2018; Tanaka和Kunin, 2021),以及平滑损失景观以实现更稳定的优化(Santurkar等人, 2018)。这些早期工作主要集中在研究批量归一化。最近的研究(Lyu等人, 2022; Dai等人, 2024; Mueller等人, 2024)进一步强调了归一化层与锐度降低之间的联系,这有助于更好的泛化。
Transformer中的归一化。随着Transformer(Vaswani等人, 2017)的兴起,研究越来越多地集中在层归一化(Ba等人, 2016)上,它在自然语言任务中的序列数据方面已被证明特别有效(Nguyen和Salazar, 2019; Xu等人, 2019; Xiong等人, 2020)。最近的研究(Ni等人, 2024)揭示了层归一化引入的强烈非线性,增强了模型的表示能力。此外,研究(Loshchilov等人, 2024; Li等人, 2024)表明,修改Transformer中归一化层的位置可以改善收敛特性。
移除归一化。许多研究探索了如何在没有归一化层的情况下训练深度模型。一些工作(Zhang等人, 2019; De和Smith, 2020; Bachlechner等人, 2021)探索了替代的权重初始化方案以稳定训练。开创性的工作由Brock等人(2021a,b)完成,他们表明可以通过结合初始化技术(De和Smith, 2020)、权重归一化(Salimans和Kingma, 2016; Huang等人, 2017; Qiao等人, 2019)和自适应梯度裁剪(Brock等人, 2021b)来训练高性能的ResNet,无需归一化(Smith等人, 2023)。此外,他们的训练策略还包括广泛的数据增强(Cubuk等人, 2020)和正则化(Srivastava等人, 2014; Huang等人, 2016)。上述研究均基于各种ConvNet模型。
在Transformer架构中,He和Hofmann(2023)探索了对Transformer块的修改,以减少对归一化层和跳跃连接的依赖。另一种方法是Heimersheim(2024)提出的,通过在移除每个归一化层后对模型进行微调,逐步从预训练的网络中移除LN。与先前的方法不同,DyT几乎不需要对架构和训练配方进行修改。尽管其简单性,DyT能够实现稳定训练并达到相当的性能。
9. 限制
我们在使用LN或RMSNorm的网络上进行了实验,因为它们在Transformer和其他现代架构中非常流行。初步实验(见附录C)表明,DyT在经典网络(如ResNet)中直接替换BN时存在困难。需要进一步深入研究DyT是否以及如何适应具有其他类型归一化层的模型。
10. 结论
在本研究中,我们证明了现代神经网络,特别是Transformer,可以在没有归一化层的情况下进行训练。这是通过动态双曲正切(DyT)实现的,它是传统归一化层的简单替代品。它通过可学习的缩放因子α调整输入激活值范围,然后通过S形的tanh函数压缩极端值。尽管是一个更简单的函数,但它有效地捕捉了归一化层的行为。在各种设置中,带有DyT的模型达到了与归一化对应模型相当甚至更好的性能。这些发现挑战了归一化层在训练现代神经网络中不可或缺的传统理解。我们的研究还为深度神经网络中最基本的构建块之一——归一化层的机制提供了见解。