探秘Transformer系列之（14）--- 残差网络和归一化

[探秘Transformer系列之（14）--- 残差网络和归一化](#探秘Transformer系列之（14）--- 残差网络和归一化)
- [0x00 概述](#0x00 概述)
- [0x01 残差连接](#0x01 残差连接)
  - [1.1 问题](#1.1 问题)
  - [1.2 相关知识](#1.2 相关知识)
    - [shortcut connections](#shortcut connections)
    - 恒等映射
  - [1.3 网络结构](#1.3 网络结构)
    - 论文V1
    - 论文V2
  - [1.4 功用](#1.4 功用)
    - 梯度消失
    - 缓解退化
    - 层间修正
    - [掩码 VS 残差](#掩码 VS 残差)
- [0x02 归一化](#0x02 归一化)
  - [2.1 问题](#2.1 问题)
  - [2.2 定义](#2.2 定义)
  - [2.3 类型](#2.3 类型)
- [0x03 BatchNorm](#0x03 BatchNorm)
  - [3.1 公式](#3.1 公式)
  - [3.2 作用](#3.2 作用)
  - [3.3 PyTorch使用](#3.3 PyTorch使用)
  - [3.4 问题](#3.4 问题)
- [0x04 layerNorm](#0x04 layerNorm)
  - [4.1 解决方案](#4.1 解决方案)
  - [4.2 公式](#4.2 公式)
  - [4.3 作用](#4.3 作用)
  - [4.4 LN和BN的差异](#4.4 LN和BN的差异)
    - 作用对象
    - 作用方向
    - 业务选择
      - CV
      - NLP
    - 具体实现
  - [4.5 Post-Norm VS Pre-Norm](#4.5 Post-Norm VS Pre-Norm)
    - 概念
    - 论文实现
    - Post-Norm
      - 难以训练
      - 需要热身
    - Pre-Norm
    - 小结
- [0x05 扩展比对](#0x05 扩展比对)
  - [5.1 Instance Norm](#5.1 Instance Norm)
  - [5.2 GroupNorm](#5.2 GroupNorm)
  - [5.3 比对](#5.3 比对)
    - 类比
    - 细节
- [0x06 实现](#0x06 实现)
  - [6.1 LayerNorm](#6.1 LayerNorm)
  - [6.2 残差](#6.2 残差)
- [0x07 优化与演进](#0x07 优化与演进)
  - [7.1 RMSNorm](#7.1 RMSNorm)
  - [7.2 Deep Norm](#7.2 Deep Norm)
  - [7.3 PRepBN](#7.3 PRepBN)
  - [7.4 RealFormer](#7.4 RealFormer)
  - [7.5 nGPT](#7.5 nGPT)
    - 研究背景&动机
    - 核心贡献
    - 对比
      - [层和块（LAYERS AND BLOCKS）](#层和块（LAYERS AND BLOCKS）)
      - 自注意力块
      - [MLP 块（MLP BLOCK）](#MLP 块（MLP BLOCK）)
    - 切换
  - [7.6 DyT](#7.6 DyT)
    - 动机
    - 实现
- [0xFF 参考](#0xFF 参考)

0x00 概述

可以把 Transformer 看作是一个通用的可微计算机，自注意力机制给这个计算机配置了类似消息传递的架构，提供了通用且强大，能够涵盖众多现实世界的算法，仅需少量计算步骤即可完成任务；残差连接、层归一化等设计，使其具备高效的学习和优化能力。在Transformer架构中，Add & Norm模块用于在多头自注意力机制和前馈神经网络之间添加残差连接和归一化操作。具体由Add和Norm两部分组成。

Add指X+MultiHeadAttention(X)操作（残差连接），其来源于论文Deep Residual Learning for Image Recognition。残差连接把网络的输入和输出相加，得到网络的输出为F(x)+x。即将本层的输入张量与输出张量相加，从而得到一个残差向量。
Norm是Layer Normalization操作（层归一化），其来源于论文Layer Normalization。归一化对残差向量进行归一化，以便更好地传递信息和控制梯度，避免梯度消失或爆炸的问题，从而提高模型的训练效率和性能。归一化可以采用不同的方法，如Layer Normalization或Batch Normalization。Transformer采用的是层归一化。

残差连接和归一化层的结合使得Transformer能够更有效地进行训练，并提高了模型的性能和稳定性。两者的结合是现代深度学习架构成功的重要因素之一。我们接下来通过几个问题来引导，在学习完本篇之后，我们会对这几个问题做出回答。

Add & Norm在Transformer的位置是哪里？
为什么transformer块使用LayerNorm而不是BatchNorm？
LayerNorm究竟应该放前面还是后面。

0x01 残差连接

我们首先给出将Transformer LM展开后，具有注意力和前馈网络块的扩展视图，包括模型权重（灰色）和残差流状态（绿色）。从中可以看出来残差连接在Transformer中的重要性。

1.1 问题

根据泛逼近定理（universal approximation theorem），只要给定足够的容量，单层的前馈网络也足以表示任何函数。但是，该层可能非常庞大，网络和数据易出现过拟合。因此，研究界普遍认为网络架构需要更多层。理论上，神经网络的层数越多，越能学习到复杂的特征，越能提升模型的性能。但是，事实上，当层数增加到一定程度之后，模型的性能会急剧下降。导致上述情形的主要是三个问题：

梯度消失（Vanishing Gradient）：梯度消失是指在训练深度神经网络时，反向传播过程中的梯度逐渐变小，最终导致梯度接近于零的现象。梯度消失是在反向传播过程中使用的链式法则引起的。因为目前训练主要用的是基于梯度的优化器，梯度消失意味着我们没有很好的信号去调整优化前面的层。换句话说，前面的层也许几乎没有得到更新，一直保持随机初始化的状态。只有比较靠近输出的层才更新得比较好。但在前向传播中，这些更新较好层的输入是前面没有更新好的层的输出，所以更新较好层的输入质量可能会很糟糕（因为经过了一个近乎随机的变换），因此哪怕后面的层更新好了，总体效果也不好。最终的现象是：模型越深，效果越差。这就是权重退化，即在模型层数加深后新加入的层就学习不到什么东西了。
梯度爆炸（gradient exploding）：梯度爆炸指的是在反向传播过程中，梯度值呈指数级增长，导致网络权重的大幅更新，从而使得网络变得不稳定，训练效果下降。在深层神经网络中，由于链式法则的应用，梯度需要通过多层进行反向传播。如果每一层的梯度都稍微增大一点，那么经过多层传播后，梯度值就会变得非常大，从而导致梯度爆炸。因此，网络层数的增加会加剧梯度爆炸的风险。
网络退化：理论上网络存在一个最优层数，如果模型超过这个层数，则多余层数带来的效果并不超过该最优层数下的模型效果，即这些冗余层数会带来网络退化。

何凯明大神提出的残差连接（Residual connection/skip-connect）可以有效缓解这些这些问题。从而使得拥有几十上百层的深度学习模型更加容易训练，在增加模型深度的同时还能保持甚至提高准确度。

1.2 相关知识

shortcut connections

残差连接的思想起源于中心化。在神经网络系统中，对输入数据等进行中心化转换，即将数据减去均值，被广泛验证是有利于加快系统的学习速度的。论文"Accelerated gradient descent by factor-centering decomposition"将这样的思想拓展到了梯度的反向传播中，不仅是输入和隐藏层单元的激活值要中心化，梯度误差以及权重的更新也可以中心化，这种将输入输出进行连接的操作被称为shortcut connection（为跳层连接）技术。具体而言，该论文提出了将网络分解为biased和centered两个子网络的思想，通过并行训练两个子网络，分别学习线性和非线性变换部分，不仅降低了各个网络的学习难度，也大大提升了梯度下降算法的训练速度。

而论文"Highway Networks" 提出了高速路网络（Highway Networks），其通过使用门控单元来学习管理信息，对深层神经网络使用了跳层连接，允许信息高速无阻碍的通过深层神经网络的各层，这样有效的减缓了梯度的问题。

恒等映射

Identity Mapping翻译成中文就是恒等映射，简单来说，就是输入和输出相同。恒等映射在神经网络中的作用很微妙。ResNet论文中者给出了例子如下图，56层的深层网络在训练集和测试集上的表现都不如20层的浅层网络。如果20层的浅层网络的性能已经足够好，那么就算后面36层网络什么都不做，也不至于比20层网络效果更差。但是，事实并非如此。为何会有如此退化的现象？论文判断是因为神经网络并不容易学习恒等映射。

1.3 网络结构

ResNet 的核心思想是引入一个「恒等快捷连接」（identity shortcut connection），直接跳过一个或多个层。从表象上来看，残差连接就是把网络的输入和输出相加，得到网络的输出为F(x)+x。具体来说，如果我们将子层的操作表示为F(x)，把子层的输入记作x，那么残差连接的输出就是x + F(x)，也是下一层网络模块的输入。ResNet有两篇论文，我们来看看论文的推导。

论文V1

首先，论文选取一个准确率饱和的浅层（shallower architecture）网络$y=F(x)$，再加上一些恒等映射结构（identity mapping，𝐻(𝑥)=𝑥 ,即输入输出相等），生成一个新网络。论文认为这种网络效果肯定不会比"shallower architecture"差，然而但是实验结果显示这种做法并没有比浅层网络 $𝑦=𝐹(𝑥)$更好，即下面效果反而不好。

因此论文提出了一个解决方案：使用深度残差网络，其具体结构如下。

残差网络将恒等映射结构H(x)拆分为两部分： H(x)=x+(H(x)−x)=x+F(x) ，这样就改优化 H(x) 为优化 F(x) 。其目的是学习输入 x 和输出$\mathcal{H}(x)$ 之间的残差，即 $\\mathcal{F}(x) := \\mathcal{H}(x) - x$ 。这里的 $\\mathcal{F}(x)$ 是通过两个权重层和非线性激活函数ReLU构成的子网络来学习的。这样设计的好处是：

用一堆非线性网络去仿真恒等映射是很困难的，但仿真 F(x)=H(x)−x=0 这个网络相对简单。
identity shortcuts（恒等快捷映射）既不增加额外的参数也不增加计算复杂度，使得网络易于优化，提高了泛化性能。
如果 $\\mathcal{F}(x)$ 的输出是0，即不存在残差，那么 shortcut connection 使得恒等映射成为可能，即直接输出 x，这对于网络训练的稳定性和加速有很大帮助。

我们接下来看几个细节。首先，从上图可以看到 $\\mathcal{F}(x)$ 有两个隐藏层，这样可以使得表达式中至少存在一个非线性激活函数ReLU。论文中也阐述了这一点，如下图。

其次，论文中指出可能会把多个非线性层权重设置为接近于0，见下图红色部分。这是为何？加深网络只是为了加深而加深吗？这是因为在实际情况中，恒等映射不可能永远是最优的。那些被跳过的层并不是真正的被跳过，其权重也不是永远为0。当神经网络能够学习到残差F(x)时，输出为F(x)+x；当神经网络无法学习到残差F(x)时，输出最差也是恒等映射x。通过残差学习的重构，如果恒等映射是最优的，求解器可能简单地将多个非线性连接的权重推向零来接近恒等映射。

我们再把Resnet和highway networks进行比对。highway networks的实现与ResNet相反。

highway networks的gated shortcut（门控快捷映射）是依赖数据的，且具备参数。而ResNet的identity shortcuts没有参数。
当highway network的gated shortcut"关闭"（接近零）时，highway networks 中的层代表非残差函数。相反，ResNet始终在学习残差函数，其identity shortcuts永远不会关闭，所有信息都会通过，且需要学习额外的残差函数。

论文V2

继ResNet论文之后，原版作者又发表了第二版论文"Identity Mappings in Deep Residual Networks" ，提出了一个改进版本ResNetV2。第二版论文把ResNet分成三个部分h(skip connection)，f(after-addition activation)，F(residual function)，然后做了进一步的分析，具体如下图。

研究团队还探索了多种捷径（残差连接）设计，如下图所示。

最终明确了一个核心结论：当网络深度超过 100 层时，简单的恒等残差连接是唯一有效的方法。

1.4 功用

梯度消失

从梯度的视角理解一下残差连接。这里以梯度消失为例。

残差连接对于权重退化的解决思路非常简单：不是担心输入的梯度会消失吗？那么就给它补上一个梯度为常数的项。作者将上一层学习到的结果直接与新加入层的学习结果相加，如果每层的输出都加上一个x的时候，输出就变成了F(x)+x。这样就允许网络中的信息直接跨过一个或多个层进行传播，即使某个子层没有学到有用的信息（或者学到了错误的信息），残差连接至少保证了不会丢失上一层学习到的东西。这有助于保持前向传播过程中信息的完整性。当反向传播时，对x求导结果为1，所以就相当于每一层求导时都加上了一个常数项'1'，即有一路的梯度不会经过梯度F(x) 计算，直接经过后续的处理（传播）。这样就算 F(x) 中的 x 梯度消失了，"直通"路 x 的梯度基本上也能得以保留。这可以保存更多的梯度信息，有效解决了梯度消失问题。

缓解退化

论文"SKIP CONNECTIONS ELIMINATE SINGULARITIES"指出，残差网络会减轻神经网络的退化，这才是有效性的关键所在。论文指出，训练深度网络的困难部分是由于模型的不可识别性导致的奇异性，比如如下：

由给定层中节点的排列对称引起的重叠奇异性（overlap singularities）。
与节点的消除相对应的消除奇异性（elimination singularities）。
由节点的线性依赖性产生的奇异性。

这些奇点会导致模型退化。退化发生时，学习在参数空间中沿退化方向大幅减慢，因此降低了模型的有效维数（使得网络的表达能力没有看起来那么强大），导致模型的"自由度"大大降低。并且随着网络层数的增加，连乘后使得整个秩变的更低。

skip connection则可以恢复网络的表达能力。下面展示了三种全连接层中的奇异性，以及跳层连接恢复网络表达能力的案例，分别是消除输入和权重零奇点，打破对称性，以及降低节点的线性依赖性。

消除奇异性。此时输入权重为零，使得输出权重w不可识别（红色）。跳过连接（蓝色）确保神经元至少有时处于活动状态，因此输出权重是可识别的（绿色）。相反，对于零输出权重w=0，跳过连接可以恢复一定的可识别性。
重叠奇点。重叠的输入权重（Ja=Jb）使输出权重不可识别；跳过连接再次打破退化。
线性相关奇点。隐藏神经元的一个子集变得线性相关，使得它们的输出权重无法识别。跳过连接打破了线性依赖

层间修正

我们知道每一层的输出是其值向量分布的聚合，但同时一个前馈网络的输出还包括了来自前馈网络之前的残差信息。那么这些残差链接起到了怎样的作用呢？论文"Transformer Feed-Forward Layers Are Key-Value Memories"的几个分析可以为我们揭示部分机理。

下图1号标签探究每一层的残差向量所对应的预测分布的最大概率词与整个模型的预测分布的最大概率词的重合率。大部分时候，残差的最高预测词是模型的最终预测词。
下图2号标签探究模型的预测分布中，最大概率词在每一层的残差向量所对应的预测分布里的概率值。即模型根据每一层的残差来输出token的概率。

1号标签和2号标签可以观察到类似的趋势是：即随着模型层层递进，模型的最终预测越来越确定，且概率值越来越大。这说明在各层之间，残差使得模型的输出被逐步改善，承担了模型层间修正的作用。

上图3号标签是随机挑出了100个模型最后一层里残差分布和前馈分布不同的例子。结果发现聚合的新分布的最高预测词在百分之六十六的情况下语义发生了剧变，只有百分之三十四的情况新预测词语义与原预测词接近。说明前馈网络是持续不断地更新残差的分布，即使在最后一层。

掩码 VS 残差

有读者会问：解码端的残差结构有没有把后续未被看见的token信息添加进来，造成信息的泄露。答案是不会的。因为残差连接本身并不涉及屏蔽或注意力计算。在解码器的自注意力层之后，残差连接确实会添加到应用掩码之后的输出上，但这不会引入任何额外的信息，因为掩码已经确保了模型不能"看到"后续的词。

0x02 归一化

2.1 问题

Internal Covariate shift（ICS，内部协变量偏移）是指在训练过程中，由于网络权重的不断更新，激活分布发生了变化。这种偏移会导致每个后续层接收到的输入分布各不相同，从而阻碍稳定学习。在深度神经网络中，一个神经层的输入是之前神经层的输出。Internal Covariate shift从具体现象来看就是，给定一个神经层，它之前的所有神经层的参数变化会导致本层的输入分布发生较大的改变。越高的层，其输入分布会改变得越明显。就像一栋高楼，低楼层发生一个较小的倾斜，可能会导致高楼层较大的倾斜。

该现象会导致下面几个问题：

在训练的过程中，网络需要不断适应新的输入数据分布（即梯度更新迭代步数增多），会大大降低学习速度。
随着网络层数的增加，通过多层的计算后参数可能开始出现过大或过小的情况，这样可能会导致学习过程出现异常，模型可能收敛非常的慢。
由于参数的分布不同，可能导致很多数据落入梯度饱和区，使得学习过早停止。
某些参数分布偏离太大，对其他层或者输出产生了巨大影响。

总结来说，输入特征的尺度会影响梯度下降算法的迭代步数以及梯度更新的难度，从而影响训练的收敛性。因此，我们需要找一个方法让各个特征有相似的尺度。白化是在传统机器学习中有效缓解 ICS 现象的一种手段，但是其存在一些问题，比如计算成本高、损害输入数据原本的表达能力等。因此我们需要一种更简化的解决思路，该方案应该：

对每一维特征进行归一化，使其满足均值为 0，方差为 1。
增加线性变换操作，让数据能够尽量恢复本身表达能力。

这就是归一化方案。

注：对于归一化方案是否可以有效缓解ICS，业界也有不同意见，我们后续会分析。

2.2 定义

归一化（normalization，有些文献也称为"正则化"或"规范化"）就是把输入数据X，在输送给神经元之前先对其进行平移和伸缩变换，将X的分布规范化成在固定区间范围的标准分布，简单的说就是控制每一层输入数据（激活值）的分布，把一部分不重要的复杂信息过滤掉，将数据拉回标准正态分布，使得数据的分布稳定下来（降低各维度数据的方差，减少不同层之间的输入分布变化，尽可能让输入的数据/特征变得独立同分布），以此来降低拟合难度和过拟合的风险，加速模型的收敛。归一化的作用如下：

稳定训练过程：归一化层通过对激活值进行标准化处理，减少了不同层之间的输入分布变化，从而提高了训练的稳定性。
加速收敛：通过减少Internal Covariate Shift，归一化层能够加速训练过程，使得模型更快地收敛。
提高模型性能：通过提供更稳定的训练信号，归一化层能够提高模型的最终性能，尤其是在较深的网络中。

2.3 类型

根据标准化操作的维度，Normalization可分为很多种，不同正则化的区别只是操作的信息维度不同，即选择损失信息的维度不同。如下图所示，Transformer 有两种常用归一化类型: LayerNorm 和 RMSNorm。

LayerNorm会调整神经网络每一层的值。它在SLM中得到了广泛的应用。
RMSNorm有助于调整和稳定神经网络每一层的值。

从下图可以看到，在过去的几年中，LayerNorm是最常用的技术。然而，如今已经向采用 RMSNorm的方向转变了。

我们接下来从深度学习最常见的BatchNorm入手开始介绍。

0x03 BatchNorm

BatchNorm是最早出现，也是最常见的归一化方案，而Layer Normalization是Transformer经常采用的方案。无论是BatchNorm还是LayerNorm，所做都是"减均值，除标准差"的操作，即两种归一化方法都是将数据在相应的维度上转换为标准正态分布，以解决上一层传递的数值的分布问题。那么二者区别何在？NLP领域又为何钟情LayerNorm？我们接下来就从BatchNorm入手进行分析。

3.1 公式

BatchNorm将不同样本相同维度的特征处理为相同的分布。我们接下来看看BatchNorm的公式。

总体算法包括：准备（mean and variance），归一化（normalize），还原（scale and shift/缩放和偏移）这三步。

下图给出了具体算法。

我们使用实例来进行算法解读。下图示意了BatchNorm针对三维样本的工作模式（注：没有展示还原这一步）。在该示例中，一个批次中包含了 5 个具有 4维的样本，记作 $x^1,...,x^5$ ，我们将其中第 k 个样本的第 i 维特征记作 $x^k_i$ 。BatchNorm将同一维度的所有特征视为一个分布并将其标准化。

准备。首先针对当前batch内，计算每个维度的均值和标准差。对于第 i(i=1,...,m) 维特征，得到的均值和标准差分别记作 $\\mu_i$ 和 $\sigma_i$ 。均值和标准化是由输入决定的，不需要学习。
归一化。接下来对每个数据减去均值再除以标准差的操作，使其均值为0方差为1。 $y_i^k=(x^k_i - \mu_i)/\sigma_i$ ，其中 $y_i^k$即为BatchNorm输出的第 k 个样本的第 i 维特征。这样可以将不同样本相同维度的特征处理为相同的分布。
还原。在完成标准化之后，再学习一个线性映射（$\gamma x + \beta$）中的$\gamma$和$\beta$，再把标准化的输出进行映射，具体怎么映射由BatchNorm的可学习参数自己学习得到。加入这一步的原因是因为：
- 深度神经网络每层的输入分布和权重要相互协调，强制把分布限制在zero mean unit variance并不见得是最好的选择，加入参数 𝛾 和 𝛽 ，对输入进行scale and shift，有利于分布与权重的相互协调。
- 标准化破坏了数据的原始分布，可能导致输入给下游非线性函数比如激活函数的时候产生负面效果，因此加入还原线性变换进行适度还原（无需担心数据的原始分布被破坏，导致影响网络训练的问题），而具体如何还原则可以在训练过程中决定什么样的分布是适合的。

假设这个5个样本都是人，4个维度分别是：身高，体重，薪资，年龄。身高数据可能是175cm、182cm等，体重数据往往形如63kg、72kg。这四个数据是完全不同的四个分布，经过BatchNorm的加工，这四个分布均成为标准正态分布。

3.2 作用

BatchNorm最开始被提出的时候宣称能缓解 ICS（Internal Covariate Shift），即通过对BatchSize这个维度归一化来让分布稳定下来。这种解释主要是基于概率分布的：通过对当前小批次的均值和方差进行归一化，使得每一层的输入分布保持稳定，这样就减少了ICS，从而稳定乃至加速了训练。

BN认为相同维的特征具有相同分布，因此针对不同样本的同一特征以跨样本的方式开展归一化，并不会破坏不同样本同一特征之间的关系，归一化的结果保持样本之间的可比较性。毕竟"减均值，除标准差"就是一个平移加缩放的线性操作。在上面例子中，这就意味着"归一化前是高个子的归一化后仍然是高个子，归一化前年纪大的归一化后也不会变小"。这一性质进而决定了经过归一化操作后，样本之间仍然具有可比较性。

不过目前"缓解 ICS"这个观点饱受质疑。因为不管哪一层的输入都不可能严格满足正态分布，从而单纯地将均值方差标准化无法实现标准分布 N (0, 1) ；其次，就算能做到把输入缩放到 N (0, 1) ，这种诠释也无法进一步解释其他归一化手段（如 Instance Normalization、Layer Normalization）起作用的原因。

论文"How Does Batch Normalization Help Optimization?"就明确地提出了上述质疑，否定了原来的一些观点，并提出了关于 BN 的新理解：BN 主要作用是使得整个损失函数的 optimization landscape（神经网络的高维损失函数构成的损失曲面）更为平滑，从而使得我们可以更平稳地进行训练。

或者说，BN缓解的是Variance Shift。这主要是运算 "+" 操作产生的：两个随机变量相加，结果的方差是两者的累积和。在CNN中，卷积操作在特征图的每个channel上进行了大量的叠加操作。在NLP中常用的注意力操作在每个token的embedding上也进行大量的叠加操作。逐层累积的方差如果不加以抑制，那数值的发散足以成为灾难。

3.3 PyTorch使用

下图给出了PyTorch中BN的示例。

BatchNorm1d：因为Batch Normalization是在C维度上计算（N，L）个切片的统计数据，所以一般称之为Temporal Batch Normalization。
BatchNorm2d：因为Batch Normalization是在C维度上计算（N，H，W）切片的统计数据，所以一般称之为Spatial Batch Normalization。即，假设feature map是$𝑥∈𝑅^{𝑁×𝐶×𝐻×𝑊}$，包含 N 个样本，每个样本通道数为 C，高为 H，宽为 W。对其求均值和方差时，将在 N、H、W上操作，而保留通道 C 的维度。具体来说，就是把第1个样本的第1个通道，加上第2个样本第1个通道 ...... 加上第 N 个样本第1个通道来求平均，得到通道 1 的均值（注意是除以 N×H×W 而不是单纯除以 N，最后得到的是一个代表这个 batch 第1个通道平均值的数字，而不是一个 H×W 的矩阵）。求通道 1 的方差也是同理。对所有通道都施加一遍这个操作，就得到了所有通道的均值和方差。
BatchNorm3d：因为Batch Normalization是在C维上计算（N，D，H，W）切片的统计数据，所以一般称之为Volumetric Batch Normalization或Spatio-temporal Batch Normalization。

我们再来看看测试代码。

python 复制代码

import torch.nn as nn

batchnorm = nn.BatchNorm1d(3) # 参数为Tensor最里面一维的长度
# 将1，6视为一个整体，2，4视为一个整体，3，5视为一个整体，分别进行归一化
input_data = torch.Tensor([[1,2,3],[6,4,5]])
output = batchnorm(input_data)
print (output)

batchnorm = nn.BatchNorm1d(2) # 参数为Tensor倒数第二维的长度,，相当于把最里面一维变成了一个数
# 将第一个的12和第二个的56视为一个整体，第一个的34和第二个的78视为整体进行归一化处理
input_data = torch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
output = batchnorm(input_data)
print(output)

batchnorm = nn.BatchNorm2d(2)
# 格式为(batch_size, num_channels,height,width)，在num_channels维度上进行操作
input_data = torch.Tensor([[[[1,2],[3,4]],[[100,102],[4,3]]],[[[2,3],[4,5]],[[101,103],[5,3]]]])
output = batchnorm(input_data)
print (output)

输出如下

python 复制代码

tensor([[-1.0000, -1.0000, -1.0000],
        [ 1.0000,  1.0000,  1.0000]], grad_fn=<NativeBatchNormBackward0>)
tensor([[[-1.2127, -0.7276],
         [-1.2127, -0.7276]],
        [[ 0.7276,  1.2127],
         [ 0.7276,  1.2127]]], grad_fn=<NativeBatchNormBackward0>)
tensor([[[[-1.6330, -0.8165],
          [ 0.0000,  0.8165]],
         [[ 0.9691,  1.0100],
          [-0.9947, -1.0151]]],
        [[[-0.8165,  0.0000],
          [ 0.8165,  1.6330]],
         [[ 0.9896,  1.0305],
          [-0.9742, -1.0151]]]], grad_fn=<NativeBatchNormBackward0>)

3.4 问题

BN并非是万能良药，也存在问题。导致问题的根源是：BN利用一个 mini-batch 的数据来计算一阶统计量和二阶统计量，这是按照样本数计算归一化统计量的。除此之外，Batch Norm 还对 Batch size 有一定要求，如何设置不当会失去统计意义，比如：

因为BatchNorm在批次中执行跨样本的归一化操作，这就意味着批次的构成和规模会直接影响BatchNorm的效果。当样本数很少时，这些样本的均值和方差不能反映全局的统计分布息，基于少量样本的BN的效果会变得很差。所以BatchNorm需要平衡小批次统计量和整体样本统计量之间的关系，还需要考虑利用批次统计量更新全局统计量的方法。
BN 层在训练和推理时的行为可能不一致。在推理时，BN 层使用训练过程中数据均值和方差得到的参数，而非由推理数据均值和方差得到的参数。而且BN在训练的时候可以根据Mini-Batch里的若干数据进行调整，但是在推理的输入只有一个实例，看不到其它样本，一个样本是无法求出均值和方差的。因此实际中有使用滑动平均来计算均值和方差，也有从所有训练样本中获得的统计量来代替推理数据的均值和方差。这样，当训练数据和测试数据分布存在差异时，预训练的均值和方差并不能真实反映测试集，这就导致在训练，验证，测试这三个阶段存在不一致性。

另外，在序列模型的背景下，Batch Norm 有一个天然的硬伤：同一批量内，输入序列的长度(n)不一致。这使得Batch Norm在所有序列模型中都不吃香。在NLP领域中，一个批次的输入包括若干个独立的句子，每个句子又由不定长度的词构成，句子中的每个词又被表达为一个定长向量。因此，每个句子即被视为批次中的一个样本，句子中的多个词特征构成句子特征，只不过不同句子的特征长度不一定相同。考虑到在以张量表示的样本批次要求样本必须具有相同长度，故需要对于长度不足的句子进行补齐（padding）操作。下图即为上述数据组织形式的示意，其中的空白圈即表示补齐的零向量。在该示例中，一个批次中包括了五句话，故批次中的样本个数为5；假设序列的最大长度为4；词特征的维度由模型自行确定，在Transformer原文中即为 $d_{model}$。

按照BatchNorm的思路，将批次中不同样本同一维特征视为相同分布，即对所有句子处于相同位置的词向量执行"减均值，除标准差"的操作。显然，在NLP领域中这是极不合理的，具体如下：

首先是无法计算。NLP的文本序列本质上是一个时间序列，而时间序列不定长。原则上，长度不同的序列属于不同的统计对象，很难得到稳定的统计量。而BN需要依靠滑动平均来获得预测用的统计量，因此BN无法成立。
其次是没意义。归一化的目的是将具有相同性质的数据转化为标准正态分布，其结果不应该破坏数据之间的可比较性。例如在前面"身高体重"的例子中，参与归一化的数据分别是所有的身高数据和体重数据，二者中任何一组数据都具有明确且相同的物理意义。但是，不同句子相同位置的词汇不具有相同性质，也不需要进行比较。例如在上图示例中，以第3个位置为例，五个句子在该位置的字分别是"太"、"香"、"网"、"杯"和空白，这些字者不具备相同属性或相同物理意义，根本无法执行归一化计算，会使得统计量不置信。另外，三个句子是独立处理的，三者自然也不会涉及比较的问题。显然，对于句子来说，基于BatchNorm的归一化无从谈起。

上面的论述最终指向一点：batch size这个维度使得BN在NLP场景下存在问题，导致计算出的平均值和标准差无法真正代表数据的分布。因此，batch normalization往往用于CV中，在NLP中就会出现问题。

论文"Rethinking Batch Normalization in Transformers"对于Transformer中BN表现不好的原因做了一定的分析，认为其主要问题是在前向传播和反向传播中，batch统计量和其贡献的梯度都会呈现一定的不稳定性（在使用BN的Transformer训练过程中，每个batch的均值与方差一直震荡，偏离全局的running statistics）。

既然明确了问题，解决起来就简单了，归一化的时候避开batch这个维度。于是就出现了layer normalization和instance normalization等工作。

0x04 layerNorm

因为Layer normalization不会受到不同长度句子的影响，所以非常适合NLP领域。

4.1 解决方案

首先我们看看归一化的目的。归一化的目的是将具有相同性质的数据转化为标准正态分布，其结果不应该破坏数据之间的可比较性。因此我们看看NLP中，什么数据需要彼此比较。

从上图可以看到，"太"、"香"、"网"、"杯"和空白没有比较的意义。我们即无法凭借一个字来确定这些字的含义，也无法通过这五个元素之间的比较来确定。我们需要做的是在每个句子中分析其所在句子中上下文，这样才能准确确定其含义，即样本内的特征关系更为重要。进一步说，某个词的含义是其与同一句中其他词"比"出来的，同一句中的词之间有关联，比较才有意义。除此之外，Transformer在构建句内关系时使用的自注意力机制，也就是将某个词与同句中其他词"比"出相似度。

所以，我们要保留句子中词与词之间的可比较性，进而需要对单个句子进行归一化操作，这也正是LayerNorm的归一化方式------计算一个句子的均值和标准差，然后对句中的每个词做归一化操作。LayerNorm所做的操作，类似于在一个句子中找到一个"语义中心"，然后将句子中的所有词都聚集在这个中心的周围，而句中词与词之间的比较关系不会遭到破坏。

注：上述表达与实际思路有所出入，并非实际NLP的LN版本。NLP中，LN实际是针对单个 token 的所有 features 做标准化，而非对单个句子内所有 tokens、所有 features 一起做标准化。所以和句子长度以及batch 大小无关。我们在后续会逐步细化。

4.2 公式

LN的输入是每一层神经元的输入，这些输入构成了矩阵，矩阵行是token的embedding，矩阵列是句子。我们就以行为为单位求均值和方差，然后用每一行的每一个元素减去这行的均值，再除以这行的标准差，从而得到归一化后的数值。之后引入两个可训练参数α，β来弥补归一化过程中损失掉的信息（一般初始化α为全1，而β为全0）。

4.3 作用

为什么 Layer Normalization 能够帮助 Transformer 更好地训练呢？常见的解释是layer normalization 重要的两个特性是缩放不变性（rescaling invariance）和平移不变性（recentering invariance）。

缩放不变性。指的是归一化过程能够适应输入数据的缩放，使得网络对这种缩放不敏感。
平移不变性。如果输入数据的均值发生了变化，但数据的分布形状和范围保持不变，那么具有平移不变性的算法或函数的输出不会受输入变化的影响。

因为这两个优点，LN对模型训练的影响也有两点:

前向计算时对数值进行归一化，可以使得前向传播的输入分布变得稳定。
归一化操作也可以使反向传播的梯度更加稳定。

论文"On the Expressivity Role of LayerNorm in Transformers' Attention" 也提出了独到见解。

没有LayerNorm情况下，keys和queries没有什么明显的几何结构；LayerNorm却可以将keys投影到同一超平面上，因此通过学习，模型可以将query与keys正交对齐。
没有LayerNorm，存在一些"无法选择"的key向量，即这些向量无法通过获取最大注意力得分进行选择（下图上深色区域）；而LayerNorm消除了"无法选择"的key向量，任何key向量都可能得到最高的注意力分数。

论文"On the Nonlinearity of Layer Normalization"则指出，LN以及其计算退化版本 RMSNorm 具有非线性表达能力，并详细讨论了 LN 的万能近似分类能力。该论文对 LN 的非线性进行了数学证明。并且提出了仅含线性层和 LN 的简单神经网络 LN-Net，在足够深的情况下，LN-Net理论上可以对给定的样本和样本类别进行任意分类。这一发现打破了人们将各种 Normalization 视为不具有拟合能力的线性变换的惯性认知，而非线性层和标准化层也不再是互不相交的神经网络模块。

为了进一步研究，论文将 LN 拆分为两个步骤：中心化（centering）和尺度缩放 (scaling)。中心化从数学上是一个线性变换，因此 LN 的非线性主要存在于尺度缩放操作当中（论文中也称之为球面投影，是 RMSNorm 执行的操作）。论文以最为简单的线性不可分的异或数据为例，通过线性变换和球面投影将这四个点进行了正确分类。

上图给出了Xor分类的解决方案。首先，我们将它们旋转45°，如图(b)所示。然后，我们将它们垂直投影到y=0.5上，如图(c)所示。接下来，我们将它们球形投影到圆$x^2+y^2=1$上，如图(d)所示。最后，我们将它们水平投影到x=0上，如图(e)所示。现在我们已经做出了正确分类。

4.4 LN和BN的差异

深度学习中的正则化方法就是"通过把一部分不重复的复杂信息损失掉，以此来降低拟合难度以及过拟合的风险，从而加速了模型的收敛"，其目的就是为了让分布稳定下来（降低各维度数据的方差）。

理解了为什么要做归一化之后，我们再来探讨要做什么样的归一化（能在哪些维度做归一化）。数据信息的重要性决定了归一化的方式。不同正则化方法的区别只是操作的信息维度不同，即选择损失信息的维度不同。所以选择不同的Normalization方法，就是针对具体问题来选择在对应维度上进行处理。

因此，两种归一化方法的差异主要体现在二者因面对任务的不同而引起的作用对象差异，以及其由此引发的一系列使用方法的不同。下图中，BN通过计算在batch维度上的通道平均值和方差来对激活（特征图）进行归一化。LN则在通道/特征（channel/feature）维度上进行归一化。

作用对象

以批次将多个样本组织起来送入神经网络能够减少数据吞吐量、提升计算效率，也为模型训练提供了更多的统计信息。Batch Normalize是考虑本batch内所有数据，Layer Normalize则是不考虑其他数据，只考虑自己。

BN是对一个batch-size样本内的每个特征维度上做归一化，特点如下：
- BatchNorm针对同一特征，以跨样本的方式开展归一化。
- BatchNorm在批方向计算均值和方差，因此不会破坏不同样本同一特征之间的关系，毕竟"减均值，除标准差"就是一个平移加缩放的线性操作。这一性质进而决定了经过归一化操作后，样本之间仍然具有可比较性。
LN是在每个样本内部，对每个样本的所有特征做归一化，特点如下：
- Layer normalization不是对一个batch内的特征维度进行处理（不需要batch的均值和方差），而是对样本维度进行处理。LN 针对单个训练样本进行上计算均值和方差，不依赖于其他数据，因此可以避免 BN 中受 mini-batch 数据分布影响的问题。
- LayerNorm 针对每个样本自己算均值和方差，不需要存全局的均值和方差。因为LN 不需要保存 mini-batch 的均值和方差，所以节省了额外的存储空间。

作用方向

其实，此处是作用对象的另一种解读，这里单独阐释，会更有利于大家理解。

BN是以行为单位（每一行减去每一行的均值然后除以每一行的标准差）。每一行就是一个batch中的相关维度，就是对这一个batch中的每一个对应维度做归一化。如果在NLP领域，BN会对一个 batch 内的所有句子的同一位置的token的embedding做归一化。这显然这不符合NLP的规律。因为语言文本的复杂性是很高的，任何一个词都有可能放在初始位置，且词序可能并不影响我们对句子的理解。这导致同一位置的token并不是一一对应的。这种归一化会导致语义发生改变。
LN是以列为单位，是对每一条句子内部做归一化。LN 在 NLP 中相当于对一个 batch 内的每个句子内所有位置的词做单独归一化。这样，归一化既保留了原来的语义，又能够有效防止值过大导致的梯度问题。从样本角度的合理性来看，就是每个句子毕竟差距还是比较大的，既然是融合全局信息，并归一化，最好还是各个句子归一化自己的。

具体参见下图。注意，下图只是初步演示，并非实际NLP的LN版本，实际是针对单个 token 的所有 features 做标准化，而非对单个样本内所有 tokens、所有 features 一起做标准化。

业务选择

前面是从表象来看，接下来就要从深层次来看。

CV

CV使用BN为主，所以我们看看为何CV选择BN。

batch Normalization 实际是对特征的每一维统计所有样本的均值和方差。一个 CNN 的常见的输出 shape 是：(N, C, H, W) ，BN要操作的这个特征维度是 Channel，也就是我们可以在一个 batch 中"$所有样本 \times 所有 H \times所有 W$ "上进行统计。最后得到的均值和方差向量都是 (C)，因此两个参数也是 C 维的向量。

为什么CNN仅在channel维度上统计呢？这是因为CV使用卷积核来输出特征图（h1，w1），一个卷积核可以得到一个通道的特征图数据，多个卷积核可以得到多个特征图。所以卷积核数=通道数，且不同图像，图像不同位置使用的卷积核的参数都共享的。因为我们希望不同图像，图像不同位置用这个卷积核执行卷积以后的数据分布是稳定的，所以需要在通道维度执行归一化。因此，对于CNN的特征图（B，C，H，W），我们可以在（B，H，W）三个维度做归一化，其中（H，W）是必须有的，因为在每个channel上，也就是每个特征图（HxW）都是经过了反复的叠加产生的，所以（HxW）沿着这两个维度是一定需要进行归一化操作的。最终，图像常见的归一化操作有BN，IN，GN，万变不离其宗。

NLP

NLP中使用LN，我们来看看为何研究人员做出这个选择。

LayerNorm认为每个样本内的特征具有相同分布，因此LayerNorm的归一化操作只在样本内部（单独一句话）独立开展，因此实际可以完全忽略批次的存在。因为LN不涉及其他样本，所以推理过程的行为和训练一致。

LN一般用在第三维度，即在（Batch size, seq len, embedding size）中的embedding size（词向量的维度）上计算统计量。这一维度各个特征的量纲应该相同。因此也不会遇到因为特征的量纲不同而导致的缩放问题。

对于对NLP（transformer）而言，每个token的embedding也是经过注意力操作，叠加其他embedding产生的。transformer每一层都在不断调整每个词在空间中的位置。比如某个token $E_i$可以看成所有环境词的加权和。这个过程会有一个潜在风险：$E_i$是不可控的，环境词向量和环境词的数量对$E_i$都有影响，这就可能造成偏移量会很大，相加后的"尺度"也有可能发生较大变化，尤其是transformer的维度较高，每个维度一个很小的变化也可能引起很大的"尺度"变化。所以对embedding进行归一化是必要操作。当$E_i$较小时，LN基本不起作用，但是当$E_i$较大时，layer normalization 可以将相加后得到的向量拉回到原来向量的附近，让它不至于跑太远。

为何 Transformer 中的 LN 层只针对最后一个维度（n_features）标准化？不对后两个维度（seq_len, n_features）进行正则化？其原因也许是因为 padding 的存在。一个样本序列中可能会有一些 token，如果对后两个维度做标准化，那么平均值和方差会受到 padding max length 的影响。相比之下，对同一个 token 的各个 feature 之间做标准化就更加合理。另外，无法在token维度进行统计，也是因为一个句子中不同的token之间相关性极强，彼此独立的假设显然不成立。

总结下，Transformer 中的 Layer Norm 比较特殊，是针对单个 token 的所有 features 做标准化，并不是对单个样本内所有 tokens、所有 features 一起做标准化。所以LN和句子长度、batch 大小无关。

下面是pytorch官网给出的在NLP和CV领域使用LayerNorm的例子。

python 复制代码

# NLP Example
batch, sentence_length, embedding_dim = 20, 5, 10
embedding = torch.randn(batch, sentence_length, embedding_dim)
layer_norm = nn.LayerNorm(embedding_dim)
# Activate module
layer_norm(embedding)

# Image Example
N, C, H, W = 20, 5, 10, 10
input = torch.randn(N, C, H, W)
# Normalize over the last three dimensions (i.e. the channel and spatial dimensions)
# as shown in the image below
layer_norm = nn.LayerNorm([C, H, W])
output = layer_norm(input)

其实，在Transformer中使用BN也是可以的，之所以效果没有LN好，更多的是工程上的原因，例如句子的长短造成的影响。"Rethinking Batch Normalization in Transformers"中就提到了这点。

具体实现

有几种非线性操作，如Softmax、LayerNorm和GELU，需要特殊的支持或芯片外计算。与线性运算相比，当使用Transformer网络进行推断时，这些非线性运算在总体运算中所占的比例相对较小。然而，与矩阵计算相比，在典型硬件上计算这些非线性运算更具挑战性，如果处理不当，它们可能会产生大量开销。

如图概述Softmax、LayerNorm和BatchNorm操作。由于它们依赖于运行时统计信息，LayerNorm和Softmax都需要对输入进行多次传递才能计算非线性运算。在Softmax的情况下，需要第一次通过输入来计算分母。对于LayerNorm，需要对输入进行三次传递：一次用于计算平均值；一个用于计算标准偏差；以及一个用于应用规范化。与LayerNorm和Softmax不同，BatchNorm只使用在训练过程中学习的统计数据，因此它只需要对输入进行一次传递。

4.5 Post-Norm VS Pre-Norm

概念

Transformer模型结构中每层都包含着残差结构，"Pre"和"Post"指的是Norm所处的位置。考虑到LayerNorm放置的位置不同，可以分为PreLayerNorm和PostLayerNorm。

Post-Norm：在Add操作后进行Norm操作。最初的 Transformer block 中 Layer Norm 层是放在残差连接之后的，也被称为 post-LN。
Pre-Norm：Norm之后再Add。就是把 Layer Norm 层放在多头自注意力层或者全连接层之前。有研究表明"Pre-LN"对梯度下降更加友好，收敛更快，更易于超参优化，但其性能总差于"Post-LN"。

两种方式的公式化如下。

\[Pre\ Norm: x_{t+1} = x_t + F_t(Norm(x_t)\\Post\ Norm: x_{t+1} = Norm(x_t + F_t(x_t)) \]

另外，下图中的Sandwich-Norm是在Pre-Norm的基础上，在sublayer之后新增了一个层归一化，在优先保证残差连接有效性的同时，试图更好地控制网络层输出值的方差。

论文实现

论文（vanilla Transformer）中，Encoder架构如下，这是Post-Norm。

\[Attention \rightarrow Add\ \&\ Norm \rightarrow Feed\ Forward \rightarrow Add\ \&\ Norm \]

Add & Norm指的是残差连接后使用Layer Norm，Sublayer表示经过的变换，X表示 Multi-Head Attention 或者 Feed Forward 的输入，MultiHeadAttention(X) 和 FeedForward(X) 表示相关模块的输出。

\[Add\ \&\ Norm : LayerNorm(X + Sublayer(X)) \]

对应代码是：

python 复制代码

return self.norm(x + self.dropout(sublayer(x)))

对于架构图如下。

最终，输入矩阵X经过Encoder之后，输出为：

\[O = LayerNorm(X + MultiHeadAttention(X)) \\ O = LayerNorm(O + FeedForward(O)) \]

Post-Norm

Post LayerNorm 是在每个神经网络层的输出上进行归一化操作（即在残差之后进行归一化）。这样可以使得每个神经网络层的输出都在相似的尺度上，保持了每个模块的一致性，稳定了前向传播的方差，避免了梯度消失和梯度爆炸的问题，提高了网络的稳定性和训练效果。因此对参数归一化的效果更好，模型的鲁棒性更强。一旦训练好效果会更优。

Post LayerNorm 存在一些问题，比如难以训练，需要仔细初始化等，导致其在目前的大模型中较少使用。

难以训练

Post LayerNorm会突出残差分支，失去了残差网络"易于训练"的优点。Post LayerNorm的公式如下：$x_{t+1} = Norm(x_t + F_t(x_t))$。假设 x 的方差是 1，F (x) 的方差是$σ^2$，Norm 操作就相当于除以$\sqrt{1+σ^2}$。如果 σ 比较小，那么残差中的 "直路" 权重就越接近于 1，模型在初始阶段就越接近一个恒等函数，越不容易梯度消失，从而更利于优化。

但是，由于所有参数都是随机初始化的，所以我们可以认为 x 与 F (x) 是两个相互独立的随机向量。假设它们各自的方差是 1，则 x+F (x) 的方差是 2，而 Norm 操作负责将方差重新变为 1，那么在初始化阶段，Norm 操作就相当于" 除以$\sqrt{2}$ "。Post-Norm每Norm一次就削弱一次残差的恒等分支的权重，本来残差的意思是给前面的层搞一条"绿色通道"，让梯度可以更直接地回传。但是在Post Norm中，这条"绿色通道"被严重削弱了，越靠近前面的通道反而权重越小（越靠近输入，削弱得越严重），残差"名存实亡"，故此不容易训练。

需要热身

相比Pre-Norm，Post-Norm的训练更加不稳定。出现训练不稳定（梯度消失/爆炸）的主要原因是Post-Norm对参数非常敏感。

当在残差块之间进行层归一化时，输出层附近的参数的梯度会较大。如果没有预热阶段，直接对这些参数使用较大的学习率可能不会导致模型的改进，反而会使优化过程不稳定。因此， Post Norm需要很仔细地调参才能取得好的结果，比如必备的warm-up学习率策略。warmup会留给模型足够多的时间进行 "预热"，这个过程中会抑制后面的层的学习速度，并且给前面的层更多的优化时间，以促进每个层的同步优化。最终使得更新的参数很小，LN相对稳定，减轻了梯度消失的可能性，使得训练更加稳定。

另外，当前 NLP 中主流的优化器是 Adam 及其变种。理论上只要梯度的绝对值大于随机误差，那么对应的参数都会有常数量级的更新量；而SGD 的更新量是正比于梯度的，只要梯度小，更新量也会很小，如果梯度过小，那么参数几乎没有被更新。所以，Post Norm 的残差虽然被严重削弱，但是在 base、large 级别的模型中，它还不至于削弱到小于随机误差的地步，配合 Adam 等优化器，它还是可以得到有效更新的，也就有可能成功训练了。当然只是有可能，事实上越深的 Post Norm 模型确实越难训练。

梯度消失也不全是 "坏处"，其实对于 Finetune 阶段来说，它反而是好处。在 Finetune 的时候，我们通常希望优先调整靠近输出层的参数，不要过度调整靠近输入层的参数，以免严重破坏预训练效果。而梯度消失意味着越靠近输入层，其结果对最终输出的影响越弱，这正好是 Finetune 时所希望的。所以，预训练好的 Post Norm 模型，往往比 Pre Norm 模型有更好的 Finetune 效果。

Pre-Norm

在论文ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE中提到, Transformer中Layer Norm的位置加的有问题（即Post Norm）。Layer Norm如果调整到Sub Layer前称为Pre Norm，这会对训练有很大帮助（即Pre Norm）。

Pre-Norm主要思路是："要用的时候才去标准化"和在"非线性之前单独处理各个矩阵"。其公式为$x_{t+1} = x_t + F_t(Norm(x_t))$。对于 Transformer，主要的非线性部分在 FFN（ReLU）和 Self-Attention（Softmax）中，这些逐个叠加的同构子层像极了 GRU 和 LSTM 等 RNN 单元。信息的流动由沿着时序穿过子层。把 LN 设置在每个子层的输出位置，意义上已经不再是"落入sigmoid 的梯度敏感空间来加速训练"了，更重要的是让每个词的向量化数值更加均衡，以消除极端情况对模型的影响，获得更稳定的深层网络结构 ------ 就像这些词从 Embdding 层出来时候那样，彼此只有信息的不同，没有能量的多少。LN 正如 LSTM 中的tanh，它为模型提供非线性以增强表达能力，同时将输出限制在一定范围内。因此，对于 Transformer 来说，LN 的效果已经不是"有多好"的范畴了，而是"不能没有"。

苏建林大神这篇文章中，为什么Pre Norm的效果不如Post Norm？表明，在同一训练设置下，Pre Norm（也就是Norm and add）往往更容易训练，其训练效果是要优于Post Norm（Add and Norm）的。但是单独调整的话（采用适合自己的训练配置）且可以训练完全，Post Norm训练出来的效果是更好的。

为什么会出现这种情况？

知乎唐翔昊大神给出的答案是：Pre Norm 的深度有 "水分"！Pre Norm 结构无形地增加了模型的宽度而降低了模型的深度。一个 L 层的 Pre Norm 模型，其实际等效层数不如 L 层的 Post Norm 模型。我们知道深度通常比宽度更重要，所以降低深度导致无形中层数变少、最终效果变差了。

我们把Pre-Norm的公式进一步展开得到如下：$x_t = x_0 + F_0(x_0) + F_1(\frac{x_1}{\sqrt 2}) + F_2(\frac{x_2}{\sqrt 3}) + ... + F_{t-1}(\frac{x_{t-1}}{\sqrt t})$。这样一来，起码每一条残差通道都是平权的，残差的作用会比Post Norm更加明显，所以Pre-Norm也更好优化。当然，这样最后的$x_t$方差将会很大，所以Pre-Norm在接预测层之前$x_t$也还要加个Normalization。

然而，当 t 比较大时，$x_t$和$x_{t+1}$ 相差较小，所以 $F_{t+1} (Norm (x_{t+1}))$ 与 $F_{t+1} (Norm (x_t))$ 很接近，因此原本一个 t 层的模型与 t+1 层之和，近似等效于一个更宽的 t 层模型。所谓 "the gradients of Pre-LN at bottom layers tend to be larger than at top layers"，就是指 Pre Norm 结构会过度倾向于恒等分支（bottom layers），从而使得 Pre Norm 倾向于退化（degradation）为一个 "浅而宽" 的模型，最终不如同一深度的 Post Norm。在 Pre Norm 中多层叠加的结果更多是增加宽度而不是深度，层数越多，这个层就越"虚"。

另外，Pre-Norm由于并不是所有的参数都参与正则化（因为有一部分参数直接加在了后面，不需要对这部分参数进行正则化），因此整体来说更不容易发生梯度爆炸或梯度消失的问题，模型训练的稳定性更强。

小结

下图给出了Pre-Norm和Post-Norm两者模型效果的数学推导。Pre-Norm结构往往过度依赖恒等分支（底层），导致其容易退化为一个"浅而宽"的模型，这使得Pre-Norm的表现不如同样深度的Post-Norm。从直观上看，表现为L层的堆叠效果未能达到预期。而Post-Norm则恰恰相反，它通过每次归一化操作逐渐削弱恒等分支的影响，反而更加突出残差分支，因此Post-Norm的层次结构更为"充分"，层数更加 "足秤"，在训练完成后通常表现更优。

下图给出了两者的架构区别。Pre-Norm在同样参数配置下没有Post-Norm能达到的上界高，但是它容易训练，在训练刚开始的阶段，每层的梯度范数近似相等，这有助于我们抛弃 warm-up 阶段，在提升训练稳定性的同时，还可以加速模型的收敛。而Post-norm虽然理论上行可以达到更好的算法上限，但是其在训练的时候也需要时常关注梯度，收敛起来会比较难，但一旦训练好的好，效果肯定更优。

下图给出了一些LLM的配置。这里，PE表示位置嵌入，#L表示层数，#H表示注意力头的数量，$d_{model}$表示隐藏状态的大小，MCL表示训练期间的最大上下文长度。在Bert时代由于层数较浅，往往采用的是Post-Norm，而到了大模型时代，由于transformer的层数开始加深，尽管 Pre-LN 可能会带来一定程度上的性能损失，为了训练稳定性，大多数 LLM 还是选择使用 Pre-LN。

0x05 扩展比对

5.1 Instance Norm

为了更好的比对，我们也介绍下Instance Norm。Instance Normalization (IN) 最初用于图像的风格迁移。其作者发现，在生成模型中，feature map 的各个 channel 的均值和方差会影响到最终生成图像的风格，因此可以先把图像在 channel 层面归一化，然后再用目标风格图片对应 channel 的均值和标准差"去归一化"，以期获得目标图片的风格。IN 操作也在单个样本内部进行，不依赖 batch。

IN 对每个样本的 H、W 维度的数据求均值和标准差，保留 N 、C 维度，也就是说，它只在 channel 内部求均值和标准差。因此，instance norm对CNN而言是在归一化时在特征维度上做选择的最小集，如果不选择这个维度，优化都很难进行下去。同理，layer norm对transformer而言也是如此。

对CNN而言，只要保证batch之间的图像是独立的，那么我们可以把instance norm拓展到B这个维度，也就是BN。但如果batch之间有很强的相关性，例如有大量的重复图像，因为feature map之间是有非常强的相关性的，每个特征图之间存在大量pattern的共现性，强行施加独立同分布假设并不会有好结果。

5.2 GroupNorm

Group Normbalization（GN）是另一种深度学习归一化方式，可以替代BN。GN 也是独立于 batch 的，它是 LN 和 IN 的折中。

Group Normalization优化了BN在比较小的mini-batch情况下表现不太好的劣势。从批量维度进行归一化会带来一些问题------批量统计估算不准确导致批量变小时，BN 的误差会迅速增加。在训练大型网络和将特征转移到计算机视觉任务中（包括检测、分割和视频），内存消耗限制了只能使用小批量的BN。而小的Batch Size 则会导致Batch Normalization 失效，因为没办法通过几个样本的数据量，来近似总体的均值和标准差。

Group Normalization则更适用于占用显存比较大的任务。Group Normalization首先将 Channels 划分为多个 groups，再计算每个 group 内的均值和方法，以进行归一化。从深度学习上来讲，可以认为卷积提取的特征是一种非结构化的特征或者向量。每一层有很多的卷积核，这些核学习到的特征并不完全是独立的，在同一张图像上学习到的特征应该是具有相同的分布，那么，具有相同的特征可以被分到同一个group中。

GN 的计算与 Batch Size 无关，且对于不同的 Batch Size ，精度都比较稳定，这样就避开了batch size对模型的影响。

5.3 比对

我们使用一个例子来进行比对。本例来自网络，但没有找到原始出处，遗憾。如果哪位读者知道出处，还请帮忙指出。

类比

如果把 $𝑥∈𝑅^{𝑁×𝐶×𝐻×𝑊}$类比为一摞书：总共有 N 本，每本有 C 页，每页有 H 行，每行有 W 个字符。以下是各种归一化求均值的方法，求标准差时也是同理。

BN：把这些书按页码一一对应地加起来（例如第1本书第6页，第2本书第6页...），再除以每个页码下的字符总数：N×H×W，因此可以把 BN 看成求"平均书"的操作（注意这个"平均书"每页只有一个字）。
LN：把每一本书的所有字加起来，再除以这本书的字符总数：C×H×W，即求整本书的"平均字"。
IN：把一页书中所有字加起来，再除以该页的总字数：H×W，即求每页书的"平均字"。
GN：把一本 C 页的书平均分成 G 份，每份成为有 C/G 页的小册子，求每个小册子的"平均字"。

细节

下图形象的表示了四种归一化的工作方式，也可以看出来CV和NLP领域使用BN和LN的不同区别。

对于CV，输入为 (B, C, H, W) 格式。N是batch size，H/W是feature的高/宽，C是feature的通道数，压缩H/W至一个维度，其三维的表示如下图。假设单个方格的长度是1，那么其表示的是[6, 6，*, * ]。一个C,H,W是一个面，这个面里所有数据叫一个batch。

BN是在batch上，对N、H、W做归一化（归一化维度为[N，H，W]），而保留通道 C 的维度。
LayerNorm避开了batch size维度，在通道方向上，对C、H、W归一化（归一化的维度为[C，H，W]），最后输出(B, 1, 1, 1)的mean和std。
InstanceNorm是对H、W做归一化（归一化的维度为[H，W]），最后输出(B, C, 1, 1)的mean和std。
GN将channel分组，然后再做归一化（归一化的维度为[C//G , H, W]）。事实上，GN的极端情况就是LN和IN，分别对应G等于C和G等于1。

另外，还需要注意它们的映射参数λ 和β的区别：对于 BN，IN，GN，其λ 和β都是维度等于通道数 C 的向量。而对于 LN，λ 和β 都是维度等于 normalized_shape 的矩阵。

在NLP里面，对于输入为(N, L, D)的feature：

BN会固定每句话的第n个位置，这个切片是（N，D）维的矩阵。
LN会固定一句话，则切片是（L, D）维。LayerNorm是沿着D做reduce，最后输出(N, L, 1)的mean和std。并且最后的gamma和beta是(1, 1, D)维度的。针对 NLP 的例子，只在最后一个维度（n_features）内计算均值和方差。即归一化形状为（n_features）。和 BatchNorm 有所不同（BN 归一化形状之内的所有元素上共用一套缩放、平移参数），LayerNorm 在其归一化形状内的每一个元素上分别进行缩放以及平移。由上可见，NLP语境下的LayerNorm，其实是CV里面的InstanceNorm。

0x06 实现

6.1 LayerNorm

下面是LayerNorm的代码，其功能与torch.nn.BatchNorm2d的作用一致。从形式上来说，用$\mathbf{x} \in \mathcal{B}$表示一个来自小批量$\mathcal{B}$的输入，批量规范化$\mathrm{BN}$根据以下表达式转换$\mathbf{x}$：

\[\mathrm{BN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x} - \hat{\boldsymbol{\mu}}\mathcal{B}}{\hat{\boldsymbol{\sigma}}\mathcal{B}} + \boldsymbol{\beta}. \]

$\hat{\boldsymbol{\mu}}\mathcal{B}$是小批量$\mathcal{B}$的样本均值，$\hat{\boldsymbol{\sigma}}\mathcal{B}$是小批量$\mathcal{B}$的样本标准差。应用标准化后，生成的小批量的平均值为0和单位方差为1。由于单位方差（与其他一些魔法数）是一个主观的选择，因此我们通常包含拉伸参数（scale）$\boldsymbol{\gamma}$和偏移参数 （shift）$\boldsymbol{\beta}$，它们的形状与$\mathbf{x}$相同。请注意，$\boldsymbol{\gamma}$和$\boldsymbol{\beta}$是需要与其他模型参数一起学习的参数。

python 复制代码

"""构建一个层归一化（layernorm）模块，对应论文原图中 "Add & Norm"中"Norm"的部分，其实与torch.nn.BatchNorm2d的作用一致。"""
class LayerNorm(nn.Module):   
    # 初始化函数，接收features（特征维度大小）和eps（防止除以零的微小值）作为输入参数
    def __init__(self, features, eps=1e-6):
        """
        features: int类型，含义为特征数。也就是一个词向量的维度。该值一般和d_model一致。
        eps: 在规范化公式的分母中出现的微小值，防止分母为零，默认是1e-6
        """
        super(LayerNorm, self).__init__() # 调用父类nn.Module的构造函数
        """
        这两个参数其实对应BatchNorm的参数，a_2对应gamma(γ), b_2对应beta(β)。
        而nn.Parameter的作用就是将这个两个参数作为模型参数，之后要进行梯度下降。
        """
        self.a_2 = nn.Parameter(torch.ones(features)) # 定义一个大小为features的一维张量，初始化为全1，并将其设置为可训练参数
        self.b_2 = nn.Parameter(torch.zeros(features))  # 定义一个大小为features的一维张量，初始化为全0，并将其设置为可训练参数
        self.eps = eps # 将防止除以零的微小值eps保存为类实例的属性

    # 定义前向传播函数，参数x是输入张量
    def forward(self, x):
        """
        x代表来自上一层（注意力层或者FFN层）的输出。x的形状和解码器的输入一样，其实在整个过程中，x的形状都不会改变。例如，x的shape为(1, 7, 128)，即batch_size为1，7个单词，每个单词是128维度的向量。
        """

        # 计算输入x在最后一个维度（最后一维才是样本的维度）上的均值，保持维度与输入一致
        mean = x.mean(-1, keepdim=True) # mean的形状为 (1, 7, 1)
        # 计算输入x在最后一个维度上的标准差，保持输出结果的维度，保持维度与输入一致
        std = x.std(-1, keepdim=True) # std的形状为 (1, 7, 1)
        # 进行归一化
        # 对输入x进行层归一化，使用可训练参数a_2和b_2进行缩放和偏移，最后返回归一化后的结果
        # 对应BN的操作：output = (gamma * (x - mean) / (std + eps)) + beta
        # *号代表同型点乘，即对应位置进行乘法操作
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

Encoder中使用方式如下。

python 复制代码

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"

    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

Decoder中使用方式如下。

python 复制代码

class Decoder(nn.Module):
    "Generic N layer decoder with masking."

    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

6.2 残差

哈佛教程之中的实现方式是Pre-Norm。

\[Norm \rightarrow Attention \rightarrow Add \rightarrow Norm \rightarrow Feed\ Forward \rightarrow Add \]

不管是 Self-Attention 还是全连接层，都首先是接一个 LayerNorm，然后是 Self-Attention/Dense 和 Dropout，最后是残差连接。这里面有很多可以重用的代码，因此哈佛代码把这些可以重复的代码封装成 SublayerConnection类。为了简单，代码是按照下图这个顺序实现的：

我们先看看如何使用SublayerConnection类。EncoderLayer类的forward()函数中，调用SublayerConnection时候，传入的分别是self_attn和self.feed_forward。其中调用self.sublayer[0] 就调用到self_attn的功能。

我们把x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) 简化为self.sublayer[0] (x, z) 。self.sublayer[0] 是个 callable，self.sublayer[0] (x, z) 会调用 self.sublayer[0].__call__(x,z)，然后会调用 SublayerConnection.forward(x, z)。从后续代码中可以看到，SublayerConnection.forward(x, z)有两个参数，一个是输入 Tensor，一个是一个 callable，并且这个 callable 可以用一个参数来调用。最终会调用 sublayer(self.norm(x))。sublayer 就是传入的参数 z，因此就是 z(self.norm(x))。

我们再看看self_attn如何转换为z。self_attn 有 4 个参数，但是我们知道在 Encoder 里，前三个参数都是输入 x，第四个参数是 mask。这里 mask 是已知的，因此我们可以用 lambda 的技巧它变成一个参数的函数 z = lambda x: self.self_attn(x, x, x, mask)，这个函数的输入是 x。

可以看到，在向量经过Position Encoding后到进入Encoder前是有一个Norm动作的，Pre-Norm 在此处发挥了作用。

python 复制代码

class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"

    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 传入参数
        return self.sublayer[1](x, self.feed_forward) # 传入参数

SublayerConnection代码如下。这个类会构造 LayerNorm 和 Dropout，但是 Self-Attention 或者 Dense 并不在这里构造，还是放在了 EncoderLayer 里，在 forward 的时候由 EncoderLayer 传入。这样写代码的好处是更加通用，比如 Decoder 也是类似的需要在 Self-Attention、Attention 或者 Dense 前面后加上 LayerNorm 和 Dropout 以及残差连接，我们就可以复用代码。此处要求传入的 sublayer 可以使用一个参数来调用的函数 (或者有 __call__函数)。

python 复制代码

class SublayerConnection(nn.Module):
    """
    本类实现的功能是Layer Norm层之后跟随一个残差连接（residual connection），即LayerNorm + sublayer(Self-Attenion/Dense) + dropout + 残差连接。为了实现简单，此处把LayerNorm放到了前面，这和原始论文稍有不同，原始论文LayerNorm在最后。
    """
   
    # 初始化函数，把接收size（层的维度大小）和dropout（dropout率）作为输入参数
    def __init__(self, size, dropout):
        # size是d_model，即词向量的维度
        super(SublayerConnection, self).__init__() # 调用父类nn.Module的构造函数
        self.norm = LayerNorm(size)   # 定义一个层归一化(Layer Normalization)操作，使用size作为输入维度
        self.dropout = nn.Dropout(dropout) # 定义一个dropout层

    # 定义前向传播函数，输入参数x是输入张量，sublayer是待执行的子层操作
    def forward(self, x, sublayer):
        """
        Apply residual connection to any sublayer with the same size.
        x是前一层的输出，本层的输入
        sublayer是Attention层或Feed ForWard层，它可以当成函数调用，这个函数的有一个输入参数
        """
        
        # 将残差连接应用于任何具有相同大小的子层
        # 首先对输入x进行层归一化，然后将结果传给子层，执行子层操作（如self-attention或前馈神经网络）
        # 接着应用dropout，dropout是为了随机停止一些网络中神经元的作用，来防止过拟合。
        # 最后将结果与原始输入x相加。
        # 因为存在跳跃连接，所以是将输入x与dropout后的子层输出结果相加作为最终的子层连接输出。    
        return x + self.dropout(sublayer(self.norm(x)))

0x07 优化与演进

7.1 RMSNorm

目前很多LLM采用了 RMSNorm（Root Mean Square Layer Normalization）来替代传统的 LayerNorm 函数，这一改变在保持训练稳定性和提升模型收敛速度的同时，大幅提高了计算效率。RMSNorm 之所以能更高效，是因为其创造者对LN进行研究和改进。LN重要的两个部分是平移不变性和缩放不变性。RMSNorm的作者认为LN取得成功重要的是缩放不变性，而非平移不变性。基于这一发现，与 Layer Normalization 先计算均值和方差不同，RMSNorm 去除了计算过程中的平移（去除了减掉均值的部分），只保留了缩放，即直接计算输入特征的平方均值，然后用其对输入特征进行缩放。这种计算方式使得RMSNorm 具备如下优点：

更好的数值稳定性。RMSNorm 的计算方式更加稳定，可以避免 Layer Normalization 中可能出现的异常值问题，尤其是在训练大规模模型时效果更明显。
更快的训练速度。省略了归一化过程中的均值计算，使得算法更加简洁，计算量更小，可以加速模型训练过程。
更好的模型性能。在一些实验中，与LN相比，RMSNorm 效果基本相当，甚至略有提升。被证明可以提升模型的最终性能。

下图给出了层归一化（LayerNorm）与均方根归一化（RMSNorm）之间的部分差异。

Huggingface transformer 库的 Llama 对于 RMSNorm 的实现如下。

python 复制代码

from torch import nn
class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        # RMS Norm层的权重参数，和hidden states相同尺寸。在具体的计算过程中会进行相乘
        self.weight = nn.Parameter(torch.ones(hidden_size))
        # 防止根号运算出现根号0计算非法的问题，也不要太小，推荐使用1e-5到1e-7来防止数据精度溢出
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        # 输入是隐变量,它的shape是[batch_size, seq_len, hidden_size]
        # 先保存输入hidden_state的数据类型
        # 中间计算过程使用torch.float32数据类型
        # 在计算最后会将其恢复为hidden_state的数据类型
        input_dtype = hidden_states.dtype
        # 为了防止计算出错,先将隐藏层张量的数据类型提升为torch.float32
        hidden_states = hidden_states.to(torch.float32)
        # 在最后一个维度计算均方值
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        # torch.rsqrt表示对输入进行开根号求倒数
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        # 然后在除以均方根的基础上再点乘一个参数张量，回到输入的数据类型
        return self.weight * hidden_states.to(input_dtype)

Llama3的实现如下。

python 复制代码

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

7.2 Deep Norm

Deep Norm 是微软 2022 年提出的、用于提升深层 Transformer 训练稳定性的标准化方法。按论文中的描述，Deep Norm兼具Post-LN的良好性能及Pre-LN的训练稳定性，作者借此将Transformer的深度扩充到了1000层。清华大学发布的 GLM 和 ChatGLM 就采用了基于 Deep Norm 的 post LN。

Deep Norm的特点如下：

在进行LN之前，以参数 $\alpha$ 增大输入。
在Xavier初始化过程中，以参数 $\beta$ 减小部分参数（全连接层、value 投影层以及输出层的参数）的初始化范围。

其实，就是加了一个系数a，用来平衡残差分支和主干路。具体信息参加下图。

7.3 PRepBN

LayerNorm需要在每个样本的特征维度上计算均值和标准差，这可能在特征维度非常大时导致较高的计算开销，且LayerNorm可以稳定训练。BatchNorm使用训练时的统计均值和方差数据直接计算，导致较低的推理延迟，但可能导致训练崩溃和较差的性能。

论文"SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization"提出了一种PRepBN的新方法，通过使用超参数lamda来控制两个归一化层的比例，在训练中逐步用重新参数化的BatchNorm替换LayerNorm。具体参见下图。

7.4 RealFormer

论文"RealFormer: Transformer Likes Residual Attention"提出 RealFormer 模型（Residual Attention Layer Transformer），如下图所示。该模型将残差结构运用到attention层，使得模型对训练超参更具鲁棒性的同时，保证模型性能的提升。

具体来说，RealFormer相较于前面提到的两种结构（"Pre-LN"和"Post-LN"）不同在于，模型在每层中计算所有头的attention score时，加上了残差结构，即本层的attention score加上之前层的attention score。直接在attention计算时增加跳连连接并不会增加指数级的运算量，因此其效率是相对可观的。

RealFormer沿用了Post-LN的模型设计，只是在每个Transformer层计算多头注意力时，加入前一层的Attention Scores矩阵。即计算第n层的注意力矩阵时，从公式（1）变为了公式（2）。

实现以上计算方式的改变只需要在Transformer的模型代码中做很少的代码改动，并且网络中不止一种类型的attention模块时也适用。

7.5 nGPT

论文"nGPT: Normalized Transformer with Representation Learning on the Hypersphere"提出了一种新颖的神经网络架构：在超球面上进行表示学习的归一化 Transformer。

研究背景&动机

Transformer 架构是现代语言模型的基础，为了提高其训练稳定性、推理成本、上下文长度和鲁棒性等，研究人员提出了大量的修改方案。其中，应用各种归一化技术被认为是有益的，例如添加 LayerNorm 和 RMSNorm 等归一化层，以及通过权重衰减控制权重的范数。同时，也有研究表明在超球面上进行表示学习与更稳定的训练、更大的嵌入空间可分性以及在下游任务中更好的性能相关。在此基础上，论文作者提出了归一化 Transformer，旨在统一该领域的各种发现和观察结果。

核心贡献

作者提出将构成网络矩阵嵌入维度的所有向量归一化，使其位于单位范数超球面上。这样，矩阵 - 向量乘法就可以看作是表示在 [-1,1] 范围内的余弦相似度的点积，从而使权重衰减变得不必要。

归一化 Transformer 本身在超球面上执行多步优化（每层两步），其中注意力和 MLP 更新的每一步都由特征学习率（可学习的变度量矩阵的对角元素）控制。对于输入序列中的每个标记，归一化 Transformer 的优化路径从超球面上对应其输入嵌入向量的点开始，并移动到超球面上最能预测下一个标记嵌入向量的点。

在原始的仅解码器 Transformer 中，标记嵌入向量的范数不受约束，这可能导致不准确的相似度估计。在 nGPT 中，作者提出在训练算法的每一步之后，对存储在和中的嵌入向量进行归一化。同时，由于所有 nGPT 嵌入都是归一化的，原始公式中的逻辑值代表在 [-1,1] 范围内的点积，这限制了 softmax 生成的概率分布的置信度（温度）。因此，作者引入了一个可训练的缩放参数来调整。

对比

层和块（LAYERS AND BLOCKS）

基线 Transformer：对隐藏状态应用层变换，包括交替的自注意力（ATTN）和多层感知器（MLP）块，并使用 RMSNorm 进行归一化。
归一化 Transformer：对于超球面上的任意两点和，可以使用 SLERP 或其近似的 LERP 来计算沿着测地线的插值。作者将其改写为 nGPT 中的更新方程，其中涉及到注意力和 MLP 块的更新方程，通过可学习的参数和以及归一化函数 Norm 来控制更新过程。与基线 Transformer 不同，nGPT 在最后一层之后不需要额外的归一化。

自注意力块

基线 Transformer：注意力机制是 Transformer 的关键组件，它允许每个标记关注序列中的其他标记。在基线 Transformer 中，首先使用 RMSNorm 对输入隐藏状态进行归一化，然后将其投影为查询、键和值，并应用旋转位置嵌入（RoPE）。通过计算查询和键向量的点积，缩放后应用 softmax 函数得到注意力权重，最后计算值向量的加权和。
归一化 Transformer：作者提出对$W_k$、$W_q$、$W_v$和$W_o$沿着其嵌入维度进行归一化，使得与计算的点积可以解释为单位范数向量之间的余弦相似度。此外，还对q和k进行额外的归一化，以确保每个查询和键的点积在控制范围内。同时，调整了 softmax 缩放因子。

MLP 块（MLP BLOCK）

基线 Transformer：MLP 块的输入隐藏状态首先使用 RMSNorm 进行归一化，然后通过两个单独的线性投影产生两个中间向量和，使用 SwiGLU 激活函数进行组合，最后通过一个最终的线性变换得到输出。
归一化 Transformer：作者提出对矩阵$W_u$和$W_v$沿着嵌入维度进行归一化，使得u和v向量分别代表h与存储在$W_u$和$W_v$中的向量之间的余弦相似度。为了控制它们的影响，引入了缩放因子$S_u$和$S_v$。

切换

将基线 Transformer 转换为归一化 Transformer 的步骤包括：移除所有归一化层；在每次训练步骤后，对所有矩阵沿着其嵌入维度进行归一化；替换更新方程；改变注意力中的 softmax 缩放因子并对和进行重新缩放和归一化；对 MLP 块的中间状态进行重新缩放；对逻辑值进行重新缩放；移除权重衰减和学习率预热。

7.6 DyT

论文"Transformers without Normalization" 提出了动态Tanh（DyT）的逐元素操作，旨在让无需归一化层的Transformer 模型也能达到相同甚至更好的性能，且大多数情况下无需调整超参数。

动机

论文的探索始于一个观察：LN层通过类似于tanh的S形曲线将其输入映射到输出，对输入激活进行缩放的同时压缩极端值。受这一启发，论文提出了一种称为动态Tanh（Dynamic Tanh, DyT）的逐元素操作，其定义为：DyT(x) = tanh(αx)，其中α是一个可学习的参数。该操作旨在通过α学习适当的缩放因子，并通过有界的tanh函数压缩极端值，从而模拟LN的行为。值得注意的是，与归一化层不同，DyT无需计算激活统计量即可实现这两种效果。

作者选择了三个模型进行分析，对于这三个训练好的网络，我们采样了一个小批量的样本，并通过网络进行前向传播。然后，我们测量了归一化层的输入和输出，即在归一化操作之前和之后的张量，且不包含可学习的仿射变换。由于LN保留了输入张量的维度，我们可以在输入和输出张量元素之间建立一一对应关系，从而直接可视化它们之间的关系。我们在下图中绘制了这些映射关系。

我们发现输入和输出之间的关系大多是线性的，类似于x-y图中的一条直线。然而，在更深的LN层中，这些曲线的形状非常类似于完整的或部分的S形曲线，类似于tanh函数。

实现

受到归一化层与缩放后的tanh函数形状相似的启发，我们提出了动态Tanh（DyT）作为归一化层的替代方案。给定输入张量x，DyT层的定义如下图标号1所示。

其中，alpha是一个可学习的标量参数，允许根据输入范围的不同进行缩放，从而适应不同的x尺度（如图2所示）。这也是我们将整个操作命名为"动态"Tanh的原因。gamma 和 beta 是可学习的、每个通道的向量参数，与所有归一化层中使用的参数相同------它们允许输出缩放到任意尺度。有时这被认为是一个独立的仿射层；但在我们的设计中，我们将它们视为DyT层的一部分，就像归一化层也包含这些参数一样。

DyT并不是一种新的归一化层，因为它在前向传播过程中独立地对张量中的每个输入元素进行操作，而不计算统计量或其他类型的聚合。然而，它确实保留了归一化层在非线性压缩极端值方面的效果，同时对输入的中心部分进行几乎线性的变换。