PaperlayerView exact revision on arXiv

Source-linked AI summary

Muon-C: Operator-Aligned Muon for Convolutional Kernels

Jiaxin Qing, Lexin Li

arXiv:2609.09676v1cs.LGstat.ML

TL;DR

Muon-C addresses the ambiguity of applying Muon to convolutional kernels by aligning its geometry with frequency-wise convolution operator maps. The method has a no-weaker worst-case oracle guarantee than unfolding, stronger guarantees for 3 × 3 kernels, and improved optimization efficiency across varied settings.

  • Problem

    Standard kernel unfolding makes Muon act on a local patch-to-output map rather than directly on the convolution operator, selecting one non-neutral optimizer geometry.

  • Method

    Muon-C applies independent polar updates to frequency-wise channel-transfer matrices on a critical Fourier grid, preserving finite kernel support.

  • Results

    Muon-C is an exact linear minimization oracle for the critically sampled convolution norm, with a worst-case guarantee no weaker than unfolding and strictly stronger for 3 × 3 kernels.

  • Takeaways & Limitations

    The results support translation-frequency organization as a useful operator-aligned geometry for improving convolutional optimization efficiency across data scales, training objectives, and architectures.

  • Takeaways & Limitations

    Muon-C adds optimizer-side computation, while the strongest empirical evidence is concentrated in relatively small-resolution settings.

Abstract

from arXiv · show

Muon replaces matrix momentum with an approximately orthogonal polar direction, but its geometry depends on the matrix representation. For convolution, standard unfolding describes a local patch map rather than the convolution operator. We introduce Muon-C, an operator-aligned optimizer that represents kernel momentum as frequency-wise channel-transfer matrices, polarizes these blocks independently, and uses a critical Fourier grid to return updates exactly to the original finite kernel support. We show that the new geometry arises from combining the block partition and Fourier coordinates. The exact-polar direction is a linear minimization oracle under the critically sampled convolution norm. Its worst-case guarantee relative to the continuous convolution-operator norm is never weaker than unfolding and is strictly stronger for $3\times3$ kernels. On CIFAR-10 flow matching with matched applied-update RMS, Muon-C reaches 9.87 FID at 40k iterations, compared with 22.26 for unfolded Muon and 51.31 for Adam. It reaches their final quality using $0.62\times$ and $0.64\times$ their model FLOPs, respectively. Under equal tuning budgets, Muon-C achieves 3.42 FID. Gains persist across data scales and transfer to classification across convolutional architectures.

1 Introduction

Muon-C addresses the ambiguity of applying Muon to convolutional kernels by aligning its independently polarized blocks with frequency-wise channel-transfer maps, while critical sampling preserves finite kernel support. It provides stronger theoretical guarantees than unfolding and improves optimization efficiency across generative and classification settings.

  • Operator-aligned geometry: Muon-C represents convolutional momentum as frequency-indexed channel maps and independently polarizes each map, rather than applying Muon to an arbitrary tensor unfolding.The design follows translation equivariance, which makes Fourier-domain channel-transfer matrices the natural linear maps of convolution.
  • Theory: Muon-C is the exact linear minimization oracle under the critically sampled convolution norm, with a worst-case guarantee never weaker than unfolding and strictly stronger for 3 × 3 kernels.For 3 × 3 kernels, the guaranteed fraction of the continuous-oracle optimum improves from 1/3 for unfolding to 9/25 for Muon-C.
  • Critical sampling: Critical frequency sampling combines with blockwise polarization to map updates exactly back to the original finite kernel support.The two ingredients are distinct: blockwise polarization changes the optimizer geometry, while critical sampling preserves support.
  • Empirical evaluation: Muon-C improves optimization efficiency across generative and classification settings, reaching comparable loss or sample quality earlier than unfolded Muon and Adam/AdamW.The paper reports persistence across data scales and transfer to classification across convolutional architectures.

2 Related Work

Prior work interprets Muon through spectral-norm geometry, structured tensor optimization, and convolutional spectral analysis. Muon-C differs by deriving independently updated matrices from convolution’s frequency-wise operator structure and preserving finite support through critical sampling.

  • Convolutional geometries: Muon-S independently polarizes spatial-offset channel matrices, whereas Muon-C independently polarizes translation-frequency channel maps.The matched spatial-block construction isolates the contribution of frequency organization.
  • Structured optimization: Muon-C complements K-FAC, Shampoo, and Tensorion by deriving its geometry from the parameterized linear operator rather than from tensor modes or adaptively selected unfoldings.This distinction makes operator semantics, not coordinate structure alone, determine the blocks seen by Muon.
  • Relation to spectral analysis: Unlike prior convolution-spectrum methods that use spectra as constraints or diagnostics, Muon-C uses transfer-matrix structure to design optimizer updates.The update blocks correspond to the linear maps implemented on translation modes.
  • Operator-aligned optimization: Muon-C derives its independently updated matrices from frequency-wise channel-transfer maps of the convolution operator, rather than selecting an unfolding from tensor structure alone.Critical sampling provides a dimension-matched representation for the finite stored kernel.

3 Muon as a Spectral-Norm Oracle

Muon treats matrix momentum through a polar update whose spectral-norm interpretation makes representation and block structure part of the optimizer geometry. This motivates analyzing structured parameters through the matrices and independently constrained blocks on which the update acts.

  • Convolutional notation: For convolutional kernels, the notation represents each spatial coefficient as a Cout × Cin channel matrix over a kh × kw kernel grid.The kernel is equivalently a four-dimensional tensor of channel and spatial dimensions.
  • Muon update: Muon replaces matrix momentum with its hard polar factor, approximated in practice using Newton–Schulz iterations on normalized momentum.The polar factor is defined from a compact singular value decomposition.
  • Spectral-norm oracle: The polar direction is a linear minimization oracle for a spectral-norm trust region, solving min∥U∥op≤ρ⟨M, U⟩F with U⋆ = −ρ polar(M).This gives Muon a steepest linearized-descent interpretation rather than treating it as heuristic momentum normalization.
  • Representation-dependent geometry: The matrix representation determines Muon’s constraint set because row and column assignments define the spectral norm and independently constrained blocks define a product of spectral-norm balls.Thus, different representations of the same coefficients can induce different optimizer geometries.

4 From Muon to Operator-Aligned Geometry for Convolution

The section argues that convolutional Muon geometry should act on frequency-wise channel-transfer maps, not only on an unfolded local patch map. It shows that independent block partitioning—not Fourier reparameterization alone—creates the distinct operator-aligned direction used by Muon-C.

  • 4.1 From Local Patch Geometry to the Convolution Operator: Standard unfolding gives an exact spectral-norm oracle for a local patch-to-output map, but its row and column choices impose one particular kernel geometry.The full convolution operator is resolution dependent if formed explicitly, motivating its decomposition into smaller invariant channel maps.
  • 4.2 Translation Equivariance Identifies Fourier Channel Maps: Frequency-wise Cout×Cin channel-transfer matrices are resolution-independent maps implemented by convolution on its translation modes.Translation equivariance makes the Fourier basis block diagonal, leaving only channel mixing within each frequency.
  • 4.3 Changing Basis Is Not Enough: Blocks Define the Geometry: Fourier reparameterization alone does not change Muon’s global direction: unitary row and column transformations preserve the polar map.A single polar map followed by undoing the transformation returns the unfolded polar direction, including the descent sign.
  • 4.3 Changing Basis Is Not Enough: Blocks Define the Geometry: Independent block partitions change the constraint set from one spectral-norm ball to a product of spectral-norm balls, yielding genuinely different linear minimization oracles.Thus global unfolding, spatial partitioning, and critical-Fourier partitioning can produce pairwise non-collinear Muon directions for the same coefficients.
  • 4.4 Critical-Fourier Geometry: Muon-C independently polarizes one channel map per critical translation frequency, combining operator alignment with critical sampling that preserves finite kernel support.The operator view identifies frequency-wise maps—not Fourier coordinates alone—as the appropriate geometry.

5 Muon-C: Critical Fourier Geometry for Finite Convolutional Kernels

Muon-C aligns Muon with convolution operators by independently polarizing frequency-wise channel-transfer matrices on a critical Fourier grid that preserves finite kernel support. This construction yields an exact oracle for a sampled convolution norm, with no weaker worst-case control than unfolding and strictly stronger control for 3 × 3 kernels.

  • 5.1 From Continuous Operator Geometry to Critical Sampling: The continuous convolution norm is defined by the supremum of the operator norms of the kernel’s frequency-dependent transfer matrices.It is the induced ℓ2 norm of multichannel convolution and is resolution-independent.
  • 5.1 From Continuous Operator Geometry to Critical Sampling: Independent polarization of full-grid frequency blocks can violate finite support, whereas critical sampling makes independent frequency-wise polarization feasible.The critical grid matches the number of stored coefficient matrices and keeps every inverse-transformed update within the kernel support.
  • 5.1 From Continuous Operator Geometry to Critical Sampling: Muon-C uses a critical kh × kw Fourier grid whose frequency-wise channel-transfer blocks map bijectively back to the original finite kernel support.The grid has exactly n = khkw samples, and its inverse transform preserves support without cropping or projection.
  • 5.3 Theoretical Properties of Muon-C: Muon-C is the exact linear minimization oracle under the critically sampled convolution norm, with real-valued updates preserved through conjugate symmetry.The characterization follows from Parseval’s identity and spectral-nuclear norm duality applied independently to each critical-frequency block.
  • 5.3 Theoretical Properties of Muon-C: For 3 × 3 kernels, Muon-C’s worst-case approximation factor is 25/9, versus 3 for global unfolding, corresponding to a stronger continuous-oracle guarantee.The general Muon-C distortion is no larger than unfolding for any finite kernel size and is strictly smaller for 3 × 3 kernels.
  • 5.4 Practical Muon-C Algorithm: Muon-C adds kernel-axis FFTs and multiple channel-block polar operations, creating optimizer-side overhead excluded from model-FLOP accounting.The overhead is independent of feature-map resolution and is separately evaluated using wall-clock and peak-memory measurements.

6 Numerical Experiments

Across generative and discriminative tasks, Muon-C improves early optimization and generally maintains an endpoint advantage over unfolded Muon and Adam/AdamW. Controlled experiments, robustness studies, and ablations associate these gains with frequency-wise operator alignment and reduced directional curvature.

  • Optimization efficiency: 9.87 FID at 40k on CIFAR-10 flow matching compares with 22.26 for unfolded Muon and 51.31 for Adam under matched applied-update RMS.Muon-C reaches the baselines’ 400k-step final quality using 0.62× and 0.64× their model FLOPs, respectively.
  • Robustness and tuning: 3.42 FID under equal tuning budgets compares with 3.54 for unfolded Muon and 3.83 for Adam on CIFAR-10.Muon-C also has the lowest FID at every reported budget fraction on CIFAR-10 and ImageNet-1k-32.
  • Classification transfer: 80.22% top-1 accuracy on ResNet-34, 79.88% on ResNet-50, and 84.88% on ConvNeXtV2-T show endpoint gains over the compared optimizers.The gains are larger early in training and smaller at the 100-epoch endpoint.
  • Mechanism ablation: 166.83 FID versus 211.25 for Muon-S at 20k iterations shows that frequency-wise blocks outperform matched spatial-offset blocks early.The ablation supports translation-frequency organization as a contributor to optimization efficiency.
  • Theoretical comparison: For 3×3 kernels, Muon-C’s worst-case approximation factor is 25/9 versus 9 for spatial-block Muon-S, with corresponding continuous-oracle guarantees of 9/25 versus 1/9.Both methods are exact linear minimization oracles for their respective finite-support geometries.

7 Discussion

Muon-C combines frequency-wise operator alignment with independent block polarization, yielding a geometry with theoretical and empirical support. Its guarantees are no weaker than unfolding, while practical advantages persist across settings, alongside clear open questions about scope and convergence.

  • Muon-C realizes operator-aligned geometry by independently polarizing frequency-wise channel-transfer maps while returning updates to the original finite kernel support.The construction combines translation-frequency organization with a critical Fourier grid rather than relying on a coordinate change alone.
  • Muon-C’s worst-case guarantee is no weaker than global unfolding for finite kernels and is strictly stronger for 3 × 3 kernels.The critical DFT is unitary, bijective, and support-preserving, and Muon-C is an exact linear minimization oracle for the critically sampled convolution norm.
  • Empirical gains persist across stochastic replication, data scales, discriminative training, and architectures with different routing coverage.A matched Muon-S comparison supports translation-frequency organization as a contributor to early trajectory gains.
  • Muon-C adds optimizer-side computation, while strongest evidence remains concentrated in relatively small-resolution generative models and ImageNet-100 classification.The paper identifies broader wall-clock studies, full-resolution ImageNet-1k, larger kernels, and additional convolution-heavy tasks as future tests.
  • Whether improved operator-oracle alignment yields faster convergence remains open, with reduced directional curvature offered as a possible mechanism.Further questions include finite-iteration Newton–Schulz analysis and extensions to strided and dilated convolution.

Appendix A. Proofs and Operator-Norm Analysis of Muon-C

The proofs establish Muon-C’s Fourier-block formulation, its exact oracle property under the critically sampled norm, and its comparison with the continuous convolution-operator norm. The analysis also derives the strict 3 × 3 improvement over unfolding.

  • At every critical frequency, Parseval separates the objective and constraint by block, allowing spectral–nuclear norm duality to produce independent polar updates.The resulting Fourier oracle uses the polar factor of each frequency block and preserves reality under conjugate symmetry.
  • The critical Fourier grid provides a unitary, bijective, support-preserving representation of stored convolution coefficients.This representation enables frequency-wise block analysis while preserving real spatial tensors through conjugate-symmetric choices.
  • Muon-C is an exact linear minimization oracle for the critically sampled convolution norm.The full-grid reference defines a frequency-block norm proportional to the induced norm of the corresponding circular-convolution operator.
  • For 3 × 3 kernels, the proof obtains Λ3 = 5/3 from the extremal frequency comparison.The bound follows by analyzing three absolute-value expressions whose sum reaches equality at θ = π.
  • The continuous norm bounds imply that the surrogate oracle gap is at least as favorable as the continuous convolution-operator gap.The theorem gives ∥U∥r ≤ ∥U∥conv ≤ κr∥U∥r, yielding Δr ≥ Δconv.

Appendix B. Implementation and Experimental Details

The implementation applies FFT-based frequency-block polarization with a fixed Newton–Schulz approximation and routes only spatial convolutions through Muon-C. The appendices document training schedules, tuning, evaluation, and compute accounting.

  • Muon-C uses a two-dimensional orthonormal FFT and five float32 complex Newton–Schulz iterations to approximate the hard-polar direction.Momentum is 0.95 with Nesterov momentum, and no support-constraint iteration is used.
  • Spatial Conv2d kernels with area greater than one use Muon-C, while 1 × 1 convolutions, biases, normalization parameters, and residual matrices use AdamW.This routing rule is shared across the reported training configurations.
  • CIFAR-10 and ImageNet-1k-32 use 400k-step U-Net training with warmup-constant schedules, mixed precision, gradient clipping, and EMA weights.ImageNet-100 trains ResNet-34, ResNet-50, and ConvNeXtV2-T for 100 epochs at 224 × 224 resolution.
  • Learning rates are selected from ten values between 10^-4 and 10^-3 using the lowest final FID-50k.The selected rates differ across Adam, global unfolded Muon, and Muon-C and across datasets.
  • FID-50k uses 50k generated samples, 100 Dopri5 sampling steps, fixed real-data statistics, and EMA checkpoints.Figure 7 counts cumulative U-Net forward–backward model FLOPs while excluding optimizer, evaluation, and data-pipeline computation.
  • A 3 × 3 critical grid with 256 input and output channels stores six complex frequency blocks and occupies approximately 3 MiB in complex64.Larger feature grids require substantially more storage, with estimates reaching approximately 12.4 GiB at 224 × 224 before Newton–Schulz temporaries.

Appendix C. Additional Empirical Results and Diagnostics

The diagnostics measure applied parameter displacement under matched RMS scaling and show that Muon-C generally uses smaller route-level and layerwise RMS than globally unfolded Muon.

  • Muon-C’s global route RMS is 0.888 times global unfolded Muon on average across measured steps.The comparison uses AdamW-referenced RMS scaling with shared initialization, mini-batches, and stochastic draws.
  • The median layerwise RMS ratio is 0.744, with 43 of 52 layers at ratios no greater than one.These diagnostics indicate that the smaller global route RMS is distributed across most measured layers.

C.2.1 Replication across Training Seeds

Across three paired CIFAR-10 training seeds, Muon-C consistently achieves the best endpoint while preserving a controlled common-learning-rate comparison.

  • C.2.1 Replication across Training Seeds: The replication uses three paired runs with shared training seeds, base learning rate 2 × 10^-4, and the same final FID-50k protocol.
  • C.2.1 Replication across Training Seeds: Muon-C achieves the best endpoint in every paired seed, improving mean FID by 0.252 over global unfolded Muon and 0.338 over Adam.Seedwise gains over global unfolded Muon range from 0.080 to 0.354.
  • C.2.1 Replication across Training Seeds: Table 10 reports final CIFAR-10 FID-50k means and sample standard deviations, while Table 19 summarizes applied-update RMS ratios for eligible spatial-Conv2d layers.

C.2.2 Performance across Training Budgets

Muon-C maintains the strongest reported performance across training budgets, while diagnostics attribute its finite-support behavior to the interaction between frequency-wise polarization and stored kernel support.

  • C.2.2 Performance across Training Budgets: At every reported budget fraction on both datasets, Muon-C has the lowest FID; at 0.1B it reaches 10.37 on CIFAR-10 and 21.14 on ImageNet-1k-32.At 0.1B, global unfolded Muon reaches 24.97 and 28.02, while Adam reaches 51.37 and 58.87, respectively.
  • C.2.2 Performance across Training Budgets: The selected runs preserve the same optimizer ordering at 0.1B and 0.5B as in the final tuned comparison.
  • C.2.2 Performance across Training Budgets: The Figure 10 diagnostic compares global unfolded Muon, Muon-S, and Muon-C directions from shared momentum states using matched convolution-operator-norm budgets and exact SVD-based polar factors.States are sampled at 40k, 200k, and 400k iterations on a CIFAR-10 U-Net trajectory.
  • C.2.2 Performance across Training Budgets: The finite-support diagnostic finds substantial, stable obstruction from independently polarizing the full feature-frequency grid and then restricting updates to stored kernel support.The discrepancy is attributed to incompatibility between full-grid frequency polarization and the original finite support, not inaccurate full-grid optimization.
  • C.2.2 Performance across Training Budgets: Exact-SVD checks solve the idealized full-grid oracle accurately, with maximum relative objective gap 9.27 × 10^-8 and feasibility error 7.15 × 10^-7.The finite-support identity has maximum residual 1.79 × 10^-7.
Loading 2609.09676v1…