长上下文的军备竞赛里,线性注意力是省钱的那条路:把 softmax 注意力的二次复杂度压成线性,代价是记忆状态变成一个不断被改写的矩阵。问题在于"怎么改"——写入新东西时,旧记忆被扰动多少。NX-AI 团队的新论文 KronLSTM(arXiv:2610.26834,10 月 19 日提交,六位作者包括 Sepp Hochreiter)给出了一个更"外科手术"的答案。
写入即干扰:线性注意力的老毛病
NX-AI 在仓库 README 里用一张动画把三种写入行为摆在一起:传统 mLSTM 每写一个新值,会扰动矩阵里所有已存的值;Gated DeltaNet 会把与新 key 共享维度的旧值整批抹掉;而 KronLSTM 只修改与新 key-value 对重叠的那一小块,其余记忆原样保留。对长上下文来说这不是细节——上下文越长,写入次数越多,被反复扰动的旧记忆越难保真。
Kronecker Delta Rule:只改该改的
论文的核心是给关联矩阵记忆设计的新更新规则。KronLSTM 把 xLSTM 家族扩展出一个记忆修正机制:更正目标关联时,只在矩阵空间的单一方向上修改记忆,相比 delta-rule 更新对已存信息的干扰更小。理论上,这套记忆动力学可以描述为一条收敛到稳态分布的马尔可夫链,记忆矩阵的期望范数在整个输入序列上保持一致有界——换句话说,不需要给记忆状态外挂额外的归一化或稳定化机制,它自己就是稳的。实验侧的论文口径是:与现有线性注意力架构(包括 (Gated) DeltaNet 家族)相比有竞争力、且常常胜出,长上下文基准上尤其明显。
开源清单:FLA 兼容,四个 1.5B 模型
工程侧给得相当完整。代码是 Triton 内核加 PyTorch 层,直接兼容 flash-linear-attention(FLA)生态,训练用 chunk_kronlstm(chunkwise WY/UT 形式,与 GDN-2、KDA 同套路),解码用 fused_recurrent_kronlstm;仓库本身就是 FLA 的一个快照加上 KronLSTM 实现,作者表示未来会合入 FLA 主线。预训练模型有四个,全部 1.5B 参数、100B token 训练、Llama-2 tokenizer:两个纯 KronLSTM(24 层,8K 与 32K 上下文两档),两个混合版(20 层 KronLSTM + 4 层滑窗注意力,窗口 2048、不用位置编码)。训练配置也交代得细:8K 模型在 FineWeb-Edu 上练;32K 版从 75B token 处分支,再用 ProLong-64k 长文档混合喂 25B token。许可证 MIT,在 H100/A100 上以 PyTorch 2.13 + Triton 3.7.1 开发测试。
所以呢
两点判断。其一,线性注意力的竞争已经从"能不能压成线性"进入"记忆改写精度"的深水区:今年 6 月 NVIDIA 的 Gated DeltaNet-2 把擦与写拆开,这次的 Kronecker Delta Rule 更进一步,连修正都限定在单一方向——架构层的微创新在回潮,谁的记忆更接近"只动该动的",谁的长上下文就站得住。其二,对从业者这是即插即用的:FLA 生态用户改个层名就能跑,作者在 HF 社区里预告了更快的内核正在路上。需要保持清醒的是,1.5B/100B token 是研究级配置,"常常胜出"是论文自报口径,离生产级验证还有距离——但它给了长上下文这条赛道一个值得盯住的新基线。
参考:arXiv:2610.26834;github.com/NX-AI/KronLSTM;huggingface.co/collections/NX-AI/kronlstm