论文浅读之《Active Sampling for Min-Max Fairness》
前言
我也是从去年年末因工作需要开始接触min-max相关的内容，所谓min-max，简单来说便是将数据进行分组之后， 设计一种学习方法，使得所有组别的损失的最大值尽可能的小—这个过程是一个动态的过程，随着学习过程的进行，损失最大的组别可能会发生改变。 此类方法通常在小样本学习或是样本不均衡的时候起到一些作用—最小化所有组别的损失的最大值，从而避免出现某些组效果很好，而某些组效果非常差的情况，尽可能的保证了学习的“公平性”，这也是文章标题所谓的“Fairness”。
而具体到实际做法，最符合直觉思考的便是在采样时，选择当前损失最大的组别进行采样学习，从而降低损失最大组别的损失。 类似的思路早在2016年DeepMind的论文《Prioritized Experience Replay》就被使用过—根据样本的损失大小进行优先级排序，优先级越高的样本被采样到的概率越大。 而Active Sampling for Min-Max Fairness中所做的要更加简单—每次更新时我们直接在损失最大的组别中进行采样学习。
1. Min-Max
min-max的问题设定总是从数据分组开始的。
我们在我们的数据空间$\mathcal{Z}$中有着$g$个组别， 每个组别中数据的分布为$D_1, D_2, \ldots, D_g$。
我们所学习的模型（model）可以看做是参数为$\theta \in \Theta \subsetneq R^d$的函数$f_\theta: \mathcal{Z} \rightarrow R^c$。 对于给定的数据样本$z \in \mathcal{Z}$， 我们使用损失函数$l(f_\theta; z)$来衡量模型的优劣。
对于不同组别中不同的数据分布，我们使用
\[v(\theta; D) = \mathbb{E}_{z\sim D}l(f_\theta; z)\]来表示分布$D$下的期望损失。 那么所谓的min-max，即寻找一个最优的模型$f_{\theta^\star}$， 使得所有组别的损失的最大值$\text{max}_{i \in [g]} v(\theta^\star; D_i)$最小，即
\[\max_{i\in [g]}v(\theta^\star; D_i) = \inf_{\theta \in \Theta} \max_{i \in [g]} v(\theta; D_i)\]其中$[g] = {1, 2, \ldots, g}$。
文章所给出的求解上述min-max问题的方法非常简单—每次从损失最大的组中进行采样学习，最后将整个学习过程中的参数进行平均即可
Algorithm 1: Min-max Stochastic Gradient Descent
1: Init: $\theta_1 \in \Theta$ arbitrary
2: for $t = 1, \ldots, T-1$ do
3: compute $i_t = \text{argmax}_{i\in [g]} v(\theta_t; \hat{D}_i)$
4: sample $z_t \sim D_{i_t}$
5: compute $\nabla_t \leftarrow \nabla_\theta l(f_{\theta_t}; z_t)$
6: update $\theta_{t+1} \leftarrow Proj_{\Theta}(\theta_t - \eta \nabla_t)$
7: end for
8: return $\overline \theta_T = \frac{\sum_{t=1}^T \theta_t}{T}$
这里的$\hat D_i$是指的是从$D_i$中采样得到的若干样本所构成的经验分布。 而为了保证算法的收敛性，作者假设了参数空间$\Theta$是一个紧致凸集（compact convex set）。同时，他也要求（假设）损失函数$l(f_\theta; z)$ 在参数$\theta$上是凸的。 因此参数空间$\Theta$是有边界的，那么如果我们在更新过程中超出了参数空间的边界，我们使用一个$l_2-$投影将其投影到参数空间中，即
\[Proj_\Theta(\theta_1) = \text{argmin}_{\theta \in \Theta} \lVert\theta - \theta_1 \rVert_2\]1.1. 收敛性
Theorem 1. Assume we have a function $R_\delta$ which guarantees that \(\sup_{\theta \in \Theta} \max_{i \in [g]} \lvert v(\theta; D_i) - v(\theta; \hat D_i) \rvert \le R_\delta\) with probability at least $1 - \delta$. Let $W=\sup_{\theta \in \Theta} \lVert\theta - \theta_1\rVert_2$ and $L= \sup_{\theta \in \Theta} \max_{i \in [g]} \lVert \nabla_\theta v(\theta; D_i) \rVert_2$. With $\eta=\frac{W}{L\sqrt{T}}$, Algorithm 1 ensures that \(\mathbb{E}_{z_{1:T}}[\max_{i\in [g]} v(\overline \theta_T;D_i)] \le \inf_{\theta \in \Theta} \max_{i \in [g]} v(\theta; D_i) + \frac{WL}{\sqrt{T}}+2R_\delta\) with probability at least $1-\delta$.
定理中首先假设了我们每一组数据的经验分布与其真实分布的期望损失的上界为$R_\theta$，考虑到在实际学习过程中，我们不太可能获取数据的真实分布，因此更多时候，我们用到的是Theorem 1的如下推论
Corollary 1. Consider a version of Algorithm 1 that, on line 4, draws samples IID from the empirical distribution $\hat D_{i_t}$ as opposed to fresh samples from $D_{i_t}$. Then, with $W, L$ defined as in Theorem 1, we have \(\mathbb{E}_{z_{1:T}}[\max_{i\in [g]} v(\overline \theta_T;\hat D_i)] \le \inf_{\theta \in \Theta} \max_{i \in [g]} v(\theta; \hat D_i) + \frac{WL}{\sqrt{T}}\)
Theorem 1告诉我们在执行$T$次更新后，模型的损失与最优模型的损失可以控制在$\frac{WL}{\sqrt{T}}$以内， 即算法的收敛速度为$\mathcal{O}(\frac{1}{\sqrt{T}})$。 注意到这个误差项中包含了参数空间的直径$W$与梯度的上界$L$，因此我们必须保证参数空间是有界的，否则$W$便可能取到无穷大。 这便是作者需要假设参数空间为紧致凸集的主要原因。
Theorem 1的证明并不困难，其主要是在Online Gradient Descent (OGD)的基础上做了一些小小的变换，如下所示
\[\begin{align*} & \mathbb{E}_{z_{1:T}}[\max_{i\in [g]} v(\overline \theta_T;D_i)] \\ \text{(Jensen's inequality)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \max_{i\in [g]} \sum_{t=1}^T v(\theta_t;D_i)] \\ \text{(max sum $\le$ sum max)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \sum_{t=1}^T \max_{i\in [g]} v(\theta_t;D_i)] \\ \text{(deviation between $D, \hat D$ + union bound)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \sum_{t=1}^T \max_{i\in [g]} v(\theta_t; \hat D_i)] + R_\delta \\ \text{(definition of $i_t$)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \sum_{t=1}^T v(\theta_t;\hat D_{i_t})] + R_\delta \\ \text{(additional deviation between $\hat D, D$)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \sum_{t=1}^T v(\theta_t;D_{i_t})] + 2R_\delta \\ \text{(since $z_t \sim D_{i_t}$ + outer expectation)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \sum_{t=1}^T l(f_{\theta_t};z_t)] + 2R_\delta \\ \text{(apply OGD regret bound)} & \le \mathbb{E}_{z_{1:T}}[\frac{1}{T} \sum_{t=1}^T l(f_{\theta^\star};z_t) + \frac{\eta L^2}{2} + \frac{W^2}{2T\eta}] + 2R_\delta \\ \text{($\eta=\frac{W}{L\sqrt{T}}$)} & = \mathbb{E}_{z_{1:T}}[v(\theta^\star;\{z_1, z_2, \ldots, z_T\})] + \frac{WL}{\sqrt{T}} + 2R_\delta \\ \text{(average $\le$ maximum)} & \le max_{i \in [g]} v(\theta^\star; D_i) + \frac{WL}{\sqrt{T}} + 2R_\delta \end{align*}\]式中的$\theta^\star$为最优参数。
1.2. Online Gradient Descent
接下来让我们来看看1.1节的证明中作者使用了OGD算法的regret bound究竟是什么东西。
所谓的OGD看起来是下面这个样子的
Algorithm 2: Online Gradient Descent
Input: compact convex set $K$, $x_1 \in K$, step sizes ${\eta_t}$
1: for $t=1,2, \ldots, T$ do
2: Play $x_t$ and observe cost $f_t(x_t)$
3: Update: $y_{t+1} = x_t - \eta_t \nabla f_t(x_t)$
4: Project: $x_{t+1} = Proj_K (y_{t+1})$
采样，计算损失，梯度下降，投影一气呵成。所谓Regret，即累计最优损失与学习过程中的累计损失之差
\[\text{Regret}_T = \sum_{t=1}^T f_t(x_t) - \min_{x^\star \in K}\sum_{t=1}^T f_t(x^\star)\]OGD的Regret Bound由以下定理给出
Theorem 2. Onilne gradient descent with step sizes ${\eta_t=\frac{W}{L\sqrt{t}} }$ guarantees the following for all $T \ge 1$: \(\text{Regret}_T = \sum_{t=1}^T f_t(x_t) - \min_{x^\star \in K}\sum_{t=1}^T f_t(x^\star) \le \frac{3}{2} WL\sqrt{T}.\) where $L > 0$ is an upper bound on the norm of the subgradients of $f$ over $K$, i.e., $\lVert \nabla f(x) \rVert \le L$, and $W$ is an upper bound on the diameter of $K$, i.e., $\forall x,y \in K, \lVert x-y \rVert \le W$.
证明并不难，我们首先使用勾股定理，将投影进行放缩，则有
\[\lVert x_{t+1} - x^\star \rVert^2 = \lVert Proj_K(x_t -\eta_t \nabla f_t(x_t)) - x^\star \rVert^2 \le \lVert x_t - \eta_t \nabla f_t(x_t) - x^\star \rVert^2\]展开右边的平方项，得到
\[\lVert x_{t+1} - x^\star \rVert^2 \le \lVert x_t - x^\star \rVert^2 + \eta_T^2 \lVert \nabla f_t(x_t) \rVert^2 - 2\eta_t \nabla^T f_t(x_t) (x_t - x^\star)\]用上界$L$替换掉$\lVert \nabla f_t(x_t) \rVert^2$，即有
\[2\nabla^T f_t(x_t) (x_t - x^\star) \le \frac{\lVert x_t - x^\star \rVert^2 - \lVert x_{t+1} - x^\star \rVert^2}{\eta_t} + \eta_t L^2\]式子左边的项实际上是非常眼熟的，它经常出现在凸集的定义当中。由于我们的空间$K$是凸的，因此有
\[\nabla^T f_t(x_t) (x_t - x^\star) \ge f_t(x_t) - f_t(x^\star)\]将上述式子利用$\nabla^T f_t(x_t) (x_t - x^\star)$串起来，得到
\[2\sum_{t=1}^T (f_t(x_t) - f_t(x^\star)) \le \sum_{t=1}^T \frac{\lVert x_t - x^\star \rVert^2 - \lVert x_{t+1} - x^\star \rVert^2}{\eta_t} + L^2 \sum_{t=1}^T \eta_t\]考虑到$\lVert x_{T+1} - x^\star \rVert^2 \ge 0$，我们将$\frac{\lVert x_{T+1} - x^\star \rVert^2}{\eta_T}$从和式中拿出，得到
\[2\sum_{t=1}^T (f_t(x_t) - f_t(x^\star)) \le \sum_{t=1}^T \lVert x_t - x^\star \rVert^2 (\frac{1}{\eta_t} - \frac{1}{\eta_{t-1}}) + L^2 \sum_{t=1}^T \eta_t\]并用直径$W$替换到$\lVert x_t - x^\star \rVert^2$
\[2\sum_{t=1}^T (f_t(x_t) - f_t(x^\star)) \le \sum_{t=1}^T D^2 (\frac{1}{\eta_t} - \frac{1}{\eta_{t-1}}) + L^2 \sum_{t=1}^T \eta_t \le \sum_{t=1}^T D^2 \frac{1}{\eta_T} + L^2 \sum_{t=1}^T \eta_t\]由于$\eta_t = \frac{W}{L\sqrt{t}}$，因此我们有
\[\sum_{t=1}^T \eta_t = \frac{W}{L} \sum_{t=1}^T \frac{1}{\sqrt{t}} \le \frac{W}{L} \sum_{t=1}^T \frac{2}{\sqrt{t} + \sqrt{t - 1}} = \frac{2W}{L} \sum_{t=1}^T (\sqrt{t} - \sqrt{t-1}) = \frac{2W\sqrt{T}}{L}\]最终得到
\[2\sum_{t=1}^T (f_t(x_t) - f_t(x^\star)) \le 3DG\sqrt{T}\]命题得证。
1.3. Accelerated Min-max Gradient Descent
上面给出的Algorithm 1收敛速度实在是太慢了，$\mathcal{O}(\frac{1}{\sqrt{T}})$。 作者尝试对其进行了改进，使用一种基于最大熵（maximum entropy）的权重来决定当前从哪个组别进行采样，如Algorithm 3所示
Algorithm 1: Min-max Stochastic Gradient Descent
1: Init: $q_0 = (\frac{1}{g}, \ldots, \frac{1}{g})$, $\theta_1 \in \Theta$ arbitrary
2: for $t = 1, \ldots, T$ do
3: Compute $u_t(i) = v(\theta_t; \hat D_i)$ for $i = 1, \ldots, g$
4: Update $q_t(i) = q_{t-1}(i) \exp (\gamma u_t(i))$ for $i = 1, \ldots, g$
5: Normalize $q_t = \frac{q_t}{\lVert q_t \rVert_1}$
6: Compute $\nabla_t = \nabla_\theta v (\theta_t; \hat D_{q_t}$)
7: Update $\theta_{t+1} = Proj_\Theta (\theta_t - 2\eta \delta_t + \eta_{t-1})$
8: end for
9: return $\overline \theta_T = \frac{\sum_{t=1}^T \theta_t}{T}$
其中$\hat D_{q_t}$指的是 先采样组别$i \sim q_t$，之后从$\hat D_i$中采样数据样本。
Algorithm 3的收敛性由以下定理给出
Theorem 2. Algorithm 3, with parameters $\eta = \frac{W}{L \sqrt{\log g}}$ and $\gamma = \frac{\sqrt{\log g}}{WL}$, outputs $\overline{\theta}_T$ that satisfies \(\max_{i\in [g]} v(\overline \theta_T;\hat D_i) \le \inf_{\theta \in \Theta} \max_{i \in [g]} v(\theta; \hat D_i) + \frac{2WL \sqrt{\log g}}{T}\) where $W$ and $L$ are defined as in Theorem 1.
这个定理的证明核心是一个被称之为Optimistic Mirror Descent (OMD)的算法上界证明，这不是什么新东西，但比起OGD还是需要更多的背景的，因此在此处我们先挖一个坑，Theorem 2的证明就留在未来OMD的文章之中吧。
Theorem 2同时给出了算法的收敛速度，$\mathcal{O}(\frac{1}{T})$，比起Algorithm 1就要快得多了。
2 后记
今天我们简单的讲述了min-max的问题设定以及一种非常简单的min-max算法—尽管这个简单的算法是发布在ICML上的，显得有些不可思议。 不过遗憾的是，ICML这篇文章并没有过多体现min-max可能的应用场景，在实验部分也主要是展示了其优越的收敛速度，即其第二个主要贡献Algorithm 3。 实际上min-max是有可能在应对不均匀分布的数据时发挥一些作用的—或许在下次min-max相关的内容中会为大家进行讲述。