Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Aaron Zolnai-Lucas, Jack Boylan, Chris Hokamp, Parsa Ghaffari

TL;DR

STAGE tackles TAG node classification by using off-the-shelf LLM embeddings as fixed node features and training an ensemble of GNNs on the enriched graphs. It proposes a one-stage, finetuning-free pipeline and investigates diffusion-based GNNs to scale to larger graphs while keeping the approach simple. Empirical results across multiple TAG benchmarks show competitive performance with notable reductions in training complexity compared to SoTA methods like TAPE and SimTeG. The study highlights the robustness of zero-shot LLM embeddings for node features and demonstrates that diffusion operators can further reduce training time on large graphs.

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 pre-trained 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.

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

TL;DR

STAGE tackles TAG node classification by using off-the-shelf LLM embeddings as fixed node features and training an ensemble of GNNs on the enriched graphs. It proposes a one-stage, finetuning-free pipeline and investigates diffusion-based GNNs to scale to larger graphs while keeping the approach simple. Empirical results across multiple TAG benchmarks show competitive performance with notable reductions in training complexity compared to SoTA methods like TAPE and SimTeG. The study highlights the robustness of zero-shot LLM embeddings for node features and demonstrates that diffusion operators can further reduce training time on large graphs.

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 pre-trained 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.
Paper Structure (38 sections, 9 equations, 2 figures, 8 tables)

This paper contains 38 sections, 9 equations, 2 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Graph Learning and LLMs
  3. Simplifying Node Representation Generation
  4. Scalable GNN Architectures
  5. Background
  6. Text-Attributed Graphs
  7. LLMs as Text Encoders
  8. Language Models and GNNs
  9. Approach
  10. Text Embedding Retrieval
  11. GNN Training
  12. Diffusion-based GNNs
  13. Parameter-efficient Finetuning LLM
  14. Experiments
  15. Instruction-biased Embeddings
  16. ...and 23 more sections

Figures (2)

  • Figure 1: Our proposed approach to node classification. Firstly, the textual attributes of the input graph nodes are encoded using an off-the-shelf LLM. The text embeddings will be used alongside the graph adjacency matrix as input to train a downstream ensemble of GNNs. GNN predictions are then mean-pooled to obtain the final prediction.
  • Figure 2: The performance trade-off between node classification accuracy and total training time on ogbn-arxiv for SoTA LM-GNN methods. The STAGE model uses text embeddings generated from Salesforce-Embedding-Mistral and an ensemble of GNNs (GCN, SAGE and RevGAT) and MLP. The size of each marker indicates the total number of trainable parameters. Figure adapted from he2024harnessing.