多模态大模型笔记：LLaVA-UHD

前言

多模态模型目前存在一个挑战——如何处理高分辨的图像。目前已经开源的模型包括 Qwen-VL、Yi-VL、Llava-1.6，处理图像的方式都显得不够好。常见的方式有两种：一是不论图像原始分辨率多大，直接图像缩放到同一的分辨率 448x448；二是将图像进行切片，然后把多个patch拼接到一起，组合成图像 Embedding。第一种方式会丢失很多图像细节，比如对于包含密集文本的图像，就不能很好的处理图片中的文字。第二种方式虽然保留的图像的高分辨率特征，但是图像的token会很长，训练起来比较费劲。

今天要介绍的论文是 LLaVA-UHD: an LMM Perceiving Any Aspect Ratio and High-Resolution Images. 顾名思义，这篇论文是为了解决高分辨率图像如何编码的问题，主要创新点有：

• 图像切片编码：将原始分辨率的图像分割为多个可变大小的切片进行编码，从而避免了常规模型中对图像进行强制填充或缩小造成的视觉信息丢失和变形问题。
• 视觉 Token 压缩：在视觉编码器输出的基础上添加了压缩层，将视觉 Token 进一步精简为适当长度，大大降低了大规模语言模型处理这些视觉信息所需要的计算资源。
• 图像切片的组织方式：-由于图像切片是动态生成的，论文用了一个非常简单的方式来区分图片切片的行列，使用特殊字符“,”和“\n”来标记切片在行内的位置和不同行之间的切换。

方法

处理高分辨率且具有不同宽高比的图像时，一种朴素的方法是对ViT的位置嵌入进行插值以适应目标形状，然后整体直接编码。然而，这种方法由于计算成本呈二次增长以及因超出训练分布而导致的性能下降问题而不够理想。为了解决这一挑战，论文提出了模块化视觉编码策略。基本思路是将原生分辨率的图像分割成多个大小可变的小块切片，确保每个切片的形状与ViT预训练的标准设置相差不大。通过使用大小可变的切片，LLaVA-UHD能够在无需填充或形状扭曲调整的情况下，实现对原生分辨率图像的完全适应性。

高分辨率图像的切分策略

图像分区是LLaVA-UHD模型处理高分辨率和任意宽高比图像的核心步骤之一，旨在将原生分辨率的输入图像分割成多个适配视觉编码器的小型可变大小切片。具体做法如下：

1. 确定切片数量： 给定一个高分辨率图像（宽度$W_I$，高度 $H_I$）和一个预训练的Vision Transformer（ViT）的标准分辨率（宽度 $W_v$，高度 $H_v$），首先计算所需的切片总数 $N$，即$N$等于图像总像素面积除以ViT标准分辨率下的像素面积向上取整得到的值 $N=\lceil \frac{W_I\times H_I}{W_v\times H_v}\rceil$
2. 选择最优分区： 接下来寻找最佳的图像分割方式，即将N分成m行和n列的子块，记作集合 ${\mathbb{C}}_N=\{(m,n)|m\times n=N,m\in \mathbb{N},n\in \mathbb{N}\}$ 定义一个评分函数S用来衡量特定分割方式（m, n）与ViT标准设置之间的差异，公式为：
   $$S(W_I,H_I,W_v,H_v,m,n)=-\left|\log \frac{W_I\times n}{H_I\times m}-\log \frac{W_v}{H_v}\right|$$
   评分函数S值越高，表示该分割方式越接近ViT的标准分辨率比例，因而更优。通过最大化评分函数找到最佳的行数m_和列数n来进行图像分割。
3. 保持切片分辨率接近标准比例： 通过上述算法选择合适的m和n，保证分割出来的每个图像切片的宽高比尽可能接近ViT预训练时的标准比例，这样有助于在处理非标准分辨率图像时，仍能保持模型对视觉信息的有效编码。
4. 二维位置嵌入差值： 将ViT预训练得到的一维位置嵌入序列按照其预训练顺序重新排列为二维格式 $P\in R^{q\times q\times l}$，其中M=q×q，l代表位置嵌入的维度。随后，我们对二维的P进行插值以适应由视觉编码分区策略给出的切片分辨率。论文的实验证明了在预训练阶段可以保持ViT和位置嵌入参数冻结，而在指令微调阶段更新这些参数足以获得良好的性能表现。

视觉Token压缩

如果直接采用ViT的输出作为图像的token，图像部分 embedding 就会很长。论文使用类似于Q-Former的 resampler，用一个可学习的 Query与图像特征计算 cross attention，这样就能保证无论多长的图片，输出图片特征长度都可以保持一致。

实验

实验中图片分辨率最大为 $672\times 1008$ ，视觉Token最长为 $64\times (6+1)$ .
模型训练分2个阶段。

• 阶段1：预训练。在这个阶段，仅对 resampler 进行微调，使用的数据集为CC-595K，训练周期为1个epoch，采用AdamW优化器，学习率为1e-3，并采用余弦学习率调度方案。全局批次大小设置为256。本阶段的训练成本约为使用8台A100 GPU进行5小时的训练。
• 阶段2：指令微调。在这个阶段，视觉编码器被冻结，对 resampler 和大型语言模型进行微调，所用的数据集为包含LLaVA-Instruct、TextVQA、GQA、OCR-VQA和Visual Genome的656K混合数据集。学习率设为2e-5，批次大小为128。其他设置与阶段1相同。本阶段的训练成本约为使用8台A100 GPU进行18小时的训练。

实验结果如下。

看结果而言，效果确实比大多数模型好，但 TextVQA 数据集上准确率依然只有67.7，对文字识别的准确率仍然不够。

代码

论文最大的改进应该是视觉重采样 resampler，前向计算代码如下

    def forward(self, x, tgt_size=(24,24), attn_mask=None):
        pos_embed = get_abs_pos(self.pos_embed, tgt_size)
        x = self.kv_proj(x)
        x = self.ln_kv(x).permute(1, 0, 2)
        
        N = x.shape[1]
        q = self.ln_q(self.query) # query 长度为64
        # self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        out = self.attn(
            self._repeat(q, N) + self.pos_embed.unsqueeze(1), # q
            x + pos_embed.unsqueeze(1), # k
            x, # v
            attn_mask=attn_mask)[0]
        x = out.permute(1, 0, 2)

        x = self.ln_post(x)
        x = x @ self.proj
        return x

参考

论文：LLaVA-UHD: an LMM Perceiving Any Aspect Ratio and High-Resolution Images

代码： https://github.com/thunlp/LLaVA-UHD