Memory-Efficient Adaptive Optimization

目录

• 概
• 符号说明
• SM3
  • 区间的划分
• 代码

Anil R., Gupta V., Koren T., Singer Y. Memory-efficient adaptive optimization. NeurIPS, 2019.

概

本文提出了一种 memory-efficient 的优化器: SM3.

符号说明

• $t = 1,\ldots, T$, optimization rounds;
• $w_t \in \mathbb{R}^d$, paramter vector;
• $\ell_t$, convex loss function;
• $g_t = \nabla \ell_t (w_t)$, 梯度;
• $w^* \in \mathbb{R}^d$, optimal paramter.

SM3

• 自适应的优化器 (AdaGrad) 形式如下:

  $$\gamma_t (i) = \sum_{s=1}^t g_s^2 (i), \quad \forall i \in [d],$$

  然后每一步按照如下的规则更新:

  $$w_{t+1}(i) = w_t(i) - \eta \frac{g_t(i)}{\sqrt{\gamma_t(i)}}, \quad \forall i \in [d].$$

• $\gamma_t$ 的存在意味着我们需要 $\mathcal{O}(d)$ 的额外存储. 作者提出的 SM3 将这个额外的存储消耗降低为 $\mathcal{O}(k)$.

• 首先, 通过某种方式确定 $k$ 个非空子集:

  $$\{S_r\}_{r=1}^k, \quad \bigcup_{r=1}^k S_r = [d].$$

• 然后按照如下的方式更新:

• 可以注意到, $S_r$ 的存在相当于指定 $i \in S_r$ 的参数共享一个自适应的学习率. 特别的, 由于 $S_r$ 不一定是互斥的, 所以每一次我们从中挑选一个最好的. 作者证明了这个方法的收敛性.

• 进一步的, 作者提出了一个更加的稳定的版本, 具有更好一点的 bound:

区间的划分

• 现在的问题是, 如何确定 $\{S_r\}_{r=1}^k$, 作者给出的建议是, 对一个 $m \times n$ 的权重, 可以分别按行共享和按列共享, 从而需要 $m + n$ 个缓存.

代码

[PyTorch-SM3]