How Flatiron Institute Mathematicians Helped Make AI More Efficient

How Flatiron Institute Mathematicians Helped Make AI More Efficient The data centers that train and run large language models (LLMs) such as ChatGPT now emit as much carbon dioxide as New York City and use more water than the global bottled water industry. And those numbers keep rising as demand for LLMs skyrockets. Making artificial intelligence more efficient, therefore, has become a critical component of curbing the industry’s resource demands. Over recent years, a subculture has grown around making the most resource-intensive component of building an LLM — the training of these models — faster and cheaper. Even a modest efficiency boost of a percent or two could reduce electricity demand by the equivalent of an entire country’s usage, saving tens of millions of dollars. That magnitude of savings is a major motivation fueling competition among deep learning hackers and scientists alike. Out of this competition, a new hyper-efficient AI training tool called Muon emerged, and researchers at the Simons Foundation’s Flatiron Institute developed one of its key components, called Polar Express. The leaps in progress have been staggering: In 2019, it took four days to train the predecessor of ChatGPT. In May 2024, the record was 45 minutes for a scaled-down version of ChatGPT called nanoGPT. Now, with advancements such as Polar Express, Muon’s latest version can train nanoGPT in just 1.23 minutes. Polar Express optimized Muon’s core operation — adjusting the strength of the connections between an AI model’s artificial neurons as it learns — and now powers the record-setting entries in the field’s most closely watched training contests. Beyond its ability to break speedrunning benchmarks, its ubiquity and the method’s mathematical optimality are points of pride for its creators. “Our Polar Express code embedded within Muon is so optimal that if you want to do better, you are going to have to reframe the question,” says Robert Gower, a research scientist at the Flatiron Institute’s Center for Computational Mathematics (CCM). Training the Machine Muon is an optimizer, a tool that trains AI models. It’s a mathematical algorithm that adjusts a model’s internal settings, improving its performance. Once training is done, the optimizer steps aside, and the model is ready for use. Gower explains that models start out as “baby language models that know absolutely nothing.” During training, the model examines data sequences, such as strings of words with a missing piece, and tries to fill in the blank. “If it guesses correctly, you give it a ‘lollipop,’” says Gower, meaning the model learns to do more of what worked. When it guesses wrong, the optimizer steps in and adjusts the model’s internal representation by slightly nudging its parameters. Optimization is repeated millions to billions of times until the model is fully trained. The mechanism for those parameter adjustments is called gradient descent. The model’s mistakes can be pictured as a mountainous landscape, and the optimizer always tries to go ‘downhill’ toward fewer mistakes. To know which way is downhill, the model needs to know the slope at that exact spot, what mathematicians call the gradient. The challenge is that precisely computing that slope would require processing the entire training dataset (essentially “the entire internet,” Gower says) just to take one step. Instead, models are trained on small, random batches from the training dataset. For over a decade, the Adam optimizer has been the standard. However, Adam is nonoptimal, because it is ‘blind’ to the model’s structure. Adam adjusts the model’s internal parameters — which are organized as two-dimensional grids of numbers called matrices — by flattening the matrices into a giant list called a vector. Every parameter is treated equally, regardless of its position or relationship to the others. Many researchers have attempted to end Adam’s dominance over the last 10 years and come up with something markedly better. “All those papers have formed this giant graveyard of algorithms that nobody uses,” says incoming CCM postdoc Michael Crawshaw. Optimizing the Optimizer Muon, the optimizer challenging Adam’s long reign, emerged from a competition. OpenAI co-founder Andrej Karpathy released the original nanoGPT code and challenged the world to train it as simply and efficiently as possible. Keller Jordan (now a researcher at OpenAI) and others created Muon as a branch of Karpathy’s code, which optimized the model’s internal ‘hidden’ layers, with Adam handling the input and output. Muon’s secret is to treat the model’s parameters in their native matrix form rather than flattening them into vectors. Each matrix works like a little machine: it takes in a string of words, transforms it and passes the result on. To improve the model, the optimizer makes small adjustments to these machines, and every adjustment has two parts — direction plus magnitude, meaning which way to move and by how much. The optimizer needs to account for both to ensure that the model grows steadily more accurate without becoming unstable. The direction Muon moves, explains Crawshaw, is what makes it special. Because the Adam optimizer flattens those matrices into vectors, ignoring their geometry, it can take only very small steps without destabilizing the whole training process. In contrast, Muon, which works with that geometry, can take much bigger steps, making major updates to the model at each step. That difference is why Muon trains in fewer overall steps, or more efficiently, and thus much faster. The Polar Express Gower saw an opportunity for CCM scientists to lend their expertise to the problem. To take one of those big update steps, Muon needs the gradient’s orientation. The direction-only version of the gradient is called the polar factor (which is where Polar Express derives its name from), and computing it is the computational heart of Muon. The problem is that doing this rigorously on a GPU is expensive. So Muon’s creators came up with a clever way to approximate the polar factor using polynomials, which GPUs excel at computing. A polynomial is a mathematical function composed of variables raised to some power, where each variable’s impact is defined by a number called a coefficient. The original implementation used a single set of coefficients, a reasonable approximation but only a ‘good enough’ one. Nobody knew how to find an optimal approximation. Gower recruited linear algebra expert David Persson, now a joint postdoc at the Flatiron Institute and New York University, along with NYU doctoral student Noah Amsel, to ask a more precise question: What is the best polynomial for the job? By reformulating the problem in rigorous mathematical terms, the team solved for the optimal expression, which recalculates the polynomial coefficients at every iteration to find the best approximation at each step. Their algorithm is the Polar Express, and those optimal coefficients are now a core component of the Muon optimizer. Polar Express Success The Polar Express paper earned an honorable mention at the 2026 International Conference on Learning Representations — one of only a handful awarded out of nearly 20,000 submissions. The next day, DeepSeek released a new model trained with Muon. Many Muon-trained models are outperforming larger and more expensive ones, pushing the boundary of model accuracy as a function of computing time. In addition to being integrated into the leading entries of the field’s premier training contests, including the nanoGPT speedrun, Amsel says that Polar Express is also being used by Microsoft Research and in larger ‘frontier-scale’ LLMs. The deep learning hackers haven’t stopped. Karpathy’s newest challenge is to retrain OpenAI’s original chatbot, nanochat, a task that once cost $43,000. With Muon and Polar Express, it now costs just $48. “I think Muon is a hopeful story for the field and for science in general,” says Amsel. “It was just thinking about the problem in a slightly different way that no one thought to do before.” Adopting decades-old applied mathematics into modern AI is exactly the kind of collaboration the Flatiron Institute fosters, Persson says. “This project wouldn’t have been possible without the cross-pollination of different perspectives,” he says. “My background is in classical numerical analysis and linear algebra. It wasn’t until Rob [Gower] approached me with this problem that I realized I could contribute to machine learning as well.”