Muon-C improves convolution learning with operator-aligned optimization

Muon-C: Operator-Aligned Muon for Convolutional Kernels

Machine Learning

Summary

Training convolutional neural networks involves updating kernels, or filters, to recognize patterns in images. The authors propose Muon-C, a method that aligns updates with how convolution operators work, improving efficiency and performance. This approach better matches the geometry of convolution, leading to faster learning and higher quality results on image tasks. Their tests show it outperforms standard optimizers like Adam and previous Muon methods on common benchmarks.

What this means in practice

  • For machine learning engineers: Optimize convolutional network training by using operator-aligned kernel updates for improved accuracy and efficiency in image tasks.
  • For computer vision developers: Incorporate Muon-C optimization to achieve better image generation and classification performance with lower computational cost.

Authors

Jiaxin Qing, Lexin Li

Abstract

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.