跳过正文
  1. 日常记录、技术札记与转载收藏。/

LLaMA 2 解析

目录

我们通过学习LLama2模型,来逐步的将这个模型搭建起来,包含各个组件,以及做预训练,SFT训练,模型推理对话,等全套流程。

模型架构
#

总体概述
#

LLaMA 2 是典型的 Decoder-only Transformer(自回归语言模型),核心思想与 GPT 系列类似,但在细节上有所优化,特别针对大规模训练和高效推理。

  • 输入:tokenized 文本序列

  • 输出:预测下一个 token 的概率分布

总体架构可以概括为:

Input tokens –> Token Embedding + Positional Encoding –> N x Transformer Block –> LM Head –> Output logits

下面,我将从输入到输出,数据经过各个组件,一一讲解,多讲解为什么,为什么要这样设计

Embeddings组件
#

将输入的 token(数字化)映射为向量表示,这里需要了解一下Tokenizer这个模型。

Tokenizer 的任务是:

  • 将文本转成 token(数字索引)

  • 将 token 转回文本

注意:这个映射必须覆盖模型可能遇到的所有单词、子词甚至符号,否则模型训练和推理都会出现未知 token。

在Llama2中使用的是BPE(Byte Pair Encoding),在我们的实际项目中使用transformers库里面的AutoTokenizer,以及tokenizers库。

这里我们举例子看看BPE是怎么工作的。

使用:中文例子 “我爱南巷的花猫”“我爱南巷” 来演示 BPE 编码过程。


原始文本
#

我爱南巷的花猫
我爱南巷

初始化词表
#

  • 每个汉字作为初始 token

  • 加上结束符 <\w>

初始词表:

我, 爱, 南, 巷, 的, 花, 猫, <\w>

文本拆分为 token 序列:

我 爱 南 巷 的 花 猫 <\w>
我 爱 南 巷 <\w>

统计频率并合并(BPE 训练(统计))
#

步骤 1:统计最频繁相邻 token

  • 合并最频繁 pair → 我 爱

  • 更新文本:

我爱 南 巷 的 花 猫 <\w>
我爱 南 巷 <\w>

步骤 2:继续合并

  • 下一轮最频繁 pair → 南 巷 → 南巷
我爱 南巷 的 花 猫 <\w>
我爱 南巷 <\w>
  • 再合并可能的 pair(选择高频或符合词典):

    • 的 花 → 的花(低频,可保留拆分)

    • 花 猫 → 花猫

最终文本:

我爱 南巷 的花 猫 <\w>
我爱 南巷 <\w>
  • 词表示例:
["我爱", "南巷", "的", "花猫", "<\w>"]

编码示例
#

ps:南巷,南巷的花猫都是我

  • "我爱南巷的花猫" → 拆分为 token:
我爱 + 南巷 + 的 + 花猫 + <\w>
token IDs: [0, 1, 2, 3, 4]
  • "我爱南巷" → 拆分为 token:
我爱 + 南巷 + <\w>
token IDs: [0, 1, 4]

可以看到:

  • 高频组合 "我爱""南巷" 被合并为一个 token

  • 稀有组合 "的花猫" 拆分或部分合并

  • 序列长度比按单字符编码短,训练效率更高

这个需要注意,因为在训练这个固定长度的词表的过程,是与我们的训练集是息息相关的,所以我们在使用数据分布差异比较大的数据的时候,需要重新训练一个新的编码器。不然效果会有的折扣。

RMSNorm(Root Mean Square Normalization)
#

在 Transformer 里,**残差连接(Residual Connection)**会导致激活值在层与层之间不断累积。

如果不做归一化,可能出现:

  • 梯度爆炸/消失

  • 训练不稳定

  • 不同层数的值分布不一致

传统解决办法是 LayerNorm

好问题,什么叫做激活值在层与层之间不断累积?

我们现在从数值的分布角度上结合公式来说明这个问题

1、残差连接
#

在 Transformer 的每一层都有 Residual Connection,公式是这样的:

$x_{l+1} = x_{l} + F(x_{l}) $

  1. $x_{l+1}$代表当前层的输出,下一层的输入

  2. $x_{l}$代表当前层的输入

  3. $F(x_{l})$代表当前层经过设计的网络的变换

Image

2、为什么会累积
#

因为每一层都在 叠加

第一层:

$x_{1}=x_{0} + f(x_{0})$

第二层:

$x_{2}=x_{0} + f(x_{0}) + f(x_{1})$

第三层:

$x_{2}=x_{0} + f(x_{0}) + f(x_{1}) + f(x_{2})$

……

可以看到:输入经过多层残差后,前面的值不断累加进去。这样就会导致数值分布变的分散,不集中,反向传播时,梯度也会被放大或变得不稳定,每一层的输入分布不同,导致训练难以收敛。

3、None (无归一化)、LayerNorm、RMSNorm
#

通过一个简单的例子可以说明这一点–在一个小的残差网络里,取把采样层改成 每隔 8 层取一次激活值(比如一个 32 层的小网络,就取第 8、16、24、32 层)的激活值,然后画直方图对比不同归一化方法 (None / LayerNorm / RMSNorm) 的分布。

Image

最直观的,从图中可以非常明显的看到,随着模型的层数增长,激活值开始慢慢变的分散。最明显的在没有任何归一化方法的None实验组,它的分散程度是最大的。加入LayerNormRMSNorm的会相对集中一点。

代码如下:

import torch, torch.nn as nn, matplotlib.pyplot as plt
import numpy as np

# --- RMSNorm 实现 ---
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-8):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
        return self.scale * (x / rms)

# --- 残差模块 ---
class ResidualBlock(nn.Module):
    def __init__(self, dim, norm_type=None):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim)
        self.fc2 = nn.Linear(dim, dim)
        self.act = nn.GELU()
        if norm_type == "layernorm":
            self.norm = nn.LayerNorm(dim)
        elif norm_type == "rmsnorm":
            self.norm = RMSNorm(dim)
        else:
            self.norm = None

    def forward(self, x):
        y = self.norm(x) if self.norm is not None else x
        y = self.fc1(y)
        y = self.act(y)
        y = self.fc2(y)
        return x + y

# --- 网络 ---
class ResidualNetwork(nn.Module):
    def __init__(self, dim=64, n_layers=32, norm_type=None):
        super().__init__()
        self.blocks = nn.ModuleList([ResidualBlock(dim, norm_type=norm_type) for _ in range(n_layers)])
    def forward(self, x, record_layers=None):
        activations = {}
        for i, block in enumerate(self.blocks, start=1):
            x = block(x)
            if record_layers and i in record_layers:
                activations[i] = x.detach().cpu().numpy().flatten()
        return activations

# --- 绘制直方图 ---
def plot_histograms(norm_types=["none", "layernorm", "rmsnorm"], n_layers=32, step=8):
    dim, batch_size = 64, 64
    x = torch.randn(batch_size, dim)

    record_layers = list(range(step, n_layers+1, step))  # 每隔 step 采样
    fig, axes = plt.subplots(len(record_layers), len(norm_types), figsize=(12, 3*len(record_layers)), sharey=True)

    for j, norm in enumerate(norm_types):
        norm_type = None if norm=="none" else norm
        net = ResidualNetwork(dim=dim, n_layers=n_layers, norm_type=norm_type)
        acts = net(x, record_layers=record_layers)
        for i, layer in enumerate(record_layers):
            ax = axes[i,j] if len(record_layers)>1 else axes[j]
            ax.hist(acts[layer], bins=30, density=True, alpha=0.7, color="steelblue")
            ax.set_title(f"{norm} - Layer {layer}")
    plt.tight_layout()
    plt.show()

# 运行绘图
plot_histograms(n_layers=32, step=8)

为什么传统的使用LayerNorm 而不是使用BatchNorm

  1. 输入数据的特点不同
  • BatchNorm

    • 假设一个 batch 内的样本分布相似

    • 用整个 batch 的均值/方差来归一化

    • 图像任务里效果好,因为图片统计分布稳定

  • LayerNorm

    • 单个样本 内的特征维度做归一化

    • 和 batch 大小无关

在 NLP / Transformer 中,输入是序列化的 token 表示,不同 batch 的分布差异很大,BatchNorm 得到的统计值可能不稳定;LayerNorm 不依赖 batch,稳定得多。


  1. 序列长度变化
  • NLP 的输入长度变化大(句子长短不一)

  • BatchNorm 在变长序列上不好用,需要额外对 padding 做 mask,否则统计量会被污染

  • LayerNorm 不受序列长度影响 → 更方便


  1. 小 batch 训练问题
  • 大模型(尤其在分布式训练中)常常 batch 很小(甚至 1~4 个样本)

  • BatchNorm 在小 batch 下统计量非常不稳定,训练会发散

  • LayerNorm 与 batch size 无关,即使 batch=1 也能正常工作


  1. 并行和分布式训练
  • BatchNorm 需要在 GPU/TPU 上 跨设备同步均值/方差(很耗通信)

  • LayerNorm 不需要跨 batch 计算,全在单样本内完成,更易于大规模分布式训练

4、RMSNorm
#

我们来看LayerNorm 和 RMSNorm两者的公式对比:

对于LayerNorm来说:

$y = \frac{x-\mu }{\sqrt{\sigma ^{2} + \epsilon } } \times \gamma + \beta $

其中$\mu$为均值,$\sigma ^{2}$为方差

对于RMSNorm来说,计算方式非常的相似,并且更加简单:

$y = \frac{x }{\sqrt{mean(x^{2}) + \epsilon } } \times \gamma$

只用 均方根 (RMS) 来归一化

不再减去均值$\mu$

只保留一个缩放参数 γ,没有偏置 β

这样做的好处为:

(1) 减少计算量 & 提升速度

  • LayerNorm 需要计算均值、方差、标准差,代价高

  • RMSNorm 只计算均方值,少一次减法和一次方差运算

  • 在大模型里,每一步都能省下显著计算时间和显存带宽


(2) 对均值中心化不敏感

  • 在 Transformer 里,每一层有残差连接(residual connection),输入往往已经被归一化或分布平衡

  • 去掉均值归一化影响不大,模型仍能稳定训练

  • 实际上,均值项经常被 residual shortcut 抵消


(3) 更稳定的训练

  • RMSNorm 只保证向量的“尺度(norm)”一致,不强制分布均值归零

  • 实践表明,它能让梯度传播更稳定,尤其在 超大模型 里表现更好


(4) 节省内存

  • LayerNorm 需要额外存储均值和方差的梯度

  • RMSNorm 计算简单,节省内存,训练时更高效

Grouped Multi Query Attention
#

关于注意力机制,我推荐几个视频帮助大家更好的理解。

self-Attention

https://www.bilibili.com/video/BV1sw4m1k7Gt/?spm_id_from=333.337.search-card.all.click&vd_source=d6bdf037117ac9f8dc86da23d37ee748

1、Self-Attention
#

用上面表格最简单的例子来表述,要找到词与词直接相似度高的,相似度高,那么最后加权也高,相似度低的,加权就低,灰色的部分为mask值,一般取负无穷。这样做的原因是,模型输出我这个字,不直接输出到后面的字,后面的字应该被遮挡。就是说,模型只可以看见前面的内容去预测下一个被遮挡的字。

依旧,我们回归到最原始的,从公式入手,这个是最直接,最直观,最美丽。代码也不过是公式的表达方式而已。

$Q = XW_{Q}$ $K = XW_{K}$ $V = XW_{V}$

$Attention(Q,K,V) = Softmax(\frac{QK^{T} }{\sqrt{d_{k} } })V$

  • 其中我们可以这样理解$QK^{T}$得到每个 token 对其他 token 的相似度,$\sqrt{d_{k}}$是为了进行数值缩放,不要让梯度拉的太大或者太小。

  • Softmax → 转成概率

  • 与 V 相乘 → 得到加权信息聚合

核心思想:每个 token “看”其他 token 的信息,并按相关性加权汇总。

用一个非常简单的代码来表示:

Image

在代码中进行了torch.manual_seed(0)固定,所以生成的X每次运行都是一样的,所以我们的矩阵图每次也都是一样的。

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# ---- 生成随机输入 ----
L, d = 6, 8  # 序列长度6,特征维8
torch.manual_seed(0)
X = torch.randn(L, d)  # 每一行是一个 token 的向量

# ---- 定义 Self-Attention ----
class SimpleSelfAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)

    def forward(self, X):
        Q = self.W_Q(X)  # (L,d)
        K = self.W_K(X)  # (L,d)
        V = self.W_V(X)  # (L,d)

        scores = Q @ K.T / np.sqrt(self.d_model)  # (L,L)
        weights = torch.softmax(scores, dim=-1)   # 注意力权重
        out = weights @ V                         # 加权求和
        return out, weights

# ---- 前向计算 ----
attn = SimpleSelfAttention(d_model=d)
out, weights = attn(X)

# ---- 可视化注意力权重 ----
plt.figure(figsize=(6,5))
sns.heatmap(weights.detach().numpy(), annot=True, cmap='Blues', cbar=True)
plt.title("Self-Attention Weights Heatmap")
plt.xlabel("Key Index")
plt.ylabel("Query Index")
plt.show()

图中 横轴 Key Index纵轴 Query Index

每个元素表示 query token 对 key token 的注意力权重

颜色越深 → 权重越大 → 当前 token “越关注”这个 token

2、Multi-Head Attention
#

这个解释其实多多少少也有问题,更好的方法就是看看下面我们给出的代码。这个在论文中也是作者所做实验得到的一些证明,当然肯定是有其他的划分方式,就比如我是不是可以给每个头加一个可学习的权重参数呢?!!!通过不断的做实验,可视化展示热力图,我们总会得到一个较好的实验结果(絮絮叨叨,哈哈哈哈哈)。

https://www.bilibili.com/video/BV1sN4y1g768/?spm_id_from=333.337.search-card.all.click&vd_source=d6bdf037117ac9f8dc86da23d37ee748

在上一个过程中,我们一直讲的都是单头,但是单头有一些显而易见的缺陷:

注意力模式单一(还是从公式出发)

一个 token 只能对整个序列产生一套注意力分布

这套分布可能同时需要兼顾:

  • 最近邻依赖(语法)

  • 远距离语义相关(核心概念或上下文信息)

单头往往难以同时兼顾近距离与远距离依赖

信息容量有限

  • 单个 head 的子空间维度 $d_{k} = d_{model }$

  • 难以在同一向量空间里表示多种语义关系

长文本累积干扰

  • 随着序列长度增加,softmax 注意力容易分散(权重稀释)

  • 单头可能把注意力分散在很多 token 上,无法突出关键信息

怎么解决呢?

  • 将$d_{model}$拆分为多个子空间,每个子空间作为一个 head

  • 每个 head 有独立的 Q/K/V 参数 → 学习不同注意力模式

  • 多头 attention:多个人分工合作,每个人专注不同方面 → 综合后得到完整理解

依旧回归到数学公式上表达:

$head_{i} = Softmax(\frac{Q_{i}K_{i}^{T} }{\sqrt{d_{k} } })V_{i}$ $i = 1...h$

每个 head 的权重矩阵独立

拼接后:

$[head_{1},...,head_{h}]W_{o}$

一个 token 最终表示融合了 h 种注意力模式 → 可以同时捕获近距离/远距离/语义/上下文等多种依赖。

(可以简单的说就是把一个大矩阵划分为多个大小相等的小矩阵做同样的事情,最后再讲结果按照最后一个维度平均起来)

  • 每个头只处理 64 维子空间

  • 不是随机选择 64 维,而是通过 独立的线性变换 W_Q, W_K, W_V 投影出来

我们也可以用具体的数值来表示:

假设:

  • Transformer 的隐藏维度 $d_{model}=768 $

  • 多头数量 h=12

那么每个 head 的维度:

$d_{k} = \frac{d_{model}}{h} = \frac{768}{12} = 64$

这里有一个点需要注意,我们不是讲一个768维度的信息切分为每个头64。更好的说法是投影变换而来(可以认为是逻辑切分,但是在物理空间上还是连续的)。这个可以看代码也更好的理解,下面用一个矩阵来表述一下:

Image
Image
Image

我们用一个序列长度为6,特征维度为16的特征矩阵,用4个头,代码表述,每个头的注意力得分:

Image

这里在展示一下四个头合并后取平均的注意力得分,注意这个地方仅仅是展示探讨一下,在实际中,4个矩阵合并在一起heads_concat = torch.cat(heads, dim=-1),再经过一个线性投射层就行out = self.W_O(heads_concat) 。

Image
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# ---- 输入数据 ----
L, d_model = 6, 16  # 序列长度6,特征维16
torch.manual_seed(0)
X = torch.randn(L, d_model)
X = X.unsqueeze(0)  # 增加 batch 维度,shape=(1, L, d_model)
batch_size = X.shape[0]

# ---- Multi-Head Attention ----
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_k = d_model // n_heads

        # 每个 head 的 Q/K/V 映射
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, X):
        # X: (batch, L, d_model)
        Q = self.W_Q(X)  # (batch, L, d_model)
        K = self.W_K(X)
        V = self.W_V(X)

        # 拆分多头
        Q = Q.reshape(batch_size, L, self.n_heads, self.d_k).permute(0, 2, 1, 3)  # (batch, n_heads, L, d_k)
        K = K.reshape(batch_size, L, self.n_heads, self.d_k).permute(0, 2, 1, 3)
        V = V.reshape(batch_size, L, self.n_heads, self.d_k).permute(0, 2, 1, 3)

        all_weights = []
        heads = []

        for h in range(self.n_heads):
            q = Q[:, h, :, :]  # (batch, L, d_k)
            k = K[:, h, :, :]  # (batch, L, d_k)
            v = V[:, h, :, :]  # (batch, L, d_k)

            # 注意力分数
            scores = q @ k.transpose(-2, -1) / np.sqrt(self.d_k)  # (batch, L, L)
            weights = torch.softmax(scores, dim=-1)               # (batch, L, L)
            out = weights @ v                                     # (batch, L, d_k)

            heads.append(out)
            all_weights.append(weights.squeeze(0).detach().numpy())  # 去掉 batch 维度方便可视化

        # 拼接多头输出
        heads_concat = torch.cat(heads, dim=-1)  # (batch, L, d_model)
        out = self.W_O(heads_concat)            # (batch, L, d_model)
        return out, all_weights

# ---- 前向计算 ----
n_heads = 4
mha = MultiHeadAttention(d_model=d_model, n_heads=n_heads)
out, weights = mha(X)

# ---- 可视化每个头的注意力权重 ----
fig, axes = plt.subplots(1, n_heads, figsize=(4*n_heads,4))
for i, w in enumerate(weights):
    sns.heatmap(w, annot=True, cmap='Blues', ax=axes[i])
    axes[i].set_title(f"Head {i+1}")
    axes[i].set_xlabel("Key Index")
    axes[i].set_ylabel("Query Index")
plt.tight_layout()
plt.show()

Multi-Head Attention = “多个人看同一句话,每个人关注不同信息 → 汇总后得到丰富 token 表示”

核心优势:分工协作 + 自由组合原始特征 → 捕捉复杂语义关系

实现关键:线性映射 + reshape/permute + 多头注意力计算 + 拼接映射

3、Grouped Multi-Query Attention (GQA)
#

讲到这个,这里要和大家区分一下误区:很多教程会直接告诉你(GQA可以将 head 分组,每组 head 共享同一组 K/V,但仍然保留各自的 Q,这样的方式来节约显存)但是这个说法不是很准确。这里需要严格的区分训练和推理。

在训练时,每个 step 都要重新计算完整的 Q/K/V,不会缓存 Key/Value。

MHA (Multi-Head Attention)

  • Q、K、V 都是独立的线性层,维度 (B,L,$d_{model}$)

  • Q: (B,L,h,$d_{k }$)

  • K/V: (B,L,h,$d_{k }$)

  • 计算量:和 head 数量 h 成正比

  • 显存:主要花在存储 Q/K/V(训练时没缓存概念)

MQA (Multi-Query Attention)

  • Q:仍然每个 head 独立

  • K/V:只有 1 份(共享所有 head)

  • 计算量:Q 计算不变,K/V 映射减少了 h 倍

  • 显存:K/V 显存占用减少,但整体相比 MHA 的训练显存节省不算非常大

GQA (Grouped MQA)

  • Q:每个 head 独立

  • K/V:分组共享,g 份,而不是 h 份

  • 计算量:K/V 计算量介于 MHA 和 MQA 之间

  • 显存:K/V 显存也减少,但仍需存 Q(大头还是在梯度和激活)

在推理时(尤其是自回归生成),情况完全不同:

  • 每生成一个新 token,Q 是新的,但 K/V 来自历史 token,需要缓存 (KV Cache)

  • KV Cache 占用显存是推理的主要瓶颈

对比:

  • MHA

    • 每个 head 都存一份 K/V

    • KV cache 大小:O(L x h x $d_{k}$)

    • 随着上下文长度 L 增加,显存急剧膨胀

  • MQA

    • 所有 head 共用一份 K/V

    • KV cache 大小:O(L x $d_{k}$)

    • 比 MHA 小 h 倍,推理速度更快,显存占用极低

  • GQA

    • 分组共享 K/V,假设 g 组

    • KV cache 大小:O(L x g x $d_{k}$)

    • 介于 MHA 和 MQA 之间,既节省显存又保留一定多样性

这里相信细心的你就会有疑问了?

“既然不同的 head 表示不同类型的注意力模式(语法、语义、长程依赖等),为什么还能让它们共享 Key/Value,而不让表达能力塌缩?”

我们依然回到公式出发:

$head_{i} = Softmax(\frac{Q_{i}K_{i}^{T} }{\sqrt{d_{k} } })V_{i}$ $i = 1...h$

每个 head 的输出之所以不同,是因为它们的 Q/K/V 投影矩阵不同:

$Q = XW^{i}_{Q}$ $K = XW^{i}_{K}$ $V = XW^{i}_{V}$

代码中直观的体现就是为

self.W_Q = nn.Linear(d_model, d_model)
self.W_K = nn.Linear(d_model, d_model)
self.W_V = nn.Linear(d_model, d_model)
self.W_O = nn.Linear(d_model, d_model)

Query 决定了“看哪里”(要关注的模式)

Key/Value 决定了“被看见的是什么”

而实验表明:

对于同一个输入 X,不同 head 的 K/V 表征往往高度相似,差异主要来自 Q 的查询方向。

也就是说,不同 head 虽然计算不同的注意力模式,但它们访问的“信息库”(K/V)非常接近。

换句话说,多个 head 在 K/V 空间中的“感知视角”高度重叠。

也就是说在高维度的比如一个token 768维度的空间信息上,里面有大量的维度信息是重合的(冗余信息非常多),这里并不是说这些特征信息不重要,而是说从降低算力开销的角度上来思考,在平衡效果的前提下可以这样做。正如我们之前提到到LayerNorm与RMSNorm一样本质上两个公式非常相似,但从实验分析上来说可以进一步去除增益效果不大的部分。这就是RMSNorm的来源。

怎么进行分组呢?

这里推荐观看原文:GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

https://arxiv.org/abs/2305.13245

Image

GQA方案是基于前面我们提到的MHA,MQA的中间折中方案。但是为了可以方便矩阵乘法可以直接运算,需要将K V矩阵的数量复制到和Q矩阵相等的数量,如下图:

我们从代码上来看实现是非常简单的,后面也会一步一步带着写代码‘’

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    # 获取输入张量的形状:批量大小、序列长度、键/值对头的数量、每个头的维度大小
    bs, slen, n_kv_heads, head_dim = x.shape

    # 如果重复次数为1,则不需要重复,直接返回原始的张量
    if n_rep == 1:
        return x

    # 对张量进行扩展和重塑操作用以重复键值对

    return (
        x[:, :, :, None, :]  # 在第四个维度(头的维度前)添加一个新的维度
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)  # 将新添加的维度扩展到n_rep大小,实现重复的效果
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)  # 重新塑形,合并键/值对头的数量和重复次数的维度
    )

旋转位置编码
#

1、为什么要用位置编码
#

Transformer 天然不懂“顺序”,无法想图像那样,天然具体空间特征知道前后左右像数值是那个(本质上可以当做是一个二维矩阵,只要知道分辨率多少,天然就可以得到所有像素的索引。)下面我们举一个例子:

Transformer 的输入是一串向量(token embedding),例如:

这些向量被送进注意力机制(Self-Attention),但注意力只看到 一组向量的集合,它不会知道 e₁ 是第一个,e₃ 是第三个。这个从数学公式上可以很明显的看出,每个向量只知道和所有包括自己相互点乘,这个阶段不会知道任何有关于空间位置的信息存在。

对模型来说,「我喜欢南巷的花猫」和「花猫喜欢南巷的我」如果没有位置信息,是一样的!

举几个变体:

这些句子里的词一模一样,但顺序不同。 人类理解句子主要依靠语序(句法结构), 而 Transformer 不像 CNN/RNN 沿时间逐步扫描,它是同时看到所有词, 如果不告诉它“第一个词是谁”第二个词是谁”, 它无法知道谁修饰谁、谁是主语谁是宾语。

位置编码就是在告诉模型“谁在哪儿”

位置编码(Positional Encoding)做的事非常简单: 在每个词的 embedding 上加入一个位置向量, 让模型区分:

  • “我”在第 1 个位置

  • “喜欢”在第 2 个

  • “花猫”在第 5 个

这就像每个词带了一个“坐标标签”。

举个简化示意:

这样解释是不是很清晰易懂了,在原来特征向量中,加入位置向量的信息。这个做法是当前主流大模型非常通用的做法,但是仍然有一些差异,有一些模型只是在输入的时候加上位置信息,有的模型在每一个注意力层都加入了位置编码,LLama2模型正是这样做的

Image

这里如果大家仔细的话,就会发现,仅仅只有Q,K矩阵加入了位置编码的信息,那么V矩阵不需要吗,不是是QKV三个矩阵,里面的特征信息都是一样的X吗?

这里就是需要我们再次强调一下,QKV三个矩阵到底在做什么了,我们再次回到公式:

$Attention(Q,K,V) = Softmax(\frac{QK^{T} }{\sqrt{d_{k} } })V$

我们在做什么呢,我们通过QK点积,获取相似度得分,通过softmax转化为概率分布,然后再去乘V矩阵。相当于是softmax转化QK为注意力得分,任何再去V里面让数值大的更加突出,数值小的更加小,相当于是加强和抑制的做用。

更加通俗的来说:

Q(Query):当前词想“去找”谁的信息–》需要知道自己的位置

K(Key):每个词的“被找”标签–》需要知道自己的位置

V(Value):被取回的内容(语义特征)–》不需要知道位置,只要内容

Q、K 负责匹配,它们之间的点积决定“谁和谁相关”;

V 只是携带信息,根据 QK 的匹配结果被加权平均后输出。

碎碎念—-致自己

这里应该就懂了,说的很详细了。这个公式是不是看吐了,感觉很简单。那当然不是了,这个也是经过很多实验,验证出来的,如果没有验证出来,那必然会有很多小黑粉喷,虽然现在也有很多人喷,但这又何尝不是前进的动力呢?!!!人生就是这样,在否定之否定中旋转上升前行。曾经我们深信不疑的,比如爱情,年少时,以为深爱就能长长久久幸福过一辈子,可却是匆匆忙忙破碎的结局收场。后来再遇到喜欢的人,心性也更加稳定了,也懂得一点经营感情,却也需要面临更多现实的原因。失败常有,但促使自己变得更加坚韧不屈的部分也成为自己的一部分。可正因为如此,我们才慢慢明白,所谓成长,从来不是获得什么,而是逐渐学会接受什么。接受不完美的自己,接受有些事情无论多努力都无法改变,接受有些人注定只是人生路上的过客。可接受并不等于妥协,而是学会带着清醒去热爱,带着伤痕去前行。就像那些公式、那些实验——看似枯燥重复,却在一次次失败和修正中逼近真理。人生何尝不是如此?我们不断试错、修正,再试错,再修正,最后也许并没有得到一个“标准答案”,但却在过程中,成为了更完整的自己。所以啊,不怕被质疑,不怕被误解。只要方向不偏,哪怕走得慢一点,也是在往前走。那些挫折和坎坷,终有一天都会变成你眼底的光。

Image

2、为什么要用旋转位置编码
#

之前,在RoPE之前用的是什么位置编码,为什么不用了

绝对位置编码(Absolute Positional Encoding)

先上公式:

$PE_{i,2k} = \sin (\frac{i}{10000^{\frac{2k}{d_{model} } } } )$

$PE_{i,2k+1} = \cos (\frac{i}{10000^{\frac{2k}{d_{model} } } } )$

也就是说:

  • 每两个维度(2k 和 2k+1)组成一个二维平面;

  • 在这个平面上,随着位置 i 增加,这个点会“绕圆转动”;

  • 不同的 k(不同维度对)旋转速度不同(频率不同)。

换句话说:

每一对维度 (2k, 2k+1) 代表一个不同频率的“旋转编码器”, 整个 768 维的向量就是由 384 个不同频率的旋转组合而成。

这里举一个简单的例子(假设d=6):

我们有维度编号:0, 1, 2, 3, 4, 5

→ 共 3 对维度:

  • 第 1 对:(0,1)

  • 第 2 对:(2,3)

  • 第 3 对:(4,5)

对第 i 个位置:

结果是:

  • 前面的维度变化得快(对近距离敏感);

  • 后面的维度变化得慢(对远距离敏感);

  • 组合起来,模型既能感知局部顺序,也能感知整体位置。

为什么要用旋转位置编码

  • Transformer 的 self-attention 本质上靠 Q·K 来决定“谁关注谁”。如果把位置信息放到 Q 与 K 上,那么注意力分数(点积)可以直接包含相对位置信息。

  • 旋转(而非加法)实现了这样一个漂亮的代价很低的性质:两个向量分别在位置 p 和 q 旋转后,它们的点积仅依赖 q−p(相对位移)。也就是说注意力权重是相对位置敏感的。

  • 旋转使用三角函数(cos, sin),是连续的、参数少且容易向更长序列外推(相比纯可学习位置向量更友好)。

举例:

“我(1) 喜欢(2) 南巷(3) 的(4) 花猫(5)” —— RoPE 让“喜欢”(2) 和 “我”(1) 的 Q·K 得分只与距离 1 有关,而不是“它们分别在哪里”的绝对索引。

先来看公式证明吧:

设 q 和 k 是同一 head 下未旋转的向量(维度被成对看作 2D 块)。对第 i 对(block)规定一个位置依赖的旋转角:

$\theta _{p,i} = p\cdot \omega _{i} $

其中 $\omega _{i} $ 是第 i个频率(通常 $\omega _{i}=\frac{1}{10000^{\frac{2i}{d} } } $ 的某种变体)。

对每个 2D 块用旋转矩阵 R(θ):

$R(\theta) =\begin{bmatrix} \cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{bmatrix}$

把位置 P 的 Q、K 分别旋转:

$Q_{p} = R(\theta_{p})Q$, $K_{q} = R(\theta_{q})K$

关键代数性质(旋转矩阵的正交性):

$Q^T_{p}K_{q} = Q^TR(\theta_{p})^TR(\theta_{q})K = Q^TR(\theta_{q} - \theta_{p})K$

因为,$\theta_{q}-\theta_{p} = (q-p)w$所以点积只依赖于 q−p(相对位移)——这正是 RoPE 的“魔法”根源。

(对多个 2D 块分别成立,整体向量的点积是各块总和。)

复数视角(更直观):把每对维度视为复数 z=a+ib旋转对应乘以$e^{i\theta}$.两个复数各乘$e^{i\theta_{p}}$与$e^{i\theta_{q}}$。它们的内积只涉及 $e^{i(θ_{q}−θ_{p})}$。

好吧,我也看的懵懵懂懂。

RoPE 的核心思想就是:

把每个词向量按照它在句子里的“位置”旋转一点角度,

这样模型在计算两个词的相似度(Q·K)时,

自然能知道它们之间相隔多远(相对位置)。

举例子:

于是:

  • “我”的向量朝正右;

  • “喜欢”的向量朝右下;

  • “南巷”的向量朝下;

  • “花猫”的向量朝左下。

你现在看它们的夹角, 就知道它们在句子里相隔多远了。 这就是 RoPE 的妙处: 两个词的角度差,就是它们的相对位置!

当然,怎么计算就是靠上面这些公式了,得拿起线性代数好好研究一番了。

这样是不是懂了,哈哈哈哈哈,这个不比纯+好很多吗,天然自带。当然还是实验证明了,这个做法是有效果的。回过头去看看公式,发现学好数学真重要,有idea,却实现不出来。

前馈网络(Feed Forward Network, FFN)
#

1、Feed Forward
#

这个就相较于之前的内容,就非常的简单了。

Image

但是大家不知道想过没有,为什么要一个FFN呢,可以不要吗?

当然可以啊,模型依然可以跑通,但是效果肯定是会大大折扣的!!!注意力擅长“跨 token 信息交互”,而 FFN 擅长“单 token 内部特征处理”

如果单纯只放注意力机制,那么

问题

  • 注意力本身是线性的加权求和

  • 如果没有非线性变换,模型只能做“线性组合”,无法生成复杂的非线性特征

  • 对单个 token 内部的特征组合能力有限

在后面接一个FFN可以做到:

  • 逐 token 处理:对每个 token 独立计算,不依赖其他 token

  • 非线性增强:通过激活函数(GeLU、SwiGLU 等)引入非线性

这里同时也说明了一点,比如一个token维度是768,是不是768这些维度信息也会交换处理。让信息之间充分的交互。这里可以回顾一下拆分多头,这里是不是就可以很好的解释,为什么KV矩阵里面信息是会有重合,冗余的,这里是非常关键的一点。

从代码上也可以很好的看出这一点:

class MLP(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
        super().__init__()
        # 如果没有指定隐藏层的维度,我们将其设置为输入维度的4倍
        # 然后将其减少到2/3,最后确保它是multiple_of的倍数
        if hidden_dim is None:
            hidden_dim = 4 * dim
            hidden_dim = int(2 * hidden_dim / 3)
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        # 定义第一层线性变换,从输入维度到隐藏维度
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        # 定义第二层线性变换,从隐藏维度到输入维度
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        # 定义第三层线性变换,从输入维度到隐藏维度
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
        # 定义dropout层,用于防止过拟合
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 前向传播函数
        # 首先,输入x通过第一层线性变换和SILU激活函数
        # 然后,结果乘以输入x通过第三层线性变换的结果
        # 最后,通过第二层线性变换和dropout层
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))

扩维 → 每个 token 的特征被“拆开”成更多维度

非线性 + 门控 → 不同维度信息进行选择性增强或屏蔽

降维 → 特征重新混合回原来的维度空间

2、 SwiGLU是Gated Linear Unit (GLU) 的变种
#

上公式(其实以前我也不喜欢看公式,但是慢慢发现公式其实还更好理解):

$SwiGLU(x)=(W_{a}x)⊙SiLU(W_{b}x)$

$SiLU(x)=x⋅\frac{1}{1+e^{-x} } $

线性分支 $a=W_{a}x$ → 线性组合原 token 内部 768 维

门控分支 $b=SiLU(W_{b}x)$ → 非线性 + 门控

逐元素相乘 $a⊙b$ → 不同维度信息被选择性放大或抑制

乘法就像给每个维度加了一个开关:

  • 如果 $b_{i}≈0$ → 这个维度被屏蔽

  • 如果 $b_i \approx 1$ → 这个维度保留

不同于普通激活,SwiGLU 可以“控制信息流”,让网络学会选择哪些维度重要

这个比较简单,就不过多解释了。

Linear层与softmax层
#

1、最后的 Linear 层(输出投影)
#

公式:

$ Z = H W_\text{out} + b$

  • $W_\text{out} \in \mathbb{R}^{d_\text{model} \times V}$,V = 词表大小

  • $b \in \mathbb{R}^{V}$

  • 作用:把每个 token 的隐藏向量 投影到词表空间

换句话说:每个 token 的向量 $h_t \in \mathbb{R}^{d_\text{model}}$ → 生成 $v_t \in \mathbb{R}^{V}$,每个维度对应一个词在词表中的“未归一化得分(logits)”。

  • Linear 层就是把高维语义向量映射到每个词的分数空间

2、 Softmax 层(生成概率)
#

公式:

$P(\text{token}=i \mid \text{context}) = \frac{\exp(z_i)}{\sum_{j=1}^{V} \exp(z_j)}$

  • 将 Linear 输出的 logits 转换成 概率分布

  • Softmax 确保所有 token 的概率之和 = 1

  • 这个概率就是 语言模型预测下一个 token 的依据

自此我们学会了,所有组件是怎么来的,怎么设计的,为什么要这样设计。我们目前仅仅是学会了模型框架的设计,我们其实更加重要的是知道这个任务形式是怎么样的,我们怎么训练,训练数据是怎么样的,怎么去做loss。对于网络架构来说,可以简单抽象为一个特征提取器。将内容抽象为计算机可以理解的特征。如此而已。

后续我们详细介绍,每个组件代码是怎么编写的,数据到底是长什么样子,怎么去处理的。

后面章节就是详细的实际操作,我们的CodeLab

相关文章