提出一种新型注意力机制,能更好适应动态数据并提升模型泛化能力。
Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression
- 基于测试时回归思想设计,融合线性与Softmax注意力优势。
- 在多种任务中表现优于现有方法,尤其在非平稳数据上效果显著。
- 支持高效计算,适合大规模模型部署,适合需要动态适应的场景。
Transformer架构在多个领域取得显著成功。尽管高效的Softmax注意力替代方案已广泛研究,但基于理论洞察、更具表达力的机制(即使计算成本更高)仍相对未被探索。本文提出局部线性注意力(LLA),源自非参数统计学,从测试时回归的角度出发。首先,通过偏差-方差权衡分析,证明LLA在关联记忆任务中优于线性与Softmax注意力。其次,针对其计算复杂度Θ(n²d)和Θ(nd²),提出两种内存高效原语;进一步设计FlashLLA,一种硬件友好的分块算法,实现现代加速器上的可扩展并行计算。此外,还实现并优化了专用推理内核,显著降低内存开销。最后,在测试时回归、上下文回归、关联回忆和状态跟踪任务上进行实证验证,结果表明LLA能有效适应非平稳性,在测试时训练和上下文学习中超越强基线,展现出良好的可扩展性与大规模模型应用前景。代码已公开于https://github.com/Yifei-Zuo/Flash-LLA。
原文摘要 · Abstract (English)
Transformer architectures have achieved remarkable success in various domains. While efficient alternatives to Softmax Attention have been widely studied, the search for more expressive mechanisms grounded in theoretical insight-even at greater computational cost-has been relatively underexplored. In this work, we bridge this gap by proposing Local Linear Attention (LLA), a novel attention mechanism derived from nonparametric statistics through the lens of test-time regression. First, we show that LLA offers theoretical advantages over Linear and Softmax Attention for associative memory via a bias-variance trade-off analysis. Next, we address its computational challenges and propose two memory-efficient primitives to tackle the $Θ(n^2 d)$ and $Θ(n d^2)$ complexity. We then introduce FlashLLA, a hardware-efficient, blockwise algorithm that enables scalable and parallel computation on modern accelerators. In addition, we implement and profile a customized inference kernel that significantly reduces memory overheads. Finally, we empirically validate the advantages and limitations of LLA on test-time regression, in-context regression, associative recall and state tracking tasks. Experiment results demonstrate that LLA effectively adapts to non-stationarity, outperforming strong baselines in test-time training and in-context learning, and exhibiting promising evidence for its scalability and applicability in large-scale models. Code is available at https://github.com/Yifei-Zuo/Flash-LLA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。