用可学习的中心点控制注意力范围,实现高效又不降性能的长序列处理。
Why Attend to Everything? Focus is the Key
- 通过少量可学习中心点划分令牌组,仅组内令牌互注意力。
- 在124M至70B模型上均零退化,124M时困惑度30.3优于全注意力的31.4。
- 适合想提升长序列效率且不想重训练的用户,兼容所有主流架构。
标准注意力随序列长度呈二次方增长。现有高效方法虽降低复杂度,但迁移到预训练模型时常导致困惑度或下游准确率下降。本文提出Focus,通过少量可学习中心点(每层仅148K参数)作为门控机制:仅同组令牌间进行长程注意力。Focus可组合使用——仅训练中心点,冻结原有权重。实验表明,在124M至70B参数规模、五种注意力架构的预训练模型上,使用Focus均无性能退化。令人惊讶的是,124M规模下稀疏注意力困惑度为30.3,优于全注意力的31.4;7B规模从头训练时达到13.82,接近全注意力的13.89。此外,top-k分组带来2倍加速,且质量更优。结合FlashAttention分解,100万令牌下实现8.6倍加速,无需自定义核函数。
原文摘要 · Abstract (English)
Standard attention scales quadratically with sequence length. Efficient attention methods reduce this O(n^2) cost, but when retrofitted into pretrained models, they often degrade perplexity, downstream accuracy, or both. We introduce Focus, a method that learns which token pairs matter. Focus adds a small set of learnable centroids--as few as 148K parameters per layer--that act as gates: only token pairs belonging to the same centroid group attend to each other over long ranges. Focus is composable: it can be added to any pretrained model by training only the centroids while keeping all original weights frozen. Experiments show that composing Focus onto pretrained models yields zero degradation on downstream benchmarks across model sizes from 124M to 70B parameters and five attention architectures. Surprisingly, sparse Focus attention outperforms full attention at 124M scale (30.3 vs. 31.4 perplexity) and matches full attention when trained from scratch at 7B scale (13.82 vs. 13.89). Focus is also fast: top-k group membership gives a 2x speedup with better quality than the original pretrained model. Using our FlashAttention decomposition, Focus achieves an 8.6x speedup at 1M tokens without custom kernels.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。