Llama 模型架构解析：从 Pre-RMSNorm 到 GQA 的技巧演进

前言

Llama（Large Language Model Meta AI）是由 Meta AI 团队研发的一系列开源大语言模型。与许多闭源商业模型（如 GPT 系列）不同，Llama 系列在发布时给出了模型权重和训练细节（尽管早期版本有使用限制），使得研究人员、企业和个人开发者都能在其基础上进行微调、部署和创新。

架构

Llama的架构整体上遵循了Transformer的Decoder-only设计，这与GPT系列类似。然而，它在一些关键组件上进行了重要的改进和优化。

Llama模型架构的核心组成部分：就是以下

Pre-RMSNorm

Llama采用了**前置层归一化（Pre-Layer Normalization）**策略，即在每个子层（自注意力层或前馈网络）之前进行归一化。同时，它启用了一种更高效的归一化方式——RMSNorm，来代替传统的LayerNorm。

• 前置层归一化：将归一化层（如LayerNorm）放在子层（自注意力层或前馈网络）之前，而非传统后置归一化（Post-normalization）的输出后。这种策略能缓解梯度消失或爆炸问题，提升深层网络稳定性。

• RMSNorm：一种轻量级归一化途径，仅通过输入特征的均方根（Root Mean Square, RMS）进行缩放，省略均值计算，提升了计算效率。公式如下：
  $$\text{RMSNorm}(\mathbf{x})=\gamma \odot \frac{\mathbf{x}}{\text{RMS}(\mathbf{x})+ϵ}$$
  其中：

  • $\text{RMS}(\mathbf{x})=\sqrt{\frac{1}{d}{\sum }_{i=1}^d{x}_i^2}$（输入向量的均方根）。
  • $\gamma$为可学习参数向量。
  • $\odot$表示逐元素相乘。

位置编码：旋转位置编码（RoPE）

• LLaMA 没有使用原始 Transformer 中的绝对位置编码或可学习的位置嵌入。
• 而是采用 RoPE，将 token 的绝对位置信息编码为旋转矩阵，并在计算注意力分数时以相对位置的方式作用于查询和键向量，从而实用建模长距离依赖。

分组查询注意力 （Grouped-Query Attention，GQA）

从 Llama 2 开始，模型引入了GQA来优化推理性能。GQA 通过将查询头分为若干组，每组内的所有查询头共享同一对键头和值头，实现计算量与模型质量的折中。具体而言：

• MHA（多头注意力）：每个查询头独立计算键值，捕捉丰富的语义信息，但参数量和计算成本过高。
• MQA（多查询注意力）：所有查询头共享同一套键值矩阵，显著提升推理速度，但因过度简化导致精度损失。
• GQA（分组查询注意力）：介于两者之间，通过分组共享键值，既减少计算量，又保留不同头之间的语义多样性。

SwiGLU激活函数

• 在两层全连接网络中间，使用激活函数SwiGLU（而非 ReLU 或 GELU）

• 公式如下：
  $$\text{SwiGLU}(\mathbf{x})=(\mathbf{x}{\mathbf{W}}_1)\odot \text{Swish}(\mathbf{x}{\mathbf{W}}_2)$$
  其中，$\text{Swish}(\mathbf{x})=x\cdot \sigma (\beta \mathbf{x})$，$\beta$为标量，通常为1，$\sigma$表示 Sigmoid 函数。

• 优势：相比ReLU，能捕捉更复杂的特征模式，提升模型表达能力。