How Lion and Sophia Optimizers Speed Up AI Pretraining

Training massive artificial intelligence models requires immense computational power and memory, often creating bottlenecks that slow down development and skyrocket hardware costs. Advanced optimizers like Lion and Sophia resolve these issues by rethinking how neural networks update their weights during the pretraining phase. While the Lion optimizer drastically cuts memory usage by simplifying momentum tracking and only using the sign of gradients, Sophia leverages second-order curvature information to reach optimal performance in half the steps of traditional methods. Together, these memory-efficient optimizers accelerate the pretraining of massive AI architectures by freeing up hardware memory for larger batch sizes and navigating the training landscape much faster than the standard Adam optimizer.

The Bottleneck of Traditional Optimizers

For years, Adam and its variant, AdamW, have been the default optimizers for training large language models and vision transformers. However, Adam is memory-intensive. To update a model's parameters effectively, Adam must store two additional values for every single parameter in the network: the first moment (momentum) and the second moment (variance of the gradients).

For a model with 100 billion parameters, Adam requires enough memory to store an extra 200 billion values. This massive memory footprint limits the size of the model that can fit onto a GPU and restricts the batch size researchers can use, directly slowing down the pretraining process.

Lion: Memory Reduction Through Sign Operations

Discovered by Google researchers through an evolutionary search algorithm, the Lion (EvoLved Sign Momentum) optimizer strips away the excess memory requirements of Adam.

Lion accelerates pretraining through two main mechanisms:

  • Single Momentum Tracking: Unlike Adam, which tracks both first and second moments, Lion only tracks the momentum. This immediately cuts the optimizer’s memory overhead in half. The freed-up GPU memory can then be used to process larger batch sizes, which directly translates to faster training throughput.
  • Sign-Based Updates: Instead of calculating complex divisions and square roots to determine the exact magnitude of a weight update, Lion simply takes the sign (positive or negative) of the momentum. It applies a uniform update magnitude to every parameter. This reduces computational complexity and acts as a natural form of regularization, helping the model generalize better.

Sophia: Second-Order Speed at First-Order Cost

Developed by researchers at Stanford, Sophia (Second-order Clipped Stochastic Optimization) takes a different approach to acceleration. While Lion focuses on reducing memory per step, Sophia focuses on drastically reducing the number of steps required to train a model.

AI training can be visualized as trying to find the bottom of a complex valley. First-order optimizers like Adam only look at the steepness of the slope (the gradient) to take a step. This causes problems in uneven terrain, where the optimizer might take steps that are too small in flat areas and too large in steep areas.

Sophia introduces second-order information, known as the Hessian, which measures the curvature of the landscape. Historically, second-order optimizers were never used for massive AI models because calculating the curvature required an impossible amount of memory and compute. Sophia bypasses this limitation:

  • Lightweight Curvature Estimates: Sophia uses a cheap, randomized method to estimate the diagonal of the Hessian, rather than calculating the whole thing.
  • Infrequent Updates: Instead of calculating the curvature at every single training step, Sophia only updates its curvature estimate every 10 to 20 steps. This keeps the computational overhead down to around 5%, making it nearly as fast per-step as Adam.
  • Adaptive Step Sizes: By knowing the curvature, Sophia divides the gradient by the Hessian estimate. It takes massive, confident leaps in flat directions and careful, clipped steps in sharp, steep directions.

The Impact on Massive AI Architectures

When applied to massive architectures like Large Language Models (LLMs), the benefits of Lion and Sophia compound.

By using Lion, engineers can fit larger models onto fewer GPUs, reducing the need for slow cross-GPU communication. By using Sophia, models reach the same level of accuracy and low loss in roughly 50% fewer steps—cutting both the time and the financial cost of pretraining in half. Moving away from Adam toward these specialized optimizers allows organizations to train frontier AI models faster, cheaper, and with significantly less hardware strain.