Muon vs AdamW, Lion and Sophia: a fair race

Intro

0:00 Hello and welcome back to episode twenty-three, where we're going to race AdamW, the optimizer we've used since episode two, against three newer ones on our stories model. They're Lion, Sophia and Muon, and we're going to see how much each of them actually saves. Let's begin with a quick look at where we're going to be by the end of the episode.

0:21 By the end of today, our stories model trains with four optimizers, each at its own best learning rate. With AdamW, tuned, its held-out loss after one thousand steps is about one point nine three. With Muon, it's about one point seven six, a loss AdamW needs about eighteen hundred steps to reach.

0:40 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and today, the docs get a page on optimizers, with the race and its results. And each episode has a page here as well, with its chapters and its code.

1:00 Now, before we get to today's work, let's quickly go over what we did last time. We gave our web model, trained on windows of one thousand and twenty-four tokens, a window four times as long, by stretching how its rotary embeddings turn. With YaRN and no training at all, its loss stayed about flat to four thousand and ninety-six tokens, a little above a sliding window's. After fifteen minutes of fine-tuning on long pages, it went below the sliding window, so it used the tokens far back. Each minute of those runs, and of each run before them, came down to the optimizer's step, and that's what we're going to look at today.

1:38 We're going to write AdamW from scratch and check it against PyTorch's. Then we set up a race that's fair to each optimizer: one model, one set of batches, a thousand steps, and a sweep of the learning rate for each of them. Papers often compare a new optimizer against an AdamW at a learning rate that doesn't suit it, and the sweep is there to avoid that. The first challenger is Lion, which Google found with a search over programs, and which moves each weight by the sign of a blend of its momentum and its gradient. The second is Sophia, which divides by an estimate of the loss's curvature, and the third is Muon, which treats each weight matrix as a whole, and is what Andrej Karpathy's nanochat trains with. Then we measure the steps each needs, the time a step takes, and the memory each keeps. All right, let's get started with what an optimizer does.

What an optimizer does

2:28 All our training scripts make their optimizer in one line, and grep, which prints the lines that match a pattern, with dash n for their line numbers, finds ten of them. There's pretraining and the run on several GPUs, supervised fine-tuning, LoRA, DPO, PPO and GRPO, the calculator's and the agent's fine-tunes, and last episode's long pages. All of them use PyTorch's AdamW, and since episode two, we haven't asked whether another optimizer would train our model faster.

2:57 After each backward pass, each weight w has a gradient, g, which says how the loss changes when the weight grows. Gradient descent moves each weight against its gradient, by the learning rate, eta, times the gradient. Momentum keeps a running average of the gradients, m, so that the noise of each batch partly cancels out, and a direction that keeps coming back builds up. Beta one sets how quickly the average forgets, and we use zero point nine. Adam keeps a second running average, v, of the gradient squared, with beta two at zero point nine five. Then it divides the momentum by the square root of v, the typical size of that weight's gradient. So whatever the scale of a weight's gradients, its step is about the learning rate in size, which lets one learning rate suit all the weights. The hats are a correction for the first steps, as both averages start at zero. And weight decay shrinks each weight by a little, the learning rate times lambda, apart from the update, which is the W in AdamW.

3:59 The W comes from Decoupled Weight Decay Regularization, by Ilya Loshchilov and Frank Hutter, from 2017. Its abstract says that for an adaptive method like Adam, adding the weights to the gradient, which libraries called weight decay, isn't equivalent to weight decay. The way I see it, Adam then divides that term by v along with the gradient. So they decouple it, and shrink the weights apart from the update. Algorithm two is the result, with the two versions marked in colour. Lines seven to ten are the running averages and their corrections. Line twelve is the update, with AdamW's decay added outside the division.

4:39 We're going to write our optimizers in a new file, optimizers dot py, starting with AdamW. It subclasses PyTorch's optimizer class, which keeps the weights in groups, each group with its own settings, and gives each weight a dictionary of state, for what the optimizer remembers between steps. Over the class, a comment says where it comes from: Adam, by Kingma and Ba, from 2014, with decoupled weight decay, by Loshchilov and Hutter, and what it does in a sentence. The defaults are the settings our training scripts use: a learning rate of six times ten to the minus four, and betas of zero point nine and zero point nine five. Epsilon is ten to the minus eight, and the weight decay is zero point one. The base class stores them as each group's settings. Step runs with no gradient, as nothing in it should become part of the graph that backward goes through. For each group and each weight with a gradient, it looks up the weight's state. The first time, the state is empty, so it gets a step count and two tensors of zeros, shaped like the weight. The comment says what they are: two numbers per weight, the momentum, m, and the average of the gradient squared, v. So AdamW keeps twice our model's size in memory, which we're going to come back to. Then the two averages are updated in place, as the functions ending in an underscore change their tensor instead of making a new one. So m becomes beta one times m, plus one minus beta one times the gradient, and v does that with the gradient squared. Both start at zero, so the first averages are too small, and dividing by one minus beta to the power of the step corrects them. On the first step, m hat is the gradient, and v hat its square. Then the decay, which the comment calls decoupled: the weights shrink a little each step, apart from the update. And the update is the learning rate times m hat, over the square root of v hat plus epsilon, which keeps the division safe where a gradient is zero. The tests go in tests slash test optimizers dot py. Tiny builds a small model with two blocks, from a seed. Train takes the model, a list of optimizers and a number of steps. It seeds torch, so each call gets one batch of random tokens and random targets, and trains on it again and again. Each step runs the forward pass, zeroes each optimizer's gradients, runs the backward pass and steps each optimizer, and the losses on the way come back as a list. The comment over it says that, and why the optimizers come as a list: a model can have more than one optimizer, each for some of its weights, and one of today's needs that. The first test is the one that counts: two copies of the small model, one trained with our AdamW and one with PyTorch's, with our betas, epsilon and weight decay, for five steps. Each weight has to match to within ten to the minus five, which allows for the order PyTorch does its sums in. The second checks something worth knowing about Adam. After one step, m is the gradient and v its square, so the update is the gradient over its own size, plus or minus one, as its comment says. So each weight moves by the learning rate, up or down, unless its gradient is so small that epsilon counts. Lion, later on, is going to do that at each step.

7:52 Let's run them with pytest, where dash v lists each case. Both pass, so our AdamW takes the steps that PyTorch's takes, and its first step moves each weight by the learning rate.

A fair race

8:04 To compare optimizers fairly, everything else has to stay put. So each run trains our stories model, with rotary embeddings and the BPE tokenizer, from one starting point. A seed fixes the batches and their order, so each optimizer reads the stories the others read, in that order. Each run takes a thousand steps, with the warmup and the cosine from episode two, and its loss is measured on one set of held-out batches, for all the runs. Each optimizer gets its own sweep of the peak learning rate, by factors of two, then at the square root of two either side of its best. The best rate for one optimizer can be far from the best for another. Picking one rate for all of them would favour whichever it suits. And the best rate of each runs again with two more seeds, so we know how much a result moves by chance alone.

8:53 The race is bake off dot py, which trains our stories model with an optimizer and a peak learning rate, for a fixed number of steps. Next to argparse, json, math, os and time, it imports numpy and torch, and the summary writer, so that each run logs to TensorBoard, like our other runs. From our own files, it takes the runs folder, the data folder and the tokenizer, the results folder, the model and its config, and our AdamW. Each run takes a thousand steps, with a warmup of a hundred and batches of forty-eight, as in pretraining. A fixed number of steps, where pretrain dot py had a number of minutes, means each optimizer reads as many tokens as the others, however fast its steps are. It measures the held-out loss each hundred steps, on forty batches, and the comment works the tokens out: forty times forty-eight times two hundred and fifty-six, about half a million. Names lists the optimizers we have so far, AdamW alone. Make optimizers turns a name and a learning rate into a list of optimizers for the model. As in pretrain dot py, weight decay applies to the weight matrices, the parameters with two dimensions or more. The biases and the norms' gains get a weight decay of zero, in a group of their own, as shrinking a gain towards zero would fight what the norm is there for. AdamW gets the two groups, and a name we don't have raises a value error that lists the names we do. Learning rate is the schedule from episode two, by step instead of by time. It climbs to the peak over the warmup, then follows a cosine down to a tenth of the peak by the last step. So the peak is the number each sweep varies, and the shape of the schedule doesn't change. Batches gives random windows of the tokens, one batch at a time, as a Python generator, a function that yields a value each time we ask for the next one. The starts come from a torch generator, and the comment says why: a seed fixes the batches, whatever trains, so two runs with one seed read the stories in one order. Each window becomes the inputs, and the window one token along becomes the targets, both moved to the device. Sync waits for the GPU, and its comment explains it. PyTorch queues the GPU's work and goes back to Python straight away, so without it, a clock would time how long Python took to queue the work. Held out loss averages the loss over the held-out batches, in evaluation mode, with no gradients, and goes back to training mode at the end. Main builds the stories model from episode eleven, with rotary embeddings and the BPE tokenizer, and the comment explains the one change, no dropout. A run reads the stories about once, and dropout is for a model that sees its data many times. The seed fixes the starting weights, so each optimizer starts from one model. The training batches come from a generator seeded with the run's seed, and the held-out ones from a seed of their own, so all the runs are measured on one set of tokens, forty batches made up front. Each run is named after the optimizer, the learning rate and the seed, with the steps when they aren't a thousand. It logs to TensorBoard under bake off, and prints what it's training with, and how many weights the model has. Evaluate is a function inside main, so it can use main's model, held-out batches and writer without having them passed in. It measures the held-out loss, keeps it with its step, logs it to TensorBoard, and prints it, with the batch's loss, the learning rate and the milliseconds a step has taken so far. It runs once before any training, so each run's curve starts from the untrained model. The loop sets the learning rate of each group of each optimizer, which is how a schedule changes it from step to step in PyTorch. Then it takes the next batch and times the step, with a sync before the clock starts, so the work queued before it doesn't count. Then the forward pass, zero grad, the backward pass and clipping, as in pretraining, and the clock is read again in between, before the optimizers step. So trained adds up the time of whole steps, and stepped the optimizers' share of it, as the comment says. The batch's loss goes to TensorBoard at each step, and each hundred steps, and at the last one, the loop runs an evaluation. At the end, kept counts the numbers each optimizer keeps between steps, all the tensors in their state, which is what an optimizer costs in memory. The result is a dictionary, written as JSON for the report. It has the run's settings, its name, learning rate, seed, steps and tokens, then the final held-out loss and each evaluation along the way. Then the time per step and the optimizers' part of it, in milliseconds, the number of weights, and what the optimizers keep. It goes to a JSON file in results, bake off, which makedirs creates the first time, and the last line prints the loss, the tokens, the time, and the numbers kept per weight. On the command line, argparse reads the arguments. The optimizer's name comes first, and choices limits it to the names we have, so a typo stops with a list of them. Then one or more learning rates, which nargs plus collects into a list, with a seed, a number of steps and a device as options. A loop runs main for each learning rate, one after the other, so one command runs a whole sweep. Bake off report dot py reads the results back. Runs gathers the results of all the runs. Glob finds the files in results, bake off, whose names end in dot json, sorted, so the order doesn't depend on the file system, and each file is read into a dictionary. Sweep picks an optimizer's full-length runs with seed zero, and maps each learning rate to its final held-out loss, as its comment says. The other seeds and lengths come in later. Short writes a learning rate like one point two e minus three, where Python would write zero point zero zero one two. It formats the number in scientific notation and splits it at the e, which gives the digits and the power. Main gets each optimizer's sweep, and the learning rates any of them has a run at, sorted. Then it prints a table, with a row for each learning rate and a column for each optimizer, with a blank where that optimizer has no run. On the command line, it takes the names of the optimizers to show, and the comment over them says that none means all of them. And the tests get one more, for each optimizer in the bake off. Pytest's parametrize mark runs the test once for each name in names, so each optimizer we add gets it without our writing it again. Twenty steps on our batch of random tokens, at a learning rate of ten to the minus three, have to lower the loss by more than zero point two, so an optimizer that's wired up wrong fails here first.

15:30 Let's run it with AdamW at a peak learning rate of two point four times ten to the minus three, which, as we're going to see, is its best. Our stories model has about thirteen point eight million weights. At step zero, its held-out loss is about nine, the loss of a guess spread evenly over the eight thousand one hundred and ninety-two tokens of our vocabulary. It falls fast in the first hundred steps, as the warmup ends, and then more slowly, as the cosine brings the learning rate down. After a thousand steps, about eight and a half minutes on the laptop, it's at about one point nine three. And AdamW keeps two numbers per weight, m and v.

16:09 I ran AdamW's sweep off camera, a run for each peak learning rate, each about eight and a half minutes, and bake off report dot py reads their results. At six times ten to the minus four, the learning rate pretrain dot py uses, AdamW ends at about two point one three. Half of it is worse, at two point three four, and four times it is the best of the sweep, at one point nine three. The two points at the square root of two either side, about one point nine five and one point nine four, are close to it, so the best is in that range. At twice the best, the loss climbs again, to about two point zero five. Our pretraining ran for about ten minutes too, with the learning rate a quarter of the best for a run this long. Picking it again costs us about zero point two of loss, and that's the size of difference we're going to be looking for between optimizers.

Lion

17:04 The first challenger comes from Symbolic Discovery of Optimization Algorithms, by Xiangning Chen and others at Google, from 2023. Its abstract says they searched over programs, small pieces of code that update the weights, and found Lion, short for evolved sign momentum. It keeps track of the momentum and nothing else, so it needs less memory than Adam, and its update has one size for all the parameters, through the sign operation. Program one is all of Lion, in eight lines. The update is a blend of the gradient and the momentum, and its sign. Then the momentum is blended with the gradient with a second beta, and the weight decay is added to the update. Section five gives the betas: the search found zero point nine and zero point nine nine, and for language models, the paper used zero point nine five and zero point nine eight, so we're going to use those. And in that section, the paper says a learning rate for Lion is typically three to ten times smaller than for AdamW, with a weight decay three to ten times larger, so that the decay keeps a similar strength.

18:06 Lion blends the momentum and the newest gradient into c, with a first beta. Then each weight moves by the learning rate times the sign of c, plus one or minus one, whatever the size of its gradient. Where Adam's first step moved each weight by the learning rate, Lion does it at each step. The momentum is updated after the step, with a second beta. For language models, the paper used zero point nine five and zero point nine eight, where its search had found zero point nine and zero point nine nine. A running average with a beta of zero point nine five remembers about the last twenty steps, and one with zero point nine eight about the last fifty. So the momentum averages over a longer stretch, while each update leans more on the newest gradient. Each weight's step is the full learning rate, which is larger than most of Adam's steps, and that's why its learning rate has to be smaller.

18:58 Lion goes under AdamW in optimizers dot py. Its comment says where it comes from, and what it does: each weight moves by the learning rate, up or down, by the sign of a blend of the momentum and the gradient. Like our AdamW, it subclasses PyTorch's optimizer, so it gets the groups of weights and their state from the base class. Its betas are the ones the paper used for language models, zero point nine five and zero point nine eight, and the comment says the search had found zero point nine and zero point nine nine. Its weight decay is one, ten times AdamW's, as the paper's language models used learning rates ten times smaller, and the decay is the learning rate times lambda. The step loops over the groups and their weights, like AdamW's, and the state is one number per weight, the momentum, half what AdamW keeps. The update blends the momentum and the gradient with beta one, and takes the sign, which is plus one, minus one, or zero for a zero. Then the decay, and the step, which is the learning rate times the sign. And the momentum is updated after the step, with beta two, the slower one, as the comment says. In bake off dot py, the import gets Lion, the names get lion, and make optimizers gives it the two groups, as it does AdamW. Its test makes Lion with no weight decay, so that the decay doesn't move the weights, and takes three steps, one at a time. Before each, it copies the weights, and after it, it measures how far each one moved. Each moved by the learning rate, up or down, as the comment says, or not at all, where the gradient and the momentum are both zero.

20:36 Lion's test passes, and so does the bake off's for Lion. Its sweep went down to lower learning rates, as the paper advised. Lion's best is at three times ten to the minus four, eight times lower than AdamW's, at about two point zero seven. That's above AdamW's best, by about zero point one four. I first ran Lion with the betas its search found, zero point nine and zero point nine nine, and at its best it reached two point one seven. So a setting left at its default cost Lion about zero point one, about as much as running it at half its best learning rate. Around its best, at four point two times ten to the minus four, it's about two point one zero. At half its best, it's two point one eight, and at twice it, two point two one. And Lion is below AdamW at the learning rate pretrain dot py uses, which reached two point one three. So against the AdamW we've been training with, Lion would look like the better optimizer, with half the memory, and against a tuned AdamW, it isn't.

Sophia

21:39 The second challenger is Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training, by Hong Liu and others at Stanford, also from 2023. Its abstract says Sophia divides the moving average of the gradients by a moving average of an estimate of the curvature, and clips each weight's update. It estimates the curvature once in a handful of steps, so that it costs little, and on GPT models from a hundred and twenty-five million to one and a half billion parameters, it reaches Adam's loss in half the steps. Figure one shows it: the steps each needs to reach a validation loss, with Sophia at about half of AdamW, for three sizes of model. Algorithm three is Sophia, written out step by step. The momentum, m, is updated at each step, and the curvature, h, once in k steps, with one of two estimators. The weights decay, and the update is m over gamma times h, clipped to one.

22:40 Newton's method divides each gradient by the curvature, h, the second derivative of the loss, which says how quickly the gradient changes as the weight moves. Where the loss curves sharply, the step is small, and where it's flat, the step is large, which is what a learning rate can't do on its own for all the weights at once. Sophia does that with the momentum, over rho, B and h, where B is the number of tokens in a batch, and rho a setting, zero point zero five. And it clips the ratio at one, so no weight moves more than the learning rate. That keeps it safe where h is tiny or out of date, as the loss of a network isn't a bowl, and its curvature changes as the weights move. Where the ratio hits one, the step is the sign of the momentum, like Lion's. For h, SophiaG draws labels from the model's own predictions, y hat, and squares the gradient of the loss on them. Averaged over draws, and times B, that gives the diagonal of the Gauss-Newton matrix, a stand-in for the curvature that can't be negative. It does that once in ten steps, so it costs a forward and a backward pass more, a tenth of a step's work on average.

23:50 In optimizers dot py, the functional module comes in with the imports, for the softmax and the cross entropy that Sophia's estimate needs. Sophia goes under Lion, and its comment says it's SophiaG: the momentum divided by an estimate of the curvature, the loss's second derivative for each weight, with each weight's step clipped to the learning rate. Its defaults are the settings of its code for GPT-2 small: betas of zero point nine six five and zero point nine nine, and rho of zero point zero five. The weight decay is zero point two, twice AdamW's, as its README advises. Tokens is a batch's tokens, which the estimate is scaled by, as the comment says. State of makes a weight's state the first time it's asked for: the momentum, m, and the curvature's running average, h, two numbers per weight, like AdamW. It's a function of its own, as the step and the estimate can each be the first to need it. Update curvature runs after a backward pass on labels the model drew from its own predictions, as its comment says. It adds each weight's gradient squared into h's running average, with beta two, the way Adam's v averages the gradient squared, with a different gradient. Clipped counts the share of the weights whose step is clipped to the learning rate: those where the momentum's size is at least rho times the tokens times h. The paper's code logs the rest, the share that isn't clipped, to tune rho, and we're going to print it as we train. The step updates the momentum with beta one, and decays the weights, like AdamW. Then the ratio, and the comment over it says what it is: the momentum over the curvature, a Newton step, clipped at one, so no weight moves more than the learning rate. Clamp, with a max of one, does the capping. The tiny number added keeps the division safe where h is zero, before the first estimate, and there the ratio is one. So the weight moves by the learning rate, times the sign of the momentum, times the ratio. Curvature is the function that makes the estimate, and its comment names the estimator, Gauss-Newton-Bartlett. The model reads the batch, and multinomial draws a label for each of its positions from the softmax of its logits, with a generator we pass in. The loss on those labels goes backward, and the paper's code clips that gradient too, so we do. Then update curvature takes the squared gradients, and the gradients are zeroed, so the next step starts from its own. In bake off dot py, Sophia and curvature come in with the import, and the names get sophia. Curvature every is ten, as in its paper. Make optimizers takes a batch's tokens now, forty-eight times two hundred and fifty-six, and gives Sophia the two groups. Sophia's labels come from a generator of their own, seeded with the run's seed, so drawing them doesn't change the batches. At each evaluation, a Sophia run keeps the share of clipped steps, and prints it. Its comment says why we watch it: Sophia's README tunes rho so that between fifty and ninety percent of the steps are clipped. And once in ten steps, after the optimizer's step, curvature runs on that step's batch, before the clock is read, as it's the optimizer's work. Then Sophia gets three tests of its own. The first builds the small model and a Sophia with no weight decay, and tokens set to a hundred and twenty-eight, the four rows of thirty-two tokens of a test batch. For five steps, it copies the weights, takes a step, runs an estimate on a random batch, and checks that no weight moved further than the learning rate, which is what the clipping is for. The second sets up a single weight, with a curvature so large that its step isn't clipped, as its comment says. With a beta one of zero point five, the momentum after one gradient of two is one, and the step has to be one over rho times the tokens times h, the Newton step. And the third runs one estimate on the small model, and checks that the curvature it leaves is positive, and not all zero.

27:54 Sophia's three tests pass, and so does the bake off's. Its sweep found its best at three times ten to the minus four, eight times lower than AdamW's. There, it reaches about two point zero eight, above AdamW's best by about zero point one five. That's about level with Lion at its best, two point zero seven, as close as their seeds land to each other. Like Lion, Sophia is below AdamW at the learning rate pretrain dot py uses, and above AdamW tuned. At four times its best rate, Sophia learns much more slowly, and ends at about two point nine five after a thousand steps.

28:35 For Sophia's best run, the share of the weights whose step was clipped, at each evaluation. Early on, the curvature's running average is still close to zero, so almost all the steps are clipped, and Sophia moves each weight by the learning rate, like Lion. Its README tunes rho to keep this share between fifty and ninety percent, and ours comes down below ninety by step three hundred. It ends at about eighty-seven percent, so even then, most of the weights take the sign of their momentum, and the curvature sets the step of the rest. The paper's runs are a hundred times longer than ours, with ten thousand estimates of the curvature to our hundred, so its running average has much more to go on.

Muon

29:18 The third challenger came out as a blog post, not as a paper. It's by Keller Jordan, Muon: An optimizer for hidden layers in neural networks, from December 2024. Its definition is momentum, followed by a Newton-Schulz iteration on the result, for each weight matrix of the hidden layers. Its public code uses Nesterov momentum, which the post says works a bit better, and so do we. And the iteration is a few lines of PyTorch: five steps, with three coefficients, a, b and c. Its results are training speed records: CIFAR-10 to ninety-four percent, and the NanoGPT speedrun, a race to a validation loss of three point two eight on FineWeb, which Muon sped up by a factor of one point three five. And on why it works, the post observes that the updates of SGD with momentum and of Adam, for a transformer's weight matrices, are almost low-rank: a few directions dominate them. Orthogonalizing, it speculates, raises the scale of the rare directions, which are small in the update, and still important for learning.

30:21 Any matrix, M, splits into U, sigma and V transposed: two sets of directions, and in sigma, a length for each pair of them, the singular values. A momentum dominated by a few directions has a few large singular values and many small ones. Orthogonalizing keeps the directions and sets all the lengths to one, so the update moves the weights as much along each direction. Splitting the matrix that way at each step would take too long, so Muon gets close to it with Newton-Schulz: five steps of a few matrix products. On each singular value, s, a step is a polynomial with Keller Jordan's coefficients. It pushes small values up quickly, about three and a half times a step, and keeps the values near one between about zero point seven and one point two, while the directions don't change. Last comes the scale of the update, for each matrix. An orthogonal update of an A by B matrix is small in each entry, so Moonshot AI's paper on Muon at scale multiplies it by zero point two times the square root of the larger side. Its entries are then about the size of AdamW's, so the learning rate and the weight decay can be AdamW's too.

31:29 In optimizers dot py, math comes in with the imports, for a square root. Muon starts with orthogonalize, the Newton-Schulz iteration, and its comment says what it does: G's singular values pushed towards one, its singular vectors kept. A, b and c are Keller Jordan's coefficients. Dividing G by its norm, the square root of the sum of its entries squared, brings all its singular values to one or below, which the iteration needs. A tall matrix is turned on its side first, so that X times X transposed is the smaller of the two products, and the comment says that too. Each step computes A, X times X transposed, and replaces X with a times X, plus b times A plus c times A squared, all times X. On each singular value, that's the polynomial we saw, and each step is three matrix products, which the GPU does quickly, where a full decomposition would take far longer. And at the end, a tall matrix is turned back. Muon's comment: for the blocks' weight matrices, the momentum, orthogonalized, so the step moves the weights as much in all directions. Muon subclasses PyTorch's optimizer too, like the others. Its defaults are a momentum of zero point nine five, as in the post, and a weight decay of zero point one, like AdamW's, as Moonshot's scale lets the two share their settings. Its step loops over the weights like the others, and its state is the momentum, one number per weight, a running average like Adam's m. The update is Nesterov momentum: the gradient blended once more with the momentum, which, as the comment says, looks a step ahead, and the result is orthogonalized. The weights decay, and the step is scaled by zero point two times the square root of the matrix's larger side. The comment says why: an update the size of AdamW's, so the learning rate and the weight decay mean what they do there. In bake off dot py, Muon comes in with the import and the names, and in make optimizers, Muon gets the blocks' weight matrices. The comment says where the rest goes: the embedding, which is also the head, and the biases and the gains stay with AdamW, at Muon's learning rate. The post says that too, as Muon is for the hidden layers. So a Muon run has two optimizers, which is why train in our tests takes a list. Then Muon gets two tests of its own. The first builds a matrix out of known directions, U and V, which QR, a factorisation, gives us as orthogonal matrices from random ones. In between, the singular values span from zero point one to ten, a hundred times apart, as the comment says. In those directions, orthogonalize's result is diagonal, so it kept G's directions, and the values on the diagonal are between zero point six and one point two. The second trains two copies of the small model, one with our Muon and one with PyTorch's, which has Moonshot's scale as an option. PyTorch's runs the iteration in bfloat16, and ours in float32, as the comment says. So we check that each matrix moved in the direction PyTorch's did, to a cosine of zero point nine nine, and by as much, to within five percent.

34:42 Muon's tests pass: orthogonalize keeps the directions, and our Muon steps where PyTorch's does. Then bake off report dot py prints the four sweeps side by side. Muon's best is at four point eight times ten to the minus three, twice AdamW's, at about one point seven six, about zero point two below AdamW's best. That's about the size of the gap that tuning AdamW gave us, and in steps, it's worth more, as we're going to see. And Muon barely minds its learning rate: about one point seven six at three point four times ten to the minus three, and one point seven nine at twice its best.

35:22 On a chart, with the learning rate as a power of ten, so that minus three is ten to the minus three. Lion and Sophia want the lowest rates, and Sophia's curve is the steepest past its best. AdamW's best is near minus two point six, and it rises on both sides. Muon is lower than the others at each rate we tried, and flatter, so a learning rate picked roughly is close to its best.

The results

35:48 To compare the four at their best, bake off report dot py gets a second table. Math and statistics come in with the imports, for a logarithm and for averages. Best lr picks the learning rate with the lowest held-out loss in an optimizer's sweep. Seeds finds the full-length runs at a learning rate, one per seed, as its comment says, so the table can show how far they land from each other. AdamW steps turns a loss into the number of steps AdamW needs to reach it, which is how we're going to measure a speedup. It takes AdamW's runs at its best rate, with seed zero, one for each number of steps we ran, and finds the two whose losses are either side of the one we're after. Between them, it reads the steps off a straight line in the logarithm of the steps, as the loss falls smoothly with the logarithm of the steps, so a line between two neighbouring runs stays close to it. And it gives None outside the runs we have, as the comment says, as we don't guess past them. Main becomes print sweep, which takes the runs it's given, as the new table does. Print best prints a header, then a row for each optimizer. It finds the optimizer's best rate, its runs there with each seed, and the one with seed zero. The row has the best rate, the held-out loss with seed zero, the mean over the seeds, and the spread, which the comment defines: the highest held-out loss minus the lowest. Then the steps AdamW needs to reach seed zero's loss. The comment says why seed zero: with it, all the optimizers start from one model and read the batches in one order, so their runs differ by the optimizer alone, and AdamW's own row is its thousand steps. The time a step takes doesn't depend on the learning rate, so the table uses the median over all the optimizer's runs. Its comment says why: other work on the laptop slows a run down now and then, and the median, the middle value, isn't moved by a slow run. The row ends with the milliseconds a step, the optimizer's part of them, and the numbers kept per weight. On the command line, dash dash best prints the new table.

37:56 And the second table, with each optimizer at its best rate, for a thousand steps. The three seeds of each land within about a hundredth of each other, while the optimizers are about a tenth or more apart, so their order isn't down to chance. AdamW is the measuring stick, at its thousand steps. Reaching Muon's loss would take AdamW about one thousand eight hundred and forty steps, Sophia's about six hundred and eighty, and Lion's about seven hundred. Then the time, where the optimizers' parts of a step are what differ. AdamW takes about eleven milliseconds of each step, Lion about eleven, Sophia about fifty-eight, with its estimate, and Muon about seventy-two, with its iteration. The whole steps move by tens of milliseconds from run to run, with whatever else the laptop is doing. And the last column is the numbers each keeps per weight.

38:49 Along the way, each at its best rate, with seed zero. Muon is ahead from the first evaluation to the last, and AdamW is second throughout. Sophia and Lion start further behind, with their steps clipped or signed, and close part of the gap as the learning rate comes down.

39:08 To turn losses into steps, AdamW trained at its best rate for other lengths too, from five hundred to two thousand steps, each with its own warmup and cosine. The longer the run, the lower the loss, and two thousand steps read the stories about one point two times. Muon's loss after a thousand steps is the one AdamW reaches in about eighteen hundred, so here, Muon does in a thousand steps what AdamW needs about one point eight times as many for. The other optimizers' losses turn into steps like this too, between the two runs of AdamW around each of them.

39:44 So here's all of it, counted in AdamW's steps. AdamW at the rate pretrain dot py uses is worth about six hundred and ten of its tuned steps, so tuning the learning rate saved about one point six times the steps on its own. Against that AdamW, Sophia would look about one point one times as fast, and Lion about as much, and it would be easy to stop there. Against a tuned AdamW, both are worth fewer steps than its own thousand. And Muon is faster than either, at about one point eight times the tuned AdamW's steps.

40:16 Then comes what each of them costs, starting with memory. Our model has about thirteen point eight million weights, and each number an optimizer keeps for each weight takes four bytes, about fifty-five megabytes. AdamW and Sophia keep two, about a hundred and ten megabytes. Muon keeps one for the matrices and AdamW's two for the embedding, one point two three on average, and Lion keeps one. For a model of seven billion weights, with its state in float32, that's fifty-six gigabytes for AdamW against twenty-eight for Lion, the saving Lion's abstract points to. As for time, Keller Jordan's post works out that Muon's iteration adds at most T m over B to a step: five steps, times the embedding size, over the tokens a step reads, so about sixteen percent for us. On the laptop, Muon's part of a step took about seventy-two milliseconds of five hundred and forty, and AdamW's about eleven of four hundred and sixty-two. Sophia's estimate costs a forward and a backward pass more once in ten steps, so its part of a step comes to about fifty-eight milliseconds. So in time to reach a loss, Muon is still about one point six times as fast as AdamW here.

41:29 The docs get a page on optimizers, with AdamW and the three challengers, each with its paper and its update, and the rules of the race. It has the results too: each optimizer's sweep, its best rate with its seeds, the steps AdamW needs to match it, and what each costs in time and memory. The page goes in the list on the docs' home page, and at the end of the navigation. And the README has the commands, to run a sweep, another seed or another length, to print the two tables, and to run the tests.

What's real, and what's scaled down

42:00 A study from 2025 ran this kind of race at a larger scale: Fantastic Pretraining Optimizers and Where to Find Them, by Kaiyue Wen and others at Stanford. Its abstract says AdamW has stayed the optimizer of pretraining, despite claims of speedups of one point four to two times. They tuned ten optimizers with care, on models from a hundred million to one point two billion parameters, and found the speedups lower than claimed, falling to one point one times at one point two billion. The fastest, like Muon and Soap, multiply the gradient by matrices, where AdamW, Lion and Sophia scale each weight on its own. Figure one, at the top left, shows the baseline problem: for a model of a hundred million parameters, in a recipe from GPT-3, tuning the learning rate alone gives up to a two times speedup. And on Sophia, the paper notes that its original code base didn't fully shuffle the data, which lowered AdamW's best learning rate there. In their setting, Sophia doesn't offer a significant speedup over AdamW for models under half a billion parameters.

43:03 Andrej Karpathy's nanochat trains its model with Muon too. Its set up optimizer splits the weights like we do: the matrices of the blocks go to Muon, at a learning rate of zero point zero two, and the embeddings, the head and a few scalars go to AdamW. Each of those AdamW groups has a learning rate and betas of its own, and the embeddings' and the head's rates are scaled with the width of the model. Our race gave each optimizer a single learning rate to sweep. So nanochat sets more knobs than our race does, each tuned for its model.

43:39 The method is the one the study used: tune each optimizer, then compare them at their best, at the end of training. Their models went from a hundred million to one point two billion parameters, trained on one to eight times the tokens the Chinchilla paper found best for their size, about twenty a weight, with many settings tuned for each optimizer. Ours is a model of about thirteen point eight million weights, trained on twelve point three million tokens, less than one a weight, so far shorter than theirs, and we tuned one setting, the learning rate. Their study found Muon's speedup shrinks as models grow, from about one point three times at a hundred million parameters to one point one at one point two billion. Our model is ten times smaller than their smallest, where their trend gives Muon the most, and our runs are far shorter than any of theirs. So I wouldn't expect our number to carry over to larger models.

Wrap-up

44:29 So, today we wrote AdamW from scratch, checked it against PyTorch's, and raced it against three newer optimizers, each with its own sweep of the learning rate. AdamW at the rate we've trained with since episode two was about zero point two of loss behind AdamW at its best, for a run this long. Lion and Sophia came in behind a tuned AdamW, though both beat the AdamW we had, and Lion needed the betas its paper used for language models to get there. Muon was ahead at each learning rate we tried, and matched in a thousand steps what AdamW needs about eighteen hundred for. It keeps less than AdamW, and its part of a step takes about seventy-two milliseconds on the laptop.

45:12 Let me know in the comments which optimizer you train with, and whether you've seen Muon's gain hold up on your own models. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.