---
title: "Module-Aware Optimization for Graph Neural Networks"
source: "arXiv AI/CL/LG（论文）"
source_tier: "t2"
url: "https://arxiv.org/abs/2611.08379v1"
permalink: "https://feed.astralor.com/items/itm_a542b66df001bc09e9a327db"
published: "2026-11-05T17:56:33.000Z"
score: 60
channel: "tech"
tags: ["Module-Aware Optimization","Graph Neural Networks","Different","Adaptive","Adam","AdamW","Muon"]
---

# Module-Aware Optimization for Graph Neural Networks

## 速览

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 thro…

## 关键点

- 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.

## 全文

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.

---

引用前请核对原文;本文由星际流动整理:https://feed.astralor.com/items/itm_a542b66df001bc09e9a327db