← Latest papers
🤖 machine learning

On the Principles Behind Neural Network Optimizers

This paper provides a principled theoretical foundation for the Adam optimizer by resolving its convergence debate, explaining its superiority over SGD on Transformers through evolving Hessian structures, and leveraging these insights to introduce Adam-mini, a new optimizer that halves memory usage while maintaining performance.

Original authors: Yushun Zhang

Published 2026-08-18
📖 7 min read🧠 Deep dive

Original authors: Yushun Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern artificial intelligence relies on a delicate balancing act. To teach a computer to recognize a cat in a photo or to write a coherent story, researchers must guide a massive mathematical system through a landscape of errors, constantly adjusting its internal settings to find the lowest point. This process is called training, and the tool used to navigate this terrain is an optimizer. For years, the industry standard has been an algorithm called Adam. It is the default engine for the most powerful language models, the ones that can write code, translate languages, and hold conversations. Yet, despite its ubiquity, the mathematical foundation of Adam has been shaky. For nearly a decade, a famous theoretical result suggested that Adam could fail catastrophically, spiraling out of control even on simple problems. This created a paradox: the tool that powers the most advanced AI was theoretically broken, yet it worked perfectly in practice. Researchers were left wondering why the theory failed to match reality and whether the tool they relied on was truly safe.

A new doctoral thesis by Yushun Zhang at the Chinese University of Hong Kong, Shenzhen, resolves this paradox by looking at the problem from a fresh angle. The work does not just patch the theory; it re-examines the very nature of the problems these models solve. The researcher discovered that the fear of Adam failing was based on a specific, artificial setup that does not reflect how real-world training actually happens. By changing the perspective to focus on the size of the data chunks used during training, the study proves that Adam is indeed safe, provided its settings are tuned correctly for the specific task. More importantly, the thesis uncovers a hidden geometric structure within the data of modern neural networks. This structure explains why Adam outperforms its rivals on complex tasks like training large language models, while failing on simpler ones. It turns out that the internal landscape of these models is not a chaotic mess, but rather a collection of distinct, organized blocks. Recognizing this pattern allowed the researcher to design a new, more efficient optimizer called Adam-mini, which cuts the memory required to train these massive models in half without sacrificing performance.

The story begins with the long-standing debate over whether Adam is reliable. For years, a widely cited paper claimed that Adam could diverge, meaning the training process would run away to infinity rather than settling down. This claim was based on a specific mathematical example where the algorithm was tested on a problem that changed its rules depending on the algorithm's settings. In the real world, however, researchers do not change the problem to fit the tool; they fix the problem and tune the tool to fit it. Zhang's work shows that when the problem is fixed, as it is in actual training, Adam does not diverge. Instead, it exhibits a clear phase transition: if the settings are chosen poorly, it can fail, but if they are chosen correctly, it converges safely. The key to this safety lies in a specific setting that controls how much weight the algorithm gives to past information. The study proves that for larger datasets, this setting must be higher to ensure stability. This finding aligns with what engineers have observed in practice: when training large language models with small batches of data, increasing this setting prevents the training from crashing. The thesis provides the first rigorous mathematical proof that the standard version of Adam, without any modifications, is safe to use when tuned properly.

Having established that Adam is safe, the research turns to a more puzzling question: why does it work so much better than its main competitor, SGD, on complex models like Transformers, while performing worse on simpler ones? To answer this, the researcher looked at the shape of the error landscape, specifically at a mathematical object called the Hessian, which describes how the error changes in every direction. In simple problems, this landscape is dense and tangled, like a thick forest where every path is connected to every other. In such environments, Adam's strategy of adjusting each setting individually is inefficient. However, when the researcher examined the Hessian of deep neural networks and Transformers, a surprising pattern emerged. As training progressed, the complex, tangled landscape simplified into a structure of distinct, separate blocks. Imagine a vast spreadsheet where, instead of every cell influencing every other cell, the influence is confined to specific rows and columns. In these networks, the parameters that control one specific output neuron or one specific attention head form their own isolated group.

This block-like structure is the secret to Adam's success. Because the landscape is broken into independent blocks, Adam's method of assigning a unique learning rate to each parameter becomes highly effective. It can adjust the settings for one block without accidentally messing up the settings for another. In contrast, simpler optimizers like SGD apply a single learning rate to the entire system, which struggles to handle the varying speeds and scales of these different blocks. The study further revealed that this block structure is not a coincidence; it arises naturally from the way these networks are built, specifically from the consecutive multiplication of large matrices during the calculation process. As the network trains, the connections between distant parts of the system fade away, leaving behind these clean, separate blocks. This insight explains why Adam is the engine of choice for modern AI: the problems it solves have a hidden geometry that matches the optimizer's design perfectly.

Armed with this understanding of the hidden block structure, the researcher developed a new optimizer called Adam-mini. The standard Adam algorithm is memory-hungry because it keeps track of a unique learning rate for every single parameter in the model. For a massive language model, this requires storing twice as much data as the model itself, creating a bottleneck that slows down training and limits the size of models that can be run on available hardware. The new insight was that because the parameters are organized into blocks, we do not need a unique learning rate for every single one. Instead, we can assign one learning rate to each entire block. This simple change reduces the memory footprint by 50 percent. The new optimizer, Adam-mini, groups parameters by their natural block structure—grouping by rows for most layers and by attention heads for specific parts of the network—and applies a single learning rate to each group.

The results of this redesign are immediate and practical. In tests training models ranging from 39 million to one billion parameters, Adam-mini matched the performance of the standard Adam optimizer while using half the memory. This efficiency allows researchers to train larger models on the same hardware or train existing models faster. The approach has already been adopted by major AI labs, including DeepSeek and the team behind the Kimi K3 model, who use a variation of this method to train their next-generation systems. The thesis also showed that this principle of block-wise learning rates can be applied to other advanced optimizers, improving their efficiency without changing their core logic. By revealing the hidden geometry of neural network training, this work has moved the field from a place of uncertainty and trial-and-error to one of principled design. It demonstrates that the most effective tools for artificial intelligence are not just lucky guesses, but algorithms that are perfectly tuned to the specific mathematical structure of the problems they solve.

Drowning in papers in your field?

Get daily digests of the most novel papers matching your research keywords — with technical summaries, in your language.

Try Digest →