Research map / Architectures and efficiency / RNNs and normalization
SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SMMF, a memory-efficient optimizer that square-matricizes arbitrary-rank momentum tensors and applies one-shot NNMF to both momenta, cutting optimizer memory by up to 96% with a regret bound similar to AdamNC.
Problem
Adaptive optimizers like Adam store two full-size momentum tensors, which limits trainable model size on memory-constrained devices. Existing memory-efficient optimizers (Adafactor, CAME, SM3) only work well on specific rank-2 or pattern-restricted momentum tensors and their memory remains high for
Approach
Reshape any rank-d momentum tensor into the matrix closest to square (square-matricization), then factorize both 1st and 2nd momentum with a single NNMF into two vectors, using a decompression-then-compression update order so the current gradient is reflected before factorization.
Key results
- SMMF reduces optimizer memory for ResNet-50 on ImageNet from 220 MiB (Adafactor) to 3.7 MiB while reaching 73.7% top-1 accuracy versus 69.5% for Adafactor.
- On Transformer-base full training on WMT32k, SMMF uses 0.8 GiB end-to-end memory versus 1.4 GiB for Adam with comparable perplexity (6.7).
- SMMF is on average about 1.4x slower than Adam per the appendix, trading time for memory.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SMMF instead of the strongest listed baseline (SM3) | MobileNetV2 on CIFAR100 · CIFAR100 | top-1 accuracy: 70.0 -> 74.1 | Helped |
| Use SMMF instead of the strongest listed baseline (SM3) | ResNet-50 · CIFAR100 | top-1 accuracy: 73.8 -> 74.5 | Helped |
| Use SMMF instead of the strongest listed baseline (CAME) | MobileNetV2 · ImageNet | top-1 accuracy: 69.5 -> 69.5 | No clear effect |
| Use SMMF instead of the strongest listed baseline (SM3) | ResNet-50 · ImageNet | top-1 accuracy: 75.8 -> 73.7 | Hurt |
| Use SMMF instead of the strongest listed baseline (Adafactor) | YOLOv5s · COCO | mAP50: 53.6 -> 54.1 | Helped |
| Use SMMF instead of the strongest listed baseline (CAME) | YOLOv5m · COCO | mAP50: 59.3 -> 59.6 | Helped |
| Use SMMF instead of Adam | Transformer-base full training · WMT32k | test perplexity: 6.6 -> 6.7 | Hurt |
| Use SMMF instead of Adam | Transformer-big full training · WMT32k | test perplexity: 6.9 -> 6.8 | Helped |
| Use SMMF instead of Adam | BERT pre-training · BookCorpus & Wikipedia | test perplexity: 16.1 -> 20.4 | Hurt |
| Use SMMF instead of Adam | GPT-2 pre-training · BookCorpus & Wikipedia | test perplexity: 19.2 -> 19.2 | No clear effect |
| Use SMMF instead of Adafactor | GPT-2 pre-training · BookCorpus & Wikipedia | test perplexity | Helped |
| Use SMMF instead of Adam | T5 pre-training · BookCorpus & Wikipedia | test perplexity: 2.6 -> 2.6 | No clear effect |
Surprising findings
- Adafactor diverged (NaN) when pre-training GPT-2 on BookCorpus and Wikipedia, while SMMF trained stably.
Limitations
- Sign matrix for the first momentum still adds 1-bit-per-element overhead, which is not fully eliminated.
- Only evaluated on models up to LLaMA-7b with LoRA; very large models and diffusion models were not tested.
- Occasional loss spikes early in training for Transformers, as with other optimizers.
Experimental setup
Datasets: CIFAR100 ImageNet COCO WMT32k BookCorpus Wikipedia QNLI MNLI QQP STSB MRPC
Benchmarks: CIFAR100 ImageNet COCO WMT32k QNLI MNLI QQP STSB MRPC
Metrics: top-1 accuracy mAP50 perplexity optimizer memory (MiB/GiB) optimization time per step
Baselines: Adam Adafactor SM3 CAME
Builds on: Adam Adafactor CAME SM3 AdamNC
Closest papers by meaning
- A geometric framework for momentum-based optimizers for low-rank training (NeurIPS 2025 · similarity 0.88)
- MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts (NeurIPS 2024 · similarity 0.87)
- Lean and Mean Adaptive Optimization via Subset-Norm and Subspace-Momentum with Convergence Guarantees (ICML 2025 · similarity 0.87)
- Momentum-SAM: Sharpness Aware Minimization without Computational Overhead (NeurIPS 2025 · similarity 0.87)
- Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration (NeurIPS 2026 · similarity 0.86)
- Tensorion: A Tensor-Aware Generalization of the Muon Optimizer (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.