做多任务这几年,我最烦的一件事不是模型设计,而是配 loss 权重。两个任务还好说,三个以上就开始头疼:调好一组权重,训到一半发现有个任务已经完全躺平,梯度全被另一个任务霸占。手动调、grid search、写个 scheduler 慢慢试,都有点像玄学。后来接触到 GradNorm,这个思路让我有一种“早该这样”的感觉,它不需要为每个任务预设什么分布假设,也不用手工枚举权重,而是盯着共享层的梯度范数做自动平衡。这篇文章把 GradNorm 的原理、公式、PyTorch 实现和调参经验一起说清楚,适合正在做多任务学习、被 loss 加权折磨的算法工程师和研究人员。
1. 多任务学习中的梯度失衡:GradNorm 到底在解决什么问题
1.1 多任务学习为什么难:loss 加权不是拍脑袋
多任务学习(MTL)的核心假设是:多个任务之间共享底层表示可以互相促进,同时节省参数和计算量。这个假设在图像、语音、推荐等领域都得到了广泛验证,但实际训练时有一个尴尬的问题——共享网络的梯度是所有任务梯度的叠加,如果某个任务梯度大、另一个任务梯度小,小梯度的任务信号就会被淹没。
很多人第一反应是给 loss 加权,比如 A 任务 loss 乘 0.5,B 任务 loss 乘 2.0。这个思路没错,但怎么确定权重系数呢?最直接的办法是看 loss 尺度,把两个任务的初始 loss 拉平。可问题在于 loss 尺度只是一个粗略参考,两个任务损失函数不同、收敛速度不同、样本难度分布不同,初始 loss 相等不代表它们对共享层的梯度贡献相等。更麻烦的是,训练过程中各任务的相对难度一直在变化,固定权重跟不上这种变化。
我刚开始做多任务的时候,最常踩的坑就是:一个任务收敛得很快,另一个任务还很差,但共享层的梯度几乎被已收敛任务主导。因为已收敛任务的梯度范数往往也不小,甚至在分类任务里,一个接近收敛的任务梯度反而可能因为置信度不够而持续产出中等量级梯度。你在那里调权重,调到怀疑人生,最后发现换一个随机种子权重又得重调。
1.2 两种失衡:幅度失衡与速率失衡
GradNorm 论文把多任务优化中的失衡拆成两类,这个拆法非常关键。
第一类是梯度幅度失衡。不同任务的 loss 函数本身尺度就不同,比如回归任务用 MSE 可能动辄几十上百,分类任务用交叉熵通常是 0.1 到 1 的量级。即使都乘上各自权重,共享层的梯度范数也可能差一个数量级。幅度失衡的直接后果是小梯度任务几乎得不到有效更新。
第二类是收敛速率失衡。有些任务简单,训练几轮就收敛得差不多;有些任务难,loss 下降得很慢。如果一直给简单任务同样的权重,它贡献的梯度就会“挤占”难任务的学习空间。反过来,如果一开始就过度加权难任务,又可能在训练初期破坏简单任务的学习。速率失衡是动态的,固定权重解决不了。
GradNorm 的高明之处在于,它不只盯着梯度幅度,同时还引入了“相对反向训练速率”这个指标,把任务当前的收敛状态也纳入权重更新逻辑。相当于一个权重调节器同时看两个仪表盘:一个显示梯度大小,一个显示任务下降速度。
1.3 为什么手动调权和不确定性加权都不够
手动调权的问题前面说了,本质上是一个组合爆炸问题。任务数从 2 涨到 5,权重组合空间就完全没法枚举了。而且即使你枚举出一组看起来不错的权重,任务难度变化之后这组权重又可能失效。我试过在训练中动态调权,人工看曲线调,结果比训练本身还累。
另一个常见的方案是 Kendall 等人提出的不确定性加权,通过建模任务的 aleatoric uncertainty 来自动学习权重。这个思路很有意思,在回归任务上效果也不错,但它有几个前提条件没那么容易满足:需要对每个任务定义一个合理的概率分布,分类任务上的 noise parameter 容易出现梯度震荡,而且不确定性权重和网络参数交替更新的稳定性也需要额外调参。
GradNorm 走了一条更直接的路:不假设任何概率分布,直接把“共享层梯度范数之间的关系”作为优化目标。它要的是让每个任务在共享层上产生的梯度范数,维持在一个动态调整的目标值附近。谁贡献的梯度太弱就提高权重,谁贡献的梯度太强就降低权重。这个逻辑简单、直观,也不需要引入额外的建模假设。
2. GradNorm 核心公式拆解:三分钟看懂它在干什么
2.1 两个关键量:梯度范数与相对反向训练速率
GradNorm 要操作的是任务 i 的 loss 权重 w_i(t),它通过两个量来驱动权重更新。
第一个量是任务 i 在共享层参数 W 上的梯度范数:
G_W^(i)(t) = || ∇_W ( w_i(t) * L_i(t) ) ||_2
由于 w_i(t) 是标量,这个式子其实就是 w_i(t) 乘以任务 i 单独对共享层参数计算出的梯度 L2 范数。这里要注意,用来计算梯度范数的 W 不是全部网络参数,而是共享层的参数,通常是共享主干最后一层的参数。选择这一层的原因很直接:共享层是所有任务信息汇聚的接口,各任务分支从这里分叉出去,在这个位置观察每个任务对共享表示的梯度贡献最干净。如果把任务专属 head 的参数也算进来,head 参数的梯度会引入大量任务独有噪声,平衡的意义就被稀释了。
第二个量是相对反向训练速率 r_i(t):
r_i(t) = ( L_i(t) / L_i(0) ) / ( (1/T) * Σ_j ( L_j(t) / L_j(0) ) )
这个式子看着复杂,拆开说就是:先算出每个任务当前 loss 相对于训练初始 loss 的下降比例,然后用任务 i 的下降比例除以所有任务下降比例的平均值。如果任务 i 的 loss 下降得比平均水平快,r_i(t) 小于 1;如果它学得慢,r_i(t) 大于 1。r_i 越大,说明这个任务越“落后”,越需要更多梯度。
举一个直观的例子:两个任务初始 loss 都是 10,训练到某个时刻,任务 A 的 loss 降到了 2,任务 B 的 loss 降到了 8。那么任务 A 的下降比例是 0.2,任务 B 是 0.8,平均是 0.5,因此 r_A = 0.4,r_B = 1.6。显然任务 B 更需要关注。
2.2 目标梯度范数的构造与超参数 α 的作用
有了 G_W^(i)(t) 和 r_i(t),GradNorm 构造了目标梯度范数:
target_GW^(i)(t) = G_avg(t) * ( r_i(t) )^α
其中 G_avg(t) 是所有任务梯度范数的平均值,它表示当前整体梯度水平。target 的含义是:任务 i 的梯度范数应当约等于全任务平均梯度范数乘以一个由 r_i 和 α 决定的缩放系数。如果 r_i > 1(任务落后),target 就会大于平均值,促使任务 i 的权重提高;如果 r_i < 1(任务超前),target 就会小于平均值,权重会被抑制。
α 是一个关键超参数,控制着平衡的“激进程度”。α=0 时,target 对所有任务都等于 G_avg(t),相当于强行把每个任务的梯度范数拉到同一个水平。这在任务难度差异不大时可行,但如果任务难度差异明显,强行拉平反而会让简单任务被过度压制、难任务被过度抬升。α 越大,落后任务获得的梯度补偿越大,训练越倾向于“牺牲容易任务、优先难任务”。论文里的实验表明 α 取 0.5 附近效果较好,实际使用时我一般先试 0.5,然后在 0.25 到 1.0 之间搜索。
需要强调的是,G_avg(t) 本身也在随训练动态变化。这意味着 GradNorm 不是把梯度范数固定到某个绝对值,而是维持任务间梯度范数的相对平衡。这个设计比固定阈值稳健得多,因为它不需要人工估计“合理的梯度大小”。
2.3 从优化角度理解 GradLoss
GradNorm 的权重更新不直接优化总 loss,而是优化一个专门的梯度损失:
L_grad(t) = Σ_i | G_W^(i)(t) - target_GW^(i)(t) |
这是一个一范数损失。为什么选 L1 而不是 L2?因为单 batch 估计出的梯度范数方差非常大,L2 会对异常值更敏感,造成权重剧烈震荡。L1 更抗噪,训练曲线更稳。
更新 w_i 时,以 L_grad 对 w_i 求梯度,做一步梯度下降,然后重新归一化 w_i,使所有任务权重之和保持为任务总数 T。归一化是必要的,否则权重会无约束漂移。比如 T=2,初始 w=[1,1],某次更新后变成 [1.2,0.9],归一化后就是约 [1.14,0.86],总和仍是 2。这个约束保证权重只在相对关系上变化,不改变总体的梯度缩放尺度。
整个算法里存在两个优化器、两个更新尺度:主优化器按 common loss 更新网络参数,GradNorm 按梯度损失更新 w_i。这两个更新交替进行,w_i 的学习率通常比主网络学习率低一个数量级,否则权重震荡会很厉害。
3. PyTorch 实战:GradNorm 的可复现实现
3.1 算法流程速览
在写代码之前,先把完整的训练循环梳理清楚,避免被细节带偏。一个标准的 GradNorm 训练步骤长这样:
- 前向传播,得到每个任务的 loss L_i(t)
- 如果是第一个 batch,记录初始 loss L_i(0)
- 计算每个任务 loss 相对初始 loss 的下降比例,得到 r_i(t)
- 用
torch.autograd.grad计算各任务 loss 对共享层参数的梯度范数 - 计算加权梯度范数 G_W^(i)(t) = w_i(t) * ||∇_W L_i||_2
- 计算平均梯度范数 G_avg(t) 和目标梯度范数 target_GW^(i)(t)
- 构造梯度损失 L_grad,反向传播更新 w_i
- 归一化 w_i,使总和为任务数
- 用 total_loss = Σ w_i * L_i 更新主网络参数
顺序上有一个关键点:GradNorm 的权重更新必须放在主网络参数更新之前。因为第 4 步用autograd.grad(retain_graph=True)保留了计算图,是为了之后还能对主网络做 backward。如果先做了主 loss 的 backward 并释放图,再想计算各任务单独梯度就会报错——除非再做一次前向。所以循环里的顺序要严格遵守:先 GradNorm 更新,再主更新。
3.2 GradNorm 模块设计与代码实现
先写一个 GradNorm 模块,包含权重参数、梯度范数计算和更新逻辑。这里我直接给出我在项目中会用的版本,注释写清楚每个容易踩坑的地方。
import torch import torch.nn as nn class GradNorm(nn.Module): def __init__(self, num_tasks, init_w=1.0, alpha=0.5): super().__init__() self.num_tasks = num_tasks self.alpha = alpha self.w = nn.Parameter( torch.full((num_tasks,), init_w, dtype=torch.float32) ) def get_grad_norms(self, losses, shared_params): """ 计算每个任务单独对共享层参数的梯度 L2 范数。 shared_params: 共享主干最后一层的参数列表。 losses: 各任务 loss 的列表,每个 loss 都是标量张量。 """ grad_norms = [] for loss in losses: grads = torch.autograd.grad( loss, shared_params, retain_graph=True, allow_unused=True, ) # 有些参数可能与当前任务无关,返回 None,需要跳过 flat = torch.cat([ g.flatten() for g in grads if g is not None ]) grad_norms.append(torch.norm(flat)) return torch.stack(grad_norms) # shape (num_tasks,) def update(self, losses, shared_params, initial_losses, gradnorm_optimizer): """ 更新权重 w。losses 是当前 batch 的各任务 loss。 initial_losses 是训练开始时记录的初始 loss。 """ # 计算当前各任务梯度范数,shape (T,) grad_norms = self.get_grad_norms(losses, shared_params) # 计算相对反向训练速率 r_i loss_ratios = [] for i, loss in enumerate(losses): # 这里必须 detach,loss 本身带图,但 r_i 只是作为统计量使用 ratio = loss.detach() / initial_losses[i] loss_ratios.append(ratio) loss_ratios = torch.stack(loss_ratios) mean_ratio = loss_ratios.mean() r_i = loss_ratios / mean_ratio # shape (T,) # 目标梯度范数 avg_grad_norm = grad_norms.mean() target_grad_norms = avg_grad_norm * (r_i ** self.alpha) # 梯度损失,L1 grad_loss = torch.abs(grad_norms - target_grad_norms).mean() # 只更新 w gradnorm_optimizer.zero_grad() grad_loss.backward() gradnorm_optimizer.step() # 归一化:保持所有任务权重之和为任务数 with torch.no_grad(): self.w.data.mul_(self.num_tasks / self.w.data.sum()) # 防止出现负权重或极端接近 0 self.w.data.clamp_(min=1e-8) return grad_loss.item(), grad_norms.detach(), r_i.detach()这里有个细节值得单独说:第 1 步到第 4 步,我用autograd.grad而不是先对每个任务单独backward再取参数梯度,主要原因是避免污染param.grad。如果先对任务 A 的 loss 做 backward(retain_graph=True),再对任务 B 的 loss 做 backward,最终param.grad里存放的是 A 和 B 梯度的累加和,这样就没法分别拿到每个任务的梯度范数了。你当然可以在每次 backward 后立刻读取并清空 grad,但操作顺序容易出错,而且autograd.grad更干净、不会改动任何参数的.grad。
另一个细节是detach()。计算 r_i 时,如果直接用loss / initial_loss[i],这个比例会保留从 loss 到网络参数的完整计算图,后续构造grad_loss并 backward 时,可能把梯度误传到主网络参数上,造成二次求导或额外显存开销。r_i 只是一个观测指标,把它从计算图中摘掉是正确做法。
3.3 把 GradNorm 放进完整训练循环
下面给出一个完整的最小训练循环骨架。以两个任务为例,假设模型由共享主干和两个任务 head 组成。
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader class SharedEncoder(nn.Module): """示例共享主干,最后一层作为 GradNorm 观测层""" def __init__(self): super().__init__() self.features = nn.Sequential( nn.Linear(32, 64), nn.ReLU(), ) self.shared_layer = nn.Linear(64, 32) def forward(self, x): h = self.features(x) return self.shared_layer(h) class TaskHead(nn.Module): def __init__(self, in_dim, out_dim): super().__init__() self.net = nn.Linear(in_dim, out_dim) def forward(self, h): return self.net(h) model = SharedEncoder() head_a = TaskHead(32, 1) head_b = TaskHead(32, 1) # GradNorm 观测层参数:共享主干最后一层 shared_params = list(model.shared_layer.parameters()) gradnorm = GradNorm(num_tasks=2, init_w=1.0, alpha=0.5) # 主优化器 + GradNorm 优化器 main_params = list(model.parameters()) + list(head_a.parameters()) + list(head_b.parameters()) main_optimizer = optim.SGD(main_params, lr=0.01) gradnorm_optimizer = optim.SGD(gradnorm.parameters(), lr=0.001) initial_losses = None loader = DataLoader(...) # 按你的数据组织 for step, batch in enumerate(loader): x, y_a, y_b = batch h = model(x) loss_a = torch.nn.functional.mse_loss(head_a(h), y_a) loss_b = torch.nn.functional.mse_loss(head_b(h), y_b) losses = [loss_a, loss_b] # 第一个 batch 记录初始 loss if initial_losses is None: initial_losses = [l.detach() for l in losses] # 先做 GradNorm 权重更新 grad_loss_val, grad_norms, r_i = gradnorm.update( losses, shared_params, initial_losses, gradnorm_optimizer ) # 再用加权总 loss 更新主网络 total_loss = losses[0] * gradnorm.w[0] + losses[1] * gradnorm.w[1] main_optimizer.zero_grad() total_loss.backward() main_optimizer.step() if step % 100 == 0: print(f"step {step} | w: {gradnorm.w.detach().tolist()} " f"| r: {r_i.tolist()} | grad_norms: {grad_norms.tolist()}")运行起来后,你会看到 w_i 的变化趋势:如果一个任务长期落后,w_i 会慢慢上升;如果一个任务收敛过快,w_i 会下降。正常情况下 w_i 是缓慢漂移的,如果 w_i 出现大幅度锯齿震荡,说明 gradnorm 学习率太高或者 batch size 太小。
3.4 使用时的工程细节
有四个工程细节直接影响 GradNorm 是否work,很多复现翻车都翻在这里。
第一个是 initial_losses 的取值。建议用第一个 batch 的 loss 作为初始值,不要在 optimizer 已经 step 之后再记录,否则 r_i 的分母会偏小,导致后续权重调整失真。如果第一个 batch 恰好是个难样本,初始 loss 偏高没关系,它只作为相对基准。更稳的做法是取前几个 batch 的平均值,但要注意第一个 batch 后 network 参数已经更新了,严格说已经不是“初始”状态了。我实践中用第一个 batch 就够了,batch 大点更好。
第二个是 shared_params 的选择。很多人把整个 backbone 的所有参数都传进去,这是不对的。GradNorm 希望观察的是任务分支分叉之前的最后一层共享表示,因为这层参数承载了所有任务对共享信息的最终整合。如果传入整个 backbone,早期的低层特征梯度也会被计入,而这些低层梯度往往比较稳定但并不能很好反映任务间的竞争关系。选最后一层共享层,效果明显更稳定。
第三个是计算图的保留与释放。在gradnorm.update内部,torch.autograd.grad(retain_graph=True)会保留当前计算图,目的是让后续主网络 backward 可以继续用。这里有一个重要认知:gradnorm.update中grad_loss.backward()只沿着 grad_norms 到 w_i 的路径传播,不会重复计算网络参数的梯度,因为autograd.grad返回的梯度张量本身不携带网络计算图。但retain_graph=True确保原图没有被释放。所以顺序必须是 GradNorm 更新在前,主 backward 在后。如果反过来,主 backward 会把图释放掉,再调get_grad_norms就会报“Trying to backward through the graph a second time”的错误。
第四个是torch.autograd.grad的allow_unused=True参数。某些任务可能没有使用到共享层中的全部参数,返回的梯度会是 None。如果直接对所有梯度做torch.norm会报错,所以代码里做了if g is not None过滤。如果某个任务经常出现大量 None,说明这个任务没有真正走到共享层,要检查模型结构。
4. 调参实验与常见坑:我的实战记录
4.1 α 怎么选:0.5 不是万能钥匙
论文默认 α = 0.5,很多实现会直接复制这个值,但这只是一个安全起点的推荐值,不是恒定最优值。α 决定了落后任务得到多少梯度补偿,当你发现某个重要任务的 loss 长期偏高、权重已经接近下界时,可以调大 α,比如 1.0。如果发现任务间出现“为了平衡而平衡”,也就是强任务掉点明显、弱任务提升有限,说明 α 太大了,可以降到 0.25 或 0.1。
我自己的经验是:在任务难度差异确实很大的场景,α 大一点更管用;任务难度差不多时,α 取 0.3 左右更好。最简单的做法是跑一个小规模的消融实验,固定训练 20 epoch,对比 α = 0.25/0.5/1.0 三个值的任务平均准确率,选最优的。GradNorm 的权重更新本身比较便宜,消融实验代价可控。
4.2 gradnorm 学习率与 batch size 的相互作用
w_i 的学习率建议从主网络学习率的十分之一开始。我试过直接用相同学习率,结果 w_i 在前 100 步就剧烈震荡,loss 曲线出现明显毛刺。原因很好理解:w_i 是标量,它的梯度来自梯度范数差异,这个信号本身方差很大,学习率高了自然会过冲。
batch size 对 GradNorm 的影响经常被低估。梯度范数是基于一个 batch 的梯度统计量,batch 越小,估计越不稳定。如果你的训练 batch 只有 16 或者 32,建议对梯度范数做指数滑动平均(EMA)再计算目标值。我在一个目标检测多任务项目里,把 batch 从 16 提到 32,并且对梯度范数做 0.9 的 EMA 之后,w_i 曲线立刻就平稳了。
具体实现也不复杂,在 GradNorm 模块里维护两个 buffer 即可:
self.ema_grad_norms = None # 在 update 内部 if self.ema_grad_norms is None: self.ema_grad_norms = grad_norms.detach() else: self.ema_grad_norms = 0.9 * self.ema_grad_norms + 0.1 * grad_norms.detach()然后用 EMA 值代替原始 grad_norms 来计算 grad_loss。注意,EMA 值需要 detach,它不应该参与梯度传播。
4.3 三个最容易踩的坑
我把复现 GradNorm 的过程中最容易踩的三个大坑列出来,每一个都让我多烧过好几天。
第一个坑:权重变成负数。虽然理论上 w_i 是标量权重,正常更新不至于变负,但实际训练中如果某个 batch 的任务梯度范数异常大,gradnorm 更新后 w_i 真的可能越过零点。负权重会让 total_loss 出现负项,训练直接发散。最开始我在代码里没加clamp_min_,结果一个分割任务训到 2000 步 loss 变成负数。解决办法就是归一化后w.data.clamp_(min=1e-8),虽然不是论文原始实现,但工程上必须加。
第二个坑:把任务 head 参数也放进 shared_params。我在做一个二任务模型时,图方便直接把所有主干加两个 head 的参数都传给了 GradNorm,结果两个任务的梯度范数都被 head 里最后几层的参数量带偏了,权重更新毫无规律。GradNorm 的论文里明确是在共享层观测梯度,head 参数不属于共享层。
第三个坑:r_i 的计算没 detach 和初始 loss 顺序错误。如果不 detach,grad_loss 反向传播时可能沿着 r_i 的计算路径去对网络参数求二阶梯度,轻则显存暴涨,重则报错。初始 loss 如果用训练几个 batch 之后的 loss,r_i 的相对关系会被扭曲,尤其是初始 loss 已经很小的情况下,分母趋近 0,r_i 会变成巨大的数。这个问题特征很明显:训练早期 w_i 就开始剧烈变化,从几十到几百分之一,完全不受控。
4.4 GradNorm 和其他方法怎么配合
GradNorm 处理的是梯度大小失衡,但多任务优化还有一个经典问题是梯度方向冲突。两个任务在共享层上给出的梯度方向如果差异很大,即使范数差不多,更新也会互相干扰。这时候可以用 PCGrad 或者 GradVac 对梯度方向做修正,再把 GradNorm 作为第一层平衡,方向修正作为第二层,两者是互补关系。我实际做过组合实验,在语义分割和深度估计的多任务模型上,GradNorm + PCGrad 比单独用任何一个都稳定,分割任务的 mIoU 和深度估计的 RMSE 都有提升。
DWA(Dynamic Weight Average)是另一种动态权重方法,它只看 loss 下降速率,不看梯度范数。GradNorm 和 DWA 的区别在于,GradNorm 同时考虑了梯度的实际贡献大小,而 DWA 只从 loss 的数值变化出发。在梯度范数异常大的任务上,DWA 可能给出错误权重,GradNorm 则能捕捉到这一点,因为它直接观测梯度。如果你的项目里任务 loss 和梯度之间的关系比较直白,DWA 也可以;但更复杂的模型里我推荐 GradNorm。
5. 实验效果参考与适用场景分析
5.1 论文中的经典实验结果
GradNorm 论文的核心实验集中在两类任务上。一类是合成回归任务,用于证明方法机制;另一类是标准的多任务分类,比如 Multi-MNIST、Multi-Fashion + MNIST。在 Multi-MNIST 上,固定权重 1:1 训练时,两个任务准确率往往一高一低;GradNorm 通过调整权重,能让两个任务的准确率同时处于一个比较均衡的高位。论文里有个典型的数据是:固定权重会让较弱任务的准确率掉到 70% 以下,GradNorm 可以把弱任务拉回到 85% 以上,同时强任务只损失一两个点。
在真实应用里,这种“平衡提升”的意义尤其重要。自动驾驶里同时做车辆检测和车道线分割,如果检测任务霸占梯度,车道线分割质量就会明显下降,直接影响下游决策。GradNorm 不会让两个任务都变成最优,但能避免某个任务被完全牺牲。
5.2 更适合 GradNorm 的模型结构与任务场景
从我的项目经验看,GradNorm 有几个比较明确的适用前提。
首先是共享层足够深、足够宽。如果两个任务只在最后的全连接层共享一个很薄的表示,GradNorm 观测到的梯度范数噪声会很大,平衡效果有限。共享层越深,梯度统计量越稳定,GradNorm 越容易发挥效果。
其次是任务数量不要太多。GradNorm 的权重更新机制在 2 到 5 个任务时非常有效。任务数到 10 个以上,平均梯度范数的参考意义会被稀释,r_i 的计算也变得不稳定。当然理论上是支持任意任务数的,但工程上我一般配合任务分组使用。
第三是任务之间的 loss 尺度差异确实存在且无法通过 normalization 彻底解决。如果两个任务的 loss 数值经过精心调整后已经在同一量级,GradNorm 带来的收益有限。但现实里这种精心调整很难维持,尤其数据分布变化后 loss 尺度会漂移,GradNorm 的在线调节优势就体现出来了。
下面是一个简单的适用性对照表,方便快速判断:
| 场景 | 是否适合 GradNorm | 原因 |
|---|---|---|
| 2-5 个任务共享主干 | 非常推荐 | 梯度统计稳定,平衡效果明显 |
| 任务间 loss 尺度差异大 | 推荐 | 在线自动调节,节省调参时间 |
| 任务数超过 10 | 谨慎使用 | 平均梯度范数参考意义弱化,需分组处理 |
| 共享层非常浅(1-2 层) | 不建议 | 梯度噪声大,观测不稳定 |
| 任务之间几乎没有共享表示 | 不建议 | 平衡共享层梯度没有意义 |
| 训练 batch 很小且无法增大 | 需要改造 | 建议加 EMA 或对梯度范数做平滑 |
5.3 与其它平衡方法的对比心得
我实际用过三种常用的动态权重方法:不确定性加权、DWA、GradNorm。不确定性加权需要定义每个任务的似然模型,回归任务还好,分类任务会繁琐一些;而且噪声参数的梯度很不稳定,经常需要专门限制学习率。DWA 实现最简单,用 loss 比例就能算权重,但它只关注 loss 变化,忽略了梯度范数这个更底层的信号,在复杂任务上容易失准。GradNorm 的工程成本介于两者之间,实现也就几十行代码,但它给的是一个真正和优化器交互的信号,我觉得这是它最大的价值。
如果你现在用的是固定权重,并且已经感受到调参成本,我的建议是先换成 GradNorm,而不是急着上更复杂的方法。GradNorm 不需要额外维护概率模型,也不会引入大量新超参数,最重要的一点是它的行为比较可解释:w_i 高说明任务落后,w_i 低说明任务已经稳定,你能从训练日志里直观地读懂模型在干什么。
我自己在实际项目里用 GradNorm 最深的体会是:它不是让多任务训练从“很难”变成“不难”,而是让训练过程变得可观测、可解释。以前调权重靠感觉,现在我能看着 w_i 曲线判断某个任务是缺梯度还是已经饱和。如果你是第一次接多任务项目,可以直接按这篇文章的思路把 GradNorm 先跑起来,然后把 α 和 gradnorm 学习率作为最优先调的两个超参数。初始权重的选择不用太纠结,GradNorm 会在训练早期自动修正它——这也是它和手动调权最本质的区别。