【多模态大模型】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的结构图如下:

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

  • 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的公式:

其中, 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=1nxi和 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=1nxi2 是均方根。

可见,RMSNormLayerNorm主要有如下差别:

  • RMSNorm无需计算均值 E [ x ] {\bf E}[\boldsymbol{x}] E[x]。
  • RMSNorm将有偏方差 V a r [ x ] {\bf Var[\boldsymbol{x}]} Var[x]替换为了均方根 R M S [ x ] {\bf RMS[\boldsymbol{x}]} RMS[x]。
  • RMSNorm无需平移项 γ \boldsymbol{\gamma} γ。

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

python 复制代码
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 QiKjT就会变成一个与相对位置 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代码:

python 复制代码
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)。

SwiGLUGLU做了两点改进:

  • 去掉了两个线性映射的偏执项,此时公式变成 σ ( x W 1 ) ⊗ ( x W 2 ) \sigma(\boldsymbol{xW_1})\otimes(\boldsymbol{xW_2}) σ(xW1)⊗(xW2)。
  • sigmoid替换为了Swish,此时公式变成 Swish β ( x W 1 ) ⊗ ( x W 2 ) \text{Swish}_{\beta}(\boldsymbol{xW_1})\otimes(\boldsymbol{xW_2}) Swishβ(xW1)⊗(xW2)。

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−xW1xW1⊗xW2

致谢:

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

相关推荐
韬小志1 天前
【LLaMa-Factory】监督微调训练方法
人工智能·深度学习·llama
大拨鼠2 天前
【多模态读论文系列】LLaMA-Adapter V2论文笔记
论文阅读·人工智能·llama
努力的光头强3 天前
太炸裂了,Ollama跑本地模型已成为历史,现在都在使用这个工具,而且还能集成本地知识库
人工智能·ai·pdf·产品经理·llama
AIBigModel4 天前
LLaMA系列一直在假装开源...
开源·llama
三月七(爱看动漫的程序员)5 天前
Tree of Thoughts: Deliberate Problem Solving with Large Language Models
人工智能·gpt·语言模型·自然语言处理·chatgpt·llama
F5 AI社区7 天前
【AI换脸整合包及教程】AI换脸技术新贵:Rope换脸工具全面解析
ai换脸·视频换脸·rope·ai免费工具
励志成为美貌才华为一体的女子7 天前
基于LLaMA Factory对LLama 3指令微调的操作学习笔记
llama
HyperAI超神经8 天前
对标Hugging Face?GitHub Models新增OpenAI o1/Llama 3.2等, 新功能支持模型并排比较
人工智能·机器学习·github·llama·huggingface
努力的光头强10 天前
人工智能大模型赋能医疗健康产业白皮书(2023年)|附88页PDF文件下载
人工智能·算法·ai·pdf·产品经理·llama
cv2016_DL10 天前
CLIP改进
人工智能·深度学习·机器学习·计算机视觉·llama