CLIP对比学习损失优化实战：从理论到PyTorch实现

1次阅读

没有评论

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

CLIP 的对比学习损失函数通过计算图像和文本嵌入的相似度矩阵，鼓励正样本对（匹配的图文对）相似度高，负样本对相似度低。标准实现通常使用以下公式：

$$
\mathcal{L} = -\frac{1}{N}\sum_{i=1}^N \log\frac{e^{s_{i,i}/\tau}}{\sum_{j=1}^N e^{s_{i,j}/\tau}}
$$

其中 $s_{i,j}$ 是图像 $i$ 和文本 $j$ 的相似度，$\tau$ 是温度参数。实际应用中我们发现两个主要问题：

计算复杂度高：相似度矩阵计算需要 $O(N^2)$ 内存，当 batch size 较大时显存容易爆炸
梯度贡献不均：简单随机采样导致多数负样本梯度微弱，只有少数困难负样本主导优化过程

将大 batch 拆分为 $k×k$ 子块（如 $k=8$），依次计算子块相似度后拼接。数学形式保持不变，但显存占用从 $O(N^2)$ 降至 $O((N/k)^2)$

引入基于梯度统计的自适应温度系数：

$$
\tau_t = \tau_0 \cdot \frac{|\nabla_\theta\mathcal{L}|}{\mathbb{E}[|\nabla_\theta\mathcal{L}|]}
$$

在计算损失时，对每行相似度排序后选取 top- K 负样本参与计算：

$$
\mathcal{L}{hard} = -\frac{1}{N}\sum}^N \log\frac{e^{s_{i,i}/\tau}}{e^{s_{i,i}/\tau} + \sum_{j\in \mathcal{Ni} e^{s
$$}/\tau}

import torch
import torch.nn.functional as F

class ImprovedCLIPLoss(torch.nn.Module):
    def __init__(self, chunk_size=8, topk_neg=32):
        super().__init__()
        self.chunk_size = chunk_size
        self.topk_neg = topk_neg
        # 初始化可学习温度参数
        self.logit_scale = torch.nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1/0.07)))

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

        # 分块计算相似度矩阵
        sim_matrix = []
        for img_chunk in image_features.chunk(self.chunk_size):
            chunk_row = []
            for txt_chunk in text_features.chunk(self.chunk_size):
                # 使用 einsum 高效计算余弦相似度
                chunk_sim = torch.einsum('i d, j d -> i j', img_chunk, txt_chunk)
                chunk_row.append(chunk_sim)
            sim_matrix.append(torch.cat(chunk_row, dim=1))
        sim_matrix = torch.cat(sim_matrix, dim=0)

        # 应用可学习温度系数
        sim_matrix = sim_matrix * self.logit_scale.exp()

        # 困难负样本挖掘
        pos_sim = torch.diag(sim_matrix).unsqueeze(1)
        neg_mask = ~torch.eye(len(sim_matrix), dtype=torch.bool, device=sim_matrix.device)
        neg_sim = sim_matrix[neg_mask].view(len(sim_matrix), -1)
        topk_neg = neg_sim.topk(self.topk_neg, dim=1).values

        # 计算对比损失
        numerator = pos_sim.exp()
        denominator = numerator + topk_neg.exp().sum(dim=1, keepdim=True)
        loss = -torch.log(numerator / denominator).mean()

        return loss

在 COCO 数据集上测试（RTX 3090 单卡）：