方法原理

基于Transformer架构的大型机器学习模型近期在视觉和语言任务上展现出卓越性能，但大模型通常难以满足实时应用需求。实际系统常采用知识蒸馏技术，将大模型知识提炼到更轻量、更快速的模型中。

Transformer模型的核心特性是注意力机制，它决定已处理数据对当前数据处理的影响程度。注意力机制通常组织为多个注意力头，每个头关注数据的不同方面。

传统的大模型蒸馏方法需要将训练好的大模型（教师）的注意力头与目标轻量模型（学生）的注意力头进行一一对齐。但减少注意力头数量是降低学生模型复杂度的关键方式之一。

在AAAI 2024会议上提出的新方法，将教师模型所有注意力头的知识蒸馏到学生模型的所有注意力头中。由于学生模型的注意力头数量少于教师模型，学生模型的单个注意力头最终可能编码教师模型多个注意力头包含的信息。

实验验证

该方法在两个不同的视觉语言模型上进行评估，这些模型将图像和文本映射到同一向量空间。模型在视觉问答、图像描述和基于图像上下文的翻译任务上进行了微调，与两个最先进的基线方法进行比较。新方法在所有任务上都优于基线模型。

目标任务特点

视觉语言模型通常为每种模态单独预训练子模块，然后对整个网络进行多模态表示学习的进一步预训练。最后，预训练模型在特定任务上进行微调。

实验中学生模型仅在微调任务上进行蒸馏，但考虑了教师模型没有任何多模态预训练的情况，发现蒸馏方法在很大程度上可以弥补这种缺失。

技术实现

对于给定输入或输入集，Transformer的每个注意力头构建一个注意力图矩阵，表示输入中每个元素对其他元素的影响程度。在大型语言模型中，注意力图将文本序列的单词相互映射；在视觉模型中，注意力图可能表示图像每个区域对其他区域解释的影响。

任何矩阵的行可以连接生成单个向量，该方法依赖于注意力图的向量化（“扁平化”）版本。

蒸馏过程的损失函数包含两个组件：一是最小化教师和学生输出差异的函数；二是对齐注意力图的组件。具体来说，对于给定训练样本和教师模型中的特定注意力头，注意力图对齐损失寻求最小化教师注意力图与学生所有注意力头生成图的加权和之间的距离。

加权和的权重基于扁平化教师图与学生图之间的余弦相似度。换句话说，与教师图更相似的学生图在加权和中占比更大。在训练过程的连续步骤中，这种相似度应该增加，分配给相似学生图的权重也应该增加。

如果学生模型与教师模型具有相同数量的注意力头，且教师注意力头生成的图之间没有任何相关性，该过程可能产生类似于标准蒸馏过程的一对一映射。但该方法的核心价值在于即使学生模型的注意力头数量少于教师模型，也能保持注意力图信息。

实证表明，不同头生成的注意力图通常存在一定相关性。这些相关性可能解释了方法的成功：正是因为这些相关性，教师生成的多个注意力图可以被蒸馏到学生生成的单个图中。

应用价值

该方法为视觉语言模型的压缩和加速提供了新的技术路径，特别是在需要保持模型性能同时降低计算资源的实际应用场景中具有重要意义。