Source-linked AI summary
Muon Is Theoretically Wrong For Convolutions, But Empirically Effective
Thibaut Boissin, Thomas Massena, Mathieu Serrurier, Franck Mamalet
TL;DR
Standard Muon reshapes convolution kernels into matrices, a shortcut that does not generally orthogonalize the convolution operator. The paper formulates the update in convolutional operator geometry and introduces Conv-NS, which preserves kernel support. Conv-NS and reshape-based Muon achieve comparable performance on CIFAR-10 and ImageNet, while stronger orthogonalization may overconstrain updates.
Problem
Standard Muon reshapes convolution kernels before orthogonalization, which does not generally orthogonalize the convolution operator and leaves its update misaligned with the layer's operator geometry.
Source: §1 · pp. 1–2, §2.2 · p. 3Method
Results
Takeaways & Limitations
Both approaches have limitations: reshaping does not enforce the convolutional operator-norm constraint, while Conv-NS targets a stronger orthogonality constraint, so extending Muon requires more than computing the convolutional polar factor.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Muon, an optimizer known for its efficiency, has a clear interpretation for matrix-valued updates, but convolutional kernels are stored as four-dimensional tensors. Standard implementations reshape these tensors into matrices, a shortcut which breaks the theoretical understanding behind Muon. To investigate this, we formalize the corresponding optimization objective directly in convolutional operator geometry and introduce Convolutional Newton-Schulz (Conv-NS), which approximates the polar factor in this geometry while preserving kernel support. When applied in fast training experiments, Conv-NS and reshape-based Muon are both computationally efficient and achieve comparable accuracy on CIFAR-10 and ImageNet classification tasks. However, as one could expect a theoretically aligned Conv-NS to outperform reshape-based Muon, we investigate this mismatch between practice and theoretical understanding, with the hypothesis that exact convolutional orthogonalization may overconstrain updates. These findings highlight Muon's strong practical performance while opening directions for its further development on convolutions. Our code is publicly available at \href{https://github.com/thib-s/muonconv-cifar10-airbench}{github conv-muon}.
1 INTRODUCTION
Standard Muon reshapes convolution kernels, but that does not generally orthogonalize the convolution operator. The paper introduces a convolution-aware alternative and finds that better operator-level orthogonalization does not consistently improve training.
1 INTRODUCTION
- Reshaping a convolution kernel before orthogonalization does not generally orthogonalize the equivalent linear operator, leaving standard Muon misaligned with convolutional operator geometry.The paper frames this as a mismatch between the kernel matrix surrogate and the convolution's actual spatial and channel interactions.Source: §1 · pp. 1–2
- Conv-NS and reshape-based Muon achieve comparable performance on CIFAR-10 and ImageNet, while both substantially outperforming momentum SGD in compute-constrained settings.Increasing Conv-NS iterations can degrade training performance.Source: §1 · p. 2
- Conv-NS targets the convolutional polar factor while preserving kernel support and empirically approaches convolutional orthogonality with near-minimal gradient modification.The paper presents it as a convolution-aware construction for addressing the mismatch.Source: §1 · pp. 1–2
- The paper revisits Muon's update principle as maximizing predicted loss reduction while limiting effects on layer outputs, a balance that reshaping does not enforce and convolutional orthogonality may constrain too strongly.The authors identify balancing these objectives as a direction for further convolutional Muon development.Source: §1 · p. 2
2 MUON FOUNDATIONS AND THE CONVOLUTION MISMATCH
Muon approximates a dense-layer gradient update with its polar factor, while standard convolutional implementations flatten kernels and apply that matrix procedure. This creates a convenient surrogate objective that does not capture convolutional operator geometry.
2.1 ABOUT THE MUON UPDATE RULE
- Muon approximately orthogonalizes a dense-layer gradient, using its polar factor as the update; a small number of Newton-Schulz iterations approximates that factor without computing an exact SVD at every step.The polar factor is framed as projection onto the Stiefel manifold.Source: §2.1 · p. 3
2.2 HOW STANDARD MUON HANDLES CONVOLUTIONS
- Standard convolutional Muon flattens each finite-support 4D kernel into a matrix, applies dense Newton-Schulz orthogonalization, then unflattens the result.Materializing the equivalent linear transformation is infeasible at scale.Source: §2.2 · p. 3
- This reshape-based procedure drives the flattened matrix toward semi-orthogonality, but ignores spatial and channel interactions and does not generally orthogonalize the convolution operator.It treats a 3 × 3 convolution with 8 input channels like a pointwise 1 × 1 convolution with 72 input channels when both have 8 output channels.Source: §2.2 · p. 3
3 OUR CORRECTION: THE CONVOLUTIONAL POLAR FACTOR
The convolutional polar factor adapts Muon's operator-level update to convolution kernels while requiring the resulting update to retain fixed kernel support. Directly orthogonalizing the induced operator is costly and generally loses that support.
3 OUR CORRECTION: THE CONVOLUTIONAL POLAR FACTOR
- The induced convolution operator is structured, but directly applying Newton-Schulz to it is costly and its polar factor generally cannot be converted back into a valid fixed-support kernel update.The operator is typically Toeplitz or doubly circulant, depending on boundary conditions.Source: §3 · pp. 3–4
- The convolutional orthogonality residual measures whether the induced operator is semi-orthogonal, with an equivalent kernel-level interpretation as a spatial self-convolution identity.The kernel-level distance is denoted Lorth.Source: §3 · p. 4
- The convolutional polar factor is the convolution-operator analogue of Muon's dense polar-factor update, but neither operator SVD nor reshape-based orthogonalization computes or approximates it.This motivates a convolution-aware construction.Source: §3 · p. 4
4 CONVOLUTIONAL NEWTON-SCHULZ: A CONVOLUTION-AWARE ALTERNATIVE
Conv-NS iteratively approximates convolutional orthogonality while keeping updates within the original kernel support. On real training gradients, it follows an estimated orthogonality–fidelity frontier, trading stronger orthogonality for greater gradient modification.
Convolution-aware construction
- Conv-NS applies order-3 Newton–Schulz steps followed by a structural projection that returns the kernel to its finite-support manifold.Without projection, the kernel support grows at each iteration.Source: §4 · p. 4
Empirical validation
- Conv-NS progressively concentrates singular values around one, whereas reshape-based RKO leaves them dispersed for the same random 3 × 3 kernel.The figure compares Conv-NS at increasing iteration counts with RKO.Source: Figure 3 · p. 5
Convolution-aware construction
- Block convolutions perform the Newton–Schulz scheme directly on kernels, avoiding materialization of the structured convolution matrices.The expanded result is cropped to the original kernel support.Source: §4 · pp. 4–5, Figure 2 · p. 5
Numerical reference
- The numerical reference estimates the orthogonality–fidelity frontier by trading convolutional orthogonality against fidelity to the original gradient, but does not certify global optimality.The objective is nonconvex, and the reference retains the best result across multiple trials.Source: §4 · p. 5
Empirical validation
- Across CIFAR-10 layers and three training stages, Conv-NS stays near the estimated frontier, with stronger orthogonality requiring larger departures from the gradient.Metrics are averaged across layers and stages; the frontier is numerically estimated.Source: Figure 4 · p. 6, §4 · p. 6
5 EXPERIMENTS: CIFAR-10 AND IMAGENET CLASSIFICATION
Experiments compare Conv-NS with reshape-based Muon and SGD on CIFAR-10 and ImageNet under short training budgets. Conv-NS essentially matches reshape-based Muon on CIFAR-10 and remains competitive on ImageNet, while both Muon variants outperform SGD in short-training settings.
Experimental setup
- The evaluation uses CIFAR-10 AirBench94 and 30-epoch ResNet-50 ImageNet-1k, plus an extreme one-epoch setting on both datasets.On CIFAR-10, Conv-NS receives a 1000-configuration tuning budget; on ImageNet, Conv-NS and reshape-based Muon each receive 48 configurations.Source: §5 · p. 6
Classification results
- Above 52% accuracy for both Muon variants versus 30.47% for SGD in one-epoch ImageNet shows their advantage in this extreme short-training setting.The reported gap narrows in the 8-epoch and 30-epoch settings as performance becomes increasingly capped by network capacity.Source: §5 · p. 6
- 93.98% vs. 94.02% on CIFAR-10 and 74.53% vs. 74.82% on ImageNet: Conv-NS is near parity with reshape-based Muon, which keeps a slight overall edge.Conv-NS adds marginal overhead on the small CIFAR-10 network and runs slightly faster than the reshape baseline on ResNet-50.Source: Table 1 · p. 7, §5 · p. 7
6 WHAT SHOULD CONVOLUTIONAL ORTHOGONALIZATION OPTIMIZE?
The section tests whether more accurate convolutional orthogonalization improves optimization and examines polar-factor orthogonalization against the convolutional LMO. More iterations can improve orthogonality without improving training, and the two update approaches have different constraint limitations.
Orthogonalization and training performance
- 87.56% validation accuracy is the one-epoch CIFAR-10 peak at 10 Conv-NS iterations; it falls to 86.42% at 25 and 85.13% at 100 iterations.Under training-loss selection, training loss also worsens at higher iteration counts.Source: §6 · p. 7, Table 2 · p. 7
Convolutional LMO
- The convolutional LMO constrains induced operators to the operator-norm unit ball, whereas ConvPF solves the LMO over the more constrained Stiefel set.ConvPF satisfies the unit-ball constraint by construction, unlike RKO.Source: §6 · p. 8
- Because the feasible-set inclusion is generally strict, more accurate ConvPF orthogonalization may enforce constraints beyond those required by the convolutional LMO.The paper proposes this as a possible explanation for increased gradient modification hindering early optimization.Source: §6 · pp. 8–9
Orthogonalization and training performance
- In eight-epoch CIFAR-10 runs, higher Conv-NS iteration counts lower both Lorth and Lflat, yet high counts can slow early learning without improving validation accuracy.The figure summarizes five runs and shows means with one-standard-deviation bands.Source: §6 · pp. 8–9, Figure 5 · p. 9
Convolutional LMO
- Reshape-based Muon does not enforce the convolutional operator-norm constraint, while Conv-NS targets stronger orthogonality; neither is guaranteed to solve the convolutional LMO.The paper concludes that a principled convolutional extension requires more than computing the convolutional polar factor.Source: §6 · p. 9
7 CONCLUSION
The paper formalizes convolutional orthogonalization through the induced operator and introduces Conv-NS, yet finds performance comparable to reshape-based Muon. It connects this gap to the constraints imposed by orthogonality and identifies evaluation beyond classification as future work.
7 CONCLUSION
- The standard reshape-based implementation does not generally orthogonalize the induced convolution operator.The paper frames its analysis around the geometry of that operator.Source: §7 · p. 9
- Conv-NS preserves kernel support and approaches the convolutional polar factor, but performs comparably to reshape-based Muon on CIFAR-10 and ImageNet.Increasing Conv-NS iterations could degrade performance, suggesting exact convolutional orthogonality may overconstrain updates.Source: §7 · p. 9
- The convolutional polar factor is optimal under the Stiefel constraint but not necessarily for the full operator-norm LMO.The paper suggests balancing operator-level control with sufficient freedom in convolutional updates.Source: §7 · p. 9
- Further evaluation could cover generative modeling, self-supervised representation learning, and non-2D convolutions.On classification gradients, reshape-based Muon appears to balance gradient fidelity and orthogonality; the paper notes this is not the case for synthetic gradients with different distributions.Source: §7 · p. 9
A.1 MUON AND MATRIX-WHITENING OPTIMIZATION
- Related work describes matrix-whitening optimizers including Muon and SOAP, while noting that optimizer scaling gains depend on parameter scale and hyperparameter tuning.The discussion also connects these optimizers to implicit bias and challenges in controlling weight drift.Source: §A.1 · p. 13
A.2 COMMUNITY EXPLORATIONS FOR CONVOLUTIONAL ORTHOGONALIZATION
- Community efforts include frequency-domain Freq-Muon and factorization-based Matrixizer, but the paper reports no peer-reviewed formalization or rigorous evaluation of these approaches.The authors describe both methods and provide appendix-only comparisons.Source: §A.2 · p. 13
B IMPLEMENTATION OF CONV-NS
Conv-NS performs Newton–Schulz-style updates directly on convolutional kernels while retaining their finite spatial support. Its implementation uses convolution operations, grouping support, and damping to manage stability.
B IMPLEMENTATION OF CONV-NS
- Conv-NS initializes with spectral rescaling and iteratively updates the kernel using convolution and transposed-convolution operations.The algorithm specifies N = 12, β = 0.5, α = 0.91, and ε = 10^-2.Source: B IMPLEMENTATION OF CONV-NS · p. 13
- Group packing and unpacking let the iteration process grouped convolutions while restoring the kernel's original layout.The convolution pair evaluates the cubic Newton–Schulz correction, with cropping implemented through padding choices.Source: B IMPLEMENTATION OF CONV-NS · p. 13
- Damping with α = 0.91 is intended to keep iterates within the operator-norm unit ball after support cropping.The unprojected polynomial preserves that ball for 0 < β ≤ 0.5, but cropping need not; the final division by α removes the last damping factor.Source: B IMPLEMENTATION OF CONV-NS · pp. 13–14
- Cropping preserves kernel support, giving each iteration fixed computational and memory cost for fixed kernel dimensions, stride, and group count.The loop uses real-valued kernels and avoids FFTs and activation-grid transforms.Source: B IMPLEMENTATION OF CONV-NS · p. 14
C CONVERGENCE ANALYSIS
The convergence analysis interprets support-projected Newton–Schulz updates as descent on a fixed-support orthogonality objective, with convergence guarantees under stated conditions. The accompanying sections describe operator-based grafting factors and the CIFAR-10 and ImageNet experimental settings.
C CONVERGENCE ANALYSIS
- Support projection turns the classical Newton–Schulz iteration into projected descent on a fixed-support orthogonality objective.For strided convolutions, the appropriate objective is the stride-aware Lorth surrogate, whose gradient is implemented by convolution followed by transposed convolution.Source: C CONVERGENCE ANALYSIS · pp. 14–15
- Under the stated boundedness and step-size conditions, the iterates converge to a single critical point of the support-constrained objective.When an orthogonal kernel exists in the fixed-support class, the solver can converge to one; otherwise it converges to a stationary approximation.Source: C CONVERGENCE ANALYSIS · pp. 15–16
- The convergence setup requires the induced operator norm to start at or below 1, enforced with a spectral-norm upper bound and damping for numerical drift.A tighter spectral estimate can accelerate convergence; the cited implementation computes a certifiable bound across padding modes.Source: C CONVERGENCE ANALYSIS · p. 16
D CORRECTION OF THE GRAFTING FACTOR:
- Muon grafting factors for convolutional updates should use the dimensions of the induced operator rather than the reshaped kernel.The paper compares adaptations of existing scaling conventions and notes that one cited proposal uses distinct dense and convolutional formulas.Source: D CORRECTION OF THE GRAFTING FACTOR · p. 17, Table 3 · p. 17
E CIFAR-10 AND IMAGENET CLASSIFICATION DETAILS
- CIFAR-10 uses AirBench94 with a compact CNN, an eight-epoch schedule, batch size 2000, and matched settings except for convolutional orthogonalization.The study reports 25 runs per method; Conv-NS is tuned with a 1000-configuration search budget because no reference setting existed.Source: E CIFAR-10 AND IMAGENET CLASSIFICATION… · p. 17
- The ImageNet-1k comparison uses ResNet-50 and a 30-epoch recipe, comparing reshape-based Muon with Conv-NS.Training uses global batch size 2048 across four H100 GPUs, and both strategies receive a 48-configuration grid search.Source: E CIFAR-10 AND IMAGENET CLASSIFICATION… · p. 17
F PER-STAGE PARETO ANALYSIS ON REAL TRAINING GRADIENTS
On real training gradients, reshape-based Muon lies on or near the orthogonality–fidelity Pareto front across training stages. Conv-NS traces the front as its iteration count increases.
F PER-STAGE PARETO ANALYSIS ON REAL TRAINING GRADIENTS
- Reshape-based Muon (RKO) lies on or very near the Pareto front across beginning, middle, and end training snapshots.The stage-by-stage pattern is reported as consistent rather than specific to one training phase.Source: Figure 6 · p. 18, F PER-STAGE PARETO ANALYSIS ON REAL… · p. 18
- Conv-NS traces the Pareto front as the iteration count increases.The figure computes a separate front for each of the three training stages.Source: Figure 6 · p. 18
G HIERARCHICAL OPTIMIZATION OF ORTHOGONALITY OBJECTIVES
The section compares reshape-based and convolution-aware orthogonalization across synthetic kernels, training gradients, and classification benchmarks. It also examines how iteration count and alternative convolutional methods affect the comparison.
G HIERARCHICAL OPTIMIZATION OF ORTHOGONALITY OBJECTIVES
- On synthetic kernels, RKO's Lflat residual was 0.02 but its Lorth was 14.40; ConvPF achieved 0.10 and 0.02, while ten Conv-NS iterations yielded 3.38 and 1.11.On real training gradients, ConvPF's Lflat was ≈29 versus RKO's 0.22–0.51, so the synthetic dual-objective relationship was not established there.Source: G HIERARCHICAL OPTIMIZATION OF… · pp. 18–19, Table 4 · p. 18, Table 5 · p. 19
H EFFECT OF ITERATION COUNT IN THE 8-EPOCH SETTING
- In the eight-epoch CIFAR-10 setting, the iteration-count sweep selects momentum by average TTA validation accuracy across learning rates, with five runs per learning rate and standard-deviation bands.The text reports that, at low iteration budgets, more iterations improve performance while shifting the optimal learning rate and momentum.Source: H EFFECT OF ITERATION COUNT IN THE… · p. 19, Figure 7 · p. 19
I ONE-EPOCH SELECTION CRITERIA
- Under both validation- and training-accuracy selection, Conv-NS performance improved through 10 iterations and degraded at 25 and 100, indicating that the decline also occurs at the optimization level.The result persisted despite substantial shifts in selected learning rate.Source: I ONE-EPOCH SELECTION CRITERIA · p. 20, Table 6 · p. 20
J COMPARISON WITH ONE PASS IMAGENET
- Both Muon variants exceeded all OPIN results at two effective epochs, including its 46.1% full-storage baseline, while using half the examples and about eight times fewer updates without replay.They remained below OPIN's four-effective-epoch results of 55.5–57.2%; differing training recipes make this a cross-paper, not controlled, comparison.Source: J COMPARISON WITH ONE PASS IMAGENET · p. 20, Table 7 · p. 21
K COMMUNITY BASELINE COMPARISONS
- Freq-Muon orthogonalizes per-frequency matrices, becoming exact for circular convolution only when the FFT grid matches the full input dimensions; its practical small grid is a surrogate.Matrixizer searches tensor permutations and split points for a lower-cost 2D factorization, but remains reshape-based and does not model stride, groups, dilation, or the structured convolution operator.Source: K COMMUNITY BASELINE COMPARISONS · pp. 21–22
K.1 SYNTHETIC ANALYSIS OF THE ORTHOGONALITY–FIDELITY TRADE-OFF
The synthetic analysis compares orthogonalization methods against an empirically estimated orthogonality–fidelity frontier. The section also reports CIFAR-10 community-baseline results and gives dense and convolutional linear-minimization constructions.
K.1 SYNTHETIC ANALYSIS OF THE ORTHOGONALITY–FIDELITY TRADE-OFF
- As Conv-NS iterations increase from 2 to 100, its operating points move along the empirical Pareto front toward stronger orthogonality and greater kernel modification.Five iterations already place Conv-NS near the front at Lorth ≈4.9; no single-shot baseline reaches a comparable position.Source: §K.1 · p. 22
- On synthetic kernels, reshape-based Muon modified the kernel (Lfrob ≈2.3) without improving operator orthogonality (Lorth ≈15.6 versus ∼10.2 unmodified), reflecting structural misalignment.The text notes that real training-gradient structure corrects this misalignment, bringing RKO surprisingly close to the Pareto front in practice.Source: §K.1 · p. 22
- Matrixizer's Lorth ≈157 and Freq-Muon's Lorth ≈3.7×10^4 place both far above the synthetic Pareto front.The text attributes the artifacts to their Kronecker-factorization and partial-spectrum FFT approximations, respectively.Source: §K.1 · p. 22
K.2 CIFAR-10 COMMUNITY BASELINE RESULTS
- Conv-NS and reshape-based Muon achieved near-identical CIFAR-10 AirBench94 performance and both outperformed the single-run SGD, Freq-Muon, and Matrixizer baselines.The community-method results are diagnostic context because those implementations have preliminary, non-standardized evaluations.Source: §K.2 · pp. 22–23, Table 8 · p. 23
L LMO DENSE MATRIX PROOF
- For dense matrices, the paper frames Muon's update as maximizing first-order loss decrease under an operator-norm bound on layer-output variation, yielding a direction proportional to the gradient's polar factor.The displayed first-order expansion and the accompanying prose establish the optimization objective and its polar-factor solution.Source: L LMO DENSE MATRIX PROOF · p. 23
M CONVOLUTIONAL LMO PROOF
- The convolutional LMO evaluates output variation through the structured operator, but ConvPF is optimal only on a kernel subspace and therefore overconstrains the update.Unlike the dense-matrix proof, the convex-hull argument is not an equality in the convolutional setup.Source: M CONVOLUTIONAL LMO PROOF · p. 24