PyTorch多节点多损失组合训练是否需要将损失缩放到相近范围？

核心结论 #

• 你需要将不同子损失的取值缩放到相近范围，除非你主动希望某一个子损失主导训练过程
• 你举例的场景中，和loss1关联的节点确实会在训练中拥有更高的权重，甚至会让loss2、loss3对应的优化目标完全失效

背后逻辑 #

反向传播时，模型参数收到的总梯度是所有子损失对该参数梯度的加和，损失的量级会直接传导到梯度量级上：

• 固定学习率优化器（如SGD）场景下，参数更新步长和梯度绝对值直接成正比，10k量级的loss1贡献的梯度会比1_{10量级的loss2、loss3大3}4个数量级，参数更新方向会完全被loss1的优化目标主导，模型只会优先拟合loss1对应的任务，loss2、loss3的优化信号会被完全覆盖。
• 自适应学习率优化器（如Adam、RMSprop）也无法完全抵消量级差异：自适应优化器的学习率调整是基于参数历史梯度的波动幅度做归一化，如果loss1的梯度全程都比另外两个子损失大几个量级，归一化后它的贡献仍然会占据绝对主导，只有不同损失的梯度波动幅度差异大时，自适应机制才能起到明显的平衡作用。

常用处理方案 #

最通用的做法是给不同子损失增加权重系数，加权后再求和得到总损失：
total_loss = w1 * loss1 + w2 * loss2 + w3 * loss3

• 静态权重设置：训练正式启动前先跑3~5个step的预训练，统计三个子损失的平均取值，权重取值参考w_n = 目标平均量级 / 子损失n的统计平均量级，让加权后的三个子损失平均量级基本一致。比如你举例的场景中loss1平均为10000，loss2、loss3平均为5，就可以给loss1设置权重0.0005，让加权后三个子损失的平均值都在5左右。
• 动态权重调整：如果不同子任务的收敛速度、难度差异较大，也可以随着训练轮次动态调整权重，比如前期给难度更高的子损失设置更高权重，后期再平衡各任务的优化优先级。

内容的提问来源于stack exchange，提问作者Richie