Attention Editing: A Versatile Framework for Cross-Architecture Attention Conversion

📄 arXiv: 2604.05688v1 📥 PDF

作者: Zhen Cheng, Hao-Bo Yang, Wan-Yi Huang, Jin-Long Li

分类: cs.CL, cs.AI

发布日期: 2026-04-07

💡 一句话要点

提出Attention Editing框架，用于跨架构注意力机制转换，提升长文本处理效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱六：视频提取与匹配 (Video Extraction) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 注意力机制 模型编辑 知识蒸馏 大型语言模型 长文本处理

📋 核心要点

  1. 长文本和长序列生成场景下，Key-Value缓存的内存和带宽成为大语言模型推理成本的主要瓶颈。
  2. Attention Editing框架通过可学习的注意力模块替换原有注意力，并使用渐进式蒸馏进行训练，实现跨架构的注意力转换。
  3. 实验表明，该方法在Qwen3-8B和Qwen3-30B-A3B模型上，能够在保持性能的同时显著提高效率。

📝 摘要（中文）

本文提出了一种名为Attention Editing的通用框架，用于将已训练的大型语言模型（LLMs）转换为使用新的注意力机制，而无需从头开始重新预训练。该框架通过可学习的目标模块替换原始注意力，并使用渐进式蒸馏进行训练，包括：（1）使用中间激活监督的逐层教师强制优化，以防止冷启动误差累积；（2）基于下一个token分布的模型级蒸馏，可以选择性地通过弱特征匹配进行正则化。我们在两种不同的目标架构——MLA和GateSWA（一种门控混合SWA设计）上实例化了该框架，并将其应用于Qwen3-8B和Qwen3-30B-A3B。结果模型在保持竞争性能的同时，显著提高了效率，证明了大规模注意力转换的可行性和鲁棒性。值得注意的是，实验是在Ascend 910B集群上进行的，提供了一个在国内硬件上进行实际训练的案例研究。

🔬 方法详解

问题定义：现有方法在将MLA和SWA等新型注意力架构集成到现有大型语言模型中时，对源和目标注意力模块施加了细粒度的结构要求，这在实际部署中难以满足。因此，如何高效地将已训练好的LLM迁移到新的注意力架构，成为了一个亟待解决的问题。

核心思路：Attention Editing的核心思路是通过注意力编辑，将预训练好的LLM的注意力机制替换为新的、更高效的注意力机制，而无需从头开始重新训练整个模型。这种方法利用知识蒸馏，将原始模型的知识迁移到新的注意力模块中。

技术框架：Attention Editing框架主要包含两个阶段：(1) 逐层教师强制优化：使用中间激活监督，以防止冷启动误差累积。目标注意力模块在每一层都模仿原始注意力模块的输出。(2) 模型级蒸馏：在下一个token分布上进行模型级蒸馏，可以选择性地通过弱特征匹配进行正则化，以进一步提升性能。

关键创新：Attention Editing的关键创新在于其渐进式蒸馏策略，它避免了从头开始训练的需要，并且能够有效地将知识从原始模型迁移到新的注意力模块中。此外，该框架的通用性使其能够应用于不同的目标注意力架构，例如MLA和GateSWA。

关键设计：在逐层教师强制优化阶段，使用均方误差（MSE）损失函数来最小化目标注意力模块和原始注意力模块之间的中间激活差异。在模型级蒸馏阶段，使用交叉熵损失函数来最小化目标模型和原始模型之间的下一个token分布差异。弱特征匹配正则化通过最小化目标模型和原始模型之间的特征差异来进一步提升性能。具体参数设置（如学习率、蒸馏温度等）根据具体模型和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Attention Editing框架在Qwen3-8B和Qwen3-30B-A3B模型上取得了显著的效率提升，同时保持了具有竞争力的性能。具体而言，使用MLA和GateSWA作为目标注意力架构的模型，在推理速度和内存占用方面均优于原始模型。该框架在Ascend 910B集群上的成功应用，也验证了其在国内硬件上的可行性。

🎯 应用场景

Attention Editing框架可广泛应用于各种需要长文本处理的大型语言模型，例如机器翻译、文本摘要、对话生成等。通过替换为更高效的注意力机制，可以显著降低模型的推理成本，使其更容易部署在资源受限的设备上。该技术还有助于加速新型注意力架构的探索和应用。

📄 摘要（原文）

Key-Value (KV) cache memory and bandwidth increasingly dominate large language model inference cost in long-context and long-generation regimes. Architectures such as multi-head latent attention (MLA) and hybrid sliding-window attention (SWA) can alleviate this bound, but integrating them into existing models remains difficult. Prior methods impose fine-grained structural requirements on both source and target attention modules, which cannot meet the feasible requirement in practical deployment. We present Attention Editing, a practical framework for converting already-trained large language models (LLMs) with new attention architectures without re-pretraining from scratch. Attention editing replaces the original attention with a learnable target module and trains it using progressive distillation, consisting of (1) layer-wise teacher-forced optimization with intermediate activation supervision to prevent cold-start error accumulation, and (2) model-level distillation on next-token distributions, optionally regularized by weak feature matching. We instantiate the framework on two different target--MLA and GateSWA, a gated hybrid SWA design, and apply it to Qwen3-8B and Qwen3-30B-A3B. The resulting models maintain competitive performance while delivering substantial efficiency improvements, demonstrating that large-scale attention conversion is both feasible and robust. Notably, experiments are conducted on an Ascend 910B clusters, offering a practical training case study on domestic hardware.