SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization

TL;DR

SMMF introduces a memory-efficient optimization method that significantly reduces memory usage of adaptive optimizers like Adam, while maintaining comparable performance across CNN and Transformer models.

Contribution

It presents a novel square-matricization technique for flexible, low-rank momentum tensor factorization applicable to various deep learning architectures.

Findings

Reduces memory usage by up to 96% compared to existing optimizers.

Achieves similar model performance on CNN and Transformer tasks.

Provides theoretical convergence guarantees similar to AdamNC.

Abstract

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.