GOFA: A Generative One-For-All Model for Joint Graph Language Modeling

In plain text applications, aligning the pre-training objective with the input and output format of downstream tasks is straightforward, as they are purely text-based. However, graph data from different domains vary significantly by input feature dimensions and output target. Hence, the first challenge is to define a unified input and output format for graph tasks, such that the model can do large-scale self-supervised pre-training on arbitrary graphs and downstream tasks.

For the unified input, following the previous work OFA Liu et al. (2023a), we extend the definition of Text-Attribute Graph (TAG) beyond graphs that already have text features such as citation and product networks. Specifically, we observe that any node and edge features can be represented by texts. For example, atom and bond descriptions replace the node and edge features in a molecular graph. Furthermore, we can represent non-textual features like pure numbers as texts too, just like they are represented as text strings in LLM. Formally, a TAG is a graph $G=\{V,E,X_{V},X_{E}\}$ where $V$ and $E$ are the sets of nodes and edges. Each node $v\in V$ (edge $e\in E$) corresponds to a text description $x(v)\in X_{V}$ ($x(e)\in X_{E}$). Such a format encodes almost all existing graph data and serves well as the universal input representation. For the unified output, we believe that natural language is still the most tangible format for users. Such input and output formats cover almost all commonly seen graphs-related tasks, including classification, regression, and free-form question answering.

Next, we define the generative modeling framework for graphs. Generative modeling for language is to predict the next token given its preceding tokens. However, in graph-based tasks, input sentences are distributed to different nodes and edges in a TAG. To generalize next token prediction to graphs, we propose to specify Nodes Of Generation (NOG) on graphs as starting points for output generation. Formally, we can define graph generative modeling as the likelihood of the text $y$ associated with the NOG $v$:

where $y_{l}$ is the $l$-th token of $y$, and $y_{<l}$ is its preceding tokens. The NOG $v$ is a learning target with initial corresponding text $x(v)$, and $x(v)$ can be empty. $G$ contains structural and textual information of neighbor nodes to help the model generate $y$. For example, the sentence completion task is shown on the left of Figure 2, where the text on node $v$ is incomplete, and the goal is to complete the sentence using the existing text and the neighbor information. The question answering task is shown on the right of Figure 2, where a user’s question is injected as a node in the graph. The goal is to generate the correct answer. Note that our model is not limited to node-level tasks in conventional graph learning. We discuss more details about data construction in Section 3.3. Users can also specify multiple generation starting points for the same graph. In addition, when the graph has only one node, the problem degenerates to conventional language modeling.

To solve the generative graph modeling problem proposed in Equation 1, we design the GOFA architecture, with an illustration shown in Figure 3. GOFA consists of a graph language encoder and an LLM decoder. In the graph language encoder, we interleave token-level GNN layers with frozen LLM compressor layers. The LLM compressor independently processes the graph’s node/edge text features and outputs their compressed fixed-size multi-token embeddings which are fed into a token-level GNN layer. The GNN layer updates the multi-token embeddings by incorporating the input graph structure and feeds the output into the next LLM layer. We alternate GNN and LLM layers to get the final multi-token node representations containing joint structural and semantic information. We then use an LLM decoder to generate texts from the NOG representation. The LLM compressor and decoders are all pre-trained decoder-only transformers. We describe each component in detail as follows.

LLM compressor: Because GNNs require node and edge representations to have the same input dimension, many previous works propose to pool all tokens’ output embeddings from the LLM as the node and edge vector representations and feed them to a GNN Liu et al. (2023a); Huang et al. (2023a); He and Hooi (2024). While this approach shows effectiveness in tasks of fixed form, such as classification and regression, it is insufficient in more complex tasks, such as generation tasks, as 1) the pooling process inevitably loses semantic information, 2) standard LLMs are not trained in a way such that the pooled output embedding is a summarization of the input sentence, and 3) the pooled representation space is no longer compatible with the space of the downstream LLM decoder. Hence, we adopt a sentence compressor Ge et al. (2023) that preserves as much information as possible from the original sentence in fixed-size multi-token embeddings. Specifically, the sentence compressor layer has the same architecture as a decoder-only LLM, but the initial sequences of token representations $\{q(x_{i})\}_{i=1}^{l}$ are appended by a sequence of $K$ memory tokens $\{q(m_{j})\}_{j=1}^{K}$, and the $t$-th layer of the LLM becomes:

We use $Q_{x}^{t}$ and $Q_{m,x}^{t}$ to represent the $t$-th LLM layer output corresponding to actual text tokens in $x$ and the $K$ memory tokens appended at the end of text tokens. We use $H_{x}^{t}$ to represent the $t$-th GNN layer output, which will be explained later. In Equation 2, the text tokens ($Q_{x}^{t}$) and memory tokens ($H_{x}^{t}$, processed by the previous GNN layer) are concatenated as a single sequence of embeddings, which are fed to the current LLM layer. Because the last K tokens attend to all previous tokens, they can compress all information in the sentence into the output embeddings of the K tokens. This compressor architecture is inspired by ICAE Ge et al. (2023). While the compression capability is not apparent, it is viable through auto-encoder-style fine-tuning, as discussed in Appendix A.1.

Token-level GNN: Conventional GNNs only pass messages between node vector representations, which is incompatible with the vector sequence output $Q_{m,x}$ from our LLM. Hence, we propose a simple extension of conventional GNNs to token level. For node $v\in V$, the $t$-th GNN layer is

In the GNN layer, tokens at different indices do not communicate. If we directly stack these GNN layers, they degenerate into multiple isolated GNNs for each token. Nevertheless, we interleave the GNN layers into the LLM layers such that the subsequent self-attention layers on $H^{t}$ permit information flows between tokens. This approach significantly reduces memory usage because we do not perform cross-token attention in GNN. While edge memory tokens are passed into GNN to assist message passing, their representations are not updated in the GNN layer but directly passed to the next LLM layer, and $H^{t}_{x(e)}=Q^{t}_{m,x(e)}$. The specific choice of the GNN can vary. In GOFA , we use a modified Transformer Convolutional GNN Shi et al. (2021) to be consistent with the transformer architecture of LLM (see Appendix A.3 for details).

We insert one GNN layer between two transformer layers, while the first and the last layer are always transformer layers. In GOFA , we only insert GNN to the last few transformer layers, but this can be flexible depending on the computational resources. Following previous practice, we incorporate feed-forward (FF) layers into the GNN to increase expressivity and residual connections to stabilize training. Inspired by the gating mechanism in earlier works Hochreiter and Schmidhuber (1997); Alayrac et al. (2022), we apply a $tanh$ gate, initialized at 0, to the GNN and FF layer outputs so that the initial model ignores the information from GNN layers and is equivalent to the pre-trained LLM.

LLM decoder: After applying the model to the textual graph, the memory tokens $Q_{m,x}$ of every node contain information about the text on the node, the surrounding node text, and the graph structure due to the message-passing process in the GNN layer. Then, for the NOG $v$ and its corresponding target text $y$, we insert $Q_{m,x}$ at the front of the token embeddings of the target text to generate and use teacher-forcing to maximize the standard log-likelihood for the next-token-prediction objective.

In this way, we have modeled the problem in Equation 1. The compressor, decoder, and GNN parameters can be jointly or separately optimized, potentially with parameter-efficient fine-tuning methods like LoRA Hu et al. (2022). In this paper, we use ICAE Ge et al. (2023) as our backbone LLM, but the GOFA architecture is not tied to any specific LLM. More details are discussed in Appendix A.2.

The formulation in Equation 1 provides a natural way for users to query the graph by selecting a NOG. However, tasks in the graph domain have different levels, such as node-, link-, and graph-levels, and different tasks have distinct objectives and target nodes. To cover these tasks under the same formulation, inspired by OFA Liu et al. (2023a), we convert all tasks into subgraph tasks and connect all target nodes in a graph to a virtual prompt node as the NOG for generation (i.e., $P$ nodes in Figure 1). Specifically, for node- and link-level tasks, we sample a $k$-hop rooted subgraph surrounding the target node or node pairs for node- or link-level tasks and connect the prompt NOG node to target nodes. For graph-level tasks, the prompt NOG node will connect to all nodes in the graph. This design has several advantages: (1) We can specify the prompt node as the NOG input to the LLM decoder for generative modeling. Therefore, the distribution of all tasks can be unified into a single space; (2) The text feature for the prompt node describes the task details. Connecting the prompt node to target nodes helps the prompt node query the most important knowledge from the input graph through attention and message passing. This encourages the model to dynamically learn from the message-passing process subject to prompt instead of only statically learning from the node text $x(v)$.

In this section, we describe the self-supervised pre-training of GOFA . The goal of the pre-training is to let the GOFA model obtain the ability to query graph structure and context but retain the ability to reason about plain text. Specifically, we perform the pre-training task using MAG240M Hu et al. (2021a) and Ultrachat200k Ding et al. (2023) datasets. Details about the datasets can be found in Appendix C. For each node $v$ in the MAG240M, we sample an ego-subgraph $G$ around $v$ to form a training example. We randomly generate a unique node ID (such as [Node A]) for each node in the training sample and append it to the original node text. This ID will serve as a basis for querying nodes in the graph. We design three pre-training tasks: sentence completion, structural tasks on MAG240M, and question-answering tasks on Ultrachat200k. Figure 1 shows an example of each task. We describe the details of each task below.

Sentence completion task. The objective of the sentence completion task is to train the model to reason about the rest of the text in a node given both the existing text and the information in the rest of the graph. First, we assign the root node as the target for each sampled training graph. Next, we randomly select $n$ nodes in the graph and the target node. All selected nodes’ texts are split into halves. The first half forms node text $x(v)$, and the second half becomes the target $y$ to generate. Additionally, we connect a prompt node $v_{p}$ for each selected node, as shown in Figure 1. Note that for all prompt nodes, we use directed edges to connect them with nodes in the input graph, such that only the information in the input graph will be passed to the prompt node for decoding, and information in different prompt nodes does not influence each other. The text in each prompt node asks the model to complete the sentence of the selected node it connects. This sentence-completion pre-training task adapts LLMs’ standard ability to the graph context.

Structural tasks. The objective of the structural tasks is to pre-train GOFA to understand basic graph structural properties. In light of this, we design the shortest paths and common neighbors reasoning tasks between nodes. Specifically, For each training subgraph sample, we randomly sample $n$ nodes as the selected nodes. For each selected node, we ask the model to compute the shortest path distance between it and the root node and output all possible shortest paths between them. Meanwhile, we also ask the model to output the number of common neighbors the selected node and the target node have and the node IDs of their common neighbors. The prompt node $v_{p}$ will connect to both the target and the selected nodes since our structural tasks need the model to reason about two nodes simultaneously. The text in the prompt node will be the corresponding question. Through these two tasks, the model is expected to gain the ability to identify basic graph structures, which are critical to downstream tasks.

Question answering task. Unlike language corpus, which naturally contains many question-and-answer (QA) pairs, graph data usually only contain objective descriptions of entities. Nevertheless, for the model to be fluid in tasks, we need the model to understand user prompts and be sensitive to different tasks. Hence, we synthesize a QA-chain dataset from Ultrachat200k, as shown in Figure 1. A language QA sequence is converted to a chain where nodes with question texts alternate with nodes with answer texts, which are connected by directed edges to represent the conversation order. The last question becomes the text on prompt node $v_{p}$, which is connected to every node in the chain, and the last answer is the target text $y$ (see Figure 1 QA-Chain Task for an example). This QA task provides QA pseudo-graphs missing from the common graph corpus, and we found it critical for allowing the model to be responsive to arbitrary tasks expressed in free-form questions.

To answer Q1, we pre-train two GOFA models using ICAE models on Llama2-7B Touvron et al. (2023) and Mistral-7B Jiang et al. (2023), optimizing the objective in Equation 4. The pre-training datasets include MAG240M Hu et al. (2021a), UltraChat200k Ding et al. (2023), and our proposed structural learning tasks as introduced in Section 3.4. The training details can be found in Appendix F. Figure 2 shows a training loss decrease for the number

of seen tokens. We report the perplexity in Table 2 and compare GOFA with base LLM finetuned on the MAG240M dataset. Note that during pre-training, we only update the weight of the GNN layers, and GOFA ’s lower perplexity shows that the structural and semantic information in the node’s neighbor can effectively help complete the sentence with more relevance than the original LLM. We observe that fine-tuned LLM does not improve much from the base LLM as it is already well-trained on large-scale academic datasets and cannot directly use the neighbor information to refine the generation. Besides sentence completion, another important GOFA pre-training objective is the structure learning ability; we report shortest path distance and common neighbor count prediction results in Table 2, compared with LLM models whose inputs are textualized graphs, with descriptions of edge connections. We see a significant performance improvement of GOFA over base LLM, showing that a difficult graph task for LLM can be well solved by the GNN layers with better structure modeling ability.

To answer Q2, we performed zero-shot experiments on various graph tasks. Despite using QA-chain data in the pre-training stage, the graph data does not include knowledge about task formats like classification and does not output exact matches to the answers. Hence, we first perform instruction tuning on the pre-trained GOFA in Section 5.1 using two question-template datasets, arxiv, and Pubmed-link. We only use a part of the data, and the goal of instruction fine-tuning is not to let the model learn particular knowledge from these datasets but to make the model understand the task format described in Appendix F.3. We intentionally chose these datasets because the distribution is similar to the pre-training MAG240M dataset, which is important for fair zero-shot experiments.

We perform 10-way classification fine-tuning on only 10k/40k arxiv data for one epoch and show zero-shot results in Table 4. While arxiv only contains node classification tasks, we observe that GOFA achieves very non-trivial performance on all node-level (Cora, Pubmed, WikiCS, Products), link-level (FB15K237), and graph-level (ExplaGraphs, SceneGraphs) tasks. GOFA also generalizes to different ways and even question-answering (SceneGraphs) tasks, showing its desirable fluidity. GOFA outperforms LLM and graph foundation model baselines on most datasets, and exceeds second best results by a large margin ($>5\%$) on WikiCS, Products, and ExplaGraphs, showing GOFA ’s ability to combine the advantage of both LLM and graph models. We also notice that GOFA does not perform as outstanding on the FB15K237 knowledge graph dataset, potentially because it is a link prediction task, which requires task-specific knowledge different from our arxiv instruction tuning data. Using a knowledge graph for tuning can potentially improve the performance and we leave this to future work. We also observe that GOFA is only achieving comparable performance to LLM on the SceneGraph dataset. We notice that SceneGraph contains object coordinate description which makes it easy for LLM to answer tasks about relative location. Whereas, our compressed representation might still discard relevant number information compared to directly prompting this information to LLM. Hence, the advantages of GOFA are not as apparent on tasks where text-only information is sufficient. We discuss potential solutions in Appendix F and leave this to future work.

For Q3, we observe one unique feature of GOFA is that its GNN layers are embedded into the LLM. A natural question is how important the interleaved GNN layers affect the LLM layer outputs. Hence, we compute a layer-wise Representation Change Ratio value, $\delta=||H^{t}-Q^{t}||/||Q^{t}||$ ($H^{t}$ undergoes MLP, residual, and gating modules), quantifying the change that GNN brings to the LLM outputs; a larger value means a higher dependence on graph information. We take the mean ratio of 100 data randomly sampled from each zero-shot dataset. The results in Figure 5 show that almost all GNN layers have a significant impact on the LLM output in both pre-trained and node(arxiv)/link(Pubmed-link) instruction-tuned versions. The pre-trained and arxiv versions have similar importance for each layer as they are all academic datasets and node-level tasks usually obtain enough information for neighbors within 2-hops. The link-finetuned version has larger third and fourth-layer importance, showing that the link task usually requires the model to consider a larger neighborhood to understand the graph structure.

Another important design in GOFA that enables in-context zero-shot learning is the injected prompt node with the user instruction. This allows the GNN to pass messages differently when users provide different prompts. We present a zero-shot ablation study on the connection between target and prompt nodes in Table 5. Single edge means that only the prompt node receives a message from the target nodes, and double edge means that target nodes also receive a message from the prompt node. The former results in static message passing independent of the prompt, and the latter in dynamic message passing. We simultaneously train GOFA -llama on Pubmed-link and arxiv datasets and evaluate their zero-shot performance. Without dynamic message passing, the fine-tuned model can only perform well on one dataset and fail to use its inductive knowledge when the task (Cora-link v.s. Cora-node) or the domain (Cora v.s. Products) shifts. This shows the importance of the prompt node design, which enables dynamic message-passing that can better use the pre-trained knowledge of the GNN and LLM compressor for zero-shot tasks. We answer Q4 with supervised training in Appendix B.