CLIP对比学习损失函数详解：从原理到PyTorch实战

1次阅读

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

在多模态模型训练中，核心难题是如何让视觉和语言两个模态的特征空间对齐。传统方法通常需要大量标注数据来建立模态间的映射关系，而 CLIP 通过对比学习实现了自监督的跨模态对齐。

对比学习的核心思想是：
– 正样本对（如图像和对应文本描述）在特征空间中应该相近
– 负样本对（如图像和其他无关文本）在特征空间中应该远离

CLIP 采用的损失函数是基于 InfoNCE 的改进版本，其数学形式为：

$$\mathcal{L} = -\frac{1}{N}\sum_{i=1}^N \log\frac{\exp(\text{sim}(v_i,t_i)/\tau)}{\sum_{j=1}^N \exp(\text{sim}(v_i,t_j)/\tau)}$$

其中：
– $v_i$, $t_i$ 分别是第 i 个图像和文本的特征向量
– $\text{sim}(·,·)$ 是余弦相似度
– $\tau$ 是温度系数

相比原始 InfoNCE，CLIP 做了两个重要改进：
1. 双向对称计算（图像→文本和文本→图像两个方向的损失求和）
2. 对大批次训练做了优化（使用分布式 AllGather 收集负样本）

温度系数是控制损失函数 ” 软硬度 ” 的关键参数：
– 当 τ→0 时，损失函数趋近于 hard 分类
– 当 τ→∞时，所有样本的权重趋于相同

实践中发现：
– τ 过大导致模型无法区分难样本
– τ 过小导致梯度爆炸风险
– CLIP 原论文推荐值在 0.01 到 0.1 之间

以下是支持多 GPU 训练的工业级实现（关键注释已添加）：

import torch
import torch.distributed as dist
from torch import nn, Tensor

def clip_loss(
    image_features: Tensor, 
    text_features: Tensor,
    temp: float = 0.07,
    label_smoothing: float = 0.1
) -> Tensor:
    """
    CLIP 对比损失函数的 PyTorch 实现

    参数:
        image_features: [batch_size, feat_dim] 图像特征
        text_features: [batch_size, feat_dim] 文本特征
        temp: 温度系数
        label_smoothing: 标签平滑系数

    返回:
        标量损失值
    """
    # 分布式环境下收集所有设备上的特征
    if dist.is_initialized():
        image_features = concat_all_gather(image_features)
        text_features = concat_all_gather(text_features)

    # 归一化特征
    image_features = nn.functional.normalize(image_features, dim=-1)
    text_features = nn.functional.normalize(text_features, dim=-1)

    # 计算相似度矩阵 [batch_size, batch_size]
    logits = image_features @ text_features.T / temp

    # 创建标签 [batch_size]
    batch_size = image_features.shape[0]
    labels = torch.arange(batch_size, device=image_features.device)

    # 标签平滑
    if label_smoothing > 0:
        logits = logits * (1 - label_smoothing) + label_smoothing / batch_size

    # 对称计算两个方向的损失
    loss_i = nn.functional.cross_entropy(logits, labels)
    loss_t = nn.functional.cross_entropy(logits.T, labels)
    return (loss_i + loss_t) / 2

def concat_all_gather(tensor: Tensor) -> Tensor:
    """分布式 AllGather 实现"""
    tensors_gather = [torch.ones_like(tensor) for _ in range(dist.get_world_size())]
    dist.all_gather(tensors_gather, tensor)
    return torch.cat(tensors_gather, dim=0)