Abstract
Transformer architectures deliver state-of-the-art accuracy via dense full-attention, but their quadratic time and memory complexity with respect to sequence length limits practical deployment. Linear attention mechanisms offer linear or near-linear scaling yet often incur performance degradation. Hybrid models that integrate full and linear attention layers promise a balance between efficiency and expressiveness, but face two major challenges: training such hybrid models from scratch is computationally expensive, and manually designing the optimal placement of attention types is highly nontrivial. We address both issues by first transferring weights from the pretrained full-attention modules to its linear attention counterparts through blockwise local distillation, and second, introducing a greedy layer replacement strategy that iteratively substitutes full attention blocks with linear ones while monitoring validation performance on the target task. This yields a task-specific hybrid model in a single efficient pass, without costly re-training or neural architecture search, and can be applied to any pretrained full-attention backbone for diverse downstream tasks.
Abstract (translated)
翻译如下: 基于Transformer的架构通过密集全注意力机制提供了最先进的准确性,但它们的时间和内存复杂度与序列长度呈二次关系,这限制了其实用部署。线性注意力机制则提供线性的或接近线性的缩放比例,尽管如此,往往会导致性能下降。混合模型将全注意力层和线性注意力层结合在一起,旨在平衡效率和表达能力,但面临着两大挑战:从头开始训练此类混合模型计算成本高昂,而手动设计最佳注意力类型放置方案极其复杂且不切实际。 为了解决这两个问题,我们首先通过分块本地蒸馏的方式,将预训练的全注意力模块权重转移到其线性注意力对应模块中。其次,我们引入了一种贪婪层替换策略:该策略在监测目标任务验证性能的同时,迭代地用线性注意力模块替代全注意力模块。这可以在一次高效的遍历过程中生成特定任务的混合模型,而无需耗时的重新训练或神经架构搜索,并且可以应用于任何预训练的全注意力骨干网络以应对各种下游任务。
URL
https://arxiv.org/abs/2601.11667