Distilling to Hybrid Attention Models via KL-Guided Layer Selection

TL;DR

This paper proposes an efficient method for selecting layers to convert in Transformer models, enabling effective distillation into hybrid attention architectures that improve inference efficiency without extensive pretraining.

Contribution

It introduces a simple layer importance scoring method for hybrid attention model distillation, outperforming existing heuristic and diagnostic-based approaches.

Findings

Layer importance scores guide effective layer conversion.

The method improves inference efficiency of LLMs.

Outperforms heuristic and diagnostic-based layer selection methods.

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.