1. 多任务loss加权为什么会变成一场灾难1.1 一段我亲历的调参循环我一度被多任务学习的loss加权问题搞得非常焦头烂额。去年做一个室内场景理解项目模型同时要输出目标检测框、语义分割和深度估计三个任务共用一个骨干网络。最开始我把三个任务的loss直接相加训练了十几个epoch发现分割mIoU只有52%深度估计的rmse也一直下不来唯独检测任务的mAP还能看。于是我开始手动调权重从1:1:1试到10:1:0.1又从10:1:0.1试回0.5:2:5折腾了一周始终陷入抬一个压一个的死循环——把分割权重调大分割涨了三个点检测和深度立刻掉两个点把深度权重调大深度是稳了分割又崩了。真正让我反思的是一张梯度范数曲线图。我把三个任务在共享层最后一层的梯度范数打出来发现检测任务的梯度范数长期在0.4左右分割任务只有0.05深度估计更是小到0.01量级。也就是说在共享骨干每一轮更新里检测任务用这么大的力把参数往自己的方向拉另外两个任务几乎拽不动它。当时GradNorm这篇论文我已经读过但一直觉得用起来要处理二阶梯度、要维护额外的优化器太麻烦就没上手。后来被手动调参折磨得不行才咬牙把GradNorm完整复现了一遍替换掉手动权重。第一天跑通看到权重自动从 1.0/1.0/1.0 慢慢偏移向 0.8/1.2/1.5我心里那块石头才算落地。1.2 失衡背后的三层原因多任务训练失衡这种现象几乎每个做过的人都遇见过但很少有人认真拆解它到底由哪些因素叠加而成。我后来复盘至少有三个层次的原因。第一层是loss尺度不一致。检测任务用的focal loss数值通常在1~3之间分割任务用cross entropy加aux loss大概在0.3~0.6深度估计如果走smooth-L1数值经常在0.05以下。三个数值差距几十倍直接加权求和小数值的loss在总loss里几乎不占分量这就像三个人合伙抬桌子其中一个人用了全力另外两个只是象征性搭了把手桌子的方向自然被大力气的人控制。第二层是梯度尺度与loss尺度不成比例。这是很多人忽略的点。Loss数值大的任务它的梯度不一定大loss数值小的任务梯度也不一定小。因为梯度还要过head结构、激活函数、归一化层每层都会缩放信号。有的任务head里加了个sigmoid梯度直接被压到0.x有的任务用了一个没有约束的线性输出梯度反而能冲上几个数量级。所以单纯按loss大小去调权重本质上是在猜loss和梯度之间那条看不见的换算关系猜得准是运气猜不准是常态。第三层是训练速率不一致。就算两个任务loss尺度相同、梯度尺度也恰好接近它们的学习速度还可能天差地别。比如一个任务在初期就迅速下降另一个任务要跑到一半才开始收敛。固定权重下快速任务把梯度降得很小以后慢速任务即使有不错的权重所获得的梯度也会因为共享参数被快任务带偏而难以翻身。这三个因素搅在一起就形成了我上面那种手动调参怎么调都别扭的局面。1.3 GradNorm想解决的是权重怎么动这个问题理解了上面三层原因再看GradNorm就比较清晰了。它不试图去修改网络结构也不改变任务内部的loss函数而是把注意力放在任务权重w_i本身上——而且不是给出一组静态最优权重而是让权重在训练过程中持续变化。核心思想一句话总结根据每个任务在共享参数上产生的梯度范数动态调整它的loss加权系数使各任务的梯度尺度保持在一个合理的平衡状态。换句话说它不做谁的loss大就压低谁这种经验操作而是直接盯着共享网络那一层收到的力来做事。这也解释了为什么它叫Gradient Normalization——归一化的对象不是loss也不是feature而是梯度。2. GradNorm在衡量什么梯度范数、反向训练速率与对齐目标2.1 三个关键量的数学定义要真正理解GradNorm必须先抠清楚它使用的三个量。我用 W 表示共享网络最后一层共享层的参数用 L_i(t) 表示第 i 个任务在第 t 步的loss用 w_i(t) 表示该任务当前的加权系数。第一个量是每个任务在共享层上的梯度范数G_W^{(i)}(t) || \nabla_W (w_i(t) \cdot L_i(t)) ||_2注意这里计算的是该任务加权后的loss对共享层参数的梯度的二范数。它衡量的是这个任务在当前时刻对共享参数施加的推力大小。GradNorm选择的是最后一层共享层而不是所有共享层——因为所有层的梯度会随着网络深度发生变化直接对所有层做范数求和会让目标变得不稳定而最后一层共享层既靠近任务分叉点又能代表共享表征被拉动的方向。第二个量是平均梯度范数\overline{G}_W(t) \frac{1}{T} \sum_i G_W^{(i)}(t)T是任务总数。这个量是动态的每步都会重新计算代表当前整体梯度的平均规模。第三个量最有意思叫反向训练速率inverse training rater_i(t) \frac{L_i(t) / L_i(0)}{\frac{1}{T} \sum_j L_j(t) / L_j(0)}L_i(0)是任务在训练开始时的初始loss。所以 L_i(t)/L_i(0) 表示任务i当前的loss相对起点下降到了什么比例比如从1.0降到0.4就是0.4。分母是所有任务这个比例的平均值。这样一来r_i(t)大于1说明任务i比平均速度慢小于1说明它比平均速度快。2.2 为什么对齐梯度范数而不是对齐loss数值我见过不少复现GradNorm的人上来就问为什么不直接让各任务的loss相等这个问题问得其实很好答案也值得展开。假设任务A的loss是3.0任务B的loss是0.03你为了让它们看起来公平给B乘以100让两个加权项都是3.0。但B任务loss的梯度在共享层上的范数经过head的层层缩放之后可能是0.001。你哪怕把它乘到1000它仍然是一股微弱的拉力。反过来任务A的loss虽然只有3.0但梯度范数可能高达10它对共享参数的影响远大于B。所以loss数值本身不是共享网络真正感受到的东西梯度才是。这就好比两个人搬箱子一个人看起来喊得很大声loss很大实际出的力很小另一个人闷不吭声但每一次都用腰部肌肉发力梯度很大。你只听声音去调整出力永远调不对。GradNorm等于给你在共享层装了一个力量计直接告诉你每个人实际用了多大力然后在此基础上做平衡。2.3 对齐目标的直觉与α的含义GradNorm的对齐目标并不是简单地让每个任务梯度范数等于平均值。如果那样做就会忽略训练速率差异。它的目标是G_W^{(i)}(t) \rightarrow \overline{G}_W(t) \times [r_i(t)]^\alpha其中 α 是一个超参数论文里记为alpha控制训练速率差异对目标的影响强度。如果 r_i(t) 大于1说明任务i学得慢那它的目标梯度范数就该比平均值大给共享参数施加更大的牵引力让它加速追赶。如果 r_i(t) 小于1说明这个任务已经学得比较快了目标梯度范数可以适当低于平均值避免它继续夺取过多共享表征资源。α越大慢任务被赋予的目标就越高追赶力度越猛α越小这种按速率纠偏的力度就越弱。当 α0 时所有任务的目标都退化为同一个平均梯度范数退化成只平衡梯度尺度、不区分训练快慢的简化模式。我在实践里一般默认 α1.0 起步然后看慢任务是否得到有效追赶如果追得太猛导致快任务显著下降就往下调如果慢任务还是跟不上就往上调。论文中的实验也显示α在0.5到1.5这个区间内通常都能获得不错的结果但不同的task组合对α的敏感度差异很大不能盲抄。3. GradNorm的每次更新到底改了什么算法与数学流程3.1 一次完整更新流程的逐行拆解GradNorm的完整更新流程我在复现时把它拆成了一张十步清单。每一步都有明确的动机少一步都会出问题。第一步记录初始loss。在训练正式开始前先对一小批数据做前向得到每个任务的初始loss L_i(0)。这一步非常关键后面所有r_i(t)都依赖它。不能用第一个训练batch的loss临时充当因为第一个batch往往有噪声随机波动会给后续所有步骤引入不稳定因素。我通常的做法是取前50~100个batch的平均值或者先用一个固定的随机种子对一小块验证集做前向。第二步正常前向得到当前batch下每个任务的原始loss L_i(t)。第三步计算当前任务权重下的加权总lossL_total \sum_i w_i(t) L_i(t)。第四步计算每个任务的梯度范数G_W^{(i)}(t)。这里要强调不是对总loss求梯度再分解而是分别对 w_i(t)L_i(t) 这个单独任务项求梯度。这样得到的是每个任务独立施加在共享层上的梯度而不是混合后的梯度。第五步计算平均梯度范数\overline{G}_W(t) 和反向训练速率 r_i(t)。第六步计算GradNorm lossL_{GN} \sum_i | G_W^{(i)}(t) - \overline{G}_W(t) \times [r_i(t)]^\alpha |_1这个loss的物理含义是当前任务权重下各任务真实梯度范数与期望梯度范数之间的距离。距离越小说明各任务在共享参数上的发言权越接近目标状态。第七步计算GradNorm loss对w_i的梯度并用一个独立的SGD优化器更新w_i。这一步相对复杂后面会单独讲。第八步重新归一化权重保证所有任务权重之和始终等于任务数T\sum_i w_i(t1) T这个约束是为了防止权重在长期迭代中不断漂移要么全部膨胀要么全部萎缩。保持总权重恒定的同时允许内部比例自由变化。第九步用更新后的权重重新计算总lossL_total \sum_i w_i(t1) L_i(t)并清零网络参数上的梯度。第十步用新的总loss对网络参数做反向传播和常规的参数更新。3.2 为什么对w_i的梯度要走二阶传播第七步是整个GradNorm里最容易让实现者栽跟头的地方。表面上看我们要更新的是w_i而GradNorm loss里也直接含有 w_i 吗其实没有直接显式出现。GradNorm loss里的 G_W^{(i)}(t) 是通过用 w_i L_i 对共享参数求梯度得到的。也就是说w_i 藏在梯度计算的过程里是一个「先对参数求梯度、再对权重求梯度」的两步关系。用反向传播的术语说计算 G_W^{(i)} 时涉及对 W 的一阶导数而 G_W^{(i)} 对 w_i 的导数是在这个一阶导数基础上再求一次关于 w_i 的导数所以是二阶导数。这要求在第一次求梯度时保留计算图也就是通常所说的 create_graphTrue、retain_graphTrue。如果不保留图w_i 的梯度会直接变成None整个GradNorm静默失败。另外一个关键细节目标值 \overline{G}_W(t) \times [r_i(t)]^\alpha 在设计上被当作常量处理不参与对w_i的导数传播。因为在GradNorm的设计里target是希望达到的目标而不是可以随权重变化而变化的量。如果不把它detach即从计算图里脱离那么它在对w_i求导时也会贡献梯度这会让算法变成另一个含义完全不同的东西。绝大多数实现里用 target.detach() 来切断。3.3 权重归一化的顺序为什么不能调换我踩过的一个小坑是先归一化再更新还是先更新再归一化论文里的正确顺序是先更新权重后做归一化。这背后的逻辑是GradNorm loss对w_i的梯度是在当前归一化状态下的梯度更新完后再归一化相当于对这次更新的结果做了一次模长调整不影响更新方向只影响步长。反过来如果先归一化再更新下一轮的权重可能立刻破坏掉刚维持好的总和约束导致归一化形同虚设。这个顺序虽然只影响一行的代码位置但在多任务loss尺度差距很大的情况下会让权重曲线变得极其不稳定。4. PyTorch复现GradNorm核心代码、优化技巧与常见坑4.1 一次前向拿到所有任务独立梯度的小技巧复现GradNorm,第一步是解决如何高效地拿到每个任务的独立梯度。朴素做法是对每个任务轮流 backward再统计梯度范数。但这样会有两个问题一是每次backward都会累加梯度你需要先记住上一次的结果再清零导致周期的梯度清零操作很多二是计算图反复被销毁重建效率很低。更稳妥的做法是使用 torch.autograd.grad它对每个任务的加权loss分别求关于共享参数的梯度。关键参数有三个retain_graphTrue因为要在一个计算图上多次求梯度默认的backward会释放图所以必须保留。create_graphTrue这是为了保留计算图使后续GradNorm loss能够对w_i求二阶梯度。allow_unusedTrue防止某些任务没有用到某层参数时报错。import torch def compute_task_gradient_norms(task_losses, shared_params, weights): task_losses: list of per-task losses, each shape [] shared_params: list of shared layer parameters weights: tensor of shape [T], requires_gradTrue grads [] for i, loss in enumerate(task_losses): weighted_loss weights[i] * loss grad torch.autograd.grad( weighted_loss, shared_params, retain_graphTrue, create_graphTrue, allow_unusedTrue, )[0] grads.append(torch.norm(grad)) return torch.stack(grads) # [T]这里只取了共享参数列表里的第一个参数来计算梯度范数实际使用中可以选择最后一层共享层。如果你希望用所有共享层的平均并不推荐需要把每层的梯度范数取出来后再走一次平均。4.2 一个可直接嵌入训练的GradNorm模块下面这段代码是我在多任务训练里实际用过多轮的版本。它把权重初始化、更新、归一化全部封装成一个容器外部只需要每个iteration传入各任务的loss和初始loss即可。class GradNorm(torch.nn.Module): def __init__(self, num_tasks, alpha1.0, lr0.025, init_lossesNone): super().__init__() self.num_tasks num_tasks self.alpha alpha self.lr lr self.weights torch.ones(num_tasks, requires_gradTrue) self.init_losses init_losses # shape [T]记录L_i(0) def set_init_losses(self, init_losses): self.init_losses torch.as_tensor(init_losses, dtypetorch.float32) def forward(self, task_losses, shared_params): if self.init_losses is None: raise ValueError(必须先提供初始loss) # 1. 计算每个任务的梯度范数 G self._compute_grad_norms(task_losses, shared_params) # 2. 计算平均梯度范数 G_avg G.mean() # 3. 计算反向训练速率 ratios torch.stack(task_losses) / self.init_losses ratio_avg ratios.mean() r ratios / ratio_avg # 4. 计算target并detach target G_avg * (r ** self.alpha) # 5. GradNorm loss grad_norm_loss torch.sum(torch.abs(G - target.detach())) # 6. 更新权重 (grad_norm_loss).backward(retain_graphTrue) with torch.no_grad(): self.weights - self.lr * self.weights.grad # 7. 归一化权重保持和等于任务数 with torch.no_grad(): self.weights.data self.weights.data / self.weights.data.sum() * self.num_tasks with torch.no_grad(): self.weights.grad.zero_() return self.weights def _compute_grad_norms(self, task_losses, shared_params): grads [] for i, loss in enumerate(task_losses): weighted_loss self.weights[i] * loss grad torch.autograd.grad( weighted_loss, shared_params, retain_graphTrue, create_graphTrue, allow_unusedTrue, )[0] grads.append(torch.norm(grad)) return torch.stack(grads)注意这个模块里的 self.weights 是带梯度的叶子张量它本身是不是 nn.Parameter 不重要只要有requires_gradTrue即可。更新权重时用的是纯SGD没有momentum、没有weight decay这也是论文的设定。我在工程中发现GradNorm的更新器学习率不宜与主网络学习率等同一般取主网络学习率的0.1到0.01倍稳定性会好很多。4.3 我在实现中踩过的三个坑第一个坑忘记对target做detach。最初的实现里我直接写成 G - target结果训练了200步后发现权重几乎没有变化各任务梯度差还是老样子。排查了半天才发现target里包含 w_i 对 G_avg 和 r_i 的间接依赖这部分propagate回去的梯度和真实GradNorm目标相抵触把我想要的更新方向给抵消掉了。解法就是那行 target.detach()一行代码解决。第二个坑初始loss直接用第一个batch。我当时图省事把第一个step的loss当作 L_i(0)没有做任何平滑。后果是后续每个batch的 r_i 都在剧烈震荡权重曲线像心电图一样跳来跳去训练非常不稳定。换成了前100个batch的平均loss作为 L_i(0) 之后权重曲线明显平滑下来。如果你等不了100步50步也可以但不要少于20步。第三个坑权重更新频率。论文的算法里是每步都更新权重我也曾照着做结果在任务数较多4个以上且batch较小的场景里权重几乎抖到没法看。后来我把GradNorm的更新频率改成每5步一次也就是每5个batch才根据当前梯度范数计算一次权重并更新网络参数更新倒是保持每步都做。效果稳定了很多而且因为GradNorm的额外反传频繁更低训练时间也省了。这一点论文里不太强调但实际工程里我觉得值得注意。5. GradNorm的适用边界什么时候上什么时候别动它5.1 判断一个多任务模型是否适合GradNorm不是所有多任务模型都适合上GradNorm。我总结下来适合的情况有几个特征模型有一个比较深的共享骨干任务head在最后一层共享层之后才分流任务之间的loss尺度差异大或者训练速率差异大你手动调权重已经调了一到两轮仍然不满意。反过来如果你的模型共享层很薄只有一两个卷积层各任务的梯度在共享层上几乎没有交互那么GradNorm能做的就非常有限。它只调整共享参数感受到的推力平衡如果共享参数本身只占模型很小一部分你真正需要解决的是每个任务分支内部的收敛问题而不是加权问题。另外GradNorm处理的是任务之间的平衡处理不了任务内部的问题。一个任务自身的label噪声很大、数据量很少GradNorm再怎么调权重也没办法凭空增加它的有效信号。我在实际使用中还见过有些人把GradNorm当作万能药结果发现分割任务质量问题出在类别不平衡上和任务间权重没有半毛钱关系。定位问题发生的位置比直接上工具更重要。5.2 和不确定度加权相比怎么选多任务加权方法里另一个绕不开的baseline是Kendall等在2018年提出的不确定度加权Uncertainty Weighting。它的思路和GradNorm完全不同为每个任务引入一个可学习的噪声参数 σ_i然后按 1/(2σ_i²) 缩放loss再加一个 log σ_i 作为正则项。实践中的效果这两个方法经常接近。区别在于对比项GradNorm不确定度加权更新对象任务权重w_i噪声参数σ_i核心信号共享层梯度范数loss不确定性homoscedastic uncertainty)需要额外计算每个任务对共享参数的梯度成本较高极小只多一个参数对loss类型的假设无需要loss符合特定概率模型如MSE、交叉熵可解释性直观对齐梯度相对抽象数值稳定性需小心处理二阶梯度相对稳定我的选择习惯是如果任务数少2~4个且loss类型差别大优先GradNorm因为它不关心loss分布形式直接对标梯度如果任务很多超过5个或者只想快速看一个baseline优先不确定度加权因为它的额外计算基本可以忽略。如果你的loss函数不是标准的MSE或CE不确定度加权推导出来的缩放形式可能不再成立这时候GradNorm反而更稳妥。5.3 GradNorm不擅长的场景GradNorm并不擅长处理梯度方向冲突。它按范数对齐了每个任务的力的大小但完全不碰力的方向问题。如果任务A想把语义特征往左拉任务B想把同样的特征往右拉即使它们范数大小完全一致共享参数还是会被反复撕扯收敛缓慢。这类冲突需要的是PCGrad、CAGrad这类基于梯度方向投影/聚合的方法而不是GradNorm。我见过一些人把GradNorm吹成全场景适配其实它解决的是力量大小失衡不是方向矛盾。另外GradNorm对任务数量超过一定规模后稳定性会变差。任务达到7~10个以上每个任务梯度的范数波动本身就很大再叠加训练速率的比值权重更新会明显抖动。这时候建议要么降低更新频率要么对梯度范数做滑动平均再输入给GradNorm。6. 用GradNorm调通一个三任务模型的完整记录6.1 案例背景与初始困境回到文章开头提到的那个室内场景理解项目。模型是SegNet风格的共享encoder三个head分别输出目标检测框focal loss、语义分割cross entropy dice、深度估计smooth-L1。数据集是自己采集的5万张室内图。最初手动权重是1:1:1训练20个epoch之后三个任务的loss分布相差悬殊检测loss 1.2左右分割loss 0.4左右深度估计0.06左右。我统计了三个任务在共享层最后一层的梯度范数。检测任务是0.45分割任务是0.06深度估计是0.008。这意味着深度估计任务对共享参数的拉动力只有检测任务的一百八十分之一左右它的深度估计分支在自己的head里即使能优化也无法通过共享骨干反向影响整个模型的表征结构。手动把深度权重调到20倍以后深度任务梯度范数是起来了但检测和分割的梯度又同时被摊薄导致共享层的表征开始往深度方向偏移检测mAP掉了3个点。那段时间我试过一个土办法分阶段训练。前10个epoch只训检测和分割把共享特征训出来再解锁深度任务一起训。结果深度任务适应不了已经成型的特征分布rmse一直很差。整个调参过程持续了两周产出却越来越差。6.2 GradNorm上线后的变化轨迹后来我把手动权重全部删掉用GradNorm从1:1:1开始自动调节α取1.0GradNorm的更新学习率设为0.02每5个batch更新一次权重L_i(0)用前100个batch的平均loss。前500步权重并没有立刻大动。梯度范数差异那么大按理说GradNorm应该很快把权重拉开才对。我分析了一下原因是初始阶段所有任务的loss都在快速下降r_i(t)里 L_i(t)/L_i(0) 的比值还很小target对慢任务的放大作用还没到位。到第1000步左右权重开始明显分化分割任务权重升到1.25附近深度任务升到1.5附近检测任务降到0.75左右。这个趋势基本符合直觉检测任务本身学得快、梯度大被自动降权分割和深度学得慢、梯度小被自动增产。跑完30个epoch最终结果比两周手动调出来的还要好分割mIoU从52.1%提到57.6%深度估计rmse减少了约11%检测mAP只掉了0.4个点。手动最好的一组权重我是按检测1.0、分割2.0、深度12.0配的表面上看着合理实际上因为深度任务head内部的梯度路径很长把它权重加到12效果反而不如GradNorm自动给出的1.5。这件事让我对靠直觉配权重彻底失去了信心。6.3 几个值得固化的使用习惯经过这个项目我养成了几个使用GradNorm的习惯在这里分享出来。第一永远先画梯度范数曲线再决定要不要上GradNorm。先把各任务在共享层最后一层的梯度范数打出来如果差异在10倍以内手动权重就够了没必要为了提一点点效果付出额外计算。只有差异超过一个数量级GradNorm的价值才体现得明显。第二GradNorm的warm-up要保留。我的默认做法是前200步正常训练权重冻结在1:1:1让各任务都进入一个稳定的下降轨道200步之后才让GradNorm开始调整权重。如果一上来就让GradNorm介入初始loss还没稳定L_i(0)参考价值很低权重很容易被带偏。第三验证集上的梯度范数比训练集更可靠。训练集上的梯度范数受batch噪声影响很大尤其是有难样本采样的时候。我的习惯是每10个epoch在固定验证集上算一次各任务的梯度范数看它们是否真的收敛到了同一尺度附近。如果验证集上梯度范数差异仍然很大就要检查是不是某个任务的head本身存在梯度消失或结构设计问题。第四GradNorm不是终点而是一面镜子。它把各任务在共享层上的拉力可视化出来让你能很快判断某个任务是否在共享表征层面失去了话语权。即使最终你不采用GradNorm产出的权重那几条梯度范数曲线也已经告诉了你该往哪个方向调参。我个人在实际操作中还有一个体会GradNorm最常被误解的地方是自动两个字。它确实自动调了权重但它调的是让各任务梯度范数逼近同一目标的权重至于这个目标是不是你业务上真正想要的它并不知道。如果你的业务核心是检测那该保护检测指标的还是要单独留validation监控不要因为分割任务梯度小就放任GradNorm给分割疯狂加权最终检测掉了你也不一定受得了。我在后续的项目里会先按业务重要性设定一个基准权重范围再让GradNorm在这个范围内做小幅调节而不是完全放开。这个方法不算论文正统但工程上很实用。