当前位置：首页 > news >正文

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

news 来源：原创 2025/7/30 0:57:49

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的性能。

![[Pasted image 20241123173726.png]]

图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

![[Pasted image 20241123184930.png]]

图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_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\phi_l(y_m|y_{<m})\tag{1}$
其中 $M$ 表示预测的响应tokens的长度， $\phi_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{aligned} \mathcal{L}_\text{res}&=\sum_{m=1}^M \text{KLD}(\phi_l(y_m|y_{<m}),\phi_s(y_m|y_{<m})) \\ &=\sum_{m=1}^M \sum_{j=1}^V \phi_l(Y_j|y_{<m})\log (\frac{\phi_l(Y_j|y_{<m})}{\phi_s(Y_j|y_{<m})})\tag{2} \end{aligned}$
其中 $M$ 表示响应tokens的长度， $V$ 表示词汇空间。 $\phi_l$ 和 $\phi_s$ 表示l-MLLM和s-MLLM的参数， $\phi_l(Y_j|y_{<m})$ 和 $\phi_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 \phi_l(Y_j|y_{<k})\log (\frac{\phi_l(Y_j|y_{<k})}{\phi_s(Y_j|y_{<k})})\tag{3}$
其中 $K$ 表示视觉token的长度， $\phi_l(Y_j|y_{<k})$ 和 $\phi_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{equation} \left\{ \begin{aligned} 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{aligned} \right.\tag{4} \end{equation}$
其中 $\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||}\tag{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}\tag{6}$

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

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'\mathcal{L}_\text{res}+\beta'\mathcal{L}_\text{vis}+\gamma'\mathcal{L}_\text{rel}\tag{7}$
其中 $\mathcal{L}_\text{reg}$ 表示自回归预测loss。