AutoGFM: Automated Graph Foundation Model with Adaptive Architecture Customization

AutoGFM: Automated Graph Foundation Model with Adaptive Architecture Customization

Speakers: Haibo Chen, Xin Wang, Zeyang Zhang, Haoyang Li, Ling Feng, Wenwu Zhu

Conference: ICML 2025

YouTube: https://slideslive.com/39044093

Overview

This article delves into "AutoGFM: Automated Graph Foundation Model with Adaptive Architecture Customization," a significant contribution presented at ICML 2025 by Haibo Chen and his co-authors. The talk addresses a critical limitation in the rapidly evolving field of Graph Foundation Models (GFMs): the reliance on fixed, hand-designed Graph Neural Network (GNN) architectures. While GFMs are designed to generalize across diverse graph datasets and tasks by pre-training on a wide array of graph knowledge, their underlying GNN architectures often struggle to adapt optimally to the unique characteristics of different downstream applications. This architectural inconsistency, where a single optimal architecture does not exist across varying data, hampers the full potential of GFMs.

The researchers introduce AutoGFM as a novel solution that transitions GFM architecture design from a manual, static process to an automated, adaptive one. Their work is motivated by empirical evidence demonstrating that different graph datasets indeed necessitate distinct GNN architectures for optimal performance. AutoGFM tackles this challenge by pioneering Graph Neural Architecture Search (GNAS) specifically tailored for GFMs, overcoming inherent limitations of mainstream GNAS methods when applied to this domain. By enabling adaptive architecture customization, AutoGFM promises to unlock superior generalizability and performance for graph machine learning tasks across a myriad of domains.

The importance of AutoGFM lies in its ability to significantly enhance the utility and applicability of GFMs. By automating the architecture selection process, it not only reduces the considerable manual effort and expert knowledge required to design effective GNNs for diverse graph data but also pushes the boundaries of what GFMs can achieve. The proposed framework, with its theoretical underpinnings and empirical validation, marks a crucial step towards more intelligent, self-optimizing graph machine learning systems, offering substantial benefits for practitioners, researchers, and anyone working with complex graph-structured data.

Background

▶ Watch: Introduction to GFMs and architecture inconsistency problem (0:00)

The landscape of machine learning has seen a rapid evolution, particularly with the advent of foundation models that learn generalizable representations from vast amounts of data. In the realm of graph-structured data, Graph Neural Networks (GNNs) have emerged as powerful tools for tasks like node classification, link prediction, and graph classification. However, traditional GNNs are typically trained and optimized for a specific domain or task, limiting their generalizability when confronted with new, unseen graph data or different problem types. This domain-specific nature necessitates retraining or significant architectural modifications for each new application, a process that is both resource-intensive and often suboptimal.

To address this limitation, the concept of Graph Foundation Models (GFMs) has gained prominence. GFMs aim to parallel the success of large language models (LLMs) by pre-training on diverse graph datasets spanning multiple domains and tasks. The goal is to distill a broad "graph knowledge" base, enabling the GFM to learn more robust and generalizable representations that can be effectively adapted (fine-tuned) to a wide spectrum of downstream graph tasks. For instance, a GFM could be pre-trained on biological networks, social graphs, and citation networks, then fine-tuned for drug discovery, fraud detection, or academic recommendation systems. GNN-based GFMs are a particularly promising direction as they inherently unify various levels of graph data – nodes, edges, and entire graphs – allowing a single GNN backbone to handle multi-level tasks across diverse domains. An example mentioned is OFA, which leverages an LLM to unify input features from diverse datasets, converting various classification tasks into a unified subgraph classification task.

Despite the promise of GFMs, a critical challenge has emerged: architecture inconsistency. Existing GFMs typically rely on a single, hand-designed, and fixed GNN architecture. This fixed architecture is assumed to be universally optimal or sufficiently versatile for all downstream tasks and domains. However, as the authors highlight, this assumption is flawed. Data from different tasks and domains often exhibit unique structural properties, statistical distributions, and underlying relationships, which inherently require different optimal GNN architectures to capture effectively. For example, a GNN architecture effective for a dense social network might be suboptimal for a sparse molecular graph.

To validate this problem, the researchers conducted an experiment, testing various GNN architectures built upon a GFM across a range of datasets from diverse domains and tasks. Their findings, visualized as a heatmap, clearly demonstrated that the best-performing architecture varied significantly across different datasets. This empirical evidence underscores the "architecture inconsistency" problem: the optimal GNN architecture is not static but highly dependent on the specific graph data and task at hand. This realization forms the core motivation for exploring Graph Neural Architecture Search (GNAS) for GFMs, a direction that, prior to this work, had not been thoroughly investigated. The inherent variability in optimal architectures across datasets presents a fundamental hurdle for the widespread adoption and maximal performance of GFMs.

Key Findings

▶ Watch: Module 1: Disentangled Contrastive Graph Encoder details (3:59)

The AutoGFM project makes several pivotal contributions to the field of Graph Foundation Models and Graph Neural Architecture Search:

Firstly, the researchers rigorously identified and formalized the problem of architecture inconsistency in GFMs. Through empirical validation using a heatmap visualization, they demonstrated that different graph datasets and tasks indeed necessitate distinct optimal GNN architectures, challenging the prevailing assumption of a universally applicable, fixed architecture in current GFM designs. This finding highlights a fundamental limitation hindering the generalizability and performance of existing GFMs.

Secondly, the work provides a crucial theoretical analysis demonstrating the inadequacy of mainstream GNAS methods when applied directly to GFMs. Specifically, Assumption 3.1 posits that optimal architectures can differ across datasets. Building upon this, Proposition 3.2 mathematically proves that if two datasets require different optimal architectures, mainstream GNAS methods will encounter significant optimization conflicts during training. These conflicts lead to suboptimal architecture search results, rendering conventional GNAS approaches unsuitable for the GFM paradigm where a single model aims to serve diverse data. This theoretical insight provides a strong justification for developing a specialized GNAS approach for GFMs.

Thirdly, AutoGFM introduces a novel and comprehensive framework for Automated Graph Foundation Model with Adaptive Architecture Customization. The core innovation lies in learning a sophisticated mapping function that dynamically predicts an optimal GNN architecture based on the characteristics of the input graph data. This adaptive approach ensures that the GFM can customize its architecture on-the-fly, addressing the architecture inconsistency problem directly.

Fourthly, the proposed framework disentangles graph data into invariant patterns (stable features crucial for architecture prediction) and variant patterns (unstable or irrelevant information for architecture). This disentanglement is central to ensuring that the predicted architecture is robust and generalizable, depending only on the stable characteristics of the graph. The learning objective is meticulously formulated to achieve this disentanglement and robust mapping.

Finally, the experimental results unequivocally demonstrate the superior performance of AutoGFM. Across pre-training, fine-tuning, and few-shot learning experiments, AutoGFM consistently outperforms vanilla GNNs, self-supervised methods, GFMs with manually designed GNN architectures, and GFMs enhanced with conventional GNAS methods. Notably, its strong performance in few-shot learning highlights the rapid adaptability of its customized architectures. The visualization of customized architectures further corroborates the initial hypothesis, showing dataset-specific and even layer-specific preferences for different operations (e.g., GCN in the first layer and GAT in the second layer for the Arxiv dataset), confirming the necessity and effectiveness of fine-grained, automated architecture customization.

Technical Deep Dive

▶ Watch: Module 2: Invariant-Guided Architecture Customization (4:59)

The technical foundation of AutoGFM is built upon a sophisticated framework designed to address the architecture inconsistency problem in GFMs by learning an adaptive architecture customization mechanism. The core idea revolves around learning a mapping function that translates the characteristics of input graph data into an optimal GNN architecture.

At the heart of this mapping is the concept of decomposing graph data into two distinct components: an invariant pattern and a variant pattern. The invariant pattern is conceptualized as capturing stable, generalizable features of the graph that are crucial for predicting its optimal architecture. Conversely, the variant pattern contains unstable or dataset-specific information that should not influence architecture prediction, as doing so would lead to suboptimal and less generalizable designs. The theoretical underpinning for this decomposition is laid out in Assumption 3.3, which defines what constitutes an invariant pattern for architecture prediction.

To achieve this disentanglement and ensure a robust architecture prediction, a multi-objective learning formulation is proposed:

1. Minimize Mutual Information (Invariant Pattern, Variant Pattern): This objective ensures that the invariant and variant patterns are maximally disjoint, preventing the architecture prediction from being influenced by irrelevant, unstable features.
2. Maximize Mutual Information (Invariant Pattern, Architecture): This objective guarantees that the invariant patterns are sufficiently rich and informative for accurately predicting the optimal architecture.
3. Minimize Mutual Information (Architecture, Variant Pattern | Invariant Pattern): This crucial objective acts as a "shield," ensuring that the invariant pattern completely mediates the relationship between the architecture and the variant pattern. In essence, it guarantees that once the invariant pattern is known, the variant pattern provides no additional information for predicting the architecture, thus making the architecture prediction robust against variant influences.

The AutoGFM framework comprises two primary modules: the Disentangled Contrastive Graph Encoder and the Invariant-Guided Architecture Customization module.

1. Disentangled Contrastive Graph Encoder:

This module is responsible for extracting the disentangled invariant and variant patterns from diverse graph data. It leverages a GNN to learn a two-channel graph representation for any given graph. To guide these representations to reflect architecture requirements, an any-graph-level contrastive learning method is employed.

• Disentanglement Loss: This component explicitly encourages the learned invariant and variant patterns to be disentangled from each other, aligning with the first objective of minimizing their mutual information.
• Instance Discriminative Task: An any-graph-level instance discriminative task is introduced. Its purpose is to encourage the invariant pattern to capture the distinct architecture requirements of different data instances. This helps ensure that the invariant pattern is genuinely informative for architecture prediction, aligning with the second objective.
• The parameters of the graph encoder are optimized by minimizing this combined disentanglement loss.

1. Invariant-Guided Architecture Customization:

This module takes the disentangled patterns and uses them to customize the GNN architecture adaptively. It utilizes a weight-sharing supernetwork, a common technique in GNAS, where a vast space of possible GNN architectures is implicitly represented by a single, over-parameterized network containing all candidate operations (e.g., GCN, GAT, GraphSAGE) at each layer.

• Architecture Predictor: Given the graph representation Z (specifically, the invariant pattern component), an architecture predictor maps Z to a set of probabilities, alpha_li, which represent the likelihood of selecting the i-th operation for the l-th layer of the GNN.
• Learnable Prototypes: To facilitate this mapping, learnable prototypes are introduced as parameters of the predictor. If an input representation Z is "close" to a particular prototype in the embedding space, then the operation associated with that prototype is more likely to be selected. This creates a mechanism for grouping similar graph patterns with preferred operations.
• Auxiliary Predictor Guidance: A novel guidance mechanism is introduced using an auxiliary predictor.
• An invariant predictor directly predicts an "invariant architecture" based solely on the extracted invariant pattern.
• An auxiliary predictor predicts another architecture based on a fused representation that combines both invariant and variant patterns.
• The system then guides the customization process by minimizing the difference between the architectures predicted by these two predictors. This effectively pushes the architecture derived from the fused representation to align with the more robust "invariant architecture," indirectly enforcing the third mutual information objective (minimizing MI(Architecture, Variant | Invariant)).
• Crucially, the task loss of the GFM is calculated using only the architecture predicted by the invariant predictor, ensuring that the final GFM performance is driven by the robust, invariant architecture.
• Curriculum Architecture Customization Constraint: A key challenge in optimizing GFMs with multiple datasets is that different datasets may exert varying influences on the architecture search, potentially causing a few "easy-to-fit" datasets to dominate the selection of operations early in training. To counteract this, a curriculum learning approach is designed.
• It calculates the average choice weights for each operation in each layer across all datasets.
• It then computes the coefficient of variation of this average.
• This coefficient is used to define a constraint loss that encourages diversity in early-stage architecture customization. By penalizing low variation, this constraint loss actively mitigates the undue influence of any single data source, promoting a more balanced and comprehensive exploration of the architecture search space.

Through the synergistic operation of these modules, AutoGFM can dynamically adapt its GNN architecture to the specific demands of diverse graph data, moving beyond the limitations of fixed designs and conventional GNAS methods.

Experimental Setup & Results

▶ Watch: Experimental results: pre-training and fine-tuning performance (6:59)

The experimental validation of AutoGFM was comprehensive, designed to showcase its effectiveness across various learning paradigms crucial for Graph Foundation Models. The researchers conducted experiments in pre-training, fine-tuning, and few-shot learning settings, comparing AutoGFM against a range of strong baselines.

The baselines included:

• Vanilla GNNs: Standard GNN architectures without any foundation model pre-training or architecture search.
• Self-supervised methods: GNNs trained with self-supervised objectives, often used for pre-training, but without adaptive architecture.
• GFM with manually designed GNNs: Foundation models pre-trained on diverse graph data but utilizing a fixed, hand-designed GNN architecture (representing the state-of-the-art GFM approach that AutoGFM aims to improve).
• GFM with GNAS methods: Foundation models integrated with existing, mainstream Graph Neural Architecture Search techniques, demonstrating the limitations of applying off-the-shelf GNAS to GFMs as discussed in the theoretical analysis.

The experiments aimed to evaluate AutoGFM's ability to customize architectures for different domains and data, its generalizability, and its adaptability to novel tasks with limited data.

The headline results were consistently positive:

• Superior Performance: AutoGFM consistently outperformed all baselines across all datasets and experimental settings. This robust superiority demonstrates its effectiveness in customizing architectures and achieving better performance for diverse graph data and tasks.
• Fast Adaptability in Few-Shot Learning: In few-shot learning scenarios, where models must quickly adapt to new tasks with very limited labeled examples, AutoGFM achieved the best performance across almost all datasets. This particular result highlights the rapid adaptability of its architecture customization mechanism, suggesting that the invariant patterns learned are highly effective at guiding efficient architecture selection for novel tasks.

To provide clear evidence of the adaptive architecture customization, the researchers visualized the chosen architectures for different datasets. They presented a heatmap illustrating the choice weights of each operation (e.g., GCN, GAT, GraphSAGE) at each layer of the GNN for various datasets. This visualization yielded several critical observations:

• Dataset-Specific Architectures: The heatmap visually confirmed the initial hypothesis of architecture inconsistency: different datasets indeed prefer distinct GNN architectures. This directly validates the core problem AutoGFM sets out to solve.
• Layer-Specific Operations: Beyond overall architecture, the visualization revealed that many datasets exhibit fine-grained preferences, favoring varying operations across different layers within the same architecture. For example, the Arxiv dataset was observed to prefer a GCN operation in its first layer but a GAT operation in its second layer. Such nuanced, layer-specific preferences are exceptionally challenging, if not impossible, to meet through traditional hand-design methods. This observation powerfully underscores the advantage and necessity of AutoGFM's automated and customized architecture search capabilities.

The experimental setup and compelling results provide strong empirical evidence that AutoGFM successfully addresses the architecture inconsistency problem in Graph Foundation Models, leading to significant performance gains and enhanced adaptability across a wide range of graph machine learning tasks.

Practical Implications

▶ Watch: Few-shot learning and visualized customized architectures (7:59)

The AutoGFM framework carries significant practical implications for various stakeholders involved in the development and deployment of graph machine learning solutions.

For practitioners and model builders, AutoGFM offers a powerful mechanism to overcome the arduous and often suboptimal process of manually designing GNN architectures for diverse graph datasets. By automating architecture customization, it frees up valuable time and resources, allowing researchers and engineers to focus on higher-level problem formulation and domain-specific insights rather than GNN architecture engineering. This leads to the deployment of more generalizable and performant GFMs, which can adapt on-the-fly to the unique characteristics of new graph data, thereby improving the overall robustness and reliability of graph-based AI systems. The ability to achieve superior performance in few-shot learning scenarios also means faster prototyping and deployment of GFM-powered solutions for emerging tasks with limited data.

For infrastructure teams and deployers, AutoGFM enables a new paradigm of adaptive GFM deployment. Instead of deploying a monolithic, fixed-architecture GFM that might be inefficient for certain data types, AutoGFM allows the GFM to dynamically adjust its internal architecture to best suit the incoming graph data. This could lead to more efficient resource utilization, as the model can choose simpler or more complex operations as needed, potentially optimizing computational costs and latency. While the architecture search itself adds complexity during training, the resulting adaptive GFM offers a more flexible and potentially more performant inference environment, especially in multi-tenant or highly dynamic graph data environments.

Tradeoffs and Limitations are also important to consider. The primary tradeoff is the increased complexity and computational cost during the pre-training phase due to the integrated architecture search. Training a weight-sharing supernetwork and an architecture predictor, along with the disentanglement and curriculum constraints, is inherently more resource-intensive than training a fixed GNN. However, this upfront investment is justified by the subsequent gains in generalizability, performance, and reduced manual effort during deployment across diverse tasks. Another limitation, as acknowledged by the speakers, is that the current work primarily focuses on GNN-based GFMs. The exploration of adaptive architectures for other types of GFMs (e.g., those based on Transformers or other graph neural architectures) is left for future work. Furthermore, while the theoretical analysis elegantly demonstrates the limitations of mainstream GNAS methods, the practical deployment of AutoGFM would require careful consideration of the specific GNN operation candidates within the supernetwork to ensure broad applicability to target domains. Despite these considerations, AutoGFM represents a significant step towards more autonomous and intelligent graph machine learning systems, moving closer to the vision of truly adaptive foundation models.

Key Takeaways

• Architecture Inconsistency is a Core Problem: Existing Graph Foundation Models (GFMs) suffer from "architecture inconsistency," meaning that a single, fixed GNN architecture is suboptimal, as different graph datasets and tasks require distinct optimal architectures.
• Mainstream GNAS is Insufficient: Traditional Graph Neural Architecture Search (GNAS) methods are ill-suited for GFMs due to optimization conflicts that arise when a single GNAS process attempts to find architectures for diverse datasets with varying optimal requirements.
• AutoGFM: Adaptive Architecture Customization: The proposed AutoGFM framework addresses this by learning a dynamic mapping function that predicts an optimal GNN architecture based on the input graph data's characteristics.
• Disentangled Patterns for Robustness: AutoGFM disentangles graph data into invariant patterns (stable for architecture prediction) and variant patterns (irrelevant), ensuring that the predicted architecture is robust and generalizable.
• Modular Design with Novel Components: The framework employs a Disentangled Contrastive Graph Encoder to learn patterns and an Invariant-Guided Architecture Customization module with a weight-sharing supernetwork, an architecture predictor, and an auxiliary predictor for guidance.
• Curriculum Constraint for Diverse Search: A novel curriculum architecture customization constraint is introduced to mitigate the undue influence of specific datasets during early training, promoting a more diverse and comprehensive exploration of the architecture search space.
• Superior Performance and Adaptability: AutoGFM significantly outperforms vanilla GNNs, self-supervised methods, manually designed GFM architectures, and GFMs with conventional GNAS across pre-training, fine-tuning, and few-shot learning tasks, demonstrating superior generalizability and fast adaptability.

About the Speaker(s)

The work on AutoGFM was presented by Haibo Chen, Xin Wang, Zeyang Zhang, Haoyang Li, Ling Feng, and Wenwu Zhu. As researchers in the field of AI/ML, their collective expertise has contributed to identifying and addressing a critical challenge in Graph Foundation Models, pushing the boundaries of automated and adaptive machine learning architectures for graph-structured data.