Can Muon, the optimizer behind Kimi K2, work for SDXL and Anima LoRA training?

I keep baking character LoRAs for Anima and WAI-Illustrious, and I have never once picked the optimizer myself — every run has used whatever the training tool defaults to.
Meanwhile on the LLM side, Muon keeps getting adopted in frontier pretraining runs like Kimi K2 and DeepSeek V4.
Before writing any code, I checked how Muon works, where it can be applied in image models, and how it interacts with LoRA.

Muon orthogonalizes the update matrix

AdamW treats parameters one element at a time.
It keeps first and second gradient moments per element and adjusts the step size per element.
Where an element sits in the matrix, and how it lines up with its neighbors, has no effect on the update.

Muon (MomentUm Orthogonalized by Newton-Schulz) treats the update to a weight matrix as one whole matrix.
It accumulates the gradient matrix into momentum, pushes it toward an orthogonal matrix with the Newton–Schulz iteration (a sequence of pure matrix multiplications), and only then applies it to the weights.

flowchart TD
    G[Mini-batch gradient G] --> M[Accumulate into momentum]
    M --> NS[Orthogonalize via Newton–Schulz<br/>push update toward UVᵀ]
    NS --> S[Scale to layer shape]
    S --> W[Update weight matrix]

Orthogonalization evens out the singular values of the update matrix.
Muon’s author observed that raw update matrices tend to be dominated by a few directions, and orthogonalization cancels that skew.
The explanation that feeding step size into the weak directions is what improves training is, by the author’s own account, still a hypothesis.

Keller Jordan introduced it in 2024 for the NanoGPT speedrun, and Moonshot AI stress-tested it at scale in Muon is Scalable for LLM Training.
Since then, Kimi K2 trained on 15.5T tokens with the MuonClip variant, Kimi K3 adopted Per-Head Muon, and DeepSeek V4 uses Muon for most modules in pretraining.

Muon’s main targets are the 2D hidden-layer weight matrices; 4D convolutions can also be handled by reshaping them to 2D.
Embedding weights are 2D too, but the original setup keeps them on AdamW as an empirical choice, and 1D parameters like LayerNorm/RMSNorm scales and biases also stay on AdamW.
In other words, the design assumes from the start that some layers remain on AdamW.

The other optimizer that keeps showing up next to Muon is SOAP.
It is a Shampoo-family method: from the gradient matrix $G$ it maintains row-side $GG^\top$ and column-side $G^\top G$ statistics, and runs Adam in their eigenbasis.
It uses more preconditioning information than Muon, but the optimizer state costs correspondingly more memory.

	AdamW	Muon	SOAP
Update granularity	per element	per matrix	per element in eigenbasis
Target parameters	everything	mainly hidden-layer matrices	mainly matrices
Extra state	1st/2nd moments	momentum only	eigenbases + moments
Extra compute per step	none	Newton–Schulz iteration	periodic eigendecomposition

Since its state is just momentum, Muon’s optimizer memory is actually lighter than AdamW’s.

Anima and SDXL expose different layers to Muon

I started out assuming Anima was SDXL-based, and that premise was simply wrong — as my own first-look article says, Anima is built on NVIDIA’s Cosmos-Predict2 2B, a separate lineage from SDXL.
Cosmos-Predict2 is a DiT (Diffusion Transformer) stacking self-attention, cross-attention, and feed-forward blocks, with timestep conditioning injected through adaptive LayerNorm.

The Cosmos-Predict2 DiT block does contain layers an LLM never has, like cross-attention and the adaptive LayerNorm conditioning.
Still, the heavy weights — Q/K/V/output projections and feed-forward matrices — are 2D linear layers, exactly the shape Muon expects.
The LLM-style routing that leaves embeddings, norms, and biases on AdamW carries over as the initial plan.

SDXL models like WAI-Illustrious are different: the U-Net mixes in Conv2D.
A convolution weight is a 4D tensor $[C_{out}, C_{in}, k_h, k_w]$, so applying Muon means reshaping it to $[C_{out},\ C_{in} k_h k_w]$, orthogonalizing, and reshaping back.
Muon’s original implementation does exactly this for ConvNets, for every convolution except the first.

But being reshapeable to 2D and benefiting from having that matrix’s singular values evened out are two different things.
Once the two kernel axes and the input-channel axis are packed into one dimension, what orthogonalization is actually equalizing changes.
For SDXL I would start with Muon on attention and linear layers only, leaving convolutions on AdamW.

Here is the initial routing plan. Every row is a design guess, not something I have measured.

Parameter	Anima (DiT)	SDXL (U-Net)
Attention Q/K/V/output	Muon	Muon
Feed-forward / MLP	Muon	Muon
Timestep/condition projections	Muon candidate	Muon candidate
Conv2D	mostly absent	AdamW
Embeddings	AdamW	AdamW
Norms and biases	AdamW	AdamW

What measurements exist for diffusion and ViT

For diffusion-style objectives, Optimization Benchmark for Diffusion Models on Dynamical Systems has a same-conditions comparison.
The target is small diffusion/flow models learning dynamical-system trajectories, and after 1024 epochs Muon and SOAP landed 18% lower in final loss than AdamW.
On the other hand, one epoch takes about 1.45x longer with Muon and 1.72x longer with SOAP, and the optimal learning rate is roughly 2x AdamW’s.
This is not the same thing as image quality improving 18% on a Stable Diffusion-class model, but it is evidence that Muon-family methods have room to push diffusion losses lower.

On the vision side, Muon in Vision Transformers reports Muon beating AdamW on ViT classification.
The caveat: the gap widens with heavier data augmentation, so the effect depends on the combination with the training recipe.

How Muon interacts with LoRA’s two-factor split

My use case is LoRA training, and this is where Muon’s fit is least clear.

LoRA freezes the base weight $W$ and trains only the small low-rank pair $\Delta W = BA$.
$A$ and $B$ are each 2D matrices, so applying Muon to them individually is mechanically possible.
But what reaches the output is the product $\Delta W$, and it is not obvious that orthogonalizing the two factors separately produces a good update for the product.

LoRA meets Riemannion starts from LoRA’s parametrization ambiguity — infinitely many factorizations satisfy $BA = B'A'$ — and proposes Riemannion, which optimizes directly on the fixed-rank matrix manifold.
It generalizes Muon to the Riemannian setting and reports better convergence and final performance than plain LoRA on both LLMs and diffusion models.
The diffusion experiments teach Stable Diffusion 2 a specific subject from a few reference images, comparing LoRA and Riemannion at ranks 4/8/16 — the 400-step quantitative comparison is at rank 4, the per-rank visual comparison at 600 steps.
A setup built on a few reference images, ranks 4/8/16, and one specific subject overlaps heavily with character LoRA training.

Uniform Spectral Growth and Convergence of Muon in LoRA-Style Matrix Factorization is a theory-side result: even when Muon is applied to $A$ and $B$ separately, the singular values of the product $\Delta W$ grow nearly uniformly across the spectrum, shown in a simplified setting.
That is a partial answer to the concern above about the product update having no guarantee.

Still, I found no report testing exactly what I do: Anima-family character LoRAs baked from a few dozen images over a few thousand steps.
The only way to find out is to run it.

What swapping the optimizer can and cannot change

A LoRA training step is dominated by the model’s forward and backward passes, so swapping the optimizer should not shorten wall-clock training time.
Newton–Schulz adds per-step compute if anything, and in the diffusion benchmark above, Muon’s epochs took longer.

What I am hoping for is landing at a lower loss after the same number of steps.
The 18% figure is exactly that kind of number — final losses compared after the same 26.6K steps.
Whether any of it survives under image-LoRA conditions — rank 4–128, batch 1–4, a few thousand steps — is not something the literature answers.

Implementation plan

The training path I have in mind is AnimaLoraToolkit.
Back in April I described it as sd-scripts-based, but the current public source is its own anima_train.py (built on FHfanshu/Anima_Trainer), and the config YAML has no optimizer field at all.
sd-scripts itself can load an arbitrary optimizer class from a module path, but that route is unreachable from the Toolkit. The optimizer construction code in the Toolkit gets modified directly instead.

Putting the swap, the pairing, and the orthogonalization in all at once means not knowing which one broke. They go in one at a time.

First, read the Toolkit’s optimizer construction code and swap in a custom class that is still plain AdamW inside, just to confirm the path works.
Next, make lora_down and lora_up addressable as pairs. LoRA matrices of identical shape exist across many layers, so matching by shape or registration order risks pairing across layers — the layer name or an explicit pair record from LoRA injection time has to be carried along.
Last, add the orthogonalization. Not Newton–Schulz at first: compute the polar factor $UV^\top$ directly via SVD to confirm correctness, then swap in the iteration once it runs.

A minimal experiment needs only a handful of images, rank 4 or 8, and around 100 steps.
Quality is not the question yet. The checks are only that the loss runs to the end without diverging, the saved LoRA loads in ComfyUI, and the weights have moved from their initial values.
Once that passes, the next run is a few hundred steps against AdamW under identical conditions, comparing loss curves and wall-clock time.