1. 前置基础

1.1 注意力机制

1.1.1 普通注意力机制

  • 输入: (N,D) 的二维数组,其中 N 是序列的长度(token 数量),D 是嵌入维度(token 的嵌入向量维度)
    img
  • 注意力权重矩阵(WkW_kWqW_qWvW_v)
    img
    • 形状: (D, d{}d_{\{\cdot\}})
    • 除以 HS\sqrt{HS} 是因为使得点积保持在可控的范围内, 否则他们会随着维度的大小而增长
    • dk=dqd_{k}=d_{q}dvd_v 可以不同
    • dkd_{k}的作用是什么? 是为了使得权重矩阵变为一个高斯分布!!, 我们假设QQ, KK中的元素qijq_{ij}kijk_{ij}都是服从标准高斯分布的, 那么进行注意力计算后得到的乘积是服从N(0,dk)\mathcal{N}(0,\sqrt{d_k})
  • Softmax
    • 对矩阵的每一行进行操作

1.1.2 批注意力机制

img

  • 这里的 Wk ...W_k\ ... 的形状和上面是相同的
Danger

注意力矩阵中 (R, C) 单元格的值越高, 就意味着这两个对应索引的 Token 之间的关系越强

1.1.3 多头注意力

img

  • 现在的模型中, 通常由多个头, 每个头分别执行其注意力任务, 最终将这些结果连接起来, 并通过一个线性层
  • 注意, 多头注意力的参数量和单头注意力的参数量是一样的, 因为实际上多头注意力的权重矩阵相当与把单头注意力的权重矩阵给切开了, 所以多头注意力的dd维度更小

1.1.4 掩码注意力

注意力模块可以用于编码器和解码器

  • 对于编码器的部分(适用于语言理解和翻译任务), 每个 Token 关注序列中的其他 token 是有意义的
  • 对于解码器的部分(生成模型), 每个 token 关注序列中的其他 token 是作弊的
  • 我们希望使用掩码来进行防止作弊, 总的来说就是注意力矩阵中的上三角部分全部应该置零
    • 也即在 softmax 操作之前把这些位置全部置为 -np.inf

1.1.5 交叉注意力

img

  • 一个序列的元素会关注另一个序列中的元素

1.1.6. 跨头维度向量化

img
img

1.1.7. 稀疏注意力

625

  1. 减少计算复杂度:通过减少参与注意力计算的位置数,稀疏注意力显著降低了计算复杂度,使得模型能够处理更长的序列。
    • 传统的注意力计算的时间复杂度是O(N2dk)\mathcal{O}(N^2d_{k})级别的, 这样的计算量实际上是相当大的
  2. 减少内存需求:稀疏操作减少了需要存储的注意力权重的数量,从而降低了模型的内存需求。
  3. 提高长距离依赖学习能力:某些稀疏模式(如分层或跳跃连接)可以帮助模型更有效地学习序列中的长距离依赖关系。
  • 标准计算中,计算 QKTQK^T 需要 N×NN \times N 次内积操作,复杂度是 O(N2dk)O(N^2 \cdot d_k)。为了破除这个魔咒,作者引入了一个稀疏掩码集合 SiS_i。对于第 ii 个 Token(qiq_i),它不再与所有的 j{1,,N}j \in \{1, \dots, N\} 计算点积,而是只与属于特定子集 SiS_i 中的 kjk_j 计算。数学定义转变为:

Aij={exp(qikj/dk)mSiexp(qikm/dk),if jSi0,if jSiA_{ij} = \begin{cases} \frac{\exp(q_i \cdot k_j / \sqrt{d_k})}{\sum_{m \in S_i} \exp(q_i \cdot k_m / \sqrt{d_k})}, & \text{if } j \in S_i \\ 0, & \text{if } j \notin S_i \end{cases}

  • 此时,只要保证集合的大小 Si=wN|S_i| = w \ll N,整体计算复杂度就严格降为了 O(Nwdk)O(N \cdot w \cdot d_k),也就是关于 NN 的线性复杂度 O(N)O(N)
  • 这里我们先举第二张图的例子, 对于QQKRN×dkK\in\mathbb{R}^{N\times d_{k}}, 对于QQ的每一行ii只需要对KTK^Tjj, 其中Sij=1S_{ij}=1的做矩阵乘法, 也就是 N×dkN\times d_k 次乘法, 而集合的大小为ww
  • 这样就极大的减少了attention的计算量

1.2 Transformer 架构

img
img

1.2.1 输入

  • 输入是一个词嵌入: 由单词 Embedding 和位置 embedding 相加得到
    • 其中单词 embedding 可以通过 Word2Vec, Glove 等预训练得到, 也可以在 Transformer 中训练得到
    • 位置 embedding, 其中 pos 表示单词在句子中的位置, d 表示 pe 的维度, 2i 表示偶数维度, 2i+1 表示奇数维度img

1.2.2. Encoder 结构

可以看到是由 Multi-Head Attention, Add & Norm, Feed Forward, Add & Norm 组成的

1.2.2.1 Add&Norm

img
Layer Normalization 会将每一层神经元的输入都转成均值方差都一样的,这样可以加快收敛。

1.2.2.2. Feed Forward

两层全连接
img

1.2.3. Decoder 结构

  • Masked-Multi-Head Attention, Add&Norm, Multi-Head Attention, Add&Norm, Feed Forward, Add&Norm

1.2.3.1. 第一个multi-Head Attention

  • 采用masked操作
  • 输入: 一个mask的词嵌入矩阵

1.2.3.2. 第二个multi-Head Attention

  • 输入: Encoder的编码信息
  • 根据 Encoder 的输出 C计算得到 K, V,根据上一个 Decoder block 的输出 Z 计算 Q (如果是第一个 Decoder block 则使用输入矩阵 X 进行计算),后续的计算方法与之前描述的一致。(交叉注意力)

1.2.3.3. Softmax 预测输出单词

img
img

1.3. MoE架构

  • 稀疏MoE层: 代替传统Transformer架构中的前馈神经网络. MoE层包含若干专家, 每个专家本身是一个独立的神经网络. 在实际应用中通常是前馈神经网络, 也可能是更复杂的网络

  • 门控网络或者路由: 这个部分用于决定哪些Token (token) 被发送到哪个专家。例如,在下图中,“More”这个Token可能被发送到第二个专家,而“Parameters”这个Token被发送到第一个专家。有时,一个Token甚至可以被发送到多个专家。Token的路由方式是 MoE 使用中的一个关键点,因为路由器由学习的参数组成,并且与网络的其他部分一同进行预训练。
    img

  • 训练挑战: 虽然 MoE 能够实现更高效的计算预训练,但它们在微调阶段往往面临泛化能力不足的问题,长期以来易于引发过拟合现象。

  • 推理挑战: MoE 模型虽然可能拥有大量参数,但在推理过程中只使用其中的一部分,这使得它们的推理速度快于具有相同数量参数的稠密模型。然而,这种模型需要将所有参数加载到内存中,因此对内存的需求非常高。以 Mixtral 8x7B 这样的 MoE 为例,需要足够的 VRAM 来容纳一个 47B 参数的稠密模型。之所以是 47B 而不是 8 x 7B = 56B,是因为在 MoE 模型中,只有 FFN 层被视为独立的专家,而模型的其他参数是共享的。此外,假设每个Token只使用两个专家,那么推理速度 (以 FLOPs 计算) 类似于使用 12B 模型 (而不是 14B 模型),因为虽然它进行了 2x7B 的矩阵乘法计算,但某些层是共享的。

1.3.1门控函数

为了专家的负载均衡以及稀疏性的平衡
img
img
img

1.3.2. 辅助损失

鼓励所有专家相同的重要性

1.3.2. MoE拓展的Transformer

当扩展到多个设备时,MoE 层在不同设备间共享,而其他所有层则在每个设备上复制。
img

  • 随机路由: 在 Top-2 设置中,始终选择排名最高的专家,但第二个专家是根据其权重比例随机选择的。
  • 专家容量: 我们可以设定一个阈值,定义一个专家能处理多少Token。如果两个专家的容量都达到上限,Token就会溢出,并通过残差连接传递到下一层,或在某些情况下被完全丢弃

1.4 LoRA

LORA(Low-Rank Adaptation,低秩适应)是一种模型的轻量微调方法,通过向模型的部分层添加可训练的低秩矩阵模块,实现模型在特定任务上的能力调整,同时保持原模型参数不变。
img
img

1.4.1. 原理

LORA 的核心思想是: 在不修改模型参数的情况下,通过添加一个额外的可训练模块,对模型行为进行小范围、有针对性的调整。

LORA 的工作机制包括两个核心阶段

  • 微调阶段: 仅对低秩矩阵进行参数调整在大模型中,像注意力层的权重矩阵通常规模很大。如果直接对这些矩阵进行微调,不仅计算开销大,还容易破坏原有能力。
    • 保持原始参数矩阵W不变
    • 在旁边引入两个较小的矩阵
    • 将他们的乘积 ΔW=B A\Delta W=B\ \cdot A 作为对于W的一个附加调整项, W^=W+ΔW\hat{W}=W+\Delta W
    • 在训练过程中仅更新A和B的参数, 保持W的冻结状态
  • 推理阶段: 在计算时叠加 LORA 模块
    • 当x经过线性层时使用带有LoRA调整项的权重 W^\hat{W}

img

1.5 KV cache

本质上KV cache是用存储来换计算速度, KV cache会额外缓存一些KV的输出, 并且会随着输入的长度变大.

  • 整个KV cache 建立的过程大概是这样:
    1. 初始输入: [t0, t1, t2]
    2. 首次计算: K=[K0,K1,K2], V=[V0,V1,V2] → 生成t3
    3. 缓存状态: K=[K0,K1,K2], V=[V0,V1,V2]
    4. 第二次计算: 新Q=Q3
    5. 注意力计算: Attention(Q3, [K0,K1,K2]) → 生成t4
    6. 更新缓存: K=[K0,K1,K2,K3], V=[V0,V1,V2,V3]
    7. 第三次计算: 新Q=Q4
    8. 注意力计算: Attention(Q4, [K0,K1,K2,K3]) → 生成t5
    9. 更新缓存: K=[K0,K1,K2,K3,K4], V=[V0,V1,V2,V3,V4]

alt textalt textalt textalt text

KV Cache 核心节约的时间有三大块:

  • 前面 n-1 次的 Q 的计算,当然这块对于一次一个 token 的输出本来也没有用;
  • 同理还有 Attention 计算时对角矩阵变为最后一行,和 b 是同理的,这样 mask 矩阵也就没有什么用了;
  • 前面 n-1 次的 K 和 V 的计算,也就是上图紫色部分,这部分是实打实被 Cache 过不需要再重新计算的部分。