寻找多目标优化问题的帕累托最优解.

paper：Multi-Task Learning as Multi-Objective Optimization

多目标优化是指同时优化多个相关任务的目标，多任务学习是一个典型的多目标优化问题，其总目标函数是每个任务的目标函数的加权求和式： $\mathcal{L}_{total} = \sum_{i}^{n} w_i\mathcal{L}_i$ 。为使得每个任务在训练时都获得有益的提升，需要合理的设置任务权重 $w_i$ ，使得每一次更新时目标损失函数 $\mathcal{L}_1,\mathcal{L}_2,\cdots, \mathcal{L}_n$ 都下降或保持不变。

对于参数 $\theta^{*}$ ，若该参数任意变化都会导致某个目标的损失函数 $\mathcal{L}_i$ 上升，则称 $\theta^{*}$ 为帕累托最优解(Pareto Optimality)。帕累托最优意味着不能通过牺牲某个任务来换取其他任务的性能提升。

1. 建模多目标优化问题

本文主要讨论基于梯度下降法的多目标优化问题。记参数 $\theta$ 在任务 $i$ 上的梯度为 $g_i = \nabla_{\theta} \mathcal{L}_i$ ，则寻找帕累托最优的过程等价于构造参数变化量 $\Delta \theta$ 满足：

$\quad \langle g_i, \Delta \theta \rangle \leq 0, \quad \forall i$

注意到上述方程组存在平凡解 $\Delta \theta=0$ ，我们主要关心可行域中是否存在非零解。若能求得非零解，则将其作为参数更新方向；否则有可能已经达到了帕累托最优（必要不充分条件），此时的状态称为帕累托稳定点 (Pareto Stationary)。

上述问题等价于寻找一个向量 $u$ ，使得对所有任务 $i$ 的梯度 $g_i$ 都满足 $\langle g_i,u \rangle \geq 0$ ，此时可取 $\Delta \theta = -\eta u$ 作为参数更新量。

将该问题转化为最优化问题：

$\forall i, \langle g_i,u \rangle \geq 0 \quad \Leftrightarrow \quad \mathop{\min}_{i} \langle g_i,u \rangle \geq 0$

问题可进一步转化为最大化内积 $\langle g_i,u \rangle$ 的最小值：

$\mathop{\max}_{u} \mathop{\min}_{i} \langle g_i,u \rangle$

注意到上述优化目标可能会趋于无穷大，因此为了结果的稳定性，为 $u$ 的模长增加正则项：

$\mathop{\max}_{u} \mathop{\min}_{i} \langle g_i,u \rangle - \frac{1}{2}||u||^2$

注意到 $u=0$ 时上述目标取值也为 $0$ ，因此最优解 $u^{*}$ 满足：

$\mathop{\min}_{i} \langle g_i,u^* \rangle - \frac{1}{2}||u^*||^2 \geq 0 \quad \Leftrightarrow \quad \mathop{\min}_{i} \langle g_i,u^* \rangle \geq \frac{1}{2}||u^*||^2 \geq 0$

即加入正则项之后的最优解 $u^{*}$ 也必然是满足原最优化问题 $\mathop{\min}_{i} \langle g_i,u \rangle \geq 0$ 的解。

2. 求解多目标优化问题

定义 $\Bbb{P}^n$ 为所有 $n$ 元离散分布的集合：

$\Bbb{P}^n = \{ (w_1,w_2,\cdots,w_n) | w_1,w_2,\cdots,w_n\geq 0,\sum_{i} w_i = 1 \}$

则原优化目标等价于：

$\mathop{\min}_{i} \langle g_i,u \rangle = \mathop{\min}_{w \in \Bbb{P}^n} \langle \sum_i w_ig_i,u \rangle = \mathop{\min}_{w \in \Bbb{P}^n} \langle \tilde{g}_i(w),u \rangle$

因此加入正则项之后的目标也等价于：

$\mathop{\max}_{u} \mathop{\min}_{w \in \Bbb{P}^n} \langle \tilde{g}_i(w),u \rangle - \frac{1}{2}||u||^2$

该目标函数是关于 $u$ 的凹函数，关于 $w$ 的凸函数，且 $u,w$ 的可行域都是凸集，根据Minimax定理可以将上述问题转化为对偶问题：

$\mathop{\min}_{w \in \Bbb{P}^n} \mathop{\max}_{u} \langle \tilde{g}_i(w),u \rangle - \frac{1}{2}||u||^2$

注意到 $\max$ 部分是关于 $u$ 的无约束的二次函数最大值问题，其最大值在 $u=\tilde{g}_i(w)$ 处取，则目标函数进一步写作：

$\mathop{\min}_{w \in \Bbb{P}^n} \frac{1}{2}||\tilde{g}_i(w)||^2 = \mathop{\min}_{w \in \Bbb{P}^n} \frac{1}{2}||\sum_i w_ig_i||^2$

上式等价于求各任务损失梯度 $g_1,g_2,\cdots g_n$ 的所有凸组合中模长最小的情况，约束条件为 $\sum_i w_i =1$ 。

⚪ 无约束的梯度下降

$\mathop{\min}_{w \in \Bbb{P}^n} \frac{1}{2}||\sum_i w_ig_i||^2$

该目标可以通过去约束的方式直接用梯度下降求解。常用方法是设置可学习参数 $\beta_1,\beta_2,\cdots \beta_n \in \Bbb{R}$ ，使得：

$w_i = \frac{e^{\beta_i}}{\sum_i e^{\beta_i}}$

此时隐式地包含约束 $\sum_i w_i =1$ ；且目标转化为 $\beta$ 的无约束优化问题：

$\mathop{\min}_{\beta} \frac{1}{2(\sum_i e^{\beta_i})^2}||\sum_i e^{\beta_i}g_i||^2$

⚪ 带约束的梯度下降

$\mathop{\min}_{w \in \Bbb{P}^n} \frac{1}{2}||\tilde{g}_i(w)||^2$

原问题也可以通过Frank-Wolfe算法求解。Frank-Wolfe算法可以看作一种带约束的梯度下降算法，适合于参数的可行域为凸集的情形。该算法的求解在当前 $w$ 与位置 $\tau$ 为 $1$ 的one-hot向量 $e_{\tau}$ 之间进行插值搜索，找出最优者作为迭代结果。

当 $n=2$ 时，问题等价于求两个向量 $g_1$ （图中为 $\theta$ ）和 $g_2$ （图中为 $\overline{\theta}$ ）的模长最小凸组合：

当 $n>2$ 时，Frank-Wolfe算法将问题转化为多个 $n=2$ 的情形进行迭代，迭代过程为：

$\begin{aligned} \tau &= \mathop{\arg \min}_i \langle g_i,\tilde{g}(w^{(k)}) \rangle \\ \gamma &= \mathop{\arg \min}_{\gamma} \tilde{g}((1-\gamma)w^{(k)}+\gamma e_{\tau}) \\ &= \mathop{\arg \min}_{\gamma} ||(1-\gamma)\tilde{g}(w^{(k)})+\gamma g_{\tau}||^2 \\ w^{(k+1)} &= (1-\gamma)w^{(k)}+\gamma e_{\tau} \end{aligned}$

其中 $\gamma$ 的求解即为 $n=2$ 的特例。

3. 优化求解过程

前两节讨论了寻找帕累托稳定点的参数更新方向的方法，实现过程为在每一步的训练中，先通过另外的多步迭代来确定每个目标的权重，然后再更新模型参数。实际计算时计算量比较大，因此通过一些额外的技巧降低计算量。

⚪ 梯度内积

在优化过程中需要多次计算梯度的内积 $\langle g_i,\tilde{g}(w^{(k)}) \rangle$ ，由于梯度 $g_i$ 的维度与模型参数量相同且通常比较大，因此内积运算的计算量很大。考虑展开式：

$\langle g_i,\tilde{g}(w^{(k)}) \rangle = \langle g_i,\sum_j w_jg_j \rangle =\sum_j w_j \langle g_i,g_j \rangle$

在实现时 $\langle g_i,g_j \rangle$ 只需要计算一次并存储下来，可以减少大维度向量内积的重复计算。

⚪ 共享编码

若假设多个目标任务共享同一个特征提取网络，则还可以进一步近似地简化算法。具体地，假设批量为 $b$ 的样本中第 $j$ 个样本的特征编码输出为 $h_j$ ，则第 $i$ 个目标损失的梯度计算为：

$\begin{aligned} g_i & = \nabla_{\theta} \mathcal{L}_i = \sum_j (\nabla_{h_j} \mathcal{L}_i ) (\nabla_{\theta} h_j) \\ &=(\nabla_{h_1} \mathcal{L}_i,\cdots,\nabla_{h_b} \mathcal{L}_i) \begin{pmatrix} \nabla_{\theta} h_1 \\ \vdots \\ \nabla_{\theta} h_b \end{pmatrix} \end{aligned}$

若记 $H=(h_1,\cdots,h_b)$ ，则梯度 $g_i=(\nabla_{H} \mathcal{L}_i ) (\nabla_{\theta} H)$ 。优化目标为：

$||\sum_i w_ig_i||^2 = ||\sum_i w_i(\nabla_{H} \mathcal{L}_i ) (\nabla_{\theta} H)||^2 \leq ||\sum_i w_i\nabla_{H} \mathcal{L}_i ||^2 || \nabla_{\theta} H||^2$

因此如果不最小化 $||\sum_i w_ig_i||^2$ （需要求所有参数的梯度），而是最小化 $||\sum_i w_i\nabla_{H} \mathcal{L}_i||^2$ （只需要求编码向量的梯度），则计算量会明显减少。后者是前者的一个上界。

该上界也具有局限性，一般只适用于每一个样本都有多种标注信息的多目标优化问题，不适用于不同目标任务的训练数据无交集的场景。因为不同任务的数据无交集表示 $\nabla_{H} \mathcal{L}_i$ 是相互正交的，该上界过于宽松，没有体现出任务之间的相关性。

4. 主次型多目标优化

在前面几节的讨论中，多目标优化的目的是尽可能在所有目标任务上都取得较好的表现，即平等地处理每一个目标。在实际需求中，有时我们希望模型主要在一个或几个主目标上表现较好，并额外地增加一些辅助目标，通过学习辅助目标来提升在主目标上的表现。此时多目标优化也称为主次型多目标优化。

若记 $\mathcal{L}_0$ 为主目标的损失函数， $\mathcal{L}_1，\cdots \mathcal{L}_n$ 为 $n$ 个辅助目标的损失函数。则主目标仍然是寻找一个向量 $u$ ，使得对主目标的梯度 $g_0$ 满足 $\langle g_0,u \rangle \geq 0$ ，并采用梯度下降法更新参数 $\Delta \theta = -\eta u$ 。根据前面的讨论，优化目标为：

$\mathop{\max}_{u} \langle g_0,u \rangle - \frac{1}{2}||u||^2$

对于辅助目标 $\mathcal{L}_1，\cdots \mathcal{L}_n$ ，我们希望更新方向 $u$ 不要使得其中的任何一个目标损失增大即可(不一定最小化损失)，因此引入约束条件：

$\begin{aligned} \mathop{\max}_{u} & \langle g_0,u \rangle - \frac{1}{2}||u||^2 \\ \text{s.t.} & \langle g_1,u \rangle \geq 0,\cdots \langle g_n,u \rangle \geq 0 \end{aligned}$

建立拉格朗日方程，将上述问题转化成min-max问题：

$\mathop{\max}_{u} \mathop{\min}_{\lambda_i \geq 0} \langle g_0,u \rangle - \frac{1}{2}||u||^2 + \sum_i^n \lambda_i \langle g_i,u \rangle$

在该问题中，首先需要最小化 $\lambda_i$ 的目标函数。假设 $\langle g_i,u \rangle \geq 0$ ，则应使 $\lambda_i \langle g_i,u \rangle = 0$ ；假设 $\langle g_i,u \rangle < 0$ ，则应使 $\lambda_i \langle g_i,u \rangle \to - \infty$ 。而该目标还包含 $u$ 的最大化，因此取前一种情况，此时min-max问题的优化结果与原带约束的max问题等价。

若定义 $\Bbb{Q}^n$ 为所有 $n$ 元离散分布的集合：

$\Bbb{Q}^n = \{ (\lambda_1,\lambda_2,\cdots,\lambda_n) | \lambda_1,\lambda_2,\cdots,\lambda_n\geq 0 \}$

则优化问题记为：

$\mathop{\max}_{u} \mathop{\min}_{\lambda \in \Bbb{Q}^n} \langle g_0+ \sum_i^n \lambda_i g_i,u \rangle - \frac{1}{2}||u||^2 \\ = \mathop{\max}_{u} \mathop{\min}_{\lambda \in \Bbb{Q}^n} \langle g_0+ \tilde{g}(\lambda),u \rangle - \frac{1}{2}||u||^2$

根据Minimax定理，如果min问题和max问题的参数可行域都是凸集，并且目标函数关于min问题的参数是凸的、关于max问题的参数是凹的，则min和max的次序可以交换：

$\begin{aligned} & \mathop{\max}_{u} \mathop{\min}_{\lambda \in \Bbb{Q}^n} \langle g_0+ \tilde{g}(\lambda),u \rangle - \frac{1}{2}||u||^2 \\ &= \mathop{\min}_{\lambda \in \Bbb{Q}^n} \mathop{\max}_{u} \langle g_0+ \tilde{g}(\lambda),u \rangle - \frac{1}{2}||u||^2 \\ &= \mathop{\min}_{\lambda \in \Bbb{Q}^n} \frac{1}{2}||g_0+ \tilde{g}(\lambda)||^2 \end{aligned}$

上式等价于求主目标损失梯度 $g_0$ 与各辅助目标损失梯度 $g_1,g_2,\cdots g_n$ 的所有加权组合之和中模长最小的情况。

当 $n=1$ 时，问题等价于 $\mathop{\min}_{\lambda \geq 0} \frac{1}{2}\|g_0+ \lambda g_1\|^2$ ，该问题具有明确的几何意义和简单的解析解：

当 $n>1$ 时，采用Frank-Wolfe算法将问题转化为多个 $n=1$ 的情形进行迭代，迭代过程为：

$\begin{aligned} \tau &= \mathop{\arg \min}_i \langle g_i,g_0+\tilde{g}(\lambda^{(k)}) \rangle \\ \gamma &= \mathop{\arg \min}_{\gamma} ||\tilde{g}(\lambda^{(k)}-\lambda^{(k)}_{\tau}e_{\tau}+\gamma e_{\tau})+g_0||^2 \\ &= \mathop{\arg \min}_{\gamma} ||\tilde{g}(\lambda^{(k)})-\lambda^{(k)}_{\tau}g_{\tau}+\gamma g_{\tau})+g_0||^2 \\ \lambda^{(k+1)} &= \lambda^{(k)}-\lambda^{(k)}_{\tau}e_{\tau}+\gamma e_{\tau} \end{aligned}$

对照上述讨论，具有 $m$ 个主目标、 $n$ 个辅助目标的混合型多目标优化，应具有的对偶目标函数为：

$\mathop{\min}_{w \in \Bbb{P}^n, \lambda \in \Bbb{Q}^n} \frac{1}{2}||\tilde{g}(w)+ \tilde{g}(\lambda)||^2 \\ = \mathop{\min}_{w \in \Bbb{P}^n, \lambda \in \Bbb{Q}^n} \frac{1}{2}||\sum_i^m w_ig_i+ \sum_j^n \lambda_j g_j ||^2$

⚪ 主次型多目标优化的应用

一个典型的主次型多目标优化案例是向任务损失中加入正则项，如L2正则化：

$\mathcal{L}(\theta) + \frac{\lambda}{2} ||\theta||^2$

此时正则化损失并不需要越小越好，而是希望能够提高原损失 $\mathcal{L}(\theta)$ 的泛化性能。此时两个目标损失的梯度分别为：

$g_0 = \nabla_{\theta}\mathcal{L}(\theta), g_1= \theta$

当 $\langle g_0,g_1 \rangle>0$ 时取 $\lambda=0$ ；当 $\langle g_0,g_1 \rangle<0$ 时取与 $g_1$ 正交的 $g_0+\lambda g_1$ ：

$\langle g_1,g_0+\lambda g_1 \rangle =0 \Leftrightarrow \lambda = -\frac{\langle g_0,g_1 \rangle}{||g_1||^2}$

将上述两种情况统一写作：

$\lambda = \frac{\text{ReLU}(-\langle g_0,g_1 \rangle)}{||g_1||^2}$

代入 $g_0,g_1$ 的表达式得：

$\lambda = \frac{\text{ReLU}(-\langle \nabla_{\theta}\mathcal{L}(\theta),\theta \rangle)}{||\theta||^2}$