Attentive Mask CLIP - 郑之杰的个人网站

注意力掩码对比语言-图像预训练.

paper：Attentive Mask CLIP

0. TL; DR

本文介绍了一种名为 Attentive Mask CLIP（A-CLIP）的视觉-语言预训练框架，旨在通过引入注意力掩码（attentive mask）技术提升 CLIP 模型的训练效率和性能。A-CLIP 通过保留与文本描述语义相关性高的图像 token，解决了随机掩码导致的语义信息丢失问题，并通过多视图对比学习进一步优化性能。实验表明，A-CLIP 在 ImageNet-1K 零样本分类和多个图像-文本检索任务中显著优于 CLIP 和其他改进方法，同时训练效率更高。

1. 背景介绍

近年来，大规模视觉-语言预训练模型（如 CLIP 和 ALIGN）在零样本图像分类和多模态检索任务中取得了显著进展。然而，这些模型通常需要大量的训练数据和计算资源。例如，CLIP 使用了 4 亿图像-文本对进行训练，这使得训练成本高昂。

为了提高训练效率，研究者们尝试引入图像 token 移除技术，通过减少图像编码器的计算量来加速训练。然而，随机移除图像 token 的方法会导致与文本描述相关的语义信息丢失，从而降低模型性能。为了解决这一问题，本文提出了 Attentive Mask CLIP（A-CLIP），通过注意力掩码技术选择性保留与文本描述语义相关性高的图像 token，从而在提高训练效率的同时保持或提升模型性能。

2. A-CLIP 模型

CLIP 是一种通过对比学习对齐图像和文本表征的模型，使用大量图像-文本对进行训练。其训练目标是将图像和对应的文本描述视为正样本对，而将图像与其他文本的组合视为负样本对。CLIP 使用 InfoNCE 损失函数来优化这一目标，具体公式如下：

\[L_{vl} = 0.5 \cdot (L_v + L_l)\]

其中，

\[L_v=-\frac{1}{B} \sum_{i=1}^B \log \frac{\exp(\text{sim}(e^I_i, e^T_i)/\tau)}{\sum_{j=1}^B \exp(\text{sim}(e^I_i, e^T_j)/\tau)} \\ L_l=-\frac{1}{B} \sum_{i=1}^B \log \frac{\exp(\text{sim}(e^T_i, e^I_i)/\tau)}{\sum_{j=1}^B \exp(\text{sim}(e^T_i, e^I_j)/\tau)}\]

其中 $B$ 是批次大小，$sim(⋅)$ 是余弦相似度函数，$τ$ 是可学习的温度参数，用于缩放 logits。

为了提高 CLIP 的训练效率，本文提出了一种注意力掩码方法，通过选择性保留与文本描述语义相关性高的图像 token 来避免随机掩码带来的问题。

A-CLIP 的整体框架如图所示，支持多视图输入，并通过注意力掩码技术选择性保留与文本描述相关的图像 token。此外，A-CLIP 还引入了辅助自监督学习任务，以进一步提升模型性能。

具体来说，A-CLIP 使用以下步骤实现注意力掩码：

计算相关性分数：使用 CLIP 的文本编码器提取文本特征，并计算每个图像 token 与文本特征的相关性分数：

\[S_P=\frac{1}{HL}\sum_{h=1}^H\sum_{l=1}^L\text{Softmax}(\frac{f_{q}^{lh}(\text{CLS})\cdot f_{k}^{lh}(P)}{\sqrt{C}})\]

其中，$l$ 和 $h$ 分别表示层数和注意力头索引，$f_{q}^{lh}(\text{CLS})$是 [CLS] token 的查询嵌入，$f_{k}^{lh}(P)$是位置 $P$ 的图像 token 的键嵌入，$C$ 是嵌入通道数。

选择保留的 token：根据相关性分数，选择保留与文本描述语义相关性最高的图像 token，而移除其他 token。

为了生成注意力掩码，本文引入了一个指数移动平均（EMA）版本的视觉编码器。EMA 网络不仅用于生成掩码，还可以作为图像的另一个视图，用于引入辅助对比学习任务，如 BYOL。

A-CLIP 支持多视图输入，通过为每个视图生成不同的掩码，进一步丰富输入信号。此外，A-CLIP 还引入了辅助自监督学习任务，如在线到 EMA 的对比学习任务和在线到在线的对比学习任务，以进一步提升模型性能。具体来说：

在线到 EMA 对比学习：通过 BYOL 实现，鼓励掩码视图的特征与 EMA 特征保持一致。
在线到在线对比学习：通过 SimCLR 或 SimSiam 实现，引入掩码视图之间的对比学习。

3. 实验分析

A-CLIP使用 YFCC-15M 数据集进行训练，包含 1500 万图像-文本对。在 ImageNet-1K 数据集上评估模型的零样本分类性能。在 Flickr30K 和 MS COCO 数据集上评估图像到文本（I2T）和文本到图像（T2I）检索任务的性能。

实验表明，随机掩码会导致 CLIP 性能显著下降，尤其是在零样本分类任务中。例如，随机移除 50% 的图像 token 会导致 ImageNet-1K 零样本分类准确率下降 2.6%。相比之下，A-CLIP 的注意力掩码方法不仅解决了这一问题，还在多个任务中显著提升了性能。

A-CLIP 通过引入辅助自监督学习任务（如 SimCLR 和 BYOL）进一步提升了性能。例如，添加 SimCLR 任务后，A-CLIP 在 ImageNet-1K 零样本分类任务中的准确率从 41.3% 提升到 42.8%，在 Flickr30K 和 MS COCO 检索任务中的性能也有所提升。

A-CLIP 支持多视图输入，实验表明，使用多个视图可以进一步提升性能。例如，使用 2 个视图时，A-CLIP 的性能优于使用单个视图。

A-CLIP 的注意力掩码能够有效地保留与文本描述相关的图像区域，同时过滤掉冗余的背景信息。例如，在图中可以看到 A-CLIP 保留了与文本描述相关的显著部分，而移除了无关的背景区域。