Local Linear Attention:一种用于测试时回归的线性与 Softmax 注意力的最优插值
机器学习
2025-10-03 v1 人工智能
摘要
Transformer 架构在各个领域取得了显著的成功。 虽然对 Softmax 注意力的高效替代方案已被广泛研究,但仍相对缺乏基于理论见解(即使计算成本更高)的更具表达性的机制。 本文通过测试时回归的视角,通过非参数统计学方法来弥补这一差距,我们提出了 Local Linear Attention (LLA),这是一种新型注意力机制。 首先,我们展示了 LLA 相对于线性注意力和 Softmax 注意力在关联记忆方面具有理论优势,这通过偏差-方差权衡分析得出。 下而,我们解决其计算挑战,提出两种高效原语以处理 和 的复杂度。 我们然后引入 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}
}