一、引言

论文： LLaMA: Open and Efficient Foundation Language Models
作者： Meta AI
代码： LLaMA
特点： 该方法在Transformer的基础上增加了Pre-normalization (RMSNorm)、SwiGLU activation function (SwiGLU)、Rotary Embeddings (RoPE)、FlashAttention。

⚠️ 在学习该方法前，建议补充BatchNorm、LayerNorm、位置编码、Attention的相关知识。

二、详情

Transformer和LLaMA的结构图如下：

可见，其结构差异主要体现在如下方面：

• Transformer采用了左编码器+右解码器（Encoder+Decoder）的结构，LLaMA采用了仅解码器（Decoder-only）的结构。由于仅包含解码器不需要与编码器输出交互，故LLaMA去掉了Transformer中Decoder中间的交叉Multi-Head Attention和Add & Norm。
• LLaMA采用了归一化前置（Pre-normalization）的策略，将归一化操作放在了注意力、FFN前并在线性映射前增加了一个归一化。此外，LLaMA还将LayerNorm替换为了RMSNorm。
• LLaMA将绝对位置编码替换为了旋转位置编码，即RoPE，这是一种只对Q和K进行位置编码的方式。
• 为加速训练，LLaMA引入了FlashAttention。
• LLaMA将ReLU替换为了SwiGLU。

2.1 RMSNorm

均方根归一化RMSNorm简化了LayerNorm的计算。

要了解RMSNorm，首先需回顾LayerNorm的公式：

其中，$\mathbf{x}$为输入的token序列，$\mathbf{E}[\mathbf{x}]=\frac{1}{n}{\sum }_{i=1}^n{\mathbf{x}}_i$和$\mathbf{Var}[\mathbf{x}]=\sqrt{\frac{1}{n}{\sum }_{i=1}^n({\mathbf{x}}_i-\mathbf{E}[\mathbf{x}]{)}^2}$为$\mathbf{x}$的均值和有偏方差，$\mathbf{ϵ}$用来防止分母为0，$\mathbf{\gamma }$和$\mathbf{\beta }$是可学习的参数用来缩放和平移。

RMSNorm简化了LayerNorm的计算，其公式如下：

其中，$\mathbf{RMS}[\mathbf{x}]=\sqrt{\frac{1}{n}{\sum }_{i=1}^n{\mathbf{x}}_i^2}$是均方根。

可见，RMSNorm与LayerNorm主要有如下差别：

• RMSNorm无需计算均值$\mathbf{E}[\mathbf{x}]$。
• RMSNorm将有偏方差$\mathbf{Var}[\mathbf{x}]$替换为了均方根$\mathbf{RMS}[\mathbf{x}]$。
• RMSNorm无需平移项$\mathbf{\gamma }$。

与LayerNorm一样，RMSNorm也能以句子或单词（token）为单位进行归一化，如下给出了以token为单位的代码示例。

import torch
import torch.nn as nn


class MyRMSNorm(nn.Module):
    def __init__(self, hidden_dim, eps=1e-8):
        super().__init__()
        # 防止分母计算为0
        self._eps = eps
        # 仿射变换参数，缩放norm后的数据分布
        self._gamma = nn.Parameter(torch.ones(hidden_dim))

    def forward(self, input):
        # input(N,L,C)
        ms = input.pow(2).mean(dim=-1, keepdim=True)  # 计算均方，token-wise
        input = input / torch.sqrt(ms + self._eps)  # 执行标准化
        return input * self._gamma  # 仿射变换


if __name__ == '__main__':
    batch_size = 4
    length = 2
    hidden_dim = 3
    input = torch.rand(4, 2, 3)

    myRMSN = MyRMSNorm(hidden_dim=hidden_dim)
    MyO = myRMSN(input)

    pytorchRMSN = nn.RMSNorm(normalized_shape=hidden_dim, elementwise_affine=False)  # 不使用可学习的gamma和beta
    pytorchO = pytorchRMSN(input)

    print(MyO == pytorchO)

2.2 RoPE

旋转位置编码RoPE使用绝对位置信息设计旋转规则，使旋转后的数据能够表达相对位置信息。

要了解RoPE，首先我们来了解一下二维空间的旋转。如下图：

其中，$X=[\rho \cos \varphi ,\rho \sin \varphi ]$是一个二维向量，逆时针旋转$\theta$度变成$XR(\theta )$。此时$R(\theta )=\left[\begin{array}{c}\cos \theta ,\sin \theta \\ -\sin \theta ,\cos \theta \end{array}\right]$，证明如下：

$$\begin{gathered} XR(\theta )=[\rho \cos \varphi ,\rho \sin \varphi ]\left[\begin{array}{c}\cos \theta ,\sin \theta \\ -\sin \theta ,\cos \theta \end{array}\right] \\ =\rho [\cos \varphi \cos \theta -\sin \varphi \sin \theta ,\cos \varphi \sin \theta +\sin \varphi \cos \theta ]=[\rho \cos (\varphi +\theta ),\rho \sin (\varphi +\theta )] \end{gathered}$$

可见，$X$与$XR(\theta )$仅差一个$\theta$，所以二维空间逆时针旋转$\theta$度可通过$R(\theta )$实现。

旋转只改变角度，不改变长度。

RoPE将旋转应用在了注意力模块的查询$Q$和$K$上。它将第$i$个查询$Q_i$旋转$i\theta$的角度，再将第$j$个键$K_j$旋转$j\theta$的角度，那么$Q_i{K}_j^T$就会变成一个与相对位置$i-j$相关的值。推导过程如下：

$i$和$j$是查询$Q_i$和$K_j$的绝对位置，$i-j$是它们的相对位置。

然而，$Q_i$和$K_j$的维度通常都是大于2的，我们假设它是$D$且$D$是2的整数倍，于是我们可以将$Q_i$和$K_j$分别划分为$d=\frac{D}{2}$个子空间，每个子空间都是二维的。

下图给出了一个$D=10$的例子，我们将$Q_i$和$K_j$分为5个子空间并分配1个包括5个角度的旋转序列$\Theta =({\theta }_1,{\theta }_2,\cdots ,{\theta }_5)$，每个子空间的旋转角度是在对应旋转序列的基础上乘以$i$或$j$。

将其扩展到$d$个子空间，可以得到如下信息：

其中，$X_i$代指$Q_i$或$K_j$。此时，这种旋转仍然具有相对位置的表达能力，证明如下：

显然，上面的$R(i\Theta )$过于稀疏，为了提升计算效率，通常$d$个子空间的旋转使用下式表达：

为避免token数过多，$i{\theta }_k$和$j{\theta }_k$重叠导致相对位置得不到表达（同一个子空间$k$，绝对位置$i$和$j$不同，$i{\theta }_k-j{\theta }_k=2m\pi$时重叠，$m$是一个整数），RoPE使用了一个递减的等比数列作为$\theta$序列，如下：

${\theta }_k$是递减的，这表示token中前几个子空间的旋转角度较大，越往后旋转角度越小。

事实上，为了方便我们通常不是将相邻的两个值划分至同一子空间，而是将D分为前后两个部分，前后各取一个依次组成子空间，例如[q0,q1,q2,q3]被划分为[q0,q2], [q1,q3]而不是[q0,q1], [q2,q3]。以下为使用这种方式进行子空间划分的RoPE代码：

from torch.nn import functional as F
import torch.nn as nn
import torch
import math


class Rotator:
    """根据hidden_dim，和position_ids 生成对应的旋转位置编码, 和论文中定义略有不同，一个个二维的子空间被
    分割到了前后两部分，分别进行旋转，然后拼接起来
    """
    def __init__(self, D, position_ids):
        """ position_ids: [seq_len], D 和单个头的hidden_dim对应 """
        base = 10000
        d = D / 2
        B = base ** (1/d)
        theta_base = 1.0 / (B ** (torch.arange(0, d)))    # 等比数列， $\Theta$
        thetas = position_ids.outer(theta_base)  # [seq_len, D/2]
        # 这里的子空间划分与讲解不同，[q0,q1,q2,q3] -> [q0,q2],[q1,q3]是两个子空间而不是[q0,q1],[q2,q3]
        full_thetas = torch.cat((thetas, thetas), dim=-1)  # [seq_len, D]
        self.cos = full_thetas.cos()
        self.sin = full_thetas.sin()


    def rotate(self, x):
        """
        x: [bs, num_attention_heads, seq_len, D]
        q: [bs, num_attention_heads, seq_len, D]
        cos: [seq_len, D]
        [x,y] @ [[cos, sin], [-sin, cos]] = [x*cos-y*sin, ycos+x*sin] =[x,y]*cos+[-y, x]*sin
        """
        return x * self.cos + Rotator.reverse_half(x) * self.sin


    @staticmethod
    def reverse_half(q):
        """ q: [bs, num_attention_heads, seq_len, D] trick2 """
        u = q[..., :q.shape[-1] // 2]  # 认为是各个二维子空间的第一维的向量集结
        v = q[..., q.shape[-1] // 2:]   # 认为是各个二维子空间的第二维的向量集结
        return torch.cat((-v, u), dim=-1)


if __name__ == "__main__":
    batch_size = 2
    num_heads = 3
    D = 6  # 单个头的token向量长度
    hidden_dim = D * num_heads
    seq_len = 4
    position_ids = torch.arange(seq_len)
    rotator = Rotator(D, position_ids)

    x = torch.randn((batch_size, seq_len, hidden_dim))
    # 对每个头分别进行旋转，[batch_size,seq_len,hidden_dim] -> [batch_size,seq_len,num_heads,D] -> [batch_size,num_heads,seq_len,D]
    x = x.view(batch_size, seq_len, num_heads, D).transpose(1, 2)
    x = rotator.rotate(x)

2.3 FlashAttention

FlashAttention以分块的形式进行注意力计算，避免了SRAM和HBM之间频繁读写导致的时间浪费。

详情请参考我之前的博客FlashAttention in NeurIPS 2022。

2.4 SwiGLU

激活函数SwiGLU是门控线性单元（Gated Linear Units, GLU）的变体，下图红框中表达了GLU的计算过程：

可见，GLU会先使用两个带偏执的线性层映射输入$\mathbf{x}$，分别记为$\mathbf{x}{\mathbf{W}}_1+{\mathbf{b}}_1$和$\mathbf{x}{\mathbf{W}}_2+{\mathbf{b}}_2$；其中一个线性映射后会跟一个非线性激活函数sigmoid，记为$\sigma (\mathbf{x}{\mathbf{W}}_1+{\mathbf{b}}_1)$；然后将左右两边的结果对应元素相乘即完成了GLU，记为$\sigma (\mathbf{x}{\mathbf{W}}_1+{\mathbf{b}}_1)\otimes (\mathbf{x}{\mathbf{W}}_2+{\mathbf{b}}_2)$。

SwiGLU对GLU做了两点改进：

• 去掉了两个线性映射的偏执项，此时公式变成 $\sigma (\mathbf{x}{\mathbf{W}}_1)\otimes (\mathbf{x}{\mathbf{W}}_2)$。
• 将sigmoid替换为了Swish，此时公式变成 ${\text{Swish}}_{\beta }(\mathbf{x}{\mathbf{W}}_1)\otimes (\mathbf{x}{\mathbf{W}}_2)$。

Swish的公式为${\text{Swish}}_{\beta }(a)=a\sigma (\beta a)=\frac{a}{1+e^{-\beta a}}$，在不同的$\beta$下该非线性激活函数的曲线如下：

可见，当$\beta$较大时，该曲线与ReLU十分接近；当$\beta =1$时，小于0但接近0的曲线变得更光滑且非单调。

SwiGLU则选用了$\beta =1$的Swish，于是我们得到SwiGLU的公式如下：
$$\text{Swish}(\mathbf{x}{\mathbf{W}}_1)\otimes (\mathbf{x}{\mathbf{W}}_2)=\frac{\mathbf{x}{\mathbf{W}}_1}{1+e^{-\mathbf{x}{\mathbf{W}}_1}}\otimes \mathbf{x}{\mathbf{W}}_2$$

致谢：

本博客仅做记录使用，无任何商业用途，参考内容如下：
解密旋转位置编码：数学基础、代码实现与绝对编码一体化探索
一文为你深度解析 LLaMA2 模型架构
Llama改进之——SwiGLU激活函数