内存高效的 4 比特预条件随机优化
机器学习
2025-03-13 v2 计算机视觉与模式识别
最优化与控制
摘要
以 Shampoo 为代表的预条件随机优化算法通过在大规模神经网络训练中提供理论收敛保证和实际增益，优于一阶优化器。然而，由于非对角预条件矩阵的存储需求，它们会产生巨大的内存开销。为解决这一问题，我们引入了针对 Shampoo 预条件矩阵的 4 比特量化。我们引入了两个关键方法：首先，我们对预条件矩阵进行 Cholesky 分解并对 Cholesky 因子进行量化，利用其下三角结构减少内存使用，同时更好地保留谱特性以最小化信息损失。据我们所知，这是首个应用于预条件矩阵 Cholesky 因子的量化方法。其次，我们在量化过程中引入误差反馈，将 Cholesky 因子和误差状态高效地存储在同一矩阵的下三角和上三角部分。通过大量实验，我们证明将 Cholesky 量化与误差反馈相结合，可提升大规模深度学习任务中的内存效率和算法性能。在理论上，我们还分别在平滑和非平滑随机优化设置下提供了量化 Shampoo 的收敛性证明。
引用
@article{arxiv.2412.10663,
title = {Memory-Efficient 4-bit Preconditioned Stochastic Optimization},
author = {Jingyang Li and Kuangyu Ding and Kim-Chuan Toh and Pan Zhou},
journal= {arXiv preprint arXiv:2412.10663},
year = {2025}
}