Module-Aware Optimization for Graph Neural Networks

AI summary

Overview

Research area: Machine learning — specifically optimization algorithms for Graph Neural Networks (GNNs), sitting at the intersection of optimizer design (Adam/AdamW vs. matrix-aware methods such as Muon) and graph representation learning.

Technical level: Advanced. The paper reasons about the geometry of parameter updates, orthogonalized/polar update steps, and covariance structure of layer inputs — background in optimization and linear algebra is assumed.

Scope: The paper proposes GraphMuon, a graph-aware variant of the Muon optimizer that defines its orthogonalization metric using the covariance of graph-propagated layer inputs, argues theoretically that vanilla Muon can be destabilized by graph propagation, and reports time-to-target speedups over AdamW on OGB benchmarks and across 15 dataset-architecture configurations.

What This Paper Is About

Adaptive optimizers like Adam/AdamW treat each parameter coordinate separately, while matrix-aware optimizers like Muon treat a weight matrix as a whole and orthogonalize its update so no direction dominates. The authors point out that for GNNs this matrix-level normalization is incomplete: the numbers a GNN layer actually multiplies are not raw features but graph-propagated (neighborhood-aggregated) representations, which have their own covariance structure. So an update that looks perfectly balanced as a matrix can still produce an outsized change in the layer's output once applied to propagated features. The goal is an optimizer whose update geometry accounts for that propagation, rather than for the weight matrix alone.

Key Contributions

1. A diagnosis of graph-induced update amplification. The authors identify a mismatch between how matrix-aware optimizers normalize updates and how GNN layers actually consume graph-propagated inputs, and show theoretically that vanilla Muon can suffer from this amplification.
2. GraphMuon, a graph-aware optimizer. Muon's matrix update is adapted to the covariance structure of the graph-propagated layer inputs: orthogonalization is performed in a metric defined by those representations, and step sizes are normalized according to the update's induced effect on the layer representation.
3. A theoretical result. The paper shows that GraphMuon mitigates the amplification effect, and that its update takes the form of a graph-metric analogue of Muon's polar update step.
4. An empirical evaluation across settings. GraphMuon is tested on three large-scale OGB benchmarks and on 15 dataset-architecture configurations covering node classification, inductive graph classification, and graph regression, with code released publicly.

Main Findings

• Vanilla Muon can be amplified by graph structure: the abstract states a theoretical result that graph-propagated inputs can amplify the effect of a matrix-normalized update on the layer output, undermining the normalization Muon provides.
• GraphMuon counteracts this: the proposed metric-aware orthogonalization and representation-aware step normalization are claimed to mitigate the amplification, yielding a graph-metric analogue of Muon's polar step.
• Speedups on large OGB benchmarks: reported $1.32\times$–$2.42\times$ validation-loss time-to-target speedups over AdamW on three large-scale OGB benchmarks.
• Broader win rate across configurations: across 15 dataset-architecture configurations spanning node classification, inductive graph classification, and graph regression, GraphMuon reaches lower time-to-target in 13 of 15 cases.
• Predictive quality is not sacrificed: the abstract reports competitive or superior predictive performance alongside the faster convergence. Specific per-dataset numbers, ablations, and baseline comparisons beyond AdamW are not given in the abstract.

Methodology in Plain English

Muon's idea is to take the raw gradient for a weight matrix and "orthogonalize" it — roughly, rotate it into a form where all its directions have comparable strength, so no single direction blows up the update. The authors' observation is that this balancing act is done with respect to the weight matrix, but a GNN layer's output depends on the matrix times its input, and in a GNN that input has already been mixed with neighbors' features. Neighbor aggregation makes some directions in the input much more prominent than others, so a matrix that looks balanced can still hit those prominent directions hard.

GraphMuon therefore measures the input's actual spread (its covariance) after graph propagation, and does the orthogonalization in a coordinate system shaped by that measurement rather than the default one. It also rescales each step according to how much it would actually move the layer's output representation, not just how large the raw parameter change is. The paper supports this design with an analysis of when vanilla Muon's update gets amplified, and then validates it by measuring how quickly validation loss drops relative to AdamW on standard graph benchmarks.

Why This Matters

Research impact. Optimizer research has largely been organized around either coordinate-wise adaptivity or weight-matrix structure; this paper argues that for graph models a third factor — the statistics of the inputs a module receives, after message passing — belongs in the update rule. That reframes "module-aware" optimization as conditioning on the data flowing through a layer, and gives the Muon family a principled extension to graph architectures.

Real-world applications.

• Recommender systems and social platforms, where user–item or user–user graphs are large and training cost is dominated by many-epoch convergence.
• Fraud and anomaly detection on transaction networks, where models must be retrained frequently on evolving graphs and faster convergence directly cuts compute cost.
• Molecular property prediction and drug discovery, where graph regression/classification models are trained repeatedly across many dataset–architecture combinations.
• Traffic, logistics, and citation/graph classification pipelines where large-scale graph benchmarks (the OGB setting used here) are the standard proxy for production workloads.

Industry relevance. Time-to-target is the metric practitioners care about most: it maps directly to GPU hours and experimentation throughput. A drop-in optimizer that reaches a target validation loss substantially sooner than AdamW, while matching or improving final accuracy, is attractive for teams training GNNs at scale without wanting to redesign architectures.

Future Directions

• Extending beyond Muon-style updates: whether the same input-covariance reweighting can be applied to other matrix-aware or second-order-flavored optimizers, and how it composes with adaptive methods.
• Scalability of the covariance estimate: computing and inverting a metric defined by propagated representations raises cost and memory questions for very wide layers and for mini-batch sampling in large graphs; the abstract does not address this trade-off.
• Architectures with heterogeneous propagation: graphs with multiple relation types, heterogeneous or temporal message passing, and attention-based GNNs may have layer inputs whose covariance structure changes rapidly during training.
• Theory-to-practice gap: tightening the analysis of when amplification occurs — characterizing which graph spectra and propagation depths make vanilla Muon fail — to guide optimizer selection without trial runs.

Target Audience

Optimization researchers and graduate students working on training algorithms, and GNN practitioners or ML engineers who train graph models at scale and are weighing alternatives to AdamW. Readers should be comfortable with matrix decompositions and the standard framing of optimizers as imposing geometries on parameter space; those looking for an introductory treatment of GNNs or of Muon will need background reading first.

Authors’ abstract

Different optimizers impose different geometries on parameter updates. Adaptive methods such as Adam and AdamW use coordinate-wise updates, while matrix-aware optimizers such as Muon utilize weight-matrix structure through orthogonalized updates. For Graph Neural Networks (GNNs), however, the effective geometry of its weight matrices is shaped not only by its learned channel-mixing weights, but also by the graph-propagated representations it acts on. Thus, even if an update is well normalized as a matrix, its effect on the layer output can be amplified when applied to graph-propagated features. To address this gap, we introduce GraphMuon, a graph-aware optimizer that adapts Muon's matrix updates to the covariance structure of graph-propagated layer inputs. GraphMuon performs orthogonalization in a metric defined by these propagated representations and normalizes step sizes according to the update's induced effect on the layer representation. We show that theoretically, vanilla Muon can suffer from graph-induced update amplification, while our GraphMuon can mitigate this effect. As such, our GraphMuon yields a graph-metric analogue of Muon's polar update step. Empirically, GraphMuon achieves $1.32\times$-$2.42\times$ validation-loss time-to-target speedups over AdamW on three large-scale OGB benchmarks. Across 15 dataset-architecture configurations spanning node classification, inductive graph classification, and graph regression, it achieves lower time-to-target in 13 cases, while maintaining competitive or superior predictive performance. Code is available at https://github.com/guy-hadad/GraphMuon.