TinyCLIP

ICCV2023| TinyCLIP: CLIP Distillation via Affinity Mimicking and Weight Inheritance

• 论文链接：http://export.arxiv.org/pdf/2309.12314
• 源码链接：https://github.com/microsoft/Cream/tree/main/TinyCLIP

简介

大型视觉-语言模型（如CLIP）由于其显著的零样本迁移能力和在文本到图像生成方面前所未有的性能，最近获得了极大的关注。由于视觉和语言的极大复杂性，目前的方法往往采用大量的参数来赋予模型强大的跨模态能力。这反过来导致了这些模型在存储，内存和计算时间方面的高成本，这促使了对模型压缩的需求，使其在现实世界应用中更小更快。
知识蒸馏作为一种核心模型压缩技术，在单模态领域已经得到广泛的研究和应用，然而它在多模态领域的潜力尚未被挖掘。首先类似于CLIP的多模态模型通常由两个分支组成：图像编码器和文本编码器。在提取这种多分支模型时，至关重要的是考虑教师和学生模型中不同模态分支之间的信息交互。其次，原始CLIP模型是在4亿个图像-文本对上预训练的，这使得在计算资源有限的条件下，蒸馏成为一个巨大的挑战。
为了解决这个挑战，本文给出了一个新的跨模态蒸馏算法TinyCLIP，其中介绍了两种关键技术：亲和力模拟和权重继承。与依赖图像或文本特征进行蒸馏的方法相比，本文经验表明，在图像-文本亲和空间中提取知识更有效。具体地，利用教师模型中图像和文本嵌入的余弦相似性促进学生模型的蒸馏，使学生能够模仿教师的视觉语言特征对齐。本文称其为亲和力模仿。
为了提高蒸馏效率，本文引入了权重继承机制，一种从教师模型迁移预训练权重到对应学生模型的技术。因为继承权重为学生模型提供了更好的初始化，蒸馏过程被大大加速。权重继承的关键挑战是决定哪些权重有利。为了解决这个问题，本文介绍了两种解决方案：手动继承和自动继承。本文惊讶地发现，从教师模型中简单手动选择k维或k层权重可以为CLIP蒸馏产生令人满意的结果。另一方面，本文还介绍了可学习的掩码自动识别教师模型中最重要的权重。掩码独立施加在视觉和语言分支上，使它们能够捕捉不同模式的差异。
本文将提出的权重继承扩展到多阶段渐进模式，其中每个后续阶段自动继承来自前一阶段重要权重。本文观察到，当教师模型表现出更高的性能并于学生共享相似的架构时，权重继承可以提供更好的性能。这是因为当教师模型传递给学生时，显著的架构差异会破坏学习的权重。因此，本文将继承分解为多个阶段，允许每个阶段学生模型与前任教师模型共享更相似结构，并逐渐继承权重

本文方法

基于亲和力模仿的蒸馏

语言-图像预训练模型一般由两个分支组成：一个图像编码器抽取视觉表示和一个文本编码器抽取文本特征。通过最小化对比损失${\mathcal{L}}_0$，将视觉和文本表示线性投影到跨模态公共空间中。监督信号能被视为是单位矩阵，正样本对的对角元素设置为1，对于负样本对元素设置为0。这个监督信号忽略了负样本对之间的相似性，这阻止了模型对负样本对之间复杂关系的细微理解。因此，本文引入了亲和力模仿，使得学生模型能够从教师模型中学习负样本对之间相似性。
考虑两种类型的模仿蒸馏损失：图像-语言损失${\mathcal{L}}_{I2T}$和语言-图像损失${\mathcal{L}}_{T2I}$。前者学习教师和学生模型间基于图像-语言亲和力$A_{I2T}$对齐，表示一批中图像和所有文本描述的亲和力分数。后者表示基于语言-图像亲和力$A_{T2I}$对齐，表示比较文本描述与所有图像获得匹配分数。这两个损失的组合描述了亲和力模仿蒸馏：
$${\mathcal{L}}_{distill}={\mathcal{L}}_{I2T}+{\mathcal{L}}_{T2I}=CE(A_{I2T}^s,A_{I2T}^t)+CE(A_{T2I}^s,A_{T2I}^t)$$
$A_{I2T}$和$A_{T2I}$元素定义为：
$$\begin{array}{rl}{A}_{I2T}(i,j)& =\frac{\exp (I_i\cdot T_j/\tau )}{{\sum }_{k\in B}\exp (I_i\cdot T_k/\tau )}\\ A_{T2I}(i,j)& =\frac{\exp (I_i\cdot T_j/\tau )}{{\sum }_{k\in B}\exp (I_k\cdot T_j/\tau )}\end{array}$$
亲和力模仿允许学生模型在学习视觉语言对齐时模仿大模型行为。与以前依赖图像或文本奴特征进行蒸馏的方法相比，本文表明在该图像-文本亲和力空间中蒸馏知识更有效。

基于权重继承的蒸馏

最初的CLIP模型在4亿图像-文本对进行了32周期的预训练，耗时数千GPU天。这给蒸馏带来了巨大的成本挑战。为了提高训练效率，本文引入了权重继承，这是一种将重要权重从良好训练的大教师模型集成到小学生模型技术。权重继承关键挑战是从教师的大量权重中识别重要权重，本文提出两个选择重要权重方法：手动权重继承和自动权重继承。
手动继承
首先分析已有CLIP预训练模型的冗余性。实验发现文本编码器在深度上有更多冗余性，图像编码器在宽度有更多冗余性。基于这个发现，本文采用均匀从文本分支选择k层并选择图像分支的前k通道从大教师模型选择重要权重。本文发现简单的手动继承也能很大程度加速CLIP蒸馏。
自动继承
尽管手动继承可以被认为是一种有效的模型压缩方法，但它确实存在需要先验知识才能确定要继承的权重缺点，这种限制可能使得给技术无法应用于不同的模型。受到大语言模型中结构化剪枝启发，本文引入了可学习掩码来识别权重重要性。考虑到模态之间差异，可学习掩码被独立地施加在视觉和语言分支上。引入了整体稀疏度约束，以确保选择的重要权重数量满足压缩需要。本文以Transformer架构为例，介绍了权重继承。具体的，便准Transformer块包含多头子注意力MHSA和FFN。为了在细粒度级别上捕捉权重重要性，引入了两个掩码变量 m h e a d , m i n t ∈ { 0 , 1 } m_{head},m_{int}\in \{0,1\} mhead​,mint​∈{0,1}识别MHSA冗余注意力头和FFN冗余神经元并保留重要的部分。这两种掩码被施加在注意力头的激活和FFN的中间层上，可以描述为：
$$MHSA(X)=\sum _{h=1}^{N_H}{m}_{head}^h\cdot Att{n}_{W_Q^h,W_K^h,W_V^h,W_O^h}(X)$$
$$FFN(X)=GeLU(X{W}_U)\cdot diag(m_{int})\cdot W_D$$
$W_U\in {\mathbb{R}}^{d\times d_f},W_D\in {\mathbb{R}}^{d_f\times d}$是FFN层的参数。$m_{head},m_{int}$是MHSA和FFN层的重要性指示器。为了进一步学习Transformer嵌入维度的重要性，本文引入了额外掩码 m e m b e d ∈ { 0 , 1 } m_{embed}\in \{0,1\} membed​∈{0,1}。该掩码在所有层之间共享，因为隐藏表示中每个维度都通过残差连接连接到后续层中相应维度。
在自动继承中，由两个损失联合优化掩码和模型，损失即稀疏损失和蒸馏损失。
$$\mathcal{L}={\mathcal{L}}_{distill}+{\mathcal{L}}_{sparsity}$$