多任务学习梯度失衡怎么办？GradNorm原理与PyTorch实战

做过多任务模型的同学应该都经历过这种尴尬网络结构搭好了数据也齐了训练却怎么都不对劲。我在一个共享特征层、同时预测点击率和停留时长的双目标模型里卡了两周日志里主任务的指标很好看但辅任务的损失怎么压都下不去。后来把视线从模型结构转移到梯度层面试了 GradNorm 这个自适应梯度归一化方法才明白问题根本不是模型容量不够而是两个任务的梯度在共享层上打架一个任务把另一个任务的训练节奏彻底带偏了。这篇就从问题根源讲到 PyTorch 实操把训练中踩过的坑和能直接抄的代码一起记录下来。1. 多任务学习的失衡问题到底在失衡什么1.1 一个真实场景主任务梯度“绑架”了共享层多任务学习的目标很朴素一次训练多个相关任务共享底层表示省资源还能互相提携。但真正训练起来你会发现共享层更新方向几乎被一个任务垄断。我在推荐场景的实践中感受特别明显。CTR 任务的交叉熵损失在训练初期下降很快数值在 0.5 到 0.7 之间波动而停留时长任务用的是回归损失数值动辄几百。表面上看是“量纲不同”深层问题是共享层参数每次更新时梯度方向由所有任务梯度的加权和决定。哪个任务梯度向量的 L2 范数大哪个任务就在“物理意义上”主导权重更新。举个例子。共享层某个参数收到 CTR 任务传来的梯度是 0.02收到停留时长任务传来的梯度是 2.0。哪怕你对停留时长任务只给了 0.1 的损失权重加权后它仍然是 0.2依然比 CTR 的 0.02 大一个数量级。于是整个共享层被回归任务牵着走分类任务该学的特征被淹没。反过来如果回归任务因为数值原因梯度范数很小它又会变成“隐形任务”几乎不参与学习。这种不平衡本质上和业务优先级无关它只和梯度向量的数学范数有关。1.2 为什么手动调权重调来调去还是不行很多人第一反应是给每个任务的 loss 乘一个固定系数手动调。我也这么干过。问题在于这组系数是随时间变化的。任务在不同训练阶段的难度不同一个任务可能前 1000 步学得飞快之后收敛变慢另一个任务前期几乎不动后期突然开始下降。固定的权重永远只能照顾一个时间段。你上午把权重调好下午训练一跑最优解又变了。而且手动搜索的代价是组合爆炸。两个任务算一个参数三个任务就要在二维平面里试五个任务基本只能靠玄学。我见过有人用贝叶斯优化搜任务权重搜出来的结果换一个随机种子就不复现原因就是静态权重本质上没有和训练动态挂钩。后来也试过不确定性加权它用可学习的噪声参数 σ 为每个任务动态分配权重思路很不错。但它对梯度的控制是间接的它通过“数值尺度归一化”让大 loss 任务自动降权并没有直接约束每个任务对共享层产生的梯度范数。换句话说它解决了“量纲不同”的表面问题没有解决“任务梯度贡献失衡”的本质问题。这也是我后来转向 GradNorm 的直接原因。方案是否动态控制对象存在的问题固定手动权重否损失数值搜索空间大不能适应训练阶段变化不确定性加权是损失数值间接控制梯度无法精确把握梯度平衡GradNorm是梯度范数需要额外计算梯度实现复杂度略高2. GradNorm 的核心逻辑与公式拆解2.1 三个关键定义全局梯度范数、任务梯度范数、相对学习速率GradNorm 是 2018 年提出的一种自适应损失平衡方法核心思想非常直接不同任务对共享层的梯度贡献应该和它们各自的学习难度成正比。学习越难的任务得到越大的梯度空间学习越快的任务权重被自动压下来别抢资源。要理解这个机制先要记住三个量。第一个是全局梯度范数G_W(t)。它是指当前总损失对共享网络参数 W 求梯度后把所有梯度向量拼起来算出的 L2 范数。这个量衡量的是整个多任务模型在共享层上的整体更新强度。第二个是单任务梯度范数G_i(t)。先忽略权重单独拿第 i 个任务的 loss 对同一个 W 求梯度再算 L2 范数。它反映的是“如果没有其他任务干扰这个任务想往共享层推多远”。训练初期不同任务的这个值往往差很多。第三个是相对学习速率r_i(t)。它定义为当前任务损失和它初始损失的比值r_i(t) L_i(t) / L_i(0)。如果这个值接近 1说明任务几乎没怎么学如果它掉到 0.3说明任务已经学得很快了。所以r_i是一个用损失下降速度来刻画“任务进度”的指标。这三个量合起来GradNorm 要做的事就清晰了它不想让所有任务的梯度范数完全一样那是平均主义。它想让每个任务最终贡献的梯度范数正比于G_W × r_i。学得慢的任务 r 大理应获得更大的梯度推动力学得快的任务 r 小慢慢退居幕后。同时乘上全局梯度范数 G_W是为了让整体更新强度仍然和当前训练节奏对齐不会因为权重归一化动作把总梯度撑爆或压没。2.2 目标方程每个任务的梯度贡献按难度分配用公式表达就是希望w_i * G_i ≈ G_W * r_i。这里出现权重w_i是因为总损失定义成加权和L_total Σ w_i L_i。第 i 个任务单独对共享层产生的梯度贡献就是权重乘上单任务梯度范数即w_i * G_i。GradNorm 的优化目标就是让这个值去逼近右侧的G_W * r_i。我最初读论文时有个疑惑为什么不直接让w_i * G_i G_W就算了还要乘一个 r_i后来在实验里想明白了。如果所有任务都追着同一个 G_W 跑相当于要求所有任务的梯度贡献拉齐。这个约束太强了会让难任务权重飙得非常高容易任务权重压得非常低训练过程很容易震荡。加一个 r_i 进去相当于给目标按任务难度做了一个调节简单的任务目标更低难的任务目标更高这样既平衡了梯度又保留了任务之间的差异信息。还有一个非常容易忽略的点全局梯度范数 G_W 本身也依赖权重。因为L_total是w_i L_i的加权和对 W 求梯度时w_i 被作为一个系数乘进去。所以当你对w_i求导更新权重时G_W 这个项还会产生关于 w_i 的梯度这在数学上是一个二阶导关系。这也是 GradNorm 实现里最容易写错的地方。2.3 权重更新与归一化一个很细的工程细节GradNorm 更新权重的做法是先定义一个平衡损失L_bal Σ (w_i * G_i - G_W * r_i)^2然后对这个损失关于w_i求梯度做一次梯度下降更新权重。注意这里的更新方向和直觉一致如果某个任务当前的梯度贡献w_i * G_i太大了超过了目标G_W * r_i平衡损失对这个任务的分量是正的梯度下降会把这个任务的权重往下拉。反过来如果它当前贡献太小权重就会被推上去。整个过程是自适应的。更新完权重之后还有一个强制步骤权重归一化。通常让所有权重的和保持为任务数 T或者总和为 1取决于你的初始化方式。这一步是必要的因为权重在梯度下降过程中会整体漂移。如果不做归一化所有任务权重同时变大或变小等价于改变了全局学习率模型训练速度会忽快忽慢很难控制。实现层面还有个小坑计算L_bal对w_i的梯度时必须让框架保留高阶计算图。在 PyTorch 里就是torch.autograd.grad的create_graphTrue。如果不设置这个参数G_W 对 w 的依赖会被切断更新规则退化成只看单任务梯度范数和目标的差值虽然也能跑但和论文原始推导会有偏差。3. PyTorch 落地一份可直接改的 GradNorm 实现3.1 共享层怎么选不是所有层都适合GradNorm 里那个 W指的是多个任务共享的那部分网络参数。但具体选哪一层是个影响效果的关键决策。我的经验是选择“共享主干的最后一层参数”。比如你的模型是一个共享 MLP后面接了两个任务头那就取共享 MLP 最后一层的权重和偏置。原因有两个。第一这一层紧邻任务分支它接收到的梯度信号已经经过了足够多的特征变换任务之间的差异在这一层体现得最明显。第二相比共享主干的第一层这一层参数数量少很多计算梯度范数的开销也比较低。如果选了共享主干的第一层梯度要反向传播经过很多层信号里非线性的干扰很大而且参数规模大每步计算开销很感人。如果选得太靠后比如任务头自己的独立参数那它根本不属于共享部分GradNorm 没有意义。3.2 核心代码权重更新模块和训练循环我用 PyTorch 写了一个最小实现核心就两个辅助函数加一个权重更新逻辑可以直接粘到自己的训练脚本里。import torch import torch.nn as nn import torch.optim as optim def grad_norm_of(params, loss, create_graphFalse): 计算某个 loss 对一组共享参数 params 的梯度 L2 范数。 create_graphTrue 时保留计算图用于计算对权重 w 的二阶导。 grads torch.autograd.grad(loss, params, retain_graphTrue, create_graphcreate_graph) return torch.cat([g.flatten() for g in grads]).norm() def set_task_weights(task_losses, task_weights, shared_params, init_losses, alpha0.025): task_losses: 当前每个任务的 loss 组成的列表 task_weights: 可训练参数形状为 [T]初始化通常均等 shared_params: 选定的共享层参数列表 init_losses: 训练开始时记录的每个任务初始 loss T len(task_losses) # 当前全局总损失对共享层参数的梯度范数 total_loss sum(w * loss for w, loss in zip(task_weights, task_losses)) G_W grad_norm_of(shared_params, total_loss, create_graphTrue) # 每个任务单独的梯度范数这里不含权重系数 w_i per_task_G [] for loss in task_losses: per_task_G.append(grad_norm_of(shared_params, loss, create_graphTrue)) # 相对学习速率 r_i注意要 detach不要让权重优化器反向传播到 loss 本身 r_i [(loss / init_loss).detach() for loss, init_loss in zip(task_losses, init_losses)] # 平衡损失希望 w_i * G_i 逼近 G_W * r_i balance_loss torch.zeros((), devicetask_weights.device) for i in range(T): diff task_weights[i] * per_task_G[i] - G_W * r_i[i] balance_loss balance_loss diff * diff # 对 task_weights 做一步梯度下降 task_weights.requires_grad_(True) balance_loss.backward() with torch.no_grad(): task_weights - alpha * task_weights.grad task_weights.grad.zero_() # 保证非负并归一化权重平均值为 1/T或者说总和保持为 1 task_weights.clamp_(min1e-6) task_weights task_weights / task_weights.sum() * T task_weights.requires_grad_(True)调用方式很简单初始化时给每个任务一个1/T的权重训练过程中每隔 N 步调用一次set_task_weights。# 初始化 task_weights torch.ones(T, requires_gradTrue) / T init_losses [first_loss_i.detach() for each task] # 训练循环里 for step, (x, y_list) in enumerate(train_loader): # 正常前向、计算每个任务 loss losses [loss1, loss2] # 每隔 100 步更新一次任务权重 if step % 100 0: set_task_weights(losses, task_weights, shared_params, init_losses, alpha0.025) # 计算总损失更新网络参数 total_loss sum(w.detach() * l for w, l in zip(task_weights, losses)) optimizer.zero_grad() total_loss.backward() optimizer.step()3.3 超参设置与稳定化处理GradNorm 的超参数不多但每个都很关键。第一个是权重更新学习率 α论文推荐 0.025。我在实践中的感受是0.01 到 0.025 之间是一个合理区间但如果把权重更新频率降得很低比如每 500 步才更新一次那 α 可以相应调小一些防止权重在两次更新之间剧烈突变。第二个是权重更新频率。理论上每一步都更新是最接近论文推导的但计算梯度范数的开销摆在那。我的做法是每 100 到 500 步更新一次。任务权重的变化不需要跟随每个 batch 的噪声走它应该反映一个相对稳定的学习趋势。低频率更新反而能减少震荡。第三个是数值稳定性处理。r_i 的计算要做除法初始损失如果非常接近 0会直接爆。我一般会在分母上加一个很小的 epsilon比如 1e-6。另外平衡损失里的 G_W 和 per_task_G 都来自梯度范数如果梯度在某些层恰好为 0会导致权重更新卡住。建议对梯度范数也做一次 detach 或者加一个小的光滑项避免纯零梯度。参数取值范围我的推荐备注α0.005 ~ 0.050.01 ~ 0.025更新频率低时用更小的 α权重更新间隔1 ~ 1000100 ~ 500间隔越大开销越低但收敛变慢权重归一化目标T 或 1总和为 T与初始化方式保持一致r_i 平滑必做可加 EMA防止单 batch 损失抖动影响权重4. 我在两个对比实验里观察到的变化4.1 实验设计一个可控的失衡场景为了确认 GradNorm 真的有效我搭了一个非常刻意的实验共享输入层和共享隐藏层后面接两个任务头一个是回归任务一个是分类任务。回归任务的 loss 是 MSE数值范围大概在 100 左右分类任务用交叉熵数值在 0.6 左右。两个任务的业务重要程度先设成相同也就是初始权重完全相同。实验分三组跑第一组是固定等权重w1 w2 0.5从头到尾不动。 第二组是手动调权重我先观察 10 个 epoch 的 loss 曲线然后手动设定一个“看起来公平”的权重。 第三组是 GradNorm初始权重同样是 0.5每隔 100 步更新一次。在固定等权重下回归任务几乎单方面主导了共享隐藏层的梯度。分类任务的 loss 在第一个 epoch 还能降之后基本进入平台期。手动调权重的情况稍好但我必须每隔几百个 step 就回去看一次训练曲线发现问题不对再停掉重调总共调了三四轮才找到一个勉强能用的静态权重。第三组的现象最有意思GradNorm 在一开始并没有直接给分类任务很大的权重因为初期两个任务都还在快速学习r_i 都比较接近。跑了几百步之后分类任务的 loss 开始下降变缓r_i 相对变大它的权重被自动抬升回归任务权重被压下去。整个过程我没有做任何干预。4.2 观察结果权重演化和训练曲线我记录了一组权重演化数据大致走势是这样的训练步数任务 A 权重回归任务 B 权重分类00.5000.5002000.4320.5688000.3710.62920000.3240.676这个走势非常符合直觉回归任务那个大数值 MSE 一开始会拿到很大的梯度范数所以权重被下调分类任务因为梯度范数相对偏小所以权重被上调以获得更多共享层的更新空间。对应到 loss 曲线上最显著的变化是总 loss 下降更平稳。固定权重下总 loss 曲线是一条“被回归任务主导”的陡降曲线分类任务的 loss 几乎是条平线GradNorm 组里两个任务的 loss 同步下降虽然分类任务最终的绝对 loss 不一定会比手动调权时更低但它和回归任务之间的“进度差”被明显拉近。对于需要多个指标同时达标的业务场景这个性质非常值钱。4.3 GradNorm 不是万能的边界这部分是我最想强调的。GradNorm 平衡的是梯度范数不是梯度方向。如果两个任务的梯度方向在共享层上本来就严重冲突比如夹角接近 180 度那 GradNorm 只会让它们以相似的幅值互相拉扯最终效果可能比固定权重还差。这种情况下要先去解决方向冲突问题比如用梯度投影类方法。此外GradNorm 对“共享程度”很敏感。如果共享层非常浅只有一层线性映射任务真正要学的东西基本都在各自分支里那 GradNorm 能调节的空间非常有限。我做过一个极端实验共享层只有一个 bias 参数剩下的全是独立分支GradNorm 把那个 bias 的权重调来调去对最终结果毫无影响。这提醒我GradNorm 适合共享主干较深的多任务模型不适合几乎完全独立的多头结构。任务数量也是一个限制。当任务数到十几个时每个任务的梯度范数都要单独计算一次训练开销成倍上升而且权重归一化会让每个任务的权重变得很小数值上容易不稳定。我个人的经验边界是3 到 8 个任务最舒服超过 10 个任务会更倾向用更轻量的不确定性加权或者干脆静态权重加搜索。5. 踩坑记录与排查速查5.1 权重变成负数或突然 NaN这是 GradNorm 实现里最常见的两个事故而且往往一起来。权重变成负数的原因是梯度下降越过 0。当某个任务当前梯度贡献远大于目标值时平衡损失对它求导会得到一个很大的负梯度一步就把权重推到负值。负权重在数学上还能算但业务上没有任何含义——难道要让模型“反着学”某个任务吗我的处理是在每次更新后强制clamp_min(1e-6)再做一次归一化。NaN 的原因就比较多了。最常见的是 r_i 出现0/0或inf/0。比如某个任务在初始化时 loss 恰好为 0或者初始 loss 记录到了一个小到接近 0 的值除法直接炸。要解决我推荐不要把初始 loss 当成分母而是取训练开始后第 10 个 epoch 的平均 loss 作为基准这样数值更稳定也更符合“热身之后再看进度”的直觉。5.2 梯度范数异常大或始终为零如果某个任务的梯度范数始终为 0说明共享层参数 W 对它来说是不可达的。这种问题一般出在任务分支和共享层之间的连接断了常见于某些任务头使用了 stop-gradient 操作或者输入特征被 mask 掉。排查方法很简单单独跑一次该任务的反向传播打印共享层参数的梯度值如果全是 0先回去检查数据流。如果梯度范数异常大比如突然比上个 step 大了几个数量级通常是因为某个任务的 loss 在某一步产生了尖刺或者权重更新间隔太长导致共享层参数在两次 GradNorm 更新之间跑得太远。我的处理是给 loss 加梯度裁剪同时对 r_i 做指数移动平均效果立竿见影。5.3 训练开销增加太多收敛变慢GradNorm 需要额外对共享层参数计算 T 1 次梯度如果共享层很大这部分开销肉眼可见。我实测过一个百万参数级别的共享 MLP每一步都更新权重训练时间增加了大概 60% 到 80%这是很多人试一次就放弃的原因。缓解开销的办法有两条。第一降低权重更新频率。每 200 步更新一次开销基本可以忽略而且效果几乎没有下降。第二只对选定的共享层参数做torch.autograd.grad不要对整个模型的所有参数做。因为 GradNorm 只需要最后一层共享层的梯度范数没必要在全模型上算。5.4 一个容易忽视的初始化问题初始权重设成均匀分布理论上是公平的但在实际训练里第一轮 update 时 r_i 基本上都会偏离 1 很多因为不同任务在初始阶段 loss 下降速度差异极大。如果初始化权重直接就是最终解的起点那没问题但如果想让权重演化过程更顺滑我建议先用固定权重跑 10 到 20 个 epoch记录每个任务的初始 loss 和平均梯度范数再开启 GradNorm。我个人实践里还有一个习惯把 GradNorm 的权重更新过程打到日志里。每更新一次把task_weights的值打印出来。这个信息比 loss 曲线更能反映模型对任务难度的判断它本身就是一份很好的调试材料。有时候看到某个任务的权重一路飙到 0.8 以上我就知道这个任务要么真的太难要么它的梯度在共享层被抑制了后者往往意味着数据采样或特征工程有问题。现象可能原因处理方案权重为负梯度下降越过 0clamp_min(1e-6)再归一化权重 NaN初始 loss 为 0 或 r_i 除零用第 10 个 epoch 的 loss 当基准加 epsilon某任务梯度范数恒 0共享层到任务头链路断开单独反向传播检查梯度检查数据流梯度范数突增loss 尖刺或更新间隔太长梯度裁剪 r_i 做 EMA 平滑训练时间暴增 60%每步都计算全模型梯度固定共享层参数列表每 200 步再更新权重演化剧烈震荡α 过大或更新频率过高α 降到 0.01更新间隔加大6. 个人实践心得GradNorm 对我最大的启发是多任务失衡问题的答案不在“重写损失函数”这一层而在“控制梯度”这一层。以前我调不动一个任务时第一反应是换网络结构、加正则、调数据采样现在我会先打印一张“各任务对共享层梯度范数贡献”的表看看谁真正在做决定。GradNorm 本身不是一个花哨的算法但它给了我一个非常直接的工具来观察和干预训练过程。最后分享一个小技巧你甚至可以不完全依赖 GradNorm 去自动更新权重而是只在训练早期开启它观察它学到的权重大小然后把这些权重“固化”下来再关掉 GradNorm 继续训练。这样做可以省去后期额外的梯度计算开销同时保留动态平衡的好处。我试过几次在任务相对稳定的场景下效果非常接近训练速度快了不少。如果你正被多任务权重调不明白的问题折磨不妨从打印梯度范数开始先看清楚再决定要不要上 GradNorm。