Research map / Optimization / Optimizers
Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models
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.
Addax combines in-place first-order SGD on short sequences with MeZO zeroth-order gradients on long sequences to fine-tune LMs with near-MeZO memory and faster, more accurate convergence.
Problem
Full fine-tuning of LMs with Adam or IP-SGD needs too much memory, while MeZO's zeroth-order fine-tuning is memory-cheap but converges slowly and reaches lower final accuracy.
Approach
Partition data by sequence length: compute zeroth-order gradients (MeZO-style, in-place SPSA) on long sequences and in-place first-order gradients on short sequences, then update with a mixture alpha*z*g0 + (1-alpha)*g1, which the authors also interpret as regularizing the first-order gradient via Gaussian smoothing.
Key results
- On OPT-13B with one A100 (40GB), Addax outperforms MeZO by 14% in accuracy/F1 on average and converges about 15x faster at a comparable memory footprint.
- On OPT-30B with one H100 (80GB), Addax beats MeZO by more than 16% on average and converges about 30x faster; SGD and IP-SGD OOM on several tasks.
- Addax outperforms Adam on seven of nine OPT-13B tasks while reducing memory by up to 89%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| combine zeroth- and first-order gradient estimates in Addax | OPT-30B · OPT-30B short datasets | accuracy/F1: 75.1 -> 78.7 | Helped |
| combine zeroth- and first-order gradient estimates in Addax | OPT-30B · OPT-30B long datasets | accuracy/F1: 68.7 -> 82.0 | Helped |
| combine zeroth- and first-order gradient estimates in Addax | OPT-66B · OPT-66B short datasets | accuracy/F1: 77.1 -> 80.6 | Helped |
| combine zeroth- and first-order gradient estimates in Addax | OPT-66B · OPT-66B long dataset | accuracy/F1: 61.1 -> 80.6 | Helped |
| combine zeroth- and first-order gradient estimates in Addax | Llama-2-70B · Llama-2-70B short datasets | accuracy/F1: 78.2 -> 80.1 | Helped |
| combine zeroth- and first-order gradient estimates in Addax | Llama-2-70B · Llama-2-70B long datasets | accuracy/F1: 73.3 -> 88.9 | Helped |
| increase the first-order batch fraction K^1/(K^0 + K^1) (ablation) | Addax with RoBERTa-large · six RoBERTa-large tasks | accuracy | Helped |
| vary the mixing parameter alpha (ablation) | Addax with RoBERTa-large · six RoBERTa-large tasks | accuracy | No clear effect |
| use zeroth-order estimates alongside first-order updates | Addax on OPT-13B WIC · WIC | accuracy: 66.0 -> 68.3 | Helped |
Surprising findings
- Addax with a small first-order batch (K1=4) on WIC (68.3) beats IP-SGD with batch size 12 (66.0), suggesting the zeroth-order component helps even when first-order batch is small.
- Addax-WA (no data partition) still yields better final accuracy than IP-SGD on SST-2, RTE, and WSC, so zeroth-order gradients help even when memory is not a constraint.
Limitations
- Introduces an extra hyperparameter alpha that requires tuning; no consistent trend across tasks for alpha.
- Combination with PEFT, quantization, and memory-efficient attention is not explored.
- Memory numbers depend on chosen batch sizes and implementation; in-place updates preclude gradient accumulation and gradient normalization.
Experimental setup
Datasets: SuperGLUE SST-2 SST-5 SNLI MNLI RTE TREC BoolQ CB COPA MultiRC WIC
Benchmarks: SST-2 RTE WSC WIC BoolQ MultiRC SQuAD SST-5 SNLI MNLI TREC
Metrics: accuracy F1 peak GPU memory wall-clock time to best validation
Baselines: MeZO IP-SGD SGD Adam zero-shot in-context learning
Builds on: MeZO IP-SGD SPSA ZO-SGD
Closest papers by meaning
- Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark (ICML 2024 · similarity 0.91)
- Towards Efficient Low-Order Hybrid Optimizer for Language Model Fine-Tuning (AAAI 2025 · similarity 0.89)
- ZO-AdaMU Optimizer: Adapting Perturbation by the Momentum and Uncertainty in Zeroth-Order Optimization (AAAI 2024 · similarity 0.89)
- Gradient Weight-normalized Low-rank Projection for Efficient LLM Training (AAAI 2025 · similarity 0.89)
- Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · similarity 0.89)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.89)
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.