[Paper] Distilling to Hybrid Attention Models via KL-Guided Layer Selection

Source: arXiv - 2512.20569v1

Overview

The paper proposes a lightweight recipe for turning a standard soft‑max‑based Transformer into a hybrid attention model that mixes soft‑max and linear attention layers. By using a KL‑guided layer‑importance score derived from a tiny amount of generic text data, the authors can automatically pick which layers to replace with cheaper linear‑attention variants, then distill the original model into the hybrid architecture with a proven RADLADS pipeline. The result is a faster inference model that retains most of the original performance without the cost of pre‑training a new LLM from scratch.

Key Contributions

• KL‑guided layer‑selection: Introduces a simple, data‑efficient scoring method that ranks Transformer layers by their “importance” using a small KL‑divergence‑based probe.
• Hybrid attention recipe: Shows how to interleave soft‑max and linear attention layers based on the importance scores, rather than using naïve uniform spacing.
• Integration with RADLADS distillation: Combines the layer‑selection step with an existing distillation pipeline (attention weight transfer, hidden‑state alignment, KL distribution matching, short finetune).
• Empirical superiority: Demonstrates that the KL‑guided selection outperforms uniform‑ratio heuristics and more complex diagnostic‑dataset methods on standard NLP benchmarks.
• Efficiency‑focused: Achieves comparable accuracy to the full‑softmax model while reducing inference latency and memory footprint.

Methodology

1. Layer‑importance scoring

   • Train a tiny “probe” model on a few thousand generic sentences (e.g., Wikipedia snippets).
   • For each Transformer layer, compute the KL divergence between the layer’s output distribution and a reference distribution (the original soft‑max output).
   • Higher KL indicates that the layer contributes more unique information, so it should stay soft‑max; lower KL suggests it can be safely replaced with linear attention.

2. Hybrid architecture construction

   • Sort layers by importance.
   • Replace the lowest‑scoring layers with linear‑attention equivalents, preserving the original ordering of the remaining soft‑max layers.
   • The resulting architecture alternates between the two attention types in a data‑driven pattern.

3. Distillation via RADLADS

   • Attention weight transfer: Copy the original soft‑max attention maps into the hybrid model where possible.
   • Hidden‑state alignment: Align intermediate representations using an L2 loss.
   • KL‑based distribution matching: Encourage the hybrid model’s output logits to match the teacher’s distribution (KL loss).
   • Finetuning: Run a short (often < 1 epoch) finetune on the same generic text to polish performance.

4. Evaluation

   • Benchmarks include GLUE, SQuAD, and language modeling perplexity.
   • Compare against baselines: uniform‑ratio hybrid models and diagnostic‑dataset‑driven selections.

Results & Findings

• Latency improves by ~30 % relative to the teacher while losing less than 1 % absolute GLUE performance.
• The KL‑guided selection consistently beats uniform and diagnostic baselines across all tasks, confirming that the importance scores capture the right trade‑off between speed and accuracy.
• Memory usage drops proportionally with the number of linear‑attention layers, enabling deployment on edge GPUs and CPUs.

Practical Implications

• Faster inference for LLM‑powered services – Companies can retrofit existing Transformer models (e.g., BERT, GPT‑2) with linear attention where it matters least, cutting latency without re‑training from scratch.
• Cost‑effective scaling – Linear attention reduces the quadratic cost of self‑attention, making it feasible to serve larger batch sizes or run on cheaper hardware (e.g., CPU‑only inference).
• Simplified model compression pipeline – The KL‑guided scoring requires only a few thousand unlabeled sentences, meaning teams can apply it to any proprietary model without building task‑specific diagnostic datasets.
• Compatibility with existing distillation tools – Since the method plugs into the RADLADS pipeline, developers can reuse their current distillation scripts and only add the layer‑selection step.
• Potential for on‑device NLP – Hybrid models fit better into the memory constraints of mobile or embedded devices, opening doors for offline assistants, smart‑camera text analysis, etc.

Limitations & Future Work

• Scope of linear attention variants – The study focuses on a specific linear‑attention implementation; other variants (e.g., Performer, Linformer) may behave differently.
• Small probe dataset – While efficient, the KL‑guided scores could be sensitive to the choice of generic text; more diverse probing might improve robustness.
• Task‑specific fine‑tuning – The paper evaluates mainly on general benchmarks; real‑world downstream tasks (e.g., code generation, dialogue) may require additional finetuning to close the performance gap.
• Scalability to massive LLMs – Experiments are on models up to ~350 M parameters; extending the method to multi‑billion‑parameter LLMs could surface new challenges (e.g., memory for KL scoring).

Future work could explore adaptive layer‑selection that dynamically switches attention types at inference time, integrate other efficient attention mechanisms, and test the approach on truly large‑scale LLMs used in production.

Authors

• Yanhong Li
• Songlin Yang
• Shawn Tan
• Mayank Mishra
• Rameswar Panda
• Jiawei Zhou
• Yoon Kim

Paper Information

• arXiv ID: 2512.20569v1
• Categories: cs.CL, cs.AI
• Published: December 23, 2025
• PDF: Download PDF