f of x

Optimizers Part II: Muon

Continuing on our optimizers journey, this article covers Muon, a relatively newer optimizer that works by orthogonalizing 2D gradient matrices. Instead of thinking about each one of the gradients, this beast works on controlling the unevenness in a gradient matrix caused by stretching along different directions, ensuring uniformity.

That said, use this optimizer with 2D matrices only.

But what exactly do we mean by stretch? Why does it happen? What problems does it create? How can we fix it? And...How does Muon do it efficiently?

These are the questions that this article is gonna answer. I'm gonna give you a tour of the math behind this optimizer and by the end of it, most of it will make so much sense.

What do you mean by stretch and what's causing it?

Take a simple matrix multiplication -
y = Mx
M is a matrix and x is a vector.

As a result of this matrix multiplication, M can cause a stretch or compression in some direction of x. It might rotate or shear x too.

Direction? (just to give you clearer image)

If x is [1, 3, 0], a vector with 3 dimensions (say n_dim = 3), then 1 can be seen as the x coordinate, 3 is the position in y direction and 0 is the position in the direction z. This is what we mean by directions. A vector has n_dim (here, 3) different directions.

For instance, G can stretch input direction v1 by 3 and v2 by 1. After W <- W − nG, one axis of the layer is updated three times harder than the other.

Note: Those axes are usually not the coordinate axes.

Is there a problem with that? - The reason behind Muon

I don't want to make this part complicated. Let me explain this in simple terms.

The way we're viewing the gradient is kinda off. A larger gradient in one direction is not necessarily a bad signal, but we should not use it in updating W because the idea of gradient in a particular direction being a measure of the magnitude of update is what's actually wrong.

It's better to see it as a measure of activity in that area. A batch of the training dataset might lead to a specific region of the network going nuts while many regions will feel like they're sleeping. This is how we should look at the differences in different regions.

If we keep the stretching part, we get uneven updates in different regions and the effects accumulated over time surely converge (because that's how other optimizers are working) but it's slower and kinda sloppier. Muon focuses on the matrix as a whole and updates evenly in all directions for faster convergence and evenness from the start. This is true in the sense that Muon prevents loudness (in terms of activations in certain regions) from eating the optimization-budget. For more clarity: I've said "even" updates for the same thing.

If the mode is zero in a region, there's still no way to make it move. So, it's all in terms of making the update more about the whole matrix rather than a few specific directions.

This creates the need to remove the stretching part of the gradient matrix while keeping the important things that actually carry the useful signal.

How to fix the problem?

Simple. Throw away the stretching.
It would have been so great if the heavens bestowed a gift upon us that somehow helped us dismantle a matrix into different components and then we could've just kept the valuable things and kicked out the rest.

Well, what a day it is! We have that tool. It's called Singular Value Decomposition (SVD).

I'm not gonna explain every bit of it and I'm just gonna let you know how to use it by giving you a short manual and an example.

SVD let's us express G in the form of G = UΣV'

(I've used ' to mean transposed)

Note

||Gx||² = (Gx)'Gx = x'G'Gx This tells us that the stretching in Gx comes from G'G.

Therefore, the eigenvalues and eigenvectors SVD uses are of G'G and not of G. Let σi² be the ith largest eigenvalue of G'G.

The different components are then explained as -

U is a matrix of Gvi/σi where vi is the unit eigenvector (eigenvector/magnitude) of the ith largest eigenvalue σi². So σi is the root of ith largest eigenvalue.

If you don't get that, don't worry. It'll be a lot easier with an example where we're gonna manually perform SVD for a simple 2x2 matrix to wrap our heads around this topic.

Σ is a diagonal matrix of σi in descending order. It is the thing that we wanna throw away because it's the matrix containing the stretch factors.

Note

σi² is the value of ||Gvi||²/||vi||² so σi is the stretch for ||Gvi||. Hence, Σ is a matrix of σi.

V is a matrix whose columns are the unit eigenvectors in the same order. In SVD, we use its transpose.

If we replace Σ with I (identity matrix), we get Muon(G) = UV' (let this be called Q)

The stretching is gone. Some properties of this new matrix are:

If you don't understand the reason behind any one of these, you can note stuff down and search about it later. But as I've said before, in the scope of this article, SVD is only being used as a tool so we don't need to understand or prove everything in depth.

Example - using SVD manually

Let the matrix be A = [[1, 2], [2, 1]]

A = UΣV'

As A is symmetrical, A'A = A² = [[5, 4], [4, 5]]

You can then calculate the eigenvalues using the equation det(A'A - λI) = 0

Eigenvalues for A'A are 9 and 1. So, σ1 = 3 and σ2 = 1.

Eigenvectors are [1, 1] and [1, -1] in that order. Unit eigenvectors v1 and v2 are [1, 1]/sqrt(2) and [1, -1]/sqrt(2) respectively.

U is formed by Avi/σi

u1 = Av1/σ1 = [3, 3] / (3 x sqrt(2)) = [1, 1]/sqrt(2) u2 = Av2/σ2 = [-1, 1]/sqrt(2)

U = [u1, u2] where u1 and u2 are columns = [[1, -1], [1, 1]]/sqrt(2)

Σ = [[3, 0], [0, 1]]

V = [v1, v2] where v1 and v2 are columns = [[1, 1], [1, -1]]/sqrt(2)

Thus, Muon(A) = UV' = [[0, 2], [2, 0]]/2 = [[0, 1], [1, 0]]

For A, this is the orthogonalized matrix that we've obtained using SVD.

How does Muon do it? (Plot-twist)

SVD looks tedious, compute-heavy and just long. It's nice to have an intuition about things and so it is a good idea to learn about how to do it using SVD but now that we're in the computing territory, we need to think about efficiency.

Enter Newton-Schulz orthogonalization.

We're not gonna dive into the math here because that is definitely gonna be too much for someone who wants to make a Muon or understand how it works. All we need to know is the formula and what it does.

This method has different forms and I'll tell you the cubic one. It's this one line - X = 1.5 X - 0.5 (XX' X)

Take the matrix X through this n times (the number of iterations; say 5 for example) and you'll get a close to orthogonalized matrix after those iterations.

You can look at the quintic NS yourself if you're interested.

Note that here X = G/∥G∥F which is G scaled to Frobenius norm 1.

∥G∥F is the square root of sum of absolute squares of elements of G. ∥G∥F = sqrt(∑ij Gij²)

So, before the iteration, make sure to normalize the matrix G.

BTW, we didn't talk about the full form of Muon, did we?

MomentUm Orthogonalized by Newton-Schulz

And now we have a much more efficient way to do the same thing. An implementation using this formula is a pretty close one to the standard and I've made it in my educational ML library called lantern and it is documented too. This one also decouples L2 regularization from Muon, creating a separate MuonW optimizer (similar to AdamW). If you want to understand how Muon is implemented in Pytorch, lantern might come out useful.

Also, you'll see in the code that the G that we're orthogonalizing is actually the first moment buffer and not the gradient matrix at a particular iteration because that would be too noisy otherwise.

a note
At the time when I'm writing this, I'm 17 and there's a lot that I have to learn and a lot more that I have to create and write about. If you like my work, please show your appreciation by engaging with it and sharing it if you think you can help me get more reach on this because for some reason, the algorithm feels like an opposing force right now. The more feedbacks I receive, the more I'll get motivated to write. My goal is to always keep sharing resources and to create a nice and cozy community around these.

See ya in the next one!

Till then, make sure to gimme your feedback in the replies. DMs are open too.