Uncertainty Weight Loss
这是对论文Multi Task Learning with Homoscedastic Uncertainty的一个学习。
似然定义
设置\(\mathbf{f}^{\mathbf{W}}(\mathbf{x})\)表示为输入\(x\)通过权重\(\boldsymbol{W}\)的输出特征向量。本文将回归问题和分类问题分开考虑他们的概率表示：
回归问题：\(p(\mathbf{y}|\mathbf{f}^{\mathbf{W}}(\mathbf{x}))=\mathcal{N}(\mathbf{f}^{\mathbf{W}}(\mathbf{x}),\sigma^2)\),即另模型输出作为均值，额外用一个参数表示方差，建模为高斯分布。
分类问题：\(p(\mathbf{y}|\mathbf{f}^{\mathbf{W}}(\mathbf{x}))=\text{Softmax}(\mathbf{f}^{\mathbf{W}}(\mathbf{x}))\),分类时经过\(\text{Softmax}\)即为概率分布。
接下来对于多任务结合的概率表示，简单讲就是概率乘积： \[ \begin{aligned} p\left(\mathbf{y}_{1}, \ldots, \mathbf{y}_{K} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right)=p\left(\mathbf{y}_{1} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \ldots p\left(\mathbf{y}_{K} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \end{aligned} \]
最大似然
为了最大化似然，先构造对数似然：
回归问题： \(\log p\left(\mathbf{y} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \propto-\frac{1}{2 \sigma^{2}}\left\|\mathbf{y}-\mathbf{f}^{\mathbf{W}}(\mathbf{x})\right\|^{2}-\log \sigma\)，这里\(\mathbf{y}\)就是标准数据，\(\mathbf{f}^{\mathbf{W}}(\mathbf{x})\)为均值，后面的\(-\log \sigma\)其实是他在对数化的时候把系数给省略了。
分类问题：\(p\left(\mathbf{y} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma\right)=\operatorname{Softmax}\left(\frac{1}{\sigma^{2}} \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right)\)，这里为了添加一个可学习的系数，他构造\(sigma\)作为输出特征向量的温度系数，用来调整\(\text{Softmax}\)后的熵，比如特征向量被缩放的越厉害，那么输出分布会变得均匀，那么不确定性程度就加深了。NOTE： 因为我们的分类标签是带类别的，因此对应类别的对数似然公式如下(其中\(c^{\prime}\)表示不为\(c\)的其他类别)：
\[ \begin{aligned} \log p\left(\mathbf{y}=c \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma\right) &=\frac{1}{\sigma^{2}} f_{c}^{\mathbf{W}}(\mathbf{x}) -\log \sum_{c^{\prime}} \exp \left(\frac{1}{\sigma^{2}} f_{c^{\prime}}^{\mathbf{W}}(\mathbf{x})\right) \end{aligned} \]
接下来就是多个任务联合的对数似然并进行最大化后验概率，我们最小化负的对数似然：
\[ \begin{array}{l} =-\log p\left(\mathbf{y}_{1}, \mathbf{y}_{2}=c \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \\ =-\log \mathcal{N}\left(\mathbf{y}_{1} ; \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma_{1}^{2}\right) \cdot \operatorname{Softmax}\left(\mathbf{y}_{2}=c ; \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma_{2}\right) \\ =\frac{1}{2 \sigma_{1}^{2}}\left\|\mathbf{y}_{1}-\mathbf{f}^{\mathbf{W}}(\mathbf{x})\right\|^{2}+\log \sigma_{1}-\log p\left(\mathbf{y}_{2}=c \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma_{2}\right)\\ =\frac{1}{2 \sigma_{1}^{2}} \mathcal{L}_{1}(\mathbf{W})+\frac{1}{\sigma_{2}^{2}} \mathcal{L}_{2}(\mathbf{W})+\log \sigma_{1} +\log \frac{\sum_{c^{\prime}} \exp \left(\frac{1}{\sigma_{2}^{2}} f_{c^{\prime}}^{\mathrm{W}}(\mathrm{x})\right)}{\left(\sum_{c^{\prime}} \exp \left(f_{c^{\prime}}^{\mathrm{W}}(\mathbf{x})\right)\right)^{\frac{1}{\sigma_{2}^{2}}}}\\ \approx \frac{1}{2 \sigma_{1}^{2}}\mathcal{L}_{1}(\mathbf{W})+\frac{1}{\sigma_{2}^{2}} \mathcal{L}_{2}(\mathbf{W})+\log \sigma_{1}+\log \sigma_{2} \end{array} \]
以上的\(\mathcal{L}_{1}(\mathbf{W})=\left\|\mathbf{y}_{1}-\mathbf{f}^{\mathbf{W}}(\mathbf{x})\right\|^{2}\)即为回归损失，\(\mathcal{L}_{2}(\mathbf{W})=-\log \operatorname{Softmax}\left(\mathbf{y}_{2}, \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right)\)为分类误差。
上面公式转换的时候我觉得他为了简化而简化一下，他假设\(\sigma_2\rightarrow1\),那么转换公式中\(\log \frac{\sum_{c^{\prime}} \exp \left(\frac{1}{\sigma_{2}^{2}} f_{c^{\prime}}^{\mathrm{W}}(\mathrm{x})\right)}{\left(\sum_{c^{\prime}} \exp \left(f_{c^{\prime}}^{\mathrm{W}}(\mathbf{x})\right)\right)^{\frac{1}{\sigma_{2}^{2}}}}\)上下两项将消除，\(\log1\)就消除了。但是消除之后不能保证\(\sigma_2\rightarrow1\)了，他又添加了\(\log \sigma_{2}\)作为约束，保证他将趋向于1.
实施
那么根据作者的写法，不论回归还是分类损失其实都是用相同的计算公式，然后最后组合即可： \[ loss_{w}=\frac{1}{2\sigma^2}loss+\log \sigma \]
当然为了避免出现\(\log \sigma\)数值溢出的问题，采用\(s:=\log \sigma^2\)的方法来代替：
\[ loss_{w}=\frac{1}{2}\cdot \exp(-s) \cdot loss+\frac{1}{2}s \]
代码与实验结果
代码如下，tf2.x运行：
import tensorflow as tf |
learnable True Reconstruct True |
