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

1次阅读

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

多模态模型的核心挑战是让机器理解视觉和语言之间的关联。传统方法通常需要标注数据来建立这种联系，而 CLIP 通过对比学习实现了无监督的跨模态对齐。想象一下：当模型看到一张猫的图片和 ” 狗 ” 的文本时，它应该知道这两者不匹配；看到 ” 猫 ” 的文本时则应该匹配。这就是对比学习要解决的问题。

CLIP 使用的损失函数基于 InfoNCE（Noise Contrastive Estimation），其核心思想是通过比较正样本对和负样本对的相似度来学习表征。公式看起来可能有点吓人，但其实理解起来很直观：

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

其中：
– $s_{i,j}$ 是第 i 个图像和第 j 个文本的相似度得分
– $\tau$ 是温度参数
– N 是 batch size

这个公式的本质是让匹配的图片 - 文本对 (正样本) 的相似度远高于不匹配的对(负样本)。

温度参数 $\tau$ 控制着模型对困难样本的关注程度：

τ 越小：模型更关注那些特别难的负样本
τ 越大：所有样本得到的关注更平均

实践中 τ 需要仔细调优，通常设置在 0.01 到 0.1 之间。CLIP 原文发现 τ =0.07 效果最佳。

下面是一个考虑工业级实践的完整实现，包含分布式训练、混合精度等特性：

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

def gather_tensors(tensor: Tensor) -> Tensor:
    """
    分布式环境下收集所有 GPU 上的张量
    Args:
        tensor: 当前 GPU 上的张量
    Returns:
        所有 GPU 张量的拼接结果
    """
    gathered_tensors = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())]
    dist.all_gather(gathered_tensors, tensor)
    return torch.cat(gathered_tensors, dim=0)

class CLIPLoss(nn.Module):
    def __init__(self, temperature: float = 0.07, label_smoothing: float = 0.1):
        super().__init__()
        self.temperature = temperature
        self.label_smoothing = label_smoothing
        self.cross_entropy = nn.CrossEntropyLoss(label_smoothing=label_smoothing)

    def forward(self, image_features: Tensor, text_features: Tensor) -> Tensor:
        """
        计算 CLIP 对比损失
        Args:
            image_features: 图像特征 [local_batch_size, feature_dim]
            text_features: 文本特征 [local_batch_size, feature_dim]
        Returns:
            标量损失值
        """
        # 归一化特征
        image_features = image_features / image_features.norm(dim=1, keepdim=True)
        text_features = text_features / text_features.norm(dim=1, keepdim=True)

        # 分布式环境下收集所有特征
        if dist.is_initialized():
            image_features = gather_tensors(image_features)
            text_features = gather_tensors(text_features)

        # 计算相似度矩阵 [global_batch_size, global_batch_size]
        logits_per_image = image_features @ text_features.t() / self.temperature
        logits_per_text = logits_per_image.t()

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

        # 对称损失
        loss_i = self.cross_entropy(logits_per_image, labels)
        loss_t = self.cross_entropy(logits_per_text, labels)
        return (loss_i + loss_t) / 2