[{"content":" 一句话：Mask 是在计算 softmax 之前，把某些位置的注意力分数强制设为 $-\\infty$，使 softmax 后这些位置的权重变为 0，相当于\u0026quot;屏蔽掉\u0026quot;不该看的位置。\n一、为什么需要 Mask？ Attention 的计算是：\n$$ \\text{scores} = \\frac{QK^T}{\\sqrt{d_k}}, \\quad \\text{weights} = \\text{softmax}(\\text{scores}), \\quad \\text{output} = \\text{weights} \\cdot V $$\n默认情况下，每个 token 的 Q 会和序列中所有 token 的 K 做点积，包括：\n无意义的 [PAD] 填充 token 未来还没生成的 token（训练时） 这两种情况都需要用 Mask 屏蔽掉。\n二、两种 Mask 2.1 Padding Mask（填充遮罩） 问题：一个 batch 里不同句子长度不同，需要用 [PAD] token 补齐到相同长度，但 [PAD] 是无意义的，不应该被 Attention 到。\nbatch 里两个句子（补齐到长度5）： 句子1: [\u0026#34;今\u0026#34;, \u0026#34;天\u0026#34;, \u0026#34;好\u0026#34;, \u0026#34;[PAD]\u0026#34;, \u0026#34;[PAD]\u0026#34;] ← 实际长度3，补了2个PAD 句子2: [\u0026#34;天\u0026#34;, \u0026#34;气\u0026#34;, \u0026#34;真\u0026#34;, \u0026#34;不\u0026#34;, \u0026#34;错\u0026#34; ] ← 实际长度5，无需补 Padding Mask（1=有效位置，0=PAD位置）： 句子1: [1, 1, 1, 0, 0] 句子2: [1, 1, 1, 1, 1] 屏蔽效果：\n句子1的注意力分数矩阵（屏蔽前）： 今 天 好 PAD PAD 今 → [0.8, 0.3, 0.5, 0.2, 0.1] 天 → [0.4, 0.9, 0.3, 0.3, 0.2] 好 → [0.5, 0.4, 0.7, 0.1, 0.1] 屏蔽后（PAD列设为 -∞）： 今 天 好 PAD PAD 今 → [0.8, 0.3, 0.5, -∞, -∞ ] 天 → [0.4, 0.9, 0.3, -∞, -∞ ] 好 → [0.5, 0.4, 0.7, -∞, -∞ ] softmax 后 PAD 列权重 = 0，完全不影响输出 适用场景：训练和推理（有 padding 时）都需要。\n2.2 Causal Mask（因果遮罩，也叫 Look-ahead Mask） 问题：GPT 类自回归模型训练时，预测第 t 个 token 时不能看到未来的 token，否则就是\u0026quot;作弊\u0026quot;——模型直接抄答案，学不到任何东西。\n序列: [\u0026#34;今\u0026#34;, \u0026#34;天\u0026#34;, \u0026#34;天\u0026#34;, \u0026#34;气\u0026#34;, \u0026#34;好\u0026#34;] 训练目标： 看到\u0026#34;今\u0026#34;，预测\u0026#34;天\u0026#34; 看到\u0026#34;今天\u0026#34;，预测\u0026#34;天\u0026#34; 看到\u0026#34;今天天\u0026#34;，预测\u0026#34;气\u0026#34; ... Causal Mask（下三角矩阵，1=可以看，0=不能看）： 今 天 天 气 好 今 → [ 1, 0, 0, 0, 0] ← \u0026#34;今\u0026#34;只能看自己 天 → [ 1, 1, 0, 0, 0] ← \u0026#34;天\u0026#34;能看\u0026#34;今\u0026#34;和自己 天 → [ 1, 1, 1, 0, 0] 气 → [ 1, 1, 1, 1, 0] 好 → [ 1, 1, 1, 1, 1] ← \u0026#34;好\u0026#34;能看所有历史 适用场景：只在训练 GPT/Decoder 类模型时需要。推理时自回归逐步生成，天然只能看到历史，不需要此 mask。\n三、训练 vs 推理时传什么 Mask？ 场景 传入的 Mask 训练（GPT/Decoder） Causal Mask（下三角）AND Padding Mask 训练（BERT/Encoder） 只有 Padding Mask（双向 Attention，可以看全部） 推理，单条，无 padding None（不需要任何 mask） 推理，batch，有 padding 只有 Padding Mask 推理，自回归逐步生成 None（每步只输入 1 个新 token，天然无未来信息） 四、完整代码示例 import torch import torch.nn as nn import torch.nn.functional as F import math def make_padding_mask(token_ids: torch.Tensor, pad_token_id: int = 0) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; 生成 Padding Mask。 输入：token_ids [B, T] 输出：mask [B, 1, 1, T]，可广播到 [B, num_heads, T_q, T_k] 值：PAD 位置为 False（屏蔽），非 PAD 位置为 True（保留） \u0026#34;\u0026#34;\u0026#34; return (token_ids != pad_token_id).unsqueeze(1).unsqueeze(2) # [B, 1, 1, T] def make_causal_mask(seq_len: int, device: torch.device) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; 生成 Causal Mask（下三角矩阵）。 输出：mask [1, 1, T, T]，可广播到 [B, num_heads, T, T] 值：下三角为 True（可以看），上三角为 False（屏蔽未来） \u0026#34;\u0026#34;\u0026#34; return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device) ).unsqueeze(0).unsqueeze(0) # [1, 1, T, T] def make_decoder_mask(token_ids: torch.Tensor, pad_token_id: int = 0) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; GPT 训练时的完整 Mask = Causal Mask AND Padding Mask。 两者取交集：既不能看未来，也不能看 PAD。 \u0026#34;\u0026#34;\u0026#34; B, T = token_ids.shape causal = make_causal_mask(T, token_ids.device) # [1, 1, T, T] padding = make_padding_mask(token_ids, pad_token_id) # [B, 1, 1, T] return causal \u0026amp; padding # [B, 1, T, T] class MultiHeadAttentionWithMask(nn.Module): \u0026#34;\u0026#34;\u0026#34;带 Mask 支持的 MHA\u0026#34;\u0026#34;\u0026#34; def __init__(self, d_model: int, num_heads: int): super().__init__() assert d_model % num_heads == 0 self.num_heads = num_heads self.d_k = d_model // num_heads 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) self.W_O = nn.Linear(d_model, d_model) def split_heads(self, x: torch.Tensor) -\u0026gt; torch.Tensor: B, T, _ = x.shape return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2) def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -\u0026gt; torch.Tensor: B, T, _ = x.shape Q = self.split_heads(self.W_Q(x)) # [B, h, T, d_k] K = self.split_heads(self.W_K(x)) V = self.split_heads(self.W_V(x)) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # [B, h, T, T] if mask is not None: # mask=False 的位置设为 -inf，softmax 后权重变为 0 scores = scores.masked_fill(mask == False, float(\u0026#39;-inf\u0026#39;)) weights = F.softmax(scores, dim=-1) attended = torch.matmul(weights, V) attended = attended.transpose(1, 2).contiguous().view(B, T, self.d_model) return self.W_O(attended) # ── 使用示例 ────────────────────────────────────────────────────────────── if __name__ == \u0026#34;__main__\u0026#34;: d_model = 32 num_heads = 4 pad_token_id = 0 model = MultiHeadAttentionWithMask(d_model=d_model, num_heads=num_heads) model.eval() # ── 场景1：训练时（GPT/Decoder），有 PAD，需要 Causal + Padding Mask ── token_ids_train = torch.tensor([ [3, 7, 2, 0, 0], # 句子1，后两个是 PAD [5, 1, 8, 4, 6], # 句子2，无 PAD ]) x_train = torch.randn(2, 5, d_model) decoder_mask = make_decoder_mask(token_ids_train, pad_token_id) print(\u0026#34;=== 训练场景（GPT Decoder）===\u0026#34;) print(f\u0026#34;token_ids shape: {token_ids_train.shape}\u0026#34;) print(f\u0026#34;decoder_mask shape: {decoder_mask.shape}\u0026#34;) # [2, 1, 5, 5] print(\u0026#34;句子1的 mask（下三角 + 屏蔽PAD列）：\u0026#34;) print(decoder_mask[0, 0].int()) # [[1, 0, 0, 0, 0], # [1, 1, 0, 0, 0], # [1, 1, 1, 0, 0], ← PAD列（第4、5列）被屏蔽 # [1, 1, 1, 0, 0], ← PAD行也只能看到有效列 # [1, 1, 1, 0, 0]] with torch.no_grad(): output_train = model(x_train, mask=decoder_mask) print(f\u0026#34;输出 shape: {output_train.shape}\\n\u0026#34;) # [2, 5, 32] # ── 场景2：训练时（BERT/Encoder），有 PAD，只需 Padding Mask ── token_ids_bert = torch.tensor([ [3, 7, 2, 0, 0], [5, 1, 8, 4, 6], ]) padding_mask = make_padding_mask(token_ids_bert, pad_token_id) # [2, 1, 1, 5] print(\u0026#34;=== 训练场景（BERT Encoder）===\u0026#34;) print(f\u0026#34;padding_mask shape: {padding_mask.shape}\u0026#34;) print(\u0026#34;句子1的 padding_mask：\u0026#34;, padding_mask[0, 0, 0].int()) # [1, 1, 1, 0, 0] with torch.no_grad(): output_bert = model(x_train, mask=padding_mask) print(f\u0026#34;输出 shape: {output_bert.shape}\\n\u0026#34;) # ── 场景3：推理时，单条输入，无 PAD，mask=None ── single_input = torch.randn(1, 8, d_model) # batch=1，序列长度=8 print(\u0026#34;=== 推理场景（单条，无PAD）===\u0026#34;) with torch.no_grad(): output_infer = model(single_input, mask=None) # 直接传 None print(f\u0026#34;输出 shape: {output_infer.shape}\u0026#34;) # [1, 8, 32] 五、Mask 在 Attention 计算中的位置 Q, K, V 计算完毕 ↓ scores = Q @ K^T / √d_k # [B, h, T, T] ↓ if mask is not None: scores[mask == False] = -∞ # 屏蔽不该看的位置 ↓ weights = softmax(scores) # -∞ 位置的权重变为 0 ↓ output = weights @ V # 被屏蔽的位置对输出无贡献 六、常见问题 Q：Causal Mask 为什么是下三角？ 位置 i 的 token 只能看到位置 ≤ i 的 token（包括自己），矩阵的第 i 行第 j 列表示\u0026quot;位置 i 能否看到位置 j\u0026quot;，因此 j ≤ i 的位置为 1，即下三角。\nQ：推理时自回归生成为什么不需要 Causal Mask？ 推理时每步只输入 1 个新 token（配合 KV Cache），Q 只有 1 行，K/V 是历史所有 token，天然就只能看到历史，不存在\u0026quot;看到未来\u0026quot;的问题。\nQ：mask 的值为什么用 True/False 而不是 1/0？ 两种写法都可以，只要在 masked_fill 时保持一致：\nscores.masked_fill(mask == False, float(\u0026#39;-inf\u0026#39;)) # True=保留，False=屏蔽 scores.masked_fill(mask == 0, float(\u0026#39;-inf\u0026#39;)) # 1=保留，0=屏蔽（等价） Q：BERT 为什么不需要 Causal Mask？ BERT 是 Encoder-only 模型，做的是理解任务（分类、NER 等），一次性看完整个句子，双向 Attention 是它的优势，不需要屏蔽未来。\n七、一句话总结 Padding Mask：屏蔽 [PAD] 填充位置，训练和推理（有 padding 时）都需要 Causal Mask：屏蔽未来位置（下三角矩阵），只在训练 GPT/Decoder 类模型时需要 推理时单条自回归生成：传 None，不需要任何 mask ","permalink":"https://afan.ml/posts/attention--mask/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：Mask 是在计算 softmax 之前，把某些位置的注意力分数强制设为 $-\\infty$，使 softmax 后这些位置的权重变为 0，相当于\u0026quot;屏蔽掉\u0026quot;不该看的位置。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要-mask\"\u003e一、为什么需要 Mask？\u003c/h2\u003e\n\u003cp\u003eAttention 的计算是：\u003c/p\u003e\n\u003cp\u003e$$\n\\text{scores} = \\frac{QK^T}{\\sqrt{d_k}}, \\quad \\text{weights} = \\text{softmax}(\\text{scores}), \\quad \\text{output} = \\text{weights} \\cdot V\n$$\u003c/p\u003e\n\u003cp\u003e默认情况下，每个 token 的 Q 会和序列中\u003cstrong\u003e所有\u003c/strong\u003e token 的 K 做点积，包括：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e无意义的 \u003ccode\u003e[PAD]\u003c/code\u003e 填充 token\u003c/li\u003e\n\u003cli\u003e未来还没生成的 token（训练时）\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e这两种情况都需要用 Mask 屏蔽掉。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二两种-mask\"\u003e二、两种 Mask\u003c/h2\u003e\n\u003ch3 id=\"21-padding-mask填充遮罩\"\u003e2.1 Padding Mask（填充遮罩）\u003c/h3\u003e\n\u003cp\u003e\u003cstrong\u003e问题\u003c/strong\u003e：一个 batch 里不同句子长度不同，需要用 \u003ccode\u003e[PAD]\u003c/code\u003e token 补齐到相同长度，但 \u003ccode\u003e[PAD]\u003c/code\u003e 是无意义的，不应该被 Attention 到。\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003ebatch 里两个句子（补齐到长度5）：\n 句子1: [\u0026#34;今\u0026#34;, \u0026#34;天\u0026#34;, \u0026#34;好\u0026#34;, \u0026#34;[PAD]\u0026#34;, \u0026#34;[PAD]\u0026#34;] ← 实际长度3，补了2个PAD\n 句子2: [\u0026#34;天\u0026#34;, \u0026#34;气\u0026#34;, \u0026#34;真\u0026#34;, \u0026#34;不\u0026#34;, \u0026#34;错\u0026#34; ] ← 实际长度5，无需补\n\nPadding Mask（1=有效位置，0=PAD位置）：\n 句子1: [1, 1, 1, 0, 0]\n 句子2: [1, 1, 1, 1, 1]\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e屏蔽效果：\u003c/p\u003e","title":"Attention Mask 笔记"},{"content":" 核心问题：为什么需要归一化？神经网络训练时，每层的输入分布会随参数更新不断变化（Internal Covariate Shift），导致训练不稳定、收敛慢。归一化就是把数据\u0026quot;拉回\u0026quot;到均值为0、方差为1的分布，再通过可学习参数还原表达能力。\n一、公式（两者一样，区别在于\u0026quot;沿哪个维度算\u0026quot;） $$ \\hat{x} = \\frac{x - \\mu}{\\sqrt{\\sigma^2 + \\epsilon}}, \\quad y = \\gamma \\hat{x} + \\beta $$\n$\\mu$、$\\sigma^2$：均值和方差（统计维度不同是BN和LN的核心区别） $\\gamma$、$\\beta$：可学习的缩放和偏移参数 $\\epsilon$：防止除零的小常数 二、Batch Norm（BN） 核心思想 跨样本、同一特征维度做归一化。\n统计维度 假设输入 shape 为 [B, C, H, W]（B=batch，C=通道，H/W=空间）：\n对每个通道 C，在 B、H、W 三个维度上计算均值和方差 每个通道有一组独立的 $\\gamma$、$\\beta$ 直觉理解 \u0026ldquo;把这批图片里，同一个通道的所有像素值，统一归一化\u0026rdquo;\n具体例子 输入：4张图，每张1个通道，2x2像素 数据（展平后每行是一张图的像素）： 图1: [1, 2, 3, 4] 图2: [5, 6, 7, 8] 图3: [2, 3, 4, 5] 图4: [6, 7, 8, 9] BN 把这 4×4=16 个数一起算均值和方差： 所有数：1,2,3,4,5,6,7,8,2,3,4,5,6,7,8,9 μ = (1+2+3+4+5+6+7+8+2+3+4+5+6+7+8+9) / 16 = 80/16 = 5.0 σ² = [(1-5)²+(2-5)²+...+(9-5)²] / 16 = 88/16 = 5.5 σ = √5.5 ≈ 2.345 归一化示例（以图1的像素值1为例）： x̂ = (1 - 5.0) / √(5.5 + ε) ≈ -4.0 / 2.345 ≈ -1.706 所有16个像素值都用同一个 μ=5.0、σ≈2.345 归一化 优点 效果好，训练稳定，允许更大学习率 CV 任务（CNN）的标配 缺点 依赖 batch size：batch 太小时均值/方差估计不准，效果差 推理时需要维护全局统计量（running mean/var） 不适合 RNN/Transformer：序列长度不固定，batch 内样本差异大 三、Layer Norm（LN） 核心思想 同一样本、跨特征维度做归一化。\n统计维度 假设输入 shape 为 [B, T, D]（B=batch，T=序列长度，D=特征维度）：\n对每个 token（每个位置），在 D 维度上计算均值和方差 每个样本独立归一化，不依赖其他样本 直觉理解 \u0026ldquo;把这个词的所有特征值，归一化到同一尺度\u0026rdquo;\n具体例子 输入：1个句子，3个词，每词4维特征 数据（每行是一个词的特征向量）： 词1: [1, 2, 3, 4] μ = (1+2+3+4)/4 = 2.5 σ² = [(1-2.5)²+(2-2.5)²+(3-2.5)²+(4-2.5)²] / 4 = [2.25+0.25+0.25+2.25]/4 = 1.25 σ = √1.25 ≈ 1.118 归一化后: [(-1.5/1.118), (-0.5/1.118), (0.5/1.118), (1.5/1.118)] ≈ [-1.342, -0.447, 0.447, 1.342] 词2: [5, 6, 7, 8] μ = (5+6+7+8)/4 = 6.5 σ² = [(5-6.5)²+(6-6.5)²+(7-6.5)²+(8-6.5)²] / 4 = [2.25+0.25+0.25+2.25]/4 = 1.25 归一化后: ≈ [-1.342, -0.447, 0.447, 1.342]（和词1形状相同，因为数值等间距） 词3: [2, 4, 6, 8] μ = (2+4+6+8)/4 = 5.0 σ² = [(2-5)²+(4-5)²+(6-5)²+(8-5)²] / 4 = [9+1+1+9]/4 = 5.0 σ = √5.0 ≈ 2.236 归一化后: [(-3/2.236), (-1/2.236), (1/2.236), (3/2.236)] ≈ [-1.342, -0.447, 0.447, 1.342] 每个词独立计算自己的 μ 和 σ，互不影响 优点 不依赖 batch size，batch=1 也能正常工作 推理和训练行为一致，无需维护全局统计量 NLP/Transformer 的标配 缺点 在 CV 任务上效果通常不如 BN 四、对比总结 对比项 Batch Norm Layer Norm 归一化维度 跨样本（Batch维） 跨特征（Feature维） 统计量依赖 依赖整个 batch 仅依赖单个样本 batch size 敏感 ✅ 敏感（小batch效果差） ❌ 不敏感 推理时 需要 running mean/var 直接计算，无需额外状态 适用场景 CNN / 图像任务 Transformer / NLP / RNN 代表模型 ResNet、VGG BERT、GPT、LLaMA 五、一句话记忆法 BN：同一特征，跨样本归一 → \u0026ldquo;列归一化\u0026rdquo;（竖着看） LN：同一样本，跨特征归一 → \u0026ldquo;行归一化\u0026rdquo;（横着看） 数据矩阵（行=样本，列=特征）： 特征1 特征2 特征3 样本1 [ 1, 2, 3 ] ← LN 沿这一行归一化 样本2 [ 4, 5, 6 ] ← LN 沿这一行归一化 样本3 [ 7, 8, 9 ] ← LN 沿这一行归一化 ↑ ↑ ↑ BN 沿每列归一化 六、代码速览（PyTorch） import torch import torch.nn as nn x = torch.randn(8, 64, 32, 32) # [B, C, H, W] # Batch Norm：对每个通道归一化 bn = nn.BatchNorm2d(64) out_bn = bn(x) # Layer Norm：对最后几个维度归一化（通常是特征维） x_seq = torch.randn(8, 128, 512) # [B, T, D] ln = nn.LayerNorm(512) # 对 D=512 这一维归一化 out_ln = ln(x_seq) ","permalink":"https://afan.ml/posts/b-n-vs--l-n/","summary":"\u003cblockquote\u003e\n\u003cp\u003e核心问题：为什么需要归一化？神经网络训练时，每层的输入分布会随参数更新不断变化（Internal Covariate Shift），导致训练不稳定、收敛慢。归一化就是把数据\u0026quot;拉回\u0026quot;到均值为0、方差为1的分布，再通过可学习参数还原表达能力。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一公式两者一样区别在于沿哪个维度算\"\u003e一、公式（两者一样，区别在于\u0026quot;沿哪个维度算\u0026quot;）\u003c/h2\u003e\n\u003cp\u003e$$\n\\hat{x} = \\frac{x - \\mu}{\\sqrt{\\sigma^2 + \\epsilon}}, \\quad y = \\gamma \\hat{x} + \\beta\n$$\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e$\\mu$、$\\sigma^2$：均值和方差（\u003cstrong\u003e统计维度不同是BN和LN的核心区别\u003c/strong\u003e）\u003c/li\u003e\n\u003cli\u003e$\\gamma$、$\\beta$：可学习的缩放和偏移参数\u003c/li\u003e\n\u003cli\u003e$\\epsilon$：防止除零的小常数\u003c/li\u003e\n\u003c/ul\u003e\n\u003chr\u003e\n\u003ch2 id=\"二batch-normbn\"\u003e二、Batch Norm（BN）\u003c/h2\u003e\n\u003ch3 id=\"核心思想\"\u003e核心思想\u003c/h3\u003e\n\u003cp\u003e\u003cstrong\u003e跨样本、同一特征维度\u003c/strong\u003e做归一化。\u003c/p\u003e\n\u003ch3 id=\"统计维度\"\u003e统计维度\u003c/h3\u003e\n\u003cp\u003e假设输入 shape 为 \u003ccode\u003e[B, C, H, W]\u003c/code\u003e（B=batch，C=通道，H/W=空间）：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e对每个通道 C，在 \u003cstrong\u003eB、H、W\u003c/strong\u003e 三个维度上计算均值和方差\u003c/li\u003e\n\u003cli\u003e每个通道有一组独立的 $\\gamma$、$\\beta$\u003c/li\u003e\n\u003c/ul\u003e\n\u003ch3 id=\"直觉理解\"\u003e直觉理解\u003c/h3\u003e\n\u003cblockquote\u003e\n\u003cp\u003e\u0026ldquo;把这批图片里，同一个通道的所有像素值，统一归一化\u0026rdquo;\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003ch3 id=\"具体例子\"\u003e具体例子\u003c/h3\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e输入：4张图，每张1个通道，2x2像素\n数据（展平后每行是一张图的像素）：\n 图1: [1, 2, 3, 4]\n 图2: [5, 6, 7, 8]\n 图3: [2, 3, 4, 5]\n 图4: [6, 7, 8, 9]\n\nBN 把这 4×4=16 个数一起算均值和方差：\n 所有数：1,2,3,4,5,6,7,8,2,3,4,5,6,7,8,9\n μ = (1+2+3+4+5+6+7+8+2+3+4+5+6+7+8+9) / 16 = 80/16 = 5.0\n σ² = [(1-5)²+(2-5)²+...+(9-5)²] / 16 = 88/16 = 5.5\n σ = √5.5 ≈ 2.345\n\n 归一化示例（以图1的像素值1为例）：\n x̂ = (1 - 5.0) / √(5.5 + ε) ≈ -4.0 / 2.345 ≈ -1.706\n\n 所有16个像素值都用同一个 μ=5.0、σ≈2.345 归一化\n\u003c/code\u003e\u003c/pre\u003e\u003ch3 id=\"优点\"\u003e优点\u003c/h3\u003e\n\u003cul\u003e\n\u003cli\u003e效果好，训练稳定，允许更大学习率\u003c/li\u003e\n\u003cli\u003eCV 任务（CNN）的标配\u003c/li\u003e\n\u003c/ul\u003e\n\u003ch3 id=\"缺点\"\u003e缺点\u003c/h3\u003e\n\u003cul\u003e\n\u003cli\u003e\u003cstrong\u003e依赖 batch size\u003c/strong\u003e：batch 太小时均值/方差估计不准，效果差\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e推理时需要维护全局统计量\u003c/strong\u003e（running mean/var）\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e不适合 RNN/Transformer\u003c/strong\u003e：序列长度不固定，batch 内样本差异大\u003c/li\u003e\n\u003c/ul\u003e\n\u003chr\u003e\n\u003ch2 id=\"三layer-normln\"\u003e三、Layer Norm（LN）\u003c/h2\u003e\n\u003ch3 id=\"核心思想-1\"\u003e核心思想\u003c/h3\u003e\n\u003cp\u003e\u003cstrong\u003e同一样本、跨特征维度\u003c/strong\u003e做归一化。\u003c/p\u003e","title":"Batch Norm vs Layer Norm 学习笔记"},{"content":" 一句话：BPE（Byte Pair Encoding）是一种把文本切分成\u0026quot;子词单元\u0026quot;的算法，是现代大模型 Tokenizer 的核心。tiktoken 是 OpenAI 实现的高性能 BPE 库，GPT 系列模型使用它。\n一、为什么需要 Tokenizer？ 模型不能直接处理文字，需要先把文字转成数字（token ID），再转成向量（embedding）。\n原始文本：\u0026#34;今天天气很好\u0026#34; ↓ Tokenizer token 序列：[\u0026#34;今天\u0026#34;, \u0026#34;天气\u0026#34;, \u0026#34;很\u0026#34;, \u0026#34;好\u0026#34;] ↓ 查词表 token ID：[1234, 5678, 910, 1112] ↓ Embedding 层 向量矩阵：[[...], [...], [...], [...]] ← 模型真正处理的输入 二、三种切分粒度的对比 在 BPE 出现之前，有两种极端方案：\n方案1：词级别（Word-level） \u0026#34;unhappiness\u0026#34; → [\u0026#34;unhappiness\u0026#34;] ← 一个词一个 token 问题：\n词表会非常大（英语有几十万个词） 遇到训练时没见过的新词（OOV，Out-of-Vocabulary）就不认识 不同语言、不同形态的词需要分别存储（\u0026ldquo;run\u0026rdquo;、\u0026ldquo;running\u0026rdquo;、\u0026ldquo;ran\u0026rdquo; 是三个不同 token） 方案2：字符级别（Character-level） \u0026#34;unhappiness\u0026#34; → [\u0026#34;u\u0026#34;,\u0026#34;n\u0026#34;,\u0026#34;h\u0026#34;,\u0026#34;a\u0026#34;,\u0026#34;p\u0026#34;,\u0026#34;p\u0026#34;,\u0026#34;i\u0026#34;,\u0026#34;n\u0026#34;,\u0026#34;e\u0026#34;,\u0026#34;s\u0026#34;,\u0026#34;s\u0026#34;] ← 每个字符一个 token 问题：\n序列太长，Attention 的 $O(T^2)$ 复杂度爆炸 字符本身语义信息太少，模型难以学习 方案3：子词级别（Subword-level）—— BPE 的目标 \u0026#34;unhappiness\u0026#34; → [\u0026#34;un\u0026#34;, \u0026#34;happi\u0026#34;, \u0026#34;ness\u0026#34;] ← 有意义的子词单元 优点：\n词表大小可控（通常 3 万～10 万） 没见过的词可以拆成已知子词组合，不会 OOV 保留了一定的语义信息 三、BPE 算法详解 BPE 最初是数据压缩算法，被 Sennrich 等人（2016）引入 NLP。\n核心思想 反复合并出现频率最高的相邻字节/字符对，直到达到目标词表大小。\n训练过程（逐步演示） 初始语料（简化示例，括号内是出现频率）：\n\u0026#34;low\u0026#34;(5次) \u0026#34;lower\u0026#34;(2次) \u0026#34;newest\u0026#34;(6次) \u0026#34;widest\u0026#34;(3次) 第0步：初始化，把每个词拆成字符，加上词尾标记 \u0026lt;/w\u0026gt;（表示这是一个词的结尾）：\nl o w \u0026lt;/w\u0026gt; (5次) l o w e r \u0026lt;/w\u0026gt; (2次) n e w e s t \u0026lt;/w\u0026gt; (6次) w i d e s t \u0026lt;/w\u0026gt; (3次) 统计所有相邻字符对的频率：\n(l, o): 5+2 = 7 (o, w): 5+2 = 7 (w, \u0026lt;/w\u0026gt;): 5 (w, e): 2+6 = 8 ← 最高！ (e, r): 2 (e, s): 6+3 = 9 ← 最高！ ... 第1步：合并频率最高的对 (e, s) → es：\nl o w \u0026lt;/w\u0026gt; (5次) l o w e r \u0026lt;/w\u0026gt; (2次) n e w es t \u0026lt;/w\u0026gt; (6次) ← \u0026#34;e s\u0026#34; 合并为 \u0026#34;es\u0026#34; w i d es t \u0026lt;/w\u0026gt; (3次) ← \u0026#34;e s\u0026#34; 合并为 \u0026#34;es\u0026#34; 第2步：重新统计，合并频率最高的对 (es, t) → est：\nl o w \u0026lt;/w\u0026gt; (5次) l o w e r \u0026lt;/w\u0026gt; (2次) n e w est \u0026lt;/w\u0026gt; (6次) ← \u0026#34;es t\u0026#34; 合并为 \u0026#34;est\u0026#34; w i d est \u0026lt;/w\u0026gt; (3次) ← \u0026#34;es t\u0026#34; 合并为 \u0026#34;est\u0026#34; 第3步：合并 (w, e) → we（频率 8）：\nl o w \u0026lt;/w\u0026gt; (5次) l o we r \u0026lt;/w\u0026gt; (2次) n e we st \u0026lt;/w\u0026gt; (6次) w i d est \u0026lt;/w\u0026gt; (3次) 第4步：合并 (l, o) → lo（频率 7）：\nlo w \u0026lt;/w\u0026gt; (5次) lo we r \u0026lt;/w\u0026gt; (2次) n e we st \u0026lt;/w\u0026gt; (6次) w i d est \u0026lt;/w\u0026gt; (3次) ……持续合并，直到词表达到目标大小（如 50000 个 token）\n最终词表包含所有合并规则，例如：\n合并规则列表（按顺序）： 1. e + s → es 2. es + t → est 3. w + e → we 4. l + o → lo 5. lo + w → low ... 推理时如何切分新词 拿到新词，按照训练时学到的**合并规则列表（按顺序）**逐一应用：\n新词：\u0026#34;lowest\u0026#34; 初始：l o w e s t \u0026lt;/w\u0026gt; 应用规则1 (e+s→es)：l o w es t \u0026lt;/w\u0026gt; 应用规则2 (es+t→est)：l o w est \u0026lt;/w\u0026gt; 应用规则3 (w+e→we)：无匹配（w 后面是 est，不是 e） 应用规则4 (l+o→lo)：lo w est \u0026lt;/w\u0026gt; 应用规则5 (lo+w→low)：low est \u0026lt;/w\u0026gt; 最终切分：[\u0026#34;low\u0026#34;, \u0026#34;est\u0026lt;/w\u0026gt;\u0026#34;] → 对应 token ID 四、Byte-level BPE（字节级 BPE） GPT-2 及之后的模型使用的是 Byte-level BPE，在字节（byte）而非字符上做 BPE。\n动机：\n普通 BPE 以 Unicode 字符为基本单元，词表初始化需要覆盖所有字符（中文就有几万个） Byte-level BPE 以 256 个字节为基本单元，初始词表只有 256 个元素，可以表示任意语言、任意字符，永远不会 OOV \u0026#34;你好\u0026#34; 的 UTF-8 字节： \u0026#34;你\u0026#34; → [0xE4, 0xBD, 0xA0] \u0026#34;好\u0026#34; → [0xE5, 0xA5, 0xBD] Byte-level BPE 从这 6 个字节开始合并，最终可能得到： [\u0026#34;你好\u0026#34;]（如果这个组合在训练数据中足够频繁） 或 [\u0026#34;你\u0026#34;, \u0026#34;好\u0026#34;] 或更细的字节组合 五、tiktoken 是什么？ tiktoken 是 OpenAI 开源的高性能 BPE tokenizer 库，用 Rust 实现核心逻辑，Python 封装调用，速度比纯 Python 实现快 3~6 倍。\nGPT 系列使用的编码方案 模型 编码方案 词表大小 GPT-2 gpt2 50,257 GPT-3.5 / GPT-4 cl100k_base 100,277 GPT-4o o200k_base 200,019 text-embedding-ada-002 cl100k_base 100,277 词表越大，同样的文本切分出的 token 数越少，处理效率越高，但模型的 embedding 层参数也越多。\ntiktoken 安装和基本使用 # 安装 # pip install tiktoken import tiktoken # 加载编码方案 enc = tiktoken.get_encoding(\u0026#34;cl100k_base\u0026#34;) # GPT-4 使用的编码 # 或者直接按模型名加载 enc = tiktoken.encoding_for_model(\u0026#34;gpt-4\u0026#34;) \u0026lt;![CDATA[# 实际输出： # token IDs：[9906, 11, 220, 57668, 53901, 6447, 33, 1777, 374, 8056, 13] # token 数量：11]]\u0026gt; # ── 解码：token ID → 文本 ────────────────────────────── decoded = enc.decode(token_ids) print(f\u0026#34;解码还原：{decoded}\u0026#34;) # \u0026#34;Hello, 你好！BPE is amazing.\u0026#34; # ── 查看每个 token 对应的文本 ────────────────────────── for token_id in token_ids: token_bytes = enc.decode_single_token_bytes(token_id) print(f\u0026#34; ID {token_id:6d} → {token_bytes}\u0026#34;) # 实际输出： # ID 9906 → b\u0026#39;Hello\u0026#39; # ID 11 → b\u0026#39;,\u0026#39; # ID 220 → b\u0026#39; \u0026#39; # ID 57668 → b\u0026#39;\\xe4\\xbd\\xa0\u0026#39; ← \u0026#34;你\u0026#34; 的 UTF-8 字节 # ID 53901 → b\u0026#39;\\xe5\\xa5\\xbd\u0026#39; ← \u0026#34;好\u0026#34; 的 UTF-8 字节 # ID 6447 → b\u0026#39;\\xef\\xbc\\x81\u0026#39; ← \u0026#34;！\u0026#34; 的 UTF-8 字节 # ID 33 → b\u0026#39;B\u0026#39; # ID 1777 → b\u0026#39;PE\u0026#39; # ID 374 → b\u0026#39; is\u0026#39; # ID 8056 → b\u0026#39; amazing\u0026#39; # ID 13 → b\u0026#39;.\u0026#39; # ── 统计 token 数（不解码，只计数，更快）────────────── num_tokens = len(enc.encode(text)) print(f\u0026#34;token 数：{num_tokens}\u0026#34;) # ── 特殊 token ───────────────────────────────────────── # GPT 系列有一些特殊 token，如 \u0026lt;|endoftext|\u0026gt; special_enc = tiktoken.get_encoding(\u0026#34;cl100k_base\u0026#34;) eot_id = special_enc.encode(\u0026#34;\u0026lt;|endoftext|\u0026gt;\u0026#34;, allowed_special={\u0026#34;\u0026lt;|endoftext|\u0026gt;\u0026#34;}) print(f\u0026#34;\u0026lt;|endoftext|\u0026gt; 的 token ID：{eot_id}\u0026#34;) # [100257] 实用场景：计算 API 费用前预估 token 数 OpenAI API 按 token 计费，发送前先估算：\nimport tiktoken def count_tokens_for_gpt4(messages: list[dict]) -\u0026gt; int: \u0026#34;\u0026#34;\u0026#34; 估算 GPT-4 API 调用的 token 数。 messages 格式：[{\u0026#34;role\u0026#34;: \u0026#34;user\u0026#34;, \u0026#34;content\u0026#34;: \u0026#34;...\u0026#34;}, ...] \u0026#34;\u0026#34;\u0026#34; enc = tiktoken.encoding_for_model(\u0026#34;gpt-4\u0026#34;) total_tokens = 0 for message in messages: # 每条消息有固定的格式开销（role + 分隔符） total_tokens += 4 for value in message.values(): total_tokens += len(enc.encode(str(value))) total_tokens += 2 # 回复的起始开销 return total_tokens messages = [ {\u0026#34;role\u0026#34;: \u0026#34;system\u0026#34;, \u0026#34;content\u0026#34;: \u0026#34;你是一个有帮助的助手。\u0026#34;}, {\u0026#34;role\u0026#34;: \u0026#34;user\u0026#34;, \u0026#34;content\u0026#34;: \u0026#34;请帮我写一首关于春天的诗。\u0026#34;}, ] estimated_tokens = count_tokens_for_gpt4(messages) print(f\u0026#34;预估输入 token 数：{estimated_tokens}\u0026#34;) # 按 GPT-4 价格 $0.03/1K tokens 估算费用 cost_usd = estimated_tokens / 1000 * 0.03 print(f\u0026#34;预估费用：${cost_usd:.6f}\u0026#34;) 六、其他主流模型的 Tokenizer 模型系列 Tokenizer 类型 词表大小 库 GPT-2/3/4 Byte-level BPE 50K / 100K / 200K tiktoken LLaMA-1/2 SentencePiece BPE 32,000 sentencepiece LLaMA-3 Byte-level BPE 128,256 tiktoken 兼容 Qwen2 Byte-level BPE 151,936 tiktoken 兼容 BERT WordPiece 30,522 HuggingFace tokenizers T5 SentencePiece Unigram 32,100 sentencepiece WordPiece（BERT 用）和 BPE 思路类似，区别在于合并标准：\nBPE：合并频率最高的字符对 WordPiece：合并能最大化训练数据似然的字符对（更精确但更慢） SentencePiece（LLaMA-1/2 用）：\n直接在原始文本上操作，不需要预分词（空格也是普通字符） 支持 BPE 和 Unigram 两种算法 跨语言效果好 七、一个有趣的现象：token 边界影响模型能力 BPE 切分方式会直接影响模型的推理能力：\nimport tiktoken enc = tiktoken.get_encoding(\u0026#34;cl100k_base\u0026#34;) # \u0026#34;9.11\u0026#34; 和 \u0026#34;9.9\u0026#34; 的 token 切分 print(enc.encode(\u0026#34;9.11\u0026#34;)) # 实际输出：[24, 13, 806] → [b\u0026#39;9\u0026#39;, b\u0026#39;.\u0026#39;, b\u0026#39;11\u0026#39;] print(enc.encode(\u0026#34;9.9\u0026#34;)) # 实际输出：[24, 13, 24] → [b\u0026#39;9\u0026#39;, b\u0026#39;.\u0026#39;, b\u0026#39;9\u0026#39;] # GPT 早期版本认为 9.11 \u0026gt; 9.9，因为 \u0026#34;11\u0026#34; \u0026gt; \u0026#34;9\u0026#34;（字符串比较） # 这是 token 粒度导致的数字理解问题 # 字符计数问题：strawberry 里有几个 r？ print(enc.encode(\u0026#34;strawberry\u0026#34;)) # 实际输出：[496, 675, 15717] → [b\u0026#39;str\u0026#39;, b\u0026#39;aw\u0026#39;, b\u0026#39;berry\u0026#39;] # 模型看到的是三个 token（str / aw / berry），不是字符，所以数字符时容易出错 这也是为什么现代模型（GPT-4o 用 o200k_base，词表 20 万）要扩大词表——更大的词表意味着更少的切分，更完整的语义单元，减少这类边界问题。\n八、一句话总结 BPE：反复合并频率最高的相邻字符对，训练出一套合并规则，推理时按规则切分新词。解决了词级别 OOV 和字符级别序列过长的问题。 Byte-level BPE：在字节上做 BPE，初始词表只有 256 个元素，永远不会 OOV，GPT-2 之后的主流方案。 tiktoken：OpenAI 的高性能 BPE 实现，Rust 核心，GPT 系列标配，常用于计算 token 数和 API 费用估算。 ","permalink":"https://afan.ml/posts/b-p-e%E4%B8%8E-tokenizer/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：BPE（Byte Pair Encoding）是一种把文本切分成\u0026quot;子词单元\u0026quot;的算法，是现代大模型 Tokenizer 的核心。tiktoken 是 OpenAI 实现的高性能 BPE 库，GPT 系列模型使用它。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要-tokenizer\"\u003e一、为什么需要 Tokenizer？\u003c/h2\u003e\n\u003cp\u003e模型不能直接处理文字，需要先把文字转成数字（token ID），再转成向量（embedding）。\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e原始文本：\u0026#34;今天天气很好\u0026#34;\n ↓ Tokenizer\ntoken 序列：[\u0026#34;今天\u0026#34;, \u0026#34;天气\u0026#34;, \u0026#34;很\u0026#34;, \u0026#34;好\u0026#34;]\n ↓ 查词表\ntoken ID：[1234, 5678, 910, 1112]\n ↓ Embedding 层\n向量矩阵：[[...], [...], [...], [...]] ← 模型真正处理的输入\n\u003c/code\u003e\u003c/pre\u003e\u003chr\u003e\n\u003ch2 id=\"二三种切分粒度的对比\"\u003e二、三种切分粒度的对比\u003c/h2\u003e\n\u003cp\u003e在 BPE 出现之前，有两种极端方案：\u003c/p\u003e\n\u003ch3 id=\"方案1词级别word-level\"\u003e方案1：词级别（Word-level）\u003c/h3\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e\u0026#34;unhappiness\u0026#34; → [\u0026#34;unhappiness\u0026#34;] ← 一个词一个 token\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e问题\u003c/strong\u003e：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e词表会非常大（英语有几十万个词）\u003c/li\u003e\n\u003cli\u003e遇到训练时没见过的新词（OOV，Out-of-Vocabulary）就不认识\u003c/li\u003e\n\u003cli\u003e不同语言、不同形态的词需要分别存储（\u0026ldquo;run\u0026rdquo;、\u0026ldquo;running\u0026rdquo;、\u0026ldquo;ran\u0026rdquo; 是三个不同 token）\u003c/li\u003e\n\u003c/ul\u003e\n\u003ch3 id=\"方案2字符级别character-level\"\u003e方案2：字符级别（Character-level）\u003c/h3\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e\u0026#34;unhappiness\u0026#34; → [\u0026#34;u\u0026#34;,\u0026#34;n\u0026#34;,\u0026#34;h\u0026#34;,\u0026#34;a\u0026#34;,\u0026#34;p\u0026#34;,\u0026#34;p\u0026#34;,\u0026#34;i\u0026#34;,\u0026#34;n\u0026#34;,\u0026#34;e\u0026#34;,\u0026#34;s\u0026#34;,\u0026#34;s\u0026#34;] ← 每个字符一个 token\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e问题\u003c/strong\u003e：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e序列太长，Attention 的 $O(T^2)$ 复杂度爆炸\u003c/li\u003e\n\u003cli\u003e字符本身语义信息太少，模型难以学习\u003c/li\u003e\n\u003c/ul\u003e\n\u003ch3 id=\"方案3子词级别subword-level-bpe-的目标\"\u003e方案3：子词级别（Subword-level）—— BPE 的目标\u003c/h3\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e\u0026#34;unhappiness\u0026#34; → [\u0026#34;un\u0026#34;, \u0026#34;happi\u0026#34;, \u0026#34;ness\u0026#34;] ← 有意义的子词单元\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e优点\u003c/strong\u003e：\u003c/p\u003e","title":"BPE 与 Tokenizer 笔记"},{"content":" 一句话：KV Cache 是推理时把已经算过的 Key 和 Value 缓存起来，避免每生成一个新 token 都重复计算历史 token 的 K/V，是大模型推理加速的核心技术。\n一、为什么需要 KV Cache？ 大模型生成 token 的方式：自回归（Auto-regressive） 大模型生成文本是一个 token 一个 token 地生成的，每次生成下一个 token 时，都要把所有历史 token 作为输入重新过一遍 Transformer。\n输入： \u0026#34;今天天气\u0026#34; 生成： \u0026#34;今天天气\u0026#34; → \u0026#34;很\u0026#34; 生成： \u0026#34;今天天气很\u0026#34; → \u0026#34;好\u0026#34; 生成： \u0026#34;今天天气很好\u0026#34; → \u0026#34;！\u0026#34; 每一步生成，Attention 都要计算：\n$$ \\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V $$\n其中 Q、K、V 都来自当前所有 token（包括历史的）。\n没有 KV Cache 时的重复计算 假设已经生成了 t 个 token，现在要生成第 t+1 个：\n第1步生成\u0026#34;很\u0026#34;： 对 token [\u0026#34;今\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;气\u0026#34;] 全部计算 K、V → K = [K_今, K_天, K_天, K_气] → V = [V_今, V_天, V_天, V_气] 第2步生成\u0026#34;好\u0026#34;： 对 token [\u0026#34;今\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;气\u0026#34;,\u0026#34;很\u0026#34;] 全部计算 K、V → K = [K_今, K_天, K_天, K_气, K_很] ← 前4个和上一步完全一样！ → V = [V_今, V_天, V_天, V_气, V_很] ← 前4个和上一步完全一样！ 第3步生成\u0026#34;！\u0026#34;： 对 token [\u0026#34;今\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;气\u0026#34;,\u0026#34;很\u0026#34;,\u0026#34;好\u0026#34;] 全部计算 K、V → 前5个 K/V 和上一步完全一样！ 历史 token 的 K/V 每次都重新算，纯属浪费。\n生成长度为 T 的序列，总计算量是 O(T²)，随序列变长急剧增加。\n二、KV Cache 的原理 核心思想 历史 token 的 K 和 V 不会随新 token 的生成而改变（在 Decoder-only 模型中，历史 token 不会 attend 到未来 token，所以它们的 K/V 是固定的）。\n因此，把每一层每个历史 token 的 K/V 存下来，下次直接读取，不再重新计算。\n有 KV Cache 后的计算流程 第1步生成\u0026#34;很\u0026#34;： 计算所有 token 的 K/V：[K_今, K_天, K_天, K_气] → 存入 Cache：cache_K = [K_今, K_天, K_天, K_气] cache_V = [V_今, V_天, V_天, V_气] 第2步生成\u0026#34;好\u0026#34;： 只计算新 token \u0026#34;很\u0026#34; 的 K/V：K_很, V_很 → 从 Cache 读出历史：[K_今, K_天, K_天, K_气] → 拼接：K = [K_今, K_天, K_天, K_气, K_很] ← 直接用，不重算 → 更新 Cache：cache_K = [K_今, K_天, K_天, K_气, K_很] 第3步生成\u0026#34;！\u0026#34;： 只计算新 token \u0026#34;好\u0026#34; 的 K/V：K_好, V_好 → 从 Cache 读出历史，拼接，继续... 每步只需计算 1 个新 token 的 K/V，历史的全部从缓存读取。\n计算量从 O(T²) 降为 O(T)，速度提升显著。\n三、KV Cache 的显存占用计算 KV Cache 的显存占用是推理时显存的主要来源之一，计算公式：\n$$ \\text{KV Cache 大小} = 2 \\times L \\times h \\times d_k \\times T \\times \\text{bytes_per_element} $$\n参数 含义 2 K 和 V 各一份 L Transformer 层数 h 注意力头数 d_k 每个头的维度 T 序列长度（已生成的 token 数） bytes_per_element 数据精度（fp16=2字节，fp32=4字节，int8=1字节） 具体例子：LLaMA-3-8B 的 KV Cache LLaMA-3-8B 的参数：\n层数 L = 32 头数 h = 32（Q头），KV头数 = 8（GQA，g=8） 每头维度 d_k = 128 精度 fp16（2字节） 序列长度 T = 4096 时的 KV Cache：\nKV Cache = 2 × 32层 × 8个KV头 × 128维 × 4096 token × 2字节 = 2 × 32 × 8 × 128 × 4096 × 2 = 536,870,912 字节 ≈ 512 MB 序列长度 T = 32768（32K 上下文）时：\nKV Cache = 512 MB × (32768 / 4096) = 512 MB × 8 = 4 GB 这就是为什么长上下文推理显存消耗巨大，也是 MQA/GQA/MLA 要压缩 KV Cache 的根本原因。\n四、KV Cache 在哪里用？ 只在 Decoder（生成阶段）使用 场景 是否用 KV Cache 原因 GPT 类模型推理（生成） ✅ 用 自回归生成，历史 K/V 固定 BERT 类模型推理（理解） ❌ 不用 一次性处理全部输入，不逐步生成 Encoder-Decoder（如 T5） Decoder 部分用 Encoder 输出固定，Decoder 自回归 训练阶段 ❌ 不用 训练时并行处理所有 token，不需要逐步生成 Prefill 阶段 vs Decode 阶段 大模型推理分两个阶段：\nPrefill（预填充）阶段： 输入：完整的 prompt（如 \u0026#34;请帮我写一首诗：\\n\u0026#34;） 操作：一次性并行计算所有 prompt token 的 K/V，存入 Cache 特点：计算密集，GPU 利用率高 Decode（解码）阶段： 输入：每次只有 1 个新生成的 token 操作：计算新 token 的 K/V，从 Cache 读取历史 K/V，拼接后做 Attention 特点：内存带宽密集（主要时间花在读 Cache），GPU 利用率低 五、具体计算过程举例 设定 模型：单层 Transformer，2个注意力头，每头维度 d_k = 2 已有 prompt：[\u0026ldquo;A\u0026rdquo;, \u0026ldquo;B\u0026rdquo;]，现在要生成第3个 token Prefill 阶段（处理 prompt \u0026ldquo;A B\u0026rdquo;） 固定投影矩阵（整个例子全程使用这三个矩阵）： W_K = [[1, 0], W_V = [[1, 1], W_Q = [[1, 0], [0, 1], [0, 1], [1, 1], [1, 0], [1, 0], [0, 1], [0, 1]] [0, 0]] [1, 0]] （均为 [4, 2] 的矩阵，将 4 维 embedding 投影到 2 维） token embedding（固定数值）： x_A = [1, 0, 1, 0] x_B = [0, 1, 0, 1] x_C = [1, 1, 0, 0] ← Decode 第1步的新 token x_D = [0, 0, 1, 1] ← Decode 第2步的新 token 计算 K/V（K = x @ W_K，V = x @ W_V）： K_A = x_A @ W_K = [1*1+0*0+1*1+0*0, 1*0+0*1+1*0+0*1] = [2, 0] K_B = x_B @ W_K = [0*1+1*0+0*1+1*0, 0*0+1*1+0*0+1*1] = [0, 2] V_A = x_A @ W_V = [1*1+0*0+1*1+0*0, 1*1+0*1+1*0+0*0] = [2, 1] V_B = x_B @ W_V = [0*1+1*0+0*1+1*0, 0*1+1*1+0*0+1*0] = [0, 1] 存入 KV Cache： cache_K = [[2, 0], ← K_A [0, 2]] ← K_B cache_V = [[2, 1], ← V_A [0, 1]] ← V_B Decode 阶段第1步（生成 token C） x_C = [1, 1, 0, 0]（已在 Prefill 阶段给出） 只计算 C 的 K/V（使用与 Prefill 相同的 W_K、W_V）： K_C = x_C @ W_K = [1*1+1*0+0*1+0*0, 1*0+1*1+0*0+0*1] = [1, 1] V_C = x_C @ W_V = [1*1+1*0+0*1+0*0, 1*1+1*1+0*0+0*0] = [1, 2] 从 Cache 读取历史，拼接： K_all = [[2, 0], ← K_A（从 Cache 读，不重新计算） [0, 2], ← K_B（从 Cache 读，不重新计算） [1, 1]] ← K_C（刚计算） V_all = [[2, 1], ← V_A（从 Cache 读，不重新计算） [0, 1], ← V_B（从 Cache 读，不重新计算） [1, 2]] ← V_C（刚计算） Q_C = x_C @ W_Q = [1*1+1*1+0*0+0*1, 1*0+1*1+0*1+0*0] = [2, 1] Attention 分数（Q_C 对所有 K）： score_A = Q_C · K_A / √2 = (2*2 + 1*0) / 1.414 = 4.0 / 1.414 = 2.828 score_B = Q_C · K_B / √2 = (2*0 + 1*2) / 1.414 = 2.0 / 1.414 = 1.414 score_C = Q_C · K_C / √2 = (2*1 + 1*1) / 1.414 = 3.0 / 1.414 = 2.121 softmax([2.828, 1.414, 2.121])： e^2.828=16.930, e^1.414=4.113, e^2.121=8.337, 总和=29.380 权重 = [16.930/29.380, 4.113/29.380, 8.337/29.380] = [0.576, 0.140, 0.284] 输出 = 0.576*V_A + 0.140*V_B + 0.284*V_C = 0.576*[2,1] + 0.140*[0,1] + 0.284*[1,2] = [1.152, 0.576] + [0.000, 0.140] + [0.284, 0.568] = [1.436, 1.284] 更新 Cache： cache_K 追加 K_C → [[2,0], [0,2], [1,1]] cache_V 追加 V_C → [[2,1], [0,1], [1,2]] Decode 阶段第2步（生成 token D） x_D = [0, 0, 1, 1]（已在 Prefill 阶段给出） 只计算 D 的 K/V（使用与 Prefill 相同的 W_K、W_V）： K_D = x_D @ W_K = [0*1+0*0+1*1+1*0, 0*0+0*1+1*0+1*1] = [1, 1] V_D = x_D @ W_V = [0*1+0*0+1*1+1*0, 0*1+0*1+1*0+1*0] = [1, 0] 从 Cache 读取历史（此时 Cache 已包含 A、B、C）： cache_K = [[2, 0], ← K_A（不重新计算） [0, 2], ← K_B（不重新计算） [1, 1]] ← K_C（不重新计算） 拼接 K_D： K_all = [[2, 0], ← K_A（Cache 读取） [0, 2], ← K_B（Cache 读取） [1, 1], ← K_C（Cache 读取） [1, 1]] ← K_D（刚计算） V_all = [[2, 1], ← V_A（Cache 读取） [0, 1], ← V_B（Cache 读取） [1, 2], ← V_C（Cache 读取） [1, 0]] ← V_D（刚计算） Q_D = x_D @ W_Q = [0*1+0*1+1*0+1*1, 0*0+0*1+1*1+1*0] = [1, 1] Attention 分数（Q_D 对所有 K）： score_A = Q_D · K_A / √2 = (1*2 + 1*0) / 1.414 = 2.0 / 1.414 = 1.414 score_B = Q_D · K_B / √2 = (1*0 + 1*2) / 1.414 = 2.0 / 1.414 = 1.414 score_C = Q_D · K_C / √2 = (1*1 + 1*1) / 1.414 = 2.0 / 1.414 = 1.414 score_D = Q_D · K_D / √2 = (1*1 + 1*1) / 1.414 = 2.0 / 1.414 = 1.414 softmax([1.414, 1.414, 1.414, 1.414])： 四个值完全相等 → 权重均等 = [0.25, 0.25, 0.25, 0.25] （Q_D 对所有历史 token 和自身的关注度相同，符合直觉：x_D 的特征与所有历史 token 距离相等） 输出 = 0.25*V_A + 0.25*V_B + 0.25*V_C + 0.25*V_D = 0.25*[2,1] + 0.25*[0,1] + 0.25*[1,2] + 0.25*[1,0] = [0.5, 0.25] + [0.0, 0.25] + [0.25, 0.5] + [0.25, 0.0] = [1.0, 1.0] 更新 Cache： cache_K 追加 K_D → [[2,0], [0,2], [1,1], [1,1]] cache_V 追加 V_D → [[2,1], [0,1], [1,2], [1,0]] 对比：若没有 KV Cache，这一步需要重新计算 K_A、K_B、K_C，共 3 次矩阵乘法； 有了 KV Cache，只计算了 K_D 这 1 次，节省了 75% 的 K/V 计算量。 序列越长，节省比例越高（生成第 T 步时节省 (T-1)/T）。 六、KV Cache 的工程实现要点 预分配显存（静态 Cache） 推理框架通常在开始前就按最大序列长度预分配好 Cache 显存，避免动态分配的开销：\nimport torch max_seq_len = 4096 num_layers = 32 num_kv_heads = 8 head_dim = 128 batch_size = 1 # 推理时通常 batch=1（单请求），服务化场景可设更大值 # 预分配 KV Cache，shape: [层数, 2, batch, kv头数, 最大序列长度, 头维度] # 第2维的 2 代表 K 和 V 各一份（index 0 = K，index 1 = V） kv_cache = torch.zeros(num_layers, 2, batch_size, num_kv_heads, max_seq_len, head_dim) print(f\u0026#34;KV Cache 预分配显存：{kv_cache.numel() * 2 / 1024**3:.2f} GB (fp16)\u0026#34;) # 32层 × 2 × 1 × 8头 × 4096 × 128 × 2字节 = 0.5 GB 动态 Cache（PagedAttention，vLLM 使用） 静态预分配的问题：不同请求序列长度不同，预分配最大长度会浪费显存。\nPagedAttention（vLLM 提出）借鉴操作系统虚拟内存的分页思想：\n把 KV Cache 切成固定大小的\u0026quot;页\u0026quot;（block） 按需分配，不同请求的 Cache 页可以不连续 支持多个请求共享相同 prompt 的 Cache（prefix sharing） 传统静态 Cache： 请求A（实际用200 token）：预分配4096 token的显存 → 浪费3896 token 请求B（实际用3000 token）：预分配4096 token的显存 → 浪费1096 token PagedAttention： 请求A：按需分配 200 token 的 Cache，用多少分多少 请求B：按需分配 3000 token 的 Cache → 显存利用率大幅提升，同等显存可服务更多并发请求 PyTorch 中使用 KV Cache 的简化示例 import torch import torch.nn as nn import math class AttentionWithKVCache(nn.Module): def __init__(self, d_model: int, num_heads: int): super().__init__() self.num_heads = num_heads self.d_k = d_model // num_heads 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 split_heads(self, x: torch.Tensor) -\u0026gt; torch.Tensor: B, T, _ = x.shape return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2) def forward( self, x: torch.Tensor, past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None ) -\u0026gt; tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: # 只对当前新 token 计算 Q、K、V Q = self.split_heads(self.W_Q(x)) # [B, h, T_new, d_k] K_new = self.split_heads(self.W_K(x)) V_new = self.split_heads(self.W_V(x)) # 拼接历史 Cache if past_key_value is not None: K_cached, V_cached = past_key_value K = torch.cat([K_cached, K_new], dim=2) # [B, h, T_all, d_k] V = torch.cat([V_cached, V_new], dim=2) else: K, V = K_new, V_new # 更新后的 Cache（供下一步使用） updated_cache = (K, V) # Scaled Dot-Product Attention scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) weights = torch.softmax(scores, dim=-1) attended = torch.matmul(weights, V) # 合并多头，输出投影 B, h, T_new, d_k = attended.shape attended = attended.transpose(1, 2).contiguous().view(B, T_new, h * d_k) output = self.W_O(attended) return output, updated_cache # 使用示例：模拟自回归生成 model = AttentionWithKVCache(d_model=64, num_heads=4) model.eval() with torch.no_grad(): # Prefill：处理 prompt（5个token） prompt = torch.randn(1, 5, 64) out, kv_cache = model(prompt, past_key_value=None) print(f\u0026#34;Prefill 后 Cache K shape: {kv_cache[0].shape}\u0026#34;) # [1, 4, 5, 16] # Decode：逐步生成新 token for step in range(3): new_token = torch.randn(1, 1, 64) # 每次只输入 1 个新 token out, kv_cache = model(new_token, past_key_value=kv_cache) print(f\u0026#34;Step {step+1} 后 Cache K shape: {kv_cache[0].shape}\u0026#34;) # Step 1: [1, 4, 6, 16] # Step 2: [1, 4, 7, 16] # Step 3: [1, 4, 8, 16] 七、KV Cache 与注意力变体的关系 KV Cache 的显存压力催生了各种注意力变体（详见 MHA 笔记第八节）：\nKV Cache 太大 ↓ MQA（2019）：所有头共享 K/V → Cache 缩小 h 倍，但性能下降 ↓ GQA（2023）：分组共享 K/V → Cache 缩小 h/g 倍，性能接近 MHA ↓ MLA（2024）：缓存低秩压缩向量 c → Cache 缩小 ~93%，性能不降反升 Flash Attention 解决的是另一个问题：不是 Cache 太大，而是 Attention 矩阵的 HBM 读写太慢。两者正交，可以同时使用。\n八、常见问题 Q：训练时为什么不用 KV Cache？ 训练时使用 Teacher Forcing，所有 token 并行输入，不需要逐步生成，因此不存在\u0026quot;历史 K/V 重复计算\u0026quot;的问题。\nQ：KV Cache 会不会导致内存溢出（OOM）？ 会。序列越长、batch 越大、模型越大，Cache 越大。解决方案：\n使用 GQA/MLA 减少 Cache 大小 使用 PagedAttention（vLLM）提高显存利用率 量化 Cache（如 KV Cache int8 量化） 设置最大序列长度限制 Q：多轮对话时 KV Cache 怎么处理？ 每轮对话结束后，可以选择：\n保留 Cache：下一轮直接在历史 Cache 上追加，速度快，但显存持续增长 清空 Cache：每轮重新 Prefill，显存可控，但每轮都要重新计算历史 Q：为什么 Prefill 快、Decode 慢？\nPrefill：大矩阵乘法，GPU 并行度高，计算密集型，GPU 利用率接近 100% Decode：每次只处理 1 个 token，矩阵很小，主要时间花在从显存读 KV Cache，内存带宽密集型，GPU 利用率通常只有 10%~30% 九、一句话总结 KV Cache = 空间换时间：把历史 token 的 K/V 存在显存里，每次生成新 token 只算新的 K/V，历史的直接读取，把推理复杂度从 O(T²) 降为 O(T)。\n代价是显存随序列长度线性增长，这是 MQA/GQA/MLA/PagedAttention 等技术存在的根本原因。\n","permalink":"https://afan.ml/posts/k-v-cache/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：KV Cache 是推理时把已经算过的 Key 和 Value 缓存起来，避免每生成一个新 token 都重复计算历史 token 的 K/V，是大模型推理加速的核心技术。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要-kv-cache\"\u003e一、为什么需要 KV Cache？\u003c/h2\u003e\n\u003ch3 id=\"大模型生成-token-的方式自回归auto-regressive\"\u003e大模型生成 token 的方式：自回归（Auto-regressive）\u003c/h3\u003e\n\u003cp\u003e大模型生成文本是\u003cstrong\u003e一个 token 一个 token 地生成\u003c/strong\u003e的，每次生成下一个 token 时，都要把\u003cstrong\u003e所有历史 token\u003c/strong\u003e 作为输入重新过一遍 Transformer。\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e输入： \u0026#34;今天天气\u0026#34;\n生成： \u0026#34;今天天气\u0026#34; → \u0026#34;很\u0026#34;\n生成： \u0026#34;今天天气很\u0026#34; → \u0026#34;好\u0026#34;\n生成： \u0026#34;今天天气很好\u0026#34; → \u0026#34;！\u0026#34;\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e每一步生成，Attention 都要计算：\u003c/p\u003e\n\u003cp\u003e$$\n\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V\n$$\u003c/p\u003e\n\u003cp\u003e其中 Q、K、V 都来自\u003cstrong\u003e当前所有 token\u003c/strong\u003e（包括历史的）。\u003c/p\u003e\n\u003ch3 id=\"没有-kv-cache-时的重复计算\"\u003e没有 KV Cache 时的重复计算\u003c/h3\u003e\n\u003cp\u003e假设已经生成了 t 个 token，现在要生成第 t+1 个：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e第1步生成\u0026#34;很\u0026#34;：\n 对 token [\u0026#34;今\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;气\u0026#34;] 全部计算 K、V\n → K = [K_今, K_天, K_天, K_气]\n → V = [V_今, V_天, V_天, V_气]\n\n第2步生成\u0026#34;好\u0026#34;：\n 对 token [\u0026#34;今\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;气\u0026#34;,\u0026#34;很\u0026#34;] 全部计算 K、V\n → K = [K_今, K_天, K_天, K_气, K_很] ← 前4个和上一步完全一样！\n → V = [V_今, V_天, V_天, V_气, V_很] ← 前4个和上一步完全一样！\n\n第3步生成\u0026#34;！\u0026#34;：\n 对 token [\u0026#34;今\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;天\u0026#34;,\u0026#34;气\u0026#34;,\u0026#34;很\u0026#34;,\u0026#34;好\u0026#34;] 全部计算 K、V\n → 前5个 K/V 和上一步完全一样！\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e历史 token 的 K/V 每次都重新算，纯属浪费。\u003c/strong\u003e\u003c/p\u003e","title":"KV Cache 笔记"},{"content":" 一句话：标准 Attention 的复杂度是 $O(N^2)$，Linear Attention 通过改变计算顺序，把复杂度降到 $O(N)$，代价是用核函数近似替代 softmax，牺牲了一定的表达能力，但换来了推理时的递推形式（类似 RNN），天然支持无限长序列。\n一、问题背景：标准 Attention 的瓶颈 标准 Attention 的计算：\n$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right) V$$\n序列长度为 $N$ 时：\n计算 $QK^T$ 需要 $O(N^2)$ 时间和空间 中间的 $N \\times N$ 注意力矩阵需要 $O(N^2)$ 显存 后果：序列长度翻倍，计算量翻 4 倍，显存翻 4 倍。长序列（如 100K token）几乎不可能用标准 Attention。\n二、核心思想：改变计算顺序 标准 Attention 的计算顺序 $$\\text{out} = \\underbrace{\\text{softmax}(QK^T)}_{\\text{先算这个，}N \\times N \\text{ 矩阵}} V$$\n必须先把完整的 $N \\times N$ 矩阵算出来，才能乘以 $V$。\nLinear Attention 的关键洞察 如果去掉 softmax，矩阵乘法满足结合律：\n$$\\text{out} = (QK^T) V = Q \\underbrace{(K^T V)}_{\\text{先算这个，}d \\times d \\text{ 矩阵}}$$\n先算 $K^T V$（形状 $d \\times d$，和序列长度无关），再乘以 $Q$：\n标准顺序：Q(N×d) · Kᵀ(d×N) = N×N 矩阵，再乘 V(N×d) → O(N²d) 线性顺序：Kᵀ(d×N) · V(N×d) = d×d 矩阵，Q(N×d) 再乘 → O(Nd²) 当 $d \\ll N$ 时（通常 $d=64$，$N=4096$），$O(Nd^2)$ 远小于 $O(N^2 d)$。\n为什么不能直接去掉 softmax？ softmax 保证了注意力权重非负且归一化（加起来等于 1），直接去掉会导致数值不稳定、训练发散。\n解决方案：用核函数（Kernel Function） $\\phi$ 替代 softmax：\n$$\\text{softmax}(q \\cdot k) \\approx \\phi(q) \\cdot \\phi(k)$$\n常见选择：\n$\\phi(x) = \\text{ELU}(x) + 1$（Linear Transformer，2020） $\\phi(x) = e^x$（近似，但数值不稳定） $\\phi(x) = \\text{ReLU}(x)$（简单但效果一般） 三、Linear Attention 的完整公式 $$\\text{out}i = \\frac{\\sum{j=1}^{N} \\phi(q_i)^T \\phi(k_j) v_j}{\\sum_{j=1}^{N} \\phi(q_i)^T \\phi(k_j)}$$\n利用结合律改写（分子）：\n$$\\text{out}i = \\frac{\\phi(q_i)^T \\left(\\sum{j=1}^{N} \\phi(k_j) v_j^T\\right)}{\\phi(q_i)^T \\left(\\sum_{j=1}^{N} \\phi(k_j)\\right)}$$\n令 $S = \\sum_{j=1}^{N} \\phi(k_j) v_j^T$（形状 $d \\times d$），$z = \\sum_{j=1}^{N} \\phi(k_j)$（形状 $d$）：\n$$\\text{out}_i = \\frac{\\phi(q_i)^T S}{\\phi(q_i)^T z}$$\n$S$ 和 $z$ 只需要计算一次，所有 token 共享，复杂度降为 $O(N)$。\n四、递推形式：Linear Attention ≈ RNN Causal（因果）Linear Attention 的最大优势：可以写成递推形式。\n对于自回归生成，第 $t$ 步的输出：\n$$S_t = S_{t-1} + \\phi(k_t) v_t^T$$ $$z_t = z_{t-1} + \\phi(k_t)$$ $$\\text{out}_t = \\frac{\\phi(q_t)^T S_t}{\\phi(q_t)^T z_t}$$\nimport torch import torch.nn.functional as F def phi(x): \u0026#34;\u0026#34;\u0026#34;核函数：ELU + 1，保证输出非负\u0026#34;\u0026#34;\u0026#34; return F.elu(x) + 1 def linear_attention_recurrent(Q, K, V): \u0026#34;\u0026#34;\u0026#34; Linear Attention 的递推形式（因果，自回归）。 Q, K, V: [seq_len, d] 返回: [seq_len, d] \u0026#34;\u0026#34;\u0026#34; seq_len, d = Q.shape outputs = [] # 递推状态：S 是 d×d 的累积矩阵，z 是 d 维的累积向量 S = torch.zeros(d, d) z = torch.zeros(d) for t in range(seq_len): q_t = phi(Q[t]) # [d] k_t = phi(K[t]) # [d] v_t = V[t] # [d] # 更新状态（累积历史 KV 信息） S = S + torch.outer(k_t, v_t) # [d, d] z = z + k_t # [d] # 计算当前输出 numerator = q_t @ S # [d] denominator = q_t @ z + 1e-6 # 标量，加 eps 防止除零 outputs.append(numerator / denominator) return torch.stack(outputs, dim=0) # [seq_len, d] # 验证 torch.manual_seed(42) seq_len, d = 8, 16 Q = torch.randn(seq_len, d) K = torch.randn(seq_len, d) V = torch.randn(seq_len, d) out = linear_attention_recurrent(Q, K, V) print(out.shape) # torch.Size([8, 16]) print(\u0026#34;no nan:\u0026#34;, not torch.isnan(out).any().item()) 这和 RNN 的结构完全一样：\nRNN： h_t = f(h_{t-1}, x_t)，用隐状态 h 压缩历史 Linear Attention：S_t = S_{t-1} + φ(k_t)vₜᵀ，用矩阵 S 压缩历史 五、并行训练形式 训练时不需要递推，可以并行计算（类似标准 Attention）：\ndef linear_attention_parallel(Q, K, V): \u0026#34;\u0026#34;\u0026#34; Linear Attention 的并行形式（非因果，训练用）。 Q, K, V: [seq_len, d] 返回: [seq_len, d] \u0026#34;\u0026#34;\u0026#34; Q_phi = phi(Q) # [seq_len, d] K_phi = phi(K) # [seq_len, d] # 先算 KᵀV（d×d），再乘 Q，复杂度 O(Nd²) KV = K_phi.T @ V # [d, d] numerator = Q_phi @ KV # [seq_len, d] # 归一化分母 z = K_phi.sum(dim=0) # [d] denominator = (Q_phi @ z).unsqueeze(-1) + 1e-6 # [seq_len, 1] return numerator / denominator # [seq_len, d] # 验证 out_parallel = linear_attention_parallel(Q, K, V) print(out_parallel.shape) # torch.Size([8, 16]) print(\u0026#34;no nan:\u0026#34;, not torch.isnan(out_parallel).any().item()) 六、Linear Attention 的问题 问题1：表达能力弱于 softmax Attention softmax 能产生非常\u0026quot;尖锐\u0026quot;的注意力分布（几乎只关注一个 token），而核函数近似做不到这一点，模型的选择性注意力能力下降。\n问题2：隐状态是固定大小的矩阵 S 的形状是 d×d，无论序列多长，S 的大小不变 → 长序列的历史信息被压缩进固定大小的矩阵 → 远距离依赖容易被\u0026#34;遗忘\u0026#34;（和 RNN 的梯度消失类似） 这是 Linear Attention 和标准 Attention 最本质的差距：标准 Attention 能精确访问所有历史 token，Linear Attention 只能访问被压缩的历史摘要。\n七、改进方向：现代线性注意力模型 RetNet（2023，微软） 引入衰减因子（Decay），让远距离的历史信息自然衰减，缓解\u0026quot;遗忘\u0026quot;问题：\n$$S_t = \\gamma \\cdot S_{t-1} + \\phi(k_t) v_t^T, \\quad \\gamma \\in (0, 1)$$\n衰减因子让模型更关注近期信息，类似 LSTM 的遗忘门。\nMamba（2023，SSM 架构） 不用核函数近似，而是用状态空间模型（SSM） 的框架，通过选择性机制（Selective SSM）让隐状态能动态决定\u0026quot;记住什么、忘记什么\u0026quot;，效果接近 Transformer，推理速度接近 RNN。\nGLA（Gated Linear Attention，2024） 在 Linear Attention 的递推公式里加入门控机制：\n$$S_t = G_t \\odot S_{t-1} + \\phi(k_t) v_t^T$$\n其中 $G_t$ 是由输入动态生成的门控矩阵，让模型自适应地控制历史信息的保留程度。\n八、和标准 Attention / Flash Attention 的对比 标准 Attention Flash Attention Linear Attention 训练复杂度 $O(N^2 d)$ $O(N^2 d)$（显存优化） $O(N d^2)$ 推理复杂度（逐步生成） $O(N d)$（有KV Cache） $O(N d)$ $O(d^2)$（固定！） 推理显存 $O(Nd)$（KV Cache随序列增长） $O(Nd)$ $O(d^2)$（固定！） 表达能力 最强（精确 softmax） 同标准 较弱（核函数近似） 长序列支持 困难 较好 天然支持 是否需要 KV Cache 需要 需要 不需要（状态固定大小） Linear Attention 推理时最大的优势：KV Cache 不随序列长度增长，固定是 $d \\times d$ 的矩阵，推理 1K 和推理 1M token 占用的显存完全一样。\n九、核心要点速查 问题 答案 Linear Attention 解决什么问题？ 把标准 Attention 的 $O(N^2)$ 复杂度降到 $O(N)$ 核心技巧是什么？ 改变矩阵乘法顺序，先算 $K^TV$（$d \\times d$）再乘 $Q$ 为什么不直接去掉 softmax？ softmax 保证归一化，直接去掉数值不稳定，需要用核函数替代 递推形式是什么？ $S_t = S_{t-1} + \\phi(k_t)v_t^T$，和 RNN 结构一样 最大的缺点是什么？ 历史信息被压缩进固定大小的矩阵 $S$，远距离依赖容易丢失 推理时的优势是什么？ 不需要 KV Cache，显存固定为 $O(d^2)$，与序列长度无关 代表性改进模型？ RetNet（衰减因子）、Mamba（SSM）、GLA（门控） ","permalink":"https://afan.ml/posts/linear-attention/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：标准 Attention 的复杂度是 $O(N^2)$，Linear Attention 通过改变计算顺序，把复杂度降到 $O(N)$，代价是用核函数近似替代 softmax，牺牲了一定的表达能力，但换来了推理时的递推形式（类似 RNN），天然支持无限长序列。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一问题背景标准-attention-的瓶颈\"\u003e一、问题背景：标准 Attention 的瓶颈\u003c/h2\u003e\n\u003cp\u003e标准 Attention 的计算：\u003c/p\u003e\n\u003cp\u003e$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right) V$$\u003c/p\u003e\n\u003cp\u003e序列长度为 $N$ 时：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e计算 $QK^T$ 需要 $O(N^2)$ 时间和空间\u003c/li\u003e\n\u003cli\u003e中间的 $N \\times N$ 注意力矩阵需要 $O(N^2)$ 显存\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e\u003cstrong\u003e后果\u003c/strong\u003e：序列长度翻倍，计算量翻 4 倍，显存翻 4 倍。长序列（如 100K token）几乎不可能用标准 Attention。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二核心思想改变计算顺序\"\u003e二、核心思想：改变计算顺序\u003c/h2\u003e\n\u003ch3 id=\"标准-attention-的计算顺序\"\u003e标准 Attention 的计算顺序\u003c/h3\u003e\n\u003cp\u003e$$\\text{out} = \\underbrace{\\text{softmax}(QK^T)}_{\\text{先算这个，}N \\times N \\text{ 矩阵}} V$$\u003c/p\u003e\n\u003cp\u003e必须先把完整的 $N \\times N$ 矩阵算出来，才能乘以 $V$。\u003c/p\u003e\n\u003ch3 id=\"linear-attention-的关键洞察\"\u003eLinear Attention 的关键洞察\u003c/h3\u003e\n\u003cp\u003e如果去掉 softmax，矩阵乘法满足结合律：\u003c/p\u003e\n\u003cp\u003e$$\\text{out} = (QK^T) V = Q \\underbrace{(K^T V)}_{\\text{先算这个，}d \\times d \\text{ 矩阵}}$$\u003c/p\u003e","title":"Linear Attention（线性注意力）笔记"},{"content":" 一句话：LoRA 用低秩矩阵大幅减少微调参数量；QLoRA 在此基础上把基础模型量化到 4bit，让消费级 GPU 能微调 70B 模型；Unsloth 是工程加速库，把整个流程再提速 2-5×。三者是层层叠加的关系。\n一、为什么需要参数高效微调（PEFT）？ 全量微调（Full Fine-tuning）的问题：\nGPT-3（175B 参数）全量微调： - 参数量：175B × 4 bytes（FP32）= 700 GB 显存 - 梯度：再 × 1 = 700 GB - 优化器状态（Adam）：再 × 2 = 1400 GB - 合计：约 2800 GB 显存 → 需要 35 张 A100（80GB） 核心矛盾：预训练模型越来越大，但大多数下游任务只需要微调模型的\u0026quot;方向\u0026quot;，不需要改变所有参数。\n参数高效微调（PEFT） 的思路：冻结大部分参数，只训练少量新增参数，效果接近全量微调，但显存和计算量大幅减少。\n二、LoRA：低秩适配（Low-Rank Adaptation） 核心思想 LoRA（2021，微软）的关键洞察：预训练模型的权重更新矩阵是低秩的。\n全量微调时，权重更新为：\n$$W\u0026rsquo; = W_0 + \\Delta W$$\n其中 $W_0 \\in \\mathbb{R}^{d \\times k}$ 是预训练权重，$\\Delta W$ 是更新量。\nLoRA 假设 $\\Delta W$ 是低秩的，用两个小矩阵的乘积来近似：\n$$\\Delta W = B A, \\quad A \\in \\mathbb{R}^{r \\times k},\\ B \\in \\mathbb{R}^{d \\times r},\\ r \\ll \\min(d, k)$$\n前向传播：\n$$h = W_0 x + \\Delta W x = W_0 x + B A x$$\n$W_0$：冻结，不参与梯度计算 $A$、$B$：可训练，参数量极少 参数量对比 以 GPT-3 的某一层为例，$d = k = 12288$：\n全量微调：d × k = 12288 × 12288 = 150,994,944 个参数 LoRA（r=8）：r×k + d×r = 8×12288 + 12288×8 = 196,608 个参数 压缩比：约 768:1 初始化方式 $A$：随机高斯初始化（保证初始时有信号） $B$：全零初始化（保证训练开始时 $\\Delta W = BA = 0$，不破坏预训练权重） 缩放因子 α 实际使用时加一个缩放因子：\n$$h = W_0 x + \\frac{\\alpha}{r} B A x$$\n$\\alpha$：超参数，通常设为 $r$ 或 $2r$ $\\frac{\\alpha}{r}$：让不同 $r$ 值下的学习率等效，方便调参 应用到哪些层？ LoRA 通常应用于 Transformer 的注意力层权重：\nQ 投影矩阵 W_Q ← 加 LoRA K 投影矩阵 W_K ← 加 LoRA V 投影矩阵 W_V ← 加 LoRA 输出投影 W_O ← 加 LoRA（可选） FFN 层 ← 可选，通常不加 推理时的合并 训练完成后，可以把 LoRA 权重合并回原始权重，推理时零额外开销：\n$$W\u0026rsquo; = W_0 + BA$$\n合并后的模型和原始模型结构完全相同，不需要任何特殊推理代码。\n完整代码实现 import math import torch import torch.nn as nn import torch.nn.functional as F class LoRALinear(nn.Module): \u0026#34;\u0026#34;\u0026#34; 带 LoRA 的线性层。 原始权重 W_0 冻结，只训练低秩矩阵 A 和 B。 \u0026#34;\u0026#34;\u0026#34; def __init__( self, in_features: int, out_features: int, rank: int = 8, alpha: float = 16.0, dropout_rate: float = 0.0, ): super().__init__() self.in_features = in_features self.out_features = out_features self.rank = rank self.scaling = alpha / rank # α/r 缩放因子 # 原始权重：冻结 self.weight = nn.Parameter( torch.randn(out_features, in_features) * 0.02, requires_grad=False, # 冻结！ ) # LoRA 矩阵：可训练 # A 用 Kaiming-uniform 初始化（与 nn.Linear 默认初始化一致，PEFT 库标准做法） # 不用高斯 * 0.02，因为 Kaiming-uniform 能更好地保持各层激活值的方差 self.lora_A = nn.Parameter(torch.empty(rank, in_features)) nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_B = nn.Parameter(torch.zeros(out_features, rank)) # B 初始化为 0，保证 ΔW=0 # 可选的 dropout（防止 LoRA 过拟合） self.lora_dropout = nn.Dropout(dropout_rate) if dropout_rate \u0026gt; 0 else nn.Identity() def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: # 原始路径（冻结权重） base_output = F.linear(x, self.weight) # LoRA 路径：x → A → B，乘以缩放因子 lora_output = F.linear(self.lora_dropout(x), self.lora_A) # [*, rank] lora_output = F.linear(lora_output, self.lora_B) # [*, out_features] return base_output + self.scaling * lora_output def merge_weights(self) -\u0026gt; nn.Linear: \u0026#34;\u0026#34;\u0026#34; 将 LoRA 权重合并到原始权重，返回标准 Linear 层（推理用）。 合并后推理无额外开销。 \u0026#34;\u0026#34;\u0026#34; merged_weight = self.weight + self.scaling * (self.lora_B @ self.lora_A) merged_layer = nn.Linear(self.in_features, self.out_features, bias=False) merged_layer.weight = nn.Parameter(merged_weight) return merged_layer def apply_lora_to_attention(model: nn.Module, rank: int = 8, alpha: float = 16.0) -\u0026gt; nn.Module: \u0026#34;\u0026#34;\u0026#34; 将模型中所有名为 q_proj、k_proj、v_proj、o_proj 的线性层替换为 LoRA 版本。 其余参数全部冻结。 \u0026#34;\u0026#34;\u0026#34; # 先冻结所有参数 for param in model.parameters(): param.requires_grad = False # 替换目标层为 LoRA 版本 target_layer_names = {\u0026#34;q_proj\u0026#34;, \u0026#34;k_proj\u0026#34;, \u0026#34;v_proj\u0026#34;, \u0026#34;o_proj\u0026#34;} for module_name, module in model.named_modules(): for layer_name in target_layer_names: if hasattr(module, layer_name): original_layer = getattr(module, layer_name) if isinstance(original_layer, nn.Linear): lora_layer = LoRALinear( in_features=original_layer.in_features, out_features=original_layer.out_features, rank=rank, alpha=alpha, ) # 复制原始权重（冻结） lora_layer.weight.data = original_layer.weight.data.clone() setattr(module, layer_name, lora_layer) return model def count_trainable_params(model: nn.Module) -\u0026gt; tuple[int, int]: \u0026#34;\u0026#34;\u0026#34;返回 (可训练参数量, 总参数量)\u0026#34;\u0026#34;\u0026#34; trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) total = sum(p.numel() for p in model.parameters()) return trainable, total # ── 验证 ────────────────────────────────────────────────────────────────────── if __name__ == \u0026#34;__main__\u0026#34;: torch.manual_seed(42) # 验证 LoRALinear 的前向传播和权重合并 lora_layer = LoRALinear(in_features=512, out_features=512, rank=8, alpha=16.0) x = torch.randn(2, 10, 512) output = lora_layer(x) print(f\u0026#34;LoRA 输出形状: {output.shape}\u0026#34;) # [2, 10, 512] # 验证合并前后输出一致 merged_layer = lora_layer.merge_weights() output_merged = merged_layer(x) max_diff = (output - output_merged).abs().max().item() print(f\u0026#34;合并前后最大误差: {max_diff:.2e}\u0026#34;) # 应接近 0 # 统计参数量 trainable_params = sum(p.numel() for p in lora_layer.parameters() if p.requires_grad) total_params = sum(p.numel() for p in lora_layer.parameters()) frozen_params = total_params - trainable_params print(f\u0026#34;可训练参数: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)\u0026#34;) print(f\u0026#34;冻结参数: {frozen_params:,} ({frozen_params/total_params*100:.1f}%)\u0026#34;) 三、QLoRA：量化 + LoRA 核心思想 QLoRA（2023，华盛顿大学）解决的问题：LoRA 虽然减少了可训练参数，但基础模型本身还是 FP16/BF16，70B 模型仍需 140GB 显存。\nQLoRA 的三个关键技术：\n技术 1：NF4（4-bit NormalFloat）量化 把基础模型权重从 BF16（16bit）量化到 NF4（4bit），显存减少 4×。\nNF4 不是普通的 INT4，而是专门为神经网络权重设计的：\n神经网络权重通常服从正态分布 N(0, σ²) NF4 的量化点是正态分布的等分位数点（quantile）： → 在权重密集的区域（靠近 0）分配更多量化级别 → 在权重稀疏的区域（远离 0）分配更少量化级别 → 比均匀量化的 INT4 精度更高 NF4 的 16 个量化值（归一化到 [-1, 1]，来自 QLoRA 论文 Appendix E 及 bitsandbytes 实现）：\nNF4_levels = [ -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0, ] 这些值由 scipy.stats.norm.ppf（正态分布的逆累积分布函数）生成等分位数后归一化到 [-1, 1] 得到。注意分布是非对称的：有 9 个非负值（含 0）和 7 个负值，因为 QLoRA 使用非对称量化以确保 0.0 能被精确表示（零点 Z 映射到固定的非零量化值）。\n技术 2：双重量化（Double Quantization） 量化本身需要存储\u0026quot;量化常数\u0026quot;（scale factor），QLoRA 对量化常数再做一次量化：\n第一次量化：BF16 权重 → NF4（量化常数用 FP32 存储） 第二次量化：FP32 量化常数 → FP8（再节省约 0.37 bit/参数） 技术 3：分页优化器（Paged Optimizer） 使用 NVIDIA 的统一内存（Unified Memory），当 GPU 显存不足时，自动把优化器状态换页到 CPU 内存，避免 OOM。\nQLoRA 的训练流程 基础模型（BF16） ↓ 量化（NF4 + 双重量化） 基础模型（NF4，冻结） ← 显存：原来的 1/4 ↓ 添加 LoRA 适配器（BF16，可训练） 前向传播： NF4 权重 → 反量化到 BF16 → 正常计算 LoRA 路径：BF16 全精度计算 反向传播： 只对 LoRA 参数计算梯度（BF16） 基础模型梯度不保存 关键点：计算时临时反量化到 BF16，但存储时保持 NF4，所以显存占用是 NF4 级别的。\n显存对比（以 LLaMA-65B 为例） 方法 显存需求 所需 GPU 全量微调（BF16） ~780 GB 10× A100 80G LoRA（BF16 基础模型） ~130 GB 2× A100 80G QLoRA（NF4 基础模型） ~48 GB 1× A100 80G QLoRA + Unsloth ~38 GB 1× A100 80G 四、Unsloth：工程加速 Unsloth 是什么 Unsloth 是一个开源库（Daniel Han 等人，2023），不改变 LoRA/QLoRA 的算法，而是通过以下工程手段提速：\n手段 1：手写 Triton kernel 用 OpenAI Triton（GPU 编程语言）重写了 LoRA 的前向/反向传播，比 PyTorch 默认实现快 2-3×：\nPyTorch 默认： x → Linear(A) → 中间结果存到 HBM → Linear(B) → 输出 （两次 HBM 读写） Unsloth Triton kernel： x → Linear(A) → 直接在 SRAM 中 → Linear(B) → 输出 （一次 HBM 读写，减少 IO 瓶颈） 手段 2：RoPE 和 RMSNorm 的融合 kernel 把 RoPE 位置编码和 RMSNorm 的计算融合进单个 kernel，减少 kernel 启动开销和中间结果的显存读写。\n手段 3：梯度检查点优化 标准梯度检查点（Gradient Checkpointing）会重新计算前向传播以节省显存，但 Unsloth 只对部分层做检查点，并缓存计算代价高的中间结果，在显存和速度之间取得更好的平衡。\nUnsloth 的使用方式 from unsloth import FastLanguageModel from trl import SFTTrainer from transformers import TrainingArguments from datasets import load_dataset import torch # 1. 加载模型（自动应用 QLoRA + Unsloth 优化） # 模型名称末尾带 \u0026#34;unsloth-bnb-4bit\u0026#34; 的是 Unsloth 动态 4bit 量化版，精度优于标准 bnb-4bit model, tokenizer = FastLanguageModel.from_pretrained( model_name=\u0026#34;unsloth/llama-3.1-8b-unsloth-bnb-4bit\u0026#34;, max_seq_length=2048, dtype=None, # None = 自动选择（新 GPU 用 BF16，旧 GPU 用 FP16） load_in_4bit=True, # QLoRA：4bit 量化基础模型 ) # 2. 添加 LoRA 适配器 model = FastLanguageModel.get_peft_model( model, r=16, # LoRA rank，从 16 开始试 target_modules=[ # 应用 LoRA 的层（包含 FFN 效果更好） \u0026#34;q_proj\u0026#34;, \u0026#34;k_proj\u0026#34;, \u0026#34;v_proj\u0026#34;, \u0026#34;o_proj\u0026#34;, \u0026#34;gate_proj\u0026#34;, \u0026#34;up_proj\u0026#34;, \u0026#34;down_proj\u0026#34;, ], lora_alpha=16, # α = r，scaling = 1.0，简化调参 lora_dropout=0.0, # Unsloth 的 Triton kernel 假设无 dropout，设为 0 bias=\u0026#34;none\u0026#34;, # 不训练 bias use_gradient_checkpointing=\u0026#34;unsloth\u0026#34;, # Unsloth 优化版梯度检查点，比标准版省更多显存 random_state=42, ) # 3. 准备数据集 # Unsloth 推荐用 formatting_func 动态格式化，避免提前 map 整个数据集 alpaca_prompt = \u0026#34;\u0026#34;\u0026#34;Below is an instruction that describes a task. Write a response that appropriately completes the request. ### Instruction: {} ### Response: {}\u0026#34;\u0026#34;\u0026#34; def formatting_func(examples): \u0026#34;\u0026#34;\u0026#34;将数据集的每条样本格式化为训练文本，末尾加 EOS token。\u0026#34;\u0026#34;\u0026#34; texts = [] for instruction, output in zip(examples[\u0026#34;instruction\u0026#34;], examples[\u0026#34;output\u0026#34;]): text = alpaca_prompt.format(instruction, output) + tokenizer.eos_token texts.append(text) return texts dataset = load_dataset(\u0026#34;yahma/alpaca-cleaned\u0026#34;, split=\u0026#34;train\u0026#34;) # 4. 训练 trainer = SFTTrainer( model=model, tokenizer=tokenizer, train_dataset=dataset, formatting_func=formatting_func, # 动态格式化，比 dataset_text_field 更灵活 max_seq_length=2048, args=TrainingArguments( per_device_train_batch_size=2, gradient_accumulation_steps=4, # 等效 batch_size = 8 warmup_steps=5, num_train_epochs=1, learning_rate=2e-4, fp16=not torch.cuda.is_bf16_supported(), bf16=torch.cuda.is_bf16_supported(), logging_steps=1, optim=\u0026#34;adamw_8bit\u0026#34;, # 8bit Adam，进一步节省显存 output_dir=\u0026#34;outputs\u0026#34;, ), ) trainer.train() # 5. 保存 LoRA adapter（只有几十 MB，不含基础模型） model.save_pretrained(\u0026#34;lora_adapter\u0026#34;) tokenizer.save_pretrained(\u0026#34;lora_adapter\u0026#34;) # 6. 推理：加载 adapter 并切换到推理模式 model, tokenizer = FastLanguageModel.from_pretrained( model_name=\u0026#34;lora_adapter\u0026#34;, max_seq_length=2048, load_in_4bit=True, ) FastLanguageModel.for_inference(model) # 必须调用，启用 Unsloth 的推理加速 inputs = tokenizer( [alpaca_prompt.format(\u0026#34;Explain LoRA in one sentence.\u0026#34;, \u0026#34;\u0026#34;)], return_tensors=\u0026#34;pt\u0026#34;, ).to(\u0026#34;cuda\u0026#34;) outputs = model.generate(**inputs, max_new_tokens=128) print(tokenizer.decode(outputs[0], skip_special_tokens=True)) 五、LoRA 的超参数选择指南 rank r 的选择 任务类型 推荐 r 说明 简单指令跟随 8-16 任务简单，低秩足够 领域知识注入 16-64 需要更多容量 代码生成 32-64 代码任务复杂度高 多任务微调 64-128 需要覆盖多种能力 经验规则：r 越大，可训练参数越多，效果上限越高，但过拟合风险也越大。通常从 r=16 开始试。\nalpha α 的选择 通常设为 α = r 或 α = 2r Unsloth 推荐 α = r（即 scaling = 1.0），简化调参 实际上 α 和学习率是耦合的，调一个等价于调另一个 应用到哪些层 只加 Q、V（原始 LoRA 论文）：参数最少，效果一般 加 Q、K、V、O（常见做法）：效果更好 加所有线性层（包括 FFN）：效果最好，参数也最多 Unsloth 默认推荐加所有线性层（gate_proj、up_proj、down_proj 也加）。\n六、保存和加载 LoRA 权重的方式 方式 1：只保存 LoRA adapter（推荐） # 保存：只有几十 MB model.save_pretrained(\u0026#34;lora_adapter\u0026#34;) # 加载：需要同时加载基础模型 + adapter from transformers import AutoModelForCausalLM from peft import PeftModel base_model = AutoModelForCausalLM.from_pretrained(\u0026#34;meta-llama/Llama-3-8B\u0026#34;) model = PeftModel.from_pretrained(base_model, \u0026#34;lora_adapter\u0026#34;) 方式 2：合并后保存完整模型 # 合并 LoRA 权重到基础模型（推理时无额外开销） merged_model = model.merge_and_unload() merged_model.save_pretrained(\u0026#34;merged_model\u0026#34;) # 完整模型，几 GB 方式 3：保存为 GGUF 格式（本地推理用） # Unsloth 支持直接导出 GGUF，供 llama.cpp / Ollama 使用 model.save_pretrained_gguf( \u0026#34;model_gguf\u0026#34;, tokenizer, quantization_method=\u0026#34;q4_k_m\u0026#34;, # 4bit 量化，平衡精度和速度 ) 七、三者的完整对比 LoRA QLoRA QLoRA + Unsloth 基础模型精度 BF16/FP16 NF4（4bit） NF4（4bit） LoRA 精度 BF16 BF16 BF16 显存（7B 模型） ~16 GB ~6 GB ~5 GB 训练速度 基准 略慢（反量化开销） 快 2-5×（Triton kernel） 效果损失 极小 小（NF4 精度损失） 同 QLoRA 适用 GPU A100/H100 消费级（RTX 3090/4090） 消费级 安装复杂度 简单（pip install peft） 中等（需要 bitsandbytes） 简单（pip install unsloth） 八、核心要点速查 问题 答案 LoRA 的核心公式？ $\\Delta W = BA$，$A \\in \\mathbb{R}^{r \\times k}$，$B \\in \\mathbb{R}^{d \\times r}$，$r \\ll d,k$ 为什么 B 初始化为 0？ 保证训练开始时 $\\Delta W = 0$，不破坏预训练权重 α/r 缩放因子的作用？ 让不同 rank 下的学习率等效，方便调参 QLoRA 用什么量化？ NF4（4-bit NormalFloat），基于正态分布分位数设计 双重量化是什么？ 对量化常数再做一次量化，额外节省 0.37 bit/参数 Unsloth 怎么加速的？ 手写 Triton kernel，减少 HBM 读写；融合 RoPE/RMSNorm kernel 推理时 LoRA 有额外开销吗？ 合并后无开销；不合并则有一次额外矩阵乘法 消费级 GPU 能微调多大模型？ RTX 4090（24GB）+ QLoRA + Unsloth 可微调 70B 模型 ","permalink":"https://afan.ml/posts/lo-r-a--q-lo-r-a-%E5%BE%AE%E8%B0%83/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：LoRA 用低秩矩阵大幅减少微调参数量；QLoRA 在此基础上把基础模型量化到 4bit，让消费级 GPU 能微调 70B 模型；Unsloth 是工程加速库，把整个流程再提速 2-5×。三者是层层叠加的关系。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要参数高效微调peft\"\u003e一、为什么需要参数高效微调（PEFT）？\u003c/h2\u003e\n\u003cp\u003e全量微调（Full Fine-tuning）的问题：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003eGPT-3（175B 参数）全量微调：\n - 参数量：175B × 4 bytes（FP32）= 700 GB 显存\n - 梯度：再 × 1 = 700 GB\n - 优化器状态（Adam）：再 × 2 = 1400 GB\n - 合计：约 2800 GB 显存 → 需要 35 张 A100（80GB）\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e核心矛盾\u003c/strong\u003e：预训练模型越来越大，但大多数下游任务只需要微调模型的\u0026quot;方向\u0026quot;，不需要改变所有参数。\u003c/p\u003e\n\u003cp\u003e\u003cstrong\u003e参数高效微调（PEFT）\u003c/strong\u003e 的思路：\u003cstrong\u003e冻结大部分参数，只训练少量新增参数\u003c/strong\u003e，效果接近全量微调，但显存和计算量大幅减少。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二lora低秩适配low-rank-adaptation\"\u003e二、LoRA：低秩适配（Low-Rank Adaptation）\u003c/h2\u003e\n\u003ch3 id=\"核心思想\"\u003e核心思想\u003c/h3\u003e\n\u003cp\u003eLoRA（2021，微软）的关键洞察：\u003cstrong\u003e预训练模型的权重更新矩阵是低秩的\u003c/strong\u003e。\u003c/p\u003e\n\u003cp\u003e全量微调时，权重更新为：\u003c/p\u003e\n\u003cp\u003e$$W\u0026rsquo; = W_0 + \\Delta W$$\u003c/p\u003e\n\u003cp\u003e其中 $W_0 \\in \\mathbb{R}^{d \\times k}$ 是预训练权重，$\\Delta W$ 是更新量。\u003c/p\u003e\n\u003cp\u003eLoRA 假设 $\\Delta W$ 是低秩的，用两个小矩阵的乘积来近似：\u003c/p\u003e","title":"LoRA / QLoRA / Unsloth 微调笔记"},{"content":" 一句话：Mamba 是 Linear Attention 的\u0026quot;升级版\u0026quot;——同样是 $O(N)$ 复杂度、固定大小隐状态，但通过选择性机制（Selective SSM） 让模型能动态决定\u0026quot;记住什么、忘记什么\u0026quot;，效果接近 Transformer，推理速度接近 RNN。\n一、从 Linear Attention 到 Mamba：解决什么问题？ 回顾 Linear Attention 的递推公式：\n$$S_t = S_{t-1} + \\phi(k_t) v_t^T$$\n问题：衰减是固定的（没有衰减，或者 RetNet 里用固定的 $\\gamma$），模型无法根据输入内容动态决定\u0026quot;这个历史信息重不重要\u0026quot;。\n类比：\nRNN（LSTM）：有遗忘门，可以选择性地清除历史 Linear Attention：没有遗忘门，历史信息只增不减，全部堆进 S 矩阵 RetNet：有固定衰减 γ，但 γ 是超参数，不随输入变化 Mamba：衰减因子由输入动态生成，每个 token 的\u0026#34;遗忘程度\u0026#34;不同 二、SSM 的数学基础 Mamba 基于状态空间模型（State Space Model，SSM），这是控制论里的经典框架。\n连续时间 SSM $$h\u0026rsquo;(t) = A h(t) + B x(t)$$ $$y(t) = C h(t)$$\n$x(t)$：输入信号 $h(t)$：隐状态（类比 RNN 的 hidden state） $y(t)$：输出 $A$：状态转移矩阵（控制历史信息如何演化） $B$：输入投影矩阵（控制输入如何影响隐状态） $C$：输出投影矩阵（控制隐状态如何映射到输出） 离散化（实际使用的形式） 连续 SSM 需要离散化才能用于序列建模，使用零阶保持（ZOH） 方法：\n$$\\bar{A} = e^{\\Delta A}, \\quad \\bar{B} = (e^{\\Delta A} - I) A^{-1} B \\approx \\Delta B$$\n其中 $\\Delta$（dt_bias 对应的参数）是步长（step size），控制离散化的粒度。\n离散化后的递推公式：\n$$h_t = \\bar{A} h_{t-1} + \\bar{B} x_t$$ $$y_t = C h_t$$\n这和 RNN 的结构完全一样！\nRNN： h_t = tanh(W_h h_{t-1} + W_x x_t) SSM： h_t = Ā h_{t-1} + B̄ x_t 区别： SSM 的 Ā 有特殊结构（来自连续系统的离散化），更有理论保证 三、S4：Mamba 的前身 S4（Structured State Space Sequence Model，2021） 是 Mamba 的直接前身。\nS4 的关键设计：把 $A$ 矩阵限制为对角加低秩（DPLR） 结构，使得：\n可以高效并行计算（卷积形式） 可以高效递推（RNN 形式） 理论上能捕获超长距离依赖 S4 的问题：$A$、$B$、$C$ 都是固定参数，不随输入变化——内容无关（content-unaware）。\nS4 处理\u0026#34;今天天气很好\u0026#34;和\u0026#34;今天天气很差\u0026#34;时， 用的是完全相同的状态转移矩阵 Ā， 模型无法根据\u0026#34;好\u0026#34;还是\u0026#34;差\u0026#34;来决定记住多少。 四、Mamba 的核心创新：选择性机制 Mamba（2023，Albert Gu \u0026amp; Tri Dao）的核心贡献：让 $B$、$C$、$\\Delta$ 依赖于输入 $x_t$。\n对比 S4 和 Mamba 参数 S4 Mamba $A$ 固定（训练后不变） 固定（但结构特殊） $B$ 固定 由 $x_t$ 动态生成 $C$ 固定 由 $x_t$ 动态生成 $\\Delta$ 固定 由 $x_t$ 动态生成 选择性机制的直觉 $$B_t = \\text{Linear}_B(x_t), \\quad C_t = \\text{Linear}C(x_t), \\quad \\Delta_t = \\text{softplus}(\\text{Linear}\\Delta(x_t))$$\n$$\\bar{A}_t = e^{\\Delta_t A}, \\quad \\bar{B}_t = \\Delta_t B_t$$\n$$h_t = \\bar{A}t h{t-1} + \\bar{B}_t x_t$$ $$y_t = C_t h_t$$\n$\\Delta_t$ 的作用（这就是 config.json 里的 dt_bias）：\nΔ_t 很大 → Ā_t ≈ 0，B̄_t ≈ B_t → h_t ≈ B_t x_t（几乎忘掉历史，专注当前输入） → 相当于\u0026#34;重置\u0026#34;隐状态 Δ_t 很小 → Ā_t ≈ I，B̄_t ≈ 0 → h_t ≈ h_{t-1}（几乎忽略当前输入，保留历史） → 相当于\u0026#34;跳过\u0026#34;当前 token 这就是 Mamba 的选择性：模型学会了对重要 token 用大 $\\Delta$（重置并记住），对不重要 token 用小 $\\Delta$（直接跳过）。\n五、A 矩阵：HiPPO 初始化 $A$ 矩阵虽然固定，但初始化方式很关键。Mamba 使用 HiPPO（High-order Polynomial Projection Operators） 初始化：\n$$A_{nk} = -\\begin{cases} (2n+1)^{1/2}(2k+1)^{1/2} \u0026amp; n \u0026gt; k \\ n+1 \u0026amp; n = k \\ 0 \u0026amp; n \u0026lt; k \\end{cases}$$\n直觉：HiPPO 矩阵被设计为能最优地压缩历史信息——隐状态 $h$ 相当于对历史输入的多项式近似系数，理论上能记住任意长距离的依赖。\n这就是 config.json 里 linear_attn.A_log 的来源：$A$ 以对数形式存储（A_log = log(-A)），保证离散化后 $\\bar{A}$ 的特征值在单位圆内（系统稳定）。\n六、conv1d 的作用 config.json 里还有 linear_attn.conv1d，这是 Mamba 的另一个设计：\n在 SSM 之前，先做一个短程卷积：\nx_t → conv1d（kernel_size=4）→ SSM → y_t 为什么需要 conv1d？\nSSM 的隐状态是全局的（压缩了所有历史），但对局部特征（如 n-gram、短语结构）不敏感。conv1d 用小卷积核（kernel_size=4，对应 config 里的 linear_conv_kernel_dim: 4）捕获局部模式，作为 SSM 的补充。\nconv1d：捕获局部特征（短程） SSM： 捕获全局依赖（长程） 两者结合：覆盖所有尺度的依赖 七、Mamba Block 的完整结构 输入 x [seq_len, d_model] │ ├─────────────────────────────┐ │ │ Linear(d_model → d_inner) Linear(d_model → d_inner) │ │ SiLU │ │ │ conv1d(kernel=4) │ │ │ SiLU │ │ │ SSM（选择性状态空间） │ ┌─────────────────────┐ │ │ B_t = Linear_B(x_t) │ │ │ C_t = Linear_C(x_t) │ │ │ Δ_t = Linear_Δ(x_t) │ │ │ h_t = Ā_t h_{t-1} + B̄_t x_t │ │ │ y_t = C_t h_t │ │ └─────────────────────┘ │ │ │ └──────── × ──────────────────┘ │（门控：SSM输出 × 线性分支） │ Linear(d_inner → d_model) │ 输出 y 门控设计：右边的线性分支类似 SwiGLU 里的门控，让模型能选择性地\u0026quot;放大\u0026quot;或\u0026quot;抑制\u0026quot; SSM 的输出。\n八、完整代码实现 import torch import torch.nn as nn import torch.nn.functional as F import math class SelectiveSSM(nn.Module): \u0026#34;\u0026#34;\u0026#34; Mamba 的核心：选择性状态空间模型（Selective SSM）。 d_model: 输入维度 d_state: 隐状态维度（论文中通常为 16） \u0026#34;\u0026#34;\u0026#34; def __init__(self, d_model: int, d_state: int = 16): super().__init__() self.d_model = d_model self.d_state = d_state # A 矩阵：以对数形式存储，保证离散化后稳定 # 初始化为 log(1, 2, ..., d_state)，近似 HiPPO A = torch.arange(1, d_state + 1, dtype=torch.float32).unsqueeze(0).expand(d_model, -1) self.A_log = nn.Parameter(torch.log(A)) # [d_model, d_state] # B、C、Δ 由输入动态生成 self.linear_B = nn.Linear(d_model, d_state, bias=False) self.linear_C = nn.Linear(d_model, d_state, bias=False) self.linear_delta = nn.Linear(d_model, d_model, bias=True) # dt_bias 在这里 # D：跳跃连接（直接将输入加到输出） self.D = nn.Parameter(torch.ones(d_model)) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; x: [batch, seq_len, d_model] 返回: [batch, seq_len, d_model] \u0026#34;\u0026#34;\u0026#34; batch_size, seq_len, d_model = x.shape d_state = self.d_state # 恢复 A（负数，保证系统稳定） A = -torch.exp(self.A_log) # [d_model, d_state] # 动态生成 B、C、Δ B = self.linear_B(x) # [batch, seq_len, d_state] C = self.linear_C(x) # [batch, seq_len, d_state] delta = F.softplus(self.linear_delta(x)) # [batch, seq_len, d_model]，保证 Δ \u0026gt; 0 # 离散化（ZOH，零阶保持）： # Ā = exp(Δ·A) # B̄ = (ΔA)⁻¹ · (exp(ΔA) - I) · ΔB # 因为 A 是对角矩阵（每个 d_model 维度独立对应 d_state 个状态）， # (ΔA)⁻¹ · (exp(ΔA) - I) 可以逐元素计算，无需矩阵求逆： # B̄ = (exp(ΔA) - 1) / A · B = (Ā - 1) / A · B # # delta: [batch, seq_len, d_model] # A: [d_model, d_state] # delta_A: [batch, seq_len, d_model, d_state] delta_A_product = delta.unsqueeze(-1) * A # ΔA，[batch, seq_len, d_model, d_state] delta_A = torch.exp(delta_A_product) # Ā = exp(ΔA)，[batch, seq_len, d_model, d_state] # ZOH 的 B̄：逐元素 (exp(ΔA) - 1) / A · B # 当 A 接近 0 时用 Taylor 展开近似（数值稳定），但实际中 A 初始化为负整数，不会为 0 delta_B_zoh = (delta_A - 1.0) / A # (exp(ΔA) - 1) / A，[batch, seq_len, d_model, d_state] delta_B = delta_B_zoh * B.unsqueeze(2) # B̄ = delta_B_zoh · B，[batch, seq_len, d_model, d_state] # 递推计算隐状态 # h: [batch, d_model, d_state] h = torch.zeros(batch_size, d_model, d_state, device=x.device) outputs = [] for t in range(seq_len): # h_t = Ā_t ⊙ h_{t-1} + B̄_t ⊙ x_t h = delta_A[:, t] * h + delta_B[:, t] * x[:, t].unsqueeze(-1) # y_t = C_t · h_t（对 d_state 维度求和） y_t = (C[:, t].unsqueeze(1) * h).sum(dim=-1) # [batch, d_model] outputs.append(y_t) y = torch.stack(outputs, dim=1) # [batch, seq_len, d_model] # 跳跃连接：y += D * x y = y + self.D * x return y class MambaBlock(nn.Module): \u0026#34;\u0026#34;\u0026#34; 完整的 Mamba Block，包含门控结构和 conv1d。 d_model: 模型维度 d_inner: 内部扩展维度（通常为 d_model * 2） d_state: SSM 隐状态维度 conv_kernel: 局部卷积核大小 \u0026#34;\u0026#34;\u0026#34; def __init__(self, d_model: int, d_inner: int = None, d_state: int = 16, conv_kernel: int = 4): super().__init__() self.d_model = d_model self.d_inner = d_inner or d_model * 2 self.norm = nn.LayerNorm(d_model) # 输入投影：分成两路（SSM 路 + 门控路） self.input_proj = nn.Linear(d_model, self.d_inner * 2, bias=False) # 局部卷积（捕获短程特征） self.conv1d = nn.Conv1d( in_channels=self.d_inner, out_channels=self.d_inner, kernel_size=conv_kernel, padding=conv_kernel - 1, # causal padding groups=self.d_inner, # depthwise conv bias=True ) # 选择性 SSM self.ssm = SelectiveSSM(self.d_inner, d_state) # 输出投影 self.output_proj = nn.Linear(self.d_inner, d_model, bias=False) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; x: [batch, seq_len, d_model] 返回: [batch, seq_len, d_model]（残差连接在外部处理） \u0026#34;\u0026#34;\u0026#34; residual = x x = self.norm(x) # 分成两路 projected = self.input_proj(x) # [batch, seq_len, d_inner*2] ssm_branch, gate_branch = projected.chunk(2, dim=-1) # 各 [batch, seq_len, d_inner] # SSM 路：conv1d → SiLU → SSM # conv1d 需要 [batch, channels, seq_len] 格式 ssm_branch = ssm_branch.transpose(1, 2) # [batch, d_inner, seq_len] ssm_branch = self.conv1d(ssm_branch)[..., :x.shape[1]] # causal: 截掉多余的 padding ssm_branch = ssm_branch.transpose(1, 2) # [batch, seq_len, d_inner] ssm_branch = F.silu(ssm_branch) ssm_branch = self.ssm(ssm_branch) # [batch, seq_len, d_inner] # 门控路：SiLU gate_branch = F.silu(gate_branch) # 门控融合 output = ssm_branch * gate_branch # [batch, seq_len, d_inner] # 输出投影 + 残差 output = self.output_proj(output) # [batch, seq_len, d_model] return output + residual # ── 验证 ────────────────────────────────────────────────────────────────────── if __name__ == \u0026#34;__main__\u0026#34;: torch.manual_seed(42) batch_size, seq_len, d_model = 2, 16, 64 block = MambaBlock(d_model=d_model, d_inner=128, d_state=16, conv_kernel=4) x = torch.randn(batch_size, seq_len, d_model) y = block(x) print(f\u0026#34;输入形状: {x.shape}\u0026#34;) # [2, 16, 64] print(f\u0026#34;输出形状: {y.shape}\u0026#34;) # [2, 16, 64] print(f\u0026#34;无 NaN: {not torch.isnan(y).any().item()}\u0026#34;) # 验证推理时的递推（单步生成） ssm = SelectiveSSM(d_model=64, d_state=16) single_token = torch.randn(1, 1, 64) single_out = ssm(single_token) print(f\u0026#34;单步推理: {single_out.shape}\u0026#34;) # [1, 1, 64] 九、训练时的并行化：Parallel Scan 递推形式在训练时是串行的，效率低。Mamba 的另一个贡献：当 $\\bar{A}$ 不依赖输入时（S4），可以用卷积并行计算。\n对于固定的 $\\bar{A}$，展开递推：\n$$h_t = \\bar{A}^t h_0 + \\sum_{i=0}^{t} \\bar{A}^{t-i} \\bar{B} x_i$$\n$$y_t = C h_t = \\sum_{i=0}^{t} \\underbrace{C \\bar{A}^{t-i} \\bar{B}}{\\text{卷积核 } K{t-i}} x_i$$\n这就是一个因果卷积！可以用 FFT 在 $O(N \\log N)$ 时间内并行计算。\nMamba 的选择性机制打破了这个并行化（因为 $\\bar{A}_t$ 依赖输入，不再固定），所以 Mamba 使用了 Parallel Scan（并行前缀扫描） 算法。\nParallel Scan 的核心思想 SSM 的递推 $h_t = \\bar{A}t h{t-1} + \\bar{B}_t x_t$ 是一个结合性（associative） 操作，可以用分治法并行化：\n把相邻两步的递推合并成一步：\n$$\\begin{pmatrix} h_t \\ 1 \\end{pmatrix} = \\begin{pmatrix} \\bar{A}_t \u0026amp; \\bar{B}t x_t \\ 0 \u0026amp; 1 \\end{pmatrix} \\begin{pmatrix} h{t-1} \\ 1 \\end{pmatrix}$$\n令 $e_t = (\\bar{A}_t,\\ \\bar{B}_t x_t)$ 为第 $t$ 步的\u0026quot;元素\u0026quot;，定义结合操作 $\\oplus$：\n$$(a_2, b_2) \\oplus (a_1, b_1) = (a_2 \\cdot a_1,\\ a_2 \\cdot b_1 + b_2)$$\n则 $h_t$ 等价于对 $e_1, e_2, \\ldots, e_t$ 做前缀扫描（prefix scan）。前缀扫描可以用 $O(\\log N)$ 轮并行归约完成，总复杂度 $O(N \\log N)$，在 GPU 上每轮内部完全并行。\ndef parallel_scan_ssm(delta_A: torch.Tensor, delta_B_x: torch.Tensor) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; 用 Parallel Scan（并行前缀扫描）计算 SSM 的所有隐状态。 核心思想：把 h_t = A_t * h_{t-1} + b_t 的串行递推， 转化为对 (A_t, b_t) 对的结合性前缀扫描，实现并行计算。 delta_A: [batch, seq_len, d_model, d_state]，离散化后的 Ā_t delta_B_x: [batch, seq_len, d_model, d_state]，B̄_t * x_t（已乘以输入） 返回: [batch, seq_len, d_model, d_state]，所有时刻的隐状态 h_t \u0026#34;\u0026#34;\u0026#34; batch_size, seq_len, d_model, d_state = delta_A.shape # 每个时刻的\u0026#34;元素\u0026#34;是 (a_t, b_t) 对 # a_t = Ā_t（状态转移系数），b_t = B̄_t * x_t（输入贡献） scan_a = delta_A # [batch, seq_len, d_model, d_state] scan_b = delta_B_x # [batch, seq_len, d_model, d_state] # 结合操作：(a2, b2) ⊕ (a1, b1) = (a2*a1, a2*b1 + b2) # 含义：先经历 (a1, b1) 的转移，再经历 (a2, b2) 的转移 # 前缀扫描：用 log2(seq_len) 轮 up-sweep 完成 # 每轮将步长翻倍，并行合并相邻元素对 # 为了在纯 PyTorch 中演示，这里实现 Blelloch 并行前缀扫描 # 实际 Mamba 用 CUDA kernel 实现，效率更高 num_rounds = int(math.ceil(math.log2(seq_len))) if seq_len \u0026gt; 1 else 0 # 用列表存储每轮的中间结果（实际 CUDA 实现在原地操作） current_a = scan_a.clone() current_b = scan_b.clone() # Up-sweep（归约阶段）：步长从 1 倍增到 seq_len/2 stride = 1 for _ in range(num_rounds): # 找到需要合并的位置对：(i - stride, i)，i 从 stride 开始，步长 2*stride left_indices = torch.arange(0, seq_len - stride, 2 * stride, device=delta_A.device) right_indices = left_indices + stride if right_indices.numel() == 0: break left_a = current_a[:, left_indices] # [batch, n_pairs, d_model, d_state] left_b = current_b[:, left_indices] right_a = current_a[:, right_indices] right_b = current_b[:, right_indices] # 结合操作：right ⊕ left merged_a = right_a * left_a merged_b = right_a * left_b + right_b current_a[:, right_indices] = merged_a current_b[:, right_indices] = merged_b stride *= 2 # 注意：上面的 up-sweep 只得到了部分前缀结果（类似 Blelloch scan 的归约树） # 完整的 Blelloch scan 还需要 down-sweep 阶段。 # 在实际 Mamba 实现中，使用专门的 CUDA kernel（mamba_ssm 库中的 selective_scan_cuda） # 直接在 GPU 上高效完成，避免了 Python 层的循环开销。 # 这里为了说明原理，退回到串行递推作为等价的正确实现： hidden_states = [] h = torch.zeros(batch_size, d_model, d_state, device=delta_A.device, dtype=delta_A.dtype) for t in range(seq_len): h = delta_A[:, t] * h + delta_B_x[:, t] hidden_states.append(h) return torch.stack(hidden_states, dim=1) # [batch, seq_len, d_model, d_state] # 验证 parallel_scan_ssm 和 SelectiveSSM 的递推结果一致 def verify_parallel_scan(): torch.manual_seed(0) batch_size, seq_len, d_model, d_state = 2, 8, 16, 4 delta_A = torch.rand(batch_size, seq_len, d_model, d_state) * 0.9 + 0.05 # (0.05, 0.95) delta_B_x = torch.randn(batch_size, seq_len, d_model, d_state) * 0.1 h_scan = parallel_scan_ssm(delta_A, delta_B_x) # 串行递推作为参考 h_ref = torch.zeros(batch_size, d_model, d_state) h_ref_list = [] for t in range(seq_len): h_ref = delta_A[:, t] * h_ref + delta_B_x[:, t] h_ref_list.append(h_ref.clone()) h_ref_stack = torch.stack(h_ref_list, dim=1) max_diff = (h_scan - h_ref_stack).abs().max().item() print(f\u0026#34;Parallel Scan 与串行递推最大误差: {max_diff:.2e}\u0026#34;) # 应接近 0 assert max_diff \u0026lt; 1e-5, f\u0026#34;结果不一致！误差 {max_diff}\u0026#34; print(\u0026#34;验证通过 ✓\u0026#34;) verify_parallel_scan() 十、Mamba 和 Transformer 的对比 Transformer Mamba 核心操作 Softmax Attention 选择性 SSM 训练复杂度 $O(N^2 d)$ $O(N d)$（Parallel Scan） 推理复杂度（逐步） $O(Nd)$（有 KV Cache） $O(d^2)$（固定隐状态） 推理显存 $O(Nd)$（随序列增长） $O(d \\cdot d_{state})$（固定） 长序列能力 受限于 $O(N^2)$ 天然支持超长序列 内容感知 ✅ Attention 天然内容感知 ✅ 选择性机制实现内容感知 精确历史访问 ✅ 能精确 attend 任意历史 token ❌ 历史被压缩进固定隐状态 实现复杂度 简单 纯 PyTorch 可运行；达到生产级速度需要 CUDA kernel（可直接用 mamba-ssm 库） 十一、回到 Qwen3.6：config.json 里的参数对应 现在你能完全读懂 Qwen3.6 的 linear_attn 参数了：\n\u0026#34;linear_attn.A_log\u0026#34; → SSM 的 A 矩阵（对数形式存储，保证稳定性） \u0026#34;linear_attn.conv1d\u0026#34; → 局部卷积（kernel_size=4，捕获短程特征） \u0026#34;linear_attn.dt_bias\u0026#34; → Δ（步长）的偏置项，控制离散化粒度 \u0026#34;linear_attn.in_proj_a\u0026#34; → SSM 的 A/dt 相关投影（用于生成 Δ，控制离散化步长） \u0026#34;linear_attn.in_proj_b\u0026#34; → SSM 的 B 矩阵投影（输入 x 到隐状态 h 的映射） \u0026#34;linear_attn.in_proj_ba\u0026#34; → B 和 A/dt 的联合权重（合并存储提高访存效率，推理时拆分使用） \u0026#34;linear_attn.norm\u0026#34; → SSM 内部的归一化层 \u0026#34;linear_conv_kernel_dim\u0026#34;: 4 → conv1d 的 kernel_size \u0026#34;linear_key_head_dim\u0026#34;: 128 → 隐状态维度（d_state 的变体） \u0026#34;linear_num_key_heads\u0026#34;: 16 → 多头 SSM 的头数 \u0026#34;linear_num_value_heads\u0026#34;: 32 → 输出头数（可以和 key 头数不同） Qwen3.6 的 linear_attn 是多头 SSM（Multi-head SSM），每个头独立维护一个隐状态，类似 MHA 里每个头独立做 Attention。\n十二、核心要点速查 问题 答案 Mamba 解决什么问题？ Linear Attention 的\u0026quot;遗忘\u0026quot;问题：让衰减因子随输入动态变化 SSM 的递推公式？ $h_t = \\bar{A}t h{t-1} + \\bar{B}_t x_t$，$y_t = C_t h_t$ $\\Delta$（dt）的作用？ 控制\u0026quot;记住\u0026quot;还是\u0026quot;忘记\u0026quot;：大 $\\Delta$ 重置历史，小 $\\Delta$ 跳过当前 A_log 为什么用对数？ 保证 $A \u0026lt; 0$，离散化后 $\\bar{A} = e^{\\Delta A} \\in (0,1)$，系统稳定 conv1d 的作用？ 捕获局部短程特征，补充 SSM 的全局长程建模 训练时如何并行？ Parallel Scan（并行前缀扫描），$O(N \\log N)$ 推理时的优势？ 隐状态固定大小 $O(d \\cdot d_{state})$，不随序列增长 和 Transformer 最大的差距？ 无法精确访问历史 token，只能访问被压缩的隐状态摘要 Qwen3.6 里怎么用的？ 多头 SSM，每 4 层混合一个 Full Attention 补偿精度损失 ","permalink":"https://afan.ml/posts/mamba--s-s-m/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：Mamba 是 Linear Attention 的\u0026quot;升级版\u0026quot;——同样是 $O(N)$ 复杂度、固定大小隐状态，但通过\u003cstrong\u003e选择性机制（Selective SSM）\u003c/strong\u003e 让模型能动态决定\u0026quot;记住什么、忘记什么\u0026quot;，效果接近 Transformer，推理速度接近 RNN。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一从-linear-attention-到-mamba解决什么问题\"\u003e一、从 Linear Attention 到 Mamba：解决什么问题？\u003c/h2\u003e\n\u003cp\u003e回顾 Linear Attention 的递推公式：\u003c/p\u003e\n\u003cp\u003e$$S_t = S_{t-1} + \\phi(k_t) v_t^T$$\u003c/p\u003e\n\u003cp\u003e\u003cstrong\u003e问题\u003c/strong\u003e：衰减是固定的（没有衰减，或者 RetNet 里用固定的 $\\gamma$），模型无法根据输入内容动态决定\u0026quot;这个历史信息重不重要\u0026quot;。\u003c/p\u003e\n\u003cp\u003e\u003cstrong\u003e类比\u003c/strong\u003e：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003eRNN（LSTM）：有遗忘门，可以选择性地清除历史\nLinear Attention：没有遗忘门，历史信息只增不减，全部堆进 S 矩阵\nRetNet：有固定衰减 γ，但 γ 是超参数，不随输入变化\nMamba：衰减因子由输入动态生成，每个 token 的\u0026#34;遗忘程度\u0026#34;不同\n\u003c/code\u003e\u003c/pre\u003e\u003chr\u003e\n\u003ch2 id=\"二ssm-的数学基础\"\u003e二、SSM 的数学基础\u003c/h2\u003e\n\u003cp\u003eMamba 基于\u003cstrong\u003e状态空间模型（State Space Model，SSM）\u003c/strong\u003e，这是控制论里的经典框架。\u003c/p\u003e\n\u003ch3 id=\"连续时间-ssm\"\u003e连续时间 SSM\u003c/h3\u003e\n\u003cp\u003e$$h\u0026rsquo;(t) = A h(t) + B x(t)$$\n$$y(t) = C h(t)$$\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e$x(t)$：输入信号\u003c/li\u003e\n\u003cli\u003e$h(t)$：隐状态（类比 RNN 的 hidden state）\u003c/li\u003e\n\u003cli\u003e$y(t)$：输出\u003c/li\u003e\n\u003cli\u003e$A$：状态转移矩阵（控制历史信息如何演化）\u003c/li\u003e\n\u003cli\u003e$B$：输入投影矩阵（控制输入如何影响隐状态）\u003c/li\u003e\n\u003cli\u003e$C$：输出投影矩阵（控制隐状态如何映射到输出）\u003c/li\u003e\n\u003c/ul\u003e\n\u003ch3 id=\"离散化实际使用的形式\"\u003e离散化（实际使用的形式）\u003c/h3\u003e\n\u003cp\u003e连续 SSM 需要离散化才能用于序列建模，使用\u003cstrong\u003e零阶保持（ZOH）\u003c/strong\u003e 方法：\u003c/p\u003e","title":"Mamba / SSM（状态空间模型）笔记"},{"content":" 一句话：MoE 是对 Transformer 中 FFN 层的改进——把一个 FFN 替换成多个并行的\u0026quot;专家\u0026quot;FFN，每个 token 只激活其中少数几个专家，从而在不增加计算量的前提下大幅扩展模型参数量。\n一、问题背景：Dense 模型扩参数的代价 标准 Transformer（Dense 模型）每个 token 都要经过所有参数：\n参数量翻倍 → 计算量翻倍 → 训练/推理成本翻倍 想要更强的模型，就必须付出更高的计算代价。有没有办法让参数量和计算量解耦？\n二、核心思想：条件计算（Conditional Computation） 不是每个 token 都需要用到所有知识，让不同的 token 走不同的\u0026quot;专家\u0026quot;路径。\nMoE 的核心设计：\n把 FFN 层替换成 N 个并行的专家 FFN 每个 token 由一个路由器（Router）决定激活哪 K 个专家 只有被选中的专家才参与计算，其余专家跳过 结果：\n总参数量 = N 个专家的参数之和（很大） 每次计算量 = 只激活 K 个专家（很小） 参数量和计算量不再线性绑定 三、结构对比 Dense FFN（原始） 输入 x（形状：[seq_len, d_model]） ↓ FFN（W_gate, W_up, W_down） ↓ 输出（形状：[seq_len, d_model]） MoE FFN（替换后） 输入 x（形状：[seq_len, d_model]） ↓ Router（路由器）：为每个 token 计算各专家的得分，选出 Top-K ↓ 并行调用被选中的 K 个专家 FFN ↓ 按路由权重加权求和，合并输出 ↓ 输出（形状：[seq_len, d_model]）← 形状和 Dense FFN 完全一致 注意：Attention 层不变，只有 FFN 层被替换成 MoE。\n四、路由器（Router）的工作原理 Router 是一个简单的线性层，输出每个专家的\u0026quot;得分\u0026quot;：\n$$\\text{scores} = \\text{softmax}(x \\cdot W_{\\text{router}})$$\n其中 $W_{\\text{router}}$ 形状为 $[d_model, N_experts]$，输出 N 个专家的概率分布。\n然后取 Top-K 个得分最高的专家：\nimport torch import torch.nn as nn import torch.nn.functional as F # 构造最小示例：4 个 token，d_model=8，共 4 个专家，Top-2 torch.manual_seed(0) N_tokens, d_model, num_experts, top_k = 4, 8, 4, 2 x = torch.randn(N_tokens, d_model) # 输入 token 向量 W_router = torch.randn(d_model, num_experts) # Router 权重矩阵 experts = [nn.Linear(d_model, d_model, bias=False) for _ in range(num_experts)] # 第一步：Router 计算每个 token 对各专家的得分 scores = F.softmax(x @ W_router, dim=-1) # [4, 4] # 第二步：选 Top-2 专家 topk_scores, topk_indices = scores.topk(k=top_k, dim=-1) # 各形状 [4, 2] # 第三步：对 Top-K 权重重新归一化，使每个 token 的权重之和为 1 topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True) # 第四步：对每个 token，加权求和被选中专家的输出 output = torch.zeros(N_tokens, d_model) for k_pos in range(top_k): expert_ids = topk_indices[:, k_pos] # 每个 token 在第 k_pos 位选的专家编号，形状 [4] weights = topk_scores[:, k_pos] # 对应权重，形状 [4] for token_i in range(N_tokens): expert_id = expert_ids[token_i].item() output[token_i] += weights[token_i] * experts[expert_id](x[token_i]) print(output.shape) # torch.Size([4, 8]) print(topk_indices) # 每个 token 选中的专家编号，如 tensor([[2,1],[0,3],[1,2],[3,0]]) 五、完整代码示例（PyTorch） import torch import torch.nn as nn import torch.nn.functional as F class ExpertFFN(nn.Module): \u0026#34;\u0026#34;\u0026#34;单个专家：标准的两层 FFN\u0026#34;\u0026#34;\u0026#34; def __init__(self, d_model, d_ffn): super().__init__() self.w1 = nn.Linear(d_model, d_ffn, bias=False) self.w2 = nn.Linear(d_ffn, d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x))) class MoELayer(nn.Module): \u0026#34;\u0026#34;\u0026#34;MoE 层：N 个专家 + Router，每个 token 激活 Top-K 个专家\u0026#34;\u0026#34;\u0026#34; def __init__(self, d_model, d_ffn, num_experts=8, top_k=2): super().__init__() self.num_experts = num_experts self.top_k = top_k # N 个独立的专家 FFN self.experts = nn.ModuleList([ ExpertFFN(d_model, d_ffn) for _ in range(num_experts) ]) # Router：线性层，输出每个专家的得分 self.router = nn.Linear(d_model, num_experts, bias=False) def forward(self, x): batch_size, seq_len, d_model = x.shape x_flat = x.view(-1, d_model) # [batch*seq_len, d_model] # 路由：计算每个 token 对应各专家的得分 router_logits = self.router(x_flat) # [N_tokens, N_experts] router_scores = F.softmax(router_logits, dim=-1) # 归一化为概率 # 选 Top-K 个专家 topk_scores, topk_indices = router_scores.topk(self.top_k, dim=-1) # 对 Top-K 的权重重新归一化（使权重之和为 1） topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True) # 对每个 token，加权求和被选中专家的输出 output = torch.zeros_like(x_flat) for expert_idx in range(self.num_experts): # 找出哪些 token 选了这个专家 token_mask = (topk_indices == expert_idx).any(dim=-1) if not token_mask.any(): continue expert_input = x_flat[token_mask] expert_output = self.experts[expert_idx](expert_input) # 取该专家对应的权重 expert_weight = topk_scores[token_mask][ (topk_indices[token_mask] == expert_idx) ].unsqueeze(-1) output[token_mask] += expert_weight * expert_output return output.view(batch_size, seq_len, d_model) # 验证 d_model, d_ffn = 512, 2048 moe = MoELayer(d_model=d_model, d_ffn=d_ffn, num_experts=8, top_k=2) x = torch.randn(2, 16, d_model) # batch=2, seq_len=16 out = moe(x) print(out.shape) # torch.Size([2, 16, 512]) 六、负载均衡问题：专家塌陷 MoE 有一个严重的训练问题：专家塌陷（Expert Collapse）。\nRouter 一旦发现某个专家效果好，就会一直选它 → 这个专家越训越强，其他专家越来越少被选 → 最终退化成只有 1～2 个专家在工作，其余专家形同虚设 解决方案1：辅助损失（Auxiliary Loss） 在训练损失里加一项，惩罚负载不均衡：\n$$\\mathcal{L}{\\text{aux}} = \\alpha \\cdot N \\cdot \\sum{i=1}^{N} f_i \\cdot p_i$$\n其中 $f_i$ 是专家 $i$ 被选中的频率，$p_i$ 是 Router 给专家 $i$ 的平均概率。鼓励每个专家被均匀使用。\n解决方案2：无辅助损失的负载均衡（DeepSeek 方案） DeepSeek-V3 提出了无辅助损失的方案：给每个专家加一个可学习的偏置项，动态调整各专家被选中的概率，不需要额外的损失函数。\n七、DeepSeek-V3 的 MoE 设计 DeepSeek-V3（671B）是目前最具代表性的 MoE 大模型：\n总专家数：256 个 每个 token 激活专家数：8 个（Top-8） 共享专家：1 个（每个 token 必选，保证基础能力） 总参数：671B 每次推理激活参数：~37B（约 5.5% 的参数参与计算） 共享专家是 DeepSeek 的创新：除了路由选出的 8 个专家，还有 1 个专家是所有 token 都必须经过的，用来保留通用知识，避免专家过度特化。\n每个 token 的计算路径： 共享专家（必选 1 个）+ 路由专家（Top-8）= 共 9 个专家的输出加权求和 八、MoE 的优缺点 优点 优点 说明 参数量和计算量解耦 671B 参数但只激活 37B，推理速度接近 37B 模型 专家自然分工 不同专家会自发学习不同领域的知识 扩展性强 增加专家数量几乎不增加推理计算量 缺点 缺点 说明 显存需求高 所有专家的参数都要加载进显存，即使大多数不参与计算 训练不稳定 专家塌陷问题需要额外处理 通信开销大 分布式训练时，不同 token 路由到不同设备的专家，需要大量跨设备通信 推理延迟不均 不同 batch 里激活的专家不同，难以做静态优化 九、核心要点速查 问题 答案 MoE 改的是哪一层？ FFN 层，Attention 层不变 专家是什么？ 独立的 FFN 子网络，结构和普通 FFN 一样 Router 做什么？ 为每个 token 计算各专家得分，选 Top-K 为什么参数多但计算少？ 每次只激活 K 个专家，其余专家跳过 最大的训练问题是什么？ 专家塌陷，需要负载均衡策略 DeepSeek-V3 的规模？ 671B 总参数，每次激活约 37B（Top-8 + 1 共享专家） 和 Dense 模型的本质区别？ Dense：所有参数都激活；MoE：条件激活，参数量≠计算量 ","permalink":"https://afan.ml/posts/mo-e%E6%B7%B7%E5%90%88%E4%B8%93%E5%AE%B6%E6%A8%A1%E5%9E%8B/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：MoE 是对 Transformer 中 FFN 层的改进——把一个 FFN 替换成多个并行的\u0026quot;专家\u0026quot;FFN，每个 token 只激活其中少数几个专家，从而在\u003cstrong\u003e不增加计算量\u003c/strong\u003e的前提下大幅扩展模型参数量。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一问题背景dense-模型扩参数的代价\"\u003e一、问题背景：Dense 模型扩参数的代价\u003c/h2\u003e\n\u003cp\u003e标准 Transformer（Dense 模型）每个 token 都要经过所有参数：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e参数量翻倍 → 计算量翻倍 → 训练/推理成本翻倍\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e想要更强的模型，就必须付出更高的计算代价。有没有办法\u003cstrong\u003e让参数量和计算量解耦\u003c/strong\u003e？\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二核心思想条件计算conditional-computation\"\u003e二、核心思想：条件计算（Conditional Computation）\u003c/h2\u003e\n\u003cblockquote\u003e\n\u003cp\u003e不是每个 token 都需要用到所有知识，让不同的 token 走不同的\u0026quot;专家\u0026quot;路径。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003cp\u003eMoE 的核心设计：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e把 FFN 层替换成 N 个并行的专家 FFN\n每个 token 由一个路由器（Router）决定激活哪 K 个专家\n只有被选中的专家才参与计算，其余专家跳过\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e结果\u003c/strong\u003e：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e总参数量 = N 个专家的参数之和（很大）\u003c/li\u003e\n\u003cli\u003e每次计算量 = 只激活 K 个专家（很小）\u003c/li\u003e\n\u003cli\u003e参数量和计算量\u003cstrong\u003e不再线性绑定\u003c/strong\u003e\u003c/li\u003e\n\u003c/ul\u003e\n\u003chr\u003e\n\u003ch2 id=\"三结构对比\"\u003e三、结构对比\u003c/h2\u003e\n\u003ch3 id=\"dense-ffn原始\"\u003eDense FFN（原始）\u003c/h3\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e输入 x（形状：[seq_len, d_model]）\n ↓\nFFN（W_gate, W_up, W_down）\n ↓\n输出（形状：[seq_len, d_model]）\n\u003c/code\u003e\u003c/pre\u003e\u003ch3 id=\"moe-ffn替换后\"\u003eMoE FFN（替换后）\u003c/h3\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e输入 x（形状：[seq_len, d_model]）\n ↓\nRouter（路由器）：为每个 token 计算各专家的得分，选出 Top-K\n ↓\n并行调用被选中的 K 个专家 FFN\n ↓\n按路由权重加权求和，合并输出\n ↓\n输出（形状：[seq_len, d_model]）← 形状和 Dense FFN 完全一致\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e注意\u003c/strong\u003e：Attention 层不变，只有 FFN 层被替换成 MoE。\u003c/p\u003e","title":"MoE（Mixture of Experts）混合专家模型笔记"},{"content":" 前置知识：已了解 Self-Attention（自注意力）。 MHA 的本质：把 Self-Attention 并行做多次，每次关注不同的语义子空间，最后合并结果。\n一、为什么需要多头？Self-Attention 有什么不足？ 单头 Self-Attention 每次只能学到一种\u0026quot;关注模式\u0026quot;。例如：\n句子：\u0026#34;The animal didn\u0026#39;t cross the street because it was too tired\u0026#34; 头1 可能学到：it → animal（指代关系） 头2 可能学到：tired → animal（状态描述） 头3 可能学到：cross → street（动作与地点） 单头只能同时关注一种模式，多头让模型在不同子空间里并行捕捉多种关系。\n二、MHA 整体结构 输入 X │ ├──→ 线性投影 W_Q^1 → Q1 ─┐ ├──→ 线性投影 W_K^1 → K1 ─┤→ Attention(Q1,K1,V1) → head_1 ─┐ ├──→ 线性投影 W_V^1 → V1 ─┘ │ │ │ ├──→ 线性投影 W_Q^2 → Q2 ─┐ │ ├──→ 线性投影 W_K^2 → K2 ─┤→ Attention(Q2,K2,V2) → head_2 ─┤→ Concat → 线性投影 W_O → 输出 ├──→ 线性投影 W_V^2 → V2 ─┘ │ │ │ ├──→ ...（共 h 个头） │ │ │ └──→ 线性投影 W_Q^h → Qh ─┐ │ 线性投影 W_K^h → Kh ─┤→ Attention(Qh,Kh,Vh) → head_h ─┘ 线性投影 W_V^h → Vh ─┘ 三、完整公式 单头 Scaled Dot-Product Attention（回顾） $$ \\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V $$\n多头注意力 $$ \\text{head}_i = \\text{Attention}(XW_Q^i,\\ XW_K^i,\\ XW_V^i) $$\n$$ \\text{MultiHead}(X) = \\text{Concat}(\\text{head}_1, \\text{head}_2, \\ldots, \\text{head}_h) \\cdot W_O $$\n参数说明：\n符号 含义 维度 $X$ 输入序列 $[B, T, d_{model}]$ $W_Q^i, W_K^i, W_V^i$ 第 $i$ 个头的投影矩阵 $[d_{model}, d_k]$ $d_k = d_v = d_{model} / h$ 每个头的维度 标量 $W_O$ 输出投影矩阵 $[h \\cdot d_v, d_{model}]$ $h$ 头的数量 标量（如 8、12、16） 关键设计：每个头的维度 $d_k = d_{model} / h$，所以多头的总计算量和单头相同，不增加额外开销。\n四、逐步拆解计算过程（具体例子） 设定：\n输入序列长度 $T = 3$（3个词） 模型维度 $d_{model} = 4$ 头数 $h = 2$ 每个头的维度 $d_k = d_{model} / h = 4 / 2 = 2$ 输入 X（3个词，每词4维）：\nX = [[1, 0, 1, 0], ← 词1 [0, 1, 0, 1], ← 词2 [1, 1, 0, 0]] ← 词3 第一步：线性投影，生成每个头的 Q、K、V 每个头有独立的投影矩阵 $W_Q^i, W_K^i, W_V^i$，维度为 $[4, 2]$。 以下给出头1和头2的全部投影矩阵，后续计算均基于这些固定数值。\n头1的投影矩阵：\nW_Q^1 = [[1, 0], W_K^1 = [[1, 0], W_V^1 = [[1, 0], [0, 1], [0, 0], [0, 1], [1, 0], [0, 1], [0, 0], [0, 1]] [1, 0]] [1, 0]] 头1的 Q1、K1、V1 计算（X @ W）：\nX = [[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0]] Q1 = X @ W_Q^1: 词1: [1*1+0*0+1*1+0*0, 1*0+0*1+1*0+0*1] = [2, 0] 词2: [0*1+1*0+0*1+1*0, 0*0+1*1+0*0+1*1] = [0, 2] 词3: [1*1+1*0+0*1+0*0, 1*0+1*1+0*0+0*1] = [1, 1] → Q1 = [[2, 0], [0, 2], [1, 1]] K1 = X @ W_K^1: 词1: [1*1+0*0+1*0+0*1, 1*0+0*0+1*1+0*0] = [1, 1] 词2: [0*1+1*0+0*0+1*1, 0*0+1*0+0*1+1*0] = [1, 0] 词3: [1*1+1*0+0*0+0*1, 1*0+1*0+0*1+0*0] = [1, 0] → K1 = [[1, 1], [1, 0], [1, 0]] V1 = X @ W_V^1: 词1: [1*1+0*0+1*0+0*1, 1*0+0*1+1*0+0*0] = [1, 0] 词2: [0*1+1*0+0*0+1*1, 0*0+1*1+0*0+1*0] = [1, 1] 词3: [1*1+1*0+0*0+0*1, 1*0+1*1+0*0+0*0] = [1, 1] → V1 = [[1, 0], [1, 1], [1, 1]] 头2的投影矩阵：\nW_Q^2 = [[0, 1], W_K^2 = [[0, 1], W_V^2 = [[0, 1], [1, 0], [1, 0], [1, 0], [0, 1], [0, 1], [0, 0], [1, 0]] [1, 0]] [0, 1]] 头2的 Q2、K2、V2 计算：\nQ2 = X @ W_Q^2: 词1: [1*0+0*1+1*0+0*1, 1*1+0*0+1*1+0*0] = [0, 2] 词2: [0*0+1*1+0*0+1*1, 0*1+1*0+0*1+1*0] = [2, 0] 词3: [1*0+1*1+0*0+0*1, 1*1+1*0+0*1+0*0] = [1, 1] → Q2 = [[0, 2], [2, 0], [1, 1]] K2 = X @ W_K^2: 词1: [1*0+0*1+1*0+0*1, 1*1+0*0+1*1+0*0] = [0, 2] 词2: [0*0+1*1+0*0+1*1, 0*1+1*0+0*1+1*0] = [2, 0] 词3: [1*0+1*1+0*0+0*1, 1*1+1*0+0*1+0*0] = [1, 1] → K2 = [[0, 2], [2, 0], [1, 1]] V2 = X @ W_V^2: 词1: [1*0+0*1+1*0+0*0, 1*1+0*0+1*0+0*1] = [0, 1] 词2: [0*0+1*1+0*0+1*0, 0*1+1*0+0*0+1*1] = [1, 1] 词3: [1*0+1*1+0*0+0*0, 1*1+1*0+0*0+0*1] = [1, 1] → V2 = [[0, 1], [1, 1], [1, 1]] 第二步：每个头独立做 Scaled Dot-Product Attention 头1的计算：\nscores1 = Q1 @ K1^T / √2 Q1 @ K1^T（先不除√2）： K1^T = [[1, 1, 1], [1, 0, 0]] 词1 [2,0] · K1^T: [2*1+0*1, 2*1+0*0, 2*1+0*0] = [2, 2, 2] 词2 [0,2] · K1^T: [0*1+2*1, 0*1+2*0, 0*1+2*0] = [2, 0, 0] 词3 [1,1] · K1^T: [1*1+1*1, 1*1+1*0, 1*1+1*0] = [2, 1, 1] Q1 @ K1^T = [[2, 2, 2], [2, 0, 0], [2, 1, 1]] 除以 √2 ≈ 1.414： scores1 = [[1.414, 1.414, 1.414], [1.414, 0.000, 0.000], [1.414, 0.707, 0.707]] softmax（对每行独立计算）： 词1行 [1.414, 1.414, 1.414]：三个值相等 → softmax = [0.333, 0.333, 0.333] 词2行 [1.414, 0.000, 0.000]： e^1.414=4.113, e^0=1, e^0=1, 总和=6.113 → [4.113/6.113, 1/6.113, 1/6.113] = [0.673, 0.164, 0.164] 词3行 [1.414, 0.707, 0.707]： e^1.414=4.113, e^0.707=2.028, e^0.707=2.028, 总和=8.169 → [4.113/8.169, 2.028/8.169, 2.028/8.169] = [0.503, 0.248, 0.248] A1 = [[0.333, 0.333, 0.333], [0.673, 0.164, 0.164], [0.503, 0.248, 0.248]] head_1 = A1 @ V1： V1 = [[1, 0], [1, 1], [1, 1]] 词1: [0.333*1+0.333*1+0.333*1, 0.333*0+0.333*1+0.333*1] = [1.000, 0.667] 词2: [0.673*1+0.164*1+0.164*1, 0.673*0+0.164*1+0.164*1] = [1.000, 0.327] 词3: [0.503*1+0.248*1+0.248*1, 0.503*0+0.248*1+0.248*1] = [1.000, 0.497] head_1 = [[1.000, 0.667], [1.000, 0.327], [1.000, 0.497]] 头2的计算：\nscores2 = Q2 @ K2^T / √2 Q2 @ K2^T（先不除√2）： K2^T = [[0, 2, 1], [2, 0, 1]] 词1 [0,2] · K2^T: [0*0+2*2, 0*2+2*0, 0*1+2*1] = [4, 0, 2] 词2 [2,0] · K2^T: [2*0+0*2, 2*2+0*0, 2*1+0*1] = [0, 4, 2] 词3 [1,1] · K2^T: [1*0+1*2, 1*2+1*0, 1*1+1*1] = [2, 2, 2] Q2 @ K2^T = [[4, 0, 2], [0, 4, 2], [2, 2, 2]] 除以 √2 ≈ 1.414： scores2 = [[2.828, 0.000, 1.414], [0.000, 2.828, 1.414], [1.414, 1.414, 1.414]] softmax： 词1行 [2.828, 0.000, 1.414]： e^2.828=16.93, e^0=1, e^1.414=4.113, 总和=22.04 → [16.93/22.04, 1/22.04, 4.113/22.04] = [0.768, 0.045, 0.187] 词2行 [0.000, 2.828, 1.414]： e^0=1, e^2.828=16.93, e^1.414=4.113, 总和=22.04 → [1/22.04, 16.93/22.04, 4.113/22.04] = [0.045, 0.768, 0.187] 词3行 [1.414, 1.414, 1.414]：三个值相等 → [0.333, 0.333, 0.333] A2 = [[0.768, 0.045, 0.187], [0.045, 0.768, 0.187], [0.333, 0.333, 0.333]] head_2 = A2 @ V2： V2 = [[0, 1], [1, 1], [1, 1]] 词1: [0.768*0+0.045*1+0.187*1, 0.768*1+0.045*1+0.187*1] = [0.232, 1.000] 词2: [0.045*0+0.768*1+0.187*1, 0.045*1+0.768*1+0.187*1] = [0.955, 1.000] 词3: [0.333*0+0.333*1+0.333*1, 0.333*1+0.333*1+0.333*1] = [0.667, 1.000] head_2 = [[0.232, 1.000], [0.955, 1.000], [0.667, 1.000]] 观察：头1的注意力矩阵 A1 中词2强烈关注自己（0.673），头2的 A2 中词1强烈关注自己（0.768）、词2强烈关注自己（0.768）——两个头确实学到了不同的关注模式。\n第三步：拼接所有头的输出 head_1 = [[1.000, 0.667], head_2 = [[0.232, 1.000], [1.000, 0.327], [0.955, 1.000], [1.000, 0.497]] [0.667, 1.000]] Concat(head_1, head_2) → 沿最后一维拼接，shape: [3, 4] Concat = [[1.000, 0.667, 0.232, 1.000], ← 词1 [1.000, 0.327, 0.955, 1.000], ← 词2 [1.000, 0.497, 0.667, 1.000]] ← 词3 第四步：输出线性投影 W_O 的维度：[4, 4]（h×d_v → d_model，即 4→4） 取一个具体的 W_O（实际训练中由梯度下降学习得到，这里固定数值演示融合效果）： W_O = [[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0], [0, 0, 1, 1]] Concat = [[1.000, 0.667, 0.232, 1.000], [1.000, 0.327, 0.955, 1.000], [1.000, 0.497, 0.667, 1.000]] 输出 = Concat @ W_O，逐行计算： 词1 [1.000, 0.667, 0.232, 1.000]: 列0: 1.000*1 + 0.667*0 + 0.232*1 + 1.000*0 = 1.232 列1: 1.000*0 + 0.667*1 + 0.232*1 + 1.000*0 = 0.899 列2: 1.000*1 + 0.667*0 + 0.232*0 + 1.000*1 = 2.000 列3: 1.000*0 + 0.667*1 + 0.232*0 + 1.000*1 = 1.667 词2 [1.000, 0.327, 0.955, 1.000]: 列0: 1.000*1 + 0.327*0 + 0.955*1 + 1.000*0 = 1.955 列1: 1.000*0 + 0.327*1 + 0.955*1 + 1.000*0 = 1.282 列2: 1.000*1 + 0.327*0 + 0.955*0 + 1.000*1 = 2.000 列3: 1.000*0 + 0.327*1 + 0.955*0 + 1.000*1 = 1.327 词3 [1.000, 0.497, 0.667, 1.000]: 列0: 1.000*1 + 0.497*0 + 0.667*1 + 1.000*0 = 1.667 列1: 1.000*0 + 0.497*1 + 0.667*1 + 1.000*0 = 1.164 列2: 1.000*1 + 0.497*0 + 0.667*0 + 1.000*1 = 2.000 列3: 1.000*0 + 0.497*1 + 0.667*0 + 1.000*1 = 1.497 输出 = [[1.232, 0.899, 2.000, 1.667], ← 词1 [1.955, 1.282, 2.000, 1.327], ← 词2 [1.667, 1.164, 2.000, 1.497]] ← 词3 输出 shape 为 $[3, 4]$，与输入 X 的 shape $[3, 4]$ 完全一致。 W_O 的作用是将两个头的信息线性混合融合：每个输出维度都综合了来自 head_1 和 head_2 的不同特征，而不是简单地保留各头的独立输出。\n五、维度变化总览 输入 X: [B, T, d_model] 例：[2, 10, 512] ↓ 每个头独立投影 Q_i, K_i, V_i: [B, T, d_k] 例：[2, 10, 64] （d_k = 512/8 = 64） ↓ Scaled Dot-Product Attention head_i: [B, T, d_v] 例：[2, 10, 64] ↓ Concat 所有头 Concat: [B, T, h × d_v] 例：[2, 10, 512] （8×64=512，恢复原维度） ↓ 输出投影 W_O 输出: [B, T, d_model] 例：[2, 10, 512] 六、实际模型中的头数配置 模型 $d_{model}$ 头数 $h$ 每头维度 $d_k$ Transformer（原论文） 512 8 64 BERT-Base 768 12 64 BERT-Large 1024 16 64 GPT-3 12288 96 128 LLaMA-7B 4096 32 128 规律：每头维度 $d_k$ 通常固定在 64 或 128，增大模型主要靠增加头数和 $d_{model}$。\n七、MHA 的工程实现技巧 实际代码中，不会真的创建 h 个独立矩阵，而是用一个大矩阵一次性投影，再 reshape 成多头，效率更高。\nimport torch import torch.nn as nn import torch.nn.functional as F import math class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int): super().__init__() assert d_model % num_heads == 0, \u0026#34;d_model 必须能被 num_heads 整除\u0026#34; self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads # 每个头的维度 # 用一个大矩阵同时投影所有头的 Q、K、V（等价于 h 个独立矩阵） self.W_Q = nn.Linear(d_model, d_model) # [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 split_heads(self, x: torch.Tensor) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34;将 [B, T, d_model] reshape 成 [B, h, T, d_k]\u0026#34;\u0026#34;\u0026#34; B, T, _ = x.shape x = x.view(B, T, self.num_heads, self.d_k) return x.transpose(1, 2) # [B, h, T, d_k] def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -\u0026gt; torch.Tensor: B, T, _ = x.shape # 第一步：线性投影，生成 Q、K、V Q = self.split_heads(self.W_Q(x)) # [B, h, T, d_k] K = self.split_heads(self.W_K(x)) # [B, h, T, d_k] V = self.split_heads(self.W_V(x)) # [B, h, T, d_k] # 第二步：Scaled Dot-Product Attention scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # [B, h, T, T] if mask is not None: scores = scores.masked_fill(mask == 0, float(\u0026#39;-inf\u0026#39;)) attention_weights = F.softmax(scores, dim=-1) # [B, h, T, T] attended = torch.matmul(attention_weights, V) # [B, h, T, d_k] # 第三步：合并多头 attended = attended.transpose(1, 2).contiguous() # [B, T, h, d_k] attended = attended.view(B, T, self.d_model) # [B, T, d_model] # 第四步：输出投影 output = self.W_O(attended) # [B, T, d_model] return output # 使用示例 model = MultiHeadAttention(d_model=512, num_heads=8) x = torch.randn(2, 10, 512) # batch=2, 序列长度=10, 维度=512 output = model(x) print(output.shape) # torch.Size([2, 10, 512]) 八、MHA 的变体详解 变体演进的核心驱动力：推理时 KV Cache 占用显存太大、读写太慢。 每生成一个 token，都要把所有历史 token 的 K、V 从显存读出来做 Attention，这是推理速度的主要瓶颈。\n8.1 MQA（Multi-Query Attention，2019） 核心改动：所有头共享同一组 K 和 V，只有 Q 保持多头。\nMHA：Q1 K1 V1 | Q2 K2 V2 | Q3 K3 V3 | Q4 K4 V4 ← 每头独立 K/V MQA：Q1 | Q2 | Q3 | Q4 └──────────── K V ────────────┘ ← 所有头共享一组 K/V KV Cache 变化：\nMHA：每层缓存 2 × h × d_k 个值（h 个头各自的 K 和 V） MQA：每层只缓存 2 × d_k 个值（只有 1 组 K/V） 节省比例：h 倍（如 h=32，节省 32 倍 KV Cache） 代价：K/V 表达能力下降，多个头被迫用同一个 K/V 做 Attention，模型性能有所损失。\n代表模型：PaLM、Falcon、早期 Gemini\n8.2 GQA（Grouped-Query Attention，2023） 核心改动：把 h 个头分成 g 组，同一组内的头共享一组 K/V。MHA 是 g=h 的特例，MQA 是 g=1 的特例。\nMHA（g=h=4）：Q1 K1 V1 | Q2 K2 V2 | Q3 K3 V3 | Q4 K4 V4 GQA（g=2）： Q1 Q2 | Q3 Q4 └─ K1 V1 ─┘ └─ K2 V2 ─┘ ← 每组共享一对 K/V MQA（g=1）： Q1 Q2 Q3 Q4 └────── K V ──────┘ KV Cache 变化：\n每层缓存 2 × g × d_k 个值 g 越小，节省越多；g=h 退化为 MHA，g=1 退化为 MQA 优势：在 KV Cache 节省和模型性能之间取得平衡，实践中 g 取 h/4 或 h/8 效果接近 MHA。\n代表模型：LLaMA-2/3、Mistral、Qwen2、Gemma2\n8.3 MLA（Multi-head Latent Attention，2024，DeepSeek 提出） 核心改动：不直接缓存 K/V，而是缓存一个低秩压缩向量 c，推理时从 c 还原出 K/V。\n原理：\n标准做法： X → W_K → K（维度 d_k × h） 直接缓存 K MLA 做法： X → W_下投影 → c（维度 d_c，d_c \u0026lt;\u0026lt; d_k × h）→ W_上投影 → K ↑ 只缓存这个低维向量 c KV Cache 变化：\n标准 MHA：每 token 每层缓存 2 × h × d_k 维 MLA：每 token 每层只缓存 d_c 维（DeepSeek-V2 中 d_c = 512，而 2×h×d_k = 2×128×128 = 32768） 实际节省约 93.3% 的 KV Cache 额外优势：MLA 的表达能力经论文（TransMLA）证明强于 GQA，因为低秩压缩保留了跨头的信息交互，而 GQA 的共享 K/V 是对头的硬性绑定。\n代价：推理时需要额外的上投影计算，但因为显存读写量大幅减少，整体推理速度仍然更快。\n代表模型：DeepSeek-V2、DeepSeek-V3、DeepSeek-R1\n8.4 Flash Attention（2022，工程优化，非结构变体） 重要区分：Flash Attention 不改变 Attention 的数学结果，它是一种 GPU 内存访问的工程优化，和 MHA/GQA/MLA 是正交的概念——可以叠加使用。\n问题背景：标准 Attention 计算时，注意力矩阵 scores（shape: [B, h, T, T]）需要完整写入 GPU HBM（高带宽内存），当序列长度 T=4096 时，这个矩阵有 4096×4096 = 1600 万个元素，读写极慢。\nFlash Attention 的解法：分块（Tiling）计算，把 Q/K/V 分成小块，在 GPU SRAM（片上缓存，速度比 HBM 快 10 倍以上）中完成计算，永远不把完整的注意力矩阵写回 HBM。\n标准 Attention 内存访问路径： HBM → 读 Q,K → SRAM 计算 scores → HBM 写 scores HBM → 读 scores → SRAM softmax → HBM 写 softmax(scores) HBM → 读 softmax(scores), V → SRAM 计算输出 → HBM 写输出 （3次大规模 HBM 读写） Flash Attention 内存访问路径： 分块循环：HBM → 读一小块 Q,K,V → SRAM 内完成所有计算 → HBM 写输出 （只有 1 次 HBM 写，中间结果全在 SRAM） 效果：\n速度提升 2~4 倍（Flash Attention 2），Flash Attention 3 针对 H100 进一步优化 显存占用从 O(T²) 降为 O(T)（不存完整注意力矩阵） 支持更长的序列长度 版本演进：\nFlash Attention 1（2022）：提出分块计算思路 Flash Attention 2（2023）：优化并行策略，减少非矩阵乘法运算，速度再提升约 2 倍 Flash Attention 3（2024）：针对 Hopper 架构（H100）优化，利用异步流水线，速度达到 H100 理论峰值的 75% 使用现状：几乎所有现代开源大模型训练和推理都默认启用 Flash Attention 2，它和 GQA/MLA 可以同时使用。\n8.5 各变体对比总览 变体 Q 头数 K/V 头数 KV Cache 大小 表达能力 代表模型 MHA h h 2×h×d_k ⭐⭐⭐⭐⭐ BERT、GPT-2、GPT-3 MQA h 1 2×d_k ⭐⭐⭐ PaLM、Falcon GQA h g（1\u0026lt;g\u0026lt;h） 2×g×d_k ⭐⭐⭐⭐ LLaMA-2/3、Mistral、Qwen2 MLA h 低秩压缩 c d_c（远小于 2×h×d_k） ⭐⭐⭐⭐⭐ DeepSeek-V2/V3/R1 Flash Attention — — 不变（工程优化） 不变 几乎所有现代模型 2025 年主流开源模型的选择：\nMeta LLaMA-3：GQA + Flash Attention 2 Mistral / Mixtral：GQA + Flash Attention 2 Qwen2 / Qwen2.5：GQA + Flash Attention 2 DeepSeek-V3 / R1：MLA + Flash Attention 2（MLA 是目前最先进的注意力机制） Google Gemma2：GQA + Flash Attention 2 结论：Flash Attention 2 几乎是所有先进开源模型的标配（工程层面），注意力结构层面 GQA 是主流，DeepSeek 的 MLA 是目前表达能力最强、KV Cache 最省的方案，代表最新方向。\n九、多层 Transformer 堆叠（Multi-Layer） 多头 vs 多层，两个不同维度 多头（Multi-Head） 多层（Multi-Layer） 方向 同一层内并行 层与层之间串行 解决的问题 同时捕捉多种语义关系 逐层抽象，从低层特征到高层语义 数据流 输入 → 分成 h 份 → 各自 Attention → Concat → 输出 第1层输出 → 第2层输入 → … → 第N层输出 单个 Transformer Layer 的完整结构 每一层不只有 MHA，还包含 FFN 和残差连接：\n输入 X │ ├─→ LayerNorm → MHA → + ←── 残差（直接加上输入 X） │ │ │ ↓ └─────────────────→ LayerNorm → FFN → + ←── 残差 │ ↓ 该层输出（送入下一层） 残差连接：防止梯度消失，让深层网络可训练（GPT-3 有 96 层，没有残差根本训不动） FFN：两层线性变换 + 激活函数，维度先升后降（通常 d_model → 4×d_model → d_model），负责非线性特征变换 LayerNorm：稳定训练，现代模型多用 Pre-Norm（在 MHA/FFN 之前做归一化） 多层数据流示意 输入 tokens: [\u0026#34;今\u0026#34;, \u0026#34;天\u0026#34;, \u0026#34;天\u0026#34;, \u0026#34;气\u0026#34;, \u0026#34;好\u0026#34;] ↓ Embedding + Positional Encoding X: [B, T, d_model] ━━━━━━━━━━━━━━ 第1层（有自己独立的 W_Q/W_K/W_V/W_O/FFN 参数）━━━━━━━━━━━━━━ MHA（h 个头并行）→ Concat → W_O → 残差 → LayerNorm FFN → 残差 → LayerNorm 输出 H1: [B, T, d_model] ← shape 不变，内容变了（低层特征：词法、局部语法） ━━━━━━━━━━━━━━ 第2层（参数与第1层完全独立，不共享）━━━━━━━━━━━━━━ 输入是 H1，不是原始 X 输出 H2: [B, T, d_model] ← 中层特征：句法结构、指代关系 ━━━━━━━━━━━━━━ 第N层 ━━━━━━━━━━━━━━ 输出 HN: [B, T, d_model] ← 高层特征：语义、推理、抽象概念 关键：每层 shape 完全相同，但数值不同——这就是多层的本质，维度不变，语义逐层深化。\n实际模型的层数配置 模型 层数 头数 d_model BERT-Base 12 12 768 GPT-2（小） 12 12 768 LLaMA-3-8B 32 32 4096 LLaMA-3-70B 80 64 8192 GPT-3 96 96 12288 完整代码：多头 + 多层堆叠 import torch import torch.nn as nn import torch.nn.functional as F import math class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads 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) self.dropout = nn.Dropout(dropout) def split_heads(self, x: torch.Tensor) -\u0026gt; torch.Tensor: B, T, _ = x.shape return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2) def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -\u0026gt; torch.Tensor: B, T, _ = x.shape Q = self.split_heads(self.W_Q(x)) K = self.split_heads(self.W_K(x)) V = self.split_heads(self.W_V(x)) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float(\u0026#39;-inf\u0026#39;)) weights = self.dropout(F.softmax(scores, dim=-1)) attended = torch.matmul(weights, V) attended = attended.transpose(1, 2).contiguous().view(B, T, self.d_model) return self.W_O(attended) class FeedForward(nn.Module): def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) self.linear2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: return self.linear2(self.dropout(F.relu(self.linear1(x)))) class TransformerLayer(nn.Module): \u0026#34;\u0026#34;\u0026#34; 单个 Transformer 层 = MHA + FFN + 两个残差连接 + 两个 LayerNorm。 每一层有自己独立的参数，层与层之间不共享。 \u0026#34;\u0026#34;\u0026#34; def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1): super().__init__() self.mha = MultiHeadAttention(d_model, num_heads, dropout) self.ffn = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -\u0026gt; torch.Tensor: # Pre-Norm 风格：先 LayerNorm，再 MHA，再残差 x = x + self.dropout(self.mha(self.norm1(x), mask)) x = x + self.dropout(self.ffn(self.norm2(x))) return x # shape 不变：[B, T, d_model] class TransformerEncoder(nn.Module): \u0026#34;\u0026#34;\u0026#34; 多层 Transformer 堆叠。 num_layers 个 TransformerLayer 串行连接，每层参数独立。 \u0026#34;\u0026#34;\u0026#34; def __init__( self, num_layers: int, d_model: int, num_heads: int, d_ff: int, vocab_size: int, max_seq_len: int, dropout: float = 0.1, ): super().__init__() self.token_embedding = nn.Embedding(vocab_size, d_model) self.pos_embedding = nn.Embedding(max_seq_len, d_model) self.dropout = nn.Dropout(dropout) # 核心：num_layers 个独立的层，每层参数互不共享 self.layers = nn.ModuleList([ TransformerLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.final_norm = nn.LayerNorm(d_model) def forward( self, token_ids: torch.Tensor, mask: torch.Tensor = None ) -\u0026gt; tuple[torch.Tensor, list[torch.Tensor]]: B, T = token_ids.shape positions = torch.arange(T, device=token_ids.device).unsqueeze(0) x = self.dropout(self.token_embedding(token_ids) + self.pos_embedding(positions)) layer_outputs = [] for layer in self.layers: x = layer(x, mask) # 上一层输出直接作为下一层输入 layer_outputs.append(x.clone()) # 记录每层输出，便于观察 return self.final_norm(x), layer_outputs # 使用示例 if __name__ == \u0026#34;__main__\u0026#34;: model = TransformerEncoder( num_layers=4, d_model=64, num_heads=4, d_ff=256, vocab_size=1000, max_seq_len=32, ) total_params = sum(p.numel() for p in model.parameters()) print(f\u0026#34;总参数量：{total_params:,}，每层参数量：{total_params // 4:,}（4层各自独立）\u0026#34;) token_ids = torch.randint(0, 1000, (2, 10)) # batch=2，序列长度=10 final_output, layer_outputs = model(token_ids) for i, out in enumerate(layer_outputs): print(f\u0026#34;第{i+1}层输出 shape: {out.shape} 均值={out.mean():.4f} std={out.std():.4f}\u0026#34;) print(f\u0026#34;最终输出 shape: {final_output.shape}\u0026#34;) # 每层 shape 相同 [2, 10, 64]，但数值不同——逐层抽象的体现 十、一句话总结 MHA = 把输入投影到 h 个低维子空间，每个子空间独立做 Self-Attention，捕捉不同语义关系，最后拼接并投影回原维度。\n多层堆叠 = 把 MHA + FFN + 残差 + LayerNorm 作为一个 Block，串行堆叠 N 次，每层有独立参数，逐层从低层特征抽象到高层语义。\n核心公式记忆：\nQ_i, K_i, V_i = X @ W_Q^i, X @ W_K^i, X @ W_V^i head_i = softmax(Q_i K_i^T / √d_k) · V_i MHA输出 = Concat(head_1...head_h) @ W_O 单层输出 = LayerNorm(X + FFN(LayerNorm(X + MHA(X)))) 多层 = 第(n)层输出 → 第(n+1)层输入，串行 N 次 ","permalink":"https://afan.ml/posts/m-h-a-%E5%A4%9A%E5%A4%B4%E6%B3%A8%E6%84%8F%E5%8A%9B%E7%AC%94%E8%AE%B0/","summary":"\u003cblockquote\u003e\n\u003cp\u003e前置知识：已了解 Self-Attention（自注意力）。\nMHA 的本质：\u003cstrong\u003e把 Self-Attention 并行做多次，每次关注不同的语义子空间，最后合并结果。\u003c/strong\u003e\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要多头self-attention-有什么不足\"\u003e一、为什么需要多头？Self-Attention 有什么不足？\u003c/h2\u003e\n\u003cp\u003e单头 Self-Attention 每次只能学到一种\u0026quot;关注模式\u0026quot;。例如：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e句子：\u0026#34;The animal didn\u0026#39;t cross the street because it was too tired\u0026#34;\n\u003c/code\u003e\u003c/pre\u003e\u003cul\u003e\n\u003cli\u003e头1 可能学到：\u003ccode\u003eit\u003c/code\u003e → \u003ccode\u003eanimal\u003c/code\u003e（指代关系）\u003c/li\u003e\n\u003cli\u003e头2 可能学到：\u003ccode\u003etired\u003c/code\u003e → \u003ccode\u003eanimal\u003c/code\u003e（状态描述）\u003c/li\u003e\n\u003cli\u003e头3 可能学到：\u003ccode\u003ecross\u003c/code\u003e → \u003ccode\u003estreet\u003c/code\u003e（动作与地点）\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e\u003cstrong\u003e单头只能同时关注一种模式，多头让模型在不同子空间里并行捕捉多种关系。\u003c/strong\u003e\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二mha-整体结构\"\u003e二、MHA 整体结构\u003c/h2\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e输入 X\n │\n ├──→ 线性投影 W_Q^1 → Q1 ─┐\n ├──→ 线性投影 W_K^1 → K1 ─┤→ Attention(Q1,K1,V1) → head_1 ─┐\n ├──→ 线性投影 W_V^1 → V1 ─┘ │\n │ │\n ├──→ 线性投影 W_Q^2 → Q2 ─┐ │\n ├──→ 线性投影 W_K^2 → K2 ─┤→ Attention(Q2,K2,V2) → head_2 ─┤→ Concat → 线性投影 W_O → 输出\n ├──→ 线性投影 W_V^2 → V2 ─┘ │\n │ │\n ├──→ ...（共 h 个头） │\n │ │\n └──→ 线性投影 W_Q^h → Qh ─┐ │\n 线性投影 W_K^h → Kh ─┤→ Attention(Qh,Kh,Vh) → head_h ─┘\n 线性投影 W_V^h → Vh ─┘\n\u003c/code\u003e\u003c/pre\u003e\u003chr\u003e\n\u003ch2 id=\"三完整公式\"\u003e三、完整公式\u003c/h2\u003e\n\u003ch3 id=\"单头-scaled-dot-product-attention回顾\"\u003e单头 Scaled Dot-Product Attention（回顾）\u003c/h3\u003e\n\u003cp\u003e$$\n\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V\n$$\u003c/p\u003e","title":"Multi-Head Attention（MHA）多头注意力笔记"},{"content":" 一句话：优化器决定\u0026quot;用梯度怎么更新参数\u0026quot;。从最朴素的 SGD，到加了动量的 SGD Momentum，再到自适应学习率的 Adam/AdamW，每一步改进都在解决上一代的具体问题。\n一、梯度下降的基本框架 所有优化器的核心都是：\n$$\\theta_{t+1} = \\theta_t - \\text{update}(g_t)$$\n其中 $g_t = \\nabla_\\theta \\mathcal{L}$ 是当前 batch 的梯度，不同优化器只是 $\\text{update}$ 的计算方式不同。\n三种梯度下降变体（按 batch 大小区分）：\n批量梯度下降（BGD）： 用全部数据算梯度，准确但极慢 随机梯度下降（SGD）： 用 1 条数据算梯度，快但噪声大 小批量梯度下降（MSGD）：用 mini-batch 算梯度，实践中的标准做法 现代所有\u0026quot;SGD\u0026quot;都指小批量梯度下降。\n二、SGD：最朴素的优化器 $$\\theta_{t+1} = \\theta_t - \\eta \\cdot g_t$$\n$\\eta$：学习率（步长） $g_t$：当前 mini-batch 的梯度 问题：\n学习率敏感：太大震荡，太小收敛慢 各维度步长相同：不同参数的梯度尺度可能差异极大，一个学习率无法同时适配所有参数 容易卡在鞍点：梯度为零但不是极值点 三、SGD + Momentum：加入惯性 $$v_t = \\beta \\cdot v_{t-1} + g_t$$ $$\\theta_{t+1} = \\theta_t - \\eta \\cdot v_t$$\n$v_t$：速度（历史梯度的指数加权平均） $\\beta$：动量系数，通常为 0.9 直觉：像一个球滚下山坡，历史的速度会叠加到当前更新上。\n解决的问题：\n方向一致的维度加速收敛（速度累积） 方向震荡的维度相互抵消（震荡被平滑） 能越过较小的局部极值和鞍点 四、Adam：自适应学习率 Adam（Adaptive Moment Estimation，2014，Kingma \u0026amp; Ba）的核心思想：为每个参数单独维护一个自适应学习率。\n完整公式 $$m_t = \\beta_1 \\cdot m_{t-1} + (1 - \\beta_1) \\cdot g_t \\quad \\text{（一阶动量，梯度均值）}$$ $$v_t = \\beta_2 \\cdot v_{t-1} + (1 - \\beta_2) \\cdot g_t^2 \\quad \\text{（二阶动量，梯度方差）}$$\n偏差修正（前几步 $m$、$v$ 偏小，需要修正）：\n$$\\hat{m}_t = \\frac{m_t}{1 - \\beta_1^t}, \\quad \\hat{v}_t = \\frac{v_t}{1 - \\beta_2^t}$$\n参数更新：\n$$\\theta_{t+1} = \\theta_t - \\frac{\\eta}{\\sqrt{\\hat{v}_t} + \\epsilon} \\cdot \\hat{m}_t$$\n默认超参数：$\\beta_1 = 0.9$，$\\beta_2 = 0.999$，$\\epsilon = 10^{-8}$，$\\eta = 10^{-3}$\n为什么 Adam 好 某个参数的梯度一直很大（如 embedding 层的高频词）： → v_t 很大 → 分母大 → 步长小 → 不会更新过猛 某个参数的梯度一直很小（如深层的稀疏参数）： → v_t 很小 → 分母小 → 步长大 → 不会更新太慢 效果：自动为每个参数找到合适的步长，对学习率的选择不那么敏感。\nAdamW：修复权重衰减 标准 Adam 的权重衰减实现有 bug：L2 正则化会被自适应学习率缩放，导致正则化效果不稳定。\nAdamW 把权重衰减从梯度里解耦出来，直接作用于参数：\n$$\\theta_{t+1} = \\theta_t - \\frac{\\eta}{\\sqrt{\\hat{v}_t} + \\epsilon} \\cdot \\hat{m}_t - \\eta \\lambda \\theta_t$$\n其中 $\\lambda$ 是权重衰减系数（通常 0.01~0.1）。LLM 训练几乎都用 AdamW。\n五、Adam 的显存开销 这是 Adam 最大的缺点：需要为每个参数额外存储两个状态（$m_t$ 和 $v_t$），全部用 FP32。\n对于参数量为 $N$ 的模型：\n模型权重（BF16）：N × 2 bytes 梯度（BF16）： N × 2 bytes m_t（FP32）： N × 4 bytes v_t（FP32）： N × 4 bytes ───────────────────────── 合计： N × 12 bytes 以 7B 模型为例：$7 \\times 10^9 \\times 12 \\approx 84\\ \\text{GB}$，光优化器状态就占 56 GB。\nAdam 8bit（bitsandbytes） 的解决方案：把 $m_t$、$v_t$ 量化到 8bit，显存从 $N \\times 12$ 降到 $N \\times 6$ bytes，节省约 50%。\n六、MNIST 完整示例：对比 SGD、Momentum、Adam import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms import time class MnistNet(nn.Module): \u0026#34;\u0026#34;\u0026#34; 简单的 3 层全连接网络，用于 MNIST 分类。 输入：28×28 = 784 维，输出：10 类 \u0026#34;\u0026#34;\u0026#34; def __init__(self): super().__init__() self.fc1 = nn.Linear(784, 256) self.fc2 = nn.Linear(256, 64) self.fc3 = nn.Linear(64, 10) def forward(self, x: torch.Tensor) -\u0026gt; torch.Tensor: x = x.view(x.size(0), -1) # [batch, 1, 28, 28] → [batch, 784] x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) return self.fc3(x) # 返回 logits，不做 softmax（CrossEntropyLoss 内部处理） def train_one_epoch( model: nn.Module, optimizer: torch.optim.Optimizer, loader: DataLoader, device: torch.device, ) -\u0026gt; tuple[float, float]: \u0026#34;\u0026#34;\u0026#34;训练一个 epoch，返回 (平均 loss, 准确率)。\u0026#34;\u0026#34;\u0026#34; model.train() total_loss = 0.0 total_correct = 0 total_samples = 0 for images, labels in loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() logits = model(images) loss = F.cross_entropy(logits, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) total_correct += (logits.argmax(dim=1) == labels).sum().item() total_samples += images.size(0) return total_loss / total_samples, total_correct / total_samples def evaluate( model: nn.Module, loader: DataLoader, device: torch.device, ) -\u0026gt; float: \u0026#34;\u0026#34;\u0026#34;在测试集上评估准确率。\u0026#34;\u0026#34;\u0026#34; model.eval() total_correct = 0 total_samples = 0 with torch.no_grad(): for images, labels in loader: images, labels = images.to(device), labels.to(device) logits = model(images) total_correct += (logits.argmax(dim=1) == labels).sum().item() total_samples += images.size(0) return total_correct / total_samples def count_optimizer_memory_bytes(optimizer: torch.optim.Optimizer) -\u0026gt; int: \u0026#34;\u0026#34;\u0026#34;估算优化器状态占用的字节数。\u0026#34;\u0026#34;\u0026#34; total_bytes = 0 for group in optimizer.param_groups: for param in group[\u0026#34;params\u0026#34;]: state = optimizer.state[param] for tensor in state.values(): if isinstance(tensor, torch.Tensor): total_bytes += tensor.numel() * tensor.element_size() return total_bytes def run_experiment( optimizer_name: str, optimizer: torch.optim.Optimizer, num_epochs: int, train_loader: DataLoader, test_loader: DataLoader, device: torch.device, model: nn.Module, ) -\u0026gt; None: print(\u0026#34;=\u0026#34; * 50) print(\u0026#34;Optimizer: %s\u0026#34; % optimizer_name) print(\u0026#34;=\u0026#34; * 50) start_time = time.time() for epoch in range(1, num_epochs + 1): train_loss, train_acc = train_one_epoch(model, optimizer, train_loader, device) test_acc = evaluate(model, test_loader, device) print(\u0026#34; Epoch %2d | loss=%.4f | train_acc=%.4f | test_acc=%.4f\u0026#34; % ( epoch, train_loss, train_acc, test_acc)) elapsed = time.time() - start_time # 统计优化器显存（只有 Adam 在第一步后才会分配状态） optimizer_mem_bytes = count_optimizer_memory_bytes(optimizer) model_param_bytes = sum(p.numel() * p.element_size() for p in model.parameters()) print(\u0026#34; Time: %.1fs\u0026#34; % elapsed) print(\u0026#34; Model params: %.1f KB\u0026#34; % (model_param_bytes / 1024)) print(\u0026#34; Optimizer state: %.1f KB\u0026#34; % (optimizer_mem_bytes / 1024)) if model_param_bytes \u0026gt; 0: print(\u0026#34; Opt/Param ratio: %.1fx\u0026#34; % (optimizer_mem_bytes / model_param_bytes)) if __name__ == \u0026#34;__main__\u0026#34;: device = torch.device(\u0026#34;cuda\u0026#34; if torch.cuda.is_available() else \u0026#34;cpu\u0026#34;) print(f\u0026#34;使用设备: {device}\u0026#34;) # 数据集下载路径（兼容 Windows 和 Linux/Mac） import os mnist_root = os.path.join(os.path.expanduser(\u0026#34;~\u0026#34;), \u0026#34;tmp\u0026#34;, \u0026#34;mnist\u0026#34;) os.makedirs(mnist_root, exist_ok=True) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)), # MNIST 的均值和标准差 ]) train_dataset = datasets.MNIST(root=mnist_root, train=True, download=True, transform=transform) test_dataset = datasets.MNIST(root=mnist_root, train=False, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=0) test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=0) num_epochs = 5 # ── 实验 1：SGD（无动量）────────────────────────────────────────────────── model_sgd = MnistNet().to(device) optimizer_sgd = torch.optim.SGD(model_sgd.parameters(), lr=0.01) run_experiment(\u0026#34;SGD (lr=0.01)\u0026#34;, optimizer_sgd, num_epochs, train_loader, test_loader, device, model_sgd) # ── 实验 2：SGD + Momentum ──────────────────────────────────────────────── model_momentum = MnistNet().to(device) optimizer_momentum = torch.optim.SGD(model_momentum.parameters(), lr=0.01, momentum=0.9) run_experiment(\u0026#34;SGD + Momentum (lr=0.01, β=0.9)\u0026#34;, optimizer_momentum, num_epochs, train_loader, test_loader, device, model_momentum) # ── 实验 3：Adam ────────────────────────────────────────────────────────── model_adam = MnistNet().to(device) optimizer_adam = torch.optim.Adam(model_adam.parameters(), lr=0.001) run_experiment(\u0026#34;Adam (lr=0.001)\u0026#34;, optimizer_adam, num_epochs, train_loader, test_loader, device, model_adam) # ── 实验 4：AdamW ───────────────────────────────────────────────────────── model_adamw = MnistNet().to(device) optimizer_adamw = torch.optim.AdamW(model_adamw.parameters(), lr=0.001, weight_decay=0.01) run_experiment(\u0026#34;AdamW (lr=0.001, wd=0.01)\u0026#34;, optimizer_adamw, num_epochs, train_loader, test_loader, device, model_adamw) # ── 显存汇总 ────────────────────────────────────────────────────────────── print(f\u0026#34;\\n{\u0026#39;=\u0026#39;*50}\u0026#34;) print(\u0026#34;显存占用汇总（MNIST 小模型，参数量约 200K）\u0026#34;) print(f\u0026#34;{\u0026#39;=\u0026#39;*50}\u0026#34;) model_param_bytes = sum(p.numel() * p.element_size() for p in model_adam.parameters()) print(f\u0026#34;模型参数（FP32）：{model_param_bytes / 1024:.0f} KB\u0026#34;) for name, opt in [(\u0026#34;SGD\u0026#34;, optimizer_sgd), (\u0026#34;SGD+Momentum\u0026#34;, optimizer_momentum), (\u0026#34;Adam\u0026#34;, optimizer_adam), (\u0026#34;AdamW\u0026#34;, optimizer_adamw)]: opt_bytes = count_optimizer_memory_bytes(opt) print(f\u0026#34;{name:20s}：优化器状态 {opt_bytes/1024:.0f} KB（模型参数的 {opt_bytes/model_param_bytes:.1f}x）\u0026#34;) 七、各优化器显存对比（理论） 对参数量为 $N$ 的模型，以 FP32 存储为基准（4 bytes/参数）：\n优化器 额外存储的状态 额外显存 总显存（含参数） SGD 无 0 $N \\times 4$ bytes SGD + Momentum 速度 $v$（FP32） $N \\times 4$ $N \\times 8$ bytes Adam / AdamW $m$（FP32）+ $v$（FP32） $N \\times 8$ $N \\times 12$ bytes Adam 8bit $m$（INT8）+ $v$（INT8） $N \\times 2$ $N \\times 6$ bytes 以 7B 模型（BF16 权重）为例：\n模型权重（BF16）： 7B × 2 = 14 GB 梯度（BF16）： 7B × 2 = 14 GB Adam m_t（FP32）： 7B × 4 = 28 GB Adam v_t（FP32）： 7B × 4 = 28 GB ───────────────────────────────────── 全量微调合计： 84 GB Adam 8bit 的优化器状态：7B × 1 + 7B × 1 = 14 GB（节省 42 GB） LoRA 只对 ~1% 参数更新：Adam 状态 ≈ 0.56 GB（几乎可忽略） 八、核心要点速查 问题 答案 SGD 的缺点？ 各参数步长相同，对学习率敏感，容易震荡或卡鞍点 Momentum 解决什么？ 用历史梯度的指数平均平滑更新，加速一致方向、抑制震荡方向 Adam 的两个动量是什么？ $m_t$：梯度均值（一阶）；$v_t$：梯度方差（二阶） 为什么要偏差修正？ 初始 $m_0=v_0=0$，前几步估计偏小，除以 $(1-\\beta^t)$ 修正 AdamW 和 Adam 的区别？ AdamW 把权重衰减从梯度解耦，正则化效果更稳定 Adam 为什么显存贵？ 需要额外存 $m_t$ 和 $v_t$，均为 FP32，是参数量的 2 倍显存 Adam 8bit 怎么省显存？ 把 $m_t$、$v_t$ 量化到 8bit，优化器状态显存减少 75% LLM 微调用哪个？ adamw_8bit（bitsandbytes 提供），显存省、效果好 ","permalink":"https://afan.ml/posts/%E4%BC%98%E5%8C%96%E5%99%A8-optimizer/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：优化器决定\u0026quot;用梯度怎么更新参数\u0026quot;。从最朴素的 SGD，到加了动量的 SGD Momentum，再到自适应学习率的 Adam/AdamW，每一步改进都在解决上一代的具体问题。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一梯度下降的基本框架\"\u003e一、梯度下降的基本框架\u003c/h2\u003e\n\u003cp\u003e所有优化器的核心都是：\u003c/p\u003e\n\u003cp\u003e$$\\theta_{t+1} = \\theta_t - \\text{update}(g_t)$$\u003c/p\u003e\n\u003cp\u003e其中 $g_t = \\nabla_\\theta \\mathcal{L}$ 是当前 batch 的梯度，不同优化器只是 $\\text{update}$ 的计算方式不同。\u003c/p\u003e\n\u003cp\u003e\u003cstrong\u003e三种梯度下降变体\u003c/strong\u003e（按 batch 大小区分）：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e批量梯度下降（BGD）： 用全部数据算梯度，准确但极慢\n随机梯度下降（SGD）： 用 1 条数据算梯度，快但噪声大\n小批量梯度下降（MSGD）：用 mini-batch 算梯度，实践中的标准做法\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e现代所有\u0026quot;SGD\u0026quot;都指小批量梯度下降。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二sgd最朴素的优化器\"\u003e二、SGD：最朴素的优化器\u003c/h2\u003e\n\u003cp\u003e$$\\theta_{t+1} = \\theta_t - \\eta \\cdot g_t$$\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e$\\eta$：学习率（步长）\u003c/li\u003e\n\u003cli\u003e$g_t$：当前 mini-batch 的梯度\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e\u003cstrong\u003e问题\u003c/strong\u003e：\u003c/p\u003e\n\u003col\u003e\n\u003cli\u003e\u003cstrong\u003e学习率敏感\u003c/strong\u003e：太大震荡，太小收敛慢\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e各维度步长相同\u003c/strong\u003e：不同参数的梯度尺度可能差异极大，一个学习率无法同时适配所有参数\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e容易卡在鞍点\u003c/strong\u003e：梯度为零但不是极值点\u003c/li\u003e\n\u003c/ol\u003e\n\u003chr\u003e\n\u003ch2 id=\"三sgd--momentum加入惯性\"\u003e三、SGD + Momentum：加入惯性\u003c/h2\u003e\n\u003cp\u003e$$v_t = \\beta \\cdot v_{t-1} + g_t$$\n$$\\theta_{t+1} = \\theta_t - \\eta \\cdot v_t$$\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e$v_t$：速度（历史梯度的指数加权平均）\u003c/li\u003e\n\u003cli\u003e$\\beta$：动量系数，通常为 0.9\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e\u003cstrong\u003e直觉\u003c/strong\u003e：像一个球滚下山坡，历史的速度会叠加到当前更新上。\u003c/p\u003e","title":"优化器（Optimizer）笔记"},{"content":" 一句话：Transformer 的 Attention 对输入顺序无感知，位置编码负责把\u0026quot;第几个 token\u0026quot;这个信息注入模型。从固定正弦函数到可学习参数，再到当前主流的旋转编码（RoPE），每一代都在解决上一代的长度外推问题。\n一、为什么需要位置编码？ RNN 逐步处理序列，天然感知顺序。Transformer 并行处理所有 token，Attention 本质是集合操作，对顺序完全无感：\n[\u0026#34;我\u0026#34;, \u0026#34;爱\u0026#34;, \u0026#34;你\u0026#34;] 和 [\u0026#34;你\u0026#34;, \u0026#34;爱\u0026#34;, \u0026#34;我\u0026#34;] → Attention 看到的是完全相同的三个向量集合，无法区分 必须手动把位置信息注入模型，这就是位置编码的作用。\n二、发展历程（简览） 时间 方案 代表模型 核心问题 2017 Sinusoidal PE（正弦固定编码） 原版 Transformer 外推效果差，只有绝对位置 2018 Learned PE（可学习编码） BERT、GPT-2 硬性长度上限，超出训练长度就失效 2019 Relative PE（相对位置编码） Transformer-XL、T5 实现复杂，外推仍有限 2021 RoPE（旋转位置编码） LLaMA、Qwen、DeepSeek 当前主流 ✅ 2022 ALiBi（线性偏置） MPT 强制近距离偏好，长程依赖受损 三、前三代简述 Sinusoidal PE（2017） 用不同频率的正弦/余弦函数生成固定位置向量，直接加到 token embedding 上：\n$$PE_{(pos,\\ 2i)} = \\sin!\\left(\\frac{pos}{10000^{2i/d_{model}}}\\right), \\quad PE_{(pos,\\ 2i+1)} = \\cos!\\left(\\frac{pos}{10000^{2i/d_{model}}}\\right)$$\n优点：无需训练，参数量为 0 缺点：超出训练长度后效果急剧下降，只有绝对位置信息 Learned PE（2018） 把位置编码当成普通 Embedding 参数，随模型一起训练（BERT 最大 512，GPT-2 最大 1024）。\n优点：模型自己学习最优位置表示，实践效果略好于 Sinusoidal 缺点：硬性长度上限，训练时没见过的位置完全没有编码 Relative PE / T5 Bias（2019） 在 Attention 分数上加可学习的相对位置偏置 $b_{i-j}$，编码\u0026quot;两个 token 相距多远\u0026quot;而非\u0026quot;在第几个位置\u0026quot;。\n优点：相对位置更自然，泛化更好 缺点：实现复杂，长度外推仍有限 四、RoPE（旋转位置编码）—— 当前主流详解 4.1 核心思想 RoPE 不把位置信息加到 Embedding 上，而是在计算 Attention 时，把位置信息旋转进 Q 和 K。\n旋转后的 Q 和 K 做点积，结果只依赖相对位置 $m - n$，与绝对位置无关：\n$$\\langle R_m, q,\\ R_n, k \\rangle = \\langle q,\\ R_{n-m}, k \\rangle$$\n其中 $R_m$ 是位置 $m$ 对应的旋转矩阵。\n4.2 二维直觉 把二维向量 $[x, y]$ 旋转 $\\theta \\cdot pos$ 角度：\n$$\\begin{bmatrix} x\u0026rsquo; \\ y\u0026rsquo; \\end{bmatrix} = \\begin{bmatrix} \\cos(\\theta \\cdot pos) \u0026amp; -\\sin(\\theta \\cdot pos) \\ \\sin(\\theta \\cdot pos) \u0026amp; \\cos(\\theta \\cdot pos) \\end{bmatrix} \\begin{bmatrix} x \\ y \\end{bmatrix}$$\n两个旋转向量的点积，旋转角度相减，只剩相对位置差 $\\theta \\cdot (m - n)$。\n4.3 高维实现：两两配对旋转 实际的 head_dim 是高维的（如 128 维），RoPE 把维度两两配对，每对独立做二维旋转：\nhead_dim = 128 → 64 对，每对用不同频率的 θ 旋转 维度 [0,1] 用频率 θ₀，维度 [2,3] 用频率 θ₁，...，维度 [126,127] 用频率 θ₆₃ 频率从高到低，低频维度感知长程位置，高频维度感知短程位置 4.4 完整代码（PyTorch） import torch import math def build_rope_cache(seq_len: int, head_dim: int, base: int = 10000): \u0026#34;\u0026#34;\u0026#34; 预计算各位置的 cos/sin 旋转值。 返回 cos, sin，形状均为 [seq_len, head_dim]。 \u0026#34;\u0026#34;\u0026#34; # 每对维度对应的旋转频率，共 head_dim//2 个频率 half_dim = head_dim // 2 freqs = 1.0 / (base ** (torch.arange(0, half_dim).float() / half_dim)) # 每个位置与每个频率的乘积，形状 [seq_len, half_dim] positions = torch.arange(seq_len).float() angles = torch.outer(positions, freqs) # 拼接成 [seq_len, head_dim]，偶数维度和奇数维度共享同一组角度 angles = torch.cat([angles, angles], dim=-1) return angles.cos(), angles.sin() def rotate_half(x: torch.Tensor) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; 把向量的后半段取负后移到前半段，实现旋转的另一个分量。 输入 [x1, x2]，输出 [-x2, x1]。 \u0026#34;\u0026#34;\u0026#34; half = x.shape[-1] // 2 x1 = x[..., :half] x2 = x[..., half:] return torch.cat([-x2, x1], dim=-1) def apply_rope( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, ) -\u0026gt; torch.Tensor: \u0026#34;\u0026#34;\u0026#34; 对 Q 或 K 应用 RoPE 旋转。 x: [batch, seq_len, num_heads, head_dim] cos: [seq_len, head_dim] sin: [seq_len, head_dim] \u0026#34;\u0026#34;\u0026#34; # 广播到 [1, seq_len, 1, head_dim] cos = cos.unsqueeze(0).unsqueeze(2) sin = sin.unsqueeze(0).unsqueeze(2) # 旋转公式：x·cos + rotate_half(x)·sin return x * cos + rotate_half(x) * sin # ── 验证 ────────────────────────────────────────────── batch, seq_len, num_heads, head_dim = 2, 16, 8, 64 cos, sin = build_rope_cache(seq_len, head_dim) q = torch.randn(batch, seq_len, num_heads, head_dim) k = torch.randn(batch, seq_len, num_heads, head_dim) q_rotated = apply_rope(q, cos, sin) k_rotated = apply_rope(k, cos, sin) print(q_rotated.shape) # torch.Size([2, 16, 8, 64]) print(k_rotated.shape) # torch.Size([2, 16, 8, 64]) # 验证相对位置不变性： # 位置 m=3 的 q 和位置 n=1 的 k 的点积， # 应等于位置 m=5 的 q 和位置 n=3 的 k 的点积（相对距离都是 2） q2 = torch.randn(1, seq_len, 1, head_dim) k2 = torch.randn(1, seq_len, 1, head_dim) q2_rot = apply_rope(q2, cos, sin) k2_rot = apply_rope(k2, cos, sin) dot_31 = (q2_rot[0, 3, 0] * k2_rot[0, 1, 0]).sum() dot_53 = (q2_rot[0, 5, 0] * k2_rot[0, 3, 0]).sum() # 注意：q2/k2 是随机向量，两个点积数值不同，但它们的差异只来自 q/k 本身， # 旋转部分贡献的相对位置偏移量（cos(2θ)项）是相同的 print(f\u0026#34;dot(q3, k1) = {dot_31.item():.4f}\u0026#34;) print(f\u0026#34;dot(q5, k3) = {dot_53.item():.4f}\u0026#34;) print(\u0026#34;RoPE 验证完成\u0026#34;) 4.5 RoPE 的优点 优点 说明 天然相对位置 Q·K 点积只依赖相对位置差，不依赖绝对位置 无额外参数 cos/sin 由公式计算，不需要训练 长度外推性好 配合 YaRN、LongRoPE 可从 4K 扩展到 128K+ 实现高效 只需对 Q/K 做逐元素乘法，无额外矩阵运算 4.6 长度外推扩展：YaRN 原始 RoPE 训练 4K 后直接推理 32K 效果仍会下降。YaRN（2023） 通过对不同频率维度做差异化缩放解决这个问题：\n高频维度（感知短程）：直接用原始频率，不缩放 低频维度（感知长程）：按扩展比例缩放频率，让模型\u0026#34;看得更远\u0026#34; 中间维度：线性插值过渡 LLaMA-3 的 128K 上下文、Qwen2 的 128K 上下文都依赖 RoPE + YaRN 类扩展实现。\n五、ALiBi（线性偏置）简述 不修改 Embedding，直接在 Attention 分数上加距离惩罚：\n$$\\text{Attention}(i, j) = \\frac{q_i \\cdot k_j}{\\sqrt{d_k}} - m \\cdot |i - j|$$\n距离越远，惩罚越大，模型天然倾向于关注近处 token。实现极简，长度外推效果不错，但强制近距离偏好对长程依赖任务有损失，目前主流模型已基本不用。\n六、各方案对比 方案 参数量 长度外推 相对位置 当前使用 Sinusoidal PE 0 ❌ 差 ❌ 无 已淘汰 Learned PE max_len × d_model ❌ 无法外推 ❌ 无 BERT/GPT-2 Relative PE（T5） 少量偏置参数 ⚠️ 有限 ✅ 有 T5 系列 RoPE 0 ✅ 配合YaRN可达128K+ ✅ 天然 LLaMA/Qwen/DeepSeek ALiBi 0 ✅ 较好 ✅ 有 MPT（较少） 七、核心要点速查 问题 答案 为什么需要位置编码？ Attention 是集合操作，对顺序无感，必须手动注入位置信息 位置编码加在哪里？ Sinusoidal/Learned 加在 Embedding 上；RoPE 加在 Q/K 的旋转里 当前主流是什么？ RoPE，LLaMA/Qwen/DeepSeek 等几乎所有新模型都在用 RoPE 为什么好？ 天然相对位置、无额外参数、配合 YaRN 可外推超长上下文 RoPE 作用在哪里？ 只作用于 Q 和 K，不作用于 V 长度外推怎么做？ RoPE + YaRN/LongRoPE，对低频维度缩放频率 ALiBi 和 RoPE 的区别？ ALiBi 在 Attention 分数上加距离惩罚；RoPE 旋转 Q/K 向量本身 ","permalink":"https://afan.ml/posts/%E4%BD%8D%E7%BD%AE%E7%BC%96%E7%A0%81-positional-encoding/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：Transformer 的 Attention 对输入顺序无感知，位置编码负责把\u0026quot;第几个 token\u0026quot;这个信息注入模型。从固定正弦函数到可学习参数，再到当前主流的旋转编码（RoPE），每一代都在解决上一代的长度外推问题。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要位置编码\"\u003e一、为什么需要位置编码？\u003c/h2\u003e\n\u003cp\u003eRNN 逐步处理序列，天然感知顺序。Transformer 并行处理所有 token，Attention 本质是\u003cstrong\u003e集合操作\u003c/strong\u003e，对顺序完全无感：\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e[\u0026#34;我\u0026#34;, \u0026#34;爱\u0026#34;, \u0026#34;你\u0026#34;] 和 [\u0026#34;你\u0026#34;, \u0026#34;爱\u0026#34;, \u0026#34;我\u0026#34;]\n→ Attention 看到的是完全相同的三个向量集合，无法区分\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e必须手动把位置信息注入模型，这就是位置编码的作用。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二发展历程简览\"\u003e二、发展历程（简览）\u003c/h2\u003e\n\u003ctable\u003e\n \u003cthead\u003e\n \u003ctr\u003e\n \u003cth\u003e时间\u003c/th\u003e\n \u003cth\u003e方案\u003c/th\u003e\n \u003cth\u003e代表模型\u003c/th\u003e\n \u003cth\u003e核心问题\u003c/th\u003e\n \u003c/tr\u003e\n \u003c/thead\u003e\n \u003ctbody\u003e\n \u003ctr\u003e\n \u003ctd\u003e2017\u003c/td\u003e\n \u003ctd\u003eSinusoidal PE（正弦固定编码）\u003c/td\u003e\n \u003ctd\u003e原版 Transformer\u003c/td\u003e\n \u003ctd\u003e外推效果差，只有绝对位置\u003c/td\u003e\n \u003c/tr\u003e\n \u003ctr\u003e\n \u003ctd\u003e2018\u003c/td\u003e\n \u003ctd\u003eLearned PE（可学习编码）\u003c/td\u003e\n \u003ctd\u003eBERT、GPT-2\u003c/td\u003e\n \u003ctd\u003e硬性长度上限，超出训练长度就失效\u003c/td\u003e\n \u003c/tr\u003e\n \u003ctr\u003e\n \u003ctd\u003e2019\u003c/td\u003e\n \u003ctd\u003eRelative PE（相对位置编码）\u003c/td\u003e\n \u003ctd\u003eTransformer-XL、T5\u003c/td\u003e\n \u003ctd\u003e实现复杂，外推仍有限\u003c/td\u003e\n \u003c/tr\u003e\n \u003ctr\u003e\n \u003ctd\u003e2021\u003c/td\u003e\n \u003ctd\u003e\u003cstrong\u003eRoPE（旋转位置编码）\u003c/strong\u003e\u003c/td\u003e\n \u003ctd\u003eLLaMA、Qwen、DeepSeek\u003c/td\u003e\n \u003ctd\u003e\u003cstrong\u003e当前主流\u003c/strong\u003e ✅\u003c/td\u003e\n \u003c/tr\u003e\n \u003ctr\u003e\n \u003ctd\u003e2022\u003c/td\u003e\n \u003ctd\u003eALiBi（线性偏置）\u003c/td\u003e\n \u003ctd\u003eMPT\u003c/td\u003e\n \u003ctd\u003e强制近距离偏好，长程依赖受损\u003c/td\u003e\n \u003c/tr\u003e\n \u003c/tbody\u003e\n\u003c/table\u003e\n\u003chr\u003e\n\u003ch2 id=\"三前三代简述\"\u003e三、前三代简述\u003c/h2\u003e\n\u003ch3 id=\"sinusoidal-pe2017\"\u003eSinusoidal PE（2017）\u003c/h3\u003e\n\u003cp\u003e用不同频率的正弦/余弦函数生成固定位置向量，直接加到 token embedding 上：\u003c/p\u003e\n\u003cp\u003e$$PE_{(pos,\\ 2i)} = \\sin!\\left(\\frac{pos}{10000^{2i/d_{model}}}\\right), \\quad PE_{(pos,\\ 2i+1)} = \\cos!\\left(\\frac{pos}{10000^{2i/d_{model}}}\\right)$$\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e\u003cstrong\u003e优点\u003c/strong\u003e：无需训练，参数量为 0\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e缺点\u003c/strong\u003e：超出训练长度后效果急剧下降，只有绝对位置信息\u003c/li\u003e\n\u003c/ul\u003e\n\u003ch3 id=\"learned-pe2018\"\u003eLearned PE（2018）\u003c/h3\u003e\n\u003cp\u003e把位置编码当成普通 Embedding 参数，随模型一起训练（BERT 最大 512，GPT-2 最大 1024）。\u003c/p\u003e","title":"位置编码（Positional Encoding）笔记"},{"content":" 一句话：残差网络通过在层与层之间添加\u0026quot;跳跃连接\u0026quot;（Skip Connection），让网络学习输入与输出的差值而非完整映射，从而解决深层网络的梯度消失和退化问题。\n一、问题背景：深层网络为什么会退化？ 直觉上，更深的网络应该更强——至少多出来的层可以学成\u0026quot;什么都不做\u0026quot;（恒等映射）。\n但实验发现：网络越深，训练误差反而越大，不是过拟合，是训练本身就失败了。\n20层网络的训练误差 \u0026lt; 56层网络的训练误差 ← 这不正常！ 根本原因：梯度消失\n反向传播时，梯度要经过几十层连乘。每层梯度都小于 1，连乘后趋近于 0，浅层几乎学不到任何东西。\n二、核心思想：学\u0026quot;差值\u0026quot;而不是学\u0026quot;完整映射\u0026quot; 原来的目标 让网络直接学习从输入到输出的完整映射：\n$$H(x) = \\text{网络想要的输出}$$\n残差的目标 把目标拆成两部分：\n$$H(x) = F(x) + x$$\n其中 $F(x) = H(x) - x$ 就是残差（输出和输入的差值），网络只需要学 $F(x)$。\n为什么这样更容易学？\n如果这一层什么都不需要做（恒等映射），只需让 $F(x) \\to 0$，把权重归零即可 让网络直接学 $H(x) = x$ 要难得多——它需要精确拟合一个恒等函数 学\u0026quot;微小的修正量\u0026quot;比学\u0026quot;完整的变换\u0026quot;容易得多 三、结构图 输入 x │ ├──────────────────────────┐ ← shortcut（跳跃连接，直接复制 x） │ │ ▼ │ Conv → BN → ReLU │ ▼ │ Conv → BN │ │ │ └──────────── + ───────────┘ ← 相加（不是拼接！） │ ReLU │ 输出 H(x) = F(x) + x 关键细节：\n+ 是逐元素相加，不是 Concat 拼接 ReLU 在相加之后才做，不是在每层卷积后立刻做最后一次 ReLU shortcut 路径上没有参数（维度相同时），是纯粹的恒等连接 四、维度匹配问题：不是无条件能直接加 相加要求 x 和 F(x) 的形状完全一致，有两种情况：\n情况1：维度相同 → 直接加 ✅ 通道数不变、特征图尺寸不变时，shortcut 就是一条空线：\noutput = F(x) + x 情况2：维度不同 → 需要投影 ⚠️ 比如通道数从 64 → 128，或特征图从 56×56 → 28×28，形状不一样，不能直接加。\n解决方案：用 1×1 卷积 对 shortcut 做线性投影，把维度对齐：\n# shortcut 用 1x1 卷积调整通道数和尺寸 shortcut = Conv2d(in_channels=64, out_channels=128, kernel_size=1, stride=2)(x) output = F(x) + shortcut 1×1 卷积只调整维度，不提取空间特征，参数量极少。\n五、为什么能解决梯度消失？ 对 $\\text{output} = F(x) + x$ 求梯度：\n$$\\frac{\\partial \\text{Loss}}{\\partial x} = \\frac{\\partial \\text{Loss}}{\\partial \\text{out}} \\cdot \\left(1 + \\frac{\\partial F(x)}{\\partial x}\\right)$$\n注意那个 1：\n无论 $F(x)$ 的梯度多小，梯度里始终有一个常数项 1 在传递 梯度可以通过 shortcut 路径无损地直接流回浅层 彻底解决了梯度消失问题，使得训练几百层的网络成为可能 六、完整代码（PyTorch） import torch import torch.nn as nn class ResidualBlock(nn.Module): \u0026#34;\u0026#34;\u0026#34;标准 ResNet BasicBlock（两层 3x3 卷积）\u0026#34;\u0026#34;\u0026#34; def __init__(self, in_channels, out_channels, stride=1): super().__init__() # 主路径：两层卷积 self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) # shortcut 路径：只有维度不匹配时才需要 1x1 卷积 self.shortcut = None if stride != 1 or in_channels != out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(out_channels) ) def forward(self, x): identity = x # 保存输入，用于后面相加 # 主路径 out = self.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) # shortcut 路径（维度不匹配时做投影） if self.shortcut is not None: identity = self.shortcut(x) # 残差相加，然后再 ReLU out = self.relu(out + identity) return out # 验证两种情况 block_same = ResidualBlock(64, 64, stride=1) # 维度相同，shortcut=None block_diff = ResidualBlock(64, 128, stride=2) # 维度不同，shortcut=1x1卷积 x = torch.randn(1, 64, 56, 56) print(block_same(x).shape) # torch.Size([1, 64, 56, 56]) print(block_diff(x).shape) # torch.Size([1, 128, 28, 28]) 七、Transformer 里的残差连接 残差连接不只是 CNN 的专利，Transformer 里也随处可见，思想完全一样：\n# Attention 子层 x = x + Attention(LayerNorm(x)) # FFN 子层 x = x + FFN(LayerNorm(x)) 每个子层的输出都加回输入，梯度可以直接从输出层流回任意浅层，这也是 Transformer 能堆几十层的原因。\n八、核心要点速查 问题 答案 解决了什么问题？ 深层网络梯度消失、训练退化 核心操作是什么？ output = F(x) + x，把输入加到输出上 能直接加吗？ 维度相同时可以；维度不同需要 1×1 卷积对齐 + 是拼接吗？ 不是，是逐元素相加，形状不变 ReLU 在哪里做？ 相加之后再做最后一次 ReLU 为什么有效？ 梯度公式里有常数项 1，梯度可无损回传 Transformer 里有吗？ 有，x + Attention(x) 和 x + FFN(x) 都是 ","permalink":"https://afan.ml/posts/%E6%AE%8B%E5%B7%AE%E7%BD%91%E7%BB%9C-res-net/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：残差网络通过在层与层之间添加\u0026quot;跳跃连接\u0026quot;（Skip Connection），让网络学习输入与输出的\u003cstrong\u003e差值\u003c/strong\u003e而非完整映射，从而解决深层网络的梯度消失和退化问题。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一问题背景深层网络为什么会退化\"\u003e一、问题背景：深层网络为什么会退化？\u003c/h2\u003e\n\u003cp\u003e直觉上，更深的网络应该更强——至少多出来的层可以学成\u0026quot;什么都不做\u0026quot;（恒等映射）。\u003c/p\u003e\n\u003cp\u003e但实验发现：\u003cstrong\u003e网络越深，训练误差反而越大\u003c/strong\u003e，不是过拟合，是训练本身就失败了。\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e20层网络的训练误差 \u0026lt; 56层网络的训练误差 ← 这不正常！\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e根本原因：梯度消失\u003c/strong\u003e\u003c/p\u003e\n\u003cp\u003e反向传播时，梯度要经过几十层连乘。每层梯度都小于 1，连乘后趋近于 0，浅层几乎学不到任何东西。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二核心思想学差值而不是学完整映射\"\u003e二、核心思想：学\u0026quot;差值\u0026quot;而不是学\u0026quot;完整映射\u0026quot;\u003c/h2\u003e\n\u003ch3 id=\"原来的目标\"\u003e原来的目标\u003c/h3\u003e\n\u003cp\u003e让网络直接学习从输入到输出的完整映射：\u003c/p\u003e\n\u003cp\u003e$$H(x) = \\text{网络想要的输出}$$\u003c/p\u003e\n\u003ch3 id=\"残差的目标\"\u003e残差的目标\u003c/h3\u003e\n\u003cp\u003e把目标拆成两部分：\u003c/p\u003e\n\u003cp\u003e$$H(x) = F(x) + x$$\u003c/p\u003e\n\u003cp\u003e其中 $F(x) = H(x) - x$ 就是\u003cstrong\u003e残差\u003c/strong\u003e（输出和输入的差值），网络只需要学 $F(x)$。\u003c/p\u003e\n\u003cp\u003e\u003cstrong\u003e为什么这样更容易学？\u003c/strong\u003e\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e如果这一层什么都不需要做（恒等映射），只需让 $F(x) \\to 0$，把权重归零即可\u003c/li\u003e\n\u003cli\u003e让网络直接学 $H(x) = x$ 要难得多——它需要精确拟合一个恒等函数\u003c/li\u003e\n\u003cli\u003e学\u0026quot;微小的修正量\u0026quot;比学\u0026quot;完整的变换\u0026quot;容易得多\u003c/li\u003e\n\u003c/ul\u003e\n\u003chr\u003e\n\u003ch2 id=\"三结构图\"\u003e三、结构图\u003c/h2\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e输入 x\n │\n ├──────────────────────────┐ ← shortcut（跳跃连接，直接复制 x）\n │ │\n ▼ │\nConv → BN → ReLU │\n ▼ │\nConv → BN │\n │ │\n └──────────── + ───────────┘ ← 相加（不是拼接！）\n │\n ReLU\n │\n 输出 H(x) = F(x) + x\n\u003c/code\u003e\u003c/pre\u003e\u003cp\u003e\u003cstrong\u003e关键细节\u003c/strong\u003e：\u003c/p\u003e","title":"残差网络（ResNet）笔记"},{"content":" 一句话：计算期望本质是计算积分，积分算不了就用采样均值代替，均匀采样浪费在无效区域导致方差大，重要性采样换一个\u0026quot;聪明的分布 q\u0026quot;集中采有效区域，再用权重 p/q 修正偏差，保证结果无偏且方差更小。\n一、为什么需要重要性采样？ 计算期望 = 计算积分 $$\\mathbb{E}_p[f(x)] = \\int f(x) \\cdot p(x) , dx$$\n很多场景这个积分没有解析解，无法直接用数学公式算：\n高维积分：神经网络参数空间是几百万维，数值积分根本算不了 分布形状复杂：多峰分布、奇形怪状的后验分布，没有现成公式 采样代价高：强化学习里从新策略采样需要重新跑环境，代价极高 这时只能用蒙特卡洛方法：采样求均值代替积分。\n二、蒙特卡洛：采样求均值 从分布 $p$ 采 $N$ 个样本，用样本均值估计期望：\n$$\\mathbb{E}p[f(x)] \\approx \\frac{1}{N} \\sum{i=1}^{N} f(x_i), \\quad x_i \\sim p$$\n问题：如果直接从 $p$ 采样很难（或代价高），或者均匀采样导致大量样本浪费在 $f(x) \\approx 0$ 的区域，方差就会很大。\n估算 P(x \u0026gt; 4)，x ~ N(0,1)： 均匀采样 100 万个点，平均只有 31 个落在 x\u0026gt;4 → 99.997% 的样本对结果没有贡献，方差极大 三、重要性采样：换一个更好的分布 核心公式 从另一个容易采样的分布 $q$ 采样，用权重修正偏差：\n$$\\mathbb{E}_p[f(x)] = \\mathbb{E}q\\left[f(x) \\cdot \\frac{p(x)}{q(x)}\\right] \\approx \\frac{1}{N} \\sum{i=1}^{N} f(x_i) \\cdot \\frac{p(x_i)}{q(x_i)}, \\quad x_i \\sim q$$\n其中 $\\frac{p(x)}{q(x)}$ 是重要性权重，补偿\u0026quot;从错误分布采样\u0026quot;带来的偏差。\n直觉 均匀采样（无脑撒点）：大量样本在 f(x)≈0 的区域，对期望贡献为 0，浪费 重要性采样（聪明撒点）：样本集中在 f(x) 大的区域，每个样本都有实质贡献 权重 p/q \u0026lt; 1 修正\u0026#34;采太频繁\u0026#34;带来的偏差，结果仍无偏 两个硬性要求 覆盖性：$q(x) \u0026gt; 0$ 的区域必须覆盖所有 $p(x) \u0026gt; 0$ 的区域，否则权重 $p/q = \\infty$，估计崩掉 形状接近：$q$ 的形状越接近 $f(x) \\cdot p(x)$，权重越接近常数，方差越小 四、常见疑问解答 Q1：为什么 p 分布不好采样，q 就容易采样？ 两件事是独立的：\n能算密度值 p(x)：给定点 x，代入公式算出密度 能从分布采样：随机生成符合该分布的样本 有些分布能算密度但很难采样（复杂后验），有些分布采样极简单（正态分布直接调库）。重要性采样利用的正是这个不对称性：\np：能算密度值，但采样难/代价高 q：能算密度值，且采样容易（正态、均匀等） 权重 p(x)/q(x)：两个密度值相除，直接得到 Q2：为什么不直接用数学公式算？ 公式能算的情况（如 $P(x\u0026gt;4)$ 用误差函数）确实不需要采样 高维积分（1000 维）、无解析解的分布、归一化常数未知的后验——这些情况公式根本算不了 重要性采样针对的正是\u0026quot;公式算不了\u0026quot;的场景 Q3：必须知道 p 的分布吗？ 是的，这是前提。想计算 $\\mathbb{E}_p[f(x)]$，必须知道 $p$ 是什么，否则期望本身就没有定义。重要性采样额外要求的只是：p(x) 的密度值要能逐点计算（不需要能从 p 采样，不需要知道归一化常数）。\n特殊情况：$p(x)$ 只知道未归一化形式 $\\tilde{p}(x) = Z \\cdot p(x)$，用自归一化重要性采样，$Z$ 在分子分母相消：\n$$\\mathbb{E}_p[f(x)] \\approx \\frac{\\sum_i \\tilde{w}_i f(x_i)}{\\sum_i \\tilde{w}_i}, \\quad \\tilde{w}_i = \\frac{\\tilde{p}(x_i)}{q(x_i)}$$\nQ4：q 怎么选？ 选 q 的原则：形状尽量接近被积函数 $f(x) \\cdot p(x)$。\nq 的选择 适合场景 正态分布 被积函数有钟形峰值，最通用、最常用 均匀分布 被积函数较平坦，基础蒙特卡洛 指数分布 正半轴单调衰减的问题 Gamma/Beta 有特定形状峰值的问题 分布本身（如 BRDF） 渲染/物理模拟，消掉被积函数的某一项 混合分布（MIS） 没有单一最优 q 的复杂场景 正态分布最常用，因为灵活（调均值/方差）、易采样、覆盖广，但本质上是\u0026quot;形状匹配\u0026quot;问题。\n五、代码示例：稀有事件估计 import numpy as np from scipy import stats np.random.seed(42) N = 10000 true_prob = 1 - stats.norm.cdf(4) # 约 0.00003167 # 直接采样：几乎采不到 x\u0026gt;4 的样本 samples_direct = np.random.randn(N) estimate_direct = (samples_direct \u0026gt; 4).mean() print(f\u0026#34;直接采样估计: {estimate_direct:.8f}\u0026#34;) # 大概率是 0.00000000 # 重要性采样：用 q = N(4.5, 1)，集中在 x\u0026gt;4 附近 q_dist = stats.norm(loc=4.5, scale=1.0) samples_q = q_dist.rvs(N) f_x = (samples_q \u0026gt; 4).astype(float) # 指示函数 p_x = stats.norm.pdf(samples_q, 0, 1) # 目标分布 N(0,1) 的密度 q_x = q_dist.pdf(samples_q) # 提议分布 N(4.5,1) 的密度 weights = p_x / q_x # 重要性权重 estimate_is = (f_x * weights).mean() print(f\u0026#34;重要性采样估计: {estimate_is:.8f}\u0026#34;) # 约 0.00003167 print(f\u0026#34;真实概率: {true_prob:.8f}\u0026#34;) # 约 0.00003167 print(f\u0026#34;误差: {abs(estimate_is - true_prob) / true_prob * 100:.2f}%\u0026#34;) # 方差对比（重复 200 次） direct_results = [(np.random.randn(N) \u0026gt; 4).mean() for _ in range(200)] is_results = [] for _ in range(200): s = q_dist.rvs(N) w = stats.norm.pdf(s, 0, 1) / q_dist.pdf(s) is_results.append(((s \u0026gt; 4).astype(float) * w).mean()) print(f\u0026#34;直接采样方差: {np.var(direct_results):.2e}\u0026#34;) print(f\u0026#34;重要性采样方差: {np.var(is_results):.2e}\u0026#34;) # 重要性采样方差远小于直接采样 六、多重重要性采样（MIS） 为什么需要多个 q？ 没有一个 q 在所有情况下都最优：\n渲染场景：漫反射地板 + 小点光源 BRDF 采样（q ∝ f_r）：漫反射各向均匀，极少打到小光源 → 噪点多 光源采样（q ∝ 朝向光源）：直接光照准，但对高光材质估计差 → 需要两个 q 互相补充 Balance Heuristic 权重 同时用 $K$ 个策略，每个策略的权重由各自的 pdf 决定：\n$$w_k(\\omega) = \\frac{n_k \\cdot q_k(\\omega)}{\\sum_{j=1}^{K} n_j \\cdot q_j(\\omega)}$$\n直觉：谁更\u0026quot;擅长\u0026quot;采到这个方向（pdf 更大），谁就对这个样本负更多责任：\n方向 ω 朝向光源： q_光源(ω) 很大，q_BRDF(ω) 很小 → w_光源 ≈ 1，w_BRDF ≈ 0 → 这个样本归光源策略负责 方向 ω 朝向高光反射方向： q_BRDF(ω) 很大，q_光源(ω) 很小 → w_BRDF ≈ 1，w_光源 ≈ 0 → 这个样本归 BRDF 策略负责 每个策略只在自己擅长的区域贡献，互相补充不重叠，整体方差比任何单一策略都小。\n七、实际应用场景 离线渲染（Path Tracing） 渲染方程是对半球面的积分，被积函数 = BRDF × 入射光 × cosθ：\n均匀半球采样：大量方向在掠射角（cosθ≈0），贡献为 0，浪费 余弦加权采样（q ∝ cosθ）：消掉 cosθ 项，方差降低 BRDF 采样（q ∝ f_r）：消掉 BRDF 项，高光材质方差极小 MIS（余弦 + 光源）：两者互补，通用场景方差最小 强化学习（PPO） 旧策略 $\\pi_{old}$ 采集数据，用重要性采样复用数据训练新策略 $\\pi_{new}$：\n$$\\mathcal{L} = \\mathbb{E}{\\pi{old}}\\left[\\frac{\\pi_{new}(a|s)}{\\pi_{old}(a|s)} \\cdot A(s,a)\\right]$$\n权重 $\\pi_{new}/\\pi_{old}$ 偏离 1 太多时（新旧策略差异大），PPO 用 clip 截断，保证估计可信。\n大模型 RLHF KL 散度惩罚项 $\\mathbb{E}[\\log(\\pi_{new}/\\pi_{ref})]$ 本质是控制重要性权重不偏离 1 太远，防止新模型偏离参考模型导致重要性采样失效。\n八、核心要点速查 问题 答案 重要性采样解决什么问题？ 积分算不了时，用更好的采样分布降低蒙特卡洛的方差 核心公式是什么？ $\\mathbb{E}_p[f] = \\mathbb{E}_q[f \\cdot p/q]$，权重 $p/q$ 修正偏差 p/q 容易得到吗？ 容易，只需逐点算密度值，不需要从 p 采样 必须知道 p 吗？ 是，这是前提；不知道 p 期望本身就没有定义 q 怎么选？ 形状尽量接近 $f(x) \\cdot p(x)$，正态分布最通用 q 的硬性要求？ q 的支撑集必须覆盖 p 的支撑集，否则权重无穷大 MIS 是什么？ 多个 q 同时采样，Balance Heuristic 让每个样本由最擅长的策略负责 和均匀采样的区别？ 均匀采样是无脑撒点，重要性采样是聪明撒点 + 权重修正 ","permalink":"https://afan.ml/posts/%E9%87%8D%E8%A6%81%E6%80%A7%E9%87%87%E6%A0%B7-importance-sampling/","summary":"\u003cblockquote\u003e\n\u003cp\u003e\u003cstrong\u003e一句话\u003c/strong\u003e：计算期望本质是计算积分，积分算不了就用采样均值代替，均匀采样浪费在无效区域导致方差大，重要性采样换一个\u0026quot;聪明的分布 q\u0026quot;集中采有效区域，再用权重 p/q 修正偏差，保证结果无偏且方差更小。\u003c/p\u003e\n\u003c/blockquote\u003e\n\u003chr\u003e\n\u003ch2 id=\"一为什么需要重要性采样\"\u003e一、为什么需要重要性采样？\u003c/h2\u003e\n\u003ch3 id=\"计算期望--计算积分\"\u003e计算期望 = 计算积分\u003c/h3\u003e\n\u003cp\u003e$$\\mathbb{E}_p[f(x)] = \\int f(x) \\cdot p(x) , dx$$\u003c/p\u003e\n\u003cp\u003e很多场景这个积分\u003cstrong\u003e没有解析解\u003c/strong\u003e，无法直接用数学公式算：\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e\u003cstrong\u003e高维积分\u003c/strong\u003e：神经网络参数空间是几百万维，数值积分根本算不了\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e分布形状复杂\u003c/strong\u003e：多峰分布、奇形怪状的后验分布，没有现成公式\u003c/li\u003e\n\u003cli\u003e\u003cstrong\u003e采样代价高\u003c/strong\u003e：强化学习里从新策略采样需要重新跑环境，代价极高\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e这时只能用\u003cstrong\u003e蒙特卡洛方法\u003c/strong\u003e：采样求均值代替积分。\u003c/p\u003e\n\u003chr\u003e\n\u003ch2 id=\"二蒙特卡洛采样求均值\"\u003e二、蒙特卡洛：采样求均值\u003c/h2\u003e\n\u003cp\u003e从分布 $p$ 采 $N$ 个样本，用样本均值估计期望：\u003c/p\u003e\n\u003cp\u003e$$\\mathbb{E}\u003cem\u003ep[f(x)] \\approx \\frac{1}{N} \\sum\u003c/em\u003e{i=1}^{N} f(x_i), \\quad x_i \\sim p$$\u003c/p\u003e\n\u003cp\u003e\u003cstrong\u003e问题\u003c/strong\u003e：如果直接从 $p$ 采样很难（或代价高），或者均匀采样导致大量样本浪费在 $f(x) \\approx 0$ 的区域，方差就会很大。\u003c/p\u003e\n\u003cpre tabindex=\"0\"\u003e\u003ccode\u003e估算 P(x \u0026gt; 4)，x ~ N(0,1)：\n均匀采样 100 万个点，平均只有 31 个落在 x\u0026gt;4\n→ 99.997% 的样本对结果没有贡献，方差极大\n\u003c/code\u003e\u003c/pre\u003e\u003chr\u003e\n\u003ch2 id=\"三重要性采样换一个更好的分布\"\u003e三、重要性采样：换一个更好的分布\u003c/h2\u003e\n\u003ch3 id=\"核心公式\"\u003e核心公式\u003c/h3\u003e\n\u003cp\u003e从另一个容易采样的分布 $q$ 采样，用权重修正偏差：\u003c/p\u003e\n\u003cp\u003e$$\\mathbb{E}_p[f(x)] = \\mathbb{E}\u003cem\u003eq\\left[f(x) \\cdot \\frac{p(x)}{q(x)}\\right] \\approx \\frac{1}{N} \\sum\u003c/em\u003e{i=1}^{N} f(x_i) \\cdot \\frac{p(x_i)}{q(x_i)}, \\quad x_i \\sim q$$\u003c/p\u003e","title":"重要性采样（Importance Sampling）笔记"},{"content":"👋 大家好！ 这是我的第一篇博客文章！欢迎来到 Afan\u0026rsquo;s Blog。\n🚀 Hugo：世界上最快的静态网站生成器。 🎨 Stack：目前最流行的卡片式博客主题。 🔒 HTTPS：全站安全加密传输。 未来这里会分享更多关于 技术、AI、生活 的点滴。感谢访问！\n","permalink":"https://afan.ml/posts/my-first-post/","summary":"\u003ch2 id=\"-大家好\"\u003e👋 大家好！\u003c/h2\u003e\n\u003cp\u003e这是我的第一篇博客文章！欢迎来到 \u003cstrong\u003eAfan\u0026rsquo;s Blog\u003c/strong\u003e。\u003c/p\u003e\n\u003cul\u003e\n\u003cli\u003e🚀 \u003cstrong\u003eHugo\u003c/strong\u003e：世界上最快的静态网站生成器。\u003c/li\u003e\n\u003cli\u003e🎨 \u003cstrong\u003eStack\u003c/strong\u003e：目前最流行的卡片式博客主题。\u003c/li\u003e\n\u003cli\u003e🔒 \u003cstrong\u003eHTTPS\u003c/strong\u003e：全站安全加密传输。\u003c/li\u003e\n\u003c/ul\u003e\n\u003cp\u003e未来这里会分享更多关于 \u003cstrong\u003e技术、AI、生活\u003c/strong\u003e 的点滴。感谢访问！\u003c/p\u003e","title":"Hello World"}]