何恺明等暴击Transformer9行代码砍掉归一化层!性能反而更强了

360影视 日韩动漫 2025-03-14 16:52 2

摘要:Transformer架构迎来历史性突破!刚刚,何恺明LeCun、清华姚班刘壮联手,用9行代码砍掉了Transformer「标配」归一化层,创造了性能不减反增的奇迹。

编辑:编辑部 ZYJ

【新智元导读】Transformer架构迎来历史性突破!刚刚,何恺明LeCun、清华姚班刘壮联手,用9行代码砍掉了Transformer「标配」归一化层,创造了性能不减反增的奇迹。

就在刚刚,何恺明和LeCun两大巨头联手,彻底改变了Tranformer架构!

在Meta、NYU、MIT、普林斯顿等联手的这项工作中,研究者们共同发现了一种没有归一化层的Transformer。

论文地址:https://arxiv.org/abs/2503.10622

项目地址:https://github.com/jiachenzhu/DyT

在深度学习中,归一化层(Norm layer,下图Transformer架构中的绿色块)无处不在,并且长期以来一直被认为是不可或缺的。

然而,这次何恺明、Yann LeCun等人却证明:即使不使用归一化层的Transformer模型,也就能达到相同甚至更好的性能!

而且方法也非常简单,只需要9行代码。

class DyT(nn.Module): def __init__(self, num_features, alpha_init_value=0.5): super.__init__ self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value) self.weight = nn.Parameter(torch.ones(num_features)) self.bias = nn.Parameter(torch.zeros(num_features)) def forward(self, x): x = torch.tanh(self.alpha * x) return x * self.weight + self.bias

简单来说,他们发现了一种简单得出人意料的方法,来替代归一化层——缩放的tanh函数(没错,仿佛回到了80年代)。

他们称之为动态Tanh(Dynamic Tanh,简称DyT)。

这一方法的核心,源自一个非常简单的观察:LayerNorm将输入转换为输出时,呈现出类似tanh的S形曲线。

它能够压缩极端值,同时在中心区域保持线性形态。

因此,研究者用提出的动态Tanh(DyT)层替换了归一化层,过程很简单。

具体来说,他们用DyT替换了LayerNorm/RSMNorm,并在以下 Transformer结构上进行了测试:

视觉:ViT、ConvNeXt、MAE、DINO、DiT大模型(LLM):LLaMA语音:wav2vec 2.0DNA:HyenaDNA、Caduceus

结果表明,在所有测试场景中,DyT的效果都与归一化层相当或更优。

而且,DyT比RMSNorm在H100s上更快!

这项研究的意义就在于,鉴于模型训练和推理可能需要数千万次的计算,因此DyT有极大潜力帮助降低成本,接下来就是找到应用。

网友们对此纷纷表示:非常精彩的研究!而它改变应用transformer归一化基本假设的方式,尤其令人喜欢。

或许,这项研究能从根本上解决深度学习问题?

也有人提出,tanh一直因梯度消失问题而臭名昭著。如果用tanh替换归一化层,是否会重新引入这一问题,进而可能减缓甚至阻碍深度模型的学习呢?

打破固有观念:归一化层并非必不可少!

在过去的十年里,归一化层已经稳固地成为现代神经网络最基础的组件之一。

如今,几乎所有现代神经网络都采用归一化层,其中层归一化(Layer Normalization,LN)尤其流行,特别是在占主导地位的 Transformer 架构中。

原因主要在于它们对优化过程的实证优势。因此,归一化层被普遍认为是 深度网络高效训练的关键甚至不可或缺的组成部分。

新的神经网络架构往往会尝试替换注意力层或卷积层,但几乎都会保留归一化层。

然而这项工作,却提出了一种简单的归一化层替代方案,并应用于 Transformer结构。

研究始于一个关键观察:LN层将输入映射为类似tanh的S形曲线,既缩放了输入激活值,又抑制了极端值。

因此,研究者提出了一种逐元素(element-wise)操作,称为动态 tanh(Dynamic Tanh, DyT),定义如下:

其中,α 是一个可学习参数。

DyT的设计目标是:

通过α学习合适的缩放因子,通过tanh的有界性抑制极端值。

这也是整个操作命名为「动态」Tanh的原因。

值得注意的是,与归一化层不同,DyT无需计算激活统计信息,却能够同时实现缩放和极值抑制。

γ和β是可学习的、逐通道的向量参数,与所有归一化层中使用的参数相同——它们允许输出缩放到任意尺度。

有时这被认为是一个独立的仿射层;但在新设计中,它们被视为DyT层的一部分,就像归一化层也包含它们一样。

图2:Vision Transformer(ViT)、wav2vec 2.0(用于语音的Transformer模型)和Diffusion Transformer(DiT)中部分层归一化(LN)层的输出与输入关系

有关DyT在类似Pytorch伪代码中的实现,请参见算法1。

算法1:9行代码,即出DyT

将DyT层集成到现有架构中非常简单:DyT层可以直接替换归一化层(见图1)。

这种替换适用于注意力模块、前馈网络(FFN)模块中的归一化层,以及最终的归一化层。

尽管DyT可能看起来像激活函数,或者被认为是一种激活函数,但本研究仅将其用于替换归一化层,而不会改变原始架构中的任何激活函数(如GELU或ReLU)。

网络的其他部分也保持不变。

此外,DyT在应用于原始架构时,几乎不需要调整超参数就能达到良好的性能。

归一化层有什么用?

在所有三个模型中,可以观察到较浅的LayerNorm(LN)层(图2第一列)的输入-输出关系大多是线性的,在x-y平面上接近一条直线。

然而,在更深的LN层中,可以发现了一些更具启发性的现象。

在大多数深层LN层中,输入-输出曲线高度类似于tanh函数所呈现的完整或部分S形曲线(见图3)。

按常理,可能会认为LN层只是对输入张量进行线性变换,因为减去均值和除以标准差都是线性操作。

LN的归一化方式是逐token进行的,因此它对每个token施加的变换是线性的。

然而,由于不同token具有不同的均值和标准差,LN无法整体上保持输入张量所有激活值的线性关系。

尽管如此,让人仍然感到惊讶的是,实际的非线性变换竟然高度类似于缩放后的tanh函数。

在这种S形曲线中,可以观察到:

中心部分(x取值接近0)仍然基本保持线性,大约 99% 的点都位于这一线性范围内。

然而,仍有一些点超出了这一范围,即那些被认为是「极端」值的点。例如,在ViT模型中,x>50或x

归一化层的主要作用之一,便是将这些极端值压缩到相对较温和的范围,使其更加符合大部分数据点的分布。

正是在这个方面,归一化层的行为无法被简单的仿射变换所逼近。

研究假设,归一化层对极端值的这种非线性和非均匀压缩作用,正是使其至关重要且难以被取代的核心原因。

为了理解这一点,该团队分别按token和通道对这些点进行可视化。

可以看到,在左侧两幅图中,相同token(相同颜色)在不同通道上的数据点排列成直线,这是因为LN在每个token内按通道进行线性变换。

然而,当所有token的数据点整体绘制在一张图上时,这些直线却形成了一个非线性的tanh形曲线。

右侧两幅图中,每个通道的输入值分布在x轴上的不同范围,共同构成了整体的tanh形曲线。

某些通道(如红色、绿色和粉色)包含更极端的x值,这些极端值被LN压缩到较温和的范围。

实验

为了证明DyT的有效性,在Transformer及其他几种现代架构上进行了实验,涵盖了多个不同的任务和领域。

在每次实验中,用DyT层替换原始架构中的 LN(Layer Normalization)或 RMSNorm,并遵循官方开源协议来训练和测试两种版本的模型。

值得注意的是,为了突出 DyT 适配的简便性,使用的超参数与原始归一化模型完全相同。

视觉中的监督学习

在ImageNet-1K分类任务上,训练了「Base」和「Large」规模的Vision Transformer(ViT)和ConvNeXt。

选择这些模型是因为它们的广泛流行以及不同的操作特性:ViT中的注意力机制和ConvNeXt中的卷积操作。

表1报告了Top-1分类准确率。DyT在两种架构和模型规模上均略优于LN。

进一步绘制了ViT-B和ConvNeXt-B的训练损失曲线,如图5所示。曲线表明,DyT和基于LN的模型的收敛行为高度一致。

视觉中的自监督学习

使用两种流行的视觉自监督学习方法进行基准测试:掩码自编码器(MAE)和DINO。

两者默认均使用Vision Transformers作为骨干网络,但具有不同的训练目标:MAE通过重建损失进行训练,而DINO使用联合嵌入损失。

按照标准的自监督学习流程,首先在不使用任何标签的情况下在ImageNet-1K上对模型进行预训练,然后通过附加分类层并使用标签对预训练模型进行微调来测试其性能。

微调结果如表2所示。DyT在自监督学习任务中始终表现与LN相当。

扩散模型

在ImageNet-1K上训练了三种规模的Diffusion Transformer(DiT)模型,分别为B、L和XL。其patch大小分别为4、4和2。

需要注意的是,在DiT中,LN层的仿射参数用于类别条件化(class conditioning),在DyT实验中也保留了这一点,仅将归一化变换替换为tanh(αx)函数。

训练完成后,使用标准的ImageNet「参考批次」评估Fréchet Inception Distance(FID)分数,结果如表3所示。DyT在FID上取得了与LN相当或更好的性能。

大语言模型

研究人员预训练了LLaMA 7B、13B、34B和70B模型,以评估DyT相对于RMSNorm的性能,RMSNorm是LLaMA中默认使用的归一化层。

如表4所示,DyT在所有四种模型规模上的表现均与RMSNorm相当。

图6展示了损失曲线,表明所有模型规模的训练趋势相似,训练损失在整个训练过程中高度一致。

DNA序列建模

在长程DNA序列建模任务中,预训练了HyenaDNA模型和Caduceus模型。

预训练使用了来自人类参考基因组数据,并在GenomicBenchmarks上进行评估。

结果如表6所示。DyT在该任务中保持了与LN相当的性能。

语音中的自监督学习

在LibriSpeech数据集上,预训练了两个wav2vec 2.0 Transformer模型。

表5中报告了最终的验证损失。

可以观察到,DyT在两种模型规模上的表现均与LN相当。

额外的实验

研究团队进行了额外的实验,以评估超参数调优的影响,特别是针对所有非大语言模型(non-LLM)的学习率和α初始化。

学习率调优:表13总结了使用原始学习率与调优后学习率训练的模型之间的性能对比。

结果表明,调优学习率对DyT模型的性能提升较为有限。

这表明,最初为LN模型优化的原始超参数已经非常适合DyT模型。这一观察结果进一步强调了DyT与LN模型之间的内在相似性。

此外,本次还研究了优化DyT模型初始值 α₀ 的影响,相关结果见表14。

研究发现,在某些模型上调整 α₀ 仅带来轻微的性能提升,这表明默认初始值(α₀ = 0.5)通常已经能够实现接近最优的性能。

研究者对DyT的重要特性进行了多方面分析。

首先,评估其计算效率,随后进行了两项研究,分别探讨了tanh函数的作用以及可学习缩放参数α的影响。

最后,他们与其他试图去除归一化层的方法进行了对比分析。

DyT的计算效率

研究在LLaMA 7B模型上,对RMSNorm和DyT进行了基准测试,测量了100次前向传播(推理) 和100次前向-反向传播(训练) 的总时间。

实验采用4096个token的单一序列,在 Nvidia H100(BF16 精度) 上运行。

表7报告了所有 RMSNorm或DyT层以及整个模型的运行时间。

结果表明:

DyT层的计算时间显著低于RMSNorm层,在FP32精度下也观察到了类似的趋势。

由此可见,DyT可能是面向效率优化的网络设计中的一个有前景的选择。

α的分析

训练过程中的α

分析表明,α在整个训练过程中始终紧密追踪激活值的1/std(标准差的倒数)。

如图8(左)所示,α在训练早期先下降后上升,但始终与输入激活值的标准差 保持一致的波动趋势。这一现象表明:

α在训练中起到了关键作用,帮助激活值保持在适当范围内从而促进稳定且高效的训练

训练后的α

对训练完成后的α值进行进一步分析,我们发现它与输入激活的1/std存在强相关性。

如图8(右)所示,当1/std值较大时,α也较大,反之亦然。此外,可以观察到:越深的网络层,其激活值的标准差往往越大。

这一趋势与深度残差网络的特性一致,与此前研究所揭示的一致。

α作为归一化机制的作用

上述分析表明,α在一定程度上发挥了类似归一化的功能,其学习到的值大致接近 输入激活的1/std。但需要注意的是:

与 LayerNorm(LN)不同,LN是对每个token进行归一化,而α作用于整个输入激活值的整体分布。因此,α无法像LN那样以非线性方式抑制极端值。

α的初始化

研究还发现,对α(记作α₀) 的初始化进行调整,很少会显著提升模型性能。

唯一的例外是——LLM训练,在该场景下,对α₀进行精细调整可以带来明显的性能提升。

对于非LLM模型,其性能对α₀的敏感度相对较低

如图9所示,研究人员在不同任务上对α₀进行调节后,观察其对验证集性能的影响。所有实验均遵循各自任务的原始实验设置和超参数。

实验结果表明,在较宽范围的α₀取值(0.5-1.2 之间)下,模型性能基本保持稳定,通常能取得较好的结果。

此外,研究还发现,调整α₀通常仅影响训练曲线的早期阶段。

唯一的例外是监督学习ViT-L(Vision Transformer Large) 相关实验,在该场景下,当α₀超过0.6时,训练会变得不稳定并发生发散。

在这种情况下,降低学习率可以恢复训练稳定性。

较小的α₀能够提高训练稳定性

在此前观察的基础上,研究进一步分析了导致训练不稳定的因素。

研究表明,随着模型规模或学习率的增加,需要降低α₀以确保训练稳定。相反,较高的α₀需要搭配较低的学习率,以缓解训练不稳定性。

如图10展示了,在ImageNet-1K数据集上,监督学习ViT训练稳定性的消融实验。

在实验中,研究人员调整了学习率、模型规模和α₀值。结果表明,训练更大的模型更容易失败,因此需要使用更小的α₀或更低的学习率来确保稳定性。

此外,研究人员在基于LN(LayerNorm)的模型上也观察到了类似的不稳定性模式,在相同条件下,将α₀设为0.5可获得与LN类似的稳定性表现。

将α₀=0.5设为默认值

基于以上研究结果,研究人员将α₀=0.5设为所有非LLM模型的默认值。该设定能够在保证训练稳定性的同时,维持良好的性能,并与LN的稳定性模式相当。

正如前文所述,默认设置α₀=0.5在大多数任务中表现良好。然而,研究人员发现,对α₀进行调优可以显著提高LLM的性能。

他们对LLaMA系列模型进行了α₀调优,具体方法是在30B token规模的数据上进行预训练,并对比其训练损失。

表11总结了每个模型的最优α₀值。研究中,得出了两个关键结论:

较大的模型需要较小的α₀。一旦确定了较小模型的最优α₀,就可以相应缩小更大模型的搜索空间,提高参数调优的效率。

提高注意力层(attention blocks)的α₀值有助于提升性能

研究发现,在注意力块(attention blocks)中的DyT层设定较高的α₀,而在其他位置(如FFN块或最终线性投影之前)设定较低的α₀,能够提升模型性能。

为了进一步说明α₀调优的影响,图11展示了两种LLaMA模型的训练损失热图。实验表明,在注意力块(attention blocks)中使用较高的α₀,能够有效降低训练损失,从而提升模型性能。

模型宽度主要决定α₀的选择

研究人员还进一步研究了模型宽度(width)和深度(depth) 对最优α₀的影响。研究结果表明:

模型宽度(width) 在确定最优α₀方面至关重要,更宽的网络通常需要更小的α₀以实现最佳性能。模型深度(depth) 对α₀的选择影响较小,几乎可以忽略。

表12总结了不同宽度和深度下的最优α₀值,进一步证明了宽度越大的模型,适合使用更小的α₀,而深度对α₀的影响可以忽略不计。

作者介绍

Jiachen Zhu

Jiachen Zhu是纽约大学柯朗数学研究所(CIMS)的四年级计算机科学博士生,导师是Yann LeCun。

在此之前,他于2015年在香港理工大学获得学士学位,并于2020年在纽约大学获得硕士学位。

他主要对图像或视频的自监督学习,以及各类神经网络架构的设计的研究感兴趣。

Xinlei Chen(陈鑫磊)

Xinlei Chen目前是Meta FAIR的研究科学家。

他曾获得了CMU语言技术研究所的博士学位,本科毕业于中国浙江大学计算机科学专业。目前,他的研究兴趣集中在视觉表示的预训练与理解。

何恺明

何恺明目前是麻省理工学院(MIT)电气工程与计算机科学系(EECS)的副教授,于2024年2月加入。

他本人的研究重点是,构建能够从复杂世界中学习表示并发展智能的计算机模型,研究的长期目标是用更强大的AI来增强人类智能。

Yann LeCun

Yann LeCun图灵奖三巨头之一,目前任Meta首席科学家,纽约大学终身教授。

Zhuang Liu(刘壮)

刘壮是此次项目负责人,目前在Meta FAIR实验室任研究科学家。

他曾获得了UC伯克利分校的电子工程与计算机科学系(EECS) 计算机科学博士学位,导师是Trevor Darrell教授。

本科毕业于清华姚班,此外,他曾在康奈尔大学、英特尔实验室、Adobe Research、FAIR担任访问研究员或实习生。

刘壮的主要研究领域是深度学习和计算机视觉,关注深度学习模型架构、训练方法、计算效率及模型理解。

他曾主导开发了DenseNet(获得CVPR最佳论文奖)和 ConvNeXt。这两个算法皆成为深度学习和计算机视觉领域最广泛应用的神经网络架构之一。

来源:新智元

相关推荐