Second order stochastic optimizers allow parameter update step size and direction to adapt to loss curvature, but have traditionally required too much memory and compute for deep learning. Recently, Shampoo [Gupta et al., 2018] introduced a Kronecker factored preconditioner to reduce these requirements: it is used for large deep models [Anil et al., 2020] and in production [Anil et al., 2022]. However, it takes inverse matrix roots of ill-conditioned matrices. This requires 64-bit precision, imposing strong hardware constraints. In this paper, we propose a novel factorization, Kronecker Approximation-Domination (KrAD). Using KrAD, we update a matrix that directly approximates the inverse empirical Fisher matrix (like full matrix AdaGrad), avoiding inversion and hence 64-bit precision. We then propose KrADagrad$^\star$, with similar computational costs to Shampoo and the same regret. Synthetic ill-conditioned experiments show improved performance over Shampoo for 32-bit precision, while for several real datasets we have comparable or better generalization.

翻译：二阶随机优化器允许参数更新步长和方向自适应损失曲率，但传统上需要过多内存和计算资源用于深度学习。近年来，Shampoo [Gupta et al., 2018] 引入了克罗内克分解预条件子以降低这些需求：它已用于大型深度模型 [Anil et al., 2020] 并投入生产 [Anil et al., 2022]。然而，该方法需要对病态矩阵求逆矩阵根，这需要64位精度，从而施加了严格的硬件限制。本文提出一种新颖的分解方法——克罗内克近似支配（KrAD）。利用KrAD，我们直接更新一个近似逆经验Fisher矩阵（类似于全矩阵AdaGrad）的矩阵，从而避免求逆操作及64位精度。进而提出KrADagrad$^\star$，其计算成本与Shampoo相当，且具有相同的遗憾界。在合成病态实验中，该方法在32位精度下性能优于Shampoo，而在多个真实数据集上，我们获得了可比或更好的泛化能力。