通过曲率调节查询，提升线性注意力的上下文检索能力。
Don't Read Everything: A Curvature-Conditioned Query for Linear Attention
- 在读取时引入曲率调节查询，仅聚焦高方差记忆方向。
- 在4K到20K长度外推任务中，困惑度显著降低。
- 可无缝集成至任意线性注意力模型，无需修改写入机制。
线性注意力通过维护递归快速权重状态，将Softmax注意力的二次复杂度降低，但在上下文检索和长序列任务中表现持续落后。现有方法仅从写入侧改进，如门控、增量更新或核特征映射，但读取步骤不变：所有历史键向量累加贡献，导致有效目标被大量存储向量稀释。本文借鉴Softmax几何特性，提出一种廉价的读取阶段查询收缩机制。通过对Softmax对数配分函数在各向同性注意力点进行二阶泰勒展开，得到局部二次模型，其曲率与运行中的键协方差一致，该量可通过与线性注意力状态相同的递归/分块机制维护。由此构建的线性算子在读取前沿记忆的高方差方向收缩查询。该机制称为曲率调节查询（CCQ）。CCQ仅修改读取步骤，可与任何线性注意力主干组合。应用于GLA和门控增量网络后，在训练上下文长度内及之外，均提升了困惑度、零样本下游准确率、S-NIAH检索性能，以及从4K扩展至20K的长度外推困惑度和LongBench准确率。
原文摘要 · Abstract (English)
Linear attention reduces the quadratic cost of softmax attention by maintaining a recurrent fast-weight state, but it consistently lags on in-context retrieval and long-context tasks. Existing remedies act on the write side of memory through gating, delta updates, or kernel feature maps, but the read step is left unchanged: every past key contributes additively to the output, so useful targets are diluted by the bulk of stored vectors. We borrow one specific piece of softmax's geometry to construct a cheap read-time contraction of the query. A second-order Taylor expansion of the softmax log-partition at the isotropic-attention point gives a local quadratic model whose curvature coincides with the running key covariance, a quantity that can be maintained with the same recurrent/chunkwise mechanism as the linear-attention state. The associated linear operator contracts the query along the high-variance directions of memory before it reads the state. We call this mechanism Curvature-Conditioned Query (CCQ). CCQ modifies only the read step and is composable with any linear-attention backbone. Attached to GLA and Gated DeltaNet, it improves perplexity, zero-shot downstream accuracy, S-NIAH retrieval at and beyond the training context, length-extrapolation perplexity from 4K to 20K, and LongBench accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。