Optimizing Knowledge Distillation in Transformers: Enabling Multi-Head Attention without Alignment Barriers

TL;DR

The paper tackles knowledge distillation for transformers plagued by head-count misalignment between teacher and student. It introduces Squeezing-Heads Distillation (SHD), a projector-free method that linearly compresses multiple teacher attention heads into a single, aligned supervision signal for the student, achieving $O(N^2)$ computational cost and preserving key attention information. SHD combines a temperature-softened attention map and a KL-based supervision loss to transfer fine-grained attention knowledge across layers and architectures, without modifying the model design. Empirically, SHD yields strong gains across language and vision tasks, including image generation, image classification, and language pretraining/fine-tuning, demonstrating improved efficiency and performance for distilling large transformers into compact models.

Abstract

Knowledge distillation (KD) in transformers often faces challenges due to misalignment in the number of attention heads between teacher and student models. Existing methods either require identical head counts or introduce projectors to bridge dimensional gaps, limiting flexibility and efficiency. We propose Squeezing-Heads Distillation (SHD), a novel approach that enables seamless knowledge transfer between models with varying head counts by compressing multi-head attention maps via efficient linear approximation. Unlike prior work, SHD eliminates alignment barriers without additional parameters or architectural modifications. Our method dynamically approximates the combined effect of multiple teacher heads into fewer student heads, preserving fine-grained attention patterns while reducing redundancy. Experiments across language (LLaMA, GPT) and vision (DiT, MDT) generative and vision (DeiT) discriminative tasks demonstrate SHD's effectiveness: it outperforms logit-based and feature-alignment KD baselines, achieving state-of-the-art results in image classification, image generation language fine-tuning, and language pre-training. The key innovations of flexible head compression, projector-free design, and linear-time complexity make SHD a versatile and scalable solution for distilling modern transformers. This work bridges a critical gap in KD, enabling efficient deployment of compact models without compromising performance.

Figures (3)

• Figure 1: Attention maps during inference on a random sample in Dolly Datatset by GPT2-XL from different heads of one layer. Each column is from the same head. We random selected three heads to visualize it. The attention patterns are very similar and contain much redundancy within one layer. Only Response attentions are kept and Instruction attentions are masked.
• Figure 2: Comparison of typical representation distillation methods and our method. (a) Typical distillation uses a projector to align student features with teacher features, introducing extra parameters. (b) Typical distillation uses relations like self-correlations to align feature dimensions. (c) Our method, SHD, uses attention maps and outputs to squeeze attention maps, aligning with the student and ensuring minimal loss of knowledge transfer.
• Figure 3: Distribution of $\alpha$ from different layers of GPT2-XL. Same color represents same head.