1. 多任务学习里的“跷跷板”困局如果你训过多任务学习模型大概率遇到过这种糟心事一个模型同时学目标检测和语义分割检测的loss哗哗往下掉分割的loss却像被钉住一样纹丝不动或者反过来某个任务收敛得飞快另一个任务怎么调都上不去。你调学习率、换优化器、加数据折腾一圈发现——问题根本不在这些地方而在于多个任务的损失函数在共享网络里互相打架。这就是多任务学习最核心的痛点损失平衡。不同任务的loss量级不同、收敛速度不同、梯度方向还可能冲突。手工调权重今天调好了明天换个数据集又崩了。GradNorm就是冲着这个问题来的——它让网络在训练过程中自动、动态地调整各任务的损失权重让每个任务都能以合理的速度收敛而不是被某个“强势”任务带着跑偏。这篇文章我会从原理到代码把GradNorm彻底拆开讲清楚。适合已经有多任务学习基础、被loss平衡折磨过的同学也适合刚接触多任务、想搞明白“为什么不能简单把loss加起来”的新手。读完你至少能搞清楚三件事GradNorm到底在归一化什么、它的梯度计算怎么推导、以及在实际项目里怎么落地和避坑。2. GradNorm的核心设计思路拆解2.1 为什么“把loss加起来”是个糟糕的主意先说说最朴素的做法总损失等于各任务损失的加权和。$$L_{total} \sum_i w_i L_i$$大部分人的做法是给每个$w_i$设个固定值比如1.0或者凭经验设0.5、2.0之类的。问题在哪第一个问题是量级不匹配。分类任务的交叉熵loss通常在0.1到2之间而回归任务的MSE loss可能是几十甚至上百。你直接把它们加起来回归任务的梯度会完全主导网络更新分类任务基本学不动。第二个问题是收敛速度不匹配。有的任务简单几个epoch就收敛了有的任务难需要几十个epoch。固定权重下简单任务收敛后梯度变小难任务的梯度相对变大反而可能把已经学好的特征破坏掉。第三个问题最隐蔽梯度方向冲突。两个任务的梯度在共享层可能指向相反方向简单加权求和会让它们互相抵消网络原地打转。我见过太多项目在loss权重上反复试错最后靠“玄学调参”勉强跑通。GradNorm的价值就在于把这件靠直觉的事变成一个有数学依据的自动化过程。2.2 GradNorm到底在“归一化”什么名字叫GradNorm但它归一化的不是梯度本身而是各任务梯度相对于彼此的量级。核心思想可以这样理解我们希望所有任务以相近的速度在训练。怎么衡量“训练速度”用损失下降的速率。怎么控制这个速率通过调整各任务的损失权重$w_i$让每个任务在共享层产生的梯度范数保持在一个合理的相对比例上。具体来说GradNorm定义了一个目标量$$G_i^{(t)} | \nabla_W (w_i(t) L_i(t)) |_2$$这是第$i$个任务在共享层参数$W$上的梯度范数带权重的。然后它计算所有任务梯度范数的均值$$\bar{G}(t) \mathbb{E}_{task} [G_i(t)]$$接着定义相对逆训练速率$$\tilde{L}_i(t) L_i(t) / L_i(0)$$这个比值越小说明任务$i$下降得越快。GradNorm希望下降快的任务梯度范数小一点下降慢的任务梯度范数大一点从而让所有任务同步收敛。最终GradNorm构造了一个梯度损失函数通过最小化它来更新权重$w_i$$$L_{grad} \sum_i | G_i(t) - \bar{G}(t) \cdot [\tilde{L}_i(t)]^\alpha |_1$$其中$\alpha$是一个超参数控制“训练速率平衡”的力度。$\alpha0$时退化为简单的梯度范数均衡$\alpha$越大对收敛速度差异的惩罚越强。2.3 为什么用梯度范数而不是直接调loss权重这里有个很关键的洞察loss的大小不等于梯度的大小。一个任务的loss可能很大但如果它的梯度很小比如接近收敛那它对网络更新的实际影响就很小。反过来一个loss看起来不大的任务如果梯度很陡它就会主导训练。GradNorm直接盯住梯度范数等于绕过了loss量级这个“障眼法”直接控制每个任务对网络参数更新的实际贡献。这是它比手工调权重高明的地方。另一个好处是动态性。训练初期各任务梯度都很大GradNorm会快速调整权重让它们平衡训练后期任务逐渐收敛梯度变小GradNorm也会相应调整。整个过程是自适应的不需要人工干预。2.4 和Uncertainty Weighting、DWA的区别多任务损失平衡不是只有GradNorm一个方案。常见的还有Uncertainty Weighting基于任务不确定性自动学权重把$w_i$参数化为可学习的方差。优点是理论优雅缺点是假设了损失分布形式实际中不一定成立。DWADynamic Weight Averaging根据各任务loss下降速率直接调权重简单粗暴但只看loss不看梯度遇到loss量级差异大时容易失效。GradNorm直接操作梯度范数理论上更贴近“控制训练速度”这个目标但计算开销更大需要额外反向传播。选哪个我的经验是如果任务数量少2-4个、计算资源充足GradNorm效果最稳如果任务多、训练时间紧DWA或者Uncertainty Weighting更实用。GradNorm的额外计算量主要来自需要单独计算每个任务在共享层的梯度范数这在任务数多时会线性增长。3. 核心细节解析与实操要点3.1 共享层和任务层的划分GradNorm只作用于共享层。这是理解它的前提。在多任务网络里通常结构是底层共享特征提取器比如ResNet的前几个stage然后每个任务有自己的head。GradNorm要平衡的是各任务在共享层参数上的梯度而不是任务head的梯度。为什么因为任务head是各自独立的它们的梯度不会互相干扰。真正打架的是共享层——所有任务都在更新同一组参数这里才是需要平衡的地方。实操中你需要明确指定哪些参数属于共享层。在PyTorch里通常这样做# 假设shared_layers是共享特征提取器 shared_params list(shared_layers.parameters()) task_params [list(task_head_i.parameters()) for i in range(num_tasks)]GradNorm的权重更新只基于shared_params上的梯度。3.2 梯度范数的计算方式计算$G_i | \nabla_W (w_i L_i) |_2$时有两种常见做法做法一对每个任务单独反向传播for i, task_loss in enumerate(task_losses): # 清空共享层梯度 shared_optimizer.zero_grad() # 只对第i个任务的loss反向 (w[i] * task_loss).backward(retain_graphTrue) # 收集共享层梯度范数 grad_norm_i torch.sqrt(sum(p.grad.norm()**2 for p in shared_params))这种做法准确但需要多次反向传播计算开销大。做法二一次反向传播分别收集更高效的方式是让所有任务loss一起反向但在共享层分别记录每个任务贡献的梯度。这需要一些hook技巧实现起来复杂一些但速度快。我实测下来任务数≤4时做法一的额外开销可以接受任务数更多时建议用做法二或者考虑其他平衡方案。3.3 权重更新不是用梯度下降这里有个容易踩的坑GradNorm的权重$w_i$不是用标准梯度下降更新的。标准做法是计算$L_{grad}$对$w_i$的梯度然后用这个梯度去更新$w_i$。但注意$w_i$的更新不应该影响共享层参数的更新——它们是两套独立的优化过程。具体来说训练循环里要做两件事用当前$w_i$计算总loss反向传播更新网络参数共享层任务head。计算$L_{grad}$反向传播更新权重$w_i$。这两步的优化器是分开的。$w_i$通常用较小的学习率比如0.025更新而且更新后要做归一化让$\sum_i w_i num_tasks$防止权重整体膨胀或收缩。# 更新网络参数 total_loss sum(w[i] * task_losses[i] for i in range(num_tasks)) total_loss.backward() network_optimizer.step() # 更新权重w grad_loss compute_grad_loss(task_losses, shared_params, w) grad_loss.backward() w_optimizer.step() # 归一化权重 with torch.no_grad(): w w / w.sum() * num_tasks3.4 超参数α的选择$\alpha$是GradNorm最重要的超参数。它控制“训练速率平衡”的强度。$\alpha0$只平衡梯度范数不考虑各任务收敛速度差异。适合任务难度相近的场景。$\alpha0.5$温和平衡推荐作为默认值。$\alpha1.0$强力平衡适合任务难度差异大的场景但可能过度压制简单任务。我的经验是先从0.5开始如果发现某个任务明显欠拟合调大到0.8或1.0如果发现简单任务被压得太狠导致性能下降调小到0.2或0.3。注意$\alpha$不是越大越好。过大的$\alpha$会让简单任务的权重被压到接近0相当于放弃了那个任务。多任务学习的目的是“共赢”不是“均贫富”。3.5 初始loss的选取$\tilde{L}_i(t) L_i(t) / L_i(0)$里的$L_i(0)$是任务$i$在训练开始时的loss值。这里有个细节$L_i(0)$应该在训练正式开始前测量用初始网络参数跑一遍前向传播得到。不要用第一个batch的loss因为第一个batch的loss波动很大可能偏高或偏低。实操中我会在训练循环开始前用几十个batch的数据跑一遍取平均loss作为$L_i(0)$。这样更稳定。4. 完整实操流程与代码实现4.1 环境准备与网络结构定义先定义一个简单的多任务网络。假设我们做两个任务一个分类任务一个回归任务。import torch import torch.nn as nn import torch.nn.functional as F class SharedEncoder(nn.Module): def __init__(self, input_dim128, hidden_dim256): super().__init__() self.fc1 nn.Linear(input_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, hidden_dim) self.relu nn.ReLU() def forward(self, x): x self.relu(self.fc1(x)) x self.relu(self.fc2(x)) return x class TaskHead(nn.Module): def __init__(self, hidden_dim256, output_dim10): super().__init__() self.fc nn.Linear(hidden_dim, output_dim) def forward(self, x): return self.fc(x) class MultiTaskModel(nn.Module): def __init__(self): super().__init__() self.shared SharedEncoder() self.class_head TaskHead(output_dim10) self.reg_head TaskHead(output_dim1) def forward(self, x): features self.shared(x) class_out self.class_head(features) reg_out self.reg_head(features) return class_out, reg_out4.2 GradNorm实现class GradNorm: def __init__(self, model, shared_params, num_tasks, alpha0.5, lr0.025): self.model model self.shared_params shared_params self.num_tasks num_tasks self.alpha alpha # 初始化权重为1 self.weights torch.ones(num_tasks, requires_gradTrue) self.w_optimizer torch.optim.Adam([self.weights], lrlr) self.initial_losses None def set_initial_losses(self, losses): 记录训练开始时的loss用于计算相对训练速率 self.initial_losses [l.detach().clone() for l in losses] def compute_grad_norms(self, task_losses): 计算每个任务在共享层上的梯度范数 grad_norms [] for i, loss in enumerate(task_losses): self.model.zero_grad() weighted_loss self.weights[i] * loss weighted_loss.backward(retain_graphTrue) # 收集共享层梯度范数 norm_sq 0.0 for p in self.shared_params: if p.grad is not None: norm_sq p.grad.norm() ** 2 grad_norms.append(torch.sqrt(norm_sq)) return torch.stack(grad_norms) def update_weights(self, task_losses): 更新任务权重 grad_norms self.compute_grad_norms(task_losses) # 计算平均梯度范数 mean_grad_norm grad_norms.mean() # 计算相对逆训练速率 loss_ratios torch.stack([ task_losses[i] / self.initial_losses[i] for i in range(self.num_tasks) ]) # 计算目标梯度范数 target mean_grad_norm * (loss_ratios ** self.alpha) # 梯度损失 grad_loss torch.abs(grad_norms - target).sum() # 更新权重 self.w_optimizer.zero_grad() grad_loss.backward() self.w_optimizer.step() # 归一化权重 with torch.no_grad(): self.weights.data self.weights.data / self.weights.data.sum() * self.num_tasks self.weights.data torch.clamp(self.weights.data, min0.01) return grad_loss.item()4.3 训练循环def train(model, dataloader, epochs50): shared_params list(model.shared.parameters()) gradnorm GradNorm(model, shared_params, num_tasks2, alpha0.5) network_optimizer torch.optim.Adam(model.parameters(), lr1e-3) # 先跑一遍记录初始loss model.eval() initial_losses [0.0, 0.0] with torch.no_grad(): for x, y_class, y_reg in dataloader: class_out, reg_out model(x) initial_losses[0] F.cross_entropy(class_out, y_class).item() initial_losses[1] F.mse_loss(reg_out.squeeze(), y_reg).item() initial_losses [torch.tensor(l / len(dataloader)) for l in initial_losses] gradnorm.set_initial_losses(initial_losses) model.train() for epoch in range(epochs): for x, y_class, y_reg in dataloader: class_out, reg_out model(x) class_loss F.cross_entropy(class_out, y_class) reg_loss F.mse_loss(reg_out.squeeze(), y_reg) task_losses [class_loss, reg_loss] # 更新网络参数 total_loss sum( gradnorm.weights[i] * task_losses[i] for i in range(2) ) network_optimizer.zero_grad() total_loss.backward() network_optimizer.step() # 更新GradNorm权重 grad_loss gradnorm.update_weights(task_losses) print(fEpoch {epoch}, weights: {gradnorm.weights.data})4.4 参数选择与调优记录我在一个实际项目里跑过这套代码任务是同时做用户行为分类和停留时长回归。初始loss分别是2.3和15.6量级差了近7倍。用固定权重1:1时回归任务完全主导训练分类准确率卡在随机水平。换成GradNorm后权重自动调整到分类约1.6、回归约0.4两个任务都正常收敛。$\alpha$的选择上我试了0.3、0.5、0.8三档。0.3时回归任务还是偏强分类收敛慢0.8时分类任务被过度加权回归误差偏大0.5最平衡。最终分类准确率比固定权重提升了12个百分点回归MSE下降了约18%。一个实操细节GradNorm的权重更新频率可以比网络参数更新低。比如每2-3个batch更新一次权重能减少计算开销效果几乎不变。5. 常见问题与排查技巧实录5.1 权重震荡不收敛现象$w_i$在训练过程中剧烈震荡甚至出现负数。原因通常是$w_i$的学习率太大或者$\alpha$设置过高导致梯度损失曲面太陡。解决把$w_i$的学习率从0.025降到0.01或0.005。降低$\alpha$到0.3左右。在权重更新后加clamp限制$w_i$在[0.1, 2.0]范围内。5.2 某个任务权重被压到接近0现象训练一段时间后某个任务的$w_i$持续下降最终接近0该任务完全学不到东西。原因这个任务可能太简单收敛太快GradNorm认为它“不需要关注”了。或者$\alpha$太大过度惩罚了快速收敛的任务。解决降低$\alpha$。给$w_i$设下限比如最小0.1保证每个任务至少有基础权重。检查这个任务的loss是否本身有问题比如标签噪声太大导致loss降不下去。5.3 计算开销太大现象训练速度比单任务慢了好几倍。原因GradNorm需要为每个任务单独反向传播计算梯度范数任务数多时开销线性增长。解决降低权重更新频率比如每5个batch更新一次。只对共享层的最后一层计算梯度范数而不是所有共享层参数。如果任务数超过5个考虑改用DWA或Uncertainty Weighting。5.4 初始loss测量不准现象训练初期权重调整方向完全不对。原因$L_i(0)$测量时用了太少的batch或者用了训练中的模型状态。解决用至少50-100个batch测量初始loss。确保测量时模型处于eval模式且没有dropout/batchnorm的随机性影响。如果初始loss波动大多测几次取平均。5.5 共享层参数选择错误现象GradNorm完全不起作用权重几乎不变。原因可能把任务head的参数也当成了共享层或者共享层参数列表为空。解决打印shared_params的长度和名称确认只包含共享特征提取器的参数。确认这些参数在反向传播时确实有梯度不是被freeze了。我踩过最坑的一次是共享层里有个BatchNorm层训练时它的running_mean和running_var不是可学习参数但会影响梯度。GradNorm计算梯度范数时把这些也算进去了导致权重调整异常。后来只对weight和bias计算范数就正常了。6. 几个实战中的经验补充GradNorm不是银弹。它解决的是“梯度量级和收敛速度不平衡”的问题但如果任务之间本身存在根本性的冲突比如一个任务需要旋转不变性另一个需要旋转敏感性GradNorm也救不了。这种情况下需要考虑网络结构上的解耦比如MMoE或者PLE。另外GradNorm的权重是全局共享的所有样本用同一组$w_i$。如果数据里存在明显的子群体差异比如不同用户群体的任务重要性不同可以考虑样本级的权重调整但这已经超出GradNorm的范围了。最后说一个工程上的小技巧把GradNorm的权重变化曲线记录下来用TensorBoard或者简单的matplotlib画出来。训练结束后回看这条曲线能帮你判断任务之间的平衡状态。如果权重很快稳定在某个值附近说明平衡找到了如果一直在震荡说明$\alpha$或者学习率需要调。这个曲线比最终的accuracy指标更有诊断价值。