Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining".

Jane: The gist The Monarch–Muon optimizer introduces a structured matrix-aware optimizer that reduces weight, gradient,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at the paper titled "Monarch-Muon: Structure-Aware Muon for Fast and Memory-Efficient Pretraining." It sounds like they're tackling the problem of how to train massive models without running out of memory or taking forever.

Jane: That’s right, Tom. It focuses on a structured matrix-aware optimizer called Monarch–Muon that aims to reduce how much weight, gradient, and optimizer state you have to store during training.

Lu: Essentially, they parameterize selected weights using Monarch factors and replace one big dense orthogonalization step with several smaller independent computations instead. That should inherently lower the storage requirements for the model itself >

Meng: From a deployment standpoint, if you can cut down on optimizer state memory, that’s a real win for running things at scale. I'm curious how much of the total weight it actually saves when you look at real training runs >

Lalam: The paper introduces Monarch–Muon as a structured matrix-aware optimizer built on the Monarch parameterization which reduces weight, gradient, and optimizer-state storage and replaces dense orthogonalization with smaller independent computations #pg3

Tom: So the core idea is that for a weight W in R dout by din you define it as Wb(theta) = P1L(theta)P2R(theta) where theta contains K blocks of both factors >

Jane: That factorization means instead of one big dense matrix operation, the forward pass applies R, P2, L, and P1 successively without assembling a dense weight >

Lu: And each block gets its own momentum buffer initialized to zero so the gradient for each block is accumulated into its own momentum >

Meng: I wonder how this affects the actual training speed. If you're doing independent polar computations instead of one dense step, does that translate to faster wall-clock time? >

Tom: Yeah, they do. They approximate the polar factor of this momentum using s Newton–Schulz iterations and update the block >

Jane: That leads into the theoretical analysis where they prove convergence for exact-polar Monarch–Muon under generalized smoothness >

Lu: The proof involves decomposing the minimization problem into independent block subproblems where each block solves a spectral-norm subproblem >

Title and authors: Meng: Okay, so mathematically, they show that for every T greater than or equal to one the expected sum of the gradient norms is bounded by O(one + logT)T-one/four >

Tom: That convergence result shows that for every T greater than or equal to one E

sum over j nabla Xj F(theta k) two: is O(one + logT)T-one/four >

Jane: Experimentally, they found that Monarch–Muon improves validation loss over Monarch–AdamW in both dense and MoE pretraining at fixed Monarch structure #pg2

Meng: And the systems measurements show that for a 6 point 89B parameter system, it reduces peak memory from sixty-eight point nine GB to fifty-two point nine GB and Muon optimizer state from fourteen point six zero to six point seven seven GB >

Tom: That’s a significant reduction in optimizer state overhead, and they also report the Muon optimizer step goes from one thousand two hundred eighty-eight milliseconds down to just three hundred four milliseconds >

Jane: That computational saving grows with matrix size, which makes this approach particularly effective at larger scales >

Lu: They also studied initialization strategies for dense-to-Monarch transfer and found that closed-form Monarch projection minimizes weight-reconstruction error for a fixed factorization #pg1

Meng: And they found that joint sensitivity guides conversion, showing replacing FFNs in layers two through twenty-five together produces a KL divergence of seven point four to sixteen point seven times the sum of their isolated effects >

Tom: They also found that activation-aware fitting reduces the layer-output error by forty percent and closes seventy-six percent of the pre-SFT NLL gap to the dense reference >

Jane: They also looked at low precision training and found it complements structured training really well, as FP8 GEMM compresses computation and activations >

Meng: And FP8 states reduce stored optimizer buffers by three point seven times, which lowers the final-step peak by about zero point nine five GB for Monarch–Muon > <ref:2610.17429#pg1>

Tom: So they see two complementary uses of low precision: reducing the larger activation footprint and further compressing the already smaller optimizer state >

Lu: The paper notes that their method doesn't explicitly guarantee convergence unless generalized smoothness holds, which is an assumption on the coupled factor objective F(L, R) #pg2

Title and authors: Jane: They did mention that for dense-to-Monarch conversion, joint sensitivity guides conversion and activation-aware fitting helps close a large gap to the dense reference model >

Meng: What does this mean for us in production? If we want to deploy something big, do we need to worry about all this matrix structure setup or is it just a research tool >

Tom: It seems like it's moving toward being more practical because the computational savings grow with matrix size, making it particularly effective at larger scales >

Jane: So, when you look at the Monarch–Muon paper, you see they combine a structured representation with blockwise matrix-aware updates that reduces storage and replaces dense orthogonalization with smaller independent computations #pg3

Lu: That's the main contribution: the Monarch–Muon optimizer which reduces weight, gradient, and optimizer-state storage and replaces dense orthogonalization with smaller independent computations >

Meng: If you look at their experimental validation, Monarch–Muon improves validation loss over Monarch–AdamW in dense and MoE pretraining at fixed structure >

Tom: And they match dense AdamW performance in MoE with fewer parameters, specifically the two-block MoE variant matches dense AdamW loss using fewer parameters >

Jane: So for anyone who only listens to the show, what does this mean is that we are looking at a way to make training much more memory efficient while keeping the quality high >

Lu: It suggests that future work might focus on sensitivity-guided block allocation and batched low-precision kernels because the efficiency depends heavily on matrix shape and kernel execution #pg3

Meng: I'm still focused on what breaks at scale. This paper gives us a good roadmap for how to structure things to save memory, but we'll need real traffic data to see if this holds up when things get messy >

Tom: And they also showed that dense-checkpoint conversion benefits from jointly selecting components and initializing factors to preserve layer outputs >

Jane: So, the overall implication of Monarch–Muon is that we get a structured representation with blockwise matrix-aware updates that reduces storage and replaces dense orthogonalization with smaller independent computations #pg3

The paper's summary: Tom: So, to recap, Monarch–Muon is this new optimizer that uses a structured way of looking at weights—these Monarch factors—to make things much smaller and faster during training.

Jane: Right. It’s about replacing dense math with smaller pieces and doing the updates in a blockwise way instead of one big step. It’s designed to cut down on the memory needed for the weights, gradients, and even the optimizer state itself.

Lu: The core mechanism is parameterization: you take a weight W and express it as Wb(theta) = P1L(theta)P2R(theta), where theta holds these trainable blocks. This lets you apply the factors successively without assembling one giant dense matrix, which is what saves the storage space.

Meng: So, if I'm running a six point eight nine billion parameter model, and this setup cuts the peak memory from sixty-eight point nine gigabytes down to fifty-two point nine gigabytes, that’s a real thing for deployment.

Tom: It is huge because it’s not just about storage; they showed that this approach actually improves the validation loss compared to Monarch–AdamW when you keep the Monarch structure fixed during training.

Jane: That means if you use this setup, you get better model quality for less effort, especially in those dense and Mixture-of-Experts pretraining scenarios. And they found it can match the performance of a dense AdamW system even in MoE models using fewer parameters—specifically, the two-block MoE configuration did that.

Lu: The theoretical analysis is interesting too because they proved convergence for this exact-polar Monarch–Muon under generalized smoothness, showing the gradient norm bound is O of one plus log T times one to over four times T to the negative one over four.

Meng: That’s a number we can look at. It tells me that even though we’re using these smaller computations, the math is stable enough that it still converges properly, which means I don't have to worry about it blowing up during long training runs.

Tom: And they also looked at how to get old models into this new structure—the dense-to-Monarch transfer—and found that using joint sensitivity guides and activation-aware fitting helps you close that gap with the original dense model pretty well.

Jane: It’s a systematic way to move from the old format to this new structured one, and it seems like careful initialization can actually help you maintain quality when doing that conversion.

Lu: They also showed how low precision training, like using FP8 GEMM, complements this structured training really well because it compresses both the activations and the optimizer buffers by a factor of three point seven times.

Meng: So you get two ways to save space: you structure the model itself with Monarch factors, and then you compress every single number with low precision. That’s a lot of compression layered on top of each other.

Tom: And they found that this combination cuts the full-step training time by about fourteen percent for Monarch–Muon, which is a solid saving when you factor in the complexity reduction from the blockwise updates.

Jane: So, if you’re someone who only listens to this show, what does it change? It means for large models and those who care about running costs that you can get better validation loss without sacrificing quality.

Lu: The paper suggests that future work should probably focus on sensitivity-guided block allocation and batched low-precision kernels because the performance really depends on the matrix shape of your model.

Meng: Yeah, I’m focused on what breaks at scale. This gives us a solid blueprint for structuring things to save memory, but we still need real traffic data to see if it holds up when things get messy in production environments.

The paper's improvements: Tom: So, we're talking about what they suggest as improvements to the Monarch–Muon approach itself, beyond just using it for pretraining models already built on that structure.

Jane: Right. They look at how you can make the whole system even more efficient by focusing on how you set up those Monarch factors and which ones you pick for updates.

Lu: One improvement they highlight is a systematic way to handle that dense-to-Monarch conversion, because they found that just using closed-form weight projection isn't always the best path.

Meng: What does that mean for me? It means if I have an old model, I don't just blindly project it; I can use joint sensitivity guides to figure out which parts of the network are most important to convert first.

Tom: They found that replacing those layers two through twenty-five together results in a KL divergence of seven point four to sixteen point seven times compared to doing them layer by layer separately. That’s a big number for model quality preservation.

Jane: That suggests that you don't have to overhaul the whole network at once; you can target specific groups of layers and use activation-aware fitting to get closer to the original dense reference loss, which they found closes seventy-six percent of that gap.

Lu: The other improvement is related to scaling up—the benefit really grows with matrix width because the cheaper orthogonalization operations become increasingly cheap compared to the forward and backward passes.

Meng: So if I’m training a much wider model, this method gets better just by being bigger, which is good because it means you can keep getting these speedups without needing a completely different architecture.

Tom: They also looked at how low precision training fits in here and said that it complements structured training very well because FP8 GEMM compresses both the computation and the activations by about fourteen percent for Monarch–Muon.

Jane: And they found that those FP8 states reduce the stored optimizer buffers by three point seven times, which lowers the peak memory usage by about zero point nine five gigabytes for Monarch–Muon.

Meng: That’s a solid number because it shows two layers of compression: structuring the math and then compressing every single number with low precision. It really makes sense for production where memory is tight.

Lu: They also pointed out that they need to be careful about how you allocate those blocks, suggesting that sensitivity-guided block allocation might be the next big thing for making this work even better across different matrix shapes.

Tom: So, the implication here is that Monarch–Muon isn't just a quick trick; it’s a framework where you can systematically improve model quality during conversion and scaling by being smart about how you structure those blocks.

Conclusion: Tom: So, to wrap up, Monarch–Muon is this structured optimizer that uses blockwise matrix updates to slash memory usage and speed up training for large AI models like Monarch or Muon architectures.

Jane: It’s a method for making huge AI models more efficient by using these specific structural parameters to reduce the storage needed for weights and gradients during training.

Lu: The big implication is that you can achieve better validation loss than other methods while having lower memory footprint, especially in dense and MoE pretraining setups.

Meng: For production, it means we’re looking at a way to run much larger models on the same hardware because the memory overhead is so drastically reduced.

Lalam: I'm confident in saying that Monarch–Muon provides a structured representation with blockwise matrix-aware updates that really cuts down on storage and replaces dense orthogonalization with smaller independent computations.

Tom: It’s solid work because they’re showing concrete numbers, like reducing the optimizer state from fourteen point six zero to six point seven seven gigabytes on a six point eight nine billion parameter system.

Jane: And they showed that low precision training, like FP8 GEMM, works with it to further compress those states by three point seven times.

Lu: The work also provides a systematic approach for dense-to-Monarch conversion using sensitivity guides and activation-aware fitting to help preserve quality when moving from old models to this new structure.

Meng: That helps us figure out how to deploy these optimized architectures without losing the performance we’ve already worked hard to achieve.

Lalam: It really shows that efficiency isn't just about one thing; it’s about combining structural choices with precision training for maximum savings.

Tom: So, that’s the core idea behind Monarch–Muon: a structured representation with blockwise matrix-aware updates that reduces storage and replaces dense orthogonalization with smaller independent computations.

Jane: It gives us a way to build and train bigger, better AI models without immediately running into memory walls.

Lu: The future direction seems to be focusing on sensitivity-guided block allocation because the efficiency depends heavily on those matrix shapes and kernel execution details.

Meng: I’ll keep an eye on that. If we can get a more scalable way to allocate those blocks, this becomes much more practical for our deployment pipeline.

Lalam: It’s exciting stuff because it shows how combining different techniques can lead to tangible gains in efficiency for complex models like the Monarch and Muon series.