何恺明又双叒叕发新作了，这次还是与图灵奖得主 Yann LeCun 合作。
这项研究的主题是没有归一化层的 Transformer（Transformers without Normalization），并已被 CVPR 2025 会议接收。
Meta FAIR 研究科学家刘壮的推文
过去十年，归一化层已经巩固了其作为现代神经网络最基本组件之一的地位。这一切可以追溯到 2015 年批归一化（batch normalization）的发明，它使视觉识别模型的收敛速度变得更快、更好，并在随后几年中获得迅速发展。从那时起，研究人员针对不同的网络架构或领域提出了许多归一化层的变体。
如今，几乎所有现代网络都在使用归一化层，其中层归一化（Layer Norm，LN）是最受欢迎之一，特别是在占主导地位的 Transformer 架构中。
归一化层的广泛应用很大程度上得益于它们在优化方面的实证优势。除了实现更好的结果之外，归一化层还有助于加速和稳定收敛。随着神经网络变得越来越宽、越来越深，归一化层的必要性变得越来越重要。因此，研究人员普遍认为归一化层对于有效训练深度网络至关重要，甚至是必不可少的。这一观点事实上得到了微妙证明：近年来，新架构经常寻求取代注意力层或卷积层，但几乎总是保留归一化层。
本文中，研究者提出了 Transformer 中归一化层的一种简单平替。他们的探索始于以下观察：LN 层使用类 tanh 的 S 形曲线将其输入映射到输出，同时缩放输入激活并压缩极值。
受此启发，研究者提出了一种元素级运算，称为 Dynamic Tanh（DyT），定义为：DyT (x) = tanh (αx)，其中 α 是一个可学习参数。此运算旨在通过 α 学习适当的缩放因子并通过有界 tanh 函数压缩极值来模拟 LN 的行为。值得注意的是，与归一化层不同，DyT 可以实现这两种效果，而无需计算激活数据。
论文一作 Jiachen Zhu 为纽约大学四年级博士生、二作陈鑫磊（Xinlei Chen）为 FAIR 研究科学家，项目负责人为刘壮。
论文标题：Transformers without Normalization
论文地址：https://arxiv.org/pdf/2503.10622
项目主页：https://jiachenzhu.github.io/DyT/
GitHub 地址：https://github.com/jiachenzhu/DyT
DyT 使用起来非常简单，如下图 1 所示，研究者直接用 DyT 替换视觉和语言 Transformer 等架构中的现有归一化层。实证结果表明，使用 DyT 的模型可以在各种设置中稳定训练并获得较高的最终性能。同时，DyT 通常不需要在原始架构上调整训练超参数。
DyT 模块可以通过短短几行 PyTorch 代码来实现。
该工作挑战了「归一化层对训练现代神经网络必不可少」这一观念，并提供了有关归一化层属性的实证见解。此外，初步结果表明，DyT 可以提升训练和推理速度，从而成为以效率为导向的网络设计的候选方案。
刘壮发推称，对他而言，归一化层一直是深度学习中比较神秘的内容。这项工作让他对归一化层的作用有了更深的理解。另外，考虑到模型训练和推理需要数千万的算力需求，DyT 有潜力助力成本降低。他很期待接下来 DyT 的应用。
归一化层有什么作用？
要去掉 Transformer 中的归一化层，首先要做的当然是了解归一化层有什么用。
该团队通过实证研究对此进行了分析。为此，他们使用了三个不同的经过训练的 Transformer 模型：一个 Vision Transformer（ViT-B）、一个 wav2vec 2.0 Large Transformer 和一个 Diffusion Transformer（DiT-XL）。
他们使用这三个模型采样了一小批样本，并让其前向通过整个网络。然后，他们监测了其中归一化层的输入和输出，即归一化操作前后的张量。
由于 LN 会保留输入张量的维度，因此可以在输入和输出张量元素之间建立一一对应关系，从而可以直接可视化它们的关系。这个映射关系见图 2。
具有层归一化的类 tanh 映射。对于这三个模型，该团队发现，它们的早期 LN 层（图 2 第 1 列）的输入 - 输出关系基本上是线性的。但是，更深的 LN 层却有更有趣的表现。
可以观察到，这些曲线的形状大多与 tanh 函数表示的完整或部分 S 形曲线非常相似（见图 3）。
人们可能预期 LN 层会对输入张量进行线性变换，因为减去平均值和除以标准差都是线性运算。LN 以每个 token 的方式进行归一化，仅对每个 token 的激活进行线性变换。
由于 token 具有不同的平均值和标准差，因此这种线性并不对输入张量的所有激活都成立。尽管如此，该团队表示依然很惊讶：实际的非线性变换竟然与某个经过缩放的 tanh 函数高度相似！
对于这样一个 S 型曲线，可以看到其中心部分（x 值接近零的部分）仍然主要呈线性形状。大多数点（约 99%）都属于这个线性范围。但是，仍有许多点明显超出此范围，这些点被认为具有「极端」值，例如 ViT 模型中 x 大于 50 或小于 -50 的点。
归一化层对这些值的主要作用是将它们压缩为不太极端的值，从而与大多数点更加一致。这是归一化层无法通过简单的仿射变换层近似的地方。
该团队假设，这种对极端值的非线性和不成比例的压缩效应正是归一化层的关键之处。
前段时间的一篇论文《On the Nonlinearity of Layer Normalization》同样重点指出了 LN 层引入的强非线性，并且表明这种非线性可以增强模型的表征能力。
此外，这种压缩行为还反映了生物神经元对大输入的饱和（saturation）特性，这种现象大约一个世纪前就已经被观察到。
token 和通道的归一化。LN 层如何对每个 token 执行线性变换，同时以这种非线性方式压缩极端值呢？
为了理解这一点，该团队分别按 token 和通道对这些点进行可视化。图 4 给出了 ViT 的第二和第三个子图的情况，但为了更清晰，图中使用了采样的点子集。
在图 4 左边两个小图中，使用了同一颜色标记每个 token 的激活。可以观察到，任何单个 token 的所有点确实都会形成一条直线。但是，由于每个 token 都有不同的方差，因此斜率也不同。输入 x 范围较小的 token 往往具有较小的方差，并且归一化层将使用较小的标准偏差来除它们的激活，从而让直线有较大的斜率。
总的来说，它们形成了一条类似于 tanh 函数的 S 形曲线。在右侧的两个小图中，同样使用相同的颜色标记各个通道的激活。可以看到，不同通道的输入范围往往存在巨大差异，只有少数通道（例如红色、绿色和粉色）会表现出较大的极端值 —— 而这些通道正是被归一化层压缩得最厉害的。
Dynamic Tanh（DyT）
既知根本，正当创新。基于归一化层和扩展版 tanh 函数的相似性，该团队提出了 Dynamic Tanh（DyT），并且这可以作为归一化层的直接替代。
给定一个输入张量 x，DyT 层的定义如下：
其中 α 是一个可学习的标量参数，允许根据输入的范围以不同的方式缩放输入，并会考虑不同的 x 尺度（图 2）。也因此，他们将整个操作命名为「动态」tanh。γ 和 β 是可学习的每通道向量参数，与所有归一化层中使用的参数相同 —— 它们允许输出缩放到任何尺度。这有时被视为单独的仿射层；这里，该团队将它们视为 DyT 层的一部分，就像归一化层也包括它们一样。算法 1 给出了用类 PyTorch 的伪代码实现的 DyT。
要想将 DyT 层集成到现有架构中，方法很简单：直接用一个 DyT 层替换一个归一化层（见图 1）。这适用于注意力块、FFN 块和最终归一化层内的归一化层。
尽管 DyT 可能看起来像或可被视为激活函数，但本研究仅使用它来替换归一化层，而不会改变原始架构中激活函数的任何部分，例如 GELU 或 ReLU。网络的其他部分也保持不变。该团队还观察到，几乎不需要调整原始架构使用的超参数即可使 DyT 表现良好。
尽管 DyT 可能看起来像或可被视为激活函数，但本研究仅使用它来替换归一化层，而不会改变原始架构中激活函数的任何部分，例如 GELU 或 ReLU。网络的其他部分也保持不变。该团队还观察到，几乎不需要调整原始架构使用的超参数即可使 DyT 表现良好。
关于缩放参数。在这里，总是简单地将 γ 初始化为全一向量，将 β 初始化为全零向量，后接归一化层。对于 scaler 参数 α，除了 LLM 训练外，默认初始化为 0.5 通常就足够了。除非另有明确说明，否则在后续的实验中，α 均被初始化为 0.5。
说明。DyT 并非一种新型的归一化层，因为它在前向传递过程中会独立地对张量中的每个输入元素进行操作，而无需计算统计数据或其他类型的聚合。但它确实保留了归一化层的效果，即以非线性方式压缩极端值，同时对输入的中心部分执行近乎线性的变换。
DyT 在实验中的表现
为了验证 DyT 的效果，研究团队在不同任务和领域中测试了 Transformer 及其他架构，将原始架构中的 LN 或 RMSNorm 替换为 DyT 层，并按照官方开源方案进行训练和测试。
视觉监督学习
研究团队在 ImageNet-1K 分类任务上训练了 Base 和 Large 两种规模的 Vision Transformer（ViT）和 ConvNeXt 模型。
选择 ViT 和 ConvNeXt 是因为它们既具代表性，又分别采用不同机制：ViT 基于注意力机制，ConvNeXt 基于卷积操作。从表 1 的 Top-1 分类准确率来看，DyT 在两种架构和不同规模模型上均优于 LN。图 5 中展示的 ViT-B 和 ConvNeXt-B 的训练损失曲线。
表 1：ImageNet-1K 上的监督分类准确率。DyT 在两种架构和不同模型规模上均实现了优于或等同于 LN 的性能表现。
视觉自监督学习
研究团队测试了两种流行的视觉自监督学习方法：何恺明的 MAE 和 DINO。
这两种方法都默认使用 Vision Transformer 作为骨干网络，但训练目标不同。MAE 使用重建损失进行训练，而 DINO 则使用联合嵌入损失。研究团队先在 ImageNet-1K 数据集上进行无标签预训练，然后添加分类层并用标签数据微调来测试预训练模型。表 2 展示了微调的结果。在自监督学习任务中，DyT 和 LN 的表现基本持平。
扩散模型
研究者在 ImageNet-1K 上训练了三个尺寸分别为 B、L 和 XL 的 DiT 模型。需要注意的是，在 DiT 中，LN 层的仿射参数用于类调节，DyT 实验中也保留了这一参数，只是用 tanh (αx) 函数替换了归一化迁移。训练结束，如表 3 所示，与 LN 相比，DyT 的 FID 值相当或有所提高。
LLM
这些模型是按照 LLaMA 中概述的原始配方在带有 200B tokens 的 The Pile 数据集上进行训练的。在带有 DyT 的 LLaMA 中，研究者在初始嵌入层之后添加了一个可学习的标量参数，并调整了 α 的初始值（第 7 节）。下表 4 报告了训练后的损失值，并按照 OpenLLaMA 的方法，在 lm-eval 的 15 个零样本任务上对模型进行了基准测试。如表 4 所示，在所有四种规模的模型中，DyT 的表现与 RMSNorm 相当。
图 6 展示了损失曲线，显示了所有模型大小的相似趋势，训练损失在整个训练过程中都非常接近。
语音自监督学习。研究者在 LibriSpeech 数据集上预训练了两个 wav2vec 2.0 Transformer 模型。表 5 报告了最终的验证损失。在两种模型规模下，DyT 的表现都与 LN 相当。
DNA 序列建模
在长程 DNA 序列建模任务中，研究者对 HyenaDNA 模型和 Caduceus 模型进行了预训练。结果如表 6，在这项任务中，DyT 保持了与 LN 相当的性能。
α 初始化
非 LLM 模型的 α 初始化
非 LLM 模型对 α_0 相对不敏感。图 9 展示了在不同任务中改变 α_0 对验证性能的影响。
α_0 越小，训练越稳定。图 10 展示了使用 ImageNet-1K 数据集对有监督 ViT 训练稳定性的消减。
将 α_0 = 0.5 设为默认值。根据研究结果，研究者将 α_0 = 0.5 设置为所有非 LLM 模型的默认值。这种设置既能提供与 LN 相当的训练稳定性，又能保持强大的性能。
LLM 模型的 α 初始化
调整 α_0 可以提高 LLM 性能。如前所述，默认设置 α_0 = 0.5 在大多数任务中表现良好。然而，研究者发现调整 α_0 可以大幅提高 LLM 性能。他们对每个 LLaMA 模型都进行了 30B tokens 的预训练，并比较了它们的训练损失，从而调整了它们的 α_0。
表 11 总结了每个模型的调整后 α_0 值，其中有两个重要发现：
1. 较大的模型需要较小的 α_0 值。一旦确定了较小模型的最佳 α_0 值，就可以相应地缩小较大模型的搜索空间；
2. 注意力块的 α_0 值越高，性能越好。对注意力块中的 DyT 层初始化较高的 α 值，而对其他位置（即 FFN 区块内或最终线性投影之前）的 DyT 层初始化较低的 α 值，可以提高性能。
为了进一步说明 α_0 调整的影响，图 11 展示了两个 LLaMA 模型损失值的热图。这两个模型都受益于注意力块中较高的 α_0，从而减少了训练损失。
模型宽度主要决定了 α_0 的选择。我们还研究了模型宽度和深度对最优 α_0 的影响。研究者发现，模型宽度对确定最优 α_0 至关重要，而模型深度的影响则微乎其微。表 12 显示了不同宽度和深度下的最佳 α_0 值，表明较宽的网络可以从较小的 α_0 值中获益，从而获得最佳性能。另一方面，模型深度对 α_0 的选择影响微乎其微。
从表 12 中可以看出，网络越宽，「注意力」和「其他」所需的初始化就越不均衡。研究者假设，LLM 的 α 初始化的敏感度与其他模型相比过大的宽度有关。
更多研究细节，可参考原论文。
