Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs

arXiv:2610.10520 · cs.LG · Submitted 2026-10-07 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs".

Jane: The gist The proposed Graph Geometry-aware MLP (G2MLP) framework addresses spectral failure modes in GNN-to-MLP distillation by using Ollivier–Ricci curvature to guide supervision,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we’re looking at this paper, "Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs," and it’s about trying to take a powerful Graph Neural Network teacher and get its knowledge into a regular MLP student without needing the graph structure at all during inference.

Jane: That sounds tricky, Tom, because when you strip away the graph structure for inference, you usually lose all that rich geometric information the GNN learned. What’s the main problem they’re pointing out here?

Tom: Well, they say that existing methods fail because they don't tell the student *where* to keep the teacher's geometry during distillation. They show this leads to two spectral failure modes: spectral underfit on sparse graphs, and spectral overfit on dense graphs.

Lu: That points directly at a mismatch between the GNN kernel and the MLP kernel, which is what they analyze by looking at how those kernels differ in their graph-dependent aggregation terms <ref:2610.10520#pg3>.

Meng: So if we look at that mismatch, how do we actually fix it without making inference slow? That’s the real engineering question for me.

Lalam: The paper motivates a new energy-weighted alignment objective, L*, which characterizes this teacher–student geometric mismatch by emphasizing teacher directions with high functional energy <ref:2610.10520#pg2>.

Jane: But how do you actually use that? They propose this framework called Graph Geometry-aware MLP, or G2MLP, which uses Ollivier–Ricci curvature to guide the supervision during training <ref:2610.10520#pg1>.

Tom: Exactly. G2MLP uses the curvature—specifically Ollivier–Ricci curvature—to approximate that energy-weighted target L* through local, curvature-guided alignment <ref:2610.10520#pg2>.

Lu: The paper details how low-curvature nodes get stronger prediction-level supervision to fix underfit, and high-curvature nodes get stronger representation-level alignment to reduce overfit <ref:2610.10520#pg2>.

Meng: That sounds complex to implement during training. What does that look like practically for a production system? Is it adding too much complexity?

Tom: It does introduce some components, but they simplify the problem by using local alignment, matching representation distributions with Graph Wasserstein distance <ref:2610.10520#pg3>. They use moment matching to get a scalable surrogate, reducing the per-node cost to big O of d <ref:2610.10520#pg3>.

Jane: So if we have to implement this during training, what is the actual computational cost for a system like ours? Are we looking at something that breaks at scale?

Meng: The paper says the resulting Curvature-Adaptive Wasserstein objective, LCAW, is computable in big O of n plus e times d per iteration <ref:2610.10520#pg2>. That’s a good thing because it keeps the training time manageable.

Tom: And the results show it works. Across multiple node-classification benchmarks, G2MLP consistently improves over graph-free distillation baselines and reduces the teacher–student rank gap in both regimes <ref:2610.10520#pg2>.

Lu: They also report that in a production setting, G2MLP achieves the best graph-free accuracy on all six datasets, and it surpasses the GraphSAGE teacher on five of those six datasets, with gains up to four point zero two percentage points on Citeseer <ref:2610.10520#pg2>.

Title and authors: Jane: That’s a substantial improvement over just taking the raw predictions from the GNN and feeding them into an MLP Tom that's what we were trying to avoid in the first place.

Tom: And it generalizes beyond node classification too; they show it wins on all three datasets for link prediction in both AUC and AP <ref:2610.10520#pg2>.

Meng: So, if I’m thinking about deploying this, what about inference speed? Can we actually deploy something that's fast enough for real-time scoring?

Tom: That’s a huge part of the selling point. The deployed model is still a standard MLP and requires no graph access at inference <ref:2610.10520#pg1>. They show it yields a fifty-eight times speedup over the three-layer GraphSAGE teacher under standard layer-wise full inference <ref:2610.10520#pg3>.

Jane: Fifty-eight times is a massive jump for an inference step, especially when dealing with large graphs like ogbn–Arxiv <ref:2610.10520#pg3>.

Lu: It places G2MLP right on that speed–accuracy frontier alongside the GNN teachers, which is important because it keeps the representation quality high <ref:2610.10520#pg3>.

Meng: That’s what I need to hear. A massive speedup without sacrificing accuracy because we have to be careful about the deployment cost at scale.

Lalam: The authors also noted that this recipe transfers to Graph Transformer teachers and link prediction without any modifications <ref:2610.10520#pg1>.

Tom: So, to wrap up on "Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs," the paper shows how guiding distillation with Ollivier–Ricci curvature can resolve spectral underfit and overfit issues by using local supervision based on where the geometry mismatches are most concentrated <ref:2610.10520#pg2>.

Jane: It means we can get better performance from GNN teachers while keeping the inference part as simple as a vanilla MLP without needing to maintain a full graph structure during scoring Tom.

Lu: The paper’s main contribution is showing that the energy-weighted alignment objective L* is key, and curvature provides the signal to make that objective tractable through local alignment and moment matching <ref:2610.10520#pg3>.

Meng: But they did flag a limitation, which is important for me: they say the diagonal Gaussian neighborhood surrogate used to approximate the moments is only exact if the neighborhood distribution itself is already Gaussian with a diagonal covariance <ref:2610.10520#pg4>.

Tom: That means their empirical results are strongest on node classification benchmarks where those neighborhood Gaussians are statistically reliable <ref:2610.10520#pg4>.

Jane: It also says that the theoretical results they have are upper bounds and structural statements, so they haven't established matching lower bounds for everything <ref:2610.10520#pg4>.

Lu: Future work would definitely need to systematically evaluate which graphs and which teacher kernels specifically demand this curvature-adaptive treatment <ref:2610.10520#pg4>.

Meng: For production, I’d want to see that generalization tested on more heterophilic scenarios, maybe with mixture-of-Gaussian surrogates or learned sign conventions <ref:2610.10520#pg4>.

Lalam: This paper is about using local graph signals to solve a global distillation problem, and it seems like a solid step toward making GNN knowledge accessible through simpler models <ref:2610.10520#pg1>.

The paper's summary: Tom: So, to wrap up on "Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs," this paper is essentially about how we can get the geometric knowledge from a Graph Neural Network teacher into a regular MLP student without needing the graph structure during inference.

Jane: That's right, Tom. It tackles that gap between the structured world of graphs and the simpler world of neural networks, specifically focusing on how that mismatch causes errors when you try to distill them.

Lu: The core idea is identifying two specific failure modes: spectral underfit when dealing with sparse graphs and spectral overfit when you have dense graphs.

Meng: So if we look at what’s causing the error, it seems like the GNN kernel and the MLP kernel are just not speaking the same language about graph structure.

Tom: Exactly. The GNN has these aggregation terms that depend on neighbors, which disappear in a standard MLP, and that difference causes these two opposite problems with spectral performance.

Jane: And to fix this mismatch during training, they introduce Ollivier–Ricci curvature as the guiding signal for supervision.

Lu: They derive an energy-weighted alignment objective called L*, which measures exactly how geometrically different the teacher and student kernels are in terms of functional energy.

Tom: Then G2MLP uses that curvature to decide where to put the focus during training, allocating more supervision based on whether a node is in a low-curvature boundary region or a high-curvature interior region.

Jane: So, instead of applying uniform supervision everywhere, you tailor the training based on the local geometry of the graph itself.

Lu: This creates this Curvature-Adaptive Wasserstein objective, LCAW, which they show is computable efficiently in big O of n plus e times d per iteration.

Meng: That’s good for scale. If we can keep that complexity low during training, it makes deployment much more viable when we think about production systems.

Tom: And the results are solid; G2MLP beats graph-free distillation baselines on node classification and even surpasses the GraphSAGE teacher on five out of six datasets.

Jane: It also extends to link prediction, where it wins on all three datasets in both AUC and AP metrics, showing that this geometric guidance is really versatile.

Lu: The real practical win here is maintaining performance while achieving a massive speedup at inference; the student model is just a standard MLP with no graph access required.

Tom: They show a fifty-eight times speedup over the GraphSAGE teacher under standard layer-wise full inference, putting it right on that speed–accuracy frontier.

Jane: It means we can get this enhanced knowledge transfer without needing to maintain the heavy graph structure for every prediction.

Lu: The implication is that you can use these powerful GNN teachers to train a simple MLP student, and the geometry-aware distillation process makes it work better than just blindly copying weights.

Tom: It transfers over to other architectures too; they say this same method works with Graph Transformer teachers and for link prediction without any extra modification.

Jane: So, it’s not just a niche fix for one type of graph or one task; the geometric distillation recipe seems quite general across different AI structures.

Lu: The caveat is that their moment matching uses a diagonal Gaussian approximation, which only works perfectly if the neighborhood distribution is already Gaussian with a diagonal covariance.

Meng: So, their strong empirical results are mostly on node classification where those neighborhood distributions are statistically reliable, but they aren't claiming it works universally for every graph type yet.

Tom: They also admit their theoretical results are upper bounds and structural statements, so they haven't established matching lower bounds for the performance gain yet.

Jane: The future work needs to be about figuring out exactly which graphs and which teacher kernels actually benefit most from this specific curvature-adaptive treatment.

The paper's improvements: Tom: So, focusing on what they actually improved in G2MLP, it's about making sure this distillation process is robust enough to handle different graph shapes and sizes better than previous methods.

Jane: They introduce a curvature-guided distillation framework that specifically targets those spectral failure modes we talked about earlier, spectral underfit and overfit.

Lu: The main improvement is using Ollivier–Ricci curvature to dynamically adjust the supervision, so low-curvature nodes get stronger prediction alignment and high-curvature nodes get better representation alignment.

Tom: That’s a clean mechanism for training; it means the model learns to focus its effort exactly where the geometric mismatch is most pronounced in that specific graph structure.

Jane: So, what does this mean practically for a user deploying this on a large dataset? It means you get better accuracy across the board while keeping your inference cost low because you’re still using a standard MLP.

Meng: For production, that’s huge because it means we don't have to worry about the model failing spectacularly when it encounters sparse or dense sections of the graph; the framework adapts automatically.

Lu: They also show this approach works even when moving from node classification to link prediction tasks, where they achieve wins in both AUC and AP on three different datasets.

Tom: That generalization is interesting because it suggests this geometric distillation technique isn't just a trick for one specific task; it's a way to transfer knowledge across related graph problems.

Jane: It changes how we think about knowledge transfer from GNNs to MLPs, moving away from just copying raw outputs toward an intelligent, geometrically informed distillation process.

Lu: The caveat they put forward is that their moment matching approximation is only exact if the neighborhood distribution itself has a diagonal covariance, which limits where their theoretical guarantees are strongest.

Meng: That means we have to be careful when we try to apply this to graphs where the neighbor distributions are really weird or complex, because those approximations might introduce errors.

Tom: But the empirical evidence is strong; they consistently show a reduction in that teacher–student rank gap across multiple benchmarks, which is exactly what we want in a distillation setup.

Jane: It’s about taking something highly structured—the graph geometry—and using it to intelligently shape how the student model learns from its teacher, which is a really sophisticated way to handle knowledge gaps.

Conclusion: Tom: So, to wrap up on "Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs," this paper shows how using Ollivier–Ricci curvature guides distillation to fix spectral underfit and overfit problems by tailoring supervision based on local graph geometry.

Jane: It really boils down to making sure the student model learns the right things at the right places during training, and it does that by looking at those curvature signals.

Lu: The implication is that you can get a much better performance from a GNN teacher when transferring knowledge to an MLP student without sacrificing inference speed or accuracy.

Meng: It changes how we approach distillation in production because it gives us a principled way to handle the inherent structural differences between graph-based and graph-free models.

Tom: The key part for me is that they achieve this with an LCAW objective that’s still computationally cheap, only big O of n plus e times d per iteration.

Jane: That efficiency is what makes it viable; we get better results without needing a massive increase in training time or computational resources.

Lalam: I'm pretty sure this geometric distillation recipe will be applicable to other GNN architectures, so we should expect to see this approach applied widely across different graph structures.

Lu: My only caution is that the approximation for the neighborhood moments relies on a diagonal Gaussian covariance, which means our strongest results are currently on node classification where those distributions are reliable.

Meng: That’s fair; if we want it to be robust for any arbitrary graph structure, we’ll need to investigate what happens when those assumptions break down in a more heterophilic setting.

Tom: Exactly, so the next step is testing this against harder scenarios where the geometric assumptions might not hold true.

Jane: That leads us perfectly into how we can test these models for memorization in diffusion models, which is another area where geometry plays a role.

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

Australian National University · Data61, CSIRO · University of New South Wales

cs.LG

Submitted: 2026-10-07

Updated: 2026-10-07

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 90/100

The gist: The gist The proposed Graph Geometry-aware MLP (G2MLP) framework addresses spectral failure modes in GNN-to-MLP distillation by using Ollivier–Ricci curvature to guide supervision, resulting in

Key concepts

Spectral Failure Modes
These are two problems occurring during GNN-to-MLP distillation. On sparse graphs, the student underfits because it misses important directions near boundaries. On dense graphs, it overfits by retaining spurious directions that the teacher has already aggregated away.
Ollivier–Ricci Curvature (ORC)
ORC provides a local geometric signal for the graph structure. The framework uses this curvature to decide where to focus distillation supervision. Low-curvature nodes get prediction-level guidance, while high-curvature nodes get representation-level guidance.
Curvature-Adaptive Wasserstein Objective (LCAW)
This is the mathematical objective function used during training. It combines local alignment (using Graph Wasserstein distance) with curvature stratification to create a scalable loss that matches the energy-weighted target, guiding the student's learning process effectively.

Terminology

Summary

The gist The proposed Graph Geometry-aware MLP (G2MLP) framework addresses spectral failure modes in GNN-to-MLP distillation by using Ollivier–Ricci curvature to guide supervision, resulting in improved performance across various benchmarks while maintaining a graph-free inference cost.

Identifying Spectral Failure Modes

The paper identifies two spectral failure modes arising from the mismatch between the GNN kernel and the MLP kernel during knowledge distillation On sparse graphs, the student suffers from spectral underfit, missing high-energy teacher directions concentrated near boundary regions. Conversely, on dense graphs, it suffers from spectral overfit, retaining spurious directions that the teacher has collapsed through aggregation.

Deriving the Alignment Objective

The authors derive an energy-weighted alignment objective L∗ that characterizes teacher–student geometric mismatch and motivates curvature-dependent allocation of distillation supervision. This objective is defined as L∗ = fT − fˢ2K(K) = XiλKiˆfTi − ˆfˢi2, where λKi emphasize teacher directions with high functional energy.

Curvature-Guided Distillation Framework (G2MLP)

G2MLP is proposed as a training-time distillation framework guided by Ollivier–Ricci curvature. It uses ORC to approximate the energy-weighted target L∗ through local, curvature-guided alignment. Specifically, low-curvature nodes receive stronger prediction-level supervision to correct spectral underfit, while high-curvature nodes receive stronger representation-level neighborhood alignment to reduce spectral overfit. The resulting Curvature-Adaptive Wasserstein objective is LCAW = 1/V Xv∈V βv W22,diagνeT h,v, νeS h,v + αv W22,diagνeT z,v, νeS z,v, which is computable in O((V + E) · d) per iteration.

Experimental Results and Generalization

Across multiple node-classification benchmarks, G2MLP consistently improves over graph-free distillation baselines, reducing the teacher–student rank gap in both regimes. In the production setting, G2MLP achieves the best graph-free accuracy on all six datasets. Furthermore, G2MLP further surpasses the GraphSAGE teacher on five of the six datasets, with gains up to 4.02 pp on Citeseer. The framework also generalizes to link prediction, where G2MLP wins on all three datasets in both AUC and AP.

Inference Efficiency

The deployed model remains a standard MLP and requires no graph access at inference. Figure 3 shows that G2MLP yields a 58× speedup over the 3-layer GraphSAGE teacher under standard layer-wise full inference. This places G2MLP on the speed–accuracy Pareto frontier alongside the GNN teachers.

Conclusion

G2MLP closes the embedding rank gap in both regimes and consistently outperforms graph-free baselines while preserving the inference speed of a vanilla MLP. The same recipe transfers to Graph Transformer teachers and to link prediction without modification.

How it works

The paper identifies spectral underfit and spectral overfit as two failure modes of GNN-to-MLP distillation, and connects them to boundary and interior graph regions. The teacher kernel contains graph-dependent aggregation terms absent from the MLP kernel, which mismatch appears as two opposite spectral failure modes. Ollivier–Ricci curvature (ORC) provides a local graph signal for allocating supervision where these modes concentrate.

Key Components of G2MLP

The framework utilizes several components to approximate the intractable energy-weighted target L∗.

** Local Alignment:**

The authors approximate the energy-weighted alignment in L∗ by matching local teacher and student representation distributions using Graph Wasserstein distance (LGW). This objective encourages agreement in both local means and local dispersion.

** Moment Matching:**

To obtain a scalable surrogate, the paper matches the first two marginal moments of the teacher and student neighborhood measures using a diagonal Gaussian approximation. This reduces the per-node cost to O(d).

** Curvature Stratification:**

The curvature-adaptive weights αv and βv are defined using ORC as a local proxy for allocating the alignment budget. These weights emphasize prediction-side alignment at low-curvature boundary nodes and representation-side alignment at high-curvature interior nodes.

Limitations

The diagonal Gaussian neighborhood surrogate is exact only when the underlying neighborhood distribution is itself Gaussian with diagonal covariance. Furthermore, the theoretical results are upper bounds and structural statements, meaning they do not establish matching lower bounds. The empirical evaluation is restricted to homophilic node-classification benchmarks where neighborhood Gaussians are statistically reliable.

Future Directions

A complete characterization of which graphs and which teacher kernels demand curvature-adaptive treatment is an open theoretical question. A systematic heterophilic evaluation, potentially with mixture-of-Gaussian surrogates and learned sign convention, is left to future work.

Implementation Details

The training pipeline involves two stages: Stage 1 trains a standard GNN teacher and precomputes node-level Ollivier–Ricci curvature once as a graph-level preprocessing step. Stage 2 trains the MLP student under the curvature-adaptive objective. The stop-gradient normalization sg(LCAW) + ϵsg in line 10 keeps the CAW gradient direction unchanged but rescales its magnitude to match entropy-scaled losses. The constant ϵsg = 10−8 is fixed across all experiments. The only G2MLP-specific knobs are λCAW ∈ [1e−2, 5e−1], CAW TRANSFORM ∈ [NONE, STUDENT, BOTH], set per dataset. The ORC is computed with lazy-walk parameter α = 0.5, and the detach-normalized LCAW uses ϵsg = 10−8. The architecture of the student MLP is a 2-layer MLP with hidden dimension matched to the teacher. The GNN teacher used is GraphSAGE throughout, following GLNN, NOSMOG. The experiments are conducted on a single NVIDIA GeForce RTX 3090 (24 GB GDDR6X, Ampere) on a shared SLURM cluster. The total wall-clock cost across all reported runs is approximately ≈80 GPU-hours for the node-classification experiments. The inference time on Arxiv yields a 58× speedup over the 3-layer GraphSAGE teacher under standard layer-wise full inference.

References

The paper references numerous prior works including GLNN [Zhang et al., 2022], Ollivier–Ricci curvature [Ollivier, 2009], and Wasserstein distance methods like WCoRD [Chen et al., 2011]. The GNN-to-MLP distillation paradigm is compared against baselines such as vanilla MLP, GLNN [Zhang et al.

Improvements for AI systems

  1. textbf Curvature-Guided Distillation for Spectral Alignment in Graph-Free Inference: G2MLP Implementation (Node Classification)

This improved system can achieve superior graph-free inference accuracy by explicitly addressing two spectral failure modes, namely spectral underfit on sparse graphs and spectral overfit on dense graphs. The system achieves this by using Ollivier–Ricci curvature to allocate supervision: low-curvature edges receive stronger prediction-level supervision to correct underfit, while high-curvature edges receive stronger representation-level alignment to reduce overfit.

  1. textbf Enhanced Link Prediction Performance via Curvature-Adaptive Weighting (Link Prediction)

For link prediction tasks, the system can overcome the collapse observed in uniform MSE distillation on small graphs like Cora, where the teacher’s embeddings encode label-discriminative directions but lose the geometric anisotropy needed for inner-product LP scoring. It recovers performance by using a curvature-adaptive reweighting strategy that distinguishes between a community direction and a bridge direction, leading to gains like +24.5 over KRD on Cora.

  1. textbf Scalable Inference with Minimal Latency: Vanilla MLP Deployment

The system enables deployment of the distilled knowledge into a standard MLP requiring no graph access at inference, achieving a 58× speedup over the 3-layer GraphSAGE teacher under standard layer-wise full inference on large graphs like ogbn-Arxiv, placing it on the speed–accuracy Pareto frontier alongside the GNN teachers.

Sources

Related papers