视觉语言模型知识蒸馏方法优化

大型基于Transformer架构的机器学习模型近期在视觉与语言任务中展现出卓越性能。但这类大模型往往因计算速度限制难以实时部署，实际系统通常采用知识蒸馏技术将大模型知识迁移至更轻量的学生模型。

Transformer的核心特性是其注意力机制，该机制通过多个注意力头捕捉输入数据的不同特征维度。传统蒸馏方法采用教师与学生模型注意力头一对一对齐的方式，而学生模型为降低复杂度通常会减少注意力头数量。

在AAAI 2024会议上提出的新方法突破了这一限制：将教师模型所有注意力头的知识蒸馏至学生模型全部注意力头中。当学生模型头数少于教师时，单个学生注意力头可编码多个教师头的综合信息。

实验在两个视觉语言模型（VLM）上进行，这些模型将图像和文本映射到同一向量空间，并针对视觉问答、图像描述和基于图像的翻译任务进行微调。新方法在三个任务上均超越现有基线表现。

关键技术实现包含：

1. 注意力图对齐：将每个注意力头生成的注意力矩阵展平为向量，通过余弦相似度计算教师与学生注意力图的关联权重
2. 双重损失函数：包含模型输出差异最小化和注意力图对齐两个优化目标
3. 动态权重分配：训练过程中根据相似度动态调整学生注意力头的贡献权重

研究表明，不同注意力头生成的注意力图存在内在相关性，这解释了该方法在参数压缩情况下仍能保持性能的原因。即使学生模型头数显著少于教师，也能通过加权聚合方式有效保留多维度特征信息。

该方法为部署轻量级视觉语言模型提供了新的技术路径，在保持模型效率的同时最大化知识迁移效果。相关技术细节可参考论文《No head left behind: Multi-head alignment distillation for transformers》。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）
公众号二维码