提出按需混合线性与softmax注意力,提升长文本处理效率
Neural Attention Search Linear: Towards Adaptive Token-Level Hybrid Attention Models
- 同一层内对不同token动态选择线性或softmax注意力
- 在LongBench等数据集上比纯线性模型准确率高3.2%以上
- 适合需要高效长序列建模的NLP任务,如文档摘要
Softmax Transformer 的二次计算复杂度在长上下文场景中已成为瓶颈。相比之下,线性注意力模型家族为更高效的序列建模提供了可行方向。这些模型将历史键值(KV)压缩为单一隐藏状态,从而在训练和推理阶段均显著降低复杂度。然而,其表达能力受限于隐藏状态大小。先前工作通过交错使用Softmax与线性注意力层,在保持表达力的同时降低计算开销。但模型效率仍受制于其Softmax注意力层。本文提出神经注意力搜索线性(NAtS-L)框架,允许在同一层内对不同token同时应用线性注意力与Softmax注意力操作。NAtS-L自动判断:对于仅具短期影响、可编码为固定大小隐藏状态的token,采用线性注意力;对于包含长期检索相关的信息、需保留以供未来查询的token,采用Softmax注意力。通过在token级别搜索最优的门控DeltaNet与Softmax注意力组合,我们证明NAtS-L能构建强大且高效的粒度级混合架构。
原文摘要 · Abstract (English)
The quadratic computational complexity of softmax transformers has become a bottleneck in long-context scenarios. In contrast, linear attention model families provide a promising direction towards a more efficient sequential model. These linear attention models compress past KV values into a single hidden state, thereby efficiently reducing complexity during both training and inference. However, their expressivity remains limited by the size of their hidden state. Previous work proposed interleaving softmax and linear attention layers to reduce computational complexity while preserving expressivity. Nevertheless, the efficiency of these models remains bottlenecked by their softmax attention layers. In this paper, we propose Neural Attention Search Linear (NAtS-L), a framework that applies both linear attention and softmax attention operations within the same layer on different tokens. NAtS-L automatically determines whether a token can be handled by a linear attention model, i.e., tokens that have only short-term impact and can be encoded into fixed-size hidden states, or require softmax attention, i.e., tokens that contain information related to long-term retrieval and need to be preserved for future queries. By searching for optimal Gated DeltaNet and softmax attention combinations across tokens, we show that NAtS-L provides a strong yet efficient token-level hybrid architecture.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。