大模型——理论基础——常用的Norm_大模型 norm层是什么意思

一、Layer Normalization

1.1 实现原理

Layer Normalization (LayerNorm) 是一种归一化技术，常用于深度学习模型中，特别是在 Transformer 模型中。

与 Batch normalization 不同，Layer normalization 是在特征维度上进行标准化的，而不是在数据批次维度上。

Layer normalization 的计算可以分为两步：

1. 计算均值和方差：对于给定的输入 ，其中 是批次大小， 是特征维度，可以计算每个样本的均值和方差：
   均值：
   方差：
2. 标准化和重新缩放：
   标准化后的输出：
   重新缩放和偏移后的输出：
   其中， 是一个很小的正数，用来防止除以零；和  是可学习的参数，用于重新缩放和偏移。

代码测试如下：

import torch
import torch.nn as nn
 
class LayerNorm(nn.Module):
    def __init__(self,num_features,eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        self.eps = eps
 
    def forward(self,x):
        mean = x.mean(dim=-1,keepdim=True)
        std = x.std(dim=-1,keepdim=True,unbiased=False)
        normalized_x = (x - mean) / (std + self.eps)
        return self.gamma * normalized_x + self.beta
 
if __name__ == '__main__':
    batch_size = 2
    seqlen = 3
    hidden_dim = 4
 
    # 初始化一个随机tensor
    x = torch.randn(batch_size,seqlen,hidden_dim)
    print(x)
 
    # 初始化LayerNorm
    layer_norm  = LayerNorm(num_features=hidden_dim)
    output_tensor = layer_norm(x)
    print("output after layer norm:\n,",output_tensor)
 
    torch_layer_norm = torch.nn.LayerNorm(normalized_shape=hidden_dim)
    torch_output_tensor = torch_layer_norm(x)
    print("output after torch layer norm:\n",torch_output_tensor)

1.2 Post Norm Vs Pre Norm

1.2.1 公式区别

Pre Norm 和 Post Norm 的式子分别如下：

1.2.2 在大模型的区别

Post-LN :是在 Transformer 的原始版本中使用的归一化方案。在此方案中，每个子层（例如，自注意力机制或前馈网络）的输出先通过子层自身的操作，然后再通过层归一化（Layer Normalization）

Pre-LN:是先对输入进行层归一化，然后再传递到子层操作中。这样的顺序对于训练更深的网络可能更稳定，因为归一化的输入可以帮助缓解训练过程中的梯度消失和梯度爆炸问题。

1.2.3 为什么Pre效果弱于Post 

参考上面Pre Norm的公式：

其中第二项的方差由于有 norm 是不随层数变化的，于是 x 的方差会在主干上随层数积累。
到了深层以后，单层对主干的影响可以视为小量，而不同层的统计上是相似的。
于是有 ：

          = 

           

这样训练出来的深层 ResNet or Transformer，深层部分实际上更像扩展了模型宽度，所以相对好训练，但某种意义上并不是真正的 deep.

post-norm   则保证了主干方差恒定，每层对 x 都可能有较大影响，代价则是模型结构中没有从头到尾的恒等路径，梯度难以控制。通常认为会更难收敛，但训练出来的效果更好。

1.2.4 结论

在Bert时代由于层数较浅，往往采用的是Post-Norm

而到了大模型时代，由于Transformer的层数开始加深，为了训练稳定性开始使用Pre-Norm。

1.3 RMS Norm 

RMSNorm层则是通过计算沿着最后一个维度的均方根来归一化输入

并使用可学习的权重向量对归一化后的结果进行缩放。

与RMS Norm是基于LN的一种变体，主要是去掉了减去均值的部分
比于LN，可以发现，不论是分母的方差和分子部分，都取消了均值计算，经作者在各种场景中实验发现，减少约 7%∼64% 的计算时间。

代码实现如下：

import torch
import torch.nn as nn
 
class RMSNorm(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

1.4 DeepNorm

Deep Norm是对LN的的改进，主要有两点改进，其中  和 都是根据模型的Encoder（N）和Decoder（M）层数计算出来的，通过如下方案，作者把模型的层数提升到了1000+。

• DeepNorm在进行Layer Norm之前，会以  参数扩大input输入
• 在Xavier参数初始化过程中以 减小部分参数的初始化范围

DeepNorm的表达式为：

其中， 是一个常数（ >1 ）， 是参数为 的第 个Transformer子层（即注意力或前馈网络）的函数。DeepNet还将残差内部的权重  扩展了常数参数  。

具体方案如下：

下面补充一些文章做的分析，感觉很有借鉴意义。下面这幅图，是作者做的三组翻译任务，Encoder-decoder都是18层。

通过上图，可以发现:

• 普通的Post-LN的梯度很小，在top layer中，甚至没有；
• Post-LN +4k warmup会发现，top layers的梯度比较小，down layers的梯度比较大；
• 而Post-LN -init+4k warmup的梯度一直保持较大值，当然down layers的梯度会更大一些。
• 通过上面三组实验，作者认为 Post-LN 的不稳定性部分来自于梯度消失以及初始化的时候，更新太大，陷入了局部最优，跑不出去了。

 代码实现如下：

import torch
from torch import Size
from typing import Union,List
from torch.nn import LayerNorm
 
Number_Layer = 1000 # Encoder/Decoder
 
class DeepNorm(torch.nn.Module):
    def __init__(self,normalized_shape: Union[int,List[int],Size],eps: float = 1e-5, elementwise_affine: bool = True):
        '''
          Deep Layer Normalization
        :param normalized_shape: input shape from an expected input of size
        :param eps: a value added to the denominator for numerical stability
        :param elementwise_affine:  a boolean value that when set to ``True``, this module
            has learnable per-element affine parameters initialized to ones (for weights)
            and zeros (for biases). Default: ``True``.
        '''
        super(DeepNorm,self).__init__()
 
        self.alpha = (2 * Number_Layer) ** 0.25
        self.layernorm = LayerNorm(normalized_shape,eps=eps)
 
    def forward(self,x):
        x_normed = self.layernorm(x)
        return self.alpha * x + x_normed