面向长上下文线性注意力的克罗内克德尔塔法则

用标量修正的克罗内克德尔塔法则构建KronLSTM,实现低干扰、有界稳定的长上下文线性注意力

前置知识

线性注意力 (Linear Attention)

把 softmax 注意力改写为键值外积的累积:记忆矩阵按 $S_{t+1}=S_t+v_tk_t^\top$ 更新,读出为 $h_t=S_tq_t=\sum_{i<t}v_ik_i^\top q_t$。上下文被压缩进固定大小的 $d_v\times d_k$ 矩阵,计算复杂度对序列长度线性、显存恒定,思想可追溯至 Schmidhuber 1992 的快速权重编程。

本文的矩阵记忆就是这套框架,KronLSTM 只是把'如何往 $S$ 里写'这一步从 Hebbian 规则换成新规则,不读背景无法理解三种更新规则的对比。

德尔塔法则 (Delta Rule)

Widrow–Hoff 1960 提出的纠错更新:写 $v_t$ 前先减去记忆当前读出 $S_tk_t$,即 $S_{t+1}=S_t(I-\beta_tk_tk_t^\top)+\beta_tv_tk_t^\top$,等价于对 $\frac12\|S_tk_t-v_t\|^2$ 做梯度下降;转移矩阵 $I-\beta_tk_tk_t^\top$ 是广义 Householder 矩阵,支持 chunkwise 并行训练,DeltaNet 家族即基于此。

克罗内克德尔塔法则正是它的标量版;只有理解向量版修正方向 $(S_tk_t-v_t)k_t^\top$ 依赖记忆内容,才能看懂本文要解决的'干扰'问题从何而来。

xLSTM/mLSTM 的 normalizer 与 stabilizer 状态

xLSTM 家族的矩阵记忆变体 mLSTM 用指数输入门写入、遗忘门保持接近 1 以保留长程信息,但必须额外维护 normalizer 和 stabilizer 两个辅助状态,对读出做比值归一化以控制记忆尺度;大网络中大小量之比易引发数值不稳定(论文图 2 压力测试中梯度范数剧烈尖峰)。

KronLSTM 的核心卖点之一是除记忆 $S_t$ 外不携带任何辅助状态、无需额外归一化也能稳定训练,mLSTM 的这些机制正是它要替代的对象。

克罗内克积与向量化 (Kronecker Product & vec)

$\mathrm{vec}(vk^\top)=k\otimes v$ 把矩阵拉直成向量。在 $a_t=k_t\otimes v_t$、$s=\mathrm{vec}(S)$ 变换下,矩阵更新规则(3)变成向量规则 $s_{t+1}=(I-\beta_ta_ta_t^\top)s_t+\beta_tc_ta_t$,与 Widrow–Hoff 德尔塔规则完全一致,且 $\|a_t\|_2=\|k_t\|_2\|v_t\|_2$。

这是方法名字'克罗内克德尔塔法则'的由来,也是全部理论分析(马尔可夫链、范数界、容量 $d_kd_v$)和 chunkwise 并行核实现都在该向量空间进行的关键。

马尔可夫链平稳分布 (Stationary Distribution)

若随机迭代 $s_{t+1}=f(s_t,\xi_t)$ 的转移核在有限二阶矩的分布上收缩,则由一般状态空间马尔可夫链理论(Meyn & Tweedie 1993),存在唯一平稳分布 $\pi$,状态分布从任意初值几何收敛到 $\pi$。

定理 1 用它证明记忆动力学收敛到唯一稳态、且稳态均值恰为最小二乘记忆 $S_{LS}$;定理 2 进一步给出期望范数一致有界——这是'免归一化也稳定'的理论根基。

chunkwise 并行训练

利用广义 Householder 转移的结构把序列切块:块内按矩阵乘法并行计算,块间做递归传递,在保持循环语义的同时高效利用 GPU。DeltaNet(Yang et al. 2024)首创该形式,KronLSTM 因转移结构与德尔塔法则同构而直接继承。

决定了 KronLSTM 能在单张 H100 上跑出 36.8–44.5K tokens/s 吞吐的工程基础,是理论新规则能落地的前提。

研究动机

Transformer 的 softmax 注意力计算成本随上下文长度二次增长、KV cache 线性增长,长上下文推理代价高昂;线性注意力把上下文压进固定大小的矩阵记忆 $S$,复杂度线性,却带来记忆控制难题。mLSTM 一类靠遗忘门加 normalizer/stabilizer 状态控制记忆规模,但比值归一化在大网络上数值不稳——论文图 2 的高学习率压力测试中 mLSTM 梯度范数剧烈尖峰、学习发散,而遗忘门若为保长程信息保持接近 1,又会放大累积。DeltaNet 一类的德尔塔法则虽有界稳定,但修正方向 $(S_tk_t-v_t)k_t^\top$ 由当前记忆内容决定,凡键与 $k_t$ 重叠的已存值都会被连带扰动;这种干扰随序列长度累积,在 16k 上下文的 RULER 检索上 GDN 的 S2/S3 准确率从 4k 时的 96.0/54.8 崩到 17.8/5.4。

本文的目标是本文目标是设计一种'最小侵入'的错误修正写入规则:每次更新只修正当前键值对 $(k_t,v_t)$ 对应的标量关联,修正仅沿单一矩阵方向进行,不触碰记忆中其他方向,从而把对已存信息的干扰压到最低;同时让记忆动力学可被数学证明——期望记忆范数在整个输入序列上一致有界、迭代构成收敛到唯一平稳分布的马尔可夫链,从而免去 normalizer/stabilizer 等额外稳定化机制;再将该规则嵌入带门控的线性注意力块构建 KronLSTM,使其在语言建模、长上下文检索、状态跟踪上达到或超越 Gated DeltaNet、KDA、Mamba-3、GDN-2 等最强线性注意力基线,尤其在 16k~32k 长上下文上扩大领先幅度。

与已有工作不同的是，现有方案在两个极端之间摇摆:Hebbian 加性写入(线性注意力、mLSTM)从不修正旧值,记忆随序列增长,须靠遗忘门或稳定器兜底;德尔塔法则做向量级纠错,修正方向依赖记忆内容,键重叠即受波及。本文的切入角度是把纠错对象从 $d_v$ 维残差 $S_tk_t-v_t$ 降到 1 维标量 $v_t^\top S_tk_t$,修正方向 $v_tk_t^\top$ 完全由当前输入对决定、不读取记忆内容;干扰项由此分解为键重叠与值重叠的乘积(性质 A9),任一因子正交即可保护旧关联,而 Hebbian 与德尔塔法则的干扰只受键重叠控制。这是首个把'定向修订'与'可证稳态'统一起来的矩阵记忆更新规则,并借克罗内克积形式化为 Widrow–Hoff 规则,从而直接继承 DeltaNet 的并行训练理论。

核心方法

直觉上,KronLSTM 把记忆写入从'整向量纠错'改成'单标量校准':每来一个 token,只检查记忆对当前键值对的关联强度 $v_t^\top S_tk_t$ 是否等于目标 $c_t$,不等就沿 $v_tk_t^\top$ 这一个方向擦除并重写,记忆其余部分原封不动。形式上,更新 $S_{t+1}=S_t-\beta_t(v_t^\top S_tk_t-c_t)v_tk_t^\top$ 恰是对损失 $\frac12(v_t^\top S_tk_t-c_t)^2$ 的一步梯度下降,可拆成擦除项与写入项。工程路线分三步:先对层输入 $x_t$ 做线性投影加短因果深度卷积、SiLU 激活和 $\ell_2$ 归一化得到 $q,k,v$;再用 sigmoid 门产生遗忘门 $F_t$、擦除门 $\beta_t$、写入门 $c_t$,并用目标 $\tilde c_t=c_t/\beta_t$ 解耦擦除与写入的幅度;最后按公式(10)更新记忆并经输出门与 RMSNorm 读出。

核心创新是克罗内克德尔塔法则:与德尔塔法则的 $d_v$ 维残差 $S_tk_t-v_t$ 不同,它只把标量关联 $v_t^\top S_tk_t$ 校准向目标 $c_t$,修正方向 $v_tk_t^\top$ 完全由输入对决定、不依赖记忆内容。这带来三个本质区别。(1)干扰结构:更新对已存对 $(k_i,v_i)$ 的改变为 $\Delta(v_i^\top S k_t)=-\beta_t(v_t^\top S_tk_t-c_t)(v_i^\top v_t)(k_t^\top k_i)$,是值重叠与键重叠的乘积,任一正交即免疫;旧规则的干扰只由键重叠控制。(2)向量化后 $s_{t+1}=(I-\beta_ta_ta_t^\top)s_t+\beta_tc_ta_t$($a_t=k_t\otimes v_t$)恰为 Widrow–Hoff 德尔塔规则,既是名字来源,也让 DeltaNet 的并行训练直接适用。(3)容量:每个三元组只施加一个线性约束 $a_t^\top s=c_t$,最多精确存 $d_kd_v$ 个标量关联,远超德尔塔法则的 $d_k$ 个值向量。

方法步骤详情

流程分四步。(1)输入处理:$x_t$ 经线性投影得到 $q,k,v$,各接短深度卷积与 SiLU,再做 $\ell_2$ 归一化,使 $\beta_t\|k_t\|\|v_t\|=\beta_t<1$,每步更新都落在定理 2 的非扩张区。(2)门控生成:擦除门 $\beta_t$ 与写入门 $c_t$ 是每头标量 sigmoid 门;遗忘门按 Mamba 式参数化 $f_t=\sigma(-(s_f)_t)^a$,$a=\exp(a_{\log})$ 为每头可学习正指数,使各头时间尺度相差数量级。(3)记忆更新(公式 10):先从右侧衰减 $S_tF_t$,再做擦除项 $-\beta_t(v_t^\top S_tF_tk_t)v_tk_t^\top$ 与写入项 $c_tv_tk_t^\top$,三个独立门分别控制遗忘/擦除/写入。(4)读出:$\tilde h_t=S_tq_t$ 经 sigmoid 输出门与逐头 RMSNorm,拼接后过输出投影 $W_{\text{out}}$。整块除记忆 $S_t$ 外无任何辅助状态,没有 mLSTM 的 normalizer/stabilizer。

技术新颖性

新颖性体现在理论与架构两层。理论上,定理 1 证明迭代构成 $\mathbb{R}^{d_kd_v}$ 上时间齐次马尔可夫链,转移核在有限二阶矩分布上收缩,存在唯一平稳分布 $\pi_\beta$,从任意初始化几何收敛,且平稳均值恰为最小二乘记忆 $S_{LS}$;定理 2 给出非扩张区($\kappa\le1$)内从第一步起的期望范数界 $E\|S_t\|_F^2\le 4E[c^2]/\lambda_{\min}(\bar A)$,瞬态因子几何衰减——因此无需任何额外归一化即可稳定训练,图 2 压力测试中高学习率下梯度无尖峰即为实证。架构上,遗忘/擦除/写入三门分离借鉴 GDN-2 解耦擦写的思想但作用于全新规则;因转移矩阵是不依赖 $c_t$ 的广义 Householder 矩阵,chunkwise 并行核可从 DeltaNet 平滑迁移,36.8–44.5K tokens/s 的实测吞吐证明工程可行性。作者声明规则与架构创意来自人类,LLM 仅辅助推导与核实现,全部证明经独立核验。

实验结果

在 1.5B 参数、100B token(FineWeb-Edu、8k 上下文)设置下有六项发现。(1)语言建模(表 1):线性 KronLSTM 六任务平均准确率 59.72,领先 KDA 的 58.12 约 1.6 分,LAMBADA 困惑度 11.72 全场最低;混合模型(5:1 线性:SWA)59.56、困惑度 11.18 同样最佳。(2)真实检索套件(表 2):线性 Avg-6 30.31(GDN 29.81),混合 37.40 最佳。(3)RULER NIAH(表 3):除 S1 外,线性模型在所有长度均为最优或次优,4k 时 S2/S3 达 98.2/73.4(GDN 96.0/54.8);混合模型 16k 时所有任务第一。(4)32k 扩展(表 4):Avg-6 60.19 vs GDN-2 58.75,困惑度 7.66 vs 8.36,五个自然文本检索任务全面领先,差距随 haystack 增大而拉大。(5)状态跟踪(图 5):parity/S3 与 GDN 相当,modulo 不及 Mamba-3。(6)吞吐(图 6):单张 H100 上 36.8–44.5K tokens/s,比 GDN-2 快 11–18%,距 KDA 差 4–8%。

**Table 1**: 1.5B/100B 模型的下游性能(线性与混合模型分开排名)

**Table 2**: 1.5B 模型在真实世界检索套件上的结果(2k 上下文)

**Table 3**: 1.5B 模型在 RULER NIAH 任务上 4k/8k/16k 评估的结果

**Table 4**: GDN-2 与 KronLSTM 的 32k 长上下文扩展结果

**Figure 4**: 32k 上下文扩展后 GDN-2 与 KronLSTM 在 NIAH 上的结果(线性与 SWA 混合)

**Figure 5**: Gated DeltaNet、Mamba-3 与 KronLSTM 的状态跟踪结果

**Figure 6**: 单张 H100-80GB 上 KronLSTM 与基线的训练吞吐对比

查看结构化数据

任务	指标	本文	基线	提升
语言建模 Avg-6(1.5B/100B,线性)	六任务平均准确率	59.72(LAMBADA ppl 11.72)	KDA 58.12 / GDN 57.73 / GDN-2 57.94	+1.6 分,困惑度低于最强基线约 0.8
语言建模 Avg-6(混合 SWA)	六任务平均准确率 + LAMBADA ppl	59.56,ppl 11.18	GDN-2 58.79,ppl 11.85	+0.77 分,ppl 降 0.67
真实检索套件 @2k(SWDE/SQuAD/FDA/TQA/NQ/DROP,线性)	Avg-6 准确率	30.31	GDN 29.81 / GDN-2 29.39	+0.5 分,且 SWDE/SQuAD/FDA 三项第一
RULER NIAH S3 @4k(线性)	准确率	73.4	GDN 54.8 / GDN-2 50.2	+18.6 分
32k 上下文扩展 Avg-6(线性)	准确率 + LAMBADA ppl	60.19,ppl 7.66	GDN-2 58.75,ppl 8.36	+1.44 分,ppl 降 0.70,检索均值 31.08 vs 29.95
训练吞吐(单张 H100-80GB,16,384 tokens/步)	tokens/s	36.8–44.5K	GDN-2(慢 11–18%),KDA(快 4–8%)	比 GDN-2 快 11–18%

局限与改进

作者承认的局限:在 RULER 中最人工化的 S1 任务(高频重复噪声句)上明显偏弱,线性 4k 仅 66.2 而 GDN 保持 100.0、GDN-2 保持 100.0,作者归因于重复键上的残差累积;实验止步于 1.5B/100B(附录 A6 另有 400M 消融),更大规模的缩放行为未验证;混合实验只采用 SWA 路线,未测含完整 softmax 注意力层的混合。我的补充观察:16k 外推时所有线性模型的绝对准确率都很低(如 KronLSTM S1 线性 16k 仅 11.6),优势是相对的而非解决了长上下文;每头单标量的 $\beta_t$ 与 $c_t$ 无法对不同值维度差异化控制;容量优势($d_kd_v$ 个标量)以特征 $\{a_i\}$ 线性无关为前提,真实数据下未必达到;吞吐仍落后 GDN 与 KDA,且 16K×1 形状下块间递归只按序列与头并行、明显降速。

独立分析的弱点

独立分析的弱点:(1)S1 失败暴露'重复键残差累积'——同一键反复出现时标量误差每次只收缩 $(1-\beta_t\|a_t\|^2)$ 倍,若 $\beta_t$ 长期偏小则残差缓慢同号累积,改进方向是检测重复键并用牛顿步 $\beta_t=1/\|a_t\|^2$ 一步命中目标;(2)写入门 $c_t$ 是每头标量,无法对不同值维度差异化写入,可推广为通道级写入门或低秩写入;(3)理论假设 i.i.d. 三元组流与固定步长,真实语言高度非平稳,稳态分布保证对实际训练的意义需要打折;(4)32k 扩展实验只与 GDN-2 一家对比,结论普适性存疑,缺 KDA/Mamba-3 参照;(5)同一键存多个正交值时读出 $S_tq_t$ 返回其和,下游如何解混未被讨论,可能需要稀疏化或迭代读出机制;(6)混合实验中 SWA 贡献了 5–9 分提升,纯线性优势的实际占比应更谨慎解读。

未来方向

作者明确提出的方向:扩大 KronLSTM 的训练规模;构建含完整 softmax 注意力层的混合模型(而非仅滑窗 SWA)。基于本文成果可延伸的研究:利用目标 $c_t$ 可输入依赖的性质做选择性记忆写入与定向遗忘(记忆编辑),$c_t=0$ 即精确删除单条关联,这对机器卸载很诱人;自动调度 $\beta_t\|a_t\|$ 在非扩张区与负特征值区($\in(1,2)$,状态跟踪能力的来源)之间切换的门控机制;把克罗内克德尔塔法则推广到张量记忆或其他状态空间模型;用马尔可夫链框架进一步刻画非平稳输入分布下的漂移界;设计多值读出机制以兑现单键多值的容量优势;继续优化 chunk 形状与核实现以缩小与 GDN 的吞吐差距。

复现评估

复现门槛中等偏低。论文声明代码与模型全部开源;数据与训练配方完全公开:FineWeb-Edu 预训练、Llama2 分词器、8k 上下文、1.5B/100B token,混合设定为窗口 2048、5:1 比例,评估基准(Arora et al. 检索套件、RULER、五项常识推理)均为公开套件。附录给出 400M 规模的全套消融(含 unit-target 与 coupled 目标变体的表 A3)与 32k 扩展细节(A7.1:75B 处分支、ProLong-64K 长文档混合续训 25B、token 对齐)。理论部分完整:附录 A3 含性质 A2/A6/A7/A9/A11 与定理 A6–A10 的证明,作者声明所有 AI 辅助推导均经人工独立核验。算力方面,吞吐测量在单张 H100-80GB 完成,但复现 1.5B/100B 预训练需数十张 H100 量级;主要风险在于其自定义 chunkwise 并行核的数值一致性与训练超参完整性,文中已声明核经朴素参考实现数值验证。

面向长上下文线性注意力的克罗内克德尔塔法则 Kronecker Delta Rule for Long Context Linear Attention