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

TL;DR

This paper introduces Squeezing-Heads Distillation (SHD), a novel method that allows knowledge transfer between transformer models with different attention head counts without additional parameters, improving efficiency and flexibility.

Contribution

SHD enables seamless multi-head attention distillation across models with varying head numbers by compressing attention maps through linear approximation, eliminating the need for projectors or architectural changes.

Findings

SHD outperforms existing KD methods on language and vision tasks.

It achieves state-of-the-art results in image classification and language pre-training.

The method is scalable and computationally efficient with linear-time complexity.

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.