Search papers, labs, and topics across Lattice.
Muon-C repairs the geometric mismatch of applying the Muon optimizer to unfolded convolutional kernels by reformulating momentum as block-diagonal, frequency-wise channel-transfer matrices in the Fourier domain. This operator-aligned approach provides strictly stronger theoretical worst-case operator norm guarantees for standard $3\times3$ kernels compared to spatial unfolding. On CIFAR-10 flow matching, Muon-C achieves equivalent generation quality in roughly $0.63\times$ the FLOPs of Adam and standard unfolded Muon, reaching an FID of 3.42 under equal tuning budgets.
Unfolding convolutional kernels breaks Muon's optimization geometry; aligning polar updates with the true convolution operator in the frequency domain cuts flow-matching compute by nearly 40% while radically outperforming Adam.
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.
