> For the complete documentation index, see [llms.txt](https://kerasnoone.gitbook.io/garnet/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://kerasnoone.gitbook.io/garnet/shi-jian-ji-qiao/xian-cun-you-hua/you-hua-qi/adafactor.md).

# 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}}$$

解析解也很形象, 就是行, 列分别求和, 然后相乘, 再除以全体的和. 推导过程参考[推导过程](https://kexue.fm/archives/7302#%E6%8E%A8%E5%AF%BC%E8%BF%87%E7%A8%8B).

因此我们就可以维护两组缓存变量$$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](https://arxiv.org/abs/1804.04235)
* [AdaFactor优化器浅析](https://kexue.fm/archives/7302)