Research map / Architectures and efficiency / RNNs and normalization
CoMERA: Computing- and Memory-Efficient Training via Rank-Adaptive Tensor 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 CoMERA, a rank-adaptive tensor-compressed end-to-end training method with multi-objective (scalarization) rank control and GPU-optimized TTM embedding lookup, tensor contraction paths, and CUDA Graph, achieving speedups and memory savings over GaLore and LTE.
Problem
Training large AI models (LLMs, DLRMs) is extremely costly in memory and compute, and prior tensor-compressed training lacks automatic rank selection and has not shown real wall-clock or memory reductions on GPU.
Approach
Control TT/TTM ranks during end-to-end training via diagonal matrices with l1 relaxation and multi-objective scalarization (linear early stage, achievement scalarization late stage), combined with optimized TTM embedding lookup, a near-optimal contraction path for TT forward/backward, and CUDA Graph to make tensorized training faster on GPU.
Key results
- On a six-encoder transformer (MNLI), CoMERA early stage reaches 63.3% validation accuracy vs 62.2% uncompressed with 74x compression on tensorized layers and ~3x per-epoch speedup.
- CoMERA is about 2x faster than GaLore per epoch and uses 9x less memory than GaLore in single-batch training on RTX 3090.
- On DLRM, CoMERA achieves 99x model size reduction and 7x peak memory reduction with accuracy 78.76% vs 78.68% uncompressed; CodeBERT-LARGE pre-training gets 4.23x overall compression and 1.9-2.3x speedup per phase.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CoMERA rank-adaptive tensor compression instead of uncompressed training | Six-encoder transformer · MNLI validation | accuracy: 62.2% -> 63.3% | Helped |
| Apply late-stage rank adaptation with target ratio 0.8 (ablation) | CoMERA six-encoder transformer · MNLI validation | accuracy: 63.3% -> 62.2% | No clear effect |
| Apply late-stage rank adaptation with target ratio 0.5 (ablation) | CoMERA six-encoder transformer · MNLI validation | accuracy: 63.3% -> 62.1% | Hurt |
| Apply late-stage rank adaptation with target ratio 0.2 (ablation) | CoMERA six-encoder transformer · MNLI validation | accuracy: 63.3% -> 61.5% | Hurt |
| Replace uncompressed training with CoMERA compression | DLRM · Criteo Ad Kaggle test | accuracy: 78.68% -> 78.76% | No clear effect |
| Use optimized TTM embedding lookup instead of unoptimized TTM lookup (ablation) | CoMERA DLRM · DLRM training | training time per epoch: around 2× speedup | Helped |
| Use optimized TTM embedding lookup instead of unoptimized TTM lookup (ablation) | CoMERA DLRM · DLRM training | peak memory: 4-5× | Helped |
| Enable CUDA Graph (ablation) | CoMERA six-encoder transformer · MNLI training | training time: around 3× speed-up | Helped |
| Use CoMERA instead of GaLore | Six-encoder transformer · MNLI | training time per epoch: about 2× faster | Helped |
| Use CoMERA instead of LTE | Six-encoder transformer · MNLI | training time per epoch: 3× faster | Helped |
| Use CoMERA instead of GaLore in single-batch training | Six-encoder transformer · MNLI | memory cost: 9× more memory-efficient | Helped |
| Use CoMERA instead of standard pre-training | CodeBERT LARGE · CodeSearchNet pre-training, Phase 1 | training speed: 2.3× faster | Helped |
Surprising findings
- LTE did not converge on the MNLI task under its default setting.
- CoMERA's rank-adaptive training automatically removed some whole layers (e.g., an entire encoder at target ratio 0.2), a form of architecture search.
Limitations
- Speedups are from an initial GPU implementation; the authors note small tensor contractions are not well supported by GPU kernels and HPC optimization is pending.
- CodeBERT pre-training shows a final loss gap (0.40 vs 0.28) and only preliminary downstream evidence; DLRM training is slower than uncompressed training.
- No theoretical guarantee on the number of epochs needed for convergence.
Experimental setup
Datasets: MNLI Criteo Ad Kaggle CodeSearchNet Wikipedia SST-2 MRPC SQuAD
Benchmarks: MNLI Criteo Ad Kaggle SST-2 MRPC SQuAD
Metrics: validation accuracy compression ratio training time per epoch peak memory training loss normalized cross-entropy F1
Baselines: uncompressed training GaLore LTE standard pre-training
Builds on: SVDinsTN GaLore LTE CodeBERT DLRM
Closest papers by meaning
- Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage (NeurIPS 2025 · similarity 0.85)
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (ICML 2024 · similarity 0.84)
- FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training (ICML 2025 · similarity 0.84)
- Compressible Dynamics in Deep Overparameterized Low-Rank Learning & Adaptation (ICML 2024 · similarity 0.83)
- ReLoRA: High-Rank Training Through Low-Rank Updates (ICLR 2024 · similarity 0.83)
- AdaRankGrad: Adaptive Gradient Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning (ICLR 2025 · similarity 0.83)
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.