Gradient Descent Variants for LLMs: A Practical Comparison of SGD, Adam, AdamW, and Sophia
A practical benchmark comparison of SGD, Adam, AdamW, and Sophia optimizers for LLM training, with guidance on when to use each.
Meta description: A practical benchmark comparison of SGD, Adam, AdamW, and Sophia optimizers for LLM training. Learn which optimizer delivers the best speed, memory efficiency, and downstream performance.
When you train a large language model, the optimizer is the engine under the hood. It determines how fast your model converges. It dictates your GPU memory usage. And it shapes whether your model generalizes well to new tasks or just memorizes its training data.
Most practitioners default to AdamW as the universal default for LLM pretraining. It works. But is it always the right choice? And what about Sophia, the second-order newcomer that promises 2x speedup over AdamW for large models?
This guide cuts through the noise. We benchmark the four major optimizer families across the dimensions that matter for real LLM projects: convergence speed, downstream performance, memory footprint, and stability. All claims are grounded in 2025 benchmark data from the LLM optimizer benchmark study (arxiv/2509.01440).
What Is Gradient Descent and Why It Matters for LLMs
Gradient descent is the algorithm that minimizes a model's loss function during training. At each step, it calculates how each parameter contributed to the error and nudges parameters in the direction that reduces that error.
For small models, all parameters behave roughly the same. A single global learning rate handles the job.
LLMs contain parameters with wildly different gradient magnitudes. Embedding lookup gradients are sparse and large. LayerNorm scale gradients are tiny and dense. A single learning rate for all parameters means some update too aggressively while others barely move.
This heterogeneity is why first-order optimizers with a single global learning rate often struggle with transformers. Adaptive optimizers assign a per-parameter effective learning rate. This lets all parameters move on a similar timescale regardless of their raw gradient magnitude.
Choosing the right optimizer is not an afterthought. It is a design decision that affects training time, hardware costs, and model quality.
Meet the Optimizers: Core Mechanisms
Stochastic Gradient Descent (SGD)
SGD is the foundational optimizer. It updates parameters by stepping opposite to the gradient with a single global learning rate. Adding momentum tracks an exponential moving average of past gradients. This accelerates learning in consistent directions and smooths out oscillations.
SGD is memory efficient. It stores zero to one extra tensor per parameter (the momentum buffer). For a 70B parameter model, that is roughly 280GB of optimizer state at fp32.
The downside is a global learning rate that cannot handle the gradient heterogeneity of transformers. Embedding layers and attention heads get the same treatment. This makes SGD sensitive to learning rate choice and poor at serving all parameter types effectively.
Adam (Adaptive Moment Estimation)
Adam tracks both the first moment and second moment of gradients for each parameter. It normalizes gradients by the second moment estimate and applies bias correction to account for the fact that moment estimates start at zero.
The result is per-parameter adaptive learning rates. Parameters with large gradients get smaller effective learning rates. Parameters with small gradients get larger ones. Adam converges faster than SGD in most deep learning tasks because it adapts to the local geometry of the loss landscape.
Adam stores two extra tensors per parameter. Memory usage triples compared to SGD.
AdamW (Adam with Decoupled Weight Decay)
AdamW is Adam with decoupled weight decay. The key difference is how weight decay applies. In standard Adam, weight decay is blended into the gradient update step. This causes it to interact with the adaptive learning rate mechanism in ways that under-penalize sparse parameters and over-penalize dense ones.
AdamW decouples weight decay from gradient updates. Weight decay applies directly to the parameters after the gradient step. Dense layers stop underfitting. Sparse parameters are not over-penalized. This gives better regularization and improved generalization for transformers.
AdamW is the default optimizer for GPT, BERT, LLaMA, and T5 models. It is the safest choice for most LLM pretraining and fine-tuning tasks.
Sophia (Second-Order Clipped Stochastic Optimization)
Sophia is a second-order optimizer that leverages diagonal Hessian estimates. Where Adam uses second moments of gradients, Sophia estimates actual local curvature of the loss landscape. This allows more principled parameter updates — when curvature is steep, Sophia takes smaller steps.
Sophia estimates the diagonal Hessian every N steps (typically 10). It uses element-wise clipping to bound the worst-case update size. This prevents the numerical instabilities that plagued earlier second-order methods.
Sophia adds less than 5% average per-step computational overhead compared to first-order methods. Per-step cost is close to AdamW while incorporating curvature information.
Benchmark Results: 2025 Data
Training Loss and Convergence Speed
The most comprehensive 2025 benchmark evaluated 11 optimizers across model sizes from 124M to 720M parameters and training runs up to 50B tokens (arxiv/2509.01440).
Sophia consistently reaches the same perplexity as AdamW in roughly half the optimization steps. For GPT models from 125M to 6.6B parameters, Sophia achieves a 2x speedup in steps, total compute (exaFLOPs), and wall-clock time. The speedup becomes more pronounced with larger models.
Key benchmark finding — Sophia minimizes training loss fastest across model sizes 125M to 6.6B parameters, but AdamW leads on downstream zero-shot benchmarks. The fastest optimizer to train is not the same as the best-performing final model.
AdamW is the most reliable performer. It does not always minimize training loss fastest, but it generalizes the best to downstream tasks. Its convergence is stable across a wide range of learning rates and model sizes.
SGD with momentum performs comparably to adaptive methods for smaller models (150M to 1.2B parameters). It struggles for larger models where its global learning rate cannot handle gradient heterogeneity.
Lion (Evolved Sign Momentum) shows mixed results. Some studies found it faster than AdamW in GPU-hours. Others found AdamW slightly superior for downstream tasks. The gap narrows when both use proper hyperparameter tuning with Maximal Update Parametrization (μP).
Downstream Performance
Training loss is not the whole story. What matters is how the model performs on real tasks.
AdamW consistently leads on downstream zero-shot and transfer benchmarks. Sophia minimizes training and validation loss fastest, but AdamW models often perform better on benchmarks like HellaSwag, PIQA, and ARC-C after the same training compute.
Underreported trade-off — Sophia is faster to train. AdamW produces a better final product on downstream benchmarks. Most comparison articles quote the 2x speedup without mentioning this downstream performance gap.
This is the critical insight for practitioners: optimize for training speed only if your downstream task allows it.
Reinforcement Learning for LLMs
Here is a surprising result that most optimizer comparisons miss.
For Reinforcement Learning from Verifiable Rewards (RLVR) applied to LLMs, SGD matches or surpasses AdamW. The adaptive learning rates and momentum components of AdamW are less critical in RL than in supervised fine-tuning. SGD produces significantly sparser updates.
SGD is also more memory efficient. In memory-constrained RL environments, this matters. A 70B model with SGD optimizer state fits in fewer GPUs than with AdamW.
SGD excels in reinforcement learning for LLMs. This is an emerging use case where older optimizers are finding new life.
Memory Footprint: A Hidden Cost
Optimizer memory overhead is often glossed over. It should not be.
| Optimizer | Memory Multiplier | 7B Model (fp32) | 70B Model (fp32) |
|---|---|---|---|
| SGD | 1x | ~28GB | ~280GB |
| Adam/AdamW | 3x | ~84GB | ~840GB |
| Sophia | ~2x (estimated) | ~56GB | ~560GB |
AdamW triples your optimizer state memory. For a 70B parameter model in fp32, that is an additional 560GB just for optimizer bookkeeping. This directly affects your maximum feasible batch size and the number of GPUs you need.
Sophia halves the overhead of AdamW. It stores diagonal Hessian estimates instead of two full moment tensors. This makes Sophia more practical for memory-constrained training runs.
Key insight — For a 7B parameter model, switching from AdamW to SGD reduces optimizer state memory by roughly 56GB. This is enough to double your batch size in many hardware configurations.
Gradient checkpointing and mixed precision training (fp16/BF16) reduce these numbers significantly. But the relative ratios hold. AdamW is always the most memory-hungry optimizer option.
When to Use Each Optimizer
Use AdamW When
- You are pretraining or fine-tuning any standard LLM architecture (GPT, LLaMA, BERT, T5)
- You need reliable convergence without extensive hyperparameter tuning
- Downstream task performance is your primary metric
- You have ample GPU memory or can use gradient checkpointing
AdamW is the universal default for LLM pretraining and fine-tuning. It is the optimizer used to train GPT-3, GPT-4 (reported), BERT, LLaMA 1 and 2, and most instruction-tuned models.
Use Sophia When
- Training speed matters more than peak downstream accuracy
- You are training very large models where 2x speedup translates to significant compute savings
- Your training objective is stable and you can implement proper warmup
- You have infrastructure to handle its slightly higher implementation complexity
Sophia is not a plug-in replacement for AdamW everywhere. It requires warmup to stabilize Hessian estimates. It can diverge with large learning rates. But for large-scale pretraining runs where time is money, it is increasingly attractive.
Use SGD When
- You are applying RL to LLMs (RLVR setup)
- Memory is constrained and you need the minimal footprint
- You are training computer vision models and fine-tuning vision-language models
- You want a simple baseline for debugging
SGD is not dead for LLMs. It is finding new relevance in the RL pipeline where its sparse updates and memory efficiency beat adaptive methods.
Use Adam (Not AdamW) When
- You are doing quick prototyping and will discard the model
- You have not yet decided on a production optimizer
Always prefer AdamW over Adam for any model that will see production use. The decoupled weight decay gives meaningfully better regularization.
The Road Ahead: Second-Order Methods and Beyond
Sophia is not the end of optimizer evolution. Research is actively pushing in several directions.
Adaptive Hessian estimation frequencies could further reduce Sophia's overhead. Currently it estimates every 10 steps. Dynamic schedules that estimate more often early in training and less often later could improve efficiency.
Distributed training variants of Sophia are under development. Second-order methods were historically considered too communication-intensive for multi-GPU setups. Recent work suggests this barrier is lower than previously thought.
Automatic hyperparameter tuning for second-order methods is a frontier. Sophia has fewer hyperparameters than AdamW (no beta1/beta2 to tune), but learning rate and clipping threshold still require calibration.
The broader trend is clear: the ML community is moving beyond pure first-order methods for large-scale training. Second-order curvature information is becoming practical at scale. Sophia is the first commercially viable implementation of this shift.
Expert Q&A
Q: Is AdamW always the best choice for LLM training? A: Not always. AdamW is the safest default for pretraining and fine-tuning, but Sophia can be 2x faster for large models when training speed matters more than peak downstream accuracy. SGD can outperform AdamW in reinforcement learning setups.
Q: Why does AdamW outperform regular Adam for transformers? A: AdamW decouples weight decay from gradient updates. Regular Adam applies weight decay indirectly through the adaptive learning rate mechanism, which can under-penalize sparse parameters and over-penalize dense ones, leading to worse generalization.
Q: What is Sophia's main limitation? A: Sophia can diverge in large-batch or long-training setups with typical learning rates. It requires careful warmup and performs best when the training objective is stable. It also lags AdamW on downstream zero-shot benchmarks.
Q: How much more GPU memory does AdamW use compared to SGD? A: AdamW stores two additional tensors (first and second moment estimates) per parameter, effectively tripling the parameter storage footprint. For a 7B parameter model, this can mean ~84GB vs ~28GB just for optimizer state.
Q: Can I switch optimizers mid-training? A: Yes. Many practitioners warm up with AdamW and switch to SGD for the RL phase. However, switching during pretraining is less common and requires careful learning rate recalibration.
Ready to go deeper on LLM training optimization? Subscribe to the Algorithmine portal for weekly technical breakdowns of ML engineering best practices.