Switching Linear Attention
作者: Hyun Dong Lee, Xavier Gonzalez, Nicolas Zucchet, E. Kelly Buchanan, Emily B. Fox, Scott W. Linderman
分类: cs.LG, cs.AI, cs.CL
发布日期: 2026-09-30
备注: COLM 2026
💡 一句话要点
提出切换线性注意力以解决序列建模效率问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 线性注意力 序列建模 机器学习 自然语言处理 动态选择机制
📋 核心要点
- 现有的softmax注意力在序列建模中表现优异，但其线性增长的内存需求限制了可扩展性。
- 切换线性注意力（SwiLA）通过增强表示能力，同时保持线性注意力的固定状态，解决了这一问题。
- 在多个基准测试中，SwiLA的性能强劲，缩小了与softmax注意力的差距，并在某些情况下超越了它。
📝 摘要（中文）
设计高效的序列层以实现表达能力与推理效率的平衡是现代机器学习中的核心挑战。标准的softmax注意力通过丰富的非线性令牌交互实现了优异的序列建模性能，但其关键-值缓存随着序列长度线性增长，限制了其可扩展性。线性注意力虽然能够实现恒定内存占用的高效递归计算，但其表达能力的降低往往导致建模性能不佳。本文提出了切换线性注意力（SwiLA），一种新颖的序列层，通过增强表示能力并保持线性注意力的固定大小递归状态，弥补了这一差距。SwiLA的递归来自于测试时回归框架，将状态更新规则视为线性回归模型中的在线期望最大化。在测试时，每个输出维度根据输入动态选择多个线性注意力组件。SwiLA在关联记忆、上下文语言学习和语言建模基准测试中表现出色，缩小了与softmax注意力的差距，甚至在多个设置中超越了它。
🔬 方法详解
问题定义：本文旨在解决现有注意力机制在序列建模中的效率与表达能力之间的矛盾。标准softmax注意力的内存需求随着序列长度线性增长，限制了其在长序列上的应用，而线性注意力虽然高效，但表达能力不足，导致建模性能下降。
核心思路：切换线性注意力（SwiLA）通过在固定大小的递归状态中引入多种线性注意力组件，增强了模型的表示能力。该设计允许模型在推理时动态选择最适合的线性组件，从而提高了建模的灵活性和准确性。
技术框架：SwiLA的整体架构包括状态更新规则的在线期望最大化过程，利用线性回归模型的混合形式来实现。模型在训练和测试阶段的状态更新机制有所不同，测试时根据输入动态选择线性组件。
关键创新：SwiLA的主要创新在于其将线性注意力与动态选择机制结合，克服了传统线性注意力的表达能力不足问题。这一设计使得模型在保持高效性的同时，能够实现更复杂的序列建模任务。
关键设计：在SwiLA中，模型的参数设置和损失函数经过精心设计，以确保在动态选择过程中能够有效地更新状态。此外，网络结构采用了多线性回归组件，以支持不同输出维度的灵活选择。通过这些设计，SwiLA在多个任务中展现了优越的性能。
🖼️ 关键图片
📊 实验亮点
在多个基准测试中，SwiLA的表现显著优于传统线性注意力，缩小了与softmax注意力的性能差距。在某些任务中，SwiLA甚至超越了softmax注意力，展示了其在关联记忆、上下文语言学习和语言建模等领域的强大能力。
🎯 应用场景
切换线性注意力（SwiLA）在自然语言处理、时间序列预测和其他需要高效序列建模的领域具有广泛的应用潜力。其高效的内存使用和增强的表达能力使其适合处理长序列数据，能够在实际应用中提高模型的性能和响应速度。未来，SwiLA可能会推动更多高效模型的设计，促进智能系统的进一步发展。
📄 摘要（原文）
Designing expressive sequence layers with efficient inference remains a central challenge in modern machine learning. Standard softmax attention achieves excellent sequence modeling performance through rich nonlinear token interactions, but it requires a key-value cache that grows linearly with sequence length, limiting its scalability. Linear attention enables efficient recurrent computation with a constant memory footprint, yet its reduced expressivity often yields inferior modeling performance. We introduce Switching Linear Attention (SwiLA), a novel sequence layer that bridges this gap by enhancing representational capacity while retaining the fixed-size recurrent state of linear attention. We derive the SwiLA recurrence from the test-time regression framework, casting the state update rule as online expectation-maximization in a mixture of linear regressions model. At test time, each output dimension dynamically selects among multiple linear attention components based on the input. Across associative recall, in-context language learning, and language modeling benchmarks, SwiLA shows strong performance and narrows the gap to softmax attention, even surpassing it in several settings.