GraphMuon retunes GNN updates to graph-propagated features, reaching AdamW's loss target 1.32x to 2.42x faster on three OGB benchmarks
Synopsis
A GNN weight matrix acts on graph-propagated features, so an update that is well normalized in parameter space can still be amplified by graph structure into a large representation-level perturbation. GraphMuon moves Muon's matrix orthogonalization into a metric defined by the covariance of the propagated layer inputs and scales each step by its induced effect on the layer representation. On three large-scale OGB benchmarks it reaches AdamW's validation-loss target in 1/2.32 to 1/1.32 of the wall-clock time, a 1.32x to 2.42x speedup.
Interpretation
GNN optimizers need to measure updates by the graph-propagated representation, not by the parameter matrix alone. AdamW normalizes gradient moments coordinate-wise, and Muon approximately orthogonalizes the weight matrix; neither distinguishes directions amplified by graph propagation from directions attenuated by it. On GCN backbones over the three large-scale OGB benchmarks OGBN-ARXIV, OGBN-MAG and OGBN-PRODUCTS, GraphMuon's wall-clock speedups to AdamW's validation-loss target are 2.32x, 2.42x and 1.32x, with 95% bootstrap confidence intervals of [2.01, 2.69], [2.29, 2.53] and [1.19, 1.37].
The method moves Muon's polar-style update into the graph-induced metric and then sets the step by representation-level perturbation. GraphMuon forms the covariance of each layer's graph-propagated input, maps momentum into that metric with its powers, applies Newton-Schulz orthogonalization and maps back, then rescales the update by the representation-level perturbation it induces under the same covariance. Across 15 dataset-architecture configurations spanning node classification, inductive graph classification and graph regression, GraphMuon reaches the target in fewer epochs than AdamW in all 15, has the lowest wall-clock time in 13, and averages a 4.62x speedup while keeping competitive or superior predictive performance.
Perspective
The evaluation covers 15 dataset-architecture configurations on nine graph-learning benchmarks, including transductive and heterophilous node classification, inductive molecular graph classification and ZINC graph regression, with GCN, GAT, GINE, GraphGPS and GraphGCN backbones, all run on a single NVIDIA RTX 6000 GPU.
On very small graphs, where the node count is small relative to the hidden dimension, the graph-induced covariance can be low-rank and noisy; controlling representation-level perturbation then need not yield faster convergence or better generalization, and AdamW can remain competitive. On Citeseer with GCN and ZINC with GraphGCN, GraphMuon needs fewer epochs but is slightly slower than AdamW in cumulative wall-clock time. The best value of the graph-awareness strength τ is dataset-dependent.