STAGE: Simplified Text-Attributed Graph Embeddings Using Pre-trained LLMs (2024)

Abstract: We present Simplified Text-Attributed Graph Embeddings (STAGE), a straightforward yet effective method for enhancing node features in Graph Neural Network (GNN) models that encode Text-Attributed Graphs (TAGs). Our approach leverages Large-Language Models (LLMs) to generate embeddings for textual attributes. STAGE achieves competitive results on various node classification benchmarks while also maintaining a simplicity in implementation relative to current state-of-the-art (SoTA) techniques. We show that utilizing pretrained LLMs as embedding generators provides robust features for ensemble GNN training, enabling pipelines that are simpler than current SoTA approaches which require multiple expensive training and prompting stages. We also implement diffusion-pattern GNNs in an effort to make this pipeline scalable to graphs beyond academic benchmarks.

heading: 1 Introduction
text: A Knowledge Graph (KG) typically includes entities (represented as nodes), relationships between entities (represented as edges), and attributes of both entities and relationships (Ehrlinger and Wöß, 2016). These attributes, referred to as metadata, are often governed by a domain-specific ontology, which provides a formal framework for defining the types of entities and relationships as well as their properties. KGs can be used to represent structured information about the world in diverse settings, including medical domain models (Koné et al., 2023), words and lexical semantics (Miller, 1995), and commercial products (Chiang et al., 2019).
Text-Attributed Graphs (TAGs) can be viewed as a subset of KGs, where some node and edge metadata is represented by unstructured or semistructured natural language text (Yang et al., 2023). Examples of unstructured data values in TAGs
*Authors contributed equally.
could include the research article text representing the nodes of a citation graph, or the content of social media posts that are the nodes of an interaction graph extracted from a social media platform. Many real-world datasets are naturally represented as TAGs, and studying how to best represent and learn using these datasets has received attention from the fields of graph learning, natural language processing (NLP), and information retrieval.
Graph Learning and LLMs With the emergence of LLMs as powerful general purpose reasoning agents, there has been increasing interest in integrating KGs with LLMs (Pan et al., 2024). Current SoTA approaches combining graph learning with (L)LMs follow either an iterative or a cascading method. Iterative methods involve jointly training an LM and a GNN for the given task. While this approach can produce a task-specific feature space, it may be complex and resource-intensive, particularly for large graphs. In contrast, cascading methods first apply an LM to extract node features which are then used by a downstream GNN model. Cascading models demonstrate excellent performance on TAG tasks (He et al., 2024; Duan et al., 2023a), although they often require multiple stages of training targeted at each pipeline component. More recent cascading techniques implement an additional step, known as text-level enhancement (Chen et al., 2024), whereby textual features are augmented using an LLM.
Simplifying Node Representation Generation To the best of our knowledge, all existing cascading approaches require multiple rounds of data generation or finetuning to achieve satisfactory results on TAG tasks (He et al., 2024; Duan et al., 2023a; Chen et al., 2024). This bottleneck increases the difficulty of applying such methods to real-world graphs. Our proposed method, STAGE, aims to simplify existing approaches by foregoing LM finetuning, and only making use of a single pre-
92
trained LLM as the node embedding model, without data augmentation via prompting. We study possible configurations of this simplified pipeline and demonstrate that this method achieves competitive performance while significantly reducing the complexity of training and data preparation.
Scalable GNN Architectures The exponentially growing receptive field required during training of most message-passing GNNs is another bottleneck in both cascading and iterative approaches, becoming computationally intractable for large graphs (Duan et al., 2023b; Liu et al., 2024). Because we wish to study approaches that can be applied in real-world settings, we also explore the implementation of diffusion-pattern GNNs, such as SimpleGCN (Wu et al., 2019) and SIGN (Frasca et al., 2020), which may enable STAGE to be applied to much larger graphs beyond the relatively small academic benchmarks. Our code is available at https://github.com/aaronzo/STAGE.
Concretely, this work studies several ways to make learning on TAGs more efficient and scalable:
• Single Training Stage: We perform ensemble GNN training with a fixed LLM as the node feature generator, which significantly reduces training time by eliminating the need for multiple large model training runs.
• No LLM Prompting: We do not prompt an LLM for text-level augmentations such as pre-
dictions or explanations. Instead, we use only the text attributes provided in the dataset.
• Direct Use of LLM as Text Embedding Model: Using an off-the-shelf LLM as the embedding model makes this method adaptable to new models and datasets. We study several alternative base models for embedding generation.
• Diffusion-pattern GNN implementation: We contribute an investigation into diffusionpattern GNNs which enable this method to scale to larger graphs.
The rest of the paper is organized as follows: section 2 gives an overview of related work, section 3 discusses our approach in detail, section 4 studies the performance of STAGE in various settings, and section 5 is a discussion of the experimental results.

heading: 2 Background
text: Text-Attributed Graphs Yan et al. (2023) suggest that integrating topological data with textual information can significantly improve the learning outcomes on various graph-related tasks. Chien et al. (2022) incorporate graph structural information into the pre-training stage of pre-trained language models (PLMs), achieving improved performance albeit with additional training overhead, while Liu et al. (2023) further adopt sentence embedding models to unify the text-attribute and
graph structure feature space, proposing a unified model for diverse tasks across multiple datasets.
LLMs as Text Encoders General purpose text embedding models, used in both finetuned and zero-shot paradigms, are a standard component of modern NLP pipelines (Mikolov et al., 2013; Pennington et al., 2014; Reimers and Gurevych, 2019). As LLMs have emerged as powerful zeroshot agents, many studies have considered generating text embeddings as an auxiliary output (Muennighoff, 2022; Mialon et al., 2023). BehnamGhader et al. (2024) introduce LLM2Vec, an unsupervised method to convert LLMs into powerful text encoders by using bidirectional attention, masked next token prediction and contrastive learning, achieving state-of-the-art performance on various text embedding benchmarks.
Language Models and GNNs Graph Neural Networks have been successfully applied to node classification and link prediction tasks, demonstrating improved performance when combined with textual features from nodes (Kipf and Welling, 2017; Li et al., 2022b). Several studies show that finetuning pre-trained Language Models (PLMs), such as BERT (Devlin et al., 2019) and DeBERTa (He et al., 2021), enhances GNN performance by leveraging textual node features (Chen et al., 2024; Duan et al., 2023a; He et al., 2024).
Recent research has explored the integration of
LLMs with GNNs, particularly for TAGs. LLMs contribute deep semantic understanding and commonsense knowledge, potentially boosting GNNs’ effectiveness on downstream tasks. However, combining LLMs with GNNs poses computational challenges. Techniques like GLEM (Zhao et al., 2023) use the Expectation Maximization framework to alternate updates between LM and GNN modules.
Other approaches include the TAPE method, which uses GPT (OpenAI, 2023; OpenAI et al., 2024) models for data augmentation, enhancing GNN performance through enriched textual embeddings (He et al., 2024). SimTeG demonstrates that parameter-efficient finetuning (PEFT) PLMs can yield competitive results (Duan et al., 2023a). (Ye et al., 2024) suggest that finetuned LLMs can match or exceed state-of-the-art GNN performance on various benchmarks.
Building on these insights, the STAGE method focuses on efficient and scalable learning for TAGs by utilizing zero-shot capabilities of LLMs to generate representations without extensive taskspecific tuning or auxiliary data generation.

heading: 3 Approach
text: Our cascading approach consists of two steps:
• A zero-shot LLM-based embedding generator is used to encode the title and abstract (or
equivalent textual attribute) of each node. We denote the generated node embeddings as X .
• An ensemble of GNN architectures are trained on X , and their predictions are mean-pooled to obtain the final node predictions.
Ensembling the predictions from multiple GNN architectures was motivated by our observation of strong performance by different models across different datasets.

heading: 3.1 Text Embedding Retrieval
text: For the text embedding model, we select a general-purpose embedding LLM that ranks highly on the Massive Text Embedding Benchmark (MTEB) Leaderboard1. Specifically, we evaluate gte-Qwen1.5-7B-instruct, LLM2Vec-Meta-Llama-3-8B-Instruct, and SFR-Embedding-Mistral. MTEB ranks embedding models based on their performance across a wide variety of information retrieval, classification and clustering tasks. This model is used out-of-the-box without any finetuning. An appealing aspect of LLM-based embeddings is the possibility to add instructions alongside input text to bias the embeddings for a given task. We empirically evaluate the effect of instruction biased embeddings is in Table 2 of section 4.
Node representations X are generated using only the title and abstract, or equivalent textual node attributes, omitting the LLM predictions and explanations provided by (He et al., 2024). X will then be used as enriched node feature vectors for training a downstream GNN ensemble.

heading: 3.2 GNN Training
text: Using the previously generated embeddings X as node features, we train an ensemble of GNN models on the node classification task:
Losscls = Lθ (ϕ(GNN(X ,A)),Y) , (1)
where ϕ(·) is the classifier, A is the adjacency matrix of the graph and Y is the label. For the GNN architectures we choose GCN (Kipf and Welling, 2017), SAGE (Hamilton et al., 2018) and RevGAT (Li et al., 2022a). We also evaluate a multi-layer perceptron (MLP) (Haykin, 1994) among our GNN models. To combine the predictions from each of the K models in the ensemble, we compute the mean prediction as follows:
1https://huggingface.co/spaces/mteb/ leaderboard
p̄ = 1
K
K∑
k=1
pk, (2)
Cross-entropy loss is used to compute the loss value.
Diffusion-based GNNs For a graph G with node features X , a diffusion operator is a matrix AOP with the same dimensions as the adjacency matrix of G. Diffused features H are then calculated via H = AOPX .
We explored Simple-GCN (Wu et al., 2019) and SIGN (Frasca et al., 2020), both of which employ adjacency-based diffusion operators to preaggregate features across the graph before training. SIGN is a generalization of Simple-GCN, to extend to Personalized-PageRank (Page et al., 1998) and triangle-based operators. This allows expensive computation to be carried out by distributed computing clusters or efficient sparse graph routines such as GraphBLAS (Davis, 2019), which do not need to back-propagate through graph convolution. The prediction head can then be a shallow MLP or logistic regression. We provide implementation specifics in appendix section C to ensure repeatability.

heading: 3.3 Parameter-efficient Finetuning LLM
text: Motivated by the node classification performance gains seen by (Duan et al., 2023a) using PEFT, we finetune an LLM on the node classification task. Concretely, we use an LLM embedding model with a low-rank adapter (LoRA) (Hu et al., 2021a) and a densely connected classifier head. The pre-trained LLM weights remain frozen as the model trains on input text T to reduce loss according to:
Losscls = L(ϕ(LLM(T )), Y ) (3)
where ϕ(·) is the classifier head and Y is the label. Again, we use cross-entropy loss to compute the loss value.

heading: 4 Experiments
text: We investigate the performance of STAGE over five TAG benchmarks: ogbn-arxiv (Hu et al., 2021b), a dataset of arXiv papers linked by citations; ogbnproducts (Hu et al., 2021b), representing an Amazon product co-purchasing network; PubMed (Sen et al., 2008), a citation network of diabetes-related scientific publications; Cora (McCallum et al.,
2000), a dataset of scientific publications categorized into one of seven classes; and tape-arxiv23 (He et al., 2024), focusing on arXiv papers published after the 2023 knowledge cut-off for GPT3.5. We use the subset of ogbn-products provided by (He et al., 2024). Further details can be found in appendix Table 7.
For each experiment using Cora, PubMed or tape-arxiv23, 60% of the data was allocated for training, 20% for validation, and 20% for testing. For the ogbn-arxiv and ogbn-products datasets, we adopted the standard train/validation/test split provided by the Open Graph Benchmark (OGB)2 (Hu et al., 2021b). Our main results can be seen in Table 1. Multiple GNN models are trained using embeddings from a pre-trained LLM as node features. We ensemble the predictions across model architectures by taking the mean prediction.
Node classification accuracy is provided for various datasets, measured across multiple methods and feature types. Each column represents a spe-
2https://ogb.stanford.edu/
cific metric or method:
• hshallow: Performance using shallow features, indicating basic attributes provided as part of each dataset
• hGIANT: Results obtained by using GIANT features as proposed by (Chien et al., 2022), designed to incorporate graph structural information into LM training
• GPT3.5: Accuracy when using zero-shot predictions from GPT-3.5-turbo, demonstrating the utility of state-of-the-art language models in a zero-shot setting
• LMfinetune: Performance metrics reported by (He et al., 2024) after finetuning the DeBERTa (He et al., 2021) model on labeled nodes from the graph, showing the benefits of supervised finetuning
• hTAPE: Shows results for the TAPE features (He et al., 2024), which includes the original textual attributes of the node, GPT-generated predictions for each node, and GPT-generated
explanations of ranked predictions to enrich node features.
• hSTAGE: Reflects the model’s performance training with node features generated by a pre-trained LLM.
Instruction-biased Embeddings Textual attributes for each node are passed to the embedding LLM together with a task description which remains constant for every text, prefixing each input with a task-specific system prompt. We evaluated 3 simple task descriptions:
1. A short prompt describing the classification task for the text, as used during the pretraining stage of the LLM.
2. A description of the types of relationships between texts to form a graph, along with the classification task description. Specific graph structure for each node is not included in the prompt, unlike the proposed method from (Fatemi et al., 2024).
3. No task description.
Our findings are summarized in Table 2. Further details of the instructions can be found in appendix Table 8.
Parameter-efficient Finetuning In Table 3 we investigate the effect of using parameter-efficient finetuning (PEFT) on the pre-trained LLM, as described in (Duan et al., 2023a). We also compare this against finetuning both the LLM (using PEFT) and the GNN in unison.
Embedding Model Type In Table 4, we compare the results when using different pre-trained LLMs as the text encoder.
Diffusion GNNs Included in Table 4, we study the performance of using SimpleGCN and SIGN models individually. Model selection and implementation details can be found in the appendix sections C and D.
Ablation Study To study the impact of each component in the GNN ensemble, we perform a detailed ablation study. The results can be found in 6.

heading: 5 Analysis
text: Main Results (Table 1) We find that ensembling GNNs always leads to superior performance across datasets when taking the STAGE approach.
Despite the reduced computational resources and
training data requirements, the STAGE method remains highly competitive across all benchmarks. The ensemble STAGE approach lags behind the TAPE pipeline by roughly 5% on Cora, 3.5% on Pubmed, 0.8% on ogbn-products, and 4% on tape-
arxiv23. This is a strong result when we consider that STAGE involves training only the GNN ensemble, whereas TAPE also requires two finetuned LMs to generate node features. We see marginally superior results on the ogbn-arxiv dataset using the ensemble STAGE approach.
Instruction-biased Embedding Results (Table 2) From our findings we conclude that varying the instructions to bias embeddings has little effect on downstream node classification performance for the models we evaluated. We note that while the authors of all embedding models recommend providing instructions along with input text in order to avoid degrading performance, we did not measure a performance improvement in our experiments.
This experiment further supports our claim that an ensemble approach improves robustness across datasets and methods of node feature generation.
PEFT Results (Table 3) Finetuning each LLM gave marginal performance improvements across all datasets to varying degrees; we see the largest improvement on pubmed (3%). It is of note that finetuning significantly increases the number of trainable parameters (see Table 5) and total training time. Specifically, PEFT for 7B embedding models has over 20 million trainable parameters. On a single A100 GPU, training runs lasted 6 hours on ogbn-arxiv.
LLM Embedding Model Comparison (Table 4) All three LLM embedding models demonstrated comparable performance on the graph tasks, with each model exhibiting marginally better results on different datasets. Notably, there was no clear winner among them. The LLM2Vec model exhibited slightly weaker performance on the larger datasets (ogbn-arxiv, ogbn-products, tape-arxiv23), while it was marginally stronger on the smaller datasets (Cora, PubMed).
Ensembling the GNN models consistently ranked among the top three models across all three LLM embedding models, delivering an average performance increase of 1%. Among the individual GNN architectures, RevGAT consistently demonstrated superior performance.
Diffusion-pattern GNN Results (Table 4) The diffusion-based GNNs yielded variable results across datasets. Specifically, SIGN emerged as the second-best performer on the Cora dataset. As expected, SIGN consistently outperformed SimpleGCN, given that it generalizes the latter. Due to
its low training time, SIGN is a viable candidate for large datasets, although careful tuning of its hyper-parameters is recommended for optimal performance.
Ablation Study Results (Table 6) From our ablation study we observe that no individual GNN model outperforms any ensemble of models on any dataset. Additionally, we find that the full ensemble of MLP, GCN, SAGE and RevGAT achieve the highest and most stable accuracy scores across datasets.
Scalability An important advantage of STAGE is the lack of finetuning necessary to achieve strong results. This lies in contrast to approaches such as TAPE (He et al., 2024) and SimTeG (Duan et al., 2023a), both of which require finetuning at least one LM. Training an ensemble of GNNs and MLP head over the ogbn-arxiv dataset can be performed on a single consumer-grade GPU in less than 5 minutes. This is illustrated in Figure 2 where we compare the relationship between training time and accuracy for a number of SoTA node classification approaches. When using SIGN diffusion, training time was under 12 seconds for the ogbn-arxiv, but this came at a performance cost. Moreover, TAPE relies on text-level enhancement via LLM API calls, which adds a new dimension of cost and rate-limiting3 to consider when adapting to other datasets.

heading: 6 Conclusions
text: This work introduces STAGE, a method to use pretrained LLMs as text encoders in TAG tasks without the need for finetuning, significantly reducing computational resources and training time. Additional gains can be achieved through parameterefficient finetuning of the LLM. Data augmentation, which is orthogonal to our approach, could improve performance with general-purpose text embedding models. However, it likely remains intractable for many large-scale datasets due to the need to query a large model for each node.
We also demonstrate the effect of diffusion operators (Frasca et al., 2020) on node classification performance, decreasing TAG pipeline training time substantially. We aim to examine the scalability of diffusion-pattern GNNs on larger datasets in later work.
3https://platform.openai.com/docs/guides/ rate-limits
Future work may aim to refine the integration of LLM encoders with GNN heads. Potential strategies include an Expectation-Maximization approach or a joint model configuration (Zhao et al., 2023). A significant challenge is the requirement for large, variable batch sizes during LLM finetuning due to current neighborhood sampling techniques, which necessitates increased computational power. We anticipate that overcoming these limitations will make future research more accessible and expedite iterations.

heading: A Appendix

heading: B Negative Results
text: Co-training LLM and GNN: In a similar approach to iterative methods, we investigated cotraining the LLM and GNN on the ogbn-arxiv node classification task to facilitate a shared representation space. This proved unfeasible due to the memory requirements exceeding the capacity of one A100 GPU.
C Implementation of Diffusion Operators
We implement diffusion operators from two methods, Simple-GCN (Wu et al., 2019) and SIGN (Frasca et al., 2020). In the case of SIGN, the authors omit implementation details of the operators, so we include them here.
Let A denote the adjacency matrix of a possibly directed graph G, X its node features, and D the diagonal degree matrix of G.
We denote the random-walk normalized adjacency ARW& := AD−1 and the GCN-normalized adjacency (Kipf and Welling, 2017)
AGCN := (D + I) −1/2 (A+ I) (D + I)−1/2 (4)
The Personalized PageRank matrix is then given by (Gasteiger et al., 2022):
APPR := α (In − (1− α)ARW)−1 (5) And we denote the triangle-based adjacency matrix by A∆, where (A∆)ij counts the number of directed triangles in G that contain the edge (i, j)
Diffusion is applied to node features X by matrix multiplication. Simple-GCN takes a power k of AGCN as its diffusion operator, whilst SIGN diffusion generalizes this to concatenate powers of AGCN, APPR and A∆.
Diffusion can be calculated efficiently if sparsematrix-sparse-matrix multiplication is avoided. For both SIGN and Simple-GCN, the order of operations for applying a power of an operator Aop should be
Aop(Aop(...(Aop(X))...)︸ ︷︷ ︸ k times
(6)
as opposed to (Akop)X , where the operator matrix Aop is feasible to calculate, since the former avoids sparse matrix multiplication. In SIGN, the recursive nature of eq.6 can be exploited to reuse results for calculating successive powers.
In the case of personalized pagerank diffusion, we first use a trick from (Gasteiger et al., 2022) to approximate the diffused features of personalized pagerank matrix APPRX in linear time and avoid calculative APPR directly, by viewing eq.5 as topicsensitive PageRank (Haveliwala, 2002). We use the random-walk normalized adjacency matrix.
The following power iteration approximates APPRX (notation from (Gasteiger et al., 2022)):
Z(0) := X
Z(k+1) := (1− α)AZ(k) + αX To compute the nth diffused power, we repeat
the process n times:
Z (0) 0 = X
Z (0) n+1 = lim k→inf Z(k)n
Lastly, for triangle-based diffusion, we count triangles using linear algebra. For unweighted A we perform a single sparse matrix multiplication to obtain A2, in which element (i, j) counts the directed paths in G for node i to node j. We then calculate
A∆ = A T ⊙A2
where ⊙ denotes the Hadamard product, which can be efficiently calculated for sparse matrices. We then normalize and diffuse features over powers of A∆ in the same fashion as for AGCN .
An implementation of these operators as GraphBLAS (Davis, 2019) code is published alongside this paper.
C.1 Parallelism of diffusion operators All operations above can be be parallelized across columns of X , either keeping A in shared memory on one machine or keeping a copy on each executor in a distributed computing infrastructure like Apache Spark.

heading: D Preprocessing & Model Selection for Diffusion Operators
text: For Simple-GCN (Wu et al., 2019), we set the degree k by selecting the highest validation accuracy from k = 2, 3, 4, of which k = 2 had the highest accuracy in each case. For SIGN (Frasca et al., 2020), we choose s, p, t from the highest validation accuracy amongst (3, 0, 0) (3, 0, 1) (3, 3, 0), (4, 2, 1) (5, 3, 0). For Cora and PubMed, (4, 2, 1) was chosen, and for ogbn-arxiv, ogbn-products, and tape-arxiv23 (3, 3, 0) was chosen. We chose the number of layers for the Inception NLP to match the number of layers in other GNNs tested, 4. We did not perform additional hyper-parameter tuning. When preprocessing the embeddings, we centered and scaled the data to unit variance for Simple-GCN and SIGN only.
E Model Trainable Parameters

heading: F Ablation Study
text: To study the effect each model has on the GNN ensemble step of STAGE, we perform a detailed ablation study. The results are shown in Table 6.

heading: G Datasets
text: In this section, we describe the characteristics of the node classification datasets we used during our work. The statistics are shown in Table 7.
H Instruction-biased Embeddings
In Table 8 we list the specific instructions used to 655 investigate the effect of biasing embeddings.

heading: Future Work
text: com/docs/guides/ rate-limits
Future work may aim to refine the integration of LLM encoders with GNN heads. Potential strategies include an Expectation-Maximization approach or a joint model configuration (Zhao et al. , 2023). A significant challenge is the requirement for large, variable batch sizes during LLM finetuning due to current neighborhood sampling techniques, which necessitates increased computational power. We anticipate that overcoming these limitations will make future research more accessible and expedite iterations