Pretraining Language Models with Text-Attributed Heterogeneous Graphs (2023)

Abstract: In many real-world scenarios (e.g., academic networks, social platforms), different types of entities are not only associated with texts but also connected by various relationships, which can be abstracted as Text-Attributed Heterogeneous Graphs (TAHGs). Current pretraining tasks for Language Models (LMs) primarily focus on separately learning the textual information of each entity and overlook the crucial aspect of capturing topological connections among entities in TAHGs. In this paper, we present a new pretraining framework for LMs that explicitly considers the topological and heterogeneous information in TAHGs. Firstly, we define a context graph as neighborhoods of a target node within specific orders and propose a topology-aware pretraining task to predict nodes involved in the context graph by jointly optimizing an LM and an auxiliary heterogeneous graph neural network. Secondly, based on the observation that some nodes are text-rich while others have little text, we devise a text augmentation strategy to enrich textless nodes with their neighbors’ texts for handling the imbalance issue. We conduct link prediction and node classification tasks on three datasets from various domains. Experimental results demonstrate the superiority of our approach over existing methods and the rationality of each design. Our code is available at https://github.com/Hope-Rita/THLM.

## introduction pretrained language models (plms) @xcite @xcite that built upon the transformer architecture @xcite have been successfully applied in various downstream tasks such as automatic knowledge base construction @xcite and machine translation @xcite . due to the design of pretraining tasks (e.g., figure 1 : as an instance of tahg, an academic network contains three types of nodes (papers, authors, and keywords) with textual descriptions as well as their multi-relational connections. in fact, texts not only carry semantic information but also are correlated with each other, which could be well represented by text-attributed heterogeneous graphs (tahgs) that include multi-typed nodes with textual descriptions as well as relations. see figure 1 for an example. generally, tahgs usually exhibit the following two challenges that are struggled to be handled by existing plms. ## preliminaries a pretrained language model (plm) can map an input sequence @xmath0, x 2 , • • • , x l ) of l tokens into their contextual representations @xmath1, h 2 , • • • , h l ) with the design of pretraining tasks like masked language modeling @xcite , next-token prediction @xcite , autoregressive blank infilling @xcite . in this work, we mainly focus on the encoder-only plms (e.g., bert @xcite , roberta @xcite ) and leave the explorations of plms based on encoder-decoder or decoder-only architecture in the future. a text-attributed heterogeneous graph (tahg) @xcite usually consists of multityped nodes as well as different kinds of relations that connect the nodes. each node is also associated with textual descriptions of varying lengths. mathematically, a tahg can be represented by @xmath2, e, u, r, x ), where v, e, u and r denote the set of nodes, edges, node types, and edge types, respectively. each node @xmath3,v has a type ψ(e u,v ) ∈ r. x is the set of textual descriptions of nodes. note that a tahg should satisfy |u| + |r| > 2. existing plms mainly focus on textual descriptions of each node separately, and thus fail to capture the correlations among different nodes in tahgs (as explained in section 1). to address this issue, we propose a new framework for pretraining lms with tahgs, aiming to obtain plms that are aware of the graph topology as well as the heterogeneous information. ## methodology figure 2 shows the overall framework of our proposed approach, which mainly consists of two components: topology-aware pretraining task and text augmentation strategy. given a tahg, the first module extracts the context graph for a target node and predicts which nodes are involved in the context graph by jointly optimizing an lm and an auxiliary heterogeneous graph neural network. it aims to enable plms to capture both first-order and high-order topological information in tahgs. since some nodes may have little textual descriptions in tahgs, the second component is further introduced to tackle the imbalanced textual descriptions of nodes, which enriches the semantics of textless nodes by neighbors' texts. it is worth notic-ing that after the pretraining stage, we discard the auxiliary heterogeneous graph neural network and only apply the plm for various downstream tasks. ## conclusion in this paper, we pretrained language models on more complicated text-attributed heterogeneous graphs, instead of plain texts. we proposed the context graph prediction task to inject the graph learning ability of graph neural networks into lms, which jointly optimizes an auxiliary graph neural network and an lm to predict which nodes are involved in the context graph. to handle imbalanced textual descriptions of different nodes, a text augmentation strategy was introduced, which enriches the semantics of textless nodes by combining their neighbors' texts. experimental results on three datasets showed that our approach could significantly and consistently outperform existing methods across two downstream tasks.