Paper Reading AI Learner

Distill-then-Replace: Efficient Task-Specific Hybrid Attention Model Construction

2026-01-16 02:01:40
Xiaojie Xia, Huigang Zhang, Chaoliang Zhong, Jun Sun, Yusuke Oishi

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

PDF

https://arxiv.org/pdf/2601.11667.pdf


Tags
3D Action Action_Localization Action_Recognition Activity Adversarial Agent Attention Autonomous Bert Boundary_Detection Caption Chat Classification CNN Compressive_Sensing Contour Contrastive_Learning Deep_Learning Denoising Detection Dialog Diffusion Drone Dynamic_Memory_Network Edge_Detection Embedding Embodied Emotion Enhancement Face Face_Detection Face_Recognition Facial_Landmark Few-Shot Gait_Recognition GAN Gaze_Estimation Gesture Gradient_Descent Handwriting Human_Parsing Image_Caption Image_Classification Image_Compression Image_Enhancement Image_Generation Image_Matting Image_Retrieval Inference Inpainting Intelligent_Chip Knowledge Knowledge_Graph Language_Model LLM Matching Medical Memory_Networks Multi_Modal Multi_Task NAS NMT Object_Detection Object_Tracking OCR Ontology Optical_Character Optical_Flow Optimization Person_Re-identification Point_Cloud Portrait_Generation Pose Pose_Estimation Prediction QA Quantitative Quantitative_Finance Quantization Re-identification Recognition Recommendation Reconstruction Regularization Reinforcement_Learning Relation Relation_Extraction Represenation Represenation_Learning Restoration Review RNN Robot Salient Scene_Classification Scene_Generation Scene_Parsing Scene_Text Segmentation Self-Supervised Semantic_Instance_Segmentation Semantic_Segmentation Semi_Global Semi_Supervised Sence_graph Sentiment Sentiment_Classification Sketch SLAM Sparse Speech Speech_Recognition Style_Transfer Summarization Super_Resolution Surveillance Survey Text_Classification Text_Generation Time_Series Tracking Transfer_Learning Transformer Unsupervised Video_Caption Video_Classification Video_Indexing Video_Prediction Video_Retrieval Visual_Relation VQA Weakly_Supervised Zero-Shot