Mochi: Aligning Pre-training and Inference for Efficient Graph Foundation Models via Meta-Learning

TL;DR

Mochi introduces a meta-learning framework for graph foundation models that aligns training with inference, improving efficiency and performance across diverse tasks and datasets.

Contribution

Mochi's novel meta-learning approach pre-trains on few-shot episodes, directly aligning training with downstream inference, reducing training time and enhancing task performance.

Findings

Mochi achieves competitive or superior results on 25 graph datasets.

Mochi requires 8 to 27 times less training time than baselines.

Mochi's approach effectively unifies various graph tasks with improved efficiency.

Abstract

We propose Mochi, a Graph Foundation Model that addresses task unification and training efficiency by adopting a meta-learning based training framework. Prior models pre-train with reconstruction-based objectives such as link prediction, and assume that the resulting representations can be aligned with downstream tasks through a separate unification step such as class prototypes. We demonstrate through synthetic and real-world experiments that this procedure, while simple and intuitive, has limitations that directly affect downstream task performance. To address these limitations, Mochi pre-trains on few-shot episodes that mirror the downstream evaluation protocol, aligning the training objective with inference rather than relying on a post-hoc unification step. We show that Mochi, along with its more powerful variant Mochi++, achieves competitive or superior performance compared to existing Graph Foundation Models across 25 real-world graph datasets spanning node classification, link prediction, and graph classification, while requiring 8$\sim$27 times less training time than the strongest baseline.