【论文笔记】LLaVA-KD: A Framework of Distilling Multimodal Large Language Models

Abstract

大语言模型(Large Language Models, LLM)的成功，使得研究者为了统一视觉和语言的理解去探索多模态大预言模型(Multimodal Large Language Models, MLLM)。
但是MLLM庞大的模型和复杂的计算使其很难应用在资源受限的环境，小型MLLM(s-MLLM)的表现又远不如大型的MLLM(l-MLLM)。

基于上述提到的问题，本文提出了全新的LLaVA-KD框架，将l-MLLM的知识转移到s-MLLM。具体地，本文提出：

• 多模态蒸馏(Multimodal Distillation, MDist)，减小l-MLLM和s-MLLM之间的视觉-文本输出分布的差异；
• 关系蒸馏(Relation Distillation, RDist)，迁移l-MLLM对视觉特征之间相关性的建模能力。
  本文还提出了三阶段训练方案，充分挖掘s-MLLM的潜力：
• 预训练时蒸馏，对齐视觉文本表示；
• 有监督地微调，使模型具备多模态理解；
• 微调时蒸馏，进一步迁移l-MLLM的能力。

本文方法在不改变s-MLLM结构的情况下显著提升了其性能。

github仓库

Introduction

本文从研究各种训练策略的角度出发，在不改变模型架构的情况下，探索提高、s-MLLM的性能。

图1：为了训练s-MLLM，(a)已有方法遵循两步训练，包括预训练(Pre-Training, PT)和监督微调(Supervised Fine-Tuning, SFT)；(b)本文的LLaVA-KD提出了三步训练，包含蒸馏预训练(Distilled Pre-Training, DPT)来对齐视觉文本表示、SFT来提升模型的多模态理解能力、蒸馏微调(Distilled Fine-Tuning, DFT)来转移l-MLLM的能力；©本文将LLaVA-KD和其他sota的MLLM在五个热门的多模态benchmark上进行比较。

如图1所示，已有的s-MLLM遵循两步训练策略，包含PT和SFT。PT阶段将视觉特征投影到文本嵌入空间；SFT阶段增强模型的理解和推离能力。
但是s-MLLM的模型容量很小，很难像l-MLLM那样捕获复杂的知识。本文将研究如何借助蒸馏提升s-MLLM的训练。

3 LLaVA-KD

图2：LLaVA-KD的总图，包含三阶段训练。1) DPT，向l-MLLM对齐视觉和文本信息。2) SFT，为s-MLLM带来多模态理解能力。3)DFT，向s-MLLM迁移l-MLLM的能力。在DPT和DFT中应用MDist，使用RDist来使得s-MLLM捕获视觉信息的复杂关系。

3.1 Composition of Distilled MLLM Architecture

图2左侧展示了MLLM的蒸馏过程，包含l-MLLM作为教师模型，和s-MLLM作为学生模型，分别包含三个部分：

Frozen Visual Encoder：用于获得强力的视觉特征。给定输入图像$X_v\in {\mathbb{R}}^{H\times W\times 3}$，排序成2D patches $P_v\in {\mathbb{R}}^{N_p\times S_p^2\times 3}$，其中$S_p$和$N_p$表示patch的大小和数量。最后的transformer层将$P_v$变成$Z_v\in {\mathbb{R}}^{N_p\times C}$，其中特征维度为$C$。教师和学生都使用相同的Frozen Visual Encoder。

Visual Projector：包含两个MLP层，带有激活函数GELU，将$Z_v$映射到文本嵌入空间$H_v\in {\mathbb{R}}^{N_p\times D}$，其中$D$是嵌入空间维度。

Large Language Model (LLM)：用于实现对视觉和语言信息的统一认识。给定视觉嵌入的多模态输入$H_v$和文本嵌入$H_t$，LLM将二者的连接$H=[H_v,H_t]$作为输入，生成输出 y = [ y p , y v , y r ] = { y t } t = 1 T y=[y_p,y_v,y_r]=\{y_t\}_{t=1}^T y=[yp​,yv​,yr​]={yt​}t=1T​，其中$y_p,y_v,y_r$分别代表prompt、视觉和响应tokens，$T$代表所有预测token的长度。本文将教师和学生的LLM分别称为l-LLM和s-LLM。

3.2 Training Scheme of Teacher Model L-MLLM

Pre-Training：Visual Encoder和l-LLM冻结，只有Projector被优化，用于对齐视觉和文本特征。训练过程中，使用图像-描述对，对应的目标公式表示为：
$${\mathcal{L}}_{\text{reg}}=-\sum _{m=1}^M\log {\varphi }_l(y_m|y_{<m})\text{\hspace{1em}(1)}$$
其中$M$表示预测的响应tokens的长度，${\varphi }_l(y_m|y_{<m})$表示响应tokens $y_m$的分布基于先前预测$y_{<m}$的条件。

Supervised Fine-Tuning：该阶段保持Visual Encoder的冻结，旨在联合优化Projector和l -LLM，以增强教师模型l-MLLM的理解和教学跟随能力。训练过程中，利用高质量的对话数据集，训练目标${\mathcal{L}}_{SFT}$如Eq.1所示。

3.3 Framework of LLaVA-KD

3.3.1 MLLM-Oriented KD Strategy

Multimodal Distillation (MDist)：考虑到MLLM本质上是利用LLM进行多模态信息理解和推理，我们沿用LLM的朴素蒸馏方法，即利用KL散度(KLD)对响应预测进行蒸馏。训练目标可以定义为：
$$\begin{array}{rl}{\mathcal{L}}_{\text{res}}& =\sum _{m=1}^M\text{KLD}({\varphi }_l(y_m|y_{<m}),{\varphi }_s(y_m|y_{<m}))\\ & =\sum _{m=1}^M\sum _{j=1}^V{\varphi }_l(Y_j|y_{<m})\log (\frac{{\varphi }_l(Y_j|y_{<m})}{{\varphi }_s(Y_j|y_{<m})})\end{array}\text{\hspace{1em}(2)}$$
其中$M$表示响应tokens的长度，$V$表示词汇空间。${\varphi }_l$和${\varphi }_s$表示l-MLLM和s-MLLM的参数，${\varphi }_l(Y_j|y_{<m})$和${\varphi }_s(Y_j|y_{<m})$表示由l-MLLM和s-MLLM预测的词汇$Y_j$出现在token $y_m$的概率。

同时，视觉表征对于LLM的多模态理解也至关重要。因此，进一步优化教师和学生输出视觉分布之间的KLD：
$${\mathcal{L}}_{\text{vis}}=\sum _{k=1}^K\sum _{j=1}^V{\varphi }_l(Y_j|y_{<k})\log (\frac{{\varphi }_l(Y_j|y_{<k})}{{\varphi }_s(Y_j|y_{<k})})\text{\hspace{1em}(3)}$$
其中$K$表示视觉token的长度，${\varphi }_l(Y_j|y_{<k})$和${\varphi }_s(Y_j|y_{<k})$分别表示由l-MLLM和s-MLLM预测的词汇$Y_j$出现在token $y_k$的概率。

本文在DPT阶段用MDist来对齐s-MLLM中的视觉和语言特征，加强了s-MLLM的理解能力。

Relation Distillation (RDist)：为了使学生模型能够捕获视觉信息中的复杂关系，本文从LLM输出的视觉tokens中构造自相关矩阵。通过优化矩阵之间的相似性，学生模型继承了教师模型理解视觉tokens之间错综复杂关系的能力。为了达到这个目的，首先计算自相关矩阵如下：
$$\begin{array}{c}\{\begin{array}{rl}{R}_v^s& =y_v^s\otimes y_v^s\in {\mathbb{R}}^{N_p\times N_p}\\ R_v^t& =y_v^t\otimes y_v^t\in {\mathbb{R}}^{N_p\times N_p}\end{array}\end{array}\text{\hspace{1em}(4)}$$
其中$\otimes$表示矩阵乘法，$y_v^s$和$y_v^t$表示学生和教师的视觉logits，$N_p$表示视觉token的数量。目标是最大化$R_v^s$和$R_v^t$的余弦相似度：
$${\mathcal{L}}_{\text{rel}}=1-\text{Cos}(R_v^s,R_v^t)=1-\frac{R_v^s\cdot R_v^t}{||R_v^s||||R_v^t||}\text{\hspace{1em}(5)}$$
用RDist可以进一步提升s-MLLM在DPT和DFT阶段的视觉表达能力。

3.3.2 Three-stage Distillation Scheme

Distilled Pre-Training (DPT)：该阶段的主要目的是将视觉特征投射到文本嵌入空间。在LLaVA-KD中，使用蒸馏过程来像l-MLLM一样更好地对齐视觉和文本信息。

具体地，冻结visual encoder和s-MLLM中的LLM，只优化projector。在训练过程中，通过MDist最小化学生模型和教师模型在视觉和反应的输出分布上的差异。
为了优化这个目标，可以进一步促进投影的视觉特征与文本嵌入的对齐。此外，我们利用RDist来增强视觉特征的质量，使学生模型能够借鉴教师模型处理复杂视觉信息的能力。

总的来说，除了优化自回归预测结果，还使用了MDist和RDist：
$${\mathcal{L}}_{\text{DPT}}={\mathcal{L}}_{\text{PT}}+\alpha {\mathcal{L}}_{\text{res}}+\beta {\mathcal{L}}_{\text{vis}}+\gamma {\mathcal{L}}_{\text{rel}}\text{\hspace{1em}(6)}$$

Supervised Fine-Tuning (SFT)：这个阶段遵循l-MLLM训练阶段的通用SFT过程(Sec.3.2)。通过联合训练Projector和l-LLM，使模型具有推理能力和指令跟踪能力。训练目标由Eq.1定义，表示为${\mathcal{L}}_{\text{SFT}}^{\prime }$。

Distilled Fine-Tuning (DFT)：该阶段的主要目标是进一步增强s-MLLM的理解和推理能力。具体来说，采用了MDist和RDist相结合的蒸馏策略，冻结了Visual Encoder，优化了Projector和sLLM。通过使用MDist，可以对s-MLLM中的小规模s-LLM进行充分优化，从而更好地模拟大规模l-LLM的推理能力。和RDist可以进一步促进s-MLLM学习l-MLLM的视觉表征。

总体训练目标可以表示为：
$${\mathcal{L}}_{DFT}={\mathcal{L}}_{\text{reg}}+{\alpha }^{\prime }{\mathcal{L}}_{\text{res}}+{\beta }^{\prime }{\mathcal{L}}_{\text{vis}}+{\gamma }^{\prime }{\mathcal{L}}_{\text{rel}}\text{\hspace{1em}(7)}$$
其中${\mathcal{L}}_{\text{reg}}$表示自回归预测loss。