Distilling to Hybrid Attention Models via KL-Guided Layer Selection
This incremental improvement addresses the inference efficiency of LLMs for developers and researchers by reducing computational costs without full retraining.
The paper tackles the problem of efficiently converting pretrained softmax attention Transformers into hybrid architectures by proposing a simple layer selection method using importance scores from generic text data, which outperforms existing heuristics and specialized dataset approaches.
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.