FIMP: Foundation Model-Informed Message Passing for Graph Neural Networks

TL;DR

This work introduces Foundation-Informed Message Passing (FIMP), a framework that repurposes the self-attention mechanisms of pretrained non-textual foundation models to drive cross-node message passing in graphs. By unifying node tokenization and performing cross-node attention over token sequences, FIMP enables messages to be generated from powerful pretrained models such as ViT, scGPT, and BrainLM, or learned from scratch in a base configuration. Across image networks, spatial transcriptomics, and brain fMRI data, FIMP demonstrates consistent performance gains over traditional GNNs and related baselines, including strong zero-shot embeddings with pretrained ViT on image graphs. The results underscore the potential of cross-domain foundation models to enhance graph-based tasks and open avenues for multimodal and zero-shot graph learning, with future work addressing broader domains and computational efficiency.

Abstract

Foundation models have achieved remarkable success across many domains, relying on pretraining over vast amounts of data. Graph-structured data often lacks the same scale as unstructured data, making the development of graph foundation models challenging. In this work, we propose Foundation-Informed Message Passing (FIMP), a Graph Neural Network (GNN) message-passing framework that leverages pretrained non-textual foundation models in graph-based tasks. We show that the self-attention layers of foundation models can effectively be repurposed on graphs to perform cross-node attention-based message-passing. Our model is evaluated on a real-world image network dataset and two biological applications (single-cell RNA sequencing data and fMRI brain activity recordings) in both finetuned and zero-shot settings. FIMP outperforms strong baselines, demonstrating that it can effectively leverage state-of-the-art foundation models in graph tasks.

Figures (5)

• Figure 1: The proposed framework for FIMP. Pre-existing foundation models, pretrained on vast amounts of unstructured data, are repurposed into message creation modules by adapting their self-attention layers for cross-node attention between node feature sequences.
• Figure 2: Graph structure present in real-world datasets. (A) In spatially resolved RNA transcriptomics, cells are connected to adjacent cells in the 2D tissue section. (B) In the Mapillary street-view image dataset antequera2020mapillary, images form a geographical proximity graph based on latitude and longitude coordinates. (C) For fMRI recordings, the brain is parcellated into 424 regions, which are connected using a K-nearest neighbors graph based on the 3D spatial coordinates of each brain region.
• Figure 3: Performance summary across different tasks for FIMP + foundation model versus the best traditional GNN baseline. FIMP improves over traditional GNNs across multiple datasets, highlighting the benefits of leveraging foundation models pretrained on unstructured data.
• Figure 4: Visualizations of FIMP feature-level attention between different functional groups in the brain. (A) Averaged attention heatmaps between functional regions of the brain for different age populations, with the difference in attention by age group visualized on the right subplot. (B) Similar heatmaps visualized for post-traumatic stress disorder (PTSD) scores, highlighting differences in attention in patients with low vs high PTSD score.
• Figure 5: (A) Averaged attention between 15 genes across edges connecting neighboring astrocyte cells in the mouse hippocampus dataset. (B) UMAP of learned gene embeddings from FIMP, colored by average expression value of each gene across astrocyte cells.