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

2025-02-11

AI Summary

To address the alignment difficulty arising from mismatched numbers of attention heads between teacher and student models in Transformer knowledge distillation, this paper proposes Squeezing-Heads Distillation (SHD)—a lightweight, architecture-agnostic distillation method requiring no projection layers or structural modifications. Its core innovation is a novel flexible head-compression mechanism that dynamically approximates multiple teacher attention heads into fewer student heads via linear attention mapping, achieving zero additional parameters, zero architectural dependency, and O(L) linear time complexity. SHD jointly distills both logits and intermediate features. Extensive experiments across LLaMA/GPT (language) and DeiT/DiT/MDT (vision) models demonstrate that SHD consistently outperforms existing baselines, attaining state-of-the-art performance on image classification, image generation, and language model fine-tuning tasks.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning on the Edge & Model Compression
• Data Mining & Knowledge Management: Data Compression

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Enabling knowledge transfer between mismatched head count transformers
• Eliminating alignment barriers without additional parameters or modifications
• Compressing multi-head attention maps via efficient linear approximation

Innovation

• Compressing multi-head attention maps efficiently
• Eliminating alignment barriers without extra parameters
• Dynamic approximation of teacher heads effect