别再手动调权重了！用PyTorch实现多任务损失自适应加权（附代码）-编程阁

多任务学习中损失权重的自动化调参实战：PyTorch实现与工程细节

当你的神经网络需要同时预测用户点击率和购买金额时，分类损失和回归损失应该如何平衡？这个困扰无数算法工程师的问题，其实有更优雅的解决方案。传统手工调整损失权重的方式不仅耗时，而且难以捕捉任务间的动态关系。2018年CVPR论文《Multi-Task Learning Using Uncertainty to Weigh Losses》提出的自适应加权方法，让我们看到了自动化解决这一难题的可能性。

本文将带你深入理解基于不确定度的自适应损失加权原理，并手把手实现一个工业级可用的PyTorch解决方案。不同于理论推导为主的论文，我们聚焦于工程落地中的三个关键问题：如何避免数值不稳定、如何处理不同量级的损失函数、以及如何验证权重自适应的实际效果。文末提供的完整代码模块可以直接整合到你的多任务学习项目中。

1. 自适应损失加权的数学本质

理解自适应权重的核心，需要从概率建模的角度重新审视多任务学习。假设我们有两个任务：预测用户年龄（回归）和预测用户性别（分类），模型需要同时输出这两个预测结果。

关键假设：每个任务的预测误差服从独立的高斯分布。对于回归任务，这个假设很自然；对于分类任务，可以理解为对logits添加高斯噪声。由此得到联合概率分布：

p(y₁, y₂|fᴹ(x)) = p(y₁|fᴹ(x)) * p(y₂|fᴹ(x))

取负对数后，总损失自然分解为各任务损失之和。但这里出现了一个重要参数——每个任务对应的噪声方差σ²。这个方差恰恰决定了该任务损失的权重：

L = 1/(2σ₁²) * L₁(回归) + 1/σ₂² * L₂(分类) + log(σ₁) + log(σ₂)

为什么这样做更合理？因为噪声大的任务（σ²大）天然更不可靠，自然应该降低其权重（1/σ²小）。而log(σ)项则防止权重无限增大，起到正则化作用。

2. PyTorch实现的关键技巧

2.1 可学习参数的实现

在PyTorch中，我们需要将log(σ²)作为可训练参数。这里使用nn.Parameter实现：

class MultiTaskLoss(nn.Module): def __init__(self, num_tasks): super().__init__() self.log_vars = nn.Parameter(torch.zeros(num_tasks)) def forward(self, losses): # losses: list of task losses total_loss = 0 for i, loss in enumerate(losses): precision = torch.exp(-self.log_vars[i]) total_loss += precision * loss + self.log_vars[i] return total_loss

为什么预测log(σ²)而不是σ²？这保证了σ²=exp(s)始终为正，且数值更稳定。实验表明，直接预测σ²容易导致训练初期梯度爆炸。

2.2 回归与分类的统一处理

对于不同类型的任务，损失函数需要做适当调整：

任务类型	损失函数	权重系数	正则项
回归任务	MSE	1/(2σ²)	log(σ)
分类任务	CrossEntropy	1/σ²	log(σ)

实际实现时，可以通过任务标志位自动选择计算方式：

def task_loss(pred, target, task_type): if task_type == 'regression': return F.mse_loss(pred, target) elif task_type == 'classification': return F.cross_entropy(pred, target)

2.3 训练稳定性的工程技巧

在多任务训练初期，我们常遇到以下问题：

损失量级差异：分类交叉熵可能在1-10之间，而MSE可能高达1000+
权重初始化敏感：初始log(σ²)设为0可能导致某些任务完全被忽略

解决方案：

对回归任务输出做标准化预处理
采用分阶段训练策略（先单独训练各任务，再联合微调）
对log_vars使用Xavier初始化

# 改进后的初始化方式 nn.init.uniform_(self.log_vars, -3, 0) # 初始σ在[0.05,1]之间

3. 完整案例：多任务推荐模型

让我们构建一个实际案例：预测用户的活跃度（回归）和付费意愿（分类）。数据集采用模拟的用户行为数据，包含：

特征：历史点击、停留时长、设备信息等
标签：次日使用时长（回归）、是否付费（分类）

3.1 模型架构设计

class MultiTaskModel(nn.Module): def __init__(self, input_dim): super().__init__() self.shared_layer = nn.Sequential( nn.Linear(input_dim, 256), nn.ReLU(), nn.Dropout(0.3) ) self.reg_head = nn.Linear(256, 1) self.cls_head = nn.Linear(256, 2) self.loss_fn = MultiTaskLoss(num_tasks=2) def forward(self, x, targets=None): features = self.shared_layer(x) reg_out = self.reg_head(features).squeeze() cls_out = self.cls_head(features) if targets is not None: reg_loss = F.mse_loss(reg_out, targets[0]) cls_loss = F.cross_entropy(cls_out, targets[1].long()) total_loss = self.loss_fn([reg_loss, cls_loss]) return total_loss, {'reg': reg_out, 'cls': cls_out} return {'reg': reg_out, 'cls': cls_out}

3.2 训练过程监控

自适应权重的优势在于训练过程中能动态调整。我们可以记录log(σ²)的变化：

for epoch in range(100): model.train() for batch in train_loader: optimizer.zero_grad() loss, _ = model(batch['features'], [batch['duration'], batch['pay']]) loss.backward() optimizer.step() # 查看当前权重 reg_weight = torch.exp(-model.loss_fn.log_vars[0]).item() cls_weight = torch.exp(-model.loss_fn.log_vars[1]).item() print(f"Epoch {epoch}: reg_weight={reg_weight:.3f}, cls_weight={cls_weight:.3f}")

典型训练过程中，我们会观察到：