深度学习论文: LLaMA: Open and Efficient Foundation Language Models

深度学习论文: LLaMA: Open and Efficient Foundation Language Models
LLaMA: Open and Efficient Foundation Language Models
PDF:https://arxiv.org/pdf/2302.13971.pdf
PyTorch: https://github.com/shanglianlm0525/PyTorch-Networks

1 概述

本文介绍了LLaMA，这是一系列基础而先进的语言模型，其参数规模横跨7亿至65亿不等，展现了强大的语言处理能力。研究表明，通过大规模公开数据的训练，LLaMA系列模型成功打破了对专有或受限数据集的依赖，达到了业界最前沿（SOTA）的性能水平。本研究的核心目标是通过显著增加训练中的token数量，开发出在不同推理场景下均能展现出卓越性能的语言模型。

LLaMA模型家族以其多样的参数配置，为语言模型领域带来了新的竞争力量。特别值得注意的是，即便是参数规模仅为GPT-3十分之一的LLaMA-13B版本，也在多数基准测试中超越了GPT-3，展现了其高效与强大。这一成果不仅提升了语言模型的性能边界，更旨在推动LLMs的普及，使得更多研究者能够在单个GPU的资源限制下，轻松访问并深入研究这些大型模型。
在这里插入图片描述
进一步地，在LLaMA系列中，65亿参数的顶级版本在性能上足以与Chinchilla、PaLM-540B等业界顶尖的大型语言模型相抗衡。尤为关键的是，LLaMA的训练完全基于公开数据，秉持开源精神，与许多依赖非公开或未详尽记录数据集的现有模型形成鲜明对比。尽管市场上已存在如OPT、GPT-NeoX、BLOOM和GLM等使用公开数据的模型，但它们在性能上尚未能与PaLM-62B或Chinchilla等顶尖模型相提并论。LLaMA的出现，无疑为语言模型领域注入了新的活力，也为未来的研究和应用开辟了更广阔的道路。

2 Approach

2-1 Pre-training Data

本训练数据集是精心构建的多元化数据集合，其数据来源广泛且覆盖多个领域，具体比例及处理方法如下：

英文CommonCrawl [67%]：

精心预处理了2017年至2020年的五个CommonCrawl数据转储，采用CCNet管道（Wenzek等人，2020），确保数据质量。
在行级别上进行去重操作，减少重复内容。
利用fastText线性分类器识别并剔除非英文页面，保持语言一致性。
通过n-gram语言模型过滤低质量内容，提升数据集质量。
训练线性模型对页面进行分类，保留与维基百科参考相关的页面，丢弃其他低质量页面。

C4 [15%]：

鉴于C4数据集在多样化预处理方面的优势，将其纳入以进一步提升模型性能。
预处理包括去重和语言识别，确保数据纯净。
质量过滤更多依赖于启发式规则，如标点符号、单词和句子数量等，筛选高质量内容。

Github [4.5%]：

从Google BigQuery获取公开Github数据集，筛选符合Apache、BSD和MIT许可的开源项目。
基于行长度和字母数字字符比例等启发式规则过滤低质量文件。
使用正则表达式去除样板文字，如标题等，清理数据。
在文件级别上进行精确匹配去重，确保数据唯一性。

维基百科 [4.5%]：

添加2022年6月至8月期间的维基百科转储，涵盖20种使用拉丁或西里尔字母的语言。
移除超链接、评论和其他格式化样板文字，使数据更纯净。

Gutenberg and Books3 [4.5%]：

包含Gutenberg 项目和ThePile的Books3部分，提供公共领域书籍资源。
在书籍级别上进行去重，移除内容重叠超过90%的书籍，避免数据冗余。

ArXiv [2.5%]：

处理arXiv Latex文件，引入科学领域高质量数据。
去除论文的引言部分和参考文献，专注于核心研究内容。
去除.tex文件中的注释，并内联扩展用户编写的定义和宏，确保内容一致性和完整性。

Stack Exchange [2%]：

引入Stack Exchange数据转储，包含多样化领域的高质量问题和答案。
保留28个最大网站的数据，去除HTML标签，并按答案得分排序，优先使用高质量答案。

分词器：

采用字节对编码（BPE）算法（Sennrich等人，2015），结合SentencePiece（Kudo和Richardson，2018）实现，对数据进行高效分词。
将所有数字拆分成单个数字，并对未知UTF-8字符进行字节级分解，确保分词准确性和灵活性。

整个训练数据集分词后大约包含1.4T个token，其中大部分token在训练中仅使用一次，但计划对维基百科和图书领域数据进行大约两个周期的训练，以充分利用资源。
在这里插入图片描述

2-2 Architecture

基于近期在大语言模型领域的进展，提出的网络架构基于Transformer（Vaswani等，2017），并融入了多项优化改进，这些改进灵感来源于不同的先进模型如PaLM和GPTNeo。

预归一化 [GPT3灵感]

为提升训练稳定性，Llama采用了预归一化技术，对Transformer子层的输入进行规范化处理，以维持参数的范数，有效避免了梯度消失或爆炸的问题。与传统的输出归一化不同，这种策略是在输入端进行的。优化措施参考了GPT-3的经验，并结合了Zhang和Sennrich（2019）提出的RMSNorm归一化函数，旨在提高模型的整体性能。

RMSNorm作为LayerNorm的一种替代方案，发表在《Root Mean Square Layer Normalization》一文中。它的设计初衷是因为LayerNorm的计算成本较高，而RMSNorm在保持性能的同时，能够减少7%到64%的运算量。RMSNorm与LayerNorm的主要区别在于，RMSNorm省略了计算均值的步骤，只专注于方差的计算，通过Root Mean Square来实现归一化，从而简化了计算过程。

在这里插入图片描述

## LlaMa 模型的源代码，网址：https://github.com/facebookresearch/llama/blob/main/llama/model.py
import torch
import torch.nn as nn

class RMSNorm(torch.nn.Module):
    # 定义一个继承自 torch.nn.Module 的 RMSNorm 类
    def __init__(self, dim: int, eps: float = 1e-6):
        """
        初始化 RMSNorm 归一化层。
        参数:
            dim (int): 输入张量的维度。
            eps (float, optional): 为了数值稳定性添加到分母的小值，默认为 1e-6。
        属性:
            eps (float): 为了数值稳定性添加到分母的小值。
            weight (nn.Parameter): 可学习的缩放参数。
        """
        super().__init__()  # 调用父类的初始化方法
        self.eps = eps  # 存储数值稳定性参数
        self.weight = nn.Parameter(torch.ones(dim))  # 初始化可学习的缩放参数

    def _norm(self, x):
        """
        对输入张量应用 RMSNorm 归一化。
        参数:
            x (torch.Tensor): 输入张量。
        返回:
            torch.Tensor: 归一化后的张量。
        """
        # 计算归一化值，使用均方根（RMS）作为分母
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        RMSNorm 层的前向传播。
        参数:
            x (torch.Tensor): 输入张量。
        返回:
            torch.Tensor: 应用 RMSNorm 后的输出张量。
        """
        # 将输入张量转换为 float 类型，应用归一化，然后恢复原始类型
        output = self._norm(x.float()).type_as(x)
        # 应用可学习的缩放参数并返回结果
        return output * self.weight

SwiGLU激活函数 [PaLM启发]

为了提升模型的非线性表达能力和整体性能，将ReLU激活函数替换为SwiGLU激活函数。SwiGLU由Shazeer（2020）提出，并在PaLM等模型中展现出优势。Llama在此基础上进行了微调，采用 $\frac{2}{3} 4d$ 的维度设置，以适应我们的网络架构需求。

SwiGLU（Swish-Gated Linear Unit）是一种结合Swish激活函数与GLU（Gated Linear Unit）机制的激活函数。其公式如下：
在这里插入图片描述
SwiGLU激活函数需要进行三次矩阵乘法运算，相比于ReLU等传统激活函数计算复杂度更高。

然而，研究发现，尽管计算量增加，SwiGLU带来的性能提升显著：

更好的泛化性能：SwiGLU在处理复杂的文本数据时表现出更好的泛化能力。
稳定的梯度流：Swish激活函数平滑的曲线特性有助于稳定梯度，减少梯度消失和爆炸的可能性。

class SwiGLU(nn.Module):
    """
    Swish-Gated Linear Unit
    https://arxiv.org/pdf/2002.05202v1.pdf
    """
    def __init__(self, size):
        super().__init__()
        self.linear_gate = nn.Linear(size, size)
        self.linear = nn.Linear(size, size)
        self.beta = torch.randn(1, requires_grad=True)

        self.beta = nn.Parameter(torch.ones(1))
        self.register_parameter("beta", self.beta)

    def forward(self, x):
        swish_gate = self.linear_gate(x) * torch.sigmoid(self.beta * self.linear_gate(x))
        out = swish_gate * self.linear(x)
        return out

旋转位置嵌入 [GPTNeo创新]

为了更有效地处理序列中的位置信息，摒弃了传统的绝对位置嵌入，转而采用Su等人（2021）提出的旋转位置嵌入（RoPE）。

RoPE的核心思想是通过一系列旋转矩阵来编码序列中不同位置的相对关系，从而使得模型能够捕捉到序列中的位置信息。RoPE 通过将每个位置的嵌入表示为旋转矩阵，实现对位置信息的编码。这种方法使得模型能够通过旋转操作来感知序列中元素的相对位置。同时旋转矩阵保留了位置之间的相对关系，使得模型在处理长距离依赖时更加有效。

在这里插入图片描述
RoPE的关键优势：

相对位置感知：RoPE能够保持位置信息的相对关系，使得相邻位置的编码具有相似性，而远离位置的编码则具有差异性。这种特性使得模型能够更好地感知和利用位置信息，从而提高对序列数据的理解。
外推能力：RoPE可以通过旋转矩阵生成超过预训练长度的位置编码，这意味着模型可以处理比训练时更长的序列，增强了模型的泛化能力和鲁棒性。
与线性注意力机制的兼容性：RoPE不需要额外的计算或参数来实现相对位置编码，这与线性注意力机制（如Transformer中的自注意力机制）兼容，从而降低了模型的计算复杂度和内存消耗。
无需额外参数：与其他需要额外参数或计算的混合位置编码方法（如Transformer-XL、XLNet）相比，RoPE不需要引入额外的参数或计算步骤，这简化了模型的架构并可能提高效率。

RoPE的这些特性使其在处理长序列数据时特别有用，因为它可以有效地捕捉到长距离的依赖关系，同时保持计算效率。然而，RoPE作为一种相对较新的技术，可能还需要进一步的研究和实验来验证其在不同任务和数据集上的表现和稳定性。

# 生成旋转矩阵
def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):
    # 计算词向量元素两两分组之后，每组元素对应的旋转角度\theta_i
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成 token 序列索引 t = [0, 1,..., seq_len-1]
    t = torch.arange(seq_len, device=freqs.device)
    # freqs.shape = [seq_len, dim // 2] 
    freqs = torch.outer(t, freqs).float()  # 计算m * \theta

    # 计算结果是个复数向量
    # 假设 freqs = [x, y]
    # 则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs) 
    return freqs_cis

# 旋转位置编码计算
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    # xq.shape = [batch_size, seq_len, dim]
    # xq_.shape = [batch_size, seq_len, dim // 2, 2]
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)
    
    # 转为复数域
    xq_ = torch.view_as_complex(xq_)
    xk_ = torch.view_as_complex(xk_)
    
    # 应用旋转操作，然后将结果转回实数域
    # xq_out.shape = [batch_size, seq_len, dim]
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(2)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2)
    return xq_out.type_as(xq), xk_out.type_as(xk)

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.wq = Linear(...)
        self.wk = Linear(...)
        self.wv = Linear(...)
        
        self.freqs_cis = precompute_freqs_cis(dim, max_seq_len * 2)

    def forward(self, x: torch.Tensor):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(batch_size, seq_len, dim)
        xk = xk.view(batch_size, seq_len, dim)
        xv = xv.view(batch_size, seq_len, dim)

        # attention 操作之前，应用旋转位置编码
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        
        # scores.shape = (bs, seqlen, seqlen)
        scores = torch.matmul(xq, xk.transpose(1, 2)) / math.sqrt(dim)
        scores = F.softmax(scores.float(), dim=-1)
        output = torch.matmul(scores, xv)  # (batch_size, seq_len, dim)
  # ......