详解LexLIP: Lexicon-Bottlenecked Language-Image Pre-Training for Large-Scale Image-Text Retrieval

LexLIP：用于大规模图文检索的词典瓶颈化的语言图像预训练模型

论文地址

[2302.02908] LexLIP: Lexicon-Bottlenecked Language-Image Pre-Training for Large-Scale Image-Text Retrieval (arxiv.org)

说在开头，首先这个论文的名字就不好理解，lexicon意思为词典，lexicon-bottlenecked加上ed作为形容词。读完这篇文章，我更倾向于把lexicon理解为多模态中更让人熟知的token，这些token组成一个词典。至于这个bottlenecked，应该是和resnet中的瓶颈不一样的，resnet的瓶颈块名字取决于它模型的形状，即两头宽中间窄。本文中的瓶颈，我个人理解为，水流过瓶颈时的加速效果。

1. LexLIP和Dense检索的区别

LexLIP检索方法和传统的基于embedding的检索方法差异如下图。传统的检索方法中，使用图像和文本的encoder分别对图像和文本编码成向量，然后计算cosine相似度进行匹配。

而LexLIP的思路是，将图像和文本离散化成token以及每个token的权重，使用倒排检索的方法，将这些token列为一个词典，比如roof这个tokem链接到N个图片，并按权重排列。然后用精确匹配的方式与文本tokne之间进行检索。由于这样检索的话，roof只会和候选样本之间对比而不是CLIP那样和所有样本之间对比，大大减少了检索延迟。

2. LexLIP网络框架

LexLIP检索的底层模型是一个双流多模态模型，一侧为文本Encoder，另一侧为图像Encoder，两个Encoder都采用Transformer的形式，需要输图像或文本每个位置的预测字典中各个token的分布。最后需要在序列维度上做maxpooling，得到整个文本或图像各个词的重要度分布。以图像侧为例，先使用Transformer得到每个位置的预测token分布，维度为patch数量（m）*字典中token数量（V）。然后在序列维度进行maxpooling，得到图像整体的维度为V的token重要度分布。这个重要度用来建立倒排索引和与文本侧进行匹配度计算。

想要实现LexLIP的检索架构，关键是能够将图像和文本离散化成token，并将两个模态的字典进行对齐。为了解决这个问题，本文使用了两阶段的图文预训练：第一阶段是Lexicon-Bottlenecked Pretraining，主要目标是学习输入图像或文本的离散化token表示；第二阶段是Momentum Lexicon-Contrastive Pretraining，主要目标是对齐图像和文本这两个模态的字典空间

3. 第一阶段预训练

第一阶段预训练主要包括4个任务： self-supervised masked language modeling、两个lexicon-bottlenecked masked language modelings、in-batch lexicon-contrastive learning。第一阶段预训练的整体结构图如下。

Self-supervised masked language modeling：基础的MLM任务，mask掉一部分token后对这部分token进行预测，主要是训练文本侧Encoder的基础能力。

Lexicon-bottlenecked masked language modelings：这部分是实现表征离散化成token的关键。采用bottlenecked的训练思路，在Encoder得到图像和文本侧的token重要度分布后（和第二节的计算逻辑类似），将这个新信息输入到一个很浅层的Transformer模型中，预测被mask掉的token，让模型关注图像和文本侧最重要的token。由于直接使用token重要度分布维度太高，会利用这个重要度和token embedding做加权融合，token embedding侧不进行梯度传播，得到一个维度为embedding维度的向量作为Decoder的输入，相当于基于这个重要度做了一个连续词袋模型CBOW。

in-batch lexicon-contrastive learning：主要目的是对齐图像和文本的字典空间，利用batch内对比学习的思路，以图像和文本的离散化表示作为输入，计算对比学习loss。

4. 第二阶段预训练

第二阶段预训练主要是进一步对齐图像和文本的字典空间。大规模负样本对于对比学习的表征学习效果至关重要，因此这里采用了MoCo的思路，使用两个队列分别缓存图像侧和文本侧已经计算好的样本，并且使用动量慢更新的两个Encoder进行表征更新，就是MoCo中的方法。这部分的输入仍然是图像和文本的离散化字典表示，进行图文字典的对齐。这里论文中用了一个词叫decouple解耦，我感觉非常巧妙，将小batch和大规模负样本这两个相互矛盾的问题同时解决。MoCo本身也很巧妙，将对比学习看做一个QKV的问题。

5. 实验部分

由于大规模图文检索没有一个统一的benchmark，本文提出来了几个指标，“索引大小”对应于嵌入所有候选图像的存储要求。“Repr Byte”表示嵌入图像的存储要求。词典加权稀疏向量中的每个激活（非零）项都需要 3 个字节（索引 2 个字节，权重 1 个字节）。“QPS”对应于每秒查询，表示检索延迟，速度越高。

消融实验