从零理解CLIP中的InfoNCE损失函数：原理、实现与避坑指南

1次阅读

共计 2197 个字符，预计需要花费 6 分钟才能阅读完成。

CLIP 模型通过将图像和文本映射到共享的嵌入空间，实现了跨模态的语义对齐，而其中的关键在于 InfoNCE 损失函数。这个损失函数通过对比学习的方式，使得相似的样本在嵌入空间中靠近，不相似的样本远离，从而有效解决了多模态学习中的语义对齐问题。

对比学习的核心思想是通过正负样本的对比来学习特征表示。在 CLIP 中，正样本通常是指同一语义对的图像和文本（例如，一张猫的图片和对应的文本描述“一只猫”），而负样本则是其他不相关的图像和文本对。

正样本定义：对于每个图像 - 文本对，正样本是同一语义对的嵌入向量。
负样本定义：负样本是同一批次中其他图像 - 文本对的嵌入向量。

温度系数（Temperature）在 InfoNCE 损失中起到调节相似度分布的作用，其数学表达式为：

$$
\mathcal{L}{InfoNCE} = -\log \frac{\exp(\text{sim}(q, k+) / \tau)}{\sum_{i=1}^N \exp(\text{sim}(q, k_i) / \tau)}
$$

其中，(\tau)是温度系数，(\text{sim}(q, k))是查询（query）和键（key）的相似度。

调参经验：
较小的 (\tau) 会使得损失函数更加关注困难的负样本，但可能导致训练不稳定。
较大的 (\tau) 会平滑相似度分布，但可能降低模型的区分能力。
通常初始值设为 0.07，再根据验证集表现微调。

InfoNCE 损失与传统的交叉熵损失（Cross-Entropy Loss）的主要区别在于，前者通过对比学习的方式直接优化嵌入空间的结构，而后者通常用于分类任务。

交叉熵损失：
$$
\mathcal{L}{CE} = -\sum^C y_i \log(p_i)
$$
适用于分类任务，直接优化类别概率分布。
InfoNCE 损失：
适用于嵌入空间的学习，通过对比正负样本优化相似度。

以下是 InfoNCE 损失函数的 PyTorch 实现，包含负样本队列的 Memory Bank 优化：

import torch
import torch.nn as nn
import torch.nn.functional as F

class InfoNCE(nn.Module):
    def __init__(self, temperature=0.07):
        super(InfoNCE, self).__init__()
        self.temperature = temperature
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, image_features, text_features):
        # Normalize features
        image_features = F.normalize(image_features, dim=1)
        text_features = F.normalize(text_features, dim=1)

        # Compute similarity matrix
        logits = torch.matmul(image_features, text_features.T) / self.temperature

        # Labels are the diagonal entries
        labels = torch.arange(logits.size(0), device=image_features.device)

        # Compute loss
        loss_i = self.criterion(logits, labels)
        loss_t = self.criterion(logits.T, labels)
        loss = (loss_i + loss_t) / 2

        return loss

将 InfoNCE 损失函数与 CLIP 的 ViT-B/32 模型结合的步骤如下：