Distilling to Hybrid Attention Models via KL-Guided Layer Selection

TL;DR

The paper tackles the challenge of preserving long-context recall when distilling pretrained softmax Transformers into hybrid attention models, by introducing a KL-guided layer-selection method that scores each layer via a one-swap procedure within a RADLADS distillation framework. The method first distills to an all-linear student, then reintroduces each layer to softmax to compute a marginal utility score I(ell) based on reductions in teacher–student KL divergence, and finally selects the top-K layers to keep as softmax. Empirical results on Qwen2.5-3B-Instruct and Llama-3.2-3B-Instruct show substantial gains over uniform or heuristic baselines on long-context recall benchmarks like RULER and SWDE, with clear gains at low softmax budgets and transferable layer sets across different linear-attention architectures. The work also reveals that selected softmax layers cluster in depth and that the selection can be stabilized early, enabling efficient practical deployment of hybrid models with improved inference efficiency and preserved recall capabilities.

Abstract

Distilling pretrained softmax attention Transformers into more efficient hybrid architectures that interleave softmax and linear attention layers is a promising approach for improving the inference efficiency of LLMs without requiring expensive pretraining from scratch. A critical factor in the conversion process is layer selection, i.e., deciding on which layers to convert to linear attention variants. This paper describes a simple and efficient recipe for layer selection that uses layer importance scores derived from a small amount of training on generic text data. Once the layers have been selected we use a recent pipeline for the distillation process itself \citep[RADLADS;][]{goldstein2025radlads}, which consists of attention weight transfer, hidden state alignment, KL-based distribution matching, followed by a small amount of finetuning. We find that this approach is more effective than existing approaches for layer selection, including heuristics that uniformly interleave linear attentions based on a fixed ratio, as well as more involved approaches that rely on specialized diagnostic datasets.

Figures (7)

• Figure 1: Performance of a sliding-window attention model (distilled from Qwen2.5-3B-Instruct) across different window sizes on RULER and commonsense reasoning tasks.
• Figure 2: Performance on recall-intensive vs. commonsense tasks as the number of full-attention layers is varied for Qwen2.5-3B-Instruct (top) and Llama-3.2-3B-Instruct (bottom). Recall ability is highly sensitive to the softmax budget, while commonsense reasoning is not.
• Figure 3: Performance comparison of various layer selection methods on RULER (top) and SWDE (bottom) for distilling Qwen2.5-3B-Instruct (left) and Llama-3.2-3B-Instruct (right) into hybrid GDN-based models. Performance is plotted against the percentage of softmax layers retained. The dashed line indicates the performance of the all-softmax teacher model.
• Figure 4: Jaccard similarity of top-K layer selections between GDN and GLA variants over the selection pass. Llama shows higher agreement, suggesting its layer importance is less student-dependent.
• Figure 5: The evolution of RULER performance during the Stage-2 selection process for Qwen2.5-3B-Instruct.