1. 背景

在多任务学习中,我们有多个任务(例如任务 A 和任务 B),每个任务有自己的损失函数  Ltk\ {L}^{tk} Ltk 。为了平衡这些任务,我们需要为每个任务分配一个权重 wtk\ {w}^{tk} wtk 。使得每个任务对总损失的贡献相对均衡。


2. 梯度归一化的核心思想

梯度归一化的目标是让每个任务的梯度大小(即损失函数对模型参数的导数)相对均衡。如果某个任务的梯度太大,说明它对模型的影响过大,我们需要降低它的权重;反之,如果梯度太小,则需要增加它的权重。


3. 简单例子

假设我们有两个任务:

  • 任务 A:预测响应时间,损失函数为  LA。\ {L}^{A} 。 LA
  • 任务 B:预测吞吐量,损失函数为 LB。\ {L}^{B} 。 LB

步骤 1:计算每个任务的梯度

  • 计算任务 A 的梯度  (∇LA)\ ( \nabla \mathcal{L}^A ) (LA)
  • 计算任务 B 的梯度  (∇LB)\ ( \nabla \mathcal{L}^B ) (LB)

步骤 2:计算每个任务的梯度范数

  • 计算任务 A 的梯度范数  (∣∣∇LA∥)\ (|| \nabla \mathcal{L}^A \|) (∣∣∇LA)(即梯度的长度)。
  • 计算任务 B 的梯度范数  (∣∣∇LB∥)\ (|| \nabla \mathcal{L}^B \|) (∣∣∇LB)

步骤 3:计算初始权重

假设初始权重为 (wA=1)( w^A = 1 )(wA=1)(wB=1)( w^B = 1 )(wB=1)

步骤 4:调整权重

根据梯度范数调整权重:

  • 如果 (∣∣∇LA∥)\ (|| \nabla \mathcal{L}^A \|) (∣∣∇LA)远大于 (∣∣∇LB∥)\ (|| \nabla \mathcal{L}^B \|) (∣∣∇LB),说明任务 A 的梯度对模型的影响更大,我们需要降低任务 A 的权重 wAw^AwA
  • 反之,如果 (∣∣∇LB∥)\ (|| \nabla \mathcal{L}^B \|) (∣∣∇LB)更大,则降低任务 B 的权重 wBw^BwB

具体公式可以简化为:
[wtk=1∥∇Ltk∥][ w^{tk} = \frac{1}{\| \nabla \mathcal{L}^{tk} \|} ][wtk=∥∇Ltk1]
这样,梯度范数越大的任务,权重越小。

步骤 5:归一化权重

为了让权重总和为 1,可以对权重进行归一化:
[wA=wAwA+wB,wB=wBwA+wB] [ w^A = \frac{w^A}{w^A + w^B}, \quad w^B = \frac{w^B}{w^A + w^B} ][wA=wA+wBwA,wB=wA+wBwB]


4. 具体计算示例

假设:

  • 任务 A 的梯度范数 (∥∇LA∥=5)( \| \nabla \mathcal{L}^A \| = 5)(∥∇LA=5)
  • 任务 B 的梯度范数(∥∇LB∥=2)( \| \nabla \mathcal{L}^B \| = 2)(∥∇LB=2)

计算初始权重:

[wA=15=0.2,wB=12=0.5][ w^A = \frac{1}{5} = 0.2, \quad w^B = \frac{1}{2} = 0.5 ][wA=51=0.2,wB=21=0.5]

归一化权重:

[wA=0.20.2+0.5=0.20.7≈0.29][ w^A = \frac{0.2}{0.2 + 0.5} = \frac{0.2}{0.7} \approx 0.29 ][wA=0.2+0.50.2=0.70.20.29]
[wB=0.50.2+0.5=0.50.7≈0.71][ w^B = \frac{0.5}{0.2 + 0.5} = \frac{0.5}{0.7} \approx 0.71 ][wB=0.2+0.50.5=0.70.50.71]

结果:

  • 任务 A 的权重为 0.29。
  • 任务 B 的权重为 0.71。

5. 总结

梯度归一化的核心思想是通过调整每个任务的权重,使得它们的梯度大小相对均衡。具体步骤包括:

  1. 计算每个任务的梯度。
  2. 计算每个任务的梯度范数。
  3. 根据梯度范数调整权重。
  4. 归一化权重,使其总和为 1。

通过这种方法,模型可以更好地平衡多个任务,避免某个任务主导训练过程。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐