Kaczmarz Linear Attention：中文精读笔记
- 作者：Jiaxuan Zou、Ruifeng Ren、Yong Liu（通讯作者）。
- 原文版本：arXiv:2605.08587v1，提交于 2026-05-09；本次研读与外部信息查询日期：2026-11-13。
- 原文：本地 PDF；作者 TeX 源码。以下页码均指从 1 开始的 PDF 页码。
- 发表：截至 2026-11-13，未查到正式发表或接收信息；第一作者主页将此工作标为 Under Review，arXiv 当前记录未列会议、期刊或出版 DOI。
- 代码：作者提供的匿名官方实现，已公开，README 与文件列表包含模型组件、训练和效率评测脚本。
论文总览
KLA 要解决的是固定容量循环记忆的“写多大一步”。Gated DeltaNet（GDN）先遗忘旧状态，再沿当前 key 方向写入预测残差；KLA 保留这套结构，把写入系数从学习得到的 $\eta_t$ 改为 $\eta_t/(\|k_t\|_2^2+\varepsilon)$。分母消除了当前 key 能量对有效纠错幅度的直接影响。在不松弛且没有稳定项时，这一步恰好是使当前键值约束成立、同时最少改变旧状态的 Kaczmarz 正交投影。
论文报告：0.4B 参数、1B token 预训练中，验证困惑度（perplexity，PPL，越低越好）由 GDN 的 8.50 降到 8.09；MQAR 在训练长度的 8 倍处由 66.81% 提升到 73.84%；S-NIAH 在测试的 1K–8K 长度均为 100%。它的贡献集中在写入步长的优化解释与小改动实现。长上下文 PPL 的评测口径、消融中的配置对应，以及相同 kernel 下的解码加速原因仍需进一步澄清，见对应结果分析。
结构地图
| 原文章节 | 在论证中的作用 | PDF 页码与关键证据 |
|---|---|---|
| 1 Introduction | 固定状态维护问题与主要结果 | pp. 1–2；Table 1 |
| 2 Background and Problem Analysis | 加法写入、遗忘与 delta rule 的在线回归视角 | pp. 2–3 |
| 3.1 Kaczmarz coefficient from constrained optimization | 最小改动投影与近端形式 | pp. 3–4；Eqs. (1)–(4)，Proposition 1 |
| 3.2 Why the Kaczmarz coefficient is natural | 投影、线搜索、归一化梯度三种解释及残差收缩 | pp. 4–5；Theorem 1，Eq. (5) |
| 3.3 Tokenwise implementation | 单 token 数据流与门控 | p. 5；Algorithm 1 |
| 4 Chunkwise Parallelism | 复用 GDN 分块三角求解器 | pp. 5–6；Eqs. (6)–(9) |
| 5.1 Setup / 5.2 Language Modeling | 固定预算预训练与长度外推 | pp. 6–7；Fig. 1，Tables 1、3 |
| 5.3 Synthetic Tasks: Controlled State Use | 检索、逆序和栈操作 | pp. 7–8；Table 4，Fig. 2 |
| 5.4 Efficiency / 5.5 Ablation Studies | 运行成本与归一化、门控、容量取舍 | pp. 7–9；Fig. 3，Table 5 |
| 6 Limitations / 7 Conclusion | 固定容量、精度与任务边界 | pp. 9–10 |
| A Proofs for the Tokenwise Theory | 投影、拉格朗日、近端、收缩与遗忘算子 | pp. 13–15 |
| B Derivations for Chunkwise Parallelism | WY 表达与分块输出推导 | pp. 15–18；Table 6 |
| C Experimental Details / D Ablation Study Experimental Setup | 架构、优化器、计算与消融协议 | pp. 18–19 |
| E Related Work | 将 KLA 定位在更新规则这一设计维度 | pp. 19–20；Table 7 |
一：动机——残差方向正确，还需要合适的步长
1.1 从存储键值对到编辑已有关联
每个 head 用矩阵 $S_t\in\mathbb R^{d_k\times d_v}$ 压缩历史。key $k_t\in\mathbb R^{d_k}$ 指定写入方向，value $v_t\in\mathbb R^{d_v}$ 指定希望存入的内容，query $q_t\in\mathbb R^{d_k}$ 用于读取。普通线性注意力用 $S_t=S_{t-1}+k_tv_t^\top$ 累加键值对；相似 key 反复出现时，已有内容不会在写入前被显式扣除。
delta rule 先问当前状态对这个 key 已经存了什么，再只写入缺失部分。GDN 另外用遗忘门控制旧内容：
$\alpha_t$ 控制整体保留多少旧状态，$\beta_t$ 控制沿当前 key 方向改多少。论文使用的 GDN 模板令 $\beta_t=\eta_t$，其中 $\eta_t$ 为输入线性投影经 sigmoid 得到的门。（Sec. 2，PDF pp. 2–3；Sec. 3.3，p. 5。）
1.2 key 范数为什么会改变纠错幅度
写入后再次用同一 key 读取，残差变为
因此真正决定纠错比例的是 $\beta_t\|k_t\|_2^2$。同一个门值面对大范数 key 可能过冲，面对小范数 key 可能改得不够。KLA 显式除以 key 能量，让学习门主要控制“本次残差纠正几成”。若 $\varepsilon=0$，KLA 就有 $e_t^+=(1-\eta_t)e_t$。
一个会影响实现判断的条件是：若进入残差计算和写入的 key 已经严格单位归一化，则 $\|k_t\|_2^2=1$，KLA 与论文的 GDN 系数模板几乎重合，仅相差稳定项。这是公式的直接推论；论文明确描述了 query 的归一化，但没有充分交代各 baseline 进入更新式时的 key 处理差别，因此不能据本文断言任意 GDN 实现都会出现同样的尺度失配。
二：方法——从最小改动投影到可并行的循环层
2.1 符号与单 token 计算
| 符号 | 维度或范围 | 用途 |
|---|---|---|
| $S_t,\widetilde S_{t-1}$ | $d_k\times d_v$ | 更新后的状态、衰减后的状态 |
| $k_t,q_t$ | $d_k$ | 写入与读取方向 |
| $v_t,e_t$ | $d_v$ | 目标值与当前预测残差 |
| $\eta_t$ | sigmoid 输出，$(0,1)$ | 逐 token 的松弛写入门 |
| $\alpha_t$ | 标量遗忘门 | 衰减旧状态 |
| $\varepsilon$ | 实验使用 $10^{-6}$ | 避免极小 key 能量带来的除法问题 |
推理时依次计算衰减状态、预测残差、归一化步长和秩一写入，再读取
残差必须由已衰减的 $\widetilde S_{t-1}$ 计算。归一化系数本身只依赖 key、门值和稳定项；更新方向还依赖状态与 value。这里的在线优化对象是快速变化的记忆状态 $S$，模型的投影参数仍通过常规预训练学习。（Eq. (3)，PDF p. 3；Algorithm 1，p. 5；Appendix C，p. 18。）
2.2 为什么系数恰好是 key 能量的倒数
作者把当前键值写入视为约束 $S^\top k_t=v_t$，希望在满足约束的状态中最少改变旧记忆：
假设 $k_t\ne0$。约束的法向方向由 $k_t$ 决定，最优改动可写为 $k_tc^\top$。代入约束得
因此
每个 value 通道对应一个线性约束，Frobenius 范数把各通道的最小改动组合起来。对任意不改变当前键值读出的方向 $H$，有 $H^\top k_t=0$，所以 $\langle k_te_t^\top,H\rangle_F=0$；这给出了正交投影的几何解释。（Sec. 3.1，PDF p. 3；Appendix A.1–A.2，pp. 13–14。）
2.3 同一步更新的三种解释
在当前键值回归损失
上，$\nabla_S\mathcal L_t(\widetilde S_{t-1})=-k_te_t^\top$。沿此负梯度方向取 $S(\tau)=\widetilde S_{t-1}+\tau k_te_t^\top$，则
当 $e_t\ne0$ 时对 $\tau$ 求导，最优步长为 $1/\|k_t\|_2^2$；若 $e_t=0$，更新已为零，沿这条退化路径任意步长都产生相同状态。因此非零 key、$\eta_t=1$、$\varepsilon=0$ 时，同一状态更新既是正交投影，也是精确线搜索结果和归一化梯度步。（Theorem 1，PDF p. 4。）
教学例子。 令 $d_k=d_v=1$、衰减后状态为 0、$k=2$、$v=1$。直接用 $\beta=1$ 得到 $S=2$，读出 $Sk=4$，残差从 1 变为 −3；Kaczmarz 用 $\beta=1/4$ 得到 $S=1/2$，读出正好为 1。若再令 $\eta=1/2$，则只纠正一半残差。这说明分母补偿的是“写入乘一次 key、读出再乘一次 key”的双重尺度。
2.4 实际门控：局部收缩与近端回归
实际模型保留 $\eta_t$，并加 $\varepsilon$：
在 $k_t\ne0$、$0<\eta_t\le1$、$\varepsilon\ge0$ 时，括号中的系数位于 $[0,1)$，当前关联的损失不增。比较基点是衰减后的状态；这一性质没有要求遗忘操作本身降低损失。若 $k_t=0$ 且 $\varepsilon>0$，秩一更新为零，残差保持不变；无稳定项的投影公式在零 key 上没有定义。
松弛更新还等于如下近端问题的解：
一阶条件给出的系数是 $\mu_t/(1+\mu_t\|k_t\|_2^2)$，代入即得 KLA 系数。原文 Proposition 1 把 $\eta_t=1,\varepsilon=0$ 也包括在范围内，但此端点的 $\mu_t$ 发散，应按 $\mu_t\to\infty$ 的约束投影极限解释。有限 $\mu_t$ 时，该目标严格凸，解唯一。（PDF pp. 3–4、14。）
对其他 key $k'$，本次写入对其读出的影响为 $\Delta(S^\top k')=\beta_t e_t(k_t^\top k')$（本次分析）。这解释了固定状态中的干扰来源：与当前 key 正交的方向不受秩一写入影响，相似方向仍会被编辑。归一化改善写入尺度，但不会增加可存关联的容量。
2.5 分块并行：把递推依赖改写成三角系统
推理按 token 递推，训练则把长度 $C$ 的 chunk 内依赖改写成矩阵运算。令 $K,Q\in\mathbb R^{C\times d_k}$、$V,U\in\mathbb R^{C\times d_v}$，$S_0$ 为输入状态，并定义
这里使用非零累计衰减，且 $Q$ 的每行应已经按实际读出规则归一化。$u_i$ 是第 $i$ 次写入中吸收了 $\beta_i$ 的有效残差。直接展开原始递推可得
这就得到正文的分块公式：
左侧矩阵为单位下三角矩阵，精确算术下可逆。KLA 仅改变 $B$ 的数值构造，保留 GDN 的求解与状态布局。固定 head 维度和 chunk 长度时，总序列计算随长度线性增长；单 token 状态维护需要 $O(d_kd_v)$ 存储，不随已处理的上下文长度增长。训练仍需保存或重算序列激活，不能把固定循环状态理解成整个训练过程只有常数内存。（Sec. 4，PDF pp. 5–6；Appendix B.2–B.4，pp. 17–18。）
附录推导中的记号问题。 Appendix B（PDF pp. 15–17）定义 $R=S^\top\in\mathbb R^{d_v\times d_k}$，随后却写 $P_iR_0$，其中 $P_i\in\mathbb R^{d_k\times d_k}$，一般维度不匹配；其 $H_i$ 也混用了 $v_jk_j^\top$ 和 $k_ju_j^\top$。此外，p. 15 定义 $\gamma_{j\to i}=\gamma_i/\gamma_{j-1}$，p. 17 又替换成 $\gamma_i/\gamma_j$，两者相差起点处的衰减。这些中间式不能直接照抄。上面的教学推导统一使用 $S$ 方向及“写入后仅经历后续衰减”的 $\gamma_i/\gamma_j$，可得到正文 Eqs. (7)–(9)。这区分了正文算法与有误的中间记号。
三：实验——各组结果在比较什么
3.1 预训练和长度外推
主实验使用 SlimPajama-672B 中的 1B token、0.4B 参数、4K 训练上下文。作者声明各模型使用相同 tokenizer、AdamW 和余弦调度。架构为 24 层、隐藏维度 1024、8 个 head、head 维度 128、FFN 扩张 4 倍。优化器参数为 $(\beta_1,\beta_2)=(0.9,0.95)$，weight decay 0.1、梯度裁剪 1.0、前 2% step 线性 warmup。实验使用 4 张 A100 80GB，作者报告每个 1B-token run 约 4 小时；未做多随机种子预训练。（Sec. 5.1，PDF p. 6；Appendix C，p. 18。）
预训练比较 KLA、DeltaNet、GLA、GDN、Longhorn、Mamba2。GDN 是控制架构改动范围的直接对照；其他模型提供不同更新机制的参照。长度外推、合成任务和效率仅比较 KLA、GDN、Mamba2。PPL 衡量 next-token 预测，不是检索准确率；在未做长上下文微调的条件下，另测 1K–65K 长度的 PPL。
3.2 四种合成任务
| 任务 | 操作与检验对象 | 指标与协议 |
|---|---|---|
| MQAR | 多个 key-value 关联后回答多个查询，检验关联记忆 | validation accuracy；2 层模型，训练长 256、10K steps，外推到 2048 |
| S-NIAH | 大量干扰中检索一个关键关联 | exact-match accuracy；6K steps，评测长 1K、2K、4K、8K |
| Palindrome | 逆序复述随机序列，检验有序信息存储 | accuracy；20K steps，评测长 256–2048 |
| Stack | 多栈 push/pop，检验后进先出状态编辑 | accuracy；20K steps，评测长 256–2048 |
合成任务采用 seed 42。原文只对 S-NIAH 明确使用 exact match，不能把其他任务的 accuracy 一律解读为整序列全对率。（Sec. 5.3、Table 4，PDF pp. 7–8；Appendix C，p. 18。）
3.3 效率与消融的独立口径
效率指标为 prefill 延迟（ms，越低越好）、decode 吞吐（token/s，越高越好）、每输出 token 时间 TPOT（ms/token，越低越好），batch size 为 1，使用 FLA 评测路径。Prefill 测到 131072 token，decode 图测到 32768 token。（Fig. 3，PDF p. 9。）
消融使用另一套较短预算：语言建模仅预训练 100M token，长 4096，学习率 $10^{-4}$；PPL 计算包含 20 次验证、每次平均 15 iterations。MQAR 消融使用 20K/2K/2K 条 train/validation/test 数据、32 个键值对、batch 32、学习率 $10^{-3}$、最多 10K steps；每 200 step 评估，早停 patience 为 10。大状态实验降低 micro-batch 以避免 OOM，全局 token 预算不变。这些 PPL 与主实验 1B-token PPL 分别解读。（Appendix D，PDF pp. 18–19。）
四：结果——收益来自哪些任务，证据有什么取舍
4.1 语言建模：固定预算下的 PPL 改进
| 模型 | Validation PPL ↓ |
|---|---|
| KLA | 8.09 |
| DeltaNet | 8.26 |
| GLA | 8.42 |
| GDN | 8.50 |
| Longhorn | 8.79 |
| Mamba2 | 9.41 |
来源：Table 1，PDF p. 2；0.4B 参数、1B token、4K 训练上下文。KLA 相对 GDN 降低 0.41 PPL，按 GDN 为分母相对下降约 4.82%；相对最接近的 DeltaNet 仅低 0.17。这个排序支持步长修正能改善该预算下的语言建模。
原图 Fig. 1，PDF p. 6；来自作者源码图件。横轴为已训练 token 数（B），纵轴为 validation PPL；六条曲线与表中模型对应。图展示约 0.4B–1B token 区间，KLA 最终最低；它不是从第一个训练 step 开始的完整曲线。
4.2 长度外推：跨长度保持较低 PPL，但与主表不宜合并
| 模型 | 1K | 2K | 4K | 8K | 16K | 32K | 65K |
|---|---|---|---|---|---|---|---|
| KLA | 39.1 | 39.8 | 40.0 | 43.6 | 45.3 | 44.5 | 43.6 |
| GDN | 42.9 | 43.5 | 43.8 | 47.4 | 49.4 | 48.4 | 47.5 |
| Mamba2 | 49.9 | 50.9 | 52.2 | 56.4 | 58.7 | 57.3 | 56.7 |
来源：Table 3，PDF p. 7；PPL 越低越好，无长上下文微调。KLA 比 GDN 在每个长度低 3.7–4.1 PPL；65K 处为 43.6 对 47.5，相对下降约 8.21%。KLA 自身从 4K 的 40.0 到 16K 的 45.3 后又下降，支持“在测试区间没有持续恶化到失控”的描述。
同为 4K，Table 1 的 KLA/GDN 是 8.09/8.50，Table 3 却是 40.0/43.8。正文没有充分说明验证语料、checkpoint 或计分流程怎样对应，不能将这两组绝对值拼成一条长度曲线。长度外推的可靠比较首先是 Table 3 内部的同长度模型排序；作者关于减少状态损坏的解释与这一趋势相符，但表中没有直接测量状态损坏量。
4.3 合成任务：关联记忆更强，有序逆转并非全面提升
| 任务与长度 | KLA ↑ | GDN ↑ | Mamba2 ↑ | KLA 相对 GDN |
|---|---|---|---|---|
| MQAR，256 | 99.00% | 98.74% | 99.08% | +0.26 个百分点 |
| MQAR，1024 | 95.64% | 94.31% | 96.54% | +1.33 个百分点 |
| MQAR，2048（8×） | 73.84% | 66.81% | 64.29% | +7.03 个百分点 |
| S-NIAH，8K | 100.00% | 98.50% | 0.05% | +1.50 个百分点 |
| Palindrome，2048 | 79.96% | 80.74% | 0.86% | −0.78 个百分点 |
| Stack，2048 | 96.94% | 95.64% | 93.82% | +1.30 个百分点 |
来源：Table 4，PDF p. 8；按各任务原指标报告。MQAR 的 7.03 是准确率的百分点差，不是相对提升 7.03%；相对 GDN 准确率增幅约 10.52%。KLA 在 1K、2K、4K、8K 的 S-NIAH 都达到 100%。MQAR 的短长度赢家是 Mamba2，优势到 8× 时才转为 KLA；不能概括成每个长度均最好。
Stack 在长序列上的收益符合“反复编辑状态”的动机。Palindrome 则要求保存顺序，KLA 在 512 和 2048 两个长度略低于 GDN，说明改善关联写入步长不自动解决有序存储需求。
原图 Fig. 2，PDF p. 8；四面板依次为 MQAR、S-NIAH、Palindrome、Stack，横轴 training steps，纵轴各任务准确率；蓝色为 KLA、绿色为 GDN、红色为 Mamba2，嵌入图放大高准确率区间。Mamba2 在 MQAR 学得较早、在 S-NIAH 近乎失败、在 Palindrome 明显较慢，而 Stack 的后期可达很高精度。该图表达训练过程，长度外推优劣仍以 Table 4 为准。
4.4 效率：prefill 接近，decode 增益需要独立看待
| 指标与上下文 | KLA | GDN | Mamba2 |
|---|---|---|---|
| Prefill 延迟，131K，ms ↓ | 1220 | 1227 | 1936 |
| Decode 吞吐，32K，token/s ↑ | 3.54 | 1.65 | 2.10 |
来源：Sec. 5.4，PDF p. 7；Fig. 3，p. 9；batch size 1。Prefill 与 GDN 相差约 0.57%，符合没有引入新的主导分块成本。吞吐比为 $3.54/1.65\approx2.15$，原文取约 2.1×；相对 Mamba2 约 1.69×。对吞吐取倒数可得到约 282、606、476 ms/token，与 TPOT 面板的量级对应；这些倒数是由报告值计算的近似值。
原图 Fig. 3b，PDF p. 9；横轴为上下文长度，纵轴 token/s。三者均随上下文增长而显著下降，KLA 在 32K 的差距扩大。固定循环状态本身并不需要随完整历史线性增加每个解码 token 的状态维护量；原文没有充分交代这里的 decode 是否包含重算、缓存路径或其他计时开销，也未解释相同 kernel 和状态布局如何产生约两倍差距。因此这项收益应表述为论文所用评测路径的实测结果，不能由步长公式直接推导为 recurrent kernel 的普遍两倍加速。
4.5 消融：归一化、门控和状态容量各有代价
归一化消融。
| 变体 | MQAR Test ↑ | MQAR 8× ↑ | 100M-token PPL ↓ |
|---|---|---|---|
| KLA 默认 | 98.39% | 73.84% | 753.65 |
| 无归一化 | 0.03% | 0.01% | — |
| 仅 key 范数，$1/\|k_t\|_2^2$ | 97.81% | 77.68% | 3809.71 |
| Learned Scalar | 97.77% | 74.46% | 22968.87 |
来源：Table 5，PDF p. 9。仅 key 范数的 8× 准确率更高，但 PPL 显著更差，说明保留学习门对语言建模有价值。Learned Scalar 在附录中被定义为初始化为 1 的可训练参数，不应等同于标准的逐 token sigmoid 门。
这里存在需要澄清的配置对应：Appendix D（p. 19）把 NoNorm 写成 $\beta_t=\sigma(b_t)$，而 Sec. 3.3 的 GDN 也使用 sigmoid 写入门；Table 5 的 NoNorm 接近随机，但 GDN reference 能达到 66.81% 的 8× 准确率。仅从论文描述无法确定二者还差在哪些 key 处理或模块配置。因此该消融支持“在所用变体管线中归一化重要”，不能据此称去掉分母就必然使标准 GDN 崩溃。表中 “—” 也没有足够信息可具体解释成哪一种失败。
门控消融。 Single Gate 为 0.01%/0.03%/2104.61；GDN 双门参考为 98.72%/66.81%/747.24；GLA 独立门参考为 0.03%/0.03%/356.01（依次为 Test、8×、PPL）。GLA 的短预算 PPL 最低但 MQAR 失败，GDN 的短预算 PPL 也略低于 KLA 默认的 753.65。因此合理结论是任务间有取舍，而非 KLA 在消融中所有指标都最好。
状态容量消融。
| $v_{\mathrm{expand}}$ | MQAR Test ↑ | MQAR 8× ↑ | 100M-token PPL ↓ |
|---|---|---|---|
| 2 | 98.09% | 71.99% | 762.33 |
| 4 | 98.88% | 91.34% | 1669.35 |
| 8 | 95.31% | 95.37% | 6945.63 |
| 16 | 99.18% | 92.63% | 1383.38 |
来源：Table 5，PDF p. 9。扩大 value 状态显著改善某些外推准确率，但 PPL 并非随容量单调改善，且这是改变状态容量的比较，不能解释成相同容量下的纯步长收益。默认行的 Test=98.39% 也不同于 Table 4 的 validation=99.00%，两者指标与协议应分别保留。Appendix D 还列出按位置衰减系数的变体，但没有给出对应结果，因而无法据此比较 $1/t$ 等方案。
五：结论——一个清楚的局部优化原则，几项值得区分的实证主张
KLA 最有启发的地方是把门控更新拆成两个角色：key 能量决定局部二次损失的自然步长，学习门决定当下关联应被写入多强。它保留 GDN 的矩阵状态与分块求解结构，在非零 key 的理想端点上有简洁的最小改动投影解释，在实际松弛设置中有明确的残差因子。
本文的实证亮点是固定预算预训练的 PPL 降低、较长 MQAR 的关联回忆，以及 S-NIAH 的准确率。固定容量仍限制可保存的关联数；相近 key 的写入干扰仍存在，Palindrome 也提供了具体的任务边界。很小的 key 范数还使 $\varepsilon$ 与数值精度相关，不能只保留投影极限而忽略实际稳定项。
若继续研究，优先价值最高的是厘清进入更新式的 key 是否已归一化、给出 NoNorm 与 GDN reference 的完整配置差异，并把 decode 计时拆成 prefill、状态更新与其他开销；这直接决定一步系数改动与已报告增益之间的归因。更大规模、多种子预训练可再判断 0.41 PPL 差异的稳定性。对角遗忘与更大状态是作者提出的后续方向，其中局部投影可适用于不同衰减基点，但实际并行算法和成本仍需分别分析。