早起读论文,当Transformers 不再需要归一化

论文地址

[2503.10622] Transformers without NormalizationAbstract page for arXiv paper 2503.10622: Transformers without Normalizationhttps://arxiv.org/abs/2503.10622

简述

归一化层

论文题目直接说明了目的,人工智能的基础Transformers模型是不需要归一化层的

听起来十分的不可思议,听起来好像不太可能,归一化在所有的深度学习框架里都出现过,算是基石一般的存在,在 Transformer 模型中,归一化层主要是指层归一化(Layer Normalization)。它的主要功能是稳定和加速模型的训练过程。 在AI 模型训练过程里归一化就像空气一样。特别是像Transformers 这样的网红模型,归一化属于标配中的标配。在没有归一化的情况下,输入这个向量的值分布不均,可能导致在训练过程中某些神经元的激活值过高或过低,这会影响梯度的传播,导致训练不稳定,训练就会像过山车一样颠簸震荡, 根本学不到什么东西。

归一化层清洗的数据,让它变得"干净"更适合模型训练, 每次有新的网络结构,都势必会考虑要选用什么归一化层。

摘自论文

大多数归一化层共享一个通用公式。给定一个形状为 $(B, T, C)$ 的输入 $x$ ,其中 $B$ 是批次大小, $T$ 是 Token 的数量, $C$ 是每个 Token 的嵌入维度

则输出通常计算为:

 

其中 $\epsilon$ 是一个很小的常数, $\gamma$ 和 $\beta$ 是形状为 $C$, 的可学习向量参数。 它们是“缩放”和“平移”仿射参数,允许输出在任何范围内。 术语 $\mu$ 和 $\sigma^2$ 表示输入的均值和方差。 不同的方法主要区别在于如何计算这两个统计量。 这导致 $\mu$ 和 $\sigma^2$ 具有不同的维度,每个维度在计算期间应用广播。

批量归一化 (BN) 是第一个现代归一化层,主要用于 ConvNet 模型。 它的引入代表了深度学习架构设计中的一个重要里程碑。 BN 计算批次和 Token 维度上的均值和方差

其他流行的归一化层,如组归一化 (group normalization) 和实例归一化 (instance normalization),最初是为诸如目标检测和图像风格化等专门任务而提出的。它们共享相同的总体公式,但在计算统计数据的轴和范围上有所不同。

层归一化 (LN和均方根归一化 (RMSNorm)是 Transformer 框架中使用的两种主要的归一化层类型。

LN 独立地为每个样本中的每个 Token 计算这些统计数据,其中 

RMSNorm通过移除均值中心化步骤并使用以下公式归一化输入来简化 LN: 

如今,由于其简单性和通用性,大多数现代神经网络都使用 LN。最近,RMSNorm 越来越受欢迎,尤其是在诸如 T5、LaMDA、Mistral、Qwen、InternLM和 DeepSeek等大语言模型中。

我们在这项工作中验证的 Transformers 都使用 LN,除了 LaMDA 使用 RMSNorm。

研究者们通过对Transformer模型中的层归一化(Layer Normalization, LN)进行深入分析,发现LN的输出与输入之间呈现出一种类似于tanh函数的S形曲线。

喔,这个发现真是太关键了,想想看如果归一化层最后干的活和Tanh函数差不多,那为什么不直接用Tanh 呢,但是普通的Tanh不够灵活,所以论文团队搞出了一个升级版的Tanh函数, 动态Tanh(DyT)。 这是一种简单的逐元素操作,能够在不计算激活统计量的情况下,模拟LN的行为。通过实验,他们证明了使用DyT的Transformer模型在多种任务中表现优异,甚至在某些情况下超越了传统的归一化模型。

DyT 不是简单的 tanh(ax), 它像归一化层一样能够缩放和偏移输入, 但是不用计算均值和方差

从而减少了计算开销,回时保密了归一化层的压缩极端值的特性。厉害的是这个操作是独立的,每个数据都是单独出理的,不像归一化层还需要计算一堆统计数据, 把群体性的统计工作变成了个人感受上的调整(所以说经济学家别老拿统计数据说事,我有没有钱我自己最清楚),效率是蹭蹭往上涨,而且DyT层代码几行就实现了

代码实现

class DynamicTanh(nn.Module):
    def __init__(self, normalized_shape, alpha_init_value=0.5):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.alpha_init_value = alpha_init_value
        self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value)
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))

    def forward(self, x):
        x = torch.tanh(self.alpha * x)
        return x * self.weight[:, None, None] + self.bias[:, None, None]

就像一个即插即用的模块,想往哪里塞就往哪里塞,模型原有使用归一化层的地方直接换成DyT

def convert_ln_to_dyt(module):
    module_output = module
    if isinstance(module, nn.LayerNorm):
        module_output = DynamicTanh(module.normalized_shape, not isinstance(module, LayerNorm2d))
    for name, child in module.named_children():
        module_output.add_module(name, convert_ln_to_dyt(child))
    del module
    return module_output

用最简单的方式破坏原有的根基,这不就是降维打击么,撼动机器学习最理所当然的东西

研究团队通过实验证明,具有DyT的模型可以稳定地进行训练,并在各种设置中实现较高的最终性能。它通常不需要需要在原始架构上调整训练超参数。我们的工作挑战了归一化层对于训练现代神经网络是不可或缺的这一概念,并提供了对归一化层属性的经验性见解。此外,初步测量表明,DyT 提高了训练和推理速度,使其成为面向效率的网络设计的候选方案。

论文里的试验结果

下图是试验者在将原始transformers的归一化层替换为DyT,  在实际代修改上DyT的引入非常简单,只需将现有的归一化层替换为DyT层即可。

为了验证DyT的有效性,研究者在多个任务和领域中进行了广泛的实验,包括图像分类、自监督学习、扩散模型、大语言模型、语音自监督学习和DNA序列建模。

实验团队为了证明 DyT 的有效性,使用 Transformer 和一些其他现代架构,在各种任务和领域中进行了实验。在每个实验中,用 DyT 层替换原始架构中的 LN 或 RMSNorm,并遵循官方开源协议来训练和测试这两个版本的模型。

值得注意的是,为了突出 DyT 适应的简易性,团队使用了与标准化对应模型相同的超参数。在 ImageNet-1K 分类任务上训练了“Base”和“Large”尺寸的 Vision Transformer (ViT) 和 ConvNeXt 。选择这些模型是因为它们的受欢迎程度和不同的操作:ViT 中的注意力机制和 ConvNeXt 中的卷积。表 1 报告了 top-1 分类准确率。在两种架构和模型尺寸上,DyT 的性能略优于 LN。我们还在图 5 中绘制了 ViT-B 和 ConvNeXt-B 的训练损失。曲线表明,基于 DyT 和 LN 的模型的收敛行为高度一致。

在ViT 和ConvNeXt模型里,使用DyT和LN 效果基本一致

研究团队也怕出错,做了超大规模的试验,简直就是要把AI界所有的"网红"模型都试个遍!

试验方式也简单粗暴,直接把原来的Layer Norm或RMSNorm换成DyT,结果你猜怎么着,

换上DyT 非但没有掉链子,反而在很多任务上表现得更好了

这些提升基本都是白给的,都不需要去寻找和修改训练参数,就好像把家里白炽灯换成LED一样

不用改电路,不用换开关但是灯更亮了,电费也更低了

再瞅一眼这个公式

DyT里面那个可学习的缔放因子alpha,就像是一个"智能调光器”,能根据不同的情况自动调节曲线的陡峭程度,这种灵活带来了意向不到的好处,它不仅保持数据流动的通畅,还让模型能更精细地控制激活值的变化,就像是给模型装了一个"智能变速箱“,能在不同"路况"下自动找到最佳的"档位"

想想看这个发现会对AI训练产生多么深远的影响,以前调参是不是特别麻烦,学习率要调整,batch size要调,各种参数调的人头皮发麻,但是DyT说, ”别调了,我天生就稳定,基本不用调参数就能玩的转”, 进一步我们是否可以悟出来归一化层的核心价值可能不在于调整数据分布,而是在于它产生的那种特殊的非线性变换。

自己上手试验

我打算在我之前的用transformer 天气预报的项目上添加DyT 模块

GitHub - chenrui2200/weather_timeseries_predictContribute to chenrui2200/weather_timeseries_predict development by creating an account on GitHub.https://github.com/chenrui2200/weather_timeseries_predict改造起来也是非常方便,定义 dynamic_tanh.py

import torch
import torch.nn as nn
from timm.layers import LayerNorm2d


class DynamicTanh(nn.Module):
    def __init__(self, normalized_shape, channels_last, alpha_init_value=0.5):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.alpha_init_value = alpha_init_value
        self.channels_last = channels_last

        self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value)
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))

    def forward(self, x):
        x = torch.tanh(self.alpha * x)
        if self.channels_last:
            x = x * self.weight + self.bias
        else:
            x = x * self.weight[:, None, None] + self.bias[:, None, None]
        return x

    def extra_repr(self):
        return f"normalized_shape={self.normalized_shape}, alpha_init_value={self.alpha_init_value}, channels_last={self.channels_last}"


def convert_ln_to_dyt(module):
    module_output = module
    if isinstance(module, nn.LayerNorm):
        module_output = DynamicTanh(module.normalized_shape, not isinstance(module, nn.LayerNorm))
    for name, child in module.named_children():
        module_output.add_module(name, convert_ln_to_dyt(child))
    del module
    return module_output

模型参数上增加一个channels_last, DynamicTanh 实现中,channels_last 是一个布尔参数,它的作用是控制输入张量 x 的维度布局格式,从而决定如何应用 weight 和 bias 参数。具体来说,它处理了两种常见的张量格式:

  1. Channels-First 格式(默认,channels_last=False)
    • 在这种情况下,输入张量的形状通常是 (batch_size, channels, height, width),即通道维度在第1维(紧跟批量维度之后)。
    • 当 channels_last=False 时,weight 和 bias 的形状是 (normalized_shape,),它们需要通过广播机制与输入张量 x 的空间维度(height 和 width)对齐。
    • 因此,在 forward 函数中,x * self.weight[:, None, None] + self.bias[:, None, None] 会将 weight 和 bias 的形状从 (normalized_shape,) 扩展为 (normalized_shape, 1, 1),以匹配输入张量的空间维度并进行逐元素操作。
  2. Channels-Last 格式(channels_last=True)
    • 在这种情况下,输入张量的形状通常是 (batch_size, height, width, channels),即通道维度在最后一维。
    • 当 channels_last=True 时,weight 和 bias 的形状已经是 (normalized_shape,),可以直接与输入张量的最后一个维度(通道维度)对齐,无需额外的广播操作。
    • 因此,在 forward 函数中,直接使用 x * self.weight + self.bias,因为 weight 和 bias 的形状与通道维度天然匹配。

好处是

  • 灵活性:它允许 DynamicTanh 模块兼容不同的输入张量格式(channels-first 或 channels-last),这在深度学习中非常有用,因为不同的模型或框架(如 PyTorch 或 TensorFlow)可能采用不同的默认格式。
  • 广播控制:通过 channels_last 的值,模块能够正确调整 weight 和 bias 的应用方式,确保与输入张量的维度匹配,避免形状不兼容的问题。
  • 与归一化层的兼容性:在你的代码中,convert_ln_to_dyt 函数会根据原始模块是否为 LayerNorm2d 来设置 channels_last,从而保持与原始归一化层行为的一致性(LayerNorm2d 通常用于 channels-first 格式,而 nn.LayerNorm 更通用)。

修改模型结构

class Transformer(nn.Module):
...
        self.encoder = Encoder(
            n_src_vocab=n_src_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            dropout=dropout)

        if self.dynamic_tanh:
            self.encoder = convert_ln_to_dyt(self.encoder)

        self.decoder = Decoder(
            n_trg_vocab=n_trg_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, pad_idx=1, n_head=n_head, d_k=d_k, d_v=d_v,
            dropout=dropout)

        if self.dynamic_tanh:
            self.decoder = convert_ln_to_dyt(self.decoder)

如果开启了dynamic_tanh  调用convert_ln_to_dyt 把LayerNorm 替换为DynamicTanh

def convert_ln_to_dyt(module):
    module_output = module
    if isinstance(module, nn.LayerNorm):
        module_output = DynamicTanh(module.normalized_shape, not isinstance(module, LayerNorm2d))
    for name, child in module.named_children():
        module_output.add_module(name, convert_ln_to_dyt(child))
    del module
    return module_output

训练使用模型

parser.add_argument('-dynamic_tanh', type=bool, default=True)
...

model_train = Transformer(
        500,
        500,
        d_k=opt.d_k,
        d_v=opt.d_v,
        d_model=opt.d_model,
        d_word_vec=opt.d_word_vec,
        d_inner=opt.d_inner_hid,
        n_layers=opt.n_layers,
        n_head=opt.n_head,
        dropout=opt.dropout,
        dynamic_tanh=opt.dynamic_tanh
    ).to(device)


再把参数调回  --dynamic_tanh  false 训练下

最后查看loss和rl 比较,绿线是使用nn.LayerNorm 训练的, 黑线是使用DynamicTanh 训练的

loss曲线基本保持都能快速收敛,rl 也重合

回过头再看看DyT论文里的其他试验结果,都显示着DyT 性能上的提升

不过 DyT 也不是免费的午餐

从论文实验结果看,DyT在Transformer架构中确实表现出与LN相当的甚至更优的性能,且无需复杂统计计算,看似“免费”。但这种优势存在明确的边界条件。例如,在视觉Transformer(ViT)和语言模型(LLaMA)中,DyT可直接替换LN且无需调参,准确率与训练稳定性均不受影响;然而,在经典CNN(如ResNet)中替换BN时性能显著下降(ImageNet准确率下降约7%),说明其并非通用解决方案。

不过进一步联想, 当Transformer在某些领域任务上都不需要层归一化层了, 是不是社会上很多看着很重要的岗位其实都都没有想象中那么重要,世界真是个巨大的草台班子

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

CHEN_RUI_2200

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值