AdaFactor

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)$始终满足条件, 解析解成立, 得到理论保证的近似效果, 进一步保证了优化的效果.

参考资料

• Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
• AdaFactor优化器浅析