Stable Diffusion微调时loss震荡问题解决

在Stable Diffusion微调过程中，loss震荡是一个常见但棘手的问题。本文将分享一个有效的解决方案。

问题现象

在使用Dreambooster或LoRA方法微调SD模型时，loss曲线出现剧烈震荡，训练不稳定，甚至出现loss为NaN的情况。

根本原因

主要原因是学习率设置不当，特别是在使用AdamW优化器时，初始学习率过高导致梯度更新不稳定。

解决方案

采用动态学习率调度策略，并配合梯度裁剪（Gradient Clipping）：

from transformers import get_linear_schedule_with_warmup
import torch.optim as optim

# 设置优化器
optimizer = optim.AdamW(model.parameters(), lr=1e-5, weight_decay=1e-2)

# 动态学习率调度
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=1000,
    num_training_steps=total_steps
)

# 训练循环中加入梯度裁剪
for batch in dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    loss.backward()
    
    # 梯度裁剪
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()

额外建议

1. 使用更小的初始学习率（1e-5到1e-6）
2. 适当增加warmup步数至1000-2000步
3. 启用梯度裁剪防止梯度爆炸
4. 监控loss曲线，必要时降低batch size

此方法已在多个SD微调项目中验证有效，可显著改善训练稳定性。