面向 Transformer 到 Dual-State Linear Attention 的即时自适应蒸馏
机器学习
2025-06-18 v3
摘要
大语言模型 (LLM) 擅长通过自注意力捕获全局 token 依赖关系，但在长输入上面临计算和内存成本的禁锢。虽然亚二次方法（如线性注意力）可降低这些成本，但常因过度强调最近 token 而导致 accuracy 下降。本文首先提出 dual-state linear attention (DSLA)，一种 novel 设计，维护两个专门的隐藏状态——一个用于保留历史语境，一个用于跟踪时效性，从而缓解线性注意力架构中常见的 short-range bias。为进一步在动态工作负载条件下平衡 efficiency 与 accuracy，我们引入 DSLA-Serve，一种在推理时逐步替换 Transformer 层为 DSLA 层的 online adaptive distillation 框架，依据基于敏感性的 layer ordering 进行引导。DSLA-Serve 使用 chained fine-tuning 策略确保每个新转换的 DSLA 层与已替换的层保持一致，从而保持整体 quality。广泛的评估在常识推理、长 context QA 和文本摘要任务上表明，DSLA-Serve 比 Llama2-7B 快 2.3 倍，比混合 Zamba-7B 快 3.0 倍，同时在下游任务中保持可比 performance。我们的 ablation 研究表明，DSLA 的 dual states 捕获了 global 与 local dependencies，解决了先前线性注意力中 historical-token underrepresented 的问题。代码已公开于 https://github.com/utnslab/DSLA-Serve。
引用
@article{arxiv.2506.09316,
title = {On-the-Fly Adaptive Distillation of Transformer to Dual-State Linear Attention},
author = {Yeonju Ro and Zhenyu Zhang and Souvik Kundu and Zhangyang Wang and Aditya Akella},
journal= {arXiv preprint arXiv:2506.09316},
year = {2025}
}