一种计算高效的稀疏化在线牛顿法
机器学习
2023-11-17 v1 计算与语言
最优化与控制
摘要
二阶方法在提升深度神经网络训练的收敛性方面前景广阔；然而，其巨大的内存与计算需求限制了其实用性。因此，亟需可扩展的二阶方法以高效训练大型模型。本文提出稀疏化在线牛顿（SONew）方法，一种内存高效的二阶算法，可生成稀疏而有效的预条件子。该算法源于对 LogDet 矩阵散度度量的新颖使用；我们将其与稀疏性约束结合，以在在线凸优化框架下最小化遗憾。在实证中，我们在规模达 1B 参数的大规模基准上测试了该方法。与包括一阶方法在内的内存高效优化器相比，我们实现了高达30%更快的收敛、3.4%验证性能的相对提升，以及80%训练损失相对改善。支撑该方法的是一个令人惊讶的事实——施加结构化稀疏模式（如三对角与带状结构）几乎不需要额外开销，使其与一阶方法一样高效且可并行。在挂钟时间上，三对角 SONew 每步仅比一阶方法慢约3%，但因收敛快得多而带来整体收益。相比之下，最先进的（SOTA）内存密集型二阶方法之一 Shampoo 无法扩展到大型基准。此外，虽然 Shampoo 需要大量工程努力才能扩展到大型基准，SONew 提供了更直接的实现，增强了其实用吸引力。SONew 代码见：https://github.com/devvrit/SONew
引用
@article{arxiv.2311.10085,
title = {A Computationally Efficient Sparsified Online Newton Method},
author = {Fnu Devvrit and Sai Surya Duvvuri and Rohan Anil and Vineet Gupta and Cho-Jui Hsieh and Inderjit Dhillon},
journal= {arXiv preprint arXiv:2311.10085},
year = {2023}
}
备注
30 pages. First two authors contributed equally. Accepted at NeurIPS 2023