STAGE : Simplified Text-Attributed Graph Embeddings using Pre-trained LLM s (2024)

Abstract: We present 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.

## introduction a knowledge graph (kg) typically includes entities (represented as nodes), relationships between entities (represented as edges), and attributes of both entities and relationships @xcite . 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 @xcite , words and lexical semantics @xcite , and commercial products @xcite . 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 @xcite . examples of unstructured data values in tags * authors contributed equally. ## approach our cascading approach consists of two steps: ensembling the predictions from multiple gnn architectures was motivated by our observation of strong performance by different models across different datasets. ## experiments we investigate the performance of stage over five tag benchmarks: ogbn-arxiv @xcite , a dataset of arxiv papers linked by citations; ogbnproducts @xcite , representing an amazon product co-purchasing network; pubmed @xcite , a citation network of diabetes-related scientific publications; cora (mccallum et al., the experiment is run over four seeds, with mean accuracy and standard deviation shown. the best results are coloured green (first), yellow (second), and orange (third). for h stage , we use sfr-embedding-mistral as the embedding model on ta features only, and the simple task instruction to bias the embeddings. we adapt the table from @xcite and include our results. 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 @xcite . our main results can be seen in node classification accuracy is provided for various datasets, measured across multiple methods and feature types. each column represents a spe-2 https://ogb.stanford.edu/ cific metric or method: 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: ## analysis main results ( instruction-biased embedding results ( peft results ( ## conclusions 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. 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 @xcite . 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. 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 a rw & := ad -1 and the gcn-normalized adjacency @xcite a gcn := (d + i) -1/2 (a + i) (d + i) -1/2 (4) the personalized pagerank matrix is then given by @xcite : 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 a gcn as its diffusion operator, whilst sign diffusion generalizes this to concatenate powers of a gcn , a ppr 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 a op should be as opposed to (a k op )x, where the operator matrix a op 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 @xcite to approximate the diffused features of personalized pagerank matrix a ppr x in linear time and avoid calculative a ppr directly, by viewing eq.5 as topicsensitive pagerank @xcite . we use the random-walk normalized adjacency matrix. the following power iteration approximates a ppr x (notation from @xcite ): to compute the nth diffused power, we repeat the process n times: lastly, for triangle-based diffusion, we count triangles using linear algebra. for unweighted a we perform a single sparse matrix multiplication to obtain a 2 , in which element (i, j) counts the directed paths in g for node i to node j. we then calculate an implementation of these operators as graph-blas @xcite code is published alongside this paper.