中文

Local Linear Attention:一种用于测试时回归的线性与 Softmax 注意力的最优插值

机器学习 2025-10-03 v1 人工智能

摘要

Transformer 架构在各个领域取得了显著的成功。 虽然对 Softmax 注意力的高效替代方案已被广泛研究,但仍相对缺乏基于理论见解(即使计算成本更高)的更具表达性的机制。 本文通过测试时回归的视角,通过非参数统计学方法来弥补这一差距,我们提出了 Local Linear Attention (LLA),这是一种新型注意力机制。 首先,我们展示了 LLA 相对于线性注意力和 Softmax 注意力在关联记忆方面具有理论优势,这通过偏差-方差权衡分析得出。 下而,我们解决其计算挑战,提出两种高效原语以处理 Θ(n2d)\Theta(n^2 d)Θ(nd2)\Theta(n d^2) 的复杂度。 我们然后引入 FlashLLA,这是一种硬件高效的块状算法,使其能够在现代加速器上实现可扩展且并行的计算。 此外,我们实现并剖析了一个定制的推断内核,显著减少了内存开销。 最终,我们在测试时回归、上下文回归、关联记忆和状态跟踪任务上经验性地验证了 LLA 的优势与局限性。 实验结果表明,LLA 能够有效地适应非平稳性,在测试时训练和上下文学习中超越强大的基线,并在大规模模型的可扩展性和适用性方面展现出有前景的证据。 代码已在 https://github.com/Yifei-Zuo/Flash-LLA 提供。

关键词

引用

@article{arxiv.2510.01450,
  title  = {Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression},
  author = {Yifei Zuo and Yutong Yin and Zhichen Zeng and Ang Li and Banghua Zhu and Zhaoran Wang},
  journal= {arXiv preprint arXiv:2510.01450},
  year   = {2025}
}