PaperlayerView exact revision on arXiv

Source-linked AI summary

Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs

Zhewei Chen, Hao Zhu, Jiaojiao Jiang, Ahad N. Zehmakan

arXiv:2610.10520v1cs.LG

TL;DR

GNN-to-MLP distillation lacks a clear way to preserve graph-induced teacher geometry, despite targeting graph-free inference. G^2MLP uses curvature-guided prediction and representation alignment; it improves over graph-free baselines across node-classification benchmarks and reduces the rank gap in both regimes.

  • Problem

    Existing GNN-to-MLP methods mainly transfer node-wise predictions or confidence-based weights without specifying where the student should preserve graph-induced geometry.

    Source: §1 · p. 1

  • Method

    G^2MLP uses Ollivier–Ricci curvature to allocate prediction-level and representation-level alignment during training while keeping the deployed model a standard MLP.

    Source: §1 · p. 2, §6 · p. 9

  • Results

    Across multiple node-classification benchmarks, G^2MLP consistently outperforms graph-free baselines and closes the embedding rank gap in both spectral regimes.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The same recipe transfers to Graph Transformer teachers and link prediction without modification.

    Source: §6 · p. 9

    Exact Ollivier–Ricci curvature precomputation scales poorly to very large graphs, though Sinkhorn-regularized transport and Forman–Ricci are alternatives.

    Source: §5.5 · p. 9

Abstract

from arXiv · show

GNN-to-MLP distillation aims to retain the predictive accuracy of a message-passing teacher while deploying a graph-free MLP at inference. Existing methods mainly transfer node-wise predictions or use confidence-based reweighting, but they do not specify where the student should preserve the teacher's graph-induced geometry. We show that this omission leads to two spectral failure modes in the student's representation space. On sparse graphs, the student suffers from spectral underfit, missing high-energy teacher directions concentrated near boundary regions. On dense graphs, it suffers from spectral overfit, retaining spurious directions that the teacher has collapsed through aggregation. Motivated by an energy-weighted teacher-student alignment objective, we propose Graph Geometry-aware MLP (G^2MLP), a training-time distillation framework guided by Ollivier-Ricci curvature. Curvature identifies where the two spectral errors concentrate and is used to allocate supervision between prediction-level and representation-level alignment. The deployed model remains a standard MLP and requires no graph access at inference. Across node-classification benchmarks, G^2MLP consistently improves over graph-free distillation baselines, reduces the teacher-student rank gap in both regimes, and transfers without architectural changes to Graph Transformer teachers and link prediction.

1 Introduction

GNN-to-MLP distillation can make inference graph-free, but uniform prediction matching does not specify where the student should preserve graph-induced geometry. G^2MLP uses Ollivier–Ricci curvature to allocate prediction- and representation-level alignment around two spectral failure modes.

1 Introduction

  • Existing distillation methods commonly match teacher logits and may add structural priors, but usually apply the same type of pressure regardless of node topology.The passage cites positional encodings, codebooks, and neighborhood alignment as examples of added structural cues.Source: §1 · p. 1
  • On sparse graphs, students miss teacher directions concentrated near boundaries; on dense graphs, they retain variation the teacher has smoothed away.The paper names these failures spectral underfit and spectral overfit, respectively.Source: §1 · p. 1, Figure 1 · p. 2
  • The proposed energy-weighted alignment rationale says pointwise KL can match logits while leaving teacher–student representation geometries different.Ollivier–Ricci curvature provides a local signal: low-curvature edges mark boundary regions, while high-curvature edges mark strongly smoothed regions.Source: §1 · p. 2
  • G^2MLP assigns stronger prediction supervision to low-curvature nodes and stronger representation alignment to high-curvature nodes, using graph-dependent quantities only during training.The deployed model remains a vanilla MLP.Source: §1 · p. 2

2 Related Work

Prior work distills GNN knowledge through logits, structural cues, and alternative objectives, while related optimal-transport and curvature methods address different parts of the problem. G^2MLP instead uses curvature to allocate distillation pressure between teacher and student.

2 Related Work

  • GLNN established soft-label KL distillation from a GNN teacher to an MLP student; later methods add positional encodings, codebooks, neighborhood alignment, or hardness-aware weighting [Zhang et al., 2022, Hinton et al., 2015].The passage says these approaches either treat nodes uniformly or weight them using teacher confidence, rather than topology-linked failure modes.Source: §2 · p. 2
  • WCoRD and EMD+IPOT use discrete Wasserstein distances with cross-instance matching, while WKD uses continuous Gaussian W2 for feature distillation [Chen et al., 2021, Lohit and Jones, 2022, Lv et al., 2024].G^2MLP extends the Gaussian-W2 template to neighborhood representations and weights it by curvature.Source: §2 · p. 2
  • Prior graph-curvature work modifies GNN structure or training, whereas G^2MLP uses curvature to stratify distillation pressure between teacher and student [Nguyen et al., 2023, Liu et al., 2023, Hevapathige et al., 2025].The cited examples include rewiring, dropout, and adapting network depth.Source: §2 · p. 3

3 Why Uniform Distillation Fails under Curvature Heterogeneity

The analysis relates teacher–student kernel mismatch to two spectral failure modes and motivates curvature-guided supervision. Its local curvature signal is a proxy for where the errors concentrate, not an exact estimate of the teacher-kernel spectrum.

3 Why Uniform Distillation Fails under Curvature Heterogeneity

  • The kernel-motivated analysis links graph-dependent aggregation in the GNN kernel, opposite spectral failures, and Ollivier–Ricci curvature as a local signal for allocating supervision.The resulting energy-weighted target is approximated with curvature-guided alignment.Source: §3 · p. 3
  • The GNN kernel depends jointly on features and graph structure, while the MLP kernel depends only on features; their difference captures message-passing contributions.The analysis is qualitative and does not claim that the full GNN kernel is sparse or that curvature estimates its eigenspectrum exactly.Source: §3 · p. 3
  • For a K-layer message-passing GNN, graph-dependent aggregation is generated through paths in K-hop receptive fields, although feature covariance can keep the full kernel nonzero at distant nodes.The contribution for a node pair depends on overlap and feature similarity of their K-hop neighborhoods.Source: §3 · p. 3

3.1 Two Failure Modes and the Ideal Distillation Objective

  • Pointwise KL does not directly control neighborhood geometry: students matching the same logits may still induce different hidden-space geometries.The proposed energy-weighted target emphasizes high-functional-energy teacher directions, but directly using it requires storing and diagonalizing the full teacher kernel.Source: §3.1 · p. 4

3.2 Curvature Stratifies the Failure Modes

  • On sparse graphs, low-curvature boundary edges associate with missing teacher directions; on dense graphs, high-curvature interior regions associate with student variation weakly supported by the smoothed teacher.The rank gap is an observable diagnostic of representation geometry, not an exact estimator of the teacher-kernel spectrum.Source: §3.2 · p. 4
  • Because the two failures concentrate in different curvature regimes, uniform weighting does not distinguish them; the analysis motivates opposite curvature-dependent emphasis for prediction- and representation-level alignment.Increasing pressure for boundary-discriminative directions also increases pressure on already-smoothed interior regions, and conversely.Source: §3.2 · p. 4

4 G2MLP: Curvature-Guided Graph Wasserstein Alignment

G^2MLP approximates an intractable energy-weighted geometry target with local neighborhood matching, a diagonal Gaussian W2 surrogate, and curvature-dependent weights. These relaxations provide tractable training-time alignment while retaining a graph-free MLP at deployment.

4 G2MLP: Curvature-Guided Graph Wasserstein Alignment

  • The intractable global target is approximated by local neighborhood matching, diagonal Gaussian W2, and Ollivier–Ricci-dependent node weighting; these are tractable relaxations, not exact equivalences.A frozen GNN teacher and an MLP student supply embeddings and logits, while curvature-derived weights allocate alignment during training.Source: §4 · p. 5, Figure 2 · p. 5

4.1 From Global Alignment to Local Neighborhood Measures

  • Local teacher and student representation distributions approximate the global alignment target because the graph-dependent teacher-kernel component is generated through message-passing neighborhoods.The local Wasserstein term responds to neighborhood geometry even when students have identical nodewise predictions.Source: §4.1 · p. 5

4.2 Efficient Neighborhood Alignment via Moment Matching

  • The neighborhood transport problem is replaced by matching first and second marginal moments with a diagonal Gaussian W2 surrogate, reducing per-node cost to O(d).The surrogate may be less accurate for multimodal or strongly heterophilic neighborhoods; splitting feature and logit terms enables curvature-dependent weighting.Source: §4.2 · p. 6

4.3 Curvature-Adaptive Weighting

  • Opposite curvature-dependent weights emphasize prediction alignment at low-curvature boundary nodes and representation alignment at high-curvature interior nodes.The resulting objective is computable in O((|V| + |E|) · d) per iteration; KL preserves probability-simplex calibration while curvature-adaptive Wasserstein alignment matches local geometry.Source: §4.3 · p. 6

5 Experiments

Across node classification, transformer-teacher distillation, and link prediction, G2MLP improves on graph-free baselines while retaining a standard MLP for inference. Its gains come with a curvature-precomputation cost that can scale poorly on very large graphs.

5.1 Experimental Setup

  • For experimental setup, six node-classification benchmarks use transductive and inductive evaluation, with GraphSAGE teachers and feature-only MLP students.The baselines include vanilla MLP, GLNN, KRD, and FF-G2M; positional-encoding and codebook variants are excluded because they depart from the setup.Source: §5.1 · p. 7

5.2 Main Results

  • G2MLP achieves the best graph-free transductive accuracy on all six datasets, beating the strongest baseline by 0.16–1.66 pp.It also exceeds the GraphSAGE teacher on five datasets, with gains up to 4.02 pp on Citeseer; on ogbn-Arxiv it falls short by 6.16 pp.Source: Table 1 · p. 7, §5.2 · p. 7
  • G2MLP has the best graph-free production accuracy on 5/6 datasets, with gains of 0.16–1.60 pp over the strongest baseline.Its inductive gains are 2.07–2.69 pp on Cora, Citeseer, Pubmed, and A-computer; it matches the SAGE teacher on A-computer and exceeds it on Citeseer and Pubmed production scores.Source: §5.2 · p. 7, Table 2 · p. 8

5.3 GT-to-MLP Distillation

  • With Graph Transformer teachers, the MLP student meets or exceeds its teacher on 13/15 cells, and G2MLP improves over GLNN on all 15.The student architecture and inference cost remain unchanged.Source: Table 3 · p. 8, §5.3 · p. 8

5.4 Extension to Link Prediction

  • In link prediction, G2MLP wins on all three datasets in both AUC and AP, with the largest gains over KRD on Cora.Its AUC gains over KRD are +4.70 on Cora, +0.30 on Citeseer, and +0.04 on Pubmed; all three G2MLP rows also exceed the SAGE teacher in AUC.Source: Table 4 · p. 9, §5.4 · p. 9

5.5 Inference Efficiency

  • On Arxiv, G2MLP runs inference 58× faster than the 3-layer GraphSAGE teacher, at 8.6 versus 499.7 ms.Curvature reweighting is training-only, so deployment uses a feature-vector-to-logit MLP pass without graph traversal.Source: Figure 3 · p. 9, §5.5 · p. 9
  • Exact Ollivier–Ricci curvature precomputation scales as O(|E|·d̄^3) and can become prohibitive on very large graphs.The paper describes Sinkhorn-regularized transport and Forman–Ricci curvature as alternatives, and notes that preprocessing is amortized across training runs and deployment.Source: §5.5 · p. 9

6 Conclusion

G2MLP addresses spectral misalignment in GNN-to-MLP distillation while retaining vanilla-MLP inference. The included notation and kernel passage distinguish the graph-dependent GNN kernel from the feature-only MLP kernel.

6 Conclusion

  • G2MLP closes the embedding rank gap in both spectral regimes, outperforms graph-free baselines, and preserves vanilla-MLP inference speed.The conclusion also reports transfer to Graph Transformer teachers and link prediction without modification.Source: §6 · p. 9

A Notation and Preliminaries

B.1 Full GNTK Recursion

  • The MLP NTK is the aggregation-free special case and depends only on the two nodes’ features.Source: §B.1 · p. 12

B.2 Proof of Lemma 3.1 (Aggregation Locality)

The GNN kernel depends on features in closed K-hop neighborhoods, unlike the MLP kernel, which depends only on the paired nodes’ features. Under the stated Lipschitz assumption, ORC bounds the difference between neighboring nodes’ mean features.

B.2 Proof of Lemma 3.1 (Aggregation Locality)

  • The GNN kernel entry depends on features in the nodes’ closed K-hop neighborhoods, whereas the MLP kernel depends only on their own features.The kernel difference therefore captures a graph-induced contribution from paths in those receptive fields.Source: §B.2 · p. 13
  • The locality claim is proved by induction over layers, extending feature dependence from each pair to their neighborhoods at the next layer.The recursion combines covariance, its derivative kernel, and the preceding-layer kernel.Source: §B.2 · p. 13

B.3 Proof of Proposition 3.2 (ORC as a Local Mismatch Proxy)

  • For adjacent nodes, the distance between mean feature embeddings is bounded by L times one minus their Ollivier–Ricci curvature.The proof takes the infimum over couplings of the lazy random-walk distributions and applies the Lipschitz assumption.Source: §B.3 · p. 13

C Derivation of the Multi-Granularity Decomposition

The derivation separates joint Gaussian Wasserstein alignment into feature-side and logit-side terms, then motivates curvature-dependent weighting of the two. The training pipeline precomputes curvature and deploys only the trained MLP.

C Derivation of the Multi-Granularity Decomposition

  • Assuming diagonal covariance within each block reduces the per-node matrix-square-root cost from O((H + C)^3) to O(H + C).Source: C Derivation of the Multi-Granularity… · p. 14
  • The joint Gaussian W2 decomposition is additive across hidden-feature and logit blocks, motivating separate alignment terms.The decomposition relies on a near-zero cross-covariance simplification, which is a modeling assumption rather than an empirical fact.Source: C Derivation of the Multi-Granularity… · p. 14
  • Curvature stratifies logit-side and feature-side alignment in opposite directions: lower curvature emphasizes logits, while higher curvature emphasizes features.The sign-opposite weighting follows the theoretical failure-mode analysis rather than being tuned.Source: C Derivation of the Multi-Granularity… · p. 14

D Training Algorithm

  • Training first fits the GNN teacher and precomputes node-level curvature, then trains the student; inference deploys only the MLP without graph access.Source: D Training Algorithm · p. 14

E.1 Datasets and Splits

The appendix documents benchmark protocols, baselines, implementation, ablations, compute, and Graph Transformer experiments. Its ablation reports that curvature stratification—not W2 alignment alone—accounts for most of the improvement over GLNN.

E.1 Datasets and Splits

  • The six node-classification benchmarks comprise three citation graphs, two Amazon co-purchase graphs, and ogbn-Arxiv, evaluated with CPF or canonical OGB splits.The five CPF datasets use 10 random splits; ogbn-Arxiv uses its time-based split with 10 initialization and dropout seeds.Source: §E.1 · pp. 14–15

E.2 Baselines

  • The baseline comparison uses released student configurations and excludes positional-encoding variants because their preprocessing and input changes prevent like-for-like comparison.Source: §E.2 · pp. 15–16

E.3 G2MLP Implementation and Hyperparameters

  • Student optimizer settings follow released GLNN configurations and are held fixed, while the implementation includes a small learnable projection on some Amazon graphs.The projection can be omitted on the Planetoid graphs.Source: §E.3 · p. 16

E.4 Component Ablation

  • Curvature stratification improves over GLNN by +1.5 to +2.8, while W2 alignment alone is neutral and the two stratified sides are complementary.Source: §E.4 · p. 16, Table 7 · p. 17

E.5 Compute Environment

  • The reported experiments use a single RTX 3090; total wall-clock cost is approximately 80 GPU-hours for node classification, 40 for Graph Transformers, and 15 for link prediction.Source: §E.5 · p. 17

F Graph Transformer Teacher Experiments

F.1 GT Architectures and Configurations

The appendix evaluates Graph Transformer teachers and finds that G2MLP can outperform them in this setting, where limited labels make dense attention difficult to train. It also compares teacher architectures and dense versus sparse GraphGPS attention.

F.1 GT Architectures and Configurations

  • The GraphGPS comparison evaluates vanilla GT, GraphGPS, and NAGphormer teachers trained for 500 epochs, with a two-layer MLP student whose hidden dimension matches its teacher.The GraphGPS teacher uses dense N × N attention by default; the appendix also evaluates a sparse edge-softmax variant.Source: §F.1 · p. 17

F.2 Why Students Can Exceed Their GT Teachers

  • The MLP student meets or exceeds its Graph Transformer teacher on 13/15 cells, which the authors attribute to limited labels and overfit attention noise.Distillation regularizes the student by transferring teacher predictions while avoiding some of the teacher’s attention noise.Source: §F.2 · p. 18

F.3 Why NAGphormer Is Often the Strongest GT Teacher

  • NAGphormer is the strongest teacher on three of five benchmarks, while vanilla GT leads on Citeseer and A-computer.The authors connect NAGphormer’s performance to hop-token summaries that regularize attention and encode graph structure.Source: §F.3 · p. 18

F.4 GraphGPS: Sparse vs. Dense Attention

  • G2MLP improves over GLNN in nine of ten GraphGPS attention-modality cells and surpasses both dense and sparse teachers on all five datasets.The sparse GraphGPS teacher is stronger than the dense teacher on four graphs; A-photo is a tie.Source: §F.4 · p. 18, Table 8 · p. 18

G Link Prediction: Pure-Distillation Ablation

The pure-distillation ablation tests link prediction without edge labels and shows that curvature-aware geometry can preserve LP-relevant structure when uniform embedding matching fails. The contrast is strongest on sparse Cora.

Setup

  • The ablation removes edge-BCE so the student receives only the teacher’s frozen embeddings through an MSE loss, with no edge labels available.The experiments use LCC-standardized CPF graphs and a two-layer GraphSAGE teacher and two-layer MLP student.Source: G Link Prediction: Pure-Distillation… · p. 19

G2MLP-LP hyperparameter search

  • G2MLP-LP tunes KD weight, CAW weight, and the feature-side curvature direction per dataset, while holding the logit-side direction fixed.Cora selects the bridge direction; Citeseer and Pubmed select the community direction.Source: G Link Prediction: Pure-Distillation… · p. 19

Pure regime: ablating the edge-BCE term

  • 81.77 AUC on Cora: G2MLP-LP’s pure student matches the teacher, while GLNN and KRD pure students fall to 56–57 AUC.On Citeseer and Pubmed, GLNN and KRD pure-KD students already exceed the teacher; adding edge-BCE changes Citeseer AUC by less than one point and adds 1.4–1.7 AUC on Pubmed.Source: G Link Prediction: Pure-Distillation… · p. 19, Table 9 · p. 20

Cross-regime gap

  • On Cora, the multi-task versus pure-regime AUC gap is +22.4 for GLNN and +22.0 for KRD, compared with +2.1 for G2MLP-LP.The authors interpret the smaller G2MLP-LP gap as evidence that curvature-aware geometry already encodes structural signal that uniform MSE needs edge-BCE to supply.Source: G Link Prediction: Pure-Distillation… · p. 20

H Limitations and Future Directions

The paper identifies limitations in its diagonal-Gaussian neighborhood surrogate, theoretical guarantees, and empirical coverage. These leave open how well the method handles multimodal neighborhoods, when curvature-adaptive alignment is necessary, and how it performs on heterophilic graphs.

Diagonal Gaussian neighborhood surrogate

  • The diagonal-Gaussian surrogate is exact only for Gaussian neighborhoods with diagonal covariance and may under-represent alignment costs for multimodal distributions.The authors report it is sufficient on evaluated benchmarks, while mixture or full-covariance surrogates remain open; full covariance would raise per-node cost from O(H + C) to O((H + C)^3).Source: H Limitations and Future Directions · p. 20

Upper-bound nature of the theoretical results

  • The kernel-locality and curvature-proxy results are upper bounds and structural statements, not lower bounds proving curvature-stratified regimes necessary to close spectral mismatch.Which graphs and teacher kernels require curvature-adaptive rather than uniform alignment remains an open theoretical question.Source: H Limitations and Future Directions · p. 20

Limited heterophilic evaluation

  • Evaluation is limited to homophilic node-classification benchmarks, leaving strongly heterophilic graphs unevaluated.The authors note that multimodal, label-mixed neighborhoods may challenge the diagonal surrogate and that curvature-stratified weighting may need sign-flipped variants.Source: H Limitations and Future Directions · p. 20

Related papers

Loading 2610.10520v1…