SCOUT：通过分段压缩优化 Transformer 效用以实现次二次方注意力
机器学习
2025-09-03 v1 人工智能
摘要
Transformer 在广泛的序列建模任务中展现出优异的性能，但其二次方注意力复杂度限制了对长序列的可扩展性。诸如 Mamba 和滑动窗口注意力（SWA）等线性模型通过在固定大小内存下以循环或局部操作混合 token 来解决这一问题，实现了高效推理。然而，由于无法保留来自远端 token 的详细信息，这些方法在长序列上面临性能下降的风险。我们提出 SCOUT（Segment Compression for Optimized Utility in Transformers，通过分段压缩优化 Transformer 效用），这是一种混合架构，它在固定大小的分段内局部压缩 token，并仅在这些压缩表示上应用注意力。每个 token 嵌入首先通过线性局部混合器（Mamba 或 SWA）进行丰富，以整合近期上下文。然后，每个 token 不再关注所有先前的 token，而是稀疏地关注少量压缩的检查点 token，这些 token 概括了输入历史。这种设计保留了全注意力的大部分表达能力，同时大幅降低了计算和内存开销。通过关注压缩的历史而非所有先前的 token，SCOUT 的内存占用略高于纯线性模型，但其增长率仍保持次二次方，且比全 Transformer 具有更好的可扩展性。我们分析了 SCOUT 的计算和内存效率，并在长上下文语言建模和推理任务上进行了实证评估。在相同的计算预算下，采用 Mamba 和 SWA 混合器的 SCOUT 均优于强大的长序列基线，并在 400M 和 1.3B 规模下的语言建模和常识推理任务上匹敌全注意力 Transformer。此外，我们的 SCOUT 实现了比 SOTA 模型更高的端到端吞吐量，同时在长序列基准测试上交付了可比的结果。
引用
@article{arxiv.2509.00935,
title = {SCOUT: Toward Sub-Quadratic Attention via Segment Compression for Optimized Utility in Transformers},
author = {Aref Jafari and Yuhe Fan and Benyamin Jamialahmadi and Parsa Farinneya and Boxing Chen and Marzieh S. Tahaei},
journal= {arXiv preprint arXiv:2509.00935},
year = {2025}
}