CLIP对比学习中的InfoNCE函数：原理剖析与实战优化

1次阅读

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

对比学习（Contrastive Learning）的核心目标是通过拉近正样本对、推远负样本对的方式学习有效表征。在 CLIP 等多模态模型中，InfoNCE（Noise Contrastive Estimation）作为损失函数直接影响着模型对图文关系的建模能力。但在实际训练中，工程师常面临三大挑战：

梯度消失 ：当温度系数（temperature τ）设置不当时，softmax 分布会过于平坦或尖锐
负样本不足 ：小 batch_size 下负样本数量有限，导致对比效果下降
数值不稳定 ：指数运算可能引发数值溢出，尤其在混合精度训练时

原始 InfoNCE 函数定义如下（假设正样本对为 $(x_i,y_i)$）：

$$\mathcal{L}{InfoNCE} = -\frac{1}{N}\sum$$}^N \log \frac{e^{sim(x_i,y_i)/\tau}}{\sum_{j=1}^N e^{sim(x_i,y_j)/\tau} + \sum_{k=1}^M e^{sim(x_i,x_k)/\tau}

其中关键参数 τ 的物理意义是：

τ 值越大，样本间差异被平滑，模型更关注 easy negatives
τ 值越小，模型更聚焦 hard negatives，但可能导致训练不稳定

# 基础实现（存在数值不稳定问题）scores = torch.matmul(query, key.t()) / tau
exp_scores = torch.exp(scores)
loss = -torch.log(exp_scores.diag() / exp_scores.sum(1))

数值稳定性优化 ：采用 log-sum-exp 技巧

max_scores = torch.max(scores, dim=1, keepdim=True)[0]
stable_exp = torch.exp(scores - max_scores)
log_sum_exp = max_scores.squeeze() + torch.log(stable_exp.sum(1))
loss = log_sum_exp - scores.diag()

分布式训练优化 ：通过 all_gather 同步跨 GPU 样本

gathered_keys = [torch.zeros_like(keys) for _ in range(world_size)]
torch.distributed.all_gather(gathered_keys, keys)
all_keys = torch.cat(gathered_keys[:rank] + [keys] + gathered_keys[rank+1:])

内存优化 ：实现负样本队列（memory bank）

self.register_buffer('queue', torch.randn(dim, queue_size))
self.queue = torch.cat([queue[:, batch_size:], new_keys.T], dim=1)

class InfoNCE(nn.Module):
    def __init__(self, tau=0.07, auto_tau=False):
        super().__init__()
        self.tau = nn.Parameter(torch.tensor(tau)) if auto_tau else tau

    def forward(self, query, key):
        # 计算相似度矩阵
        logits = query @ key.T / self.tau

        # 数值稳定计算
        max_logits = logits.max(dim=1, keepdim=True)[0]
        exp_logits = torch.exp(logits - max_logits)
        log_prob = logits.diag() - (max_logits.squeeze() + torch.log(exp_logits.sum(1)))

        return log_prob.mean()