R-iKFAD：以低秩摩擦减少 Transformer 预训练优化器内存

iKFAD 是最近提出的优化器，它用动量动力学中的自适应摩擦取代自适应学习率，但性能与 Adam 一样好。它的局限性在于，完整的摩擦张量 $\xi\in\mathbb{R}^{m\times n}$ 每层具有与 Adam 的二阶矩缓冲区相同的 $\mathcal{O}(mn)$ 内存开销。在这里，我们用根据行和列动量统计数据构建的秩一 的外积分解替换 iKFAD 的摩擦张量 $ \xi $，从而得到秩一 的 iKFAD (R-iKFAD)。这将每层的摩擦内存占用从 $\mathcal{O}(mn)$ 减少到 $\mathcal{O}(m+n)$，这大约使 iKFAD 的总优化器状态减半。尽管有所减少，R-iKFAD 仍保持与 iKFAD 的性能相当：GPT2-Nano、TinyViT、DistilBERT 和 GPT2-S 上的实验证实，它与 iKFAD 相当或超过 iKFAD，同时内存占用量几乎减半，并且对超参数保持相对稳健。我们分析了两种阻尼状态下的连续时间动力学。对于线性阻尼（$γ>0$），我们证明了强凸性下的指数收敛。对于 $γ=0$（我们实验中的首选选项），摩擦力完全由过去的动量产生，并随着动量消失而关闭，因此无法显示几何收敛。尽管如此，我们还是证明了收敛到极小化器，以及匹配能量的上限和下限：当正则化尺度 $ε_{stab}$ 为零时，阶数为 $t^{-1}$；当正则化尺度 $ε_{stab}$ 为正时，阶数为 $t^{-1/2}$。据我们所知，这是连续时间内秩一 因子优化器的第一个收敛速度，也是第一个不需要正阻尼的结果。