AdaFactor

从Adam优化器谈起

设 $t$ 为迭代步数, $\alpha_t$ 为当前步学习率, $L(\theta)$ 是损失函数, $\theta$ 是待优化参数, $\epsilon$ 是防止溢出的小正数.

Adam优化器的更新过程如下:

\left\{\begin{aligned}&g_t = \nabla_{\theta} L(\theta_t)\\ &m_t = \beta_1 m_{t-1} + \left(1 - \beta_1\right) g_t\\ &v_t = \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2\\ &\hat{m}_t = m_t\left/\left(1 - \beta_1^t\right)\right.\\ &\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t \hat{m}_t\left/\sqrt{\hat{v}_t + \epsilon}\right. \end{aligned}\right.

作为目前最常用的优化器, 有一个缺点是占用的显存较大. 要省显存, 就首先得知道显存花在哪里的. 首先最后计算得到的梯度是占用显存的, 且这部分是任何优化器都无法节省的.

除此之外, Adam优化器使用了一阶梯度 $m$ 和二阶梯度 $v$ , 且在每一步使用滑动平均计算, 需要进行缓存, 这两部分也要占用缓存, 且各自占用的大小同上面的梯度一致.

AdaFactor节省显存的思路

抛弃动量

Adam性能优秀, 很重要的一个点是每一个参数都有自适应学习率, 从上面的公式中也可以看出:

$\theta_t = \theta_{t-1} - \alpha_t \hat{m}_t / \sqrt{\hat{v}_t + \epsilon}$

从上式中可以看出, 每个参数的自适应学习率为 $\alpha_t\left/\sqrt{\hat{v}_t + \epsilon}\right.$ , 即通过SGD+二阶动量来实现的.

因此作为节省缓存的第一步, 考虑直接抛弃一阶动量 $m$ , 这样显存的占用直接节省了1/3.

低秩分解

然后继续尝试压缩二阶动量 $v$ 的大小. AdaFactor使用到了低秩分解.

Adam中每个参数都会有各自独立的学习率, 但SGD中所有的参数共用一个学习率, 且SGD在很多任务或数据集中也能取得不错的效果, 带来的一个思路是精调每一个参数自己的学习率不是特别重要, 因此启发我们将 $\hat{v}_t$ 换一种参数更少的近似可能也就足够了. 因此使用低秩分解来实现.

对于 $m \times n$ 大小的矩阵 $C$ , 希望找到大小为 $m \times k$ 的矩阵 $A$ 和大小为 $k \times n$ 的矩阵 $B$ , 使得:

$AB \approx C$

使用一个比较小的 $k$ 值, 这样矩阵 $A$ 和 $B$ 中的参数之和会远小于原来 $C$ 中参数的数量, 且仍能取得近似的效果, 这就是上面说的不再精调每个参数的学习率, 让参数之间共享部分信息.

AdaFactor中取 $k=1$ , 将显存节省到了极致, 即寻找 $\{a_i\}_{i=1}^m$ 和 $\{b_j\}_{j=1}^n$ , 使得:

$a_i b_j \approx c_{i,j}$

为了达到近似的效果, 需要一个距离度量标准来进行约束, 容易想到欧式距离 $\sum_{i,j} (a_i b_j - c_{i,j})^2$ , 但这样 $a_i,b_j$ 没有解析解, 且在优化过程中 $c_{i,j}$ , 即对应于更新中的二阶梯度 $\hat{v}_t$ 应当是非负的, 但通过上面的目标函数优化得到的 $c_{i,j}$ 无法保证非负.

因此AdaFactor使用了新的度量标准, 广义KL散度, 形式为:

$l = \sum_{i,j} c_{i,j}\log \frac{c_{i,j}}{a_i b_j} - c_{i,j} + a_i b_j$

这个度量标准来自不等式 $x\log x\geq x - 1(\forall x > 0)$ , 当且仅当 $x=1$ 时等号成立. 将 $x = p / q\,(p,q > 0)$ 带入到不等式当中, 然后两端乘以 $q$ , 则有:

$p\log \frac{p}{q} - p + q \geq 0$

当且仅当 $p=q$ 时, 等号成立. 将 $p$ 替换成 $c_{i,j}$ , $q$ 替换成 $a_i b_j$ , 并且对所有的分量进行求和, 就得到了上面的广义KL散度的公式. 由于有取最小值的条件, 在将 $a_i$ , $b_j$ , $c_{i,j}$ 带入后, 刚好有解析解:

$a_i = \sum\limits_{j}c_{i,j},\quad b_j = \frac{\sum\limits_{i}c_{i,j}}{\sum\limits_{i,j}c_{i,j}}$

解析解也很形象, 就是行, 列分别求和, 然后相乘, 再除以全体的和. 推导过程参考.

因此我们就可以维护两组缓存变量 $v^{(r)}_t\in \mathbb{R}^m,v^{(c)}_t\in\mathbb{R}^n$ , 代表 $g_t^2$ 低秩分解后的结果, 解析解保证了 $v^{(r)}_tv^{(c)}_t$ 点乘于原始二阶动量 $g_t^2$ 之间在广义KL散度度量下的最大近似性. 因此AdaFactor优化器的计算流程如下:

\left\{\begin{aligned}&g_{i,j;t} = \nabla_{\theta} L(\theta_{i,j;t})\\ &v^{(r)}_{i;t} = \beta_2 v^{(r)}_{t-1;i} + \left(1 - \beta_2\right) \sum\limits_{j}\left(g_{i,j;t}^2+\epsilon\right)\\ &v^{(c)}_{j;t} = \beta_2 v^{(c)}_{t-1;j} + \left(1 - \beta_2\right) \sum\limits_{i}\left(g_{i,j;t}^2+\epsilon\right)\\ &v_{i,j;t} = v^{(r)}_{i;t} v^{(c)}_{j;t}\left/\sum\limits_{j}v^{(c)}_{j;t}\right.\\ &\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t g_t\left/\sqrt{\hat{v}_t}\right. \end{aligned}\right.

$v^{(r)}_t\in \mathbb{R}^m,v^{(c)}_t\in\mathbb{R}^n$ 两变量的更新逻辑, 保证了两变量的非负性, 使得不等式 $x\log x\geq x - 1(\forall x > 0)$ 始终满足条件, 解析解成立, 得到理论保证的近似效果, 进一步保证了优化的效果.

参考资料

上一页Adam 下一页Normalization

最后更新于4年前

这有帮助吗？