Module-Aware Optimization for Graph Neural Networks

Guy Hadad; Haggai Roitman; Moshe Eliasof; Bracha Shapira

cs.LGNov 5, 2026

Module-Aware Optimization for Graph Neural Networks

Authors: Guy Hadad, Haggai Roitman, Moshe Eliasof, Bracha Shapira

Organizations: Ben-Gurion University of the Negev, Beer-Sheva, Israel

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.

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 345 papers

May 29, 2026

On Efficient Scaling of GNNs via IO-Aware Layers Implementations

Graph Neural Networks (GNNs) are bottlenecked by sparse, irregular memory access. Popular frameworks such as DGL and PyTorch Geometric support general message passing, but complex layers often materialize edge-wise intermediates, increasing memory traffic and limiting scalability on large graphs. We take an I/O- and arithmetic-intensity--centric view and show that widely used layers fall into three kernel families: SpMM-based convolutions, reduction-based aggregations, and attention-based layers (GATv2/Graph Transformer). For each family, we develop GPU kernels that reduce data movement, improve locality, and remain robust across realistic graphs. We also study graph reordering and find that its impact depends on the kernel mapping: it benefits neighbor-parallel (gather-dominated) kernels more consistently than feature-parallel designs. Empirically, our fused attention kernels reach up to

\textbf{3.9}\times

speedup for Graph Transformer (median

\textbf{1.6}\times

), with Tensor Core (block-sparse) variants up to

\textbf{7.3}\times

on locally dense graphs; for GATv2 we reach up to

\textbf{8.5}\times

speedup (median

\textbf{2.0}\times

) while reducing peak memory by up to

\textbf{76}\times

(median

\textbf{6}\times

). Our degree-aware reduction kernels achieve up to

\textbf{10}\times

speedup (median

\textbf{2.6}\times

). For SpMM-based layers, properly cached cuSPARSE achieves up to

\textbf{8}\times

speedup over DGL and outperforms evaluated custom baselines in the majority of evaluations. We release our implementations as drop-in replacements to support reproducible, hardware-aware GNN acceleration.

Daria Fomina, Daniil Krasylnikov, Alexey Boykov +3

1Yandex · 2HSE University · 3ITMO University

Graph Attention Networks GPU Kernel Optimization GPU Acceleration High-Performance Computing GPU Kernel Generation Graph Transformer

Ask PDF

May 18, 2025

Never Skip a Batch: Dense Learning of Temporal GNNs via Adaptive Pseudo-Supervision

Temporal graph networks suffer from irregular supervision in realworld dynamic graphs, as most minibatches contain few labeled events. The lack of labels leads to high-variance gradient updates and, consequently, slow wall-clock convergence. To constructively reduce sparsity, our Moving-Averaged Labels (MAL) assigns soft pseudo-targets based on past supervised signals using a running label distribution while leaving the loss and the model architecture unchanged. Thus, supervision gaps are replaced with informative signals independent of a temporal graph model and the message passing or memory components used. Theoretical analysis supports our insight that aggregating historical supervision into moving average targets reduces stochastic gradient variance, yielding faster convergence under mild assumptions. Experimentally, for TGNv2 and DyRepv2 (our modification of DyRep) models, MAL boosts predictive performance, establishing a new SOTA, and improves time-to-accuracy (on average 6x faster to reach the top score) for a common suite of Temporal Graph Benchmark datasets.

Alexander Panyshev, Dmitry Vinichenko, Oleg Travkin +2

Risk AI Research Lab Moscow Independent Research Institute of Artificial Intelligence Moscow, Russia · Risk AI Research Lab Moscow, Russia · Risk AI Research Lab Applied AI Institute Moscow, Russia

Temporal GNNs Temporal Link Prediction Pseudo-Labeling Semi-Supervised Learning

Ask PDF

Sep 26, 2026

Efficient Dynamic Algorithms for Graph Neural Networks with Non-Linear Propagation

Graph Neural Networks (GNNs) are widely used for representation learning on graphs, but most methods assume static topologies, making them inefficient on evolving networks where edges change over time. Existing dynamic approaches either model graph evolution through temporal GNN architectures without focusing on efficient dynamic maintenance, or are restricted to linear propagation models based on Personalized PageRank. In this work, we study how to efficiently maintain node representations for non-linear GNN propagation under edge insertions and deletions. The propagation has no learned parameters, and only a classifier applied afterward is trained. For a broad class of standard activation functions, we develop a residual-based dynamic algorithm that selectively propagates local errors via push operations, maintaining an approximation to the evolving fixed point without full recomputation. We prove that our method achieves amortized

O(1/ε)

update time per graph change under a degree-normalized error guarantee. Our approach uses a potential-based analysis in a degree-scaled norm and, in contrast to prior work on the linear case, requires no randomness assumptions on either the update sequence or the input vector. For the linear special case, we additionally provide an exact dynamic algorithm via low-rank matrix inverse updates. Experiments on benchmark datasets show that incorporating non-linearity improves accuracy while preserving efficient update performance, yielding a scalable and theoretically grounded method for maintaining this propagation on dynamic graphs.

Kiarash Banihashem, MohammadTaghi Hajiaghayi, Mahdi JafariRaviz +2

Google Research

Graph Neural Networks Dynamic Graph Learning

Ask PDF

Abstract

Neural Network Optimization

Explore similar work

On Efficient Scaling of GNNs via IO-Aware Layers Implementations

Never Skip a Batch: Dense Learning of Temporal GNNs via Adaptive Pseudo-Supervision

Efficient Dynamic Algorithms for Graph Neural Networks with Non-Linear Propagation