Module-Aware Optimization for Graph Neural Networks

Module-Aware Optimization for Graph Neural Networks. By Guy Hadad and three others, from Ben-Gurion University of the Negev. Published on the fifth of November.

Every optimizer imposes a geometry on the updates it makes. Adaptive methods such as Adam and AdamW work coordinate by coordinate. Matrix-aware optimizers such as Muon use the structure of the weight matrix, through orthogonalized updates. For graph neural networks, the authors argue, even that is not the whole picture. The effective geometry of a layer's weight matrix is shaped by two things: the channel-mixing weights the layer learns, and the graph-propagated representations those weights act on. An update can be perfectly well normalized as a matrix and still have its effect on the layer output amplified, once it is applied to features that have been propagated over the graph.

GraphMuon is a graph-aware optimizer built to close that gap. It keeps Muon's matrix updates, but adapts them to the covariance structure of the graph-propagated inputs each layer receives. The geometry the update respects is no longer only the geometry of the weights. It is also the geometry of what the weights are applied to.

Two changes carry this. First, the orthogonalization is performed in a metric defined by the propagated representations, rather than in the plain metric of the weight matrix. Second, step sizes are normalized according to the effect the update induces on the layer representation, not according to the update alone. A theoretical result backs the design in two parts: vanilla Muon can suffer from graph-induced update amplification, and GraphMuon can mitigate it. Put together, the method amounts to a graph-metric analogue of Muon's polar update step.

The empirical claims are framed around time to target: how quickly validation loss reaches a set level. On three large-scale OGB benchmarks, GraphMuon gets there faster than AdamW, with speedups ranging from one point three two times up to two point four two times. A broader sweep covers fifteen configurations, each pairing a dataset with an architecture, and spanning node classification, inductive graph classification and graph regression. GraphMuon has the lower time to target in thirteen of those fifteen, while its predictive performance stays competitive or better.

At bottom, the argument concerns where an optimizer should measure the size of a step. Muon measures it on the weight matrix. In a graph network the weights act on representations that propagation has already reshaped, so GraphMuon measures the step by what it does to the layer's output, under the metric those propagated features define. That single relocation is the contribution, and the speedups over AdamW are the evidence offered for it.