Research map / Optimization / Optimizers
Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training
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 Sophia, a light-weight second-order optimizer for LLM pre-training that uses EMA of a stochastic diagonal Hessian (refreshed every 10 steps) with per-coordinate clipping, achieving about 2x speedup over AdamW on GPT-2 and GPT NeoX models.
Problem
Pre-training large language models is extremely costly, and existing second-order optimizers have too much per-step overhead or fail to speed up LLM training in wall-clock time or compute, with Adam still dominant.
Approach
Precondition the EMA of gradients by a cheap EMA of stochastic diagonal Hessian estimates (Hutchinson or Gauss-Newton-Bartlett), refreshed every k=10 steps, and apply per-coordinate clipping of the update so that negative or rapidly changing curvature falls back to sign-momentum-like updates.
Key results
- Sophia-G reaches 2.875 validation loss on GPT-2 Small vs AdamW 2.921, and reaches the AdamW loss in about half the steps (2x speedup in steps, compute and wall-clock).
- On GPT-2 Large (770M), Sophia-G achieves 2.524 validation loss vs AdamW 2.613 at 100K steps.
- Sophia's per-step overhead is under 5% in wall-clock time and it uses the same memory as AdamW; a 540M Sophia-H model matches a 770M AdamW model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Sophia-H instead of AdamW | GPT-2 language-model pre-training · OpenWebText validation, GPT-2 Small (125M) | validation loss: 2.921 -> 2.901 | Helped |
| use Sophia-G instead of AdamW | GPT-2 language-model pre-training · OpenWebText validation, GPT-2 Medium (355M) | validation loss: 2.691 -> 2.627 | Helped |
| use Sophia-G instead of AdamW | GPT-2 language-model pre-training · OpenWebText validation, GPT-2 Large (770M) | validation loss: 2.613 -> 2.524 | Helped |
| use Sophia-G instead of AdamW | GPT NeoX language-model pre-training · Pile validation, GPT NeoX 1.5B | validation loss: 2.250 -> 2.218 | Helped |
| use Sophia-G instead of AdamW | GPT NeoX language-model pre-training · Pile validation, GPT NeoX 6.6B | validation loss: 1.992 -> 1.969 | Helped |
| update the Hessian estimate every step instead of every 10 steps (ablation) | Sophia-G · 30M GPT-2 | convergence speed per compute: 50% | Hurt |
| update the Hessian estimate every 100 steps instead of every 10 steps (ablation) | Sophia-G · 30M GPT-2 | validation loss | Hurt |
| replace the GNB preconditioner with empirical Fisher (ablation) | Sophia-G · 30M GPT-2 | validation loss | Hurt |
| replace the Sophia-H Hessian estimator with AdaHessian (ablation) | Sophia-H · 30M GPT-2 | validation loss | Hurt |
| add per-coordinate update clipping without a preconditioner (ablation) | AdamW · 30M GPT-2 | validation loss | Helped |
| add the GNB preconditioner to element-wise clipping (ablation) | element-wise clipping without a preconditioner · 30M GPT-2 | validation loss | Helped |
| remove element-wise update clipping (ablation) | AdaHessian with k=2 · 30M GPT-2 | training stability: k=2 | Hurt |
Surprising findings
- Sophia does not require the attention-logit scaling trick by layer index that AdamW and Lion need for stable training at larger scales.
- Sophia-H and Sophia-G trigger gradient clipping far less often than AdamW and Lion (under 10% of steps vs frequent triggering).
Limitations
- Not compared with AdamW or Lion on models larger than 6.6B due to limited resources.
- Evaluation on downstream tasks is limited to pre-training validation loss and four SuperGLUE subtasks.
- Evaluated only on language modeling; other domains such as vision, RL and multimodal are not tested.
Experimental setup
Datasets: OpenWebText The Pile SuperGLUE
Benchmarks: OpenWebText validation loss The Pile validation loss SuperGLUE (4 subtasks, 2-shot)
Metrics: validation log perplexity / loss steps to reach same loss total compute (TFLOPs) wall-clock time few-shot accuracy
Baselines: AdamW Lion AdaHessian Empirical Fisher with clipping
Builds on: AdamW AdaHessian Lion Hutchinson estimator
Closest papers by meaning
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.90)
- Benchmarking Optimizers for Large Language Model Pretraining (NeurIPS 2026 · similarity 0.89)
- Second-Order Fine-Tuning without Pain for LLMs: A Hessian Informed Zeroth-Order Optimizer (ICLR 2025 · similarity 0.89)
- Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training (NeurIPS 2025 · similarity 0.89)
- Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models (ICLR 2025 · similarity 0.89)
- Enhancing Large Language Model Performance with Gradient-Based Parameter Selection (AAAI 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.