SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yanchao Tan
College of Computer and
Fuzhou University
yctan@fzu.edu.cn
reviverkey@gmail.com
Hang Lv
lvhangkenn@gmail.com
Carl Yang*
j.carlyang@emory.edu
Graphs are widely used to model interconnected entities and improve downstream predictions in various real-world applications. However, real-world graphs nowadays are often associated with complex attributes on multiple types of nodes and even links that are hard to model uniformly, while the widely used graph neural networks (GNNs) often require sufficient training toward specific downstream predictions to achieve strong performance. In this work, we take a fundamentally different approach than GNNs, to simultaneously achieve deep joint modeling of complex attributes and flexible structures of real-world graphs and obtain unsupervised generic graph representations that are not limited to specific downstream predictions. Our framework, built on a natural integration of language models (LMs) and random walks (RWs), is straightforward, powerful and data-efficient. Specifically, we first perform attributed RWs on the graph and design an automated program to compose roughly meaningful textual sequences directly from the attributed RWs; then we fine-tune an LM using the RW-based textual sequences and extract embedding vectors from the LM, which encapsulates both attribute semantics and graph structures. In our experiments, we evaluate the learned node embeddings towards different downstream prediction tasks on multiple real-world attributed graph datasets and observe significant improvements over a comprehensive set of state-of-the-art unsupervised node embedding methods. We believe this work opens a door for more sophisticated technical designs and empirical evaluations toward the leverage of LMs for the modeling of real-world graphs.
Graphs are widely used to model interconnected entities, and they are critical in enhancing downstream predictions in various real-world applications. Nowadays, real-world graphs are often associated with complex attributes on multiple types of nodes and even links [27, 28], and modeling such real-world graphs is non-trivial. For example, in the schema of a clinical attributed graph constructed
(a) Schema of real-world attributed graph
(b) Attributed random walk
(c) Composed text
Figure 1: A toy example of the transformation from a real-world attributed graph to the composed text. (a) is a schema of a real-world attributed graph on MIMIC-III that delineates how nodes (e.g., patients), edges (e.g., has between patients and visits), and the associated attributes (e.g., age) are organized and interconnected. (b) is attributed random walk for capturing structural information and can be composed to text in (c).
from MIMIC-III $^2$  data (shown in Figure 1(a)), there are multiple types of nodes such as patients, visits and diseases, each with their own attributes such as age, sex, time and name; there are also multiple types of links associated with different meanings. Such complex heterogeneous attributes of nodes and links attributes can hardly be modeled in a uniform space.
Although the widely used graph neural networks (GNNs) have shown remarkable successes in the modeling of attributed graphs for various downstream applications [5, 14, 53, 53], the representation learning (a.k.a. embedding) of GNNs often requires sufficient training toward specific downstream predictions to achieve strong performance [4, 76]. While unsupervised training has also been explored for GNNs [23, 63], the generic performances of unsupervised GNN embeddings across different downstream tasks are still unsatisfactory (as we will also show in our experiments). Consequently, it is important to devise a general-purpose graph embedding method to simultaneously understand the complex node/link attributes and incorporate the flexible graph structures in an unsupervised manner.
However, two obstacles stand in the way of achieving this goal. First, the nature of the attributes can be intricate and diverse, thus understanding their semantics in a uniform space is non-trivial. Second, the graph structures need to be accurately captured and incorporated into the embedding space, which is not straightforward due to the inherent flexibility and potential complexity of entity relations.
To address these issues, in this work, we take a fundamentally different approach than GNNs, named WalkLM, which is a uniform framework to obtain unsupervised generic graph representations that are not limited to specific downstream tasks. To this end, we first draw inspiration from the recent successes of language models (LMs), and propose to leverage LMs as uniform attribute models that can capture the intricate semantics of complex heterogeneous attributes of nodes and links. Secondly, we propose to leverage the classic tool of random walks (RWs) on graphs which have been shown effective in capturing flexible graph topological structures by various studies [8, 10, 13, 21, 38].
Specifically, we first generate attributed RWs on the graph (e.g.,  $1 \xrightarrow{has} 2 \ldots \xrightarrow{for} 5$  in Figure 1(b)), and design an automated textualization program to compose roughly meaningful textual sequences directly from the attributed RWs. As shown in Figure 1(c), the composed text is a mapping from the attributed RW in Figure 1(b), where a uniform automatic program firstly textualize different types of nodes (in different colors) by properly concatenating the nodes with the names and values of different attributes, and then textualize the whole RWs by concatenating the nodes and links. Furthermore, we fine-tune an LM using the RW-based textual sequences and extract embedding vectors from the LM. The learned embeddings encapsulate both attribute semantics and graph structures, and can be flexibly utilized for arbitrary downstream tasks.
In our experiments, we take the node embeddings and evaluate them towards different downstream prediction tasks (e.g., node classification and link prediction) on multiple real-world attributed graph datasets and observe significant improvements over a comprehensive set of state-of-the-art unsupervised node embedding methods (e.g., Wa1kLM achieves an average of  $88.98\%$  improvement over the state-of-the-art baselines ranging from existing RW-based graph embedding methods to popular unsupervised GNN modes regarding both Macro-F1 and Micro-F1 metrics). We believe this
$^{2}$ https://physionet.org/content/mimiciii/1.4/
work paves the way for more sophisticated technical designs and empirical evaluations toward the leverage of LMs for the modeling of real-world graphs.
Graph Representation Learning. In recent years, a plethora of representation learning techniques are proposed for graphs [3, 15, 55, 58, 74]. In this work, we focus on the objective of learning embedding vectors for each node that characterizes its topological (and semantic) information in the graph. Among existing node embedding methods, many have analyzed and utilized the great promise of random walks (RWs) in capturing the topological structures of graphs [8, 13, 21, 38]. However, the above methods ignore the abundant attribute information surrounding the nodes and edges [32]. In recent studies, Graph neural networks (GNNs) for learning node representations through aggregating information from neighboring nodes on graphs [14, 24, 62]. However, most existing GNNs are established in a supervised learning setting, which requires abundant task-specific labeled data that may not be available in real-world applications [4, 76], and the embeddings they learn are not generalizable across different downstream tasks [70]. Although some studies tried to reduce the labeling effort by pre-training an expressive GNN model on unlabeled data with self-supervision methods (e.g., contrastive learning) [19, 22, 75], their performances in specific downstream tasks still rely much on the properly chosen self-supervision tasks and attribute encoders [47, 73]– that is, there still lack a uniform framework for generic unsupervised attributed graph representation learning.
Language Models (LMs). With the massive corpora and powerful computation resources for pretraining, modern language models (LMs) derive various families [34]. These LMs can be grouped into: (1) auto-regressive LMs (e.g., GPT [40] and GPT-2/3 [1, 41]), (2) masked LMs (e.g., BERT [7], RoBERTa [31], and XLNet [64]), and (3) encoder-decoder LMs (e.g., BART [26] and T5 [42]). LMs have been intensively studied by NLP researchers for various language-related tasks [16, 29, 43]. In our work, we innovatively utilize LMs as uniform attribute models for nodes and links in graphs for the first time. Note that, our work also readily generalizes to recent large language models (LLMs) (e.g., InstructGPT [36], ChatGPT and GPT-4 [35]) via appropriate parameter-efficient training approaches (e.g., LoRA [18] and prompt-tuning [25, 30]). However, those are orthogonal to the innovations in this work, for which we leave the exploration as immediate future works.
LMs with Knowledge Graph (KG). In recent studies, combining LMs with KG has been widely applied in various real-world applications [37]. Among existing methods, many have proposed to enhance LMs with KG for significantly improve the performance of LMs in accessing domain-specific knowledge [46, 56, 67, 69], and the others proposed to harness the power of LMs for addressing KG-related tasks [65, 66]. However, the above methods fail to effectively combine the rich semantic information of graphs with global topological information. Furthermore, as closest to us, [12] proposed to combine random walks and neural network language models to produce new word representations. However, it ignores the rich relational information between nodes and thus fails to learn richer topological information. Moreover, compared with modern LMs with extensive prior knowledge, the text-based encoder used in [12, 51] fails to extract richer semantic information.
Given an attributed graph  $\mathcal{G}$  and multiple downstream tasks  $\tau_{i}$  (e.g., node classification  $\tau_{1}$  and link prediction  $\tau_{2}$ ), the goal of WalkLM is to sufficiently model information in  $\mathcal{G}$  and improve task performances on  $\tau_{i}$ . Specifically, we denote the graph as  $\mathcal{G} = (V,E,\Phi ,\Psi)$ , where each node  $v_{i}\in V$  is associated with attributes  $\Phi (v_{i})$ , and each edge  $e_i\in E$  is associated with attributes  $\Psi (e_i)$ .
To fully exploit both the semantic information in  $\Phi$  and  $\Psi$ , and the structural information in  $V$  and  $E$ , we first design attributed random walks (RWs) based automated textualization program on  $\mathcal{G}$ , where we can transform the attributed graph to the meaningful textual sequences  $\mathbb{W} = \{W_i\}_{i=1}^N$ . Then, we fine-tune a graph-aware language model (LM) using the RW-based textual sequences  $\mathbb{W}$ , find the optimized parameters  $\Theta$  of the LM, and extract embedding vectors from the LM. Finally, we apply these embeddings to  $\tau_1$  and  $\tau_2$  for performance evaluation.
Figure 2: The overall framework of WalkLM: (a) An attributed random walk (RW) based textualization program obtains attributed RWs and composes roughly meaningful textual sequences directly from the attributed RWs. (b) A pre-trained general LM is fine-tuned in a graph-aware fashion using the generated textual sequences to produce generic graph embeddings. (c) The learned graph embeddings can be applied to various downstream tasks.
3.2 Masked Language Modeling
Masked Language Modeling (MLM) is a widely-used and highly effective technique for training large-scale language models [7]. In the MLM task,  $15\%$  of the language tokens are randomly chosen for prediction. If the  $i$ -th token in a sequence  $W$  is chosen, it can be replaced by (1) the [MASK] token  $80\%$  of the time, (2) a random token  $10\%$  of the time, and (3) the unchanged  $i$ -th token  $10\%$  of the time. Then, token  $t_i$  will be used to predict the original token with following cross-entropy loss [7, 59, 60]:
$$
\mathcal {L} _ {M L M} = - \frac {1}{| \mathbb {X} |} \sum_ {X \in \mathbb {X}} \sum_ {t _ {i} \in \mathcal {M}} \log p (t _ {i} | T _ {\backslash i}), \tag {1}
$$
where  $\mathbb{X}$  is a set of training examples,  $\mathcal{M}$  is the prediction set of the masked token,  $T_{\backslash i} = \{t_1,\dots ,t_{i - 1},t_{i + 1},\dots ,t_L\}$  is the surrounding token set of  $t_i$ , and  $|T_{\backslash i}| = L - 1$
In this section, we present the proposed method WalkLM, which comprises two major components. The first component, the attributed random walk (RW) based textualization program, captures both topological structures and attribute information of the graph and composes corresponding text automatically. The second component, the graph-aware language model (LM) fine-tuning, leverages pre-trained LM to encode the complex semantics along with the graph structure. The overall model architecture is shown in Figure 2.
4.1 Attributed RW-based Textualization Program
To model node/link attributes, traditional machine learning algorithms require a standard process of vectorization, which transforms different types of attributes into categorical or numerical features as model input. However, such a vectorization process removes the actual semantic meanings of the attributes, and it cannot unify different types of attributes (e.g., ages, sexes, time, etc.) in the same space. Inspired by the recent successes of LMs, we find it promising to leverage pre-trained LMs to understand the intricate semantics of complex heterogeneous attributes [29, 34]. The key idea is to perform textualization instead of vectorization, that is, to transform different types of attributes into texts, which can then be modeled by the LMs in a uniform space. Therefore, we first design the following process for the textualization of individual entities in the graph (i.e., nodes and edges).
Entity-level Textualization. Inspired by a wide range of NLP tasks that leverage prompts to construct informative rule templates, we propose to textualize attributed graph entities via a rule-based program function  $\mathcal{P}(\cdot)$  that automatically concatenates attribute values with the corresponding attribute names, as well as attributes and the corresponding entity id. For example, as shown in Figure 1, for a patient node  $v_{i}$  with attributes  $\Phi (v_{i}) = \{age:35$  , sex : female, pid : P246}, the textualization program will convert it into  $\mathcal{P}(v_i) = <  A$  35-year-old female patient P246 >. For edges, in this work, we only consider simple relational edges such as has and including, which are already texts, but our framework is readily extensible to edges with more complex attributes.
For attributed graphs, after modeling the complex attributes of individual entities, the next challenge would be to model the flexible graph topological structures. To this end, we propose to utilize the
powerful and efficient tool of RWs, as the foundation for the textualization of graphs. Specifically, we design the following process:
Walk-level Textualization. We first initiate an attributed RW  $W$  by randomly selecting a node  $v_{0}$  and attaching the textualized node information  $\mathcal{P}(v_0)$  to  $W$ . Then, we extend  $W$  by randomly selecting an edge  $e_1$  starting from  $v_{0}$  as in a standard RW with a uniform probability as 1 divided by the number of out-edges of  $v_{0}$ , with its terminating node  $v_{1}$ , and appending the corresponding texts  $\mathcal{P}(e_1)$  and  $\mathcal{P}(v_1)$  to  $W$ . We can keep adding edges and nodes until the random walk terminates, such as based on a termination probability  $\alpha$ . The final textual sequence  $W = \{\mathcal{P}(v_0),\mathcal{P}(e_1),\mathcal{P}(v_1),\ldots ,\mathcal{P}(v_{L - 1}),\mathcal{P}(e_L),\mathcal{P}(v_L)\}$ , which corresponds to an actual attributed random walk with the length of  $2L + 1$  on the graph, will be a roughly meaningful sentence, such as the one shown in Figure 1(c).
After performing the RW for  $N$  times, we can obtain  $N$  attributed RWs as  $\mathbb{W} = \{W_{ij}\}_{i=1}^{N}$ , which constitutes a graph-aware corpus for training LMs without any downstream task supervision.
Discussion. We believe that the above proposed attributed RW-based textualization program is effective in capturing both complex attribute information and flexible topological structures of the graph in an unsupervised manner. Such ability arises from two critical properties: (1) Random walks are known to be capable of providing characteristic graph traits and reconstructing network proximity of nodes [21, 33]. To be specific, it has been proven that the distribution of random walks starting at a particular node, which can be estimated with sufficient numbers of random walks, can sufficiently preserve the subgraph structure around the node. This means a sufficient amount of attributed RWs from different nodes can well reflect the topological structures of graphs. (2) Our textualization program completely preserves the attributes of nodes and edges, as well as the whole RWs, and it presents such information as meaningful texts, which is natural for LMs to comprehend. Moreover, RWs are known to be efficient and highly parallelizable, where numerous threads can run simultaneously to generate large amounts of RWs [9]. Note that, we only need to perform the rule-based textualization once for every node during the pre-processing stage, which is also efficient and highly amenable to parallelization.
4.2 Graph-Aware LM Fine-Tuning
Despite the robust generalizability of LMs, fine-tuning remains a necessary step [57, 68], which allows the general LM to adapt its broad language capabilities to the specificities of the different attributed graphs.
As one of the mainstream language modeling techniques, masked language modeling (MLM) is proven to sufficiently utilize textual semantic information for further fine-tuning LMs [39, 77]. To achieve the balance between effectiveness and efficiency, we propose a graph-aware LM fine-tuning mechanism with knowledge distillation [2, 17, 44]. Specifically, we adopt a general LM DistilRoBERTa (abbr. DROBERTa) $^3$  as our starting point for fine-tuning, where RoBERTa is a widely used successor of BERT [7]. Note that, DROBERTa can further reduce the size of the original RoBERTa model by  $40\%$  and achieve  $60\%$  faster training while retaining  $97\%$  capacity of RoBERTa's language understanding [44]. Then, we feed the attributed RW  $W \in \mathbb{W}$  to the LM tokenizer and obtain the corresponding token list  $\mathcal{T} = \{t_1, t_2, \ldots, t_K, < \text{MASK} >_1, < \text{MASK} >_2, \ldots, < \text{MASK} >_{|\mathcal{M}|}\}$ , where  $t_i$  denotes the unmasked token, and  $< \text{MASK} >_i$  denotes the token chosen for prediction. In this way, we create a training example  $X_i = \langle W, \mathcal{T} \rangle \in \mathbb{X}$  for fine-tuning, where  $\mathbb{X} = \{X_1, X_2, \ldots, X_N\}$  is the whole training dataset. We adopt the cross-entropy loss as the fine-tuning objective [7, 44, 59], which is formulated as follows:
$$
\mathcal {L} _ {F T} (\Theta) = - \frac {1}{| \mathbb {X} |} \sum_ {X _ {i} \in \mathbb {X}} \left[ \sum_ {t _ {k} ^ {*} \in \mathcal {M}} \log \frac {\exp (S i m \left(t _ {k} , t _ {k} ^ {*}\right))}{\sum_ {t \in \mathcal {V}} \exp (S i m \left(t _ {k} , t\right))} \right], \tag {2}
$$
where  $\Theta$  is the learnable parameters of our graph-aware LM,  $\mathcal{M}$  is the ground-truth set of the masked token,  $\nu$  is the token vocabulary,  $t_k$  is the prediction token,  $t_k^*$  is the ground-truth token, and  $Sim(t_i,t_j)$  is the similarity scoring function between  $t_i$  and  $t_j$ .
After obtaining the fine-tuned LM based on MLM, we can extract generic graph embeddings (e.g., node embeddings based on node name). For example, we can access the representation of disease embedding in Figure 1(c) via extracting the embedding of Epistaxis.
<sup>3</sup>https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation
Complexity Analysis. The time complexity of fine-tuning is  $\mathcal{O}(Iter\cdot |N|\cdot l_{avg}^2\cdot d)$ , where Iter is the number of iterations of training,  $N$  is the number of training examples,  $l_{avg}$  is the average length of input textual sequences for the LM and  $d$  is the dimension of embedding. We infuse global graph structure knowledge into the LM to distinguish similar positions instead of negative sampling, making it possible to fine-tune more efficiently [59].
Model Extension. Our framework is a fundamental approach to integrate LMs and RWs for generic attributed graph embedding, which can choose different LMs according to different tasks and domains (shown in our experiments in Sec. 5.4) and generalize to recent large language models (LLMs) (e.g., InstructGPT [36], ChatGPT and GPT-4 [35]) via appropriate parameter-efficient training approaches such as LoRA [18]. A full exploration of different LMs is orthogonal to the main contributions in this work, which is left as future work.
Datasets Extension. Our method introduces the novel process of textualization, which converts general attributed graphs into text-like sequence data. This process allows us to leverage the capabilities of pre-trained language models for graph representation learning. Note that, our method only requires some meaningful attributes on the graphs, which are available in most real-world graphs such as biological networks, social networks, and knowledge graphs. Some preliminary experimental results of graph classification and KG-related datasets are shown in Appendix A.1 and Appendix A.2.
4.3 Various Downstream Tasks
In this work, we focus on node embeddings since they are most commonly studied for graph representation learning, and it is straightforward to extract node embeddings from the fine-tuned LM based on node names (or node IDs if the node has no meaningful name). However, Wa1kLM can also easily generate edge embeddings, by adding edge names (e.g., relation names) or edge IDs to the textualization process, and even obtain subgraph/graph embeddings via appropriate embedding aggregation mechanisms. The extracted embeddings can be directly used as fixed feature vectors to train downstream prediction models for tasks such as node classification or link prediction. Alternatively, these embeddings can also serve as initialization for more learnable embeddings in complex neural network models, which can be further updated according to the specific requirements of the downstream task.
Datasets. We conduct extensive experiments on two real-world datasets, PubMed  $^{4}$  and MIMIC-III  $^{5}$ . PubMed contains a graph of genes, diseases, chemicals, and species. The nodes and edges are extracted according to [61]. A relatively small fraction of diseases are grouped into eight categories. MIMIC-III contains a graph of diseases, patients, and visits, where nodes and relations are extracted from clinical records. Diseases are classified into nineteen categories according to ICD-9-CM  $^{6}$ . The detailed statistics are shown in Table 1.
Table 1: Statistics of the datasets.
Dataset	#attribute type	#node type	#node	#link type	# link	#label	#label node
PubMed	8	4	63,109	10	244,986	8	454
MIMIC-III	10	3	32,267	4	559,290	19	4880
Competitors. We compare our proposed Wa1kLM with ten graph-oriented baselines that are designed for heterogeneous information networks (HINs) or knowledge graphs (KG), which can handle different types of nodes and edges. We divided them into four groups as follows:
(1) RW-based methods: Metapath2Vec (abbr. M2V) [8] proposes to use user-defined meta-paths as guidance, so as to learn node embeddings on HINs. HIN2Vec [10] carries out multiple prediction
<sup>5</sup>https://physionet.org/content/mimiciii/1.4/
<sup>6</sup>https://www.cdc.gov/nchs/icd/icd9cm.htm