Auxiliary Loss损失函数实战：解决多任务学习中的梯度冲突问题

1次阅读

没有评论

共计 2195 个字符，预计需要花费 6 分钟才能阅读完成。

在计算机视觉的典型多任务场景（如语义分割 + 深度估计）中，不同任务对共享特征的需求存在天然矛盾：

语义分割需要高层语义特征（如物体类别）
深度估计依赖几何特征（如物体边缘连续性）

当使用共享骨干网络时，反向传播的梯度会出现两种典型冲突：

幅度冲突：深度估计任务的梯度范数可能比语义分割大 100 倍
方向冲突：某些卷积核的梯度更新方向完全相反（余弦相似度 <-0.8）

与传统加权损失函数 $L = \sum w_iL_i$ 相比，Auxiliary Loss 的核心改进在于：

$$
L_{total} = L_{main} + \alpha(t)\cdot L_{aux}
$$

其中动态权重系数 $\alpha(t)$ 设计为：

$$
\alpha(t) = \eta\cdot\frac{\nabla L_{main}}{|\nabla L_{aux}|_2 + \epsilon}
$$

关键差异点：

常规方法：静态权重，需网格搜索调参
Auxiliary Loss：根据梯度比例动态调整，使辅助任务既帮助特征学习又不干扰主任务

class DynamicAuxLoss(nn.Module):
    def __init__(self, main_loss, aux_loss, init_alpha=0.1):
        super().__init__()
        self.main_loss = main_loss  # 主任务损失函数
        self.aux_loss = aux_loss    # 辅助任务损失函数
        self.alpha = nn.Parameter(torch.tensor(init_alpha))
        self.eps = 1e-8

    def forward(self, main_pred, aux_pred, main_gt, aux_gt):
        # 计算基础损失
        L_main = self.main_loss(main_pred, main_gt)
        L_aux = self.aux_loss(aux_pred, aux_gt)

        # 动态调整系数
        with torch.enable_grad():
            g_main = torch.autograd.grad(L_main, main_pred, retain_graph=True)[0]
            g_aux = torch.autograd.grad(L_aux, aux_pred, retain_graph=True)[0]
            alpha = self.alpha * (g_main.norm() / (g_aux.norm() + self.eps)).detach()

        return L_main + alpha * L_aux

# 在训练循环中添加梯度规范化
optimizer.zero_grad()
total_loss.backward()

# 对各任务梯度进行归一化
for param in model.shared_parameters():
    if param.grad is not None:
        # 按任务不确定性归一化（示例取 L2 范数）task_grad_norm = param.grad.norm(p=2)
        param.grad.div_(task_grad_norm + 1e-6)

optimizer.step()

分类 + 回归任务：初始 α∈[0.01, 0.1]
双分类任务：初始 α∈[0.1, 0.5]
使用 nn.Parameter 而非直接变量，确保可学习性

建议在验证集上观察：

主任务与辅助任务的 loss 比值曲线（应收敛到稳定区间）
动态权重 α 的变化轨迹（突然跳变可能预示任务冲突）
共享层梯度余弦相似度（理想值 >-0.3）

# 构造双任务数据集
class MNISTDoubleTask(Dataset):
    def __init__(self, train=True):
        self.mnist = datasets.MNIST(..., train=train)

    def __getitem__(self, idx):
        img, digit = self.mnist[idx]
        parity = digit % 2  # 新增奇偶分类任务
        return img, (digit, parity)

# 模型定义（共享卷积层）model = nn.Sequential(nn.Conv2d(1, 32, 3),
    nn.ReLU(),
    nn.MaxPool2d(2),
    # 任务特定头
    nn.ModuleDict({'digit': nn.Linear(32*13*13, 10),
        'parity': nn.Linear(32*13*13, 2)
    })
)

# 损失配置
loss_fn = DynamicAuxLoss(main_loss=nn.CrossEntropyLoss(),  # 数字识别
    aux_loss=nn.CrossEntropyLoss()    # 奇偶分类)