Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xiaojie Xia, Huigang Zhang, Chaoliang Zhong, Jun Sun, Yusuke Oishi

TL;DR

The paper tackles the high computational cost of full-attention transformers by proposing a two-stage approach to build task-specific hybrids that blend full and linear attention. It introduces blockwise local distillation to transfer per-block behavior from full-attention modules to linear counterparts, followed by a greedy, validation-driven layer replacement that yields an optimized hybrid in a single pass without re-training or architecture search. Empirical results across diverse datasets and backbones show that the resulting hybrids can match or exceed base model performance while delivering substantial throughput gains, with a backbone-agnostic workflow and favorable cost. This work offers a practical pathway to scalable, efficient deployment of foundation-model capabilities in resource-constrained or latency-sensitive settings.

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.

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

TL;DR

The paper tackles the high computational cost of full-attention transformers by proposing a two-stage approach to build task-specific hybrids that blend full and linear attention. It introduces blockwise local distillation to transfer per-block behavior from full-attention modules to linear counterparts, followed by a greedy, validation-driven layer replacement that yields an optimized hybrid in a single pass without re-training or architecture search. Empirical results across diverse datasets and backbones show that the resulting hybrids can match or exceed base model performance while delivering substantial throughput gains, with a backbone-agnostic workflow and favorable cost. This work offers a practical pathway to scalable, efficient deployment of foundation-model capabilities in resource-constrained or latency-sensitive settings.

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.
Paper Structure (20 sections, 2 equations, 5 figures, 5 tables, 1 algorithm)

This paper contains 20 sections, 2 equations, 5 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related works
  3. Linear Complexity Models
  4. Hybrid Architectures
  5. LLM Distillation
  6. Methodology
  7. Blockwise Local Distillation to Linear Attention Counterparts
  8. Task-Specific Hybrid Architecture Construction
  9. Experiment
  10. Experiment Setup
  11. Datasets, models and evaluation details
  12. Method details
  13. Results
  14. Performance of the best searched hybrid model
  15. Trade-off between performance and efficiency
  16. ...and 5 more sections

Figures (5)

  • Figure 1: Linear attention weight by blockwise local distillation. (a) Overall distillation framework from full attention to the linear attention. (b) BLD (blockwise local distillation), which are trained in parallel and independently.
  • Figure 2: Throughput comparison under context length of 512, 2,048, 16,384 and 65,536. The numbers above points indicate the speedup relative to base full-attention model.
  • Figure 3: Layer replacement trajectories on PubMedQA (PM) and CommonsenseQA (CQ) using Qwen2.5-1.5B and Llama3.2-3B-Instruct (28 layers each) with linear attention variants: Gated Linear Attention (GLA) and Jet-Block (JET).
  • Figure 4: Comparison of different replacement strategies.
  • Figure 5: Impact of supervised fine-tuning (SFT) on searched hybrid model.