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

📄 arXiv: 2604.22031v1 📥 PDF

作者: João Mattos, Arlei Silva

分类: cs.LG, cs.AI

发布日期: 2026-04-23

备注: 20 pages, 7 figures

💡 一句话要点

Mochi：通过元学习对齐预训练与推理，实现高效图基础模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 图神经网络 图基础模型 元学习 预训练 小样本学习

📋 核心要点

  1. 现有图基础模型预训练与下游任务对齐不足，依赖后处理统一，影响性能。
  2. Mochi采用元学习，在小样本episode上预训练，直接对齐训练目标与推理过程。
  3. Mochi及其变体在多个图任务上超越现有模型，训练效率显著提升。

📝 摘要（中文）

本文提出了一种名为Mochi的图基础模型，它通过采用基于元学习的训练框架来解决任务统一和训练效率问题。现有模型通常采用基于重构的目标（如链接预测）进行预训练，并假设生成的表示可以通过单独的统一步骤（如类原型）与下游任务对齐。本文通过合成和真实世界的实验证明，这种方法虽然简单直观，但存在直接影响下游任务性能的局限性。为了解决这些局限性，Mochi在模拟下游评估协议的小样本episode上进行预训练，从而将训练目标与推理对齐，而不是依赖于事后的统一步骤。实验结果表明，Mochi及其更强大的变体Mochi++在涵盖节点分类、链接预测和图分类的25个真实世界图数据集上，与现有图基础模型相比，实现了具有竞争力或更优越的性能，同时所需的训练时间比最强的基线少8到27倍。

🔬 方法详解

问题定义：现有图基础模型通常采用重构任务（如链接预测）进行预训练，然后通过诸如类原型等方法将预训练得到的表示与下游任务对齐。这种两阶段方法存在问题，即预训练目标与下游任务目标不一致，导致次优的表示学习效果。现有方法训练效率较低，难以适应大规模图数据。

核心思路：Mochi的核心思路是通过元学习，模拟下游任务的 few-shot 学习场景，在预训练阶段就将模型调整到适应下游任务的状态。通过在多个模拟任务（episode）上进行训练，模型能够学习到一种通用的图表示，这种表示能够快速适应新的、未见过的图任务。这种方法避免了事后的对齐步骤，直接优化了下游任务的性能。

技术框架：Mochi的整体框架包含以下几个主要步骤：1）图数据采样：从原始图数据中采样得到多个小的图结构，每个图结构代表一个 episode。2）任务构建：对于每个 episode，构建一个 few-shot 学习任务，例如节点分类或链接预测。3）元学习训练：使用元学习算法（例如 Model-Agnostic Meta-Learning, MAML）训练图神经网络，使其能够快速适应新的 episode。4）推理：在下游任务上，使用预训练好的图神经网络进行推理，并根据具体任务进行微调。

关键创新：Mochi的关键创新在于将元学习引入到图基础模型的预训练中，通过模拟下游任务的 few-shot 学习场景，实现了预训练目标与推理过程的直接对齐。这种方法避免了传统方法的两阶段训练模式，提高了模型的训练效率和泛化能力。Mochi++通过引入更强大的图神经网络结构和更有效的元学习算法，进一步提升了模型的性能。

关键设计：Mochi的关键设计包括：1）Episode构建方式：如何从原始图数据中采样得到具有代表性的 episode，以及如何构建 few-shot 学习任务。2）元学习算法选择：选择合适的元学习算法，例如 MAML 或 Reptile，以及如何调整算法的参数以适应图数据。3）图神经网络结构：选择合适的图神经网络结构，例如 GCN 或 GAT，以及如何设计网络的层数和隐藏层维度。4）损失函数设计：设计合适的损失函数，以优化模型的训练效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Mochi在25个真实世界图数据集上进行了广泛的实验，涵盖节点分类、链接预测和图分类任务。实验结果表明，Mochi及其变体Mochi++在多个任务上取得了优于现有图基础模型的性能，并且训练时间显著减少，比最强的基线模型减少了8到27倍。这表明Mochi在提高模型性能的同时，也显著提升了训练效率。

🎯 应用场景

Mochi具有广泛的应用前景，可应用于社交网络分析、生物信息学、化学信息学等领域。例如，在社交网络中，Mochi可以用于预测用户之间的关系或用户的兴趣爱好。在生物信息学中，Mochi可以用于预测蛋白质之间的相互作用或药物的靶点。该研究有助于推动图神经网络在实际应用中的发展，并为解决复杂图结构数据的分析问题提供新的思路。

📄 摘要（原文）

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.