arXiv:2506.09316cs.LG2025-06ICML被引 3

用双状态线性注意力替代Transformer,实现长文本高效推理。

On-the-Fly Adaptive Distillation of Transformer to Dual-State Linear Attention

  • 设计双状态线性注意力,分别保存历史与近期信息,缓解注意力偏移。
  • 在线自适应蒸馏框架使推理速度提升2.3倍(相比Llama2-7B)。
  • 适合需要长序列处理的实时应用,如智能客服、文档摘要。

大语言模型虽能捕捉全局上下文依赖,但面对长输入时计算与内存开销巨大。现有次二次复杂度方法(如线性注意力)虽降低开销,却因过度关注近期token导致精度下降。本文提出双状态线性注意力(DSLA),通过维护两个专用隐藏状态——一个保留历史上下文,一个追踪近期变化——有效缓解线性注意力的短程偏差。进一步提出DSLA-Serve,在推理过程中按敏感度排序动态替换Transformer层为DSLA层,结合链式微调确保转换后性能稳定。在常识推理、长文本问答和文本摘要任务上的实验表明,该方法比Llama2-7B快2.3倍,比Zamba-7B快3.0倍,且下游任务表现相当。消融实验验证了双状态机制能同时捕获全局与局部依赖,解决了传统线性注意力对历史信息的忽视问题。代码已开源:https://github.com/utnslab/DSLA-Serve。

原文摘要 · Abstract (English)

Large language models (LLMs) excel at capturing global token dependencies via self-attention but face prohibitive compute and memory costs on lengthy inputs. While sub-quadratic methods (e.g., linear attention) can reduce these costs, they often degrade accuracy due to overemphasizing recent tokens. In this work, we first propose dual-state linear attention (DSLA), a novel design that maintains two specialized hidden states-one for preserving historical context and one for tracking recency-thereby mitigating the short-range bias typical of linear-attention architectures. To further balance efficiency and accuracy under dynamic workload conditions, we introduce DSLA-Serve, an online adaptive distillation framework that progressively replaces Transformer layers with DSLA layers at inference time, guided by a sensitivity-based layer ordering. DSLA-Serve uses a chained fine-tuning strategy to ensure that each newly converted DSLA layer remains consistent with previously replaced layers, preserving the overall quality. Extensive evaluations on commonsense reasoning, long-context QA, and text summarization demonstrate that DSLA-Serve yields 2.3x faster inference than Llama2-7B and 3.0x faster than the hybrid Zamba-7B, while retaining comparable performance across downstream tasks. Our ablation studies show that DSLA's dual states capture both global and local dependencies, addressing the historical-token underrepresentation seen in prior linear attentions. Codes are available at https://github.com/utnslab/DSLA-Serve.

线性注意力长序列处理模型压缩推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。