多模态模型训练时的学习率调度策略踩坑
最近在设计图像+文本联合训练系统时，踩了一个关于学习率调度的坑，分享给大家。
问题描述
在训练CLIP风格的多模态模型时，我采用了统一的学习率调度策略（cosine decay），结果发现模型收敛很慢，甚至出现梯度爆炸。后来排查发现，图像分支和文本分支的特征空间差异巨大，统一调度导致文本分支学习过快而图像分支跟不上。
复现步骤
- 使用ResNet50作为图像编码器，BERT作为文本编码器
- 采用统一的cosine decay调度策略：
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
- 训练10个epoch后发现loss波动剧烈
解决方案
应该为不同模态分支设计独立的学习率调度策略。我最终采用分阶段调度：
# 图像分支使用linear decay
img_scheduler = torch.optim.lr_scheduler.LinearLR(img_optimizer, start_factor=1.0, end_factor=0.01, total_iters=epochs)
# 文本分支使用cosine decay
txt_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(txt_optimizer, T_max=epochs)
关键教训
多模态模型训练时，不能简单套用单模态的学习率策略，必须考虑不同模态的特征分布差异。
讨论