Second-order methods hold significant promise for enhancing the convergence of deep neural network training; however, their large memory and computational demands have limited their practicality. Thus there is a need for scalable second-order methods that can efficiently train large models. In this paper, we introduce the Sparsified Online Newton (SONew) method, a memory-efficient second-order algorithm that yields a sparsified yet effective preconditioner. The algorithm emerges from a novel use of the LogDet matrix divergence measure; we combine it with sparsity constraints to minimize regret in the online convex optimization framework. Empirically, we test our method on large scale benchmarks of up to 1B parameters. We achieve up to 30% faster convergence, 3.4% relative improvement in validation performance, and 80% relative improvement in training loss, in comparison to memory efficient optimizers including first order methods. Powering the method is a surprising fact -- imposing structured sparsity patterns, like tridiagonal and banded structure, requires little to no overhead, making it as efficient and parallelizable as first-order methods. In wall-clock time, tridiagonal SONew is only about 3% slower per step than first-order methods but gives overall gains due to much faster convergence. In contrast, one of the state-of-the-art (SOTA) memory-intensive second-order methods, Shampoo, is unable to scale to large benchmarks. Additionally, while Shampoo necessitates significant engineering efforts to scale to large benchmarks, SONew offers a more straightforward implementation, increasing its practical appeal. SONew code is available at: https://github.com/devvrit/SONew
翻译:二阶方法在提升深度神经网络训练收敛性方面具有显著潜力,但其巨大的内存和计算需求限制了其实用性。因此,需要发展可扩展的二阶方法,以高效训练大型模型。本文提出了一种稀疏化在线牛顿方法(Sparsified Online Newton, SONew),这是一种内存高效的二阶算法,可生成稀疏而有效的预处理器。该算法源于对LogDet矩阵散度度量的创新应用,通过结合稀疏性约束,在在线凸优化框架中最小化遗憾。在实验层面,我们在高达10亿参数的大规模基准测试中验证了该方法。与包括一阶方法在内的内存高效优化器相比,我们实现了高达30%的更快收敛、验证性能相对提升3.4%以及训练损失相对降低80%。该方法的核心支撑在于一个令人惊讶的事实:施加结构化稀疏模式(如三对角和带状结构)几乎不引入额外开销,使其在效率和并行化方面与一阶方法相当。在挂钟时间上,三对角SONew每步仅比一阶方法慢约3%,但由于收敛速度显著加快,总体收益更高。相比之下,现有最先进(SOTA)内存密集型二阶方法之一Shampoo无法扩展到大规模基准测试。此外,尽管Shampoo需要大量工程努力才能扩展到大规模基准,SONew提供了一种更简洁的实现,增强了其实用性。SONew代码可在https://github.com/devvrit/SONew获取。