When Compilation Breaks Your GNN: A Numerical Stability Perspective
Jiawei Gu, Zechao Li
Abstract
Deep learning compilers like torch.compile accelerate GNNs through operator fusion and computation reordering, yet often introduce numerical instability that causes models to diverge under hyperparameters that work in eager mode. We trace this fragility to the interaction between compiler optimizations and GNN structure. Compilers reorder floating-point additions assuming associativity, but rounding errors from different summation orders, though typically negligible, can be catastrophically amplified by ill-conditioned computations. We show that GNN neighbor aggregation is structurally prone to such ill-conditioning: power-law degree distributions create highly variable aggregation lengths, while normalized features induce message cancellation during summation. This yields a long-tailed distribution of condition numbers across nodes, where a critical minority of aggregations amplify rounding differences into macroscopic output deviations. Based on this analysis, we propose Numerics-Aware Graph Compilation (NAGC), which estimates per-aggregation numerical risk through lightweight profiling and selectively applies aggressive optimizations only to well-conditioned operations. Experiments on Open Graph Benchmark datasets show that NAGC preserves over 95% of compilation speedups while reducing numerical deviation by 10 to 100 times and eliminating compilation-induced training failures.
Ask about this paper
Ask your agent about it.
Lune has read the top-tier papers around this one, so every answer names the papers it rests on.
Related papers
- The Hidden Fragility of GNNs: How Graph Structure Amplifies Numerical ErrorsJiawei Gu, Ziyue QiaoKDD 2026
- AGO: Boosting Mobile AI Inference Performance by Removing Constraints on Graph OptimizationZhiying Xu, Hongding Peng, Wei WangINFOCOM 2023 · 1 citation
- PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph CompilationJason Ansel, Edward Z. Yang, Horace He, Natalia Gimelshein et al.ASPLOS 2024 · 693 citations
- : Aggregation-Aware Quantization for Graph Neural NetworksZeyu Zhu, Fanrong Li, Zitao Mo, Qinghao Hu et al.ICLR 2023
- GNNerator: A Hardware/Software Framework for Accelerating Graph Neural NetworksJacob R. Stevens, Dipankar Das, Sasikanth Avancha, Bharat Kaul et al.DAC 2021 · 22 citations
