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

1次阅读

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

CLIP 模型通过将图像和文本映射到共享的嵌入空间，实现了跨模态的语义对齐，而其中的关键就在于 InfoNCE 损失函数。这个函数负责衡量不同模态样本之间的相似性，是多模态对比学习的核心。本文将带你从零开始理解 InfoNCE 的原理，并手把手教你如何实现和优化它。

对比学习的核心思想是通过拉近正样本对的相似度，同时推远负样本对的相似度。在 CLIP 中：

正样本对 ：匹配的图像 - 文本对（如狗图片和“一只狗”描述）
负样本对 ：不匹配的图像 - 文本组合（如狗图片和“一辆汽车”描述）

InfoNCE（Noise Contrastive Estimation）的公式为：

$$\mathcal{L} = -\log\frac{e^{s_{i,j}/\tau}}{\sum_{k=1}^N e^{s_{i,k}/\tau}}$$

其中 $s_{i,j}$ 是样本 i 和 j 的相似度得分，$\tau$ 是温度系数，N 是负样本数量。

温度系数 $\tau$ 控制着相似度得分的分布：

$\tau$ 较小时：分布更尖锐，模型会更关注困难样本
$\tau$ 较大时：分布更平滑，学习更稳定但可能丢失细节

经验值：CLIP 原文使用 0.07，实际建议从 0.05-0.2 范围网格搜索。

特性	InfoNCE	交叉熵
样本关系	对比学习	分类任务
计算方式	所有负样本参与	仅正确类参与
适用场景	无监督 / 自监督	监督学习

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

class InfoNCE(nn.Module):
    def __init__(self, temp=0.07):
        super().__init__()
        self.temp = temp
        self.cos_sim = nn.CosineSimilarity(dim=-1)

    def forward(self, image_emb, text_emb):
        # 归一化嵌入向量
        image_emb = F.normalize(image_emb, dim=-1)
        text_emb = F.normalize(text_emb, dim=-1)

        # 计算相似度矩阵
        sim = self.cos_sim(image_emb.unsqueeze(1), text_emb.unsqueeze(0)) / self.temp

        # 对角线是正样本
        labels = torch.arange(sim.size(0)).to(sim.device)

        # 计算 loss
        loss = F.cross_entropy(sim, labels) + F.cross_entropy(sim.T, labels)
        return loss / 2

class MemoryBank:
    def __init__(self, size=65536, dim=512):
        self.size = size
        self.dim = dim
        self.memory = torch.randn(size, dim)
        self.ptr = 0

    def update(self, features):
        batch_size = features.size(0)
        end_ptr = self.ptr + batch_size

        if end_ptr <= self.size:
            self.memory[self.ptr:end_ptr] = features
        else:
            split = self.size - self.ptr
            self.memory[self.ptr:] = features[:split]
            self.memory[:batch_size-split] = features[split:]

        self.ptr = end_ptr % self.size