如何为多损失函数设置可训练权重？

1. 能否将w1/w2/w3设为可训练参数？ #

完全可以，但直接用nn.Parameter初始化不加约束会导致梯度爆炸或消失，这就是你遇到NaN的核心原因。

2. 解决NaN问题与权重约束 #

你初始化的方式本身没问题，但权重需要做归一化+非负约束，不然模型会倾向于把权重拉到极端值（比如无穷大）来最小化损失，直接触发NaN。这里给三种常用的可行方案：

方案一：Softmax归一化（最常用） #

把权重先经过Softmax处理，确保它们非负且总和为1，从根源避免单一权重主导：

# 初始化时用均匀分布的初始值，避免Softmax初始输出过于极端
self.w = nn.Parameter(torch.tensor([0.33, 0.33, 0.33], dtype=torch.float32))

# 计算损失时
weights = torch.softmax(self.w, dim=0)
loss = weights[0] * loss1 + weights[1] * loss2 + weights[2] * loss3
loss.backward()

Softmax会自动把权重映射到(0,1)区间，且总和固定为1，既解决了负数问题，也能防止权重集中到某一个损失项上。

方案二：Sigmoid+归一化 #

如果不想强制权重和为1，也可以先用Sigmoid把权重限制在(0,1)区间，再做归一化：

self.w = nn.Parameter(torch.tensor([0.33, 0.33, 0.33], dtype=torch.float32))

# 计算损失时
raw_weights = torch.sigmoid(self.w)
weights = raw_weights / raw_weights.sum()  # 归一化确保权重占比合理
loss = weights[0] * loss1 + weights[1] * loss2 + weights[2] * loss3
loss.backward()

方案三：带正则化的硬约束 #

如果非要保留原始权重范围，可以给权重加L2正则化，同时用clamp强制非负：

self.w = nn.Parameter(torch.tensor([0.33, 0.33, 0.33], dtype=torch.float32))

# 计算损失时
weights = self.w.clamp(min=1e-6)  # 加小epsilon避免权重完全归零
loss = weights[0] * loss1 + weights[1] * loss2 + weights[2] * loss3
# 加L2正则化防止权重过大
loss += 1e-4 * torch.norm(weights, p=2)
loss.backward()

这种方式灵活性高，但需要手动调整正则化系数，不如前两种省心。

3. 避免权重归零或集中的核心逻辑 #

• 非负约束：通过Softmax/Sigmoid把权重限制在正区间，同时加1e-6的下限，防止对应损失项被彻底忽略。
• 归一化：强制权重和为1，让模型必须在不同损失之间做权衡，而不是靠放大某一个权重来降低总损失。
• 正则化：给权重加L1/L2正则，惩罚极端大的权重，避免单一损失主导训练。

4. 无需手动调参的可训练权重方法 #

上面的Softmax归一化方案就是无需手动调参的，模型会自动根据训练数据调整权重。另外还有两种进阶思路：

• 动态权重预测：用小型MLP预测权重，输入是当前各损失项的数值，输出归一化后的权重。这种方式能让权重随训练过程动态调整，比如某一阶段某损失波动大，模型会自动降低它的权重占比。
• 自适应损失加权：参考Focal Loss的思路，基于损失的历史均值自动调整权重：

  # 初始化损失均值缓存，无需梯度
  self.loss_mean = torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, requires_grad=False)
  
  # 训练时更新均值（滑动平均）
  self.loss_mean = 0.9 * self.loss_mean + 0.1 * torch.tensor([loss1.item(), loss2.item(), loss3.item()])
  # 用均值的倒数作为权重，再归一化
  weights = 1.0 / (self.loss_mean + 1e-6)
  weights = weights / weights.sum()
  loss = weights[0] * loss1 + weights[1] * loss2 + weights[2] * loss3
  loss.backward()
  

  这种方式不需要训练额外参数，完全基于损失的动态情况自动调整，适合不想增加模型复杂度的场景。

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