arXiv:2502.07436cs.CV2025-02被引 5

无需对齐头数即可实现跨模型注意力知识迁移

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

  • 通过线性近似压缩多头注意力图,动态融合教师头输出
  • 在图像分类与生成任务中超越现有方法,性能达新高
  • 无额外参数、无需修改结构,适合高效部署小模型

Transformer中的知识蒸馏常因教师与学生模型注意力头数不一致而受阻。现有方法要么要求头数相同,要么引入投影层弥补维度差异,限制了灵活性与效率。本文提出挤压头蒸馏(SHD),通过高效的线性近似压缩多头注意力图,实现不同头数模型间的无缝知识传递。SHD无需额外参数或结构修改,动态将多个教师头的组合效应近似为更少的学生头,保留细粒度注意力模式的同时减少冗余。在语言(LLaMA、GPT)和视觉(DiT、MDT、DeiT)生成与判别任务上的实验表明,SHD优于基于logit和特征对齐的基线方法,在图像分类、图像生成、语言微调及预训练任务中均达到当前最优表现。其核心创新包括灵活的头压缩机制、无投影器设计以及线性时间复杂度,使SHD成为适用于现代Transformer的通用、可扩展蒸馏方案。

原文摘要 · Abstract (English)

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.

知识蒸馏注意力机制模型压缩Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。