【杨乐昆何凯明AI论文】没有归一化的Transformer模型

摘要：归一化层在现代神经网络中无处不在，长期以来一直被视为不可或缺的组成部分。本研究表明，通过使用一种极其简单的技术，没有归一化的Transformer模型可以达到相同或更好的性能。我们引入了动态Tanh（DyT），这是一种逐元素操作，表示为DyT(x) = tanh(alpha x)，作为Transformer中归一化层的即插即用替代品。DyT的灵感来源于这样一个观察：Transformer中的层归一化往往会产生类似tanh的S形输入输出映射。通过融入DyT，没有归一化的Transformer模型可以匹敌或超越其带归一化版本的性能，而且大多数情况下无需调整超参数。我们在多种设置下验证了采用DyT的Transformer的有效性，包括从识别到生成、从监督学习到自监督学习，以及从计算机视觉到语言模型等各个领域。这些发现挑战了归一化层在现代神经网络中不可或缺的传统观念，并为深入理解它们在深度网络中的作用提供了新的见解。Huggingface链接：Paper page，论文链接：2503.10622

研究背景和目的

研究背景

归一化层在现代神经网络中占据核心地位，自2015年批归一化（Batch Normalization, BN）的提出以来，归一化层在加速模型收敛、提高模型性能以及稳定训练过程等方面发挥了重要作用。此后，针对不同网络架构和领域，各种归一化层变体被相继提出，如层归一化（Layer Normalization, LN）、实例归一化（Instance Normalization, IN）和组归一化（Group Normalization, GN）等。特别是层归一化（LN），在Transformer架构中尤为流行，成为训练深度Transformer模型的标准组成部分。

尽管归一化层在优化和泛化能力上的重要性被广泛认可，但近年来的一些研究开始探索在没有归一化层的情况下训练深度神经网络的可能性。这些研究通过调整权重初始化、使用权重归一化或其他正则化技术，成功地在没有归一化层的情况下训练了深度卷积神经网络（ConvNets）。然而，这些方法往往需要对模型架构或训练策略进行显著修改，且效果并不总是令人满意。

对于Transformer模型而言，尽管它们在自然语言处理（NLP）和计算机视觉等领域取得了巨大成功，但归一化层在Transformer中的作用及其不可替代性仍未得到充分探讨。因此，本研究旨在探索一种简单有效的方法，使Transformer模型能够在没有归一化层的情况下达到或超过其带归一化版本的性能。

研究目的

本研究的主要目的是提出一种名为动态Tanh（Dynamic Tanh, DyT）的逐元素操作，作为Transformer中归一化层的替代品。通过引入DyT，我们旨在实现以下目标：

1. 性能匹配或超越：证明没有归一化层的Transformer模型在使用DyT后，能够在各种任务上达到或超过其带归一化版本的性能。
2. 简化模型架构：提供一种即插即用的解决方案，无需对Transformer的架构进行显著修改。
3. 减少超参数调优：证明使用DyT的Transformer模型在大多数情况下无需进行额外的超参数调优。
4. 深入理解归一化层：通过对比分析，提供对归一化层在深度网络中作用的新见解。

研究方法

DyT的定义与实现

动态Tanh（DyT）是一种逐元素操作，定义为DyT(x) = tanh(alpha x)，其中alpha是一个可学习的标量参数。DyT的灵感来源于观察到Transformer中的层归一化（LN）往往会产生类似tanh的S形输入输出映射。通过引入alpha参数，DyT能够学习一个适当的缩放因子，并通过有界的tanh函数对极端值进行压缩。

DyT的实现非常简单，只需将Transformer中的归一化层替换为DyT层即可。在具体实现中，我们保留了归一化层中的可学习仿射变换参数gamma和beta，使DyT层能够像归一化层一样对输出进行缩放和平移。

实验设置

为了验证DyT的有效性，我们在多种任务和领域上进行了广泛的实验，包括图像分类、自监督学习、扩散模型、大型语言模型（LLMs）以及语音和DNA序列建模等。对于每个任务，我们都使用相同的训练协议和模型架构，仅将归一化层替换为DyT层，并尽可能保持其他超参数不变。

在图像分类任务中，我们使用了Vision Transformer（ViT）和ConvNeXt模型，并在ImageNet-1K数据集上进行了训练。对于自监督学习任务，我们测试了掩码自编码器（MAE）和DINO方法，并使用了ViT作为骨干网络。在扩散模型方面，我们训练了Diffusion Transformer（DiT）模型，并评估了其在ImageNet上的图像生成质量。对于大型语言模型（LLMs），我们预训练了LLaMA模型，并在多个零样本任务上评估了其性能。此外，我们还测试了DyT在语音（wav2vec 2.0）和DNA序列建模（HyenaDNA和Caduceus）任务上的效果。

研究结果

性能匹配或超越

我们的实验结果表明，使用DyT的Transformer模型在各种任务上均能达到或超过其带归一化版本的性能。在图像分类任务中，DyT在ViT和ConvNeXt模型上均取得了与LN相当或略高的准确率。在自监督学习任务中，DyT也表现出了与LN相当的性能。在扩散模型方面，DyT在DiT模型上实现了与LN相当的或更优的FID分数。对于大型语言模型（LLMs），DyT在LLaMA模型上的训练损失和零样本任务性能均与RMSNorm相当。此外，在语音和DNA序列建模任务上，DyT也取得了与LN相当的性能。

稳定性与收敛性

我们还发现，使用DyT的Transformer模型在训练过程中表现出良好的稳定性和收敛性。在大多数任务上，DyT模型的训练损失曲线与LN模型高度一致，表明两者可能具有相似的学习动态。这一发现进一步证明了DyT作为归一化层替代品的有效性。

计算效率

初步测量表明，DyT层在推理和训练过程中均比RMSNorm层具有更低的计算开销。这一发现使得DyT成为效率导向型网络设计的一个有前景的选择。

α参数的作用

通过分析α参数在训练过程中的变化，我们发现α能够学习到一个近似于输入激活标准差倒数（1/std）的值。这一发现支持了α在维持激活值在合适范围内的重要作用，从而有助于稳定而有效的训练。

研究局限

尽管我们的研究取得了显著成果，但仍存在一些局限性。首先，我们的实验主要集中在使用LN或RMSNorm的Transformer模型上，对于其他类型的归一化层（如BN、IN和GN）的适用性尚待进一步研究。初步实验表明，DyT在直接替换经典卷积神经网络（如ResNet）中的BN层时效果不佳。这可能是由于BN层在这些网络中出现得更为频繁，而LN在Transformer中通常每几个权重层才出现一次。

其次，尽管DyT在大多数情况下无需进行额外的超参数调优，但在训练大型语言模型（LLMs）时，我们发现对α的初始化进行微调可以显著提高性能。这表明对于某些复杂任务，可能需要对DyT的参数进行更细致的调整。

未来研究方向

针对上述局限性，我们提出以下未来研究方向：

1. 探索DyT在其他归一化层上的适用性：进一步研究DyT在替换其他类型归一化层（如BN、IN和GN）时的效果，并探索是否需要对其进行修改以适应不同类型的网络架构。

2. 优化DyT的参数初始化：针对不同类型的任务和模型架构，研究更优化的α初始化策略，以减少对超参数调优的依赖。

3. 结合其他正则化技术：探索将DyT与其他正则化技术（如dropout、weight decay等）结合使用的可能性，以进一步提高模型的性能和泛化能力。

4. 分析DyT的泛化性能：在更多数据集和任务上测试DyT的泛化性能，以验证其在不同领域和场景下的有效性。

5. 理解DyT的工作机制：通过更深入的理论分析和实验验证，探讨DyT如何模拟归一化层的作用，并理解其在深度网络中的具体工作机制。

总之，本研究为Transformer模型提供了一种简单而有效的归一化层替代品，挑战了归一化层在现代神经网络中不可或缺的传统观念。未来的研究将进一步探索DyT的潜在应用和优化策略，以推动深度学习领域的发展。