【轻量化复刻大模型】线性注意力（Linear Attention）
【轻量化复刻大模型】线性注意力（Linear Attention）
0x00 概要
「线性注意力」不是一个单一论文，而是一条技术路线。核心文献如下：
| 文献 | 标题 | 作者/机构 | 年份 | 会议/arXiv |
|---|---|---|---|---|
| Linear Attention | Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention | Katharopoulos, Vyas, Pappas, Fleuret (EPFL) | 2020 | ICML / arXiv:2006.16236 |
| DeltaNet | Linear Transformers are Secretly Fast Weight Programmers | Schlag, Irie, Schmidhuber (IDSIA) | 2021 | ICML / arXiv:2102.11174 |
| Gated Linear Attention | Gated Linear Attention | Yang, Wang, Shen, Panda, Kim (MIT) | 2023 | NeurIPS / arXiv:2312.06635 |
| Parallel DeltaNet | Parallelizing Linear Transformers with the Delta Rule over Sequence Length | Yang, Wang, Zhang, Shen, Kim (MIT) | 2024 | arXiv:2406.06484 |
| Kimi Linear / KDA v1 | Kimi Linear | Moonshot AI | 2025 | arXiv:2510.26692 |
| KDA v2 / K3 | Kimi K3 | Moonshot AI | 2025 | arXiv:2607.24653 |
注：
- 项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
- 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
- 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
- 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
- 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
- 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（
plot_*.py），可生成静态图。 - 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
- 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。
0x01 演进脉络与定位
1.1 前身与痛点
标准 softmax Attention：Attention(Q,K,V) = softmax(QK^T/√d)V
- 痛点：计算和内存复杂度 O(T²)，T=100 万时 T² 矩阵难以存储。
1.2 同期竞品对比
| 方案名 | 核心机制 | 优势 | 劣势 |
|---|---|---|---|
| FlashAttention | 分块 + online softmax，不物化 O(T²) 矩阵 | 保持 softmax 精度 | 复杂度仍 O(T²)，KV cache 仍线性 |
| 稀疏 Attention | 只计算部分 token 对 | 降低常数 | 破坏全局性，需选择策略 |
| Linear Attention | 核方法替换 softmax，O(T) | 真正的线性复杂度 | 纯加性无法覆写 |
| RetNet / RWKV | 递推状态 | O(T) | 表达力受限 |
1.3 定位一句话
不同于标准 softmax 注意力通过 O(T²) 显式计算所有 token 对的做法，线性注意力通过核函数 + 递推状态（R=φ(K)^T V）将复杂度降到 O(T)，其中 DeltaNet/KDA 进一步用可修正的 Delta Rule 替代纯加性更新，实现了 O(1) 状态下的可覆写记忆。
1.4 演进总结
Linear Attention → 只能加，不能改
↓ 加门控
GLA → 能遗忘，但不能纠错
↓ 加 delta rule
DeltaNet → 能纠错，但遗忘是固定的
↓ 门控 + delta rule
Gated DeltaNet → 数据依赖遗忘 + 纠错
↓ Householder 反射
Parallel DeltaNet → 旋转替代对角衰减（更强表达力）
↓ softplus 门控 + dt_bias + A_log
Kimi Linear (KDA) → 最精细的门控 + 纠错（K3 实际使用）
每一步都是在解决上一步的局限：不能改 → 不能忘 → 遗忘不灵活 → 衰减不够激进。
0x02. 原理讲解
2.1 直觉类比视角
想象一间黑板教室。普通注意力是"老师每讲一句，就把所有旧板书重抄一遍"——随着课越讲越长，重抄的代价越来越大（O(T²)）。线性注意力换了个思路：老师只维护一块固定大小的黑板（状态矩阵 S），每讲一句新知识，就在黑板上加/改几笔（O(1)），不再重抄。
难点在于：如果两句讲的是同一件事但答案不同，怎么改（以 Delta Rule 为例）？纯加性只会越写越糊（混叠）；Delta Rule 则先"看清黑板上的旧答案"（检索），算出"差多少"（误差），再精准擦掉旧答案写下新的（修正）。
2.2 专业水平视角
线性注意力的基本想法，是不再保存所有历史 Token 的完整 Key-Value，而是把历史信息累积到一个固定大小的状态矩阵中。
普通的 Full Attention 在读取历史时，需要访问历史上的许多 Key 和 Value。上下文越长，KV Cache 通常也越大。线性注意力则把历史压缩成一个固定形状的状态 S，之后的新 Token 只需要读取和更新这个状态，而不需要逐个访问全部历史 Token。
最基本的状态更新可以写成：
这里：
k_t是当前 Token 的 Key，用来表示这条信息应该通过什么模式被寻址；v_t是当前 Token 的 Value，表示真正要存储的内容；v_t k_t^\top是一个外积矩阵，把当前的 Key-Value 关联写入状态；S_t是处理完第t个 Token 后的记忆状态。
这个更新具有一个重要特征：当前 Token 的 v_t 和 k_t 在完成外积并累积到状态后，就可以立即丢弃。之后的计算只依赖状态矩阵 S_t，不需要保留随上下文长度增长的完整 KV Cache。不过，前面的短卷积仍然需要保留固定长度的、卷积之前的 qkv 输入。也就是说，线性注意力不代表所有中间状态都完全不需要保存；它主要消除了完整历史 Key-Value 的逐 Token 存储。
输入输出：
- 输入：
Q, K, V ∈ R^{B, H, T, D}（batch, heads, 序列, 每头维度） - 输出：
O ∈ R^{B, H, T, D}
核心推导：softmax 可以近似为核函数的乘积：softmax(q·k^T) ≈ φ(q)·φ(k)^T，则
Attention(Q,K,V) = φ(Q) · (φ(K)^T · V)
= φ(Q) · S 其中 S = φ(K)^T · V ∈ R^{D×D} 固定大小
写成递推：S_t = S_{t-1} + φ(k_t)^T v_t（纯加法），o_t = φ(q_t) S_t
与标准做法的公式级对比：
| 标准 | 线性 | |
|---|---|---|
| 计算 | softmax(QK^T/√d)V | φ(Q)(φ(K)^T V) |
| 复杂度 | O(T²D) | O(TD²) |
| 状态 | 无（每步重算） | S ∈ R^{D×D} 递推 |
2.3 讨论
- 为什么理论上 work：线性注意力把注意力重写为核方法的线性形式，借助矩阵结合律把"先算 QK^T 再乘 V"改为"先算 φ(K)^T V 再乘 φ(Q)"，从而利用结合律消除对 T 的二次依赖。这揭示了 Transformer 推理时等价于一个 RNN（带 D×D 隐藏状态）。
- 为什么纯加性不 work：核方法假设单调记忆（新信息不影响旧信息表示），导致状态 S 里新旧信息混叠，无法覆写。
- 论文消融结论：DeltaNet 证明用 Delta Rule 替代加法可显著提升合成检索任务的精确性；K3 论文进一步论证了下界衰减 + 全秩门控在数值稳定性与表达力上的优势。
- 未解决理论问题：D×D 状态的信息容量上限（多少 token 可无损压缩进固定 S）仍是开放问题；如何衡量"检索精度 vs 状态容量"的权衡也未被严格刻画。
0x03. 子知识点拆解
3.1 树状结构
线性注意力
├── [1] 核特征映射 φ(x) = ELU(x)+1 (基础，所有变体依赖)
├── [2] 纯加性 Linear Attention (依赖 [1])
├── [3] Delta Rule 更新 (改进 [2])
│ └── [3a] RETRIEVE-COMPARE-CORRECT 三段式
├── [4] Gated Linear Attention (GLA) (改进 [2] 加门控)
├── [5] Gated DeltaNet (合并 [3]+[4] + L2归一化)
├── [6] Parallel DeltaNet (改进 [3] 用 Householder)
└── [7] KDA (Kimi Delta Attention) (工程化 [5]，下界衰减+全秩门控)
├── [7a] 下界衰减 g_min=-5
├── [7b] 全秩输出门控
└── [7c] ShortConv + SiLU + L2 归一化投影链
3.2 代码映射
| 子知识点 | 代码位置 | 说明 |
|---|---|---|
| [1] 核特征映射 | core.py (elu_plus_one) |
φ(x)=ELU(x)+1 |
| [2] 纯加性 | variants.py (linear_attn_step) |
S += φ(k)v^T |
| [2] chunkwise | variants.py (linear_attn_chunkwise) |
全序列前向（demo 为循环展开，非并行） |
| [3] Delta Rule | variants.py (delta_net_step) |
RETRIEVE-COMPARE-CORRECT |
| [4] GLA | variants.py (gla_step) |
门控加性 |
| [5] Gated DeltaNet | variants.py (gated_delta_net_step) |
门控+Delta+L2 |
| [6] Parallel DeltaNet | variants.py (parallel_delta_net_step) |
Householder |
| [6] Householder 工具 | core.py |
反射实现 |
| [7] KDA v1 | variants.py (kimi_linear_step) |
softplus 门控, 默认 version="v1" |
| [7a] 下界衰减 | variants.py (version="v2" 分支) |
g_min·sigmoid, v1 默认/v2 可选 |
| [7] 基类投影 | core.py (LinearAttentionBase) |
q/k/v 投影 |
| [4] GLA 层封装 | layers.py (GatedLinearAttentionLayer) |
gate_proj 从 x 投影出 gate |
| [5] GatedDeltaNet 层封装 | layers.py (GatedDeltaNetLayer) |
低秩 gate_logit + beta_proj |
0x04. 子知识点深度分析
4.0 核心概念
4.0.1 状态矩阵 S
最基本的线性注意力可以写成：
普通线性层的投影参数在推理过程中通常保持不变，而 S_t 会随着每个 Token 的到来持续改变。
S 是所有方法共享的同一个抽象：这是一个 (B, H, Dk, Dv) 的大小固定的"记忆矩阵"，替代了标准 Transformer 中随序列长度增长的 KV Cache。每一步的操作都是：读 S → 改 S → 返回输出。
S_t 又像一个普通线性层的权重矩阵：它通过矩阵乘法作用在 Query 上。因此，从 fast-weight 的角度看，S_t 可以理解为：
- 一组随上下文快速更新的权重；
- 一个动态的 Key-Value 关联记忆；
- 一个在推理过程中在线变化的线性模型。
如果初始状态为：
把递推展开可得：
如果 S_0 不为零，还需要在右侧加上初始状态 S_0。
用 Query 读取历史
将状态乘以任意 Query q：
其中：
表示第 i 个历史 Key 与当前 Query 的匹配程度。匹配程度决定对应 Value 被取出的比例。因此：
v_i k_i^\top可以看成向状态中写入一条 Key-Value 关联；S_tq可以看成根据 Query 与历史 Key 的匹配程度，对历史 Value 做加权聚合。
用当前 Key 查询旧记忆
既然状态可以接受 Query，也可以把某个 Key 直接输入状态：
它表示当前关联记忆根据 Key k 召回的 Value。
因此，在写入当前关联 (k_t,v_t) 之前，可以先计算：
这个量可以解释为：旧记忆对当前 Value 的重建、预测或近似。
Delta Rule 的本质，就是比较这个旧预测和当前真正目标 v_t 之间的差异，然后只更新没有预测正确的部分。
4.0.2 特征映射函数
elu_plus_one 是特征映射函数，其核心是 elu(x) + 1 → 保证输出 > 0。
为什么需要正数？因为 Linear Attention 把 softmax(QK^T)V 拆成了 φ(Q) · [φ(K)^T V]，中间的 φ(K)^T V 就是状态矩阵 S。如果 φ 输出有负数，S 中的正负会互相抵消，导致信息丢失。ELU+1 保证所有特征值为正，S 只做加法累积。
4.0.3 linear_attn_step - 纯加法更新
以下是最简单的线性注意力实现。
# S_new = S + φ(k) ⊗ v ← 只加不删
# out = φ(q) · S_new ← 读出
# ============================================================
# 1. Linear Attention (Katharopoulos et al., ICML 2020)
# Core idea: S_t = S_{t-1} + phi(k_t) * v_t^T
# Pure additive update, cannot overwrite
# ============================================================
def linear_attn_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Single step of Linear Attention (recurrent form).
S shape: (B, H, Dk, Dv) — state matrix
q shape: (B, H, Dk) — query
k shape: (B, H, Dk) — key
v shape: (B, H, Dv) — value
Returns: (output, new_S)
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
S_new = S + torch.einsum('bhd,bhv->bhdv', phi_k, v) # pure additive
out = torch.einsum('bhd,bhdv->bhv', phi_q, S_new)
return out, S_new
直觉：一个只能往里塞东西、永远不扔东西的笔记本。每来一个 (k,v)，就把 φ(k) × v 叠加到 S 上。查询时用 φ(q) 去匹配。
问题：只能叠加，不能覆盖。如果两个不同的 v 对应相似的 k，旧值永远不会被修正。
4.0.4 linear_attn_chunkwise - 批量训练形式
linear_attn_chunkwise 是 linear_attn_step 的循环展开版：等价地逐 token 调用 step 函数、收集全部输出，数值上与串行递推完全一致。
注意：这份 demo 代码本身并不并行——它只是把"一条 T 长序列的递推"一次性跑完，方便训练时拿到所有位置的输出。名字里的 "chunkwise（分块）"指的是并行化策略：把序列按时间切成若干 chunk，chunk 内部用矩阵运算并行（纯加法用 cumsum 前缀和即可；带 Delta Rule/门控这类递推依赖的用 UT/WY 变换），chunk 之间才做串行递推。真正的 chunk 并行实现见 （kda_chunkwise）；
此处用循环展开，是为了数值等价、便于用测试验证 "chunkwise == recurrent"。
def linear_attn_chunkwise(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
"""
Chunkwise form of Linear Attention (parallel training form).
q, k, v shape: (B, H, T, D)
Returns: output (B, H, T, D)
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
# S = phi(K)^T @ V (prefix sum for causality)
S = torch.einsum('bhkd,bhvd->bhdv', phi_k, v)
# But we need causal: o_t = phi(q_t) @ S_{:t}
# We use cumulative sum to get each prefix
S_cum = phi_k.transpose(-2, -1) @ v # (B, H, Dk, T) @ (B, H, T, Dv) -> (B, H, Dk, Dv)
# Actually simpler: just do it step by step for demo
B, H, T, Dk = q.shape
Dv = v.shape[-1]
S = torch.zeros(B, H, Dk, Dv, device=q.device)
outputs = []
for t in range(T):
out, S = linear_attn_step(S, q[:, :, t], k[:, :, t], v[:, :, t])
outputs.append(out)
return torch.stack(outputs, dim=2)
4.0.5 讨论
Q：为什么训练时用 chunkwise、推理时用 step 版？
这背后是线性注意力「推理等价 RNN」带来的两种计算形态，本质区别在于「这次前向需要多少个位置的输出」：
- 训练要"所有位置"的输出：训练时手头有完整输入序列（T 个 token 全部已知），而语言建模损失需要对每个位置 t 都算出一个输出 o_t 去预测下一个 token。所以训练走「全序列前向」——把 T 步一次性跑完、收集 o_1..o_T，才够算整条序列的 loss。
linear_attn_chunkwise就是这 "全序列前向" 的实现（这份演示代码为了等价直观直接循环展开；生产实现会用 chunk 并行 + UT/WY 变换把 chunk 内并行掉）。 - 推理只要"这一步"的输出 + 状态递推：推理/解码是自回归的，每次只生成一个 token，且它只能来自上一轮的输出（顺序拿到）。这时只需要当前这一步的输出 o_t 和更新后的状态 S_t，把它传给下一步即可，没有也不需要历史 T 步的输出。所以用 step 版：输入
(S_{t-1}, q_t, k_t, v_t)，输出(o_t, S_t)。 - 一句话：训练要"所有位置的输出"（全序列前向）；推理只要"这一步的输出 + 状态递延"（单步前向）。
这不是线性注意力的特例，而是所有可递推序列模型（LSTM / RWKV / RetNet / KDA）的通用范式：训练走全量前向（并行化机会大），推理走单步自循环（状态跨 token 传递）。chunkwise 正是回答「有 Delta Rule 递推依赖的模型，训练时怎么并行」的桥梁——纯加法可用 cumsum 前缀和并行，而带修正/门控的递推要靠分块并行来解决。
4.0.6 朴素线性注意力的缺陷：只会累加
最简单的线性注意力只有累加，没有遗忘，也没有覆盖。
如果同一个 Key 被写入两次，那么从这个 Key 读取时，得到的是两次 Value 的和，而不是第二次写入的 Value。换句话说，系统无法判断：
- 这条信息是不是已经存在；
- 新信息是不是旧信息的修正；
- 旧信息是不是已经过期；
- 当前写入是否应该覆盖以前的内容。
因此，写错的信息和已经过期的信息都会继续留在固定容量的状态矩阵里。随着 Token 不断到来，状态可能逐渐饱和；如果 Key 的范数较大，累加的数值甚至可能发生爆炸。
线性注意力虽然把历史压缩成了固定状态，但“固定大小”也意味着它必须在有限容量里不断处理新的信息。如果没有任何选择性更新机制，它只能盲目地把所有内容叠加进去。
4.1 Delta Rule
DeltaNet 的核心改进是：写入之前，先使用当前 Key 查询旧状态，看看旧状态已经记住了什么（旧状态对当前 Value 的预测）；然后只写入“目标 Value 和旧状态预测结果之间的差（预测残差）”，从而避免同一关联被无条件重复累加。
4.1.1 直觉思路
一个会"纠错"的笔记本。来新 (k,v) 时，先用 k 查一下旧记忆里存了什么，如果查出来的和真实 v 不一样，只把差值写回去。即，像记笔记时"用修正液改错"——只涂掉写错的那几笔，不动其他内容。
4.1.2 专业角度
\(S_t = Diag(alpha)*S_{t-1} + beta*k_t*(v_t - S_{t-1}^T k_t)^T\)。检索 \(S_{t-1}^T k_t\) 是记忆对当前 key 的预测值，v_t - 预测 是误差，beta*k 外积误差 只修正与当前 key 相关的行。
这是 Widrow-Hoff (LMS) 算法的引入——把状态 S 看作线性联想记忆，每次看到 (k,v) 对就用 LMS 规则更新，只修错误部分，因此能精准覆写而非常规加法混叠。
公式解读
旧状态对当前 Key 的读取结果是：
它可以理解为：旧记忆根据当前寻址模式 k_t，对当前 Value 做出的预测。
当前真正想写入的是 v_t，因此写入残差为：
这个残差表示旧状态还没有记住的部分。如果旧状态已经能够正确预测当前 Value，那么残差就是零，也就不需要重复写入。
DeltaNet 引入写门：
在本文的描述中，每个 Head 使用一个标量写门，写门由当前 Hidden State 计算得到。它控制本次残差写入的比例。
完整更新式为：
这个式子的含义可以分成三步：
- 用旧状态读取当前 Key，得到旧状态对 Value 的预测；
- 计算目标 Value 与旧预测之间的差；
- 把这部分差值沿当前 Key 的方向写回状态，并由
β_t控制写入强度。
右乘形式
将上式整理，可以得到：
这个形式更容易看出状态发生了什么变化。
其中：
β_t v_t k_t^T是当前 Value 的写入；I-\beta_t k_tk_t^T是对旧状态的变换。
如果 k_t 做了 L2 归一化，那么这个变换会让旧状态在当前 Key 方向上按 (1-β_t) 的比例缩小，而在与当前 Key 正交的方向上保持不变。
因此，DeltaNet 不再像朴素线性注意力那样无条件累加，而是在当前 Key 方向上执行一种平滑修正：
β_t较小，表示只修正一部分；β_t较大，表示更积极地覆盖；- 当
β_t=1且 Key 已归一化时，可以在当前 Key 方向上实现精确覆写。
由于 Key 已经归一化，状态转移不会因为 Key 范数过大而无条件放大旧状态。因此，它比朴素线性注意力更不容易出现状态爆炸。
损失函数与 Delta Rule
DeltaNet 把状态 S 看成一个从 Key 预测 Value 的线性模型。
给定当前样本 (k_t,v_t)，模型的预测是：
目标是让预测尽可能接近真实 Value v_t。因此，当前样本的回归损失为：
对 S 求梯度：
以 β_t 作为当前样本的学习率，在旧状态 S_{t-1} 上执行一步梯度下降：
代入梯度后得到：
这正是 Delta Rule。
因此，Delta Rule 可以理解为一次在线学习：
- 把当前 Key 当作输入；
- 用状态矩阵预测当前 Value；
- 计算预测误差；
- 用写门
β_t控制学习率； - 把误差写回状态。
如果：
说明旧记忆已经正确预测当前 Value，此时残差为零，不需要重复写入。
4.1.3 对比基线
| 维度 | 纯加性 (Linear) | Delta Rule |
|---|---|---|
| 更新 | S += k v^T | S += beta*k*(v - S^T k)^T |
| 覆写能力 | 无 (混叠) | 有 (只修错误) |
| 精确检索 | 差 | 好 (近因优先) |
4.1.4 讨论（反方视角）
下面是一些相关讨论和思考。
- a) 最强反对：Delta Rule 引入了序列依赖循环，无法像纯加法那样做前缀和并行化，训练变慢。
- b) 若失败，最可能错在：假设误差能线性传播修正，但多头、非线性下该假设不严格成立。
- c) 胜出条件：需要精确 key-value 检索的任务 (合成检索、指代消解)；失败条件：极长序列 + 强非线性。
- d) 更稳健做法：GLA (门控加法)——放弃纠错换可并行性。
- e) 受影响方：训练方 (循环难并行)、推理方 (需状态回滚)。
4.1.5 计算复杂度
| 项 | 纯加性 | Delta Rule |
|---|---|---|
| 状态更新 FLOPs/token | O(D^2) | O(D^2) (多一次 matvec) |
| 序列复杂度 | O(T) | O(T) (但循环不可并行) |
| 状态内存 | O(D^2) | O(D^2) |
4.1.6 参数敏感性
| 超参 | 默认 | 调大 | 调小 |
|---|---|---|---|
| alpha (decay) | 0.9 | 保留更多旧信息，但混叠增加 | 遗忘快，短期记忆 |
| beta (write) | 0.8 | 修正更激进，可能震荡 | 修正温和，收敛慢 |
alpha 对检索精度的影响如下图。具体是 固定 beta，扫 alpha（衰减因子）：
| alpha | 检索误差 | 解读 |
|---|---|---|
| 0.00 | 0.641 | 全遗忘：每步丢弃历史，误差最大 |
| 0.30 | 0.468 | |
| 0.60 | 0.333 | |
| 0.80 | 0.200 | |
| 0.95 | 0.059 | |
| 0.99 | 0.012 | 近无损保留，误差最小 |
关键推论：误差随 alpha 单调下降，在这个「同 key 重复写入」的合成任务上 alpha 越大越好。
但这不意味着实际应该取 alpha→1——真实长序列下 alpha→1 会导致状态无界累积、不同 key 之间开始混叠，这正是 KDA 引入 g_min = -5.0 下界衰减（等价 alpha 上界 ≈ exp(-ε)）来强制一个最小遗忘量的原因。
4.1.7 工程实现
# variants.py delta_net_step
# RETRIEVE → COMPARE → CORRECT (关键创新)
# ① RETRIEVE: retrieved = k · S ← 用 k 查旧记忆
# ② COMPARE: error = v − retrieved ← 算出"差多少"
# ③ CORRECT: S_new = α·S + β·(k ⊗ error) ← 只修正误差部分
# ④ READ: out = q · S_new ← 用 q 读出
# • α (decay)：旧记忆按比例遗忘，防止无限累积
# • β (write strength)：修正力度，太大可能震荡，太小收敛慢
# 关键区别 vs Linear Attention：能覆盖旧信息，因为 error = v − retrieved 允许"减去旧的、加上新的"
retrieved = einsum('bhd,bhdv->bhv', k, S) # RETRIEVE: S^T @ k
error = v - retrieved # COMPARE: 误差
correction = beta * einsum('bhd,bhv->bhdv', k, error) # CORRECT
S_new = alpha * S + correction # 衰减 + 修正
out = einsum('bhd,bhdv->bhv', q, S_new) # READ
4.1.8 探讨代码（验证反方 a：并行性损失）
可运行脚本：experiments\models\llm\linear_attn\discussion.py（第一段实验）。
思路：纯加性 S = Σ k_i v_i^T 满足结合律，因此「一次 matmul 并行」与「逐 token 串行」必须得到完全相同的结果；而 DeltaNet 的 S_t 依赖 S_{t-1}，无法用 cumsum 前缀和展开。
实测输出：
============================================================
[纯加性] 前缀和(并行, matmul) vs 逐token(串行) 最大差: 4.77e-07
-> 若 diff≈0, 证明纯加性可用一次 matmul 并行算出所有前缀
[DeltaNet] S_t 依赖 S_{t-1} (循环依赖)
-> 无法用 cumsum 前缀和, 只能串行递推 (时间复杂度 O(T), 墙钟更长)
============================================================
[反方 a] KDA 固定状态的信息容量瓶颈
============================================================
写入值 | 检索 S[0,0] | 说明
---------------------------------------------
+0.9 | +0.450 | 朝新值方向移动, 旧值被冲淡
+0.5 | +0.284 | 朝新值方向移动, 旧值被冲淡
+0.1 | +0.072 | 朝新值方向移动, 旧值被冲淡
-0.3 | -0.145 | 朝新值方向移动, 旧值被冲淡
+0.7 | +0.339 | 朝新值方向移动, 旧值被冲淡
结论: 固定 D×D 状态只记住最近的写入, 同一 key 的旧 value 无法保留,
验证了'状态容量有限'的反方观点 (这正是 K3 混 3 KDA + 1 MLA 的原因).
结论：讨论 a 成立。这正是 Parallel DeltaNet（arXiv:2406.06484）用 Householder / WY 表示做 chunkwise 并行的动机——不是为了降复杂度，而是为了恢复可并行性。
4.1.9 DeltaNet 仍然存在的局限
DeltaNet 会保留与当前 Key 正交的方向，但这些“方向”不一定等同于单独的特征通道。如果某些方向长期没有被新的 Key 更新，那么其中的旧关联就可能一直保留下来。即使这些内容已经过期，也没有一个独立的机制主动把它们清掉。
因此，DeltaNet 具备残差覆写能力，但它的遗忘仍然依赖当前 Key 的方向。它不能独立地说“这一部分记忆现在应该过期”，也不能对所有记忆维度进行主动、选择性的衰减。
4.2 Gated Linear Attention (GLA)
原理如下：
-
直觉：在 Linear Attention 基础上加了一个"橡皮擦"。gate 是从输入学来的，每个维度独立决定"保留多少旧记忆"。但没有 delta rule，仍然不能纠错——只能选择遗忘或保留，不能定向修正。
-
专家：\(S_t = Diag(σ(gate_t)) · S_{t-1} + φ(k_t)·v_t^T\)，
σ把门控压到 (0,1)，逐通道缩放旧记忆。gate 是数据依赖的逐通道遗忘因子，等价给每个通道一个可学习的"寿命"；代价是只能整体缩放、无法定向改写具体条目（无纠错能力）。
S_new = σ(gate) ⊗ S + φ(k) ⊗ v ← gate 控制遗忘
out = φ(q) · S_new
具体代码如下：
# ============================================================
# 3. Gated Linear Attention (Yang et al., 2023)
# Core idea: S_t = G_t * S_{t-1} + phi(k_t) * v_t^T
# Data-dependent gating controls forget/remember
# ============================================================
def gla_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
gate: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Single step of Gated Linear Attention.
gate (B, H, Dk) — per-dimension forget gate in (0, 1), learned from input
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
# Data-dependent gating: gate controls what to forget
gate = torch.sigmoid(gate) # ensure in (0, 1)
S_new = gate.unsqueeze(-1) * S + torch.einsum('bhd,bhv->bhdv', phi_k, v)
out = torch.einsum('bhd,bhdv->bhv', phi_q, S_new)
return out, S_new
4.3 Gated DeltaNet：在 Delta Rule 前增加遗忘门
GDN 在 Delta Rule 之前加入每个 Head 的标量遗忘门，让所有状态统一衰减，同时保留 Chunkwise 并行实现的可能性。
4.3.1 原理
直觉角度
DeltaNet + GLA 的合体。既有数据依赖的门控遗忘 (③)，又有 delta rule 的定向纠错 (④⑤⑥)。
与 SGLang 对齐的关键：步骤③的顺序—先衰减 S，再用衰减后的 S 算 error。这和"先算 error 再衰减"数值不同。
公式解读
在执行 Delta Rule 写入之前，GDN 先让整张状态矩阵乘上一个标量遗忘门：
然后使用衰减后的状态计算预测误差并写入：
这里：
α_t是遗忘门；β_t是写门；α_t和β_t都是每个 Head 的标量；α_t会广播到状态矩阵的所有行和所有列。
这意味着每一步都会先对整张状态矩阵做统一衰减：所有方向上的记忆都乘以同一个 α_t。随后，模型再根据当前 Key 和 Value 做一次残差修正。
伪代码实现
在 DeltaNet 基础上引入输入依赖门控：alpha_t, beta_t 由当前输入 x_t 经 sigmoid 生成，从而每个 token 自适应决定遗忘/写入强度。加入 L2 归一化稳定训练。
# gated_delta_net_step − 门控 + Delta Rule + L2归一化
# ① L2归一化： q, k ← q/||q||, k/||k|| ← 稳定训练
# ② 缩放： q ← q * Dk^−0.5 ← 对齐 softmax 注意力的 1/√d
# ③ 门控衰减： S ← σ(gate) ⊗ S ← 数据依赖的遗忘
# ④ 检索： retrieved = k · S ← 用衰减后的 S 查
# ⑤ 纠错： error = v − retrieved
# ⑥ 修正： S_new = S + σ(β)·(k ⊗ error)
# ⑦ 读出： out = q · S_new
4.3.2 对比基线
| 维度 | DeltaNet | Gated DeltaNet |
|---|---|---|
| 门控 | 固定 alpha, beta | 输入依赖 alpha_t, beta_t |
| 归一化 | 无 | L2 归一化 Q/K |
| 读出缩放 | 无 | *(Dk^-0.5) |
4.3.3 讨论
- a) 门控是一刀切：per-channel 门控虽比 per-head 精细，但仍无法表达跨通道的信息重组。
- b) 若失败，错在门控可学习到正确的遗忘策略——实际可能学到次优策略。
- c) 胜出：数据有明显的新旧更替模式；失败：需要通道间信息旋转的任务。
- d) 更稳健：KDA 的下界衰减 (防溢出)。
4.3.4 工程实现 + SGLang 对齐
# variants.py gated_delta_net_step
q = l2_normalize(q); k = l2_normalize(k) # L2 归一化
q = q * (Dk ** -0.5) # 读出缩放 (对齐SGLang)
alpha = sigmoid(gate_logit) # per-channel 门控
beta = sigmoid(beta_logit).unsqueeze(-1)
S = alpha.unsqueeze(-1) * S # 1 先衰减 (对齐SGLang时序)
retrieved = einsum('bhd,bhdv->bhv', k, S) # 2 用衰减后状态检索
error = v - retrieved
S_new = S + beta * einsum('bhd,bhv->bhdv', k, error) # 3 修正
out = einsum('bhd,bhdv->bhv', q, S_new) # 4 读出
已按 SGLang fused_recurrent.py 修正为先衰减再检索时序，并加 q-scale，与 SGLang 逐位一致 (数值验证 max diff = 0.0)。
4.3.5 缺陷
GDN 比 DeltaNet 多了主动遗忘能力，但它的表达能力仍有限。因为同一个 Head 的所有特征通道共享同一个 α_t，所以无论某个通道中的内容是否重要，它都会以相同比例衰减。
这自然引出了一个问题：如果一个标量遗忘门不够细，是否可以把 α_t 从标量扩展成 d_k 维向量，让不同 Key 通道拥有各自的遗忘速度？
4.4 parallel_delta_net_step
− Householder 反射替代对角衰减
① L2归一化： q, k
② 反射： H·S = S − 2·u·(u^T·S) ← Householder 旋转状态空间
③ 检索+纠错：同 delta_net_step
④ 读出： out = q · S_new
直觉：把"对角衰减"（每个维度独立遗忘）换成"旋转"（维度之间混合）。Householder 反射 H = I − 2uu^T 是一个保范变换，可以旋转状态空间而不缩小它。理论上表达力更强，但计算量也更大。
4.4.1 Householder 反射
几何直觉
想象状态矩阵 S 的每一列是 Dv 维空间中的一个向量。对角衰减 a ⊗ S 的效果是：把每个坐标轴独立地缩短—就像把一个方块压扁成一个更小的方块。Householder 反射的效果是：关于一个镜面翻转—形状不变，只是方向变了。
Householder 矩阵： `H = I − 2·u·u^T`
操作： `S' = H·S = S − 2·u·(u^T·S)`
逐步拆解
`u^T·S` → 把 S 投影到 u 方向上，得到一个标量（每列）
`2·u·(u^T·S)` → 沿 u 方向的分量，翻倍
`S − 2·u·(u^T·S)` → 减去这个翻倍分量 = 关于垂直于 u 的超平面做镜像翻转
举个 2D 例子
设 u = [1, 0] (x 轴方向)
H = I − 2·[1,0]·[1,0]^T = [[1,0],[0,1]] − 2·[[1,0],[0,0]] = [[-1,0],[0,1]]
对一个点 [3, 5]：H·[3, 5] = [-3, 5] ← x 分量翻转，y 不变
这就是关于 y 轴的镜像！
为什么比对角衰减好？
对角衰减 a⊗S |
Householder H·S |
|---|---|
| 每个维度独立缩放 | 维度之间混合 |
| 只能缩小(a<1) 或保持 (a=1) | 可以旋转，把信息从一个维度搬到另一个 |
| 表达力有限 | 一个反射 = 一次旋转，多个反射可以逼近任意正交变换 |
直觉：对角衰减像"把海绵挤干"—只能压缩。Householder 像"拧海绵"—可以旋转、翻转，把信息重新排列而不是简单丢弃。
4.4.2 实际局限
parallel_delta_net_step 中只用了一个 Householder 向量 u，所以只做一次反射。理论上可以堆叠多个 u 做多次反射（类似 QR 分解），但计算量会显著增加。这也是为什么 K3 最终选择了 KDA 而非 Parallel DeltaNet—KDA 的 softplus 门控在实用中已经足够灵活，且计算更高效。
4.5 KDA
KDA 把遗忘门扩展到每个 Head 的每个 Key 通道，让不同记忆通道拥有独立遗忘速度。
4.5.1 原理
KDA 将 Gated DeltaNet 工程化落地：下界衰减防溢出、全秩门控增强表达力、ShortConv+SiLU+L2 投影链。
4.5.2 两个版本
| 名称 | 位置 | 本质 |
|---|---|---|
| KDA v1 | variants.py 的 kimi_linear_step(version="v1") |
softplus 无界门控（Kimi Linear 论文原始） |
| KDA v2 | variants.py 的 kimi_linear_step(version="v2") |
sigmoid 有界门控（K3 报告最终形态） |
| kda_step | kda_ref.py 的 kda_step |
v2 的生产版单步实现（不是新版本） |
KDA 的门控在推理/串行递推与训练/chunk 并行两种形态下采用两种实现，分别对应不同论文阶段：
-
v1（Kimi Linear 原版）：门控用无界 softplus：
g = -exp(A_log) * softplus(gate_logit + dt_bias)g 无下界（可以非常负），alpha = exp(g) 可以任意接近 0。这是 Kimi Linear 技术报告中的原始形式，也是 SGLang 的默认路径。
-
v2（K3 有界版）：门控用有界 sigmoid：
g = g_min * sigmoid(exp(A_log) * (gate_logit + dt_bias))g 严格在
(g_min, 0)之间，因此 alpha 在(exp(g_min), 1)范围内。这是 K3 技术报告中的最终形态，也是本项目kimi_k3模块实际使用的版本。
两者都已在项目中实现，由 kimi_linear_step(..., version="v1"/"v2") 切换。kimi_linear_step 默认（也是 SGLang 默认）是 KDA v1；v2 是 K3 最终采用的"安全下界"版本，下方两版对比给出了各自的取舍。
数值对比（A_log=5，z=[50, -50]）：
| 版本 | g（强遗忘方向） | alpha（强遗忘方向） | g（强保留方向） | alpha（强保留方向） |
|---|---|---|---|---|
| v1 | -7420.7 | ~0.0（下溢） | ~0.0 | 1.0 |
| v2 | -5.0 | 0.00674 | ~0.0 | 1.0 |
v1 在强遗忘方向 g 极负、alpha 下溢到 0，这就是 chunkwise 并行时 1/gamma 上溢的根源。v2 则始终保持在数值安全范围内。
实测状态差异：在非零初始状态下，v1 与 v2 的输出 max diff 可达 0.19（数值量级差异明显），说明两个版本不仅是形式不同，实际行为也有显著差异。
v2 和 kda_step 的区别只是工程封装：
kimi_linear_step输入 gate_logit，自己算 g；kda_step输入预计算好的 g（上层log_decay算）。kimi_linear_step对 q 乘了1/√Dk的 scale（对齐 SGLang fused_recurrent）；kda_step不加（尺度由 HeadRMSNorm + 输出门吸收）。kda_step是 v2 的"单步 decode 形态"，即只做一次递推（不是整段序列），不是新版本。- 状态更新的数学本质完全一样（实测 2.2e-16 机器精度一致）。
4.5.3 对比基线
| 维度 | Gated DeltaNet | Kimi Linear (KDA)：两个版本 |
|---|---|---|
| 门控公式 | sigmoid(gate_logit) | v1: -exp(A)·softplus(z+dt_bias)；v2: g_min·sigmoid(exp(A)·(z+dt_bias)) |
| 衰减因子 a 的值域 | (0,1)，但实际只能学到 >0.01 | v1 可任意接近 0；v2 有界 [exp(g_min),1] |
| 遗忘极限 | a→0 时 sigmoid 梯度消失，学不到彻底遗忘 | v1 softplus 无上界、梯度不消失，可学彻底遗忘；v2 有界保证 chunkwise 数值安全 |
| 门控粒度与旋钮 | per (head, channel)，仅 gate_logit 一个旋钮 | per (head, channel) + 每通道偏置 dt_bias + 全局尺度 A_log，三个旋钮 |
| 投影链 | q/k/v 直接 | ShortConv+SiLU+L2 |
下图给出了4 种变体的状态检索误差随 token 步数变化。横轴 token 步数，纵轴末步检索误差 |S^T k - v|：
| 变体 | 末步误差 | 解读 |
|---|---|---|
| Linear（纯加性） | 1.000 | 完全无法检索——新旧信息混叠，误差不收敛 |
| DeltaNet | 0.111 | 最准：无门控干扰，纯 LMS 纠错 |
| GatedDeltaNet | 0.288 | 门控引入遗忘，牺牲部分精度换自适应能力 |
| KDA v1 | 0.365 | 下界衰减 + softplus 门控，精度再降但数值最稳 |
关键推论：误差从 DeltaNet 到 KDA 单调上升，说明工程稳定性是用检索精度换来的。这也解释了为什么 K3 不能纯用 KDA，必须配 MLA。
4.5.4 讨论
- a) 最强反对：O(D^2) 状态矩阵对任意 token 数都固定，信息容量有限，长上下文会记不住细节。
- b) 若失败，最可能错在下界衰减足够防溢出——若 A_h 学得过大，累积仍可能出问题。
- c) 胜出：1M token 上下文的 agent 场景 (固定状态省内存)；失败：需要精确长程点对点检索的任务。
- d) 更稳健：混合架构 (K3 用 3 KDA + 1 MLA, MLA 负责精确检索)。
4.5.5 探讨代码（讨论 a：信息容量瓶颈）
思路：构造一个重复相同 key 但不同 value 的序列，观察 KDA 状态 S 能否同时保住多个 value——由于 Delta Rule 会覆写同一 key 对应的那一行，较早的 value 会被冲掉。
import torch
from models.llm.linear_attn.variants import kimi_linear_step
B, H, Dk, Dv = 1, 1, 4, 1
S = torch.zeros(B, H, Dk, Dv)
ks = torch.tensor([[[1., 0, 0, 0]]]) # 始终是同一个 key
for v in [0.9, 0.5, 0.1, -0.3, 0.7]:
vs = torch.tensor([[[v]]])
out, S, _, _ = kimi_linear_step(
S, torch.zeros(B, H, Dk), ks, vs,
torch.zeros(B, H, Dk), torch.zeros(B, H), torch.ones(H))
print(f'写入 v={v:+.1f} -> 检索 S[0,0] = {float(S[0, 0, 0, 0]):+.3f}')
实测输出：
写入值 | 检索 S[0,0] | 说明
---------------------------------------------
+0.9 | +0.450 | 朝新值方向移动, 旧值被冲淡
+0.5 | +0.284 | 朝新值方向移动, 旧值被冲淡
+0.1 | +0.072 | 朝新值方向移动, 旧值被冲淡
-0.3 | -0.145 | 朝新值方向移动, 旧值被冲淡
+0.7 | +0.339 | 朝新值方向移动, 旧值被冲淡
解读（逐步推导）：
l2_normalize(k)后 k 是单位向量，写入强度beta = sigmoid(0) = 0.5；- 每步先衰减
S *= exp(g)（g<0），再用衰减后的 S 检索，误差v - retrieved； - 因此 S 每步只朝新 value 移动约 50%，从不精确等于任何一个写入过的 value；
- 更关键的是：同一 key 只占一行，第 2 次写入后第 1 次的 value 已无法恢复。
结论：反方 a 成立——固定 D×D 状态只记得住"最近写入"的近似值。这正是 K3 采用 3 KDA + 1 MLA 混合架构的原因：把需要精确点对点检索的活交给 MLA。
下图给出了信息容量瓶颈（同一 key 旧 value 被覆写）。写入同一 key 的一串不同 value，画出「真实值」与「检索值」两条曲线：检索值只是跟随真实值移动（每步约走一半距离），从不重合——状态只记得住最近一次写入，旧 value 不可恢复。
4.5.6 工程实现
kimi_linear_step 通过 version 参数在 v1 / v2 门控间切换（variants.py），其余步骤两端共用：
# kimi_linear_step − KDA（v1/v2 仅在门控计算处分歧，其余一致）
① L2归一化： q, k
② 门控计算：
v1（默认/SGLang 默认）: g = −exp(A_log) · softplus(gate_logit + dt_bias) ← 无下界
v2（K3 实际使用） : g = g_min · sigmoid(exp(A_log)·(gate_logit+dt_bias)) ← 有下界
a = exp(g) ← 衰减因子
③ 缩放： q ← q * Dk^−0.5
④ 衰减： S ← a ⊗ S ← per‑channel 精细门控
⑤ 检索+纠错：同 gated_delta_net_step
⑥ 读出： out = q · S_new
KDA v1 vs v2
| 维度 | v1（softplus，无界） | v2（sigmoid 下界，K3 采用） |
|---|---|---|
| 门控公式 | g = −exp(A)·softplus(z) |
g = g_min·sigmoid(exp(A)·z) |
| g 的范围 | 无下界（可 → −∞） | 有下界 g ∈ (g_min, 0) |
| 衰减因子 α | exp(g) 可任意接近 0（几乎彻底遗忘） |
exp ∈ (exp(g_min), 1)，被硬下界兜底（如 exp(−5)≈0.0067） |
| 梯度是否消失 | softplus 无上界 → 大 z 处梯度 ≈1，不消失 | sigmoid 在饱和处梯度 →0，但衰减范围被锁死，不会无界累积 |
| 优势 | 表达力上限更高，可学到彻底遗忘 | 数值安全：chunk 内衰减连乘（a₁a₂…a_c）有上界，防溢出 |
| 代价 | chunk 并行下衰减连乘可能累积溢出（数值不稳） | 遗忘能力被限制到 g_min，极限表达力略降 |
| 用途 | SGLang 默认；串行推理更灵活 | K3 实际使用；保证训练 chunkwise 数值安全 |
g_min 的作用
深层动机：KDA 的递推是 S_t = exp(g_t)·S_{t-1} + …，一个 token 的信息跨步保留强度是一堆衰减因子连乘 ∏exp(g)。
- v1 的 exp(g) 虽然可以任意小（接近 0），但若一个 chunk 内连续几步 g 都非常负，连乘会指数级加速衰减、数值范围不可预期；
- 而有界的 v2 把每步衰减因子钳在
[exp(g_min),1]，chunk 内累加的数值范围是可预期的，这正是 FlashKDA 分块 kernel 能放心并行、无需额外归一化兜底的工程前置条件。
kimi_linear_step 中门控三步的解释
下面以默认的 v1（softplus）为例解释三个输入的作用，v2 只替换门控函数形态（sigmoid），三个输入的角色不变。
先看 KDA 的门控公式：
g = −exp (A_log)・softplus (gate_logit + dt_bias)
a = exp (g) ← 实际衰减因子
gate_logit − 数据依赖的"遗忘信号"
gate_logit = W_gate @ x ← 从当前输入算出来，每步不同
- 这是动态的，每个 token 都不同
- 值越大 → softplus 越大 → g 越负 → a 越小 → 遗忘越多
- 相当于模型在说："这个 token 的信息应该覆盖旧记忆"
A_log − 可学习的"遗忘速度旋钮"
- A_log shape: (H,) ← 每个 head 一个值，训练中学到
- 这是静态的，训练好后固定，所有 token 共享
exp(A_log)控制 gate_logit 的放大倍数- A_log 大 → exp(A_log) 大 → 同样的 gate_logit 产生更激进的遗忘
- 相当于给每个 head 设一个"遗忘敏感度"
dt_bias − 可学习的"遗忘偏置"
- dt_bias shape: (H, Dk) 或 (Dk,) ← 每个 (head, channel) 一个偏置
- 这也是静态的，训练好后固定
- 加到 gate_logit 上，相当于给每个通道一个"默认遗忘倾向"
- dt_bias > 0 → 倾向于遗忘更多；dt_bias < 0 → 倾向于保留更多
- 名字暗示灵感来自状态空间模型的 dt（离散化步长）
三者如何协同
gate_logit ──┐
dt_bias ─────┴─> z = gate_logit + dt_bias → softplus (z) → × exp (A_log) → 取负 → g
↓
A_log ─────────────────────────────────────────────────────────────────────> a = exp (g)
| 参数 | 动态/静态 | 粒度 | 作用 |
|---|---|---|---|
| gate_logit | 动态（每 token 不同） | (B, H, Dk) | 当前输入想要遗忘多少 |
| dt_bias | 静态（训练后固定） | (H, Dk) | 每个 (head, channel) 的遗忘基线 |
| A_log | 静态（训练后固定） | (H,) | 每个 head 的遗忘敏感度 |
4.6 分析 & 讨论
4.6.1 端到端训练
gate / gate_logit / beta_logit 是怎么来的
gla_step / gated_delta_net_step 这两个接收 gate、gate_logit、beta_logit 参数的函数是纯递推算子，自身不产生门控。它们只负责"用给定的门控去更新状态"。真正产出门控的是模型里一层把它们包起来的 nn.Module，从这一层的输入 x（即上一层 hidden state）投影出来。
一句话：门控不是超参数、也不是算子自带的，而是"这一层输入 x 的可学习投影"——所以它随 token 变化（data-dependent），模型在训练中学会"什么时候该忘、该写多强"。
GLA的 gate：gate = sigmoid(W_gate @ x)，W_gate把(B,T,D)投到(B,H,Dk)每 token 每通道一个遗忘因子。GatedDeltaNet的 gate_logit：多数实现用低秩双线性投影省参数，gate_logit = Up(Down(x)) + b_alpha（对应 kimi-k3-pytorch 的alpha_down/alpha_up/b_alpha）；beta_logit：beta_logit = W_beta @ x（W_beta：D → num_heads），每头一个写入强度（对应其beta_proj）。- 反例（静态参数，不随输入变化）：KDA 的
A_log (H,)与dt_bias (H,Dk)是可学习但输入无关的参数，训练时反传更新，但所有 token 共享同一份，相当于"遗忘基线/敏感度"旋钮。
本项目代码现状：linear_attn/ 新增了把这条链路打通的上层封装——
src/models/llm/linear_attn/layers.py 里的 GatedLinearAttentionLayer 与 GatedDeltaNetLayer：
GLA 层：gate_proj = nn.Linear(dim, H*head_dim)，forward 里对每个 t 把gate_proj(x)切给gla_step；GatedDeltaNet 层：gate_down/gate_up（低秩双线性产 gate_logit）+beta_proj（产每头 beta_logit），投给gated_delta_net_step；- 两个层内部都是「对整条序列逐 token 递推 + 收集输出」，与手写循环逐位一致（测试
test_linear_attn_layers.py验证输出 diff ≈ 0，W_gate 梯度非零）。
一句话：gate 的完整来源链是 x → nn.Linear(W_gate) → sigmoid/softplus → gate，其中 W_gate 是这一层自带的、与 q/k/v 投影并列的另一组可训练权重。
训练端到端一致
训练这些层的方式与标准 Transformer 完全一样：给定输入序列，前向传播得到每个位置的输出，算 loss（如语言建模的交叉熵），反向传播更新全部参数——包括 W_gate 在内的门控投影。关键点：因为是"数据依赖的门控"，同一个模型喂不同输入会得到不同遗忘模式，模型自己学会"什么时候该忘、什么时候该记"。
4.6.2 KDA vs Gated DeltaNet: 为什么更精细？
我们用一张图把"函数值"和"梯度"两条曲线画在一起，一眼看清区别。
图例
图中 2×2 四格讲了一个完整的故事：
-
左上（sigmoid 及其导数）：蓝线是
σ(x)，虚线是导数σ'(x)=σ(1−σ)。注意当x → −∞时：函数 σ(x)→0（能遗忘），但导数也→0（双消失）。这就是"能忘但学不到彻底忘"——模型想让某个通道彻底遗忘，需要把 logit 推到很负，可一到那个区域梯度就消失了，反向传播推不动参数。 -
右上（softplus 及其导数）：绿线是
softplus(x)=log(1+e^x)，紫虚线是导数σ(x)。当x很正时，softplus ≈ x（斜率→1），导数不消失。把 softplus 取负当门控（g=−exp(A)·softplus(z)），z 越大 → g 越负 → 衰减因子a=exp(g)→0，而 g 对 z 的导数−exp(A)·σ(z)依然是有效非零梯度。 -
左下（衰减因子 a 到底能不能逼近 0）：红实线 KDA v1
a=exp(−exp(A)·softplus(z))随 z 增大无限逼近 0（彻底遗忘）；蓝虚线 GDNa=σ(z)到 z≈−7 就基本停住（σ(−7)≈0.0009之后下降极慢），且再往下梯度就没了。图中还标注了 KDA v2 的硬下界exp(g_min)=exp(−5)≈0.0067——那是"有界版"为数值安全自愿放弃的部分遗忘能力。 -
右下（梯度通路：KDA v1 多一条 learnable 参数 A 的通路）
这一格揭露一个反直觉的事实：红绿两条
|da/dz|曲线其实处处相等——因为恒等式1−σ(z) = exp(−softplus(z))让 GDN 和 KDA v1 (A=0) 的梯度公式在数学上是同一条。所以 KDA v1 的优势不在 z 这条通路——
z→±∞时两者的|da/dz|都照样趋零。真正的优势是 KDA v1 多了一个 learnable 通道参数 A（GDN 没有这个参数）：
- 红色
|da/dz|（两者共享）：峰在 z=0 处（值 0.25），两端指数衰减。 - 绿色
|da/dA|（KDA v1 独有）：峰更高（约 0.37，位于 z≈0.4），且尾部衰减更慢——按|log a|·a衰减（对数级），而不是|da/dz|的线性a衰减。
直观理解：当模型想把某个通道的 a 压到 0.001（彻底遗忘）时
- GDN 只能调 z，此时
|da/dz| ≈ 0.001（梯度几乎死透） - KDA v1 还能调 A，此时
|da/dA| ≈ |log(0.001)| × 0.001 ≈ 0.007（梯度大 7 倍）
a 越小这个差距越大（a=1e-6 时大约 14 倍）。这就是 KDA v1 能"学得动彻底遗忘"的数学根源——不是 z 通路变好了，而是多了一条对数级慢衰减的 A 通路。
图中两个色带分别标出两者的 a→0 区（GDN 在 z→−∞ 侧，KDA v1 在 z→+∞ 侧），直观展示"在各自的极端遗忘区，谁还有可用的梯度通路"。
- 红色
公式拆解
我们来拆解公式，看看为什么一个梯度消失、一个还有梯度。
| Gated DeltaNet（sigmoid） | KDA v1（softplus） | |
|---|---|---|
| 门控函数 | a = σ(z) |
g = −exp(A)·softplus(z)，a = exp(g) |
| "彻底遗忘"需要的输入 | z → −∞（σ(−∞)=0） |
z → +∞（softplus(z)≈z 无上界 → g→−∞） |
| 函数对输入的导数 | da/dz = σ(1−σ)，z→−∞ 时 →0 |
da/dz = a·(−exp(A))·σ(z)，z→+∞ 时 = a·(−exp(A))·1，仍非零 |
| 结论 | 想彻底遗忘 → 梯度消失 → 学不到 | 想彻底遗忘 → 梯度仍在 → 学得到 |
核心机理一句话：sigmoid 的"遗忘极限"和"梯度消失"发生在同一个区域（z 很负时两者都→0），所以它永远到不了 0；softplus 的"遗忘极限"（z→+∞，无上界）与它自己"导数→1"的区间重合，所以它能一路学下去直到 a→0。
数值实测（脚本回填，plot_forget_limit.py）：
| 输入 logit | sigmoid 值 | sigmoid 梯度 σ(1−σ) | KDA v1 a=exp(−softplus(z)) | KDA v1 梯度 |
|---|---|---|---|---|
| z = −8 | 0.00034 | 0.000335 | — | — |
| z = +8 | 0.9997 | 0.000335 | 0.000335 | 0.000335（非零） |
| z = −10 | 4.5e−05 | 4.5e−05 | — | — |
| z = +10 | 0.99995 | 4.5e−05 | 4.5e−05 | 4.5e−05（非零） |
注意表中的对称性：sigmoid 在 z=−8 处的值/梯度 = KDA v1 在 z=+8 处的值/梯度（都≈0.000335）——但方向不同：sigmoid 是"值→0 且梯度→0"（同时死），KDA 是"值→0 但梯度保持在同样量级"（值死梯度活）。这就是全部差异。
总结
KDA 的 softplus 门控让模型有能力学到"彻底遗忘"（a→0），而 Gated DeltaNet 的 sigmoid 门控在 a→0 时梯度消失，学不到这种极端遗忘。加上 dt_bias 和 A_log 两个静态旋钮，模型可以更精细地调节每个通道的遗忘行为。
注意：上面的"KDA"指的是 v1（softplus）。K3 最终采用的 v2 恰恰是有下界的 sigmoid 形式（
g = g_min·sigmoid(...)）——这看起来"倒退"，实则是用有限的遗忘表达能力换取 chunkwise 数值安全（见 v1/v2 对比）。换句话说，KDA 家族内部同样在权衡"表达力（能学到多极端的遗忘）"与"工程稳定性（并行下不溢出）"，v1 站表达力、v2 站稳定性。
4.6.3 KDA 作为可学习位置编码
启示：位置编码未必需要显式张量，机制性的 recency 语义可以替代。
核心思想
KDA 不需要 RoPE，因为位置信息已经编码在衰减模式中，衰减连乘 ≈ 指数衰减 ≈ 相对距离编码。
推导
考虑一个 token j 写入状态 S 的信息，经过 t 步后它的保留程度由衰减连乘决定：
weight(j -> t) = alpha_{j+1} * alpha_{j+2} * ... * alpha_t
假设 alpha_k = exp(-gamma_k), gamma_k > 0, 则:
weight(j -> t) = exp(-sum_{k=j+1}^{t} gamma_k)
~ exp(-gamma * (t - j)) # 若 gamma 近似恒定
这就是指数衰减: 距离越远 (t-j 越大), 信息保留越少
为什么这是"位置编码"？
我们用图例说话
图中 2×2 四格逐层说明了"衰减=位置"：
- 左上（保留强度 vs 距离，对数坐标）：五条不同衰减因子 α 的保留强度曲线。
α=0.9时 token 走过 10 步还剩0.9^10≈0.35、100 步只剩0.9^100≈2.7e-5；α=0.5则几步就消失。对数坐标下直线 = 指数衰减——"距离越远、保留越少"这条规律，就是位置信息的体现。 - 右上（KDA v1 门控推导的衰减）：用
a = exp(−softplus(z))、z=+0.3 推导出α≈0.43的保留曲线，验证了"门控公式 → 指数衰减"是真实的数学关系，不是比喻。 - 左下（两个通道的时间尺度）：慢衰减通道（α=0.98，蓝）能记住很久以前的 token（长期记忆），快衰减通道（α=0.75，红）只关注最近几步（局部记忆）。位置信息被按通道分解成不同的"时间尺度"——这正是 KDA 相比 RoPE 的独特之处：每个 key 维度可以有自己的"记忆长度"。
- 右下（KDA vs RoPE vs ALiBi 对比）：RoPE 是周期振荡（固定频率，超出训练范围外推退化）、ALiBi 是线性惩罚（固定斜率）、KDA 是单调指数衰减（每个通道可学习，天然处理任意长度）。三者都表达"相对距离越远、影响越弱"，但 KDA 的衰减率是学出来的。
相对距离 t - j 决定保留强度
我们来看看 与 RoPE/ALiBi 的异同。
| 机制 | 编码方式 | 是否可学习 |
|---|---|---|
| RoPE | 旋转角度 (固定频率) | 频率固定，位置外推退化 |
| ALiBi | 距离线性惩罚 | 固定斜率 |
| KDA 衰减 | 逐通道衰减率 alpha_k | 每个通道独立学习 |
KDA 的优势
- 通道级粒度：每个 key 维度有自己的衰减率，不同通道可以有不同"时间尺度"
（有的通道保留 1000 token，有的只保留 10 token，见上图右上） - 可学习：衰减率由网络学习，不是固定的（上图右下三条曲线的本质差异）
- 无外推退化：RoPE 在序列长度超出训练范围时退化，KDA 的衰减是隐式的，天然处理任意长度
数值示例
假设 alpha = 0.9 (中等衰减):
- token 0 经过 10 步保留 = 0.9^10 = 0.35
- token 0 经过 100 步保留 = 0.9^100 = 0.000027
- 所以位置 t 时，token 0 的信息几乎消失，token t-1 的信息几乎完整
这个"越远越弱"的模式就是位置信息的体现
KDA v1 门控推导的实测 (gate_logit=+0.3, A_log=0):
- alpha = exp(-softplus(0.3)) ≈ 0.426
- 经过 10 步保留 = 0.426^10 ≈ 0.0002
- 经过 30 步保留 ≈ 0（彻底遗忘）
- 不同 gate_logit → 不同 alpha → 不同"位置尺度"：
gate_logit=+0.3 → α≈0.426（快忘，短距离）
gate_logit 更负 → α→1（慢忘，长距离）
关键点：gate_logit 是逐 token、逐通道学出来的——所以"位置编码"本身也是内容相关的（某些 token 该忘、某些 token 该记），这是 RoPE/ALiBi 这种固定编码给不了的。这也是为什么说 KDA"位置信息编码在衰减模式中"。以上数值由
plot_position_decay.py实测回填。
4.6.4 分块并行实现
Chunkwise 实现 把每一步状态更新整理为仿射变换，让 Chunk 之间可以组合，Chunk 内部可以转成下三角系统和矩阵运算。这是纯 PyTorch 实现（用 UT/WY 变换在 chunk 内并行计算，chunk 间串行递推），核心在 kimi3-master/kimi3/modules/kda_ref.py:191-307和 kimi-k3-pytorch-main/src/kda/chunkwise.py。
代码体现
# kda_ref.py kda_chunkwise 的核心结构
def kda_chunkwise(q, k, v, g, beta, *, chunk_size=16, initial_state=None, ...):
n_chunks = ceil(t / chunk_size)
# 阶段 1: token 并行 —— 所有不依赖 incoming state 的计算
log_gamma = g_c.cumsum(dim=3) # 累积衰减
gamma = log_gamma.exp()
gamma_inv = (-log_gamma).exp()
q_bar = q * gamma # Q * Gamma
k_hat = k * gamma_inv # K / Gamma
k_bar = k * gamma # K * Gamma
# UT 变换: 三角求解 (I + M)^-1
uw = torch.linalg.solve_triangular(m, rhs, upper=False, unitriangular=True)
u_mat, w_mat = uw[..., :dv], uw[..., dv:]
# 阶段 2: chunk 间串行递推
for i in range(n_chunks):
v_tilde = u_mat[:, :, i] - w_mat[:, :, i] @ state
outputs.append(carry_ok * (q_bar @ state) + a @ v_tilde)
state = keep * (gamma_last * state) + k_tail @ v_tilde
return cat(outputs), state
UT/WY 变换的体现
kimi3-master：ut_transform.py（92 行）的ut_transform，用torch.linalg.solve_triangular求解kimi-k3-pytorch：src/kda/chunkwise.py（108 行），同样用 UT/WY 变换- 本项目的实现：
Model-Viz/src/models/llm/kimi_k3/modules/kda_ref.py，kda_chunkwise与 kimi3-master 一致
我们的 test_chunkwise_equals_recurrence 测试已验证 chunkwise = recurrent（误差 < 1e-15）
附录 A. 状态检索误差
"状态检索误差" 的本质是定长循环状态能否无损还原单个关联，五个变种就是五档 "编辑旋钮"—— 从只会累积干扰（Linear），到在该键上校正（DeltaNet）、按时龄遗忘（GatedDeltaNet）、按特征遗忘（KDA v1）、按特征精准擦写（KDA v2），逐级把有限容量用得更高效、把编辑的副作用压得更小。
A1. 什么是 "状态检索误差"
线性注意力把无界的 KV 缓存换成定长的循环状态 S_t ∈ R^{d×d}。读取时对 query/key 做线性读出：想取回某 key k_s 对应的值，就读 k_sᵀ S_t。状态检索误差定义为：
E(s) = ‖ k_sᵀ S_t − v_s ‖
即 "压缩后的定长状态到底能不能无损还原当初写入的单个 (key, value) 关联"。它是线性循环模型的核心误差度量。
A2. 它的一般特点（所有线性注意力共有）
- 容量受限：
S_t只有 d×d 个自由度，能忠实表示至多～d 个线性独立的 key - 方向上的关联；装不下就丢精度。 - 干扰受限：读出是 "所有已存关联按 key 相似度加权求和"，key 不正交就互相串扰，检索误差随存储条数增长。
- 遗忘是主要旋钮：怎么 "忘"（decay）直接决定哪些关联还能被检索、误差多大 —— 这是 delta 规则一族设计的核心。
- 写即扰动：写入新关联必然扰动旧关联，检索误差不只是容量问题，更是 "编辑的代价"。各变种本质上是如何管理这个编辑代价的不同方案。
A3. 各变种的状态检索误差特色
1. Linear（纯求和） S_t = S_{t-1} + k_t v_tᵀ
- 检索误差 = 纯干扰：读出把所有键混在一起，无校正、无遗忘。
- 误差随序列长度单调累积（memory overload），新旧关联互相污染程度相同，没有任何机制按重要度或时序取舍。
2. DeltaNet S_t = S_{t-1} + β(v_t − S_{t-1}k_t)k_tᵀ
- 写入是误差校正：先用当前状态检索 "记忆以为的值"
S_{t-1}k_t，算残差δ = v_t − 已检索值，再写残差。 - 效果：刚写入的那个 key 上检索误差≈0（递归最小二乘解
k_tᵀS≈v_t）；但擦除算子(I−βk kᵀ)会把旧状态里与 k_t 同方向的成分一并消掉 → 检索误差变成 "最后写入者优先"：新 key 很准，同方向的旧关联被覆盖丢失。要求‖k‖≈1才稳定。
3. GatedDeltaNet S_t = α(I − βk kᵀ)S_{t-1} + βk v_tᵀ
- 加了 per-head 的标量遗忘门 α（Mamba-2 式）。
- 检索误差变成按时龄形状化：旧关联以
α^gap指数衰减，越久越难检索（recency bias）；近期信息误差低。 - 但 α 是全局标量 → 内容无差别地衰整个状态：不能选择性遗忘、也不能保留重要的旧项。检索误差 "只认年龄，不认内容"。
4. KDA v1（逐通道衰减） S_t = Diag(α_t)(I − β_t k_t k_tᵀ)S_{t-1} + β_t k_t v_tᵀ
- 把标量 α 升级为逐通道向量
α_t ∈ (0,1]^{d_k}→ 遗忘变成特征选择性：每个特征通道独立决定保留 / 衰减。 - 效果：定长状态能只留任务关键的维度、丢弃无关维度，同样容量下有效检索误差更低，是 "内容 / 特征感知" 的遗忘。
- 局限：
β_t仍是单个标量，同时管 "擦多少旧（key 侧）" 和 "写多少新（value 侧）"，两者耦合，无法独立调。
5. KDA v2（逐通道 erase/write 解耦） S_t = (I − k_t b_tᵀ)·Diag(exp g_t)·S_{t-1} + k_t(w_t⊙v_t)ᵀ
- 把标量 β 解耦成逐通道擦除门 b_t（key 侧移除）和逐通道写入门 w_t（value 侧提交），同时保留逐通道衰减。
- 效果：状态能精准编辑—— 逐通道独立决定 "擦多少" 与 "写多少"，写入不再连带搅乱无关旧关联，检索误差里 "编辑的副作用" 被最小化。它是这一族的严格推广：两个门坍缩回共享标量→KDA v1；再坍缩 decay→GatedDeltaNet；去门→DeltaNet/Linear。
A4. 一句话主线
从 Linear 到 KDA v2，本质是把 "定长状态如何被编辑" 的旋钮从无→粗（全局标量）→细（逐通道）逐级放开：Linear 只有纯干扰，DeltaNet 学会 "在该键上校正误差"，GatedDeltaNet 学会 "按年龄忘"，KDA v1 学会 "按特征忘"，KDA v2 学会 "按特征精准擦 + 写"。检索误差的降低都来自让有限的记忆容量优先服务重要的、按需保留的关联。
示意图如下。
统一循环状态更新 S_t = Decay·(I − erase·k kᵀ)·S_{t-1} + write·k vᵀ
四个旋钮：write写 · erase擦 · decay衰 · read读（k_sᵀ S_t）
状态检索误差 E(s) = ‖ k_sᵀ S_t − v_s ‖ ← 定长压缩状态能否无损还原单个(key,value)关联
误差由：容量 · 键正交性 · 遗忘方式 · 编辑干扰 共同决定
[1] Linear S_t = S_{t-1} + k_t v_tᵀ （纯求和）
状态检索误差：键间干扰 + 容量超载，随序列长度单调累积；
无校正、无遗忘 → 新旧关联互相污染程度相同（memory overload）。
│
▼ 加入误差校正写：β(v − S·k)kᵀ
[2] DeltaNet S_t = S_{t-1} + β(v_t − S_{t-1}k_t)k_tᵀ
状态检索误差：刚写入的 key 上误差→0（递归最小二乘解 k_tᵀS ≈ v_t）；
代价：擦除算子 (I − βkkᵀ) 会覆盖与 k 同方向的旧关联 → “最后写入者优先”。
│
▼ 加入标量遗忘门 α（per-head，Mamba-2 式）
[3] GatedDeltaNet S_t = α(I − βkkᵀ)S_{t-1} + β k vᵀ
状态检索误差：按时龄指数衰减（α^gap）→ recency bias，近期信息误差低；
局限：α 是全局标量 → 内容无差别地衰减整个状态，不能选择性遗忘重要旧项。
│
▼ α 升级为逐通道向量 Diag(α_t)
[4] KDA v1 S_t = Diag(α_t)(I − β_t k kᵀ)S_{t-1} + β_t k vᵀ
状态检索误差：特征选择性遗忘 → 有限容量优先保留关键特征通道，有效误差更低；
残余：β 仍为单个标量，同时管“擦多少旧(key侧)/写多少新(value侧)”，二者耦合。
│
▼ β 解耦为逐通道 erase b_t / write w_t
[5] KDA v2 S_t = (I − k b_tᵀ)·Diag(exp g_t)·S_{t-1} + k (w_t⊙v_t)ᵀ
状态检索误差：精准编辑 → 逐通道独立决定“擦多少/写多少”，
写入不再连带搅乱无关旧关联，编辑副作用最小化（严格推广前 4 者）。
主线：把“定长状态如何被编辑”的旋钮 从 无 → 粗(全局标量) → 细(逐通道) 逐级放开
Linear 纯求和 → 干扰累积（无编辑）
DeltaNet 在该键上校正误差（代价：覆盖同向旧关联）
GatedDeltaNet 按年龄遗忘（代价：内容无差别）
KDA v1 按特征遗忘（残余：擦/写仍耦合，标量 β）
KDA v2 按特征精准擦 + 写（最小化编辑副作用）
───────────── 改进（该键 / 近时 / 相关特征）─────────────
· DeltaNet：写入键误差→0 · KDA v1：保留关键特征通道
· GatedDN ：近期信息误差低 · KDA v2：编辑不搅乱无关旧关联
───────────── 代价 / 残余误差 ─────────────────────────
· DeltaNet：覆盖同向旧关联 · KDA v1：擦/写仍耦合
· GatedDN ：旧内容无差别衰减 · 三者共同：容量、键正交性下界
0xFF 参考
48 小时啃完 Kimi K3：一个模型装得下 22580 个 GPT-2
详解 DeepSeek V4：Infra 巨鲸 “四连击”，百万上下文走进现实
原文地址: https://www.cveoy.top/t/topic/qHQ9 著作权归作者所有。请勿转载和采集!