CLIP对比学习中的InfoNCE函数：从原理到实践指南

1次阅读

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

对比学习（Contrastive Learning）在视觉 - 语言预训练（如 CLIP）中扮演着核心角色，而 InfoNCE（Noise Contrastive Estimation）损失函数是其训练过程中最关键的组成部分。许多开发者在实际应用时，常常遇到以下问题：

梯度消失或爆炸，导致模型难以收敛
温度系数（temperature）的选择缺乏明确指导
负样本构造方式对模型性能影响显著但优化困难

这些问题直接影响模型的最终表现，但现有资料往往过于理论化，缺乏具体的工程实践指导。

InfoNCE 的核心公式如下：

$$
\mathcal{L}{InfoNCE} = -\log \frac{\exp(s
$$}/\tau)}{\sum_{k=1}^N \exp(s_{i,k}/\tau)

其中：

$s_{i,j}$ 表示正样本对的相似度得分
$\tau$ 是温度系数，控制分布的尖锐程度
分母中的求和项包含一个正样本和 N - 1 个负样本

该函数实质上是最大化正样本对的互信息下界（Mutual Information Lower Bound）。

CLIP：直接使用 batch 内所有其他样本作为负样本，实现简单但受 batch 大小限制
MoCo：引入 memory bank 存储历史负样本，扩大负样本数量但增加内存开销

以下是一个带注释的 PyTorch 实现，包含 GPU 并行计算优化：

import torch
import torch.nn.functional as F

def info_nce_loss(features, temperature=0.07):
    """
    features: 归一化后的特征矩阵 [batch_size, feature_dim]
    temperature: 温度系数
    """
    device = features.device
    batch_size = features.shape[0]

    # 计算所有样本对的相似度矩阵
    similarity_matrix = torch.matmul(features, features.T)  # [batch_size, batch_size]

    # 构建正样本掩码（对角线为 1，其余为 0）mask = torch.eye(batch_size, dtype=torch.bool, device=device)

    # 提取正负样本对
    positives = similarity_matrix[mask].view(batch_size, -1)  # [batch_size, 1]
    negatives = similarity_matrix[~mask].view(batch_size, -1)  # [batch_size, batch_size-1]

    # 合并正负样本并计算 logits
    logits = torch.cat([positives, negatives], dim=1) / temperature

    # 构建标签（第一个位置为正样本）labels = torch.zeros(batch_size, dtype=torch.long, device=device)

    # 计算交叉熵损失
    loss = F.cross_entropy(logits, labels)

    return loss

温度系数 $\tau$ 的选择对模型性能至关重要：