【论文阅读】Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics

1.背景

目前，很多应用中都因为用了多任务取得了不错的效果，那么如何将这些任务的loss有机组合在一起？

一种简答粗暴方法就是手动调节多任务之间的loss的相对权重，然后加和，如下：

这种方式把权重作为超参调试的方式，往往耗时耗力。

本文参考[1]就如何动态设置多任务（分类与回归）之间的loss的相对权重做简单介绍。

2.解决方案

在贝叶斯建模中，认为模型有一些不确定性(随机变量)，可能是先验问题导致的，也有可能是后验问题导致的，大体可将其分为两种，如下：

• 认知不确定性：数据不足导致的模型学习不足的不确定性，比如：小学生去参加高考，很多知识都没学过，自然考不好。
• 偶然不确定性：噪声导致的不确定性

偶然不确定性又分为：

• 数据依赖不确定性（异方差）：输入数据的不确定性，导致输出不确定，比如：小明学习使用的教材部分有问题，导致小明考试考不到高分。
• 任务依赖不确定性（同方差）：不同任务自身的学习能力的不同，导致学习结果不确定性。比如：小明为了提高成绩，一方面努力学习知识概念，另一方面猛做练习题，他们都有各自的优势，但是也都有一定的局限性。

其中，同方差 指的是假定数据输入一定的情况下，真实的分布与任务的输出之间有一个恒定的方差。

解决方案为什么考虑的用贝叶斯NN？贝叶斯建模中的不确定性能够表示不同任务之间的难易程度，能够很好的为任务的输出f，真实分布y以及方差δ之间建模。

动态loss权重就是考虑在同方差条件下，接下来我们如何建模展开。

3 建模

假设输入为$X$，参数矩阵为$W$，输出为$f^W(x)$，真实分布为$y$。

3.1 回归任务

对于回归任务，在给定其输出$f^W(x)$的情况下$y$的概率为高斯似然：

(对于高斯似然，$\sigma$为模型的观测噪声参数，表示输出数据中的噪声量)

将$f^W(x)$和$y$代入上述高斯函数中，可得：
$p(y|f^W(x))=\frac{1}{\sqrt{2\pi }\sigma }{e}^{-\frac{(y-f^W(x){)}^2}{2{\sigma }^2}}$

两边求对数，可得其对数似然函数，如下：

$logp(y|f^W(x))\propto -\frac{(y-f^W(x){)}^2}{2{\sigma }^2}-log\sigma$

3.2 分类任务

对于分类任务，为了建立起与$\sigma$的关系，其概率描述为softmax的一种更为普遍的表示形式Boltzmann分布，也叫吉布斯分布：

系数$\sigma$可以是设定的，也可以是通过学习得到的，决定离散分布的平坦程度。该值和分布的不确定性(熵)有关。

同理，将$f^W(x)$和$y$代入上述softmax函数中，可得对数似然函数，如下：

其中，c为某一类别

3.3 多任务

对于回归与分类任务混合的多任务似然，假定$y_1$、…、$y_k$分别为回归任务和分类任务的真实输出，其似然为：

那么，假设多任务的loss记作$L(W,{\sigma }_1,{\sigma }_2)$，那么则有：
= $L(W,{\sigma }_1,{\sigma }_2)$
= $-logN(y_1;f^W(x),{\sigma }_1^2)\cdot softmax(y_2=c;f^W(x),{\sigma }_2)$
= $\frac{1}{2{\sigma }_1^2}{\Vert y_1-f^W(x)\Vert }^2+log{\sigma }_1-logp(y_2=c|f^W(x),{\sigma }_2)$
= $\frac{1}{2{\sigma }_1^2}{\Vert y_1-f^W(x)\Vert }^2+log{\sigma }_1-\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)+log{\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}$
= $\frac{1}{2{\sigma }_1^2}{\Vert y_1-f^W(x)\Vert }^2+log{\sigma }_1+\frac{1}{{\sigma }_2^2}log{\sum }_{c{}^{\prime }}{e}^{f_{c{}^{\prime }}^W(x)}-\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)+log{\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}-\frac{1}{{\sigma }_2^2}log{\sum }_{c{}^{\prime }}{e}^{f_{c{}^{\prime }}^W(x)}$
=$\frac{1}{2{\sigma }_1^2}{\Vert y_1-f^W(x)\Vert }^2+log{\sigma }_1-\frac{1}{{\sigma }_2^2}logsoftmax(y_2,f^W(x))+log{\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}-\frac{1}{{\sigma }_2^2}log{\sum }_{c{}^{\prime }}{e}^{f_{c{}^{\prime }}^W(x)}$
= $\frac{1}{2{\sigma }_1^2}{\Vert y_1-f^W(x)\Vert }^2+log{\sigma }_1-\frac{1}{{\sigma }_2^2}logsoftmax(y_2,f^W(x))+log\frac{{\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}}{({\sum }_{c{}^{\prime }}{e}^{f_{c{}^{\prime }}^W(x)}{)}^{\frac{1}{{\sigma }_2^2}}}$

由于当${\sigma }_2$->1时，有$\frac{1}{{\sigma }_2}{\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}$≈$({\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}{)}^{\frac{1}{{\sigma }_2}}$，
所以上式最后一个$log\frac{{\sum }_{c{}^{\prime }}{e}^{\frac{1}{{\sigma }_2^2}{f}_{c{}^{\prime }}^W(x)}}{({\sum }_{c{}^{\prime }}{e}^{f_{c{}^{\prime }}^W(x)}{)}^{\frac{1}{{\sigma }_2^2}}}\approx log{\sigma }_2$，
则有： $L(W,{\sigma }_1,{\sigma }_2)$ ≈ $\frac{1}{2{\sigma }_1^2}{\Vert y_1-f^W(x)\Vert }^2+log{\sigma }_1-\frac{1}{{\sigma }_2^2}logsoftmax(y_2,f^W(x))+log{\sigma }_2$

令$L_1(W)={\Vert y_1-f^W(x)\Vert }^2$为回归问题的loss，$L_2(W)=-log(softmax(y_2,f^W(x)))$为分类问题的loss，则有多任务loss为：

$L(W,{\sigma }_1,{\sigma }_2)$ ≈ $\frac{1}{2{\sigma }_1^2}{L}_1(W)+log{\sigma }_1+\frac{1}{{\sigma }_2^2}{L}_2(W)+log{\sigma }_2$

4.代码实现

实际实现中，为了防止分母出现0，提升数值的稳定性以及简化计算，回归loss的系数去掉了$\frac{1}{2}$，如下：

$L(W,{\sigma }_1,{\sigma }_2)$ ≈ $\frac{1}{{\sigma }_1^2}{L}_1(W)+\frac{1}{{\sigma }_2^2}{L}_2(W)+2log{\sigma }_1+2log{\sigma }_2$

并且，一般令$s=log{\sigma }^2$，上式可以化简为：

$L(W,{\sigma }_1,{\sigma }_2)$ ≈ $e^{-s1}{L}_1(W)+e^{-s2}{L}_2(W)+s1+s2$

pytorch代码实现如下：

class DynamicWeightedLoss(nn.Module):
    def __init__(self, num=2):
        super(DynamicWeightedLoss, self).__init__()
        params = torch.ones(num, requires_grad=True)
        self.params = torch.nn.Parameter(params)

    def forward(self, *x):
        loss_sum = 0
        for i, loss in enumerate(x):
            loss_sum += torch.exp(-self.params[i]) * loss 
            + self.params[i]
        return loss_sum

论文[2]对论文[1]做了loss的正则项做了优化，它的loss如下：
$L(W,{\sigma }_1,{\sigma }_2)=\frac{1}{2{\sigma }_1^2}{L}_1(W)+\frac{1}{2{\sigma }_2^2}{L}_2(W)+log({\sigma }_1^2+1)+log({\sigma }_2^2+1)$，
其pytroch代码实现如下：

class DynamicWeightedLoss(nn.Module):
    def __init__(self, num=2):
        super(DynamicWeightedLoss, self).__init__()
        params = torch.ones(num, requires_grad=True)
        self.params = torch.nn.Parameter(params)

    def forward(self, *x):
        loss_sum = 0
        for i, loss in enumerate(x):
            loss_sum += 0.5 / (self.params[i] ** 2) * loss 
            + torch.log(1 + self.params[i] ** 2)
        return loss_sum

5.进一步了解

如果还想对贝叶斯不确定性进一步了解，可以阅读reference部分的论文[3]，如果还想对动态loss以及多任务进一步了解可以阅读综述性文章[4,5]以及一些经典的多任务论文[6,7,8]。

reference:

[1].Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics
[2].Auxiliary Tasks in Multi-task Learning
[3].What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
[4].An Overview of Multi-Task Learning in Deep Neural Networks.pdf
[5].https://zhuanlan.zhihu.com/p/269492239
[6].Multi-Task Learning as Multi-Objective Optimization
[7].MMOE
[8].SNR