通过 Importance Sampling 来提速神经网络训练

并非所有的数据都一样重要

这几天读了几篇和加速训练方面的相关文章,一般的加速训练的做法可能是通过混合精度训练、减小网络规模、或者是购买更多的显卡🤣,但很少有人会关注数据集本身。因为在算力充足的情况下,一般的想法还是认为数据集越大越好,因此很少有人会思考数据集本身。这么说起来和 Active Learning 会有一些联系,但是 Active Learning 本身讨论的是在获取数据较为困难的情况下如何进行学习,而本文则侧重在数据集充足的情况下,为了减少训练所需的计算量应该如何对数据集中的条目进行选择。

问题定义与 Importance Sampling

首先假设我们所需要优化的神经网络的形式如下:

$$ \theta^* = \arg\min_{\theta} \dfrac{1}{n} \sum_{i = 1}^n \mathcal{L}(\Psi(x_i; \theta), y_i) + \mathcal{R}(\theta) $$

其中 $\mathcal{R}(\cdot)$ 表示 Regularization Loss;

为了表示方便,我们定义:

$$ G_i = \nabla_\theta \mathcal{L}(\Psi(x_i; \theta), y_i) $$

$$ G_\mathcal{M} = \sum_{i \in \mathcal{M}} w_iG_i $$

$$ \overline{G} = \dfrac{1}{n}\sum_{i = 1}^n G_i $$

$$ G_\mathcal{M}^\mathcal{R} = G_\mathcal{M} + \nabla_\theta \mathcal{R}(\theta) $$

$$ \overline{G}^\mathcal{R} = \dfrac{1}{n}\sum_{i = 1}^n G_i + \nabla_\theta \mathcal{R}(\theta) $$

通常情况下我们是通过梯度下降的方式来进行优化的,即我们通过如下方式迭代:

$$ \theta_{t+1} = \theta_t - \eta \left[\sum_{i = 1}^n w_iG_i + \nabla_\theta \mathcal{R}(\theta)\right] $$

其中 $w_i$ 是待设定的参数(在一般的梯度下降中,我们取 $w_i = \dfrac{1}{n}$)

当然我们一般会采取选用 batch 的方式来进行梯度的近似而非一次性计算整个数据集上的梯度,即:

$$ \theta_{t+1} = \theta_t - \eta G_\mathcal{M}^\mathcal{R} $$

其中 $\mathcal{M}$ 是在原数据集 $\mathcal{D}$ 上根据分布 $p(\cdot)$ 采样出来的一个(可重复的)子集或者是 batch(其中 $p$ 也是待设定的参数),在我们一般的梯度下降中我们认为 $p$ 就是均匀分布;

当然为了保证我们的梯度是 unbiased 的,我们需要让这个分布按照 $p(\cdot)$ 采样出来的梯度的期望和原来是一样的,即:

$$ \mathbb{E}_{\mathcal{M}\sim p(\cdot)}\left[G_\mathcal{M}\right] = \mathbb{E}_{i \sim \mathcal{U}}\left[G_i\right] $$

通过命题1我们不难得知,当 $w_i = \dfrac{1}{np_i}$ 时,我们可以得到一个无偏的采样(这样通过在另外一个分布上 $p(\cdot)$ 采样,然后通过 $w$ 的设置来达到无偏的方法被称作 Importance Sampling)

最大化神经网络的收敛速度

在我读过的两篇相关的文献 [1, 2] 中,对于神经网络的收敛速度的定义都是一样的——一个关于 $p(\cdot)$ 的函数:

$$ \begin{split} S(p) &= -\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\left\lVert \theta_{t+1} - \theta^* \right\rVert^2 - \left\lVert \theta_t - \theta^* \right\rVert^2\right]\\
&=-\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \left\lVert \theta_{t+1} \right\rVert^2 - 2\left\langle \theta_{t+1}, \theta^* \right\rangle - \left\lVert \theta_t \right\rVert^2 + 2\left\langle \theta_t, \theta^* \right\rangle \right]\\
&=-\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \left\lVert \theta_{t} - \eta G_\mathcal{M}^\mathcal{R} \right\rVert^2 - \left\lVert \theta_t \right\rVert^2 + 2\eta \left\langle G_\mathcal{M}^\mathcal{R},\theta^* \right\rangle \right]\\
&=-\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\eta^2\left\lVert G_\mathcal{M}^\mathcal{R}\right\rVert^2 + 2\eta \left\langle G_\mathcal{M}^\mathcal{R},\theta^* - \theta_t \right\rangle \right]\\
&=2\eta\left\langle \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[G_\mathcal{M}^\mathcal{R}\right], \theta_t - \theta^*\right\rangle - \eta^2\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\left\lVert G_\mathcal{M}^\mathcal{R} \right\rVert^2\right]\\
&=2\eta\left\langle \overline{G}^\mathcal{R}, \theta_t - \theta^*\right\rangle - \eta^2\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\left\lVert G_\mathcal{M}^\mathcal{R} - \overline{G} + \overline{G} \right\rVert^2\right]\\
&=2\eta\left\langle \overline{G}^\mathcal{R}, \theta_t - \theta^*\right\rangle - \eta^2\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\left\lVert G_\mathcal{M} - \overline{G} + \overline{G}^\mathcal{R} \right\rVert^2\right]\\
&= 2\eta\left\langle \overline{G}^\mathcal{R}, \theta_t - \theta^*\right\rangle - \eta^2\left(\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\left\lVert G_\mathcal{M} - \overline{G} \right\rVert^2\right] + \left\lVert \overline{G}^\mathcal{R} \right\rVert^2\right)\\
&= 2\eta\left\langle \overline{G}^\mathcal{R}, \theta_t - \theta^*\right\rangle - \eta^2\left\lVert \overline{G}^\mathcal{R} \right\rVert^2 - \eta^2\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\left\lVert G_\mathcal{M} - \overline{G} \right\rVert^2\right]\\
&= 2\eta\left\langle \overline{G}^\mathcal{R}, \theta_t - \theta^*\right\rangle - \eta^2\left\lVert \overline{G}^\mathcal{R} \right\rVert^2 - \dfrac{\eta^2}{|\mathcal{M}|}\left(\mathbb{E}_{i \sim p(\cdot)}\left[ \left\lVert \dfrac{1}{np_i}G_i \right\rVert^2\right] - \left\lVert \overline{G} \right\rVert^2\right) \end{split} $$

即“我更新了这一步离最优解的差距变小了多少”(最后一个等号来源于命题1),那么在没有其他信息帮助的前提下,我们每一步更新肯定是让这个东西越大越好(当然我们也可以定义其他形式的收敛速度)。观察上面式子不难发现和 $p(\cdot)$ 有关的只有 $\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}G_i \right\rVert^2\right]$,因此最大化 $S(p)$ 等价于最小化 $\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}G_i \right\rVert^2\right]$。

这是一个受限优化问题(因为概率加起来必须要是 $1$),我们可以通过 Lagrange Multiplier 来解决,设:

$$ \begin{split} J(p, \lambda) &= \mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}G_i \right\rVert^2\right] - \lambda \left(\sum_{i = 1}^n p_i - 1\right)\\
&= \sum_{i = 1}^n \dfrac{1}{n^2p_i} \left\lVert G_i \right\rVert^2 - \lambda \left(\sum_{i = 1}^n p_i - 1\right) \end{split} $$

令$\dfrac{\partial J}{\partial p_i} = 0$得到:

$$ \dfrac{\partial J}{\partial p_i} = -\dfrac{1}{n^2p_i^2} \left\lVert G_i \right\rVert^2 - \lambda = 0 \Rightarrow p_i \propto \left\lVert G_i \right\rVert $$

因此我们得到的结论是:我们更加倾向于去找那些梯度模长更大的样本来进行训练

梯度模长近似

直接使用上面的方法很显然是行不通的,因为这实际上违背了我们使用这个方法的初衷——我们希望通过这个方法来减少反向传播的计算(因为梯度计算很慢!),然而我们发现要运行这个方法本身却需要对整个数据集进行梯度计算。这个问题在文献 [1] 和 [2] 中都有提到,但他们的解决思路不一样。

文献 [1] 中用到的方法是说我们可以证明 $\mathcal{L}$ 对全体参数的梯度的模长是可以用 $\mathcal{L}$ 对最后一层(激活函数前)的输出的梯度来代替的,这样我们只需要一次 Inference 就可以近似计算出梯度的大小了。除此之外,文中还提到了一种进一步优化的思路,作者设定了一种指标并在训练过程中不断地去维护这个指标,只有在这个指标满足一定的条件时才会采用作者提出的算法,否则还是按照原来的 SGD 来进行训练。

(未完待续)

参考文献

  1. Katharopoulos, Angelos, and François Fleuret. “Not all samples are created equal: Deep learning with importance sampling.” International conference on machine learning. PMLR, 2018. [链接]
  2. Johnson, Tyler B., and Carlos Guestrin. “Training deep models faster with robust, approximate importance sampling.” Advances in Neural Information Processing Systems 31 (2018): 7265-7275. [链接]

附录

命题1:基于 Importance 的随机采样 Batch 的均值与平方模长的均值

设 $\mathcal{U}$ 为均匀分布,$\mathbf{a_1}, \mathbf{a_2}, \dots, \mathbf{a_n}$ 为 $n$ 个向量,现在想从这 $n$ 个向量中随机抽取(可重复)$|\mathcal{M}|$ 个向量组成一个可重集合 $\mathcal{M}$(即 Batch),其中随机抽取时的分布服从 $p(\cdot)$,那么我们有:

  • $\mathbb{E}_\mathcal{M\sim p(\cdot)}\left[\dfrac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} \dfrac{1}{np_i} \mathbf{a_i} \right] = \overline{\mathbf{a}} = \mathbb{E}_\mathcal{i \sim \mathcal{U}}\left[\mathbf{a_i}\right]$
  • $\mathbb{E}_\mathcal{M\sim p(\cdot)}\left[\left\lVert \dfrac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} \dfrac{1}{np_i} \mathbf{a_i} \right\rVert^2\right] = \dfrac{1}{|\mathcal{M}|}\left(\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] - \left\lVert \overline{\mathbf{a}} \right\rVert^2 \right) + \left\lVert \overline{\mathbf{a}}\right\rVert^2$
  • $\mathbb{E}_\mathcal{M\sim p(\cdot)}\left[\left\lVert \dfrac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} \dfrac{1}{np_i} \mathbf{a_i} - \mathbf{\overline{\mathbf{a}}} \right\rVert^2\right] = \dfrac{1}{|\mathcal{M}|}\left(\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] - \left\lVert \overline{\mathbf{a}} \right\rVert^2 \right)$

证明: 我们用 $C_\mathcal{M}^{(i)}$ 表示 $\mathcal{M}$ 里面有多少个 $i$,即:

$$ C_\mathcal{M}^{(i)} = \sum_{k = 1}^{|\mathcal{M}|} [\mathcal{M}_k = i] $$

其中 $[\cdot]$ 是指示函数,$\mathcal{M}_k$ 表示可重子集 $\mathcal{M}$ 的第 $k$ 个元素

那么我们有:

$$ \begin{split} &\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \dfrac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i}\right] \\
=& \dfrac{1}{|\mathcal{M}|}\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\sum_{i = 1}^n C_\mathcal{M}^{(i)}\dfrac{1}{np_i}\mathbf{a_i}\right]\\
=& \dfrac{1}{|\mathcal{M}|}\sum_{i = 1}^n\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[C_\mathcal{M}^{(i)}\right] \dfrac{1}{np_i}\mathbf{a_i}\\
=& \dfrac{1}{|\mathcal{M}|}\sum_{i = 1}^n\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\sum_{k = 1}^{|\mathcal{M}|} [\mathcal{M}_k = i]\right] \dfrac{1}{np_i}\mathbf{a_i}\\
=& \dfrac{1}{|\mathcal{M}|}\sum_{i = 1}^n \sum_{k = 1}^{|\mathcal{M}|}\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[[\mathcal{M}_k = i]\right] \dfrac{1}{np_i}\mathbf{a_i}\\
=& \dfrac{1}{|\mathcal{M}|}\sum_{i = 1}^n \sum_{k = 1}^{|\mathcal{M}|}p_i\dfrac{1}{np_i}\mathbf{a_i} = \dfrac{1}{n}\sum_{i = 1}^n \mathbf{a_i} = \overline{\mathbf{a}} \end{split} $$

对于平方模长的均值,我们推理得到:

$$ \begin{split} &\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \left\lVert \dfrac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] \\
=& \dfrac{1}{|\mathcal{M}|^2}\mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \left\lVert \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i}\right\rVert^2 \right]\\
=& \dfrac{1}{|\mathcal{M}|^2} \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \left\lVert \sum_{i = 1}^n C_{\mathcal{M}}^{(i)}\dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2 \right]\\
=& \dfrac{1}{|\mathcal{M}|^2} \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \sum_{i = 1}^n\sum_{j = 1}^n C_{\mathcal{M}}^{(i)}C_{\mathcal{M}}^{(j)}\dfrac{1}{n^2p_ip_j} \left\langle \mathbf{a_i}, \mathbf{a_j} \right\rangle \right]\\
=& \dfrac{1}{|\mathcal{M}|^2} \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ \sum_{i = 1}^n\sum_{j = 1}^n \sum_{k = 1}^{|\mathcal{M}|}\sum_{l = 1}^{|\mathcal{M}|} [\mathcal{M}_k = i][\mathcal{M}_l = j]\dfrac{1}{n^2p_ip_j} \left\langle \mathbf{a_i}, \mathbf{a_j} \right\rangle \right]\\
=& \dfrac{1}{|\mathcal{M}|^2} \sum_{i = 1}^n\sum_{j = 1}^n \sum_{k = 1}^{|\mathcal{M}|}\sum_{l = 1}^{|\mathcal{M}|} \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ [\mathcal{M}_k = i][\mathcal{M}_l = j] \right] \dfrac{1}{n^2p_ip_j} \left\langle \mathbf{a_i}, \mathbf{a_j} \right\rangle\\
=& \dfrac{1}{|\mathcal{M}|^2} \left(\begin{split} &\sum_{i = 1}^n\sum_{k = 1}^{|\mathcal{M}|} \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[ [\mathcal{M}_k = i]^2 \right] \dfrac{1}{n^2p_i^2} \left\langle \mathbf{a_i}, \mathbf{a_i} \right\rangle + \\
&\sum_{i=1}^n\sum_{j=1}^n \sum_{k = 1}^{|\mathcal{M}|}\sum_{l = 1}^{|\mathcal{M}|}[k \neq l] \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[[\mathcal{M}_k = i][\mathcal{M}_l = j] \right]\dfrac{1}{n^2p_ip_j} \left\langle \mathbf{a_i}, \mathbf{a_j} \right\rangle \end{split}\right)\\
=& \dfrac{1}{|\mathcal{M}|^2} \left(\sum_{i = 1}^n |\mathcal{M}|p_i \dfrac{1}{n^2p_i^2} \left\lVert \mathbf{a_i}\right\rVert^2 + \sum_{i=1}^n\sum_{j=1}^n \sum_{k = 1}^{|\mathcal{M}|}\sum_{l = 1}^{|\mathcal{M}|}[k \neq l] p_ip_j\dfrac{1}{n^2p_ip_j} \left\langle \mathbf{a_i}, \mathbf{a_j} \right\rangle\right)\\
=& \dfrac{1}{|\mathcal{M}|^2} \left(\sum_{i = 1}^n |\mathcal{M}|p_i \dfrac{1}{n^2p_i^2} \left\lVert \mathbf{a_i}\right\rVert^2 + \sum_{i = 1}^n\sum_{j = 1}^n |\mathcal{M}|(|\mathcal{M}| - 1)\dfrac{1}{n^2}\left\langle \mathbf{a_i}, \mathbf{a_j} \right\rangle \right)\\
=& \dfrac{1}{|\mathcal{M}|}\left( \mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] + (|\mathcal{M}| - 1)\left\langle \dfrac{1}{n}\sum_{i = 1}^n \mathbf{a_i}, \dfrac{1}{n}\sum_{j = 1}^n \mathbf{a_j}\right\rangle \right)\\
=& \dfrac{1}{|\mathcal{M}|}\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] + \dfrac{|\mathcal{M}| - 1}{|\mathcal{M}|}\left\lVert \overline{\mathbf{a}}\right\rVert^2\\
=& \dfrac{1}{|\mathcal{M}|}\left(\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] - \left\lVert \overline{\mathbf{a}} \right\rVert^2 \right) + \left\lVert \overline{\mathbf{a}}\right\rVert^2 \end{split} $$

因此:

$$ \begin{split} &\mathbb{E}_\mathcal{M\sim p(\cdot)}\left[\left\lVert \dfrac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} \dfrac{1}{np_i} \mathbf{a_i} - \mathbf{\overline{\mathbf{a}}} \right\rVert^2\right] \\
=& \mathbb{E}_\mathcal{M\sim p(\cdot)}\left[ \left\lVert\dfrac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2 + \left\lVert\overline{\mathbf{a}} \right\rVert^2- 2\left\langle \dfrac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i},\overline{\mathbf{a}} \right\rangle \right]\\
=& \mathbb{E}_\mathcal{M\sim p(\cdot)}\left[ \left\lVert\dfrac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2 \right] + \left\lVert\overline{\mathbf{a}} \right\rVert^2- 2\left\langle \mathbb{E}_{\mathcal{M} \sim p(\cdot)}\left[\dfrac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i}\right],\overline{\mathbf{a}} \right\rangle\\
=& \mathbb{E}_\mathcal{M\sim p(\cdot)}\left[ \left\lVert\dfrac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2 \right] - \left\lVert\overline{\mathbf{a}} \right\rVert^2\\
=& \dfrac{1}{|\mathcal{M}|}\left(\mathbb{E}_{i \sim p(\cdot)}\left[\left\lVert \dfrac{1}{np_i}\mathbf{a_i} \right\rVert^2\right] - \left\lVert \overline{\mathbf{a}} \right\rVert^2 \right) \end{split} $$

除非特殊声明,本博客采用 CC BY-NC-SA 4.0 进行知识授权
comments powered by Disqus
Built with Hugo
Theme Stack designed by Jimmy