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

(Q2) Learned vs. Reused Weights in Algorithm 1

Another great question! Further clarification on algorithm 1 is given below:

1. Learned Weights: When trained completely from scratch (FIMP-base), all attention weights ($W_Q$, $W_K$, $W_V$) are initialized randomly and learned during training. FIMP-base therefore provides performance measures based solely on the improved tokenization of input graph data as well as cross attention message passing.
2. Reused Weights: When using a pretrained foundation model on the graph-based task (FIMP-scGPT, FIMP-ViT, FIMP-BrainLM), the pretrained self-attention layers of the foundation model attention are used in place of $W_Q$, $W_K$, and $W_V$, and are finetuned along with the rest of the model for the task. We will revise the description of Algorithm 1 to distinguish between learned and reused weights and add this clarification in Section 3.3.

(Q3) Knowledge Transfer in FIMP-Base

We would like to clarify that FIMP-base represents a version of the FIMP architecture that does not leverage pretrained foundation models. FIMP-base follows the same tokenization procedure as FIMP + foundation models for each task, however all encoders and weights of FIMP-base are trained from scratch. This is done mainly to demonstrate how the improved tokenization procedure and cross-attention message passing improve performance over GNN baselines even when trained completely from scratch. The FIMP + foundation model variants then go further and leverage pretrained foundation models in the message-passing framework, demonstrating that the pretrained representations further improves performance on each task.

We hope that this clarifies the distinction between FIMP-base and FIMP + foundation model variants. We will clarify this distinction in Sections 3.2–3.3 in updated versions of the manuscript.

(Q4) Adapting Pretrained Weights for Cross-Node Attention

Great question! We would like to take this opportunity to explain in more detail how the cross-node attention works. The foundation models contain pretrained transformer layers, which each contain self-attention layers. These self-attention layers are pretrained to do query-key-value (QKV) attention on input sequences. Specifically, in each self-attention layer, the input feature sequence is projected into query tokens, key tokens, and value tokens by the $W_Q$, $W_K$, and $W_V$ matrices, respectively. The attention matrix is computed through QK, and the attention matrices are then multiplied with the value vectors to form the output of attention.

Our main observation is that cross-attention can be formulated with the same layers if the attention is done between two sequences of features rather than one: Q representing the destination node feature sequence, and K and V representing the source node feature sequence. This represents a shift from the layer performing self-attention on unstructured data to performing cross-attention between neighboring nodes in a graph. We find that this empirically works well, and allows the cross-attention between nodes to leverage the pretrained weights which learned to compute self-attention on large amounts of unstructured data.

We hope these responses clarify the motivation, methodology, and contributions of FIMP. We are conducting additional experiments to address your points which we hope to add by the end of the rebuttal period. We respectfully request that you consider these clarifications when reevaluating your score.

Thank you for your detailed review and for identifying strengths and areas for improvement in our submission. We appreciate your balanced evaluation and the opportunity to address your concerns. Below, we provide clarifications and additional context in response to your feedback:

Clarity and Simplicity of the Approach

While we acknowledge that FIMP leverages a straightforward integration of foundation models with MPNNs, we believe its simplicity is a strength rather than a limitation. This simplicity allows for:

1. Effective Generalization: By repurposing pretrained foundation models through tokenization and cross-node attention, FIMP achieves significant performance gains across different graph domains without requiring complex architectural modifications. Our main goal is to align pretrained foundation models to work on graph-structured inputs, allowing pretrained representations from foundation models on unstructured data to benefit downstream graph-based applications.
2. Interpretability: The tokenization and cross-attention mechanisms provide interpretable feature-level interactions, as detailed in Appendix E.

Use of Outdated Baselines

We appreciate the recommendation to include comparisons with more recent state-of-the-art graph learning methods!. While we focused on widely-used baselines (e.g., GCN, GAT) and strong graph transformer baselines (GPS Graph Transformer) for comparability with prior works, we agree that incorporating modern baselines such as GRIT [1] would strengthen the experimental evaluation. We are working to add more updated GNN baselines, which we hope to complete during this discussion period.

Motivation and Real-World Relevance

1. Choice of Tasks and Motivation
   • The tasks we selected are designed to benchmark FIMP in domains where node features are rich, complex, and high-dimensional, such as cell expression profiles, images, and fMRI recordings. These domains particularly benefit from the integration of pretrained foundation models, which excel at learning general representations from large-scale unstructured data. By leveraging these pretrained representations, FIMP achieves significant performance improvements over traditional GNN baselines, as shown in our experiments. We argue that such tasks are not only appropriate benchmarks for evaluating FIMP’s strengths but also highlight its utility in handling complex, real-world data scenarios.
2. Real-World Applicability
   • For all spatial genomics datasets used in our experiments, the graph structure comes as a part of the publicly available data. As noted in the introduction, there are large amounts of unstructured cell expression data publicly available compared to spatial genomic data which preserves the graph structure inherent in tissues. We argue that utilizing foundation models pretrained on unstructured expression data (e.g., scGPT) for spatial genomics tasks provides significant real-world advantage over traditional methods which do not leverage pretraining on unstructured data in spatial genomics tasks. We believe this application demonstrates the meaningful, real-world impact of FIMP.
   • For the Mapillary image dataset and fMRI recording datasets, the latitude/longitude coordinates for images and the voxel positions for fMRI recordings are publicly available as part of the datasets as well. The graph constructions we performed (e.g., K-nearest neighbors for fMRI voxels or proximity graphs for Mapillary images) are simple preprocessing steps that anyone can reproduce using the provided dataset features. We argue that these settings represent real-world graph-structured data where foundation models pretrained on unstructured data can be leveraged.

To further ensure reproducibility, we plan to release all datasets and code after the review process, including the graph structures we used and the code for reproducing the preprocessing steps.

(Q1) Tokenization Alignment with Foundation Models

Great question! For spatial genomics datasets, we directly adopt the gene tokenization encoder from scGPT to encode cell expression profiles. This keeps the gene feature representation per cell consistent with what scGPT saw during pretraining, ensuring minimal distribution shift. For image data and fMRI recordings, we again closely follow the patch-based tokenization of ViT and BrainLM, respectively, dividing images or fMRI recordings into fixed-size patches. Appendix Section B further describes how tokenization schemes are adapted for each dataset. We will further expand this section to include more details explaining the original tokenization algorithm of each foundation model used in experiments.

We would also like to address the points raised regarding the motivation for FIMP and concerns about its simplicity. There are numerous real-world datasets with rich, information-dense features that can benefit from graph representation learning. For instance, Relational Deep Learning (RDL) [1] presents a deep learning approach for relational databases, highlighting opportunities for learning on structured data stored in such systems. As noted in our response to Reviewer Gmtx, simpler GNN models perform well on graph datasets with straightforward node features [2]. However, they struggle on datasets with more complex features, such as gene expression data in spatial genomics, spatiotemporal fMRI signals, and pixel-based image data. Across these domains, FIMP consistently outperformed all GNN baselines, demonstrating the meaningful impact of aligning pretrained foundation models to encode node entities and perform message passing effectively. Furthermore, FIMP uniquely enables zero-shot capabilities, as demonstrated on the Mapillary dataset. This allows FIMP to apply non-textual foundation models to graph datasets without any additional training, something that neither GNN baselines nor recent state-of-the-art methods can achieve.

Finally, we would like to reemphasize the motivation behind integrating foundation models into our cross-node attention message passing framework. FIMP is intended to bridge the gap between non-textual foundation models and graph-based tasks. Its simple yet effective design enables easy adoption across various domains, as demonstrated in our experiments. To our knowledge, FIMP is the first general framework to showcase the advantages of non-textual foundation models in graph settings. We believe this represents a significant contribution, paving the way for future applications of foundation models in graph learning, including zero-shot scenarios.

We will update our manuscript in future revisions to include these points and additional baselines, and with more time, we will incorporate GRIT into all other benchmarking tasks as well. We hope that these additional responses clarify why we believe FIMP represents a strong contribution to the field, and we kindly ask you to reconsider your score. Thank you for your continued engagement during the discussion phase.

References
[1] Fey, Matthias and Hu, Weihua and Huang, Kexin and Lenssen, Jan Eric and Ranjan, Rishabh and Robinson, Joshua and Ying, Rex and You, Jiaxuan and Leskovec, Jure. RelBench: A Benchmark for Deep Learning on Relational Databases
[2] Yuankai Luo, Lei Shi, Xiao-Ming Wu. Classic GNNs are Strong Baselines: Reassessing GNNs for Node Classification

In addition to our previous inclusion of the GRIT baseline, we have now also evaluated Graph Substructure Network (GSN) [1], another state-of-the-art GNN baseline, on the masked gene expression prediction task. Below, we present the results for GSN and other models on the Slideseq-V2 mouse hippocampus dataset:

While time constraints during the rebuttal period prevented us from running GSN across multiple tasks, these results demonstrate that FIMP outperforms SOTA GNNs, including GSN, on this gene expression prediction task. The performance gap highlights FIMP’s ability to handle complex graph datasets and its effectiveness in leveraging pretrained foundation models for graph-based tasks.

We kindly ask the reviewer to consider these additional experimental results in their evaluation of our work. Thank you for your continued engagement during the rebuttal process.

[1] Bouritsas, Giorgos, et al. "Improving graph neural network expressivity via subgraph isomorphism counting." IEEE Transactions on Pattern Analysis and Machine Intelligence 45.1 (2022): 657-668.

Thank you for your review and for highlighting aspects of our work that could be improved. We appreciate the opportunity to address your concerns and clarify the contributions and broader impact of FIMP. We respond to each comment in detail:

Novelty of the FIMP algorithm

We agree that attention-based Graph Neural Networks (GNNs) and transformers are well-studied in graph learning, however we argue that FIMP introduces significant innovations that distinguish it from prior methods:

1. Feature-Level Message Passing: Unlike traditional GNNs, including attention-based ones like GATs, which use scalar attention at the node level, FIMP tokenizes nodes into sequences of feature vectors and applies cross-node attention between feature tokens. This enables a finer-grained representation of node interactions, analogous to how transformers revolutionized NLP by moving from bag-of-words models to sequence-based architectures. This represented a major shift in data representation and generalizability in NLP, and we believe that similar gains in performance are attainable in graph settings based on our results.
2. Foundation Model Integration: FIMP uniquely integrates pretrained non-textual foundation models as message creation modules in graph settings. This allows us to leverage representations learned on vast amounts of unstructured data, enabling performance improvements and zero-shot capabilities that are not possible with traditional GNNs or standalone foundation models that do not operate on graphs.
3. Zero-Shot Capabilities: By repurposing pretrained foundation models, FIMP enables zero-shot applications on graph-structured data—a capability not achievable with traditional GNNs. This is particularly impactful in domains where labeled graph data is scarce and unstructured data is available for pretraining, such as single-cell genomics.

A more detailed comparison of FIMP with traditional GNNs is provided in Section 2.2 and Appendix D.

Evaluation on Traditional Graph Datasets

We deliberately focused our evaluation on domains where node features are rich and complex, as these settings highlight the advantages of FIMP’s feature-level tokenization and pretrained foundation model integration. Traditional graph benchmarks like OGB often have simpler node features (e.g., categorical or numeric vectors), where the benefits of foundation models are less pronounced, and simpler GNNs already achieve strong performance. Specifically, we focused on graph-structured data that was not bag-of-words or bag-of-features, since the preprocessing of these datasets loses information which foundation models can utilize, similar to how Large Language Models model natural language at a deeper level than term frequency or other bag-of-words based methods.

We will revise our limitations section to emphasize that FIMP is not best for every graph scenario. FIMP provides larger benefits when node entities have rich features that are not easily modeled by traditional GNNs.

Dependence on Powerful Foundation Models

We agree with the reviewer that FIMP’s reliance on powerful pretrained foundation models might limit its practicality in domains lacking such models. This is an important consideration, however we would like to highlight the following:

1. Performance Without Pretrained Models: FIMP-base, which does not use pretrained foundation model weights, still outperforms traditional GNNs on several tasks. This demonstrates that FIMP’s tokenization and cross-node feature-level attention are inherently advantageous, even without pretrained weights.
2. Foundation Models Across Domains: While foundation models are not yet available for every domain, their prevalence is rapidly increasing across fields such as biology, neuroscience, and vision. As highlighted by recent surveys [1], the development of domain-specific foundation models is accelerating, providing opportunities for methods like FIMP to bridge the gap between unstructured and graph-structured data as Reviewer ZuVo noted.

References [1] Zhou, Ce, et al. "A comprehensive survey on pretrained foundation models: A history from bert to chatgpt." arXiv preprint arXiv:2302.09419 (2023).

Computational Efficiency

We agree with the reviewer that the computational burden of leveraging foundation models in graphs is an important consideration. While FIMP’s use of foundation models introduces additional computational costs, we argue that this trade-off is justified in domains with complex node features. In the Mapillary dataset, for example, traditional GNNs fail to encode the details contained in the high-dimensional images. FIMP’s integration of pretrained ViTs enables significant performance gains, as shown in Table 3.

Regarding larger graph datasets, the Mapillary dataset is already a large dataset with over 100,000 images in the network. We show that by training on subgraphs sampled from the entire graph, FIMP-ViT can still effectively learn and generalize on the image classification task. In domains where computational resources are limited, we note that FIMP-base provides a more efficient alternative while still outperforming traditional GNNs.

Regarding the runtime comparison of FIMP + ViT and ViT itself on the image classification task, we provide the comparison in the table below:

We note here that the ViT model has less training time because it does not take any graph structure as input, and only processes a single image on a forward pass for classification. FIMP + ViT in comparison processes input subgraphs of images, which takes more time but achieves significantly higher F1-score on the image classification task since it considers a neighborhood of images by geographical proximity.

Broader Contributions

We also want to reiterate the broader impact of FIMP. By bridging the gap between graph-based tasks and pretrained non-textual foundation models, FIMP introduces a new paradigm for leveraging large-scale pretraining in graph learning. Beyond its technical contributions, FIMP opens up exciting possibilities for zero-shot applications and interpretable graph-based learning.

We hope these responses address your concerns and clarify the unique contributions of FIMP. We respectfully request that you reconsider your evaluation in light of this additional context, and we welcome any further questions or feedback you may have.

We appreciate your thoughtful and enthusiastic review of our submission, and we are glad that you recognize the novelty and importance of FIMP in leveraging pretrained non-textual foundation models for graph-based tasks. Below, we address some of the concerns and suggestions you raised:

Theoretical Analysis of FIMP’s Effectiveness

Thank you for your suggestion to include a deeper theoretical analysis of why the architecture is effective! We would like to provide additional insights into the theoretical underpinnings of FIMP:

1. Token Representation: Foundation models, particularly those based on transformers, have demonstrated a strong capacity to learn hierarchical and relational patterns in data. Tokenizing nodes into sequences of feature vectors allows FIMP to capture rich structural and relational features at a finer granularity than traditional Graph Neural Networks (GNNs). This approach is analogous to how transformers capture dependencies in sequences of words or image patches.

   • It is well established in fields such as Natural Language Processing (NLP) that tokenization and pretraining on large corpora of data provides large benefits in terms of performance and generalizability. De Santis et al. (2024) directly compared transformer-based models employing word embedding against traditional bag-of-word approaches for user categorization based on online posts, demonstrating that transformer-based models utilizing word embedding outperform traditional bag-of-words approaches. The key insight is that tokenizing inputs allows transformers to preserve the original word ordering and contextualize different parts of input sentences, better capturing the meaning of input text compared to traditional bag-of-words approaches which lose positional information. This is true in Computer Vision as well, with Vision Transformers [2] demonstrating that patch tokenization can allow transformers to be pretrained effectively at large scale to achieve state-of-the-art performance on Computer Vision tasks.
   • We view FIMP’s tokenization with the same benefits in mind, since FIMP allows for feature-level tokenization and contextualization of node entities in graphs. This is not possible with traditional GNNs, which operate over node feature vectors much like traditional bag-of-word models operated in NLP before transformers and word tokenization became widespread. This bag-of-words approach can be seen in common GNN benchmark datasets such as Cora [3], OGB [4], and Amazon [5].

2. Transferability and Generalization: The pretrained representations in foundation models, learned from vast datasets, generalize well across downstream tasks. Encoding nodes as token sequences enables FIMP to utilize these pretrained representations effectively, improving performance on downstream graph-based tasks.

3. Feature-Level Interactions: FIMP’s cross-node attention operates at the level of feature sequences rather than single node embeddings. This allows for detailed feature-level interactions between neighboring nodes, which traditional GNNs cannot achieve, and provides more expressive interaction between neighboring nodes during message passing.

We will incorporate these theoretical motivations into the revised manuscript to strengthen the understanding of FIMP's principles and its distinctions from prior works.

References
[1] De Santis, Enrico, et al. "From Bag-of-Words to Transformers: A Comparative Study for Text Classification in Healthcare Discussions in Social Media." IEEE Transactions on Emerging Topics in Computational Intelligence (2024).
[2] Dosovitskiy, Alexey. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).
[3] Yang, Zhilin, William Cohen, and Ruslan Salakhudinov. "Revisiting semi-supervised learning with graph embeddings." International conference on machine learning. PMLR, 2016.
[4] Hu, Weihua, et al. "Open graph benchmark: Datasets for machine learning on graphs." Advances in neural information processing systems 33 (2020): 22118-22133.
[5] Shchur, Oleksandr, et al. "Pitfalls of graph neural network evaluation." arXiv preprint arXiv:1811.05868 (2018).

Thank you for your reply! We are pleased to share the results of additional experiment runs for the mouse embryo classification task, which aimed to improve the statistical robustness of this particular task. We increased the number of experiment repetitions from 5 to 10 for most models by adding runs with new random seeds. These updated results are summarized below:

The rebuttal time constraints prevented us from completing all 5 new repetitions for FIMP-scGPT and scGPT, however the additional runs we were able to add suggest that the relative performance relationships between FIMP and the baseline models remain consistent. This indicates that our reported results are robust.

We appreciate your engagement throughout the review process! Please let us know if there are remaining points we have failed to address.

Broader Contributions

We also wish to reiterate the broader contributions of FIMP, which may have been unclear. Beyond its technical novelty, FIMP introduces a framework that bridges the gap between non-textual foundation models and graph-structured data. By enabling these pretrained models to operate on graphs, FIMP unlocks new capabilities such as zero-shot graph tasks and interpretable feature-level message-passing

We hope these clarifications address your concerns and provide a clearer picture of the novelty, motivation, and contributions of FIMP. We respectfully request that you reconsider your evaluation in light of this additional context, and we are happy to further discuss any remaining questions. Thank you again for your time and feedback.

Thank you for taking the time to review our submission, we appreciate your constructive feedback. Below, we respond to your comments and aim to clarify potential misunderstandings about our work:

Novelty of FIMP Compared to GATs

We respectfully disagree with the characterization that FIMP is an extension of GATs to non-textual node content, and would like to clarify the key distinctions between FIMP and GATs:

1. Node Representation: Traditional GNNs, including GATs, represent nodes as single embedding vectors, while FIMP tokenizes nodes into sequences of feature vectors (tokens). This tokenization aligns with foundation model architectures, enabling feature-level interactions that GATs cannot achieve.
2. Message-Passing Mechanism: GATs use scalar attention coefficients to aggregate single-node embeddings. In contrast, FIMP employs cross-node attention at the feature level between neighboring nodes, leveraging pretrained self-attention layers from foundation models as processing layers. This allows FIMP to capture fine-grained inter-node dependencies, enhancing expressivity and interpretability.
3. Foundation Model Integration: Unlike GATs, FIMP directly incorporates pretrained non-textual foundation models into its message-passing framework, enabling zero-shot applications on graphs and leveraging pretrained representations from domains such as images, genomics, and neuroscience. These capabilities are not achievable with GATs.

Sections 2.2 and Appendix D further emphasize the differences between FIMP and existing attention-based GNNs. We argue that traditional GNNs such as GAT cannot integrate foundation models in the same way that FIMP does, due to representational differences of node entities.

Motivation for Cross-Node Attention

We appreciate your suggestion to clarify the motivation behind cross-node attention and its impact on performance. Cross-node attention between feature sequences of neighboring nodes enables feature-level interactions, offering several advantages:

1. Enhanced Interactions: By tokenizing nodes into sequences, FIMP models $O(F^2 \cdot E)$ feature-level relationships, where $F$ is the number of features per node and $E$ is the number of edges. This considers significantly more interactions than traditional GNNs, which only model $O(E)$ node-to-node relationships. This can benefit graph domains with complex node features such as genomics where genes of neighboring cells may affect cell-cell interactions.
2. Improved Representational Power: Similar to how transformers outperform bag-of-words models in NLP by tokenizing input text, FIMP’s token-based representation allows for more contextualization of node features, leading to better performance across multiple tasks.
3. Interpretability: Attention coefficients between features of neighboring nodes can be visualized, providing insights into which node features contribute most to interactions during message-passing.

Comparisons With Baselines

FIMP is compared against both traditional GNN baselines and the foundation models themselves (e.g., scGPT, ViT, BrainLM) on each task, which represents a strong non-GNN baseline in each of their respective domains (single-cell genomics, images, and fMRI recordings, respectively). This demonstrates that FIMP effectively integrates pretrained models into a graph-based framework, outperforming both types of baselines.

Node Attribute Construction

For GNN baselines, we use raw node features from each dataset, consistent with standard preprocessing steps in GNN literature. For example, in the spatial transcriptomics datasets, node features are preprocessed gene expression vectors for each cell. For FIMP, each node entity is tokenized according to the methodology detailed in Section 3.1. A more detailed description of how node attributes are constructed is provided in Appendix C (Experimental Setup Extended).

Description of Figure 3

Thank you for pointing out the lack of description for Figure 3! We will add a detailed explanation of this figure in the Results section to ensure clarity. Specifically, Figure 3 summarizes FIMP’s performance gains over GNN baselines across tasks, demonstrating its effectiveness in diverse graph-based settings.

Thank you again for reviewing our paper and providing your initial comments. We recognize your concerns regarding the novelty and motivation of FIMP, and we have taken significant steps during the rebuttal process to clarify these aspects and address the issues raised.

In particular, we have provided detailed comparisons between FIMP and other baselines, incorporated new experiments (e.g., benchmarking GRIT on key tasks), and further expanded on the key distinctions between FIMP and traditional GNN architectures like GATs. These updates are described in our responses to other reviewers, and we encourage you to review these responses which highlight the additional experiments and insights we have added to strengthen the paper.

In particular, we would like to emphasize three key points about FIMP that we believe represent meaningful contributions of our work to the field:

1. FIMP is a natural choice for graph representation learning in settings with rich node features. FIMP is particularly well-suited for graph representation learning in settings with rich, complex node features. As noted in our response to Reviewer uKm7, traditional GNNs perform strongly on graphs with simple node features. However, in graphs with more complex and information-rich features, stronger models are needed to effectively encode and process this information. FIMP leverages pretrained foundation models, which incorporate representations learned from unstructured data, to address this need. In our experiments, FIMP consistently outperforms even state-of-the-art methods such as GRIT [1] across multiple tasks. We attribute this superior performance to FIMP’s ability to align pretrained representations with graph-structured inputs, demonstrating its effectiveness in handling complex graph datasets.
2. FIMP establishes a general framework for leveraging unstructured pretraining in graph domains. FIMP bridges the gap between foundation models and graph-based tasks by introducing a general-purpose framework for leveraging pretrained representations in graph learning. To our knowledge, no other framework currently provides the flexibility to utilize pretraining from arbitrary foundation models in graph-based tasks.
3. Zero-shot Capabilities: FIMP uniquely enables zero-shot inference on graph-structured data even when the pretrained models lack native graph support, as demonstrated in our experiments on the Mapillary image dataset. This capability allows FIMP to extend the utility of pretrained foundation models to graph domains without additional training—something not achievable with traditional GNNs or existing foundation model-based methods in non-textual graph settings.

We hope this response clarifies the motivation and significance of our work. We encourage you to review the additional context and experiments provided in our responses to other reviewers, and we kindly ask you to reconsider your evaluation. Thank you for your time and feedback.

[1] Graph Inductive Biases in Transformers without Message Passing. Ma et. al. 2023

Thank you again for your initial feedback on our paper. During the rebuttal process, we have worked to clarify the contributions of FIMP and have substantially improved our experimental results by including additional baselines such as GRIT and GSN in our response to reviewer uKm7. These updates address concerns about the novelty, motivation, and significance of our work.

Since this is the final day for reviewer comments, we would greatly appreciate your additional thoughts on our clarifications and revisions. Your engagement and feedback would be valuable in helping us further improve the manuscript.