BLIP损失函数原理解析与视觉-语言模型训练实战

1次阅读

没有评论

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

视觉 - 语言预训练需要解决两大核心问题：

模态鸿沟 ：图像像素空间与文本符号空间缺乏天然对齐关系
细粒度关联 ：需要捕捉局部区域与单词级别的对应关系（如 ” 黑色汽车 ” 与图像中特定车辆的关联）

传统方法如 CLIP 仅使用对比学习，难以建立这种细粒度关联。这正是 BLIP 提出多任务损失函数的动机。

CLIP：
单一对比损失（InfoNCE）
全局图像 - 文本匹配
负样本来自不同 batch 样本
BLIP：
三组件联合训练：
1. 图像 - 文本对比损失（ITC）
2. 掩码语言建模损失（MLM）
3. 图像 - 文本匹配损失（ITM）
关键创新：
- 使用相同编码器处理 ITC 和 ITM 任务
- 动态硬负样本挖掘

数学形式：

L_itc = 1/2*(L_i2t + L_t2i)
L_i2t = -log[exp(sim(v_i,t_i)/τ)/∑_j exp(sim(v_i,t_j)/τ)]

PyTorch 实现：

def itc_loss(image_embeds, text_embeds, temp=0.07):
    # image_embeds: [batch_size, dim]
    # text_embeds: [batch_size, dim]
    sim_matrix = image_embeds @ text_embeds.T  # [bs, bs]
    targets = torch.arange(sim_matrix.size(0)).to(device)

    loss_i = F.cross_entropy(sim_matrix/temp, targets)
    loss_t = F.cross_entropy(sim_matrix.T/temp, targets)
    return (loss_i + loss_t)/2

def mlm_loss(text_encoder, image_embeds, input_ids, mlm_prob=0.15):
    # input_ids: [batch_size, seq_len]
    # 创建掩码
    mask_pos = torch.rand(input_ids.shape) < mlm_prob
    masked_ids = input_ids.clone()
    masked_ids[mask_pos] = tokenizer.mask_token_id

    # 融合视觉特征进行预测
    logits = text_encoder(masked_ids, image_embeds)  # [bs, seq_len, vocab_size]
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
        input_ids.view(-1),
        ignore_index=tokenizer.pad_token_id
    )
    return loss

def itm_loss(image_encoder, text_encoder, images, texts):
    # 正负样本构造
    pos_pairs = (images, texts)
    neg_texts = texts[torch.randperm(texts.size(0))]  # 负样本采样
    neg_pairs = (images, neg_texts)

    # 计算匹配分数
    pos_score = matching_head(image_encoder(pos_pairs[0]), text_encoder(pos_pairs[1]))
    neg_score = matching_head(image_encoder(neg_pairs[0]), text_encoder(neg_pairs[1]))

    labels = torch.cat([torch.ones(pos_score.size(0)), torch.zeros(neg_score.size(0))])
    return F.binary_cross_entropy_with_logits(torch.cat([pos_score, neg_score]), labels)

# 初始化分布式环境
torch.distributed.init_process_group(backend='nccl')

# 使用 DistributedDataParallel 包装模型
model = torch.nn.parallel.DistributedDataParallel(
    model,
    device_ids=[local_rank],
    output_device=local_rank
)

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    loss = itc_loss(image_emb, text_emb) + mlm_loss(...) + itm_loss(...)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

典型 loss 曲线应呈现：

ITC loss 快速下降（前 1 -2epoch）
MLM loss 平稳下降（需注意 perplexity 指标）
ITM accuracy 稳定上升（验证集应达到 85%+）

[Epoch 5/50] 
  ITC Loss: 1.25 → 0.98
  MLM Perplexity: 15.2 → 12.1
  ITM Acc: 72% → 79%
  Val R@1: 45% → 52%

错误做法 ：仅使用 batch 内随机负样本
正确做法 ：维护负样本队列（MoCo 风格）或使用跨 GPU 负样本

注意检查 MLM 任务的 logits 数值范围
遇到 NaN 时可尝试：
调小学习率
增加 gradient clipping
对 softmax 前 logits 进行 clamp

ITM 验证准确率与检索任务强相关
MLM perplexity 反映语言理解能力
下游任务微调时建议：
分类任务：优先看 ITM 指标
生成任务：关注 MLM 表现

# 初始化
model = BLIPModel(...).to(device)
optimizer = AdamW(model.parameters(), lr=5e-5)

# 训练循环
for batch in dataloader:
    images, texts = batch

    with torch.cuda.amp.autocast():
        # 前向传播
        image_emb = model.image_encoder(images)
        text_emb = model.text_encoder(texts.input_ids)

        # 计算多任务损失
        loss = 0.5 * itc_loss(image_emb, text_emb)
        loss += 0.3 * mlm_loss(...)
        loss += 0.2 * itm_loss(...)

    # 反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

    # 日志记录
    if step % 100 == 0:
        log_stats = {'loss': loss.item(),
            'itc_loss': itc_loss.item(),
            'mlm_ppl': torch.exp(mlm_loss).item()}
        print(json.dumps(log_stats))

BLIP 通过精心设计的多任务损失函数，在保持 CLIP 全局对齐优势的同时，增强了细粒度跨模态理解能力。实际应用中需要注意：