【大模型系列】Language-Vision Transformer(LaVIT, ICLR2024)

• Title：Unified Language-Vision Pretraining in LLM with Dynamic Discrete Visual Tokenization
• Paper：https://arxiv.org/abs/2309.04669
• Github：https://github.com/jy0205/LaVIT
• Author：Yang Jin， 北大，快手科技

1 LaVIT总结(省流版)

1.1 LaVIT是什么

LaVIT是Language-Vision Transformer的简称，作者来自于北大和快手科技。文章提出了一种将图片tokenization的方式，使图片可以与language一样，可以用一组token来表示，这样图片就可以跟language的token合并在一起作为LLM的输入，通过自回归的方式来训练网络，相当于统一了视觉和文本两种模态的输入，其中图片的token前后通过两个特殊的token[IMG]和[/IMG]来与文本的token区分。

所以文章的重点就在于如何将图片tokenization。

1.2 LaVIT能干什么？

1.2.1 文生图

通过给定文本，生成对应的图片，效果类似于SD：

1.2.2 多模图像合成(图+文->图，图+图->图)

给定一张图片，通过文本提示或者图片的提示，合成对应的图片：

1.2.3 Image Caption(图像描述)

给定一张图片，给出对于这张图的描述：

1.2.4 Visual QA(图文问答)

给定图片和文本描述的问题，给出问题的答案：

2 Motivation

当前多模态的几种架构都存在一些问题：

• Adapter-architecture：视觉特征仅被作为prompt，没有起到监督作用
• Emu：优化图像和文本目标不一致不利于统一建模

2.1 Adapter-architecture

2.1.1 Resampler(2022, DeepMind)

• Paper: https://arxiv.org/abs/2204.14198
  
  将vision-encoder输出的可变大小的时空视觉特征映射到固定数量的output-tokens（文中设定为5个）。中间使用cross-attention与learned latent queries进行交互，最后输出固定个数的token表示图片的编码。

2.1.2 Linear projection(2023, LLaVA(Large Language and Vision Assistant))

• Paper：https://arxiv.org/abs/2304.08485
  
  复用pre-trained CLIP的vision encoder，在vision-encoder之后添加一个线性映射层，将其映射到embedding空间使其维度与language model输出的维度一致。在训练时，始终冻结Vision encoder的参数，只更新映射层和LLM的权重。

2.1.3 Q-Former(2023, BLIPv2)

• Paper：https://arxiv.org/abs/2301.12597
  
  
  通过联合优化三个目标(image-text matching， image-text contrastive learning， image-grounded text generation)，使Learned queries可以提取视觉特征中与文本最相关的视觉表示。在VQA应用如下图所示，将Question也作为条件输入到Q-former中，使得提取的图像特征与问题更加相关：
  

2.2 Emu(2023, 北京人工智能研究院)

• Paper：https://arxiv.org/abs/2307.05222
  
  Emu以自回归的方式统一不同模态，图片被编码后，与文本tokens拼接在一起作为LLM的输入。训练目标是对下一个文本标记进行分类(cross entropy)和对下一个视觉特征进行回归(L2 loss, 独立的regression head)。在推理时，回归得到的视觉特征送入到微调后的SD模型中得到真实的图像。

视觉特征回归的gt是什么？ 在训练之前，先将数据(image-text pairs、带字幕的视频等)使用Causal transformer(因果变换器)将其转化为1D的causal latent embedding sequences。所以当前图片的regression的目标其实就是下一个图片的embedding

2.3 LaVIT的方式

LaVIT的方法继承了Emu的策略： 即通过自回归的方式预测下一个图像/文本的token。并且提出一个vision tokenizer将图片转化成如文本一样的序列，这样就可以统一图像和文本，使其能在同一个目标下进行优化，无需额外的框架修改。

3 LaVIT的结构

给定图像-文本对，图像被tokenize成离散的tokens，并且与文本tokens混合在一起形成多模态序列，然后LaVIT通过一个统一的生产目标进行优化（即预测下一张图片/文本的token）。

3.1 Dynamic Visual Tokenizer

给定图片，处理流程如下：

1. 先将图片划分为N(HW/ PP, P是补丁大小)个不重叠的patch；
2. N个patch先输入到一个pre-trained ViT-encoder（Vit-G/14 of EVA=CLIP）中得到N个patch features；
3. N patch-level features输入到Token selector(MLP，二分类)中，得到预测分布，并通过可微的Gumbel-softmax技巧的得到二元决策掩码M；
4. 根据决策掩码，可以将N个图像块分为drop组和retain组，并没有直接丢弃drop组，使用token merger(12个transformer blocks)，逐步将drop组信息压缩到retain组上；
5. 矢量量化：K(K=16384=2^14)大小的codebook embeddings，第4步输出的向量与codebook计算L2距离，选择最接近的codebook embeddings代替之；得到量化后的嵌入，这就是所谓的visual token。

Dynamic Visual Tokenizer的训练目标： 将量化后的嵌入输入到一个decoder中重构visual semantic features。这种设计的背后的逻辑是图像语义重建的质量取决于信息最丰富的patch，因此通过优化语义重建质量，额可以有效的更新token selector和merger，优化目标变成：

$$L_{tokenizer}=\frac{1}{N}{\Sigma }_{i=1}^N(1-cos(x_i,x_i^{rec}))+\lambda (\rho -\frac{1}{N}{\Sigma }_{i=1}^N{M}_i{)}^2$$

重建质量用余弦相似度来衡量，$\rho$是控制保留视觉标记目标的百分比，为超参数；$\lambda$设置为2，$\rho$后面的部分表示实际视觉标记保留的百分比。

3.2 基于visual token恢复图片

在训练visual tokenizer之后，基于重建的图像视觉特征，训练conditional de-noising U-Net去填充图像细节。在推理阶段，LaVIT生成的visual tokens可以被U-Net恢复为真实图像，就可以实现图片合成的功能。

3.3 Unified Generative Modeling

给定image-text pairs ，将2D的图像tokenize为有因果依赖的1D序列，并与text token组成多模态序列，通过[IMG][/IMG]来区分。为了使LaVIT能生产图像和文本，采用2种不同的串联形式：[image, text]和[text, image]。

当图像用作生成文本的条件时，直接使用token merger输出的特征，而不采用量化后特征，减少信息丢失。这对于视觉问答等细粒度多模态理解任务至关重要。

采用通用语言建模LM目标，以自回归的方式直接最大化每个多模态的序列的概率：

$$p(y)={\Sigma }_{y\in D}{\Sigma }_{i=1}^{s}log{P}_{\theta }(y_i|y_{<i})$$

4 训练细节

Step1：训练tokenizer

• image encoder frozen：ViT-G/14 of EVA-CLIP
• codebook： K=16384=2^14
• token merger和reconstruct decoder都是12个transformer blocks

数据集：LAION-400M的100M图片，batchsize=2048，p=1/3，训练50K steps。
tokenizer训练完成之后，conditonal UNet使用SDv1.5进行初始化，用相同的数据集finetune 20k steps。

Step2：Unified Vision-Language Pre-training
LLM：LLaMA-7B

• image-to-text comprehension([image, text])：93M samples（Conceptual Caption & BLIP-Capfilt）
• text-to-image synthesis([text, image])：额外100M image-text pairs（LAION-Aesthetics）

参考资料

[1] Resampler(2022, DeepMind)
[2] Linear projection(2023, LLaVA(Large Language and Vision Assistant))
[3] Emu(2023, 北京人工智能研究院)
[4] Gumbel-Softmax 完全解析