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

背景痛点

在自监督学习的场景中，对比学习（Contrastive Learning）已成为一种主流方法。然而，传统的对比学习方法在 batch 较小时会面临显著的性能下降问题。这主要是因为：

• 样本利用率低：每个样本只能与 batch 内的其他样本构成对比对，当 batch 较小时，负样本数量有限，难以充分挖掘数据间的差异
• 计算开销大：随着 batch 增大，相似度矩阵的计算复杂度呈平方级增长，给 GPU 显存带来巨大压力

技术方案

InfoNCE 损失函数解析

对比学习的核心是 InfoNCE 损失函数，其数学表达式为：

$$
L_{q} = -\log\frac{\exp(q \cdot k_{+}/\tau)}{\sum_{i=0}^{K}\exp(q \cdot k_{i}/\tau)}
$$

其中：

• $q$ 是查询样本的特征
• $k_{+}$ 是正样本特征
• $k_{i}$ 是负样本特征
• $\tau$ 是温度系数

主流范式对比

1. MoCo：使用动量编码器和记忆库存储负样本，能构建大量负样本但实现复杂
2. SimCLR：完全依赖 batch 内样本，简单直接但 batch 需求大
3. Batch 内复用：在 batch 内挖掘负样本，平衡实现难度和性能

Batch 内负样本复用策略

核心思想是将 batch 内所有其他样本作为当前样本的潜在负样本，通过以下方式提升效率：

• 对称损失计算：同时考虑样本 i 对 j 和 j 对 i 的对比损失
• 负样本共享：每个样本自动获得 batch_size- 2 个负样本

代码实现

基础实现框架

import torch
import torch.nn.functional as F

class ContrastiveLoss(nn.Module):
    def __init__(self, temperature=0.07):
        super().__init__()
        self.temperature = temperature

    def forward(self, features):
        # 特征归一化
        features = F.normalize(features, dim=1)

        # 计算相似度矩阵
        sim_matrix = torch.matmul(features, features.T) / self.temperature

        # 构造标签
        batch_size = features.shape[0]
        labels = torch.arange(batch_size).to(features.device)

        # 计算损失
        loss = F.cross_entropy(sim_matrix, labels)
        return loss

关键优化组件

1. 特征归一化：

   features = F.normalize(features, dim=1)

2. 混合精度训练：

   with torch.cuda.amp.autocast():
       features = model(inputs)
       loss = criterion(features)

3. 梯度裁剪：

   torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

性能优化

计算复杂度分析

• 相似度矩阵计算：$O(B^2D)$，B 为 batch size，D 为特征维度
• 显存占用：主要来自相似度矩阵的存储，约 $B^2 \times 4$ 字节（float32）

显存优化技巧

1. 使用梯度检查点（Gradient Checkpointing）
2. 采用混合精度训练
3. 分块计算相似度矩阵

避坑指南

温度系数调优

• 初始值建议设为 0.07
• 调整范围通常在 [0.01, 0.2] 之间
• 温度系数过小会导致梯度爆炸，过大则难以区分相似样本

预防特征坍塌

1. 添加预测头（Projection Head）
2. 使用额外的正则化项
3. 监控特征相似度矩阵的对角优势

分布式训练注意事项

• 确保各 GPU 上的 batch 统计量同步
• 使用 all_gather 收集各节点的特征
• 注意梯度同步带来的通信开销

延伸阅读

1. SimCLR 论文
2. MoCo 论文
3. 对比学习综述

思考题

1. 如何在不增加 batch size 的情况下获得更多负样本？
2. 温度系数与学习率应该如何协同调整？
3. 在跨模态对比学习中，batch 内对比策略需要做哪些调整？