Thank you for your thoughtful review! We briefly address some of the questions and comments you raised:

Why can a larger model do better on the connected nodes task?

Great question!! Actually, very recent theoretical results from [1] provide insight to this behavior, providing bounds for the kinds of tasks where Transformers can outperform (non-attention) GNNs (and vice-versa). Their results show that Transformers have an advantage on tasks which require access to the full graph in order to answer them (like connected nodes, and shortest path). It's an open research question on whether a new kind of GNN encoder can do better than Transformer-based ones in this area.

We will happily add a discussion on the connected nodes task, this related theory, and the implications for encoder design to this work!

Any runtime analysis or comparison to the baselines for Parameter Efficiency?

Another great question -- here is a runtime analysis for the PeFT experiments in Table 3!!

This Table shows that GraphToken is not only the most performant model for the task, but also the fastest! (This is by design -- GNN parameters are ‘computationally cheaper’ than Transformer parameters). Please see our discussion with who had some similar questions about efficiency.

We have a longer discussion on the efficiency of the method in a dedicated comment to Reviewer #AQrB, that highlights.

1. Runtime efficiency (shown above)
2. Asymptotic complexity benefits (no big surprises, but a MPNN is more efficient than a Transformer)
3. More about the parameter usage benefits (where we see adding $O(100k)$ parameters allows a small model (PaLM-2-S) to beat a substantially larger model (PaLM-2-L).

GraphToken Losses and Limitations

Great idea! We’re working on a more detailed analysis of where GraphToken under performs, but the analysis is taking time. We hope to complete it this week!

Requires access to an open-source LLM to compute gradients

We completely agree that the current inability to compute gradients for proprietary LLMs may prevent the application of GraphToken to closed-source models. However we would like to emphasize a few points:

1. GraphToken makes small models ‘punch above their weight’. This can really enable new capabilities for smaller open source models. For example, we have a nice experiment comparing GraphToken to Tx-LLM, a specialized text model for therapeutics for Reviewer #FJUo. We see that a smaller (and more general model) can outperform a more specialized model, even with fine tuning on the same task! We believe this to be a significant contribution.

2. As strong open source LLMs continue to be released, GraphToken can be applied with open models which have 100 billion+ parameters.

3. Finally, we believe that the strong performance on algorithm reasoning tasks will motivate developers of proprietary LLMs to provide support for graphs and other modalities.

References

[1] “Understanding Transformer Reasoning Capabilities via Graph Algorithms” - https://arxiv.org/pdf/2405.18512

Thanks for your question on the limitations of GraphToken -- we have analyzed its performance across graphs structures and will add the following discussion on it to the manuscript!

Limitations of GraphToken

Our results have shown that GraphToken is both a flexible and generalizable encoder of graph structured data for LLMs. Here we discuss some limitations of the method as inspiration for future work.

Encoder Generalizability

The main limitation of GraphToken is that its encoder might learn spurious correlations due to idiosyncrasies in the distribution of input graphs it was trained on. As such, it's important that the encoder works robustly regardless if its evaluated on a different distribution of input graphs (w.r.t. their density, number of nodes/edges, etc). We note that this is a general weakness of GNNs and not specific to GraphToken itself. As such, there is a rich literature on creating robust GNNs [1,2] that has made significant progress in creating more generalizable GNN architectures. We expect that these results will be directly able to be “plugged in” to GraphToken encoders and will greatly aid in their generalization.

Does Graph Structure Matter?

It is interesting to study how graph structure might affect a GraphToken encoder’s performance on downstream tasks. Are more complicated graphs harder? Do other structural patterns influence its results?

Experiment design:

We use the Gemma2-2B multi-task encoder we trained for reviewer #FJUo which is trained on 9 out of 10 GraphQA tasks. Then we calculated a number of graph properties that have been shown useful for analyzing GNN performance [1], and examined the correlation between these graph characteristics and whether the model was able to correctly answer its task (out of 10 GraphQA tasks) on unseen graphs. The results are as follows:

Interestingly, we find that most graph properties are uncorrelated with the downstream task’s performance. We do see a weak negative correlation between the number of nodes in the graph and correctly answering tasks. However this is expected -- as the number of nodes grows, the graph has the potential for more complexity.

We believe that this result supports the strong generalization capabilities of GraphToken!

References

[1] GraphWorld: Fake Graphs Bring Real Insights for GNNs - https://arxiv.org/abs/2203.00112

[2] A Survey of Deep Graph Learning under Distribution Shifts: from Graph Out-of-Distribution Generalization to Adaptation - https://arxiv.org/abs/2410.19265v1

[3] Robustness of Deep Learning Models on Graphs: A Survey - https://galina0217.github.io/works/robust_survey.pdf

Thank you for your time and careful consideration of our work! We want to address your concerns raised in the Weaknesses section:

GraphToken Features

As GraphToken uses a general GNN encoder, it accommodates any featurization of the nodes, edges, or graph including textual information. An example of this capability is demonstrated in section 4.2 where we consume molecular data which includes very rich information -- including node, edge, and graph level features. We do not compare to GraphLLM because, to quote the GraphLLM paper “In this paper, we operate under the assumption that nodes (or entities) can be characterized by textual features”. (We have actually tried to compare to GraphLLM and some similar works regardless, but have not succeeded in reproducing them -- in our experience so far, papers with significant architecture modifications seem more brittle.)

Task-Dependent Encoder Output

1. We do not consider this data leakage. Instead we consider this part of the prompt. In the parameter efficient fine-tuning setting, task performance is usually related to fine-tuning a model for a specific task. (Practically speaking, a retrieval system designed to use this graph information could easily select a graph-encoder which would be specifically trained for its task.)

2. While we do indeed experiment on smaller graphs, this is primarily just to benchmark against text encoding in LLMs which have very small context windows. One benefit of this approach is that it is significantly more scalable than alternative text encoding formulations which may require hundreds of tokens to encode the information about these graphs that we are compressing into a single token.

3. While we consider the parameter-efficient fine-tuning setting (improving a model on one graph task), we do agree with you that this is an interesting direction. One benefit of our design is that we can easily accommodate advances in graph encoders (either specific ones like better molecular encoders, or general graph reasoning architectures). The encoders can also be pretrained on their tasks (independent of the LLM). We believe our results show the promise of this design (utilizing graph soft tokens) and that this path will allow us to add fundamental new capabilities to large language models.

Model memory requirements

Great question! The biggest requirement of our method is computing gradients from a LLM. While this does not necessarily require loading the entire model into memory, doing so definitely decreases the training latency.

• Our experiments with GraphToken + Gemma 2B were run on a single machine with one TPU v5e (16GB VRAM).

• The experiments with GraphToken + PaLM 2-S were run using TPU v3 accelerators. Unfortunately, we are not able to provide more detailed accelerator information for PaLM 2-S (as that is strongly related to model size).

• The GNNs (with only 100k parameters) don’t require much additional VRAM headroom over the LLM requirements.

Finally, we note that the memory requirements can be pushed substantially lower using standard techniques for minimizing LLM memory use such as quantization, or by loading only part of the model into memory at a time.

We apologize for the late reply about this concern!

Additionally, the fact that, in some cases, a single token per node/edge is utilized raises questions about the scalability to larger graphs (graphs with significantly more nodes than ~12 which is the average size in GraphQA). Has the model been evaluated with larger graphs? Insights into its performance under these conditions would be valuable.

This is a great question.

• In the soft prompt tuning literature, the number of tokens to generate is a typically a hyper-parameter which can be varied to increase model capacity.
• We do not advise using just one GraphToken for larger and more complex graphs!
• However, to be frank, we have been impressed with how well just one token can do -- the multitask experiment for Reviewer FJUo uses one GraphToken and can perform the 9 tasks it was trained on, in addition to an unseen 10th task!

To answer the related question of 'How much work could one GraphToken do?' -- we can produce a larger version of GraphQA (same tasks, but with larger generated graphs) and perform some analysis on how answer quality degrades as graph size grows.

Finally, we call your attention to the correlation experiment for Reviewer Zq2s, where we see that larger graph size (#nodes) is only weakly correlated with incorrect answers. However, we agree with your intuition that one should only be able to fit 'so much' graph inside a single token.

Thanks and please let us know if you have additional concerns or would like this additional experimentation. (it would not be too hard to do)

We thank the reviewer for a very insightful response!

Scope of Work

We believe that while the overall direction the reviewer highlighted is in line with our work, the scope of our work is in parameter-efficient encoding. There is more that LLMs bring besides generalizability and interpretability – it’s their in-context learning capabilities [1]. When considering extremely significant improvements of GraphToken over LLM-only baselines (cf. Palm-2 L results in Table 1), it is clear that the GNN encoder is a better overall design of the system. In fact, we would say the core argument of our paper is that “GNNs are necessary for LLM advancement".

We believe the design ideas sketched in this review are more general, but our work solves a particular encoding problem for graphs that has not been properly addressed before.

Unified Model

To the point of training a unified model, it is certainly a useful experiment; however, in the parameter-efficient fine-tuning literature it is common to fine-tune LLMs to perform a single task that is most useful to the concrete goal the user has. We take our graph task definitions from Fatemi et al. (2023) [2], which is a standard for this fast-evolving field. While simple, they span different levels of complexity for both transformers and GNNs – for instance, GNNs alone can not solve the connectivity task without resorting to global positional encoding.

Comparison with Specialized LLMs

To put our results on property prediction in context, we compare with Tx-LLM [3], a special purpose LLM designed for medical applications (with a special emphasis on biochemical interactions). Both models are fine tuned for the task. Please note that Tx-LLM is “text-only” but has significant pre and post training for the domain. GraphToken received no additional task-specific pretraining or fine tuning.

We see that GraphToken is again able to outperform a larger model (even though it's significantly more specialized).

Better GNN models available?

We completely agree with you that better GNNs for molecular representation should improve these results further -- on both specialized tasks (and grounding/generation for explainability). We view this as a positive statement about our architecture’s benefits -- it allows LLMs to leverage the strengths of domain-specific models.

Related Work

Thank you for the references, we are happy to extend the related work with a discussion of these papers. We believe that the rapid growth of contemporary work in the area justifies the timeliness and relevance of our work!

Why use a GNN to create LLM input?

At the core, converting from an unordered object (a graph) into a sequence introduces a bias that is hard to develop heuristics for [4]. The primary benefit from our proposed architecture is that the graph encoder can preserve properties (like equivariance, etc) for as long as they are useful, and then deliver an optimized representation for the LLM. This allows graph data to be used just like “another modality” (vision, audio, etc).

References

[1] “Language Models are Few-Shot Learners” - https://arxiv.org/pdf/2005.14165

[2] "Talk like a graph: Encoding graphs for large language models" - https://arxiv.org/pdf/2310.04560

[3] “Tx-LLM: A Large Language Model for Therapeutics” - https://arxiv.org/pdf/2406.06316

[4] "Towards an Understanding of Graph Sequence Models" - https://openreview.net/pdf?id=iaHghgG8NR

Thanks for the reply! In the meantime, we have been very actively working to answer the questions you raised in the original review!

MultiTask: Yes!

To illustrate that the same GraphToken can work for multiple tasks, we have investigated multi-task models at your request. We trained a three task model (edge existence, cycle check, reachability) with Gemma 2B which yielded the following results:

We see a small drop in performance from using a single encoder, but are quite pleased with the results. (The multi-task model has not had its hyper-parameters optimized, so there is likely more performance available).

Generality

We emphasize that this model didn’t have access to the task description (exactly the same GraphToken representation was being provided for each of the three tasks). We believe that this shows the generality of the approach.

More tasks?

We have begun training even larger multitask models (i.e. with more tasks), but hit some infrastructure issues along the way and the experiments are not done yet. We may have additional results to report prior to the rebuttal prior ending.

While these initial results seem quite promising, we suspect that training a true “graph foundation model” is a research work unto itself. For example, perhaps some tasks can not be easily combined. We note that multi-task performance on current GFM models seems to be lacking [1].

GNN-only Results

We agree that these are valuable both to illustrate what headroom there might be for the combined GNN + LLM architecture, and further the field.

GraphQA GNN-Only Results

We have run the following experiments which we will add to the paper:

For each task, we ran the following additional baselines:

• a “GNN-only” architecture using the GraphToken encoder
• Gemini 1.5 Pro (a recent capable “LLM-only” model)

Our intent is to use these experimental results to add a discussion of the tradeoffs between GNN, GNN + LLM, and LLM architectures to the paper, informed by recent theoretical results in the area.

GNN-Only Molecular Property Baselines

We also investigated running GNN-only molecule encoders, but have so far found these encoders difficult to adapt to our programming environment. We hope that the GraphQA results illustrate the strength of GNNs on the task to your satisfaction.
If not, we are happy to report other’s numerical results on the molecular property prediction datasets if you wish.

Minor points

Finally we wanted to respond to two minor points that we believe came up in your review.

GNN + LLM benefit

We are actually aware of some follow up work (which cites us) that demonstrates a unique benefit for our combined GNN + LLM architecture in a few-shot learning setting. In this setting, the GNN+LLM is better able to adapt to new unseen tasks.

Novelty

Similarly, we are fairly certain that we were very early (perhaps the first) advocates of this architecture. However we acknowledge that we are not the first to get published on it.

Unfortunately we can not provide this citation (or evidence of our novelty) for reasons of anonymity. We will leave a note for the AC about this, but suspect there is no good solution.

Thanks!

Thanks for your time spent in the review process and engaging with the authors! Please let us know if there are remaining points we have failed to address.

References

[1] Text-space Graph Foundation Models: Comprehensive Benchmarks and New Insights - https://arxiv.org/pdf/2406.10727

MultiTask Generalization: Yes!

In order to illustrate MultiTask generalization, we have performed the following experiment at your request:

• We have trained a MultiTask GraphToken model on 9 of the 10 GraphQA tasks. The LLM backbone is Gemma2-2B (to minimize training time).
• As before the same encoder is used for all tasks, and the encoder has no knowledge about the task while encoding.
• We evaluate this model on a withheld task -- cycle check. It has never seen cycle check before this eval.

Results:

For context, we provide the single task performance ("GraphToken Single Task") on Gemma2-2B as well as its text-only performance.

We see that while there is a drop on performance (as might be expected) compared to optimizing a task directly:

1. GraphToken can generalize to a unseen task
2. The generalization outperforms large models (e.g. PaLM-2-L) on the task.
   1. PaLM-2-L's highest scoring result was 83.3.
   2. The generalization is much better than using a text representation of the graph with the backbone LLM.

We thank you for the suggestion and will be including this experiment in an updated version of the paper!

We believe that this is a strong rebuttal to the concerns from your original review. Please let us know if you have any additional questions.

the leave-out dataset (cycle check) seems arbitrary, have you tried all leave-out settings (leave shortest-path out, etc)?

Great question - this was the first task train/test split we tried actually. We choose cycle check as a test because it seemed to perform well in the three task experiment above.

We'd like to extend the analysis further for the camera ready, and we agree that a full leave-one-out cross-validation would be very interesting. Unfortunately we ran into some resource issues that prohibited us from pursuing this question more this week.

Thanks for the suggestion and engaging in the review process!

C4.) Efficiency

Thanks for the question! We see several angles to further the discussion.

1. Training Speed

One way to illustrate the strength of our approach is by measuring the training speed of GraphToken compared to other PeFT baselines on the molecular property prediction task. Please see the following table below:

From this table, we see that our method is not only the best performing method, but also the fastest! The strongest performing baseline (P Tuning) takes 100x more time to train than our method.

2. Runtime Complexity

The design of the GNN encoder directly influences the runtime complexity for both training and inference in our method. This is a significant strength of our proposed architecture, as GNN layers can have much better computational complexity than transformer layers.

For instance, the complexity of a MPNN layer used in the above experiment is $O(|E|d)$ [3], while a transformer layer would require at least $O(|V|^2d)$ [4,5] to represent a graph ($d$ is the hidden dimension size). In sparse graphs, $|E| << |V|^2$. Most applications of graphs in the real world (molecules, social networks, etc) are sparse.

3. Parameter Usage

While it is not the best measure of speed here (due to the architectural differences we just mentioned), we reference the parameter usage Table from Section 4.3.3, and the GraphQA experiments in Table 1. In this we see that adding a very small number of “graph aware” parameters $O(100k)$ allowed a small model (PaLM-2 S) to outperform a very large model (PaLM-2 L).

Quite frankly, GraphToken allows a model that can run on a desktop PC to outperform a supercomputer. (We sincerely apologize that we can not be more specific about the exact parameters used by each model.)

References

[3] https://arxiv.org/pdf/2405.17311v1

[4] https://arxiv.org/pdf/1706.03762

[5] https://arxiv.org/pdf/2207.02505

We thank the reviewer for their careful review of our paper. Responding to the concerns raised:

C1.) Is Text Enough?

While we appreciate that a text-only interface is indeed a benefit of other methods, we believe that alternative (and superior) methods for introducing graph structured data must exist.

In this work, we are explicitly trying to demonstrate the benefits to graph reasoning that can come from treating graphs as a separate modality, in a similar way to how LLMs currently support vision and audio data. While this does come at the cost of no longer being able to work with closed-source LLMs, we believe that our results demonstrate significant benefits for (re: Table 1, >100% performance on 5 tasks, and substantial improvement on the rest).

We anticipate these strong results may motivate closed source LLMs to add support for graphs based on these results.

C2.) Contributions

The technical contribution of the GraphToken encoder is the extension of this 'fusion' approach to graphs including the comparison of it with simpler text approaches, extension to new applications (chemical property prediction), and studying the generalization of the representations produced. GraphToken was the first work to propose this fusion (although there have been many concurrent (and subsequent works) since it was first released).

The impact of the different degrees of freedom in encoding a graph in text [1] or even images [2] are well studied in prior work. [1] found that while certain text encoding methods perform better than others, there was broad weakness in LLMs ability to reason over graphs in this form. [2] found that the image modality was insufficient for graphs (and generally did not add much over a text representation)

C3-1.) Comparisons

Comparison to text encoding methods is indeed an important baseline. In Table 1 in our paper we provide some comparisons between GraphToken and a few methods of text encoding with different prompting methods. We provide a brief description of the comparison in section 4.1. However due to length constraints we had to move a more in depth discussion of this part to the appendix. Please see Appendix A.3. We are happy to extend this discussion further!

C3-2.) GraphWiz

Thank you for calling our attention to GraphWiz! The GraphQA benchmark that we use covers some but not all of these tasks (in particular Hamiltonian and the subgraph matching problems look interesting to compare to).

C4.) Efficiency

We will respond to this in its own comment for space

C5.) Minor comments

Thanks, for the suggestion! We believe this paragraph to be related to our main goal since we seek to find the best way to augment a prompt with graph data (e.g. a knowledge graph might have fresher world knowledge than a LLM). However, to your point, it does take a while to get there 🙂. We will update this text to be relevant to the current draft. (We do ask the reviewer to note that our task differs slightly from the ‘graph foundation model’ task as many define it. Instead, we seek the best embedding of graph structure for its utilization in foundation models.)

References:

[1] Talk like a graph: Encoding graphs for large language models - https://arxiv.org/pdf/2310.04560

[2] Which Modality should I use -- Text, Motif, or Image? : Understanding Graphs with Large Language Models - https://arxiv.org/abs/2311.09862

C3-2.) Comparison to GraphWiz on harder graph problems

Thank you again for calling attention to the additional tasks presented in GraphWiz!

To analyze GraphToken’s performance on some of these more challenging tasks, we focused on the Hamilton path graph problem proposed in GraphWiz, in which models predict whether a Hamiltonian path exists in the graph which visits each node exactly once. Since GraphToken uses parameter efficient supervision, we compare the performance of GraphToken (with a GraphSAGE encoder) against the performance reported by GraphWiz for SFT:

Again, we see very strong performance from GraphToken, matching the performance of the best results with a model that only requires tuning 338,688 parameters! This is fantastically more efficient (requiring only 0.0048384% of the parameters tuned by the Llama-2 7B SFT baseline).

Once again, we see that GraphToken’s ability to make smaller LLMs outperform much larger ones through encoding structure greatly aids in graph reasoning tasks.

As we finish the first discussion phase, we'd like to call your attention to our results highlighting our efficiency, and a comparison on a NP-Complete task we performed at your request.

Please let us know if you have any additional questions, and thanks again for your thoughtful review!