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

原创

已于 2024-08-15 09:45:17 修改 · 1.5k 阅读

13 ·

CC 4.0 BY-SA版权

文章标签：

#llama #SwiGLU #RoPE #FlashAttention #RMSNorm #原理 #代码

于 2024-08-15 09:39:28 首次发布

一、引言

论文： 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的公式：

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

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

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

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

RMSNorm无需计算均值 ${\bf E}[\boldsymbol{x}]$ 。

RMSNorm将有偏方差 ${\bf Var[\boldsymbol{x}]}$ 替换为了均方根 ${\bf RMS[\boldsymbol{x}]}$ 。

RMSNorm无需平移项 $\boldsymbol{\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