SMMF大幅降低自适应优化器内存占用，适用于各类模型结构。
SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization
- 通过平方矩阵化与单次矩阵分解，灵活压缩动量张量内存
- 相比现有方法最多节省96%内存，性能接近Adam等基准
- 适合大模型训练，尤其对低秩（如Transformer）和高秩（如CNN）模型均有效
我们提出SMMF（Square-Matricized Momentum Factorization），一种内存高效的优化器，可将广泛使用的自适应学习率优化器（如Adam）的内存需求减少高达96%。SMMF基于提出的平方矩阵化与一次性单矩阵分解，实现任意阶（形状）的一阶与二阶动量张量的灵活高效分解。这使其可有效应用于各种深度模型架构中常见的不同阶动量张量，包括偏置、矩阵及任意阶张量，覆盖卷积神经网络（高阶）与Transformer（低阶）等结构，而现有内存高效优化器仅适用于特定（二阶）动量张量，如线性层。我们对SMMF进行了后悔界分析，结果表明其收敛性与非内存高效的自适应优化器（如AdamNC）相当，为其实用性提供了理论支持。实验显示，SMMF在多种CNN与Transformer任务上，相较最先进的内存优化器（如Adafactor、CAME、SM3），内存占用最高减少96%，同时保持相近的模型性能。
原文摘要 · Abstract (English)
We propose SMMF (Square-Matricized Momentum Factorization), a memory-efficient optimizer that reduces the memory requirement of the widely used adaptive learning rate optimizers, such as Adam, by up to 96%. SMMF enables flexible and efficient factorization of an arbitrary rank (shape) of the first and second momentum tensors during optimization, based on the proposed square-matricization and one-time single matrix factorization. From this, it becomes effectively applicable to any rank (shape) of momentum tensors, i.e., bias, matrix, and any rank-d tensors, prevalent in various deep model architectures, such as CNNs (high rank) and Transformers (low rank), in contrast to existing memory-efficient optimizers that applies only to a particular (rank-2) momentum tensor, e.g., linear layers. We conduct a regret bound analysis of SMMF, which shows that it converges similarly to non-memory-efficient adaptive learning rate optimizers, such as AdamNC, providing a theoretical basis for its competitive optimization capability. In our experiment, SMMF takes up to 96% less memory compared to state-of-the-art memory efficient optimizers, e.g., Adafactor, CAME, and SM3, while achieving comparable model performance on various CNN and Transformer tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。