论文阅读 [精读]-MM1: Methods, Analysis & Insights from Multimodal LLM Pre-training

最近 Apple 出了自己的 30B 多模态大模型，涌现出了多模态的 in-context learning 效果，论文里一句 "even better" 让我想到库克那个嗓音…… 作者说明了很多在训练中收获到的经验教训，这是我最近几个月看的写法最清楚的一篇论文。正好借此讲讲多模态大模型：目前学界大火的 VLM，到底是怎么跑的？

Introduction

作者团队来自苹果，看来苹果说的 "今年 wwdc 上大模型" 真有希望了？

在 pretrain models 领域，大家一直都想把之前的所有任务的数据整合到一起，把之前一堆独立模型做的事情统一到一个模型里面，由此诞生了 LLM 技术。最近大家关注的重点开始转向多模态：除了文本领域的任务，能不能同时把和图片理解相关的任务也统一进去？具体来说，如果输入同时含有图片文本，多个图片文本交错排列，模型能不能在理解图片和文本的基础上，输出文本。这就是 MLLM (multimodal LLM) 或者 VLM (Vision-Language Model)

虽然学界、业界在这方面的探索不少，但已有的工作基本都不够开放：要么就是闭源模型给个 API，要么就是开源模型只给个权重。即使有些工作会开源数据，但是没有人会开放讨论他们对于模型结构的选择、对于训练数据的选择、对于训练超参数的选择以及背后的原因。

作者在本篇工作中开放的讨论了所有的细节，结论简单来说就是以下几点：

在模型结构方面，下面几个要素的重要性递减：vision-encoder 的分辨率、vision-encoder 大小。VL-connector 选择对于最终的表现几乎没有影响
在数据方面，主要有以下几种数据：text-only 数据、image-caption 数据、互联网文本 - 图片交错数据，作者发现：
1. 互联网文本图片交错数据和 text-only 数据对于 in-context Learning、few-shot Learning 至关重要
2. image-caption 数据对于 zero-shot Performance 至关重要

最后，从一系列 insight 出发，作者 scaling 了训练过程，训练了 3B、7B、30B 的 VLM，甚至在 3B 和 7B 规模下尝试了 top2 的 MoE 架构，并在各个层面上达到了 SOTA 的效果

除了没开源 weight，基本全给你了……

VLM 长啥样？

目前的 VLM 基本都是分三个部分，如上图的左图：

image-encoder，负责把图片编码成为 embedding，同时尽可能不要损失信息
VL-connector，负责把 image-encoder 的输出做一些转换，比如 MLP，对齐到 LLM 的 word embedding 空间
decoder-only LLM：把上一步的输出直接作为多个 Token 的 Word embedding 输入，然后跑后面的 LLM next-token-prediction。会在 text 输出的部分计算 loss

Image-encoder

image-encoder 一般也是 transformer，是一个 ViT。他首先会把图片按照从左到右、从上到下切成不同的小正方形，打成多个 patch，每个 patch 的分辨率都是比如 14x14。接下来每个 patch 会使用一个卷积层编码成为一个 vector，然后多个 patch vector 拼在一起当成多个 “token” 的 word embedding，加上 position embedding 后直接进入 transformer encoder，输出会是每个 patch 一个 hidden state。

这个 encoder 有两种训练方法：

contrastive loss: 大名鼎鼎的 CLIP。找到一大堆互联网上的挨着的 (图片 - 文本) 对作为正样本，然后随机其他的对作为负样本，然后用刚才的 ViT 作为 image-encoder，用另一个比如 T5 之类的作为 text encoder，拿到两个 embedding，跑对比学习的 loss：希望正样本 - 正样本之间的 cosine 相似度大于负样本 - 正样本之间的 cosine 相似度
reconstructive loss：这个就是传说中的 VAE。用 ViT 编码完了以后会把图片变成一个 embedding，为了保证这个 embedding 尽量信息无损，会后面再接上一个 image-decoder 去预测原来的图片长什么样子。预测的越准 loss 就越小。为了防止模型去 memorize 每个 image 的样子，还会有个辅助 KL loss 去保证所有 image 的 embedding 在 embedding 空间的分布尽可能均匀

这两种方法有个特点：都是无监督的，只要有一大堆数据就能起效果。前者需要 (图片 - 文本) 对，但有个好处是数据可以自己合成，可以人工在 text-caption 上加上各种丰富准确的细节来提高训练的要求，进而增强模型。后者只需要一大堆图片，在一些 dense 的任务上表现更好

作者在这里做了消融实验，最左边那列 AIM 是 reconstruction loss，CLIP 是对比学习 loss。architecture 基本都是 ViT，H 的参数量比 L 大。image Res 代表的是训练的时候的训练数据的图片分辨率都是多少。Data 是指用什么数据做的训练。

明显可以发现：

image res 这个变量对最终的效果差距最大
在 CLIP 中加入合成数据 (VeCap)，会提升模型表现
两种 loss 基本都不咋影响效果，都差不多

作者最终选了 CLIP + 高清 encoder

Vision-Language Connector

这一部分就是输入 image-encoder 的 output，转换到 Word embedding 空间，这里面有个至关重要的问题：一张图片应该转换到 LLM 里的多少个 token？

显然 token 更多，保留的细节更多，模型的效果理应更好
然而，当 token 多的时候，尤其如果是多个图片 (8 个、16 个)，很多个 token 其实训不起来，资源消耗太恐怖了

作者在这里消融实验了一个 image 对应 64、144 token 两种情况，然后尝试了 224、336 分辨率的 image-encoder，以及不同的 pooling 策略，大致有这么几种实现：

Average Pooling: 由于 Vit 的 token 数大于目标 token 数量，就把相邻几个 patch hidden-state 取平均数，变成一个 embedding，让数量变得一样了。再给每个 token 过一个 MLP，作为 LLM 的一个 token embedding
Attention Pooling：作者觉得 vit 的输出可能和 llm 的 word-embedding 不在一个子空间，需要一个 Attention 层变换一下。于是就加一个额外的 Attention 层，初始化 k 个 query 向量，然后用 key、value 变换阵把 ViT 的输出对齐过去 (Attention 输出在 sequence-length 维度上和 query 的长度一致)。这样 Attention 的输出就是 k 个 token embedding，然后 k 可以取 64、144，就对齐过去了
C-abstractor: 把输出用某种卷积层操作 (ResNet) 转换到 word-embedding 空间

作者对比了一圈发现：vision token 和 image resolution 最重要，几种不同的 pooling 策略几乎没区别。即使在一个 test set 上好，在另一个上可能会更差

Pretrain-data

最后，作者探索了 pretrain 数据对于结果的影响。这里有个歧义："pretrain" 并不是真正的预训练。作者这里实际上会选取一个预训练好的 LLM 作为 text-only decoder，然后选取一个在 CLIP 对比学习 loss 上预训练好的 image-encoder。把他俩用一个新初始化的 VL-connector 拼在一起。把这个叫做 "start training"。所以这里探讨的 “pretrain data” 单指 VLM 启动训练后的 "pretrain"

VLM 的 pretrain-data 基本分为三种：

image-caption: 正常 VLM 的预训练数据，也就是刚才 CLIP 里的图片文本对。作者会输入图片，让模型预测 caption
interleaved image-text: 从互联网上爬下来的语料。这里面都是图文交错的数据，作者直接把文本提取出来，在正常图片的位置放上图片，让模型预测所有文本的 next-token-prediction
text-only：正常 LLM 预训练的数据，基本上是去掉图片后的互联网语料，以及 github 之类地方爬下来的各种代码，以及找到的一大堆各个学科的教材啥的

其中，image-caption 数据分为正常版和合成版本。所谓合成数据，就是用 GPT-4v 或者其他什么开源模型 (这里说的 VeCap 用的 vicuna，是一个 Llama 在 gpt-3.5 的输出上做过 SFT 的版本)，要求他给图片生成一些非常详细的 caption，包括里面所有 Object 以及他们之间的位置空间语义承接关系之类的。

如果大家看过前几天写的 DALL.E 3 的笔记 dall.e 3 阅读笔记，就知道正常的 html 里来的图片文本对的 alt-text 质量有多差，所以需要重新 clean，去重写等等

作者在这一部分做了最多的实验，因为在 LLM 领域现在大家也普遍认为数据是影响结果最重要的因素，大致发现了如下结论：

interleaved data is instrumental for few-shot and text- only performance, while captioning data lifts zero-shot performance。作者猜测是因为互联网交错数据天生具有一些 in-context Learning 的性质
text-only data helps with few-shot and text-only performance
Careful mixture of image and text data can yield op- timal multimodal performance and retain strong text performance.
Synthetic data helps with few-shot learning.

最后，作者公布了花了很多钱试出来的配方： caption /interleaved/ text-only = 5:5:1

MM1

pretrain

由上面的结论，作者把训练做了 scaling，选择了如下方案：

CLIP loss 的 ViT-H，378x378 分辨率
144 token，c-abstractor 模式的 VL-connector
5:5:1 的 VLM pretrain-data

作者在 9M, 85M, 302M, 1.2B 几个规模上做了超参数搜索，并认为学习率 $η$ 应该和模型规模的对数成反比，由此预测了 30B 模型的最优 Learning rate=2.2e-5。然后 weight-decay $λ$ 是 lr 的 1/10 $η = \exp (- 0.4214 \ln (N) - 0.5535) λ = 0.1 η$ 另一面，作者还尝试了 MoE 架构，用了 Mixtral-8x7B 那种经典的 top2 router，在 3B 和 7B LLM 规模上做了实验。3B 是每两层加一个 64 选 2 的 FFN，7B 是每四层有一个 32 选 2 的 FFN。这些 FFN 都是从最开始的 dense FFN 初始化的，随着 router 的训练逐渐变得差异化了。为了稳定训练，作者还加了一些平衡 routing 的 loss

这方面，似乎 Mixtral 8x7B report 说他们没加任何其他 loss，也可以正常训练？不知道是不是 VLM 领域有新的 bug

作者发现，这样训练出来的 MM1 系列模型，在所有规模上基本都是目前最好的 VLM。然后 MoE 可以让 Performance "even better"

SFT

目前的公认解决方案，都是在 Pretrain 完以后，把能找到的学术任务拿过来造一个大号的数据集，让模型做一下 supervised finetuning，再找一大堆 chat 数据让模型学着遵循人类乱七八糟的指令（比如 "只能回复 emoji"），最后再在测试集上测试效果

类似的，作者的 SFT 数据构成三部分：

学术数据集转换来的数据
GPT-4v 生成的给予图片的 qa chat 数据
text-only，训练正常 LLM 时使用的 SFT 数据

所有数据掺在一起，SFT 阶段随机喂数据。他们尝试了更高级的策略，发现效果基本没提升，就用了这个最简单的招

high resolution

测试的时候，还有问题：1) 测试数据很多有巨多图片输入，或者需要 8-shot in-context 样本，这就需要最少 8 个图片。2) 很多图片清晰度巨高，比如 3840x2160 分辨率。

这方面，有一个解决方案：Sub-image decomposition, 比如对于 1344×1344 的图片，作者分成 5 张图丢给模型，每个都是 672×672。第一张图是降采样的，后面四张分别是左上角到右下角的局部

这里有个 trick：这里的复杂度不是 $n^{2}$ , 而是 $n^{4}$ 。一个大图片会变成 $n^{2}$ 个 token，后面 LLM 本身又以 $n^{2}$ 的复杂度跑前向，所以这个算法把一个大图变成 5 个小图实际上是省的。

作者通过这个方法可以支持任意分辨率的图片，实际上最高到 1792x1792 还是会崩。然后作者对比了 SFT 阶段的效果，发现了几个核心：

image 分辨率对效果至关重要，如果输入的分辨率本身不高，怎么训练效果都不行
VLM 的预训练对效果很重要。作者把预训练各个阶段的 ckpt 都拿过来做了一次 SFT 实验，发现 VLM 预训练越久，对应的 SFT Performance 也越好

另外，作者发现即使 SFT 数据全都是单张图片输入的，MM1 还是具有多张图片输入时的 in-context 推理能力。并且，在 Pretrain 阶段探索到的 insight，对于 SFT 阶段的 Performance 仍然成立。最后，作者在附录里展示了很多 MM1 的 case，感觉这个模型的效果是真的很不错

要是开源就好了

我的思考

我印象中 MM1 是第一个正式描述自己的多张图片、图文交错情况下的 in-context Learning 能力的模型。我感觉这个能力对于多模态的 Agent 很重要，目前我看到的一堆 multi-modal 的 ReAct 算法，都是把 history 用文本的形式拼到 prompt 了…… 感觉不是真正的 multi-modal agent
另外最近有一些别的模型，比如 Fuyu 是一个 decoder-only 架构，图片的 patch embedding 不用走 vit，直接在一个大 decoder 里走一次。代价是要重新预训练，好处是跑的很快…… 不知道哪种才是最好的方案，我猜目前这些方案是因为大家可以直接拿一个 llama 过来当 LLM，效果比较稳？