【多模态大模型】LLaMA in arXiv 2023

一、引言

论文： 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的结构图如下：

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

2.1 RMSNorm

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

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

其中， x \boldsymbol{x} x为输入的token序列， E [ x ] = 1 n ∑ i = 1 n x i {\bf E}\boldsymbol{[x]}=\frac{1}{n}\sum_{i=1}^{n}\boldsymbol{x}_i E[x]=n1​∑i=1n​xi​和 V a r [ x ] = 1 n ∑ i = 1 n ( x i − E [ x ] ) 2 {\bf Var}\boldsymbol{[x]}=\sqrt{\frac{1}{n}\sum_{i=1}^n(\boldsymbol{x}_i-{\bf E}\boldsymbol{[x]})^2} Var[x]=n1​∑i=1n​(xi​−E[x])2 ​为 x \boldsymbol{x} x的均值和有偏方差， ϵ \boldsymbol{\epsilon} ϵ用来防止分母为0， γ \boldsymbol{\gamma} γ和 β \boldsymbol{\beta} β是可学习的参数用来缩放和平移。

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

其中， R M S [ x ] = 1 n ∑ i = 1 n x i 2 {\bf RMS}\boldsymbol{[x]}=\sqrt{\frac{1}{n}\sum_{i=1}^{n}\boldsymbol{x}_i^2} RMS[x]=n1​∑i=1n​xi2​ ​是均方根。

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

与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 = [ ρ cos ⁡ ϕ , ρ sin ⁡ ϕ ] X=[\rho\cos\phi,\rho\sin\phi] X=[ρcosϕ,ρsinϕ]是一个二维向量，逆时针旋转 θ \theta θ度变成 X R ( θ ) XR(\theta) XR(θ)。此时 R ( θ ) = [ cos ⁡ θ ,   sin ⁡ θ − sin ⁡ θ ,   cos ⁡ θ ] R(\theta)=\left[\begin{matrix}\cos\theta,~\sin\theta\\-\sin\theta,~\cos\theta\end{matrix}\right] R(θ)=[cosθ, sinθ−sinθ, cosθ​]，证明如下：

X R ( θ ) = [ ρ cos ⁡ ϕ , ρ sin ⁡ ϕ ] [ cos ⁡ θ ,   sin ⁡ θ − sin ⁡ θ ,   cos ⁡ θ ] = ρ [ cos ⁡ ϕ cos ⁡ θ − sin ⁡ ϕ sin ⁡ θ , cos ⁡ ϕ sin ⁡ θ + sin ⁡ ϕ cos ⁡ θ ] = [ ρ cos ⁡ ( ϕ + θ ) , ρ sin ⁡ ( ϕ + θ ) ] XR(\theta)=[\rho\cos\phi,\rho\sin\phi]\left[\begin{matrix}\cos\theta,~\sin\theta\\-\sin\theta,~\cos\theta\end{matrix}\right]\\=\rho[\cos\phi\cos\theta-\sin\phi\sin\theta,\cos\phi\sin\theta+\sin\phi\cos\theta]=[\rho\cos(\phi+\theta),\rho\sin(\phi+\theta)] XR(θ)=[ρcosϕ,ρsinϕ][cosθ, sinθ−sinθ, cosθ​]=ρ[cosϕcosθ−sinϕsinθ,cosϕsinθ+sinϕcosθ]=[ρcos(ϕ+θ),ρsin(ϕ+θ)]

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

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

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

i i i和 j j j是查询 Q i Q_i Qi​和 K j K_j Kj​的绝对位置， i − j i-j i−j是它们的相对位置。

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

下图给出了一个 D = 10 D=10 D=10的例子，我们将 Q i Q_i Qi​和 K j K_j Kj​分为5个子空间并分配1个包括5个角度的旋转序列 Θ = ( θ 1 , θ 2 , ⋯   , θ 5 ) \Theta=(\theta_1,\theta_2,\cdots,\theta_5) Θ=(θ1​,θ2​,⋯,θ5​)，每个子空间的旋转角度是在对应旋转序列的基础上乘以 i i i或 j j j。

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

其中， X i X_i Xi​代指 Q i Q_i Qi​或 K j K_j Kj​。此时，这种旋转仍然具有相对位置的表达能力，证明如下：

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

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

θ k \theta_k θ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会先使用两个带偏执的线性层映射输入 x \boldsymbol{x} x，分别记为 x W 1 + b 1 \boldsymbol{xW_1+b_1} xW1​+b1​和 x W 2 + b 2 \boldsymbol{xW_2+b_2} xW2​+b2​；其中一个线性映射后会跟一个非线性激活函数sigmoid，记为 σ ( x W 1 + b 1 ) \sigma(\boldsymbol{xW_1+b_1}) σ(xW1​+b1​)；然后将左右两边的结果对应元素相乘即完成了GLU，记为 σ ( x W 1 + b 1 ) ⊗ ( x W 2 + b 2 ) \sigma(\boldsymbol{xW_1+b_1})\otimes(\boldsymbol{xW_2+b_2}) σ(xW1​+b1​)⊗(xW2​+b2​)。

SwiGLU对GLU做了两点改进：

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

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

SwiGLU则选用了 β = 1 \beta=1 β=1的Swish，于是我们得到SwiGLU的公式如下：
Swish ( x W 1 ) ⊗ ( x W 2 ) = x W 1 1 + e − x W 1 ⊗ x W 2 \text{Swish}(\boldsymbol{xW_1})\otimes(\boldsymbol{xW_2})=\frac{\boldsymbol{xW_1}}{1+e^{-\boldsymbol{xW_1}}}\otimes\boldsymbol{xW_2} Swish(xW1​)⊗(xW2​)=1+e−xW1​xW1​​⊗xW2​

致谢：

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

总结