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

基于Transformer架构的大型机器学习模型最近在视觉和语言任务上展现出卓越成果。但这些大型模型通常速度较慢，难以满足实时使用需求，因此实际系统常依赖知识蒸馏技术将大型模型的知识提炼到更精简、更快速的模型中。

Transformer模型的定义特征是其对注意力机制的依赖，该机制决定了先前看到的数据应如何影响模型对当前数据的处理。注意力机制通常组织成多个头，每个头关注数据的不同方面。

传统蒸馏方法的局限

典型的大型Transformer蒸馏涉及将经过训练的大型模型（教师）的注意力头与目标模型（学生）的注意力头进行一对一对齐。但限制注意力头数量是学生模型降低复杂度的方法之一。

在某机构人工智能促进协会的会议上，我们提出了一种替代方案：将教师模型中所有注意力头的知识蒸馏到学生模型的所有注意力头中。由于学生的头数少于教师，学生模型中的单个注意力头最终可能编码教师多个注意力头包含的信息。

方法创新：注意力图对齐蒸馏

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

目标任务设置

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

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

权重分配机制

对于给定输入或输入集，Transformer的每个注意力头构建一个注意力图，这是一个矩阵，指示输入中每个元素对其他元素的影响。在大型语言模型中，注意力图将文本序列的单词映射到自身；在决定每个新输出词时，LLM使用与该词对应的矩阵列中的注意力权重。在视觉模型中，该图可能表示图像的每个区域对其他区域解释的影响。

任何矩阵的行都可以连接起来产生单个向量，我们的知识蒸馏方法依赖于注意力图的向量版本或"扁平化"版本。

蒸馏过程的损失函数有两个组成部分：一个是寻求最小化教师和学生输出差异的函数；另一个组成部分对齐注意力图。

具体来说，对于给定的训练示例和教师模型中的给定注意力头，注意力图对齐损失寻求最小化教师注意力图与学生所有注意力头生成图的加权和之间的距离。

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

方法优势

如果学生具有与教师完全相同数量的注意力头，并且教师注意力头生成的图之间没有任何相关性，这个过程可能会产生类似于标准蒸馏过程的一对一映射。但当然，该方法的意义在于即使学生的注意力头少于教师，也能保留注意力图信息。

经验上，不同头生成的注意力图之间通常存在一些相关性。实际上，这些相关性可能解释了我们方法的成功；正是因为这些相关性，教师的多个注意力图可以被蒸馏到学生生成的单个图中。

致谢：Srikar Appalaraju, Peng Tang, Vijay Mahadevan, R. Manmatha, Ying Nian Wu。