Efficient and effective training of language and graph neural network models

We employ the BERT model [9] as the transformer in the LM-GNN framework to encode the nodes textual semantics. Given a node’s text BERT encodes the textual information to a $F\times 1$ embedding vector $\mathbf{x}_{n_{t}}$. This embedding vector corresponds to the [CLS] token embedding of the BERT model and the mapping from the text to the embedding is defined as $\mathbf{x}_{n_{t}}:=g_{\text{{BERT}}}(n_{t};\mathbf{W}_{\text{{BERT}}})$. The mapping is controlled by the learnable parameters $\mathbf{W}_{\text{{BERT}}}$. These parameters can be instantiated by any language model pre-training approach, e.g., masked language modeling (MLM). Pre-training BERT models on large unlabeled text data has shown significant benefit in different LM applications. However, transferring this benefit for graph ML applications is not straightforward. We employ the techniques in Sec. 4.3 to pre-train BERT models with graph data. For different node-types $t\in\{1,\ldots,T\}$ in the graph we may consider different semantic encoders; e.g. queries and products.

Although the LM-GNN framework can utilize any GNN model as an encoder [27], in this paper LM-GNN uses a modified RGCN encoder [21]. RGCNs extend the graph convolution operation [17] to heterogeneous graphs. The $l$th self-RGCN layer computes the $n$th node representation $\mathbf{h}^{(l+1)}_{n}$ as follows

where $\mathcal{N}_{n}^{r}$ is the neighborhood of node $n$ under relation $r$, $\sigma$ the rectified linear unit non linear function, $\mathbf{W}^{(l)}_{r}$ is a learnable matrix associated with the $r$th relation, and $\mathbf{W}_{\text{self}}^{(l)}$ is a projection matrix for the nodes embedding in layer $l$. Our adaptation augments the messages at each layer with a projected transformation of the node embedding for that layer acting as a self loop with an appropriate weight matrix. This adaptation over the traditional RGCN the node’s own representation based on the semantic encoder is essential for the downstream applications and needs to be treated separately by the model. The matrix $\mathbf{H}$ in this paper represents the embedding extracted in the final layer. The node features are the input of the first layer in the model i.e., $\mathbf{h}^{(0)}_{n}=\mathbf{x}_{n}$, where $\mathbf{x}_{n}$ is the node feature for node $n$.

Structure prediction task. Link prediction is a specific type of structure prediction, which will be our focus in the following section. Consider the heterogeneous graph $\mathcal{G}$ in Sec. 3. Given the sets of links $\{\mathcal{E}_{r}\}_{r=1}^{R}$, and the node features the goal of link prediction is to predict whether a new set of node pairs are linked or not.

Typically, structure prediction models utilize a contrastive loss function that requires the model to distinguish among positive and negative examples [29]. In this context, positive examples are the set of existing links in the graph. The negative examples, which are links that the model should classify as nonexistent, are typically sampled from the missing links in the graph. For each positive triplet $q=(n_{t},{r},{n^{\prime}}_{t^{\prime}})$ a number of negative links is generated by corrupting the head and tail entities at random $(n_{t},{r},{\nu^{\prime}}_{t^{\prime}})$ and $(\nu_{t},{r},{n^{\prime}}_{t^{\prime}})$.

The minimization function for link prediction can be defined as follows

where $c$ is a scoring function that return as scalar given the head, and tail nodes and the relation such as the DistMult model [28], $\mathbb{D}^{+}$ and $\mathbb{D}^{-}$ are the positive and negative sets of triplets and $y=1$ if the triplet corresponds to a positive example and $-1$ otherwise.

Node prediction task. Each node ${n}$ has a label $y_{n}\in\{0,\ldots,P-1\}$, which in the query-product network may represent the type of a product. In semi-supervised learning, we know labels only for a subset of nodes $\{y_{{n}}\}_{{n}\in\mathcal{M}}$, with $\mathcal{M}\subset\mathcal{V}$. This partial availability may be attributed to privacy concerns (medical data); energy considerations (sensor networks); or unrated items (recommender systems). The ${N}\times P$ matrix $\mathbf{Y}$ is the one-hot representation of the true node labels; that is, if $y_{n}=p$ then $Y_{{n}p}=1$ and $Y_{{n}p^{\prime}}=0,\forall p^{\prime}\neq p$. The minimization objective in this task is the cross-entropy loss.

Edge prediction task. Each link $l$ of a certain type ${r}$ may also associated with a label of interest $\psi_{l}\in\{0,\ldots,\Pi-1\}$. For example, consider the query to product graph in Fig. 1(a), where the task becomes to predict whether a pair of connected (query, product) is an exact search match or not. Hence, here given an existing link we predict the class label, which is different from the structure prediction task in Sec. 4.3 where the objective is to predict the existence of a link.

The final predicted label for link $l$ is the output of the following edge classification decoder

where $\mathbf{W}_{\text{ec}}$ is a projection matrix of appropriate dimension and $\mathbf{W}_{\text{ec}}$ denotes concatenation. Hence, the predicted label for $l$ is a function of the entity embeddings for the nodes at the endpoints of the link ${n_{t}},{n^{\prime}}_{t^{\prime}}$.

A straightforward approach would directly use the LM encoder as a semantic encoder that feeds representations to the GNN encoder, and train such an architecture in an end-to-end fashion. However, training large scale language models and graph neural networks involves challenges relating to efficiency and effectiveness.

Effectiveness challenges stem from the fact that the pre-trained language model is well optimized in language tasks but has not trained before in graph tasks, which surfaces three main issues. (1) Using such pre-trained transfomers may not be the most appropriate initialization and may trap the GNN to a sub-optimal local minimum. (2) Further, the well optimized transformer for the text tasks, may be more resistant in parameter updates. (3) Another hurdle stems from the random initialization of the GNN weights relative to the well-attuned transformer model, which may challenge the optimization of such an end-to-end framework.

Efficiency challenges relate to the large number of neighbors required by message passing in GNNs. In mini-batch training of a $k$-layer GNN the $k$-hop ego-network of every target node is created and the target node embedding is computed as a function of all the node in the expanded ego-network (also known as source nodes). The number of source nodes in an ego-network may be very large even for shallow GNNs. Alleviating this issue, recent GNN approaches apply random sampling [12] to reduce the number of neighbors. However, even with a shallow GNN (2 layers) and modest sampling (20 neighbors per layer), there are up to 400 source nodes. This remains a serious challenge in our unique setting where the transformer model needs to make 400 forward passes to calculate the embedding of a single target node. As a consequence the size of the required GPU is quite large even for small mini-batch sizes, which is a unique challenge in our framework.

Public. Our unique setting requires graph datasets where the nodes are associated with text. We employ the arxiv, and products datesets from the OGB benchmark [14] with $N$=169,343 and $E=$1,166,243 and $N$=2,449,029 $E$=61,859,140 respectively with the standard split ratios from [14]. We further augment the data with the original text features for each node; the data are collected in [8]. In the arxiv dataset the original title and abstract is used as text feature for the node. On the other hand the product dataset represents Amazon products and the product title was crawled from the web and used as the text feature. The benchmark in [14] provides text embeddings of the original text as features for the nodes. The task here is to predict the labels on the nodes in a standard semi-supervised setting. The labels are the type the paper and the category of product for arxiv and product respectively. For these datasets we also formulate a link prediction problem with splits 80% training, 10% validation and 10% testing and the tasks are predicting paper citation and product co-purchase links. Further we also construct the Yelp dataset augmented with the text using sources provided from [2]. The following node types are included with corresponding number of nodes business $N_{1}$=160,585, category $N_{2}$=1330, city $N_{3}$=836, review $N_{4}$=8,635,403, user $N_{5}$=2,189,457. The following edges are considered (user, friendship, user) $E_{1}$=17,971,548, (business, in, city) $E_{2}$=160585, (business, in category, category) $E_{3}$=708968, (review, on, business) $E_{4}$=8,635,403, (user, write, review) $E_{5}$=8,635,403. In this dataset only the review nodes are associated with text. The task here is to predict the stars for a business and is formulated as a node prediction task

Private. Additionally, we consider the dataset provided by the recent Amazon KDD22 challenge [1]. The graph structure is indeed similar to the one depicted in Fig. 1(a). There are $N_{1}=646,640$ product and $N_{2}=33,804$ query nodes in the graph and $E=781,744$ edges that represent a match among the query and the product. Each edge in this dataset is associated with a label which corresponds to whether a match between the query and the product is an exact, substitute, complement or irrelevant. This problem is known as ESCI and is solved as an edge classification task. We create a custom split for this task by splitting the set of edges to 60% for training 10% for validation and 30% for testing. Finally, we also consider the private query-purchase-product dataset that is used to predict which product will be bought based on a query. Specifically, there are $N_{1}=130,191,253$ products and $N_{2}=5,878,377$ queries. Also there are the following behavioral data represented as edges among queries and products: Based on a query a product was added in the basked $E_{1}=$15,163,234, a product was clicked $E_{2}=$24,896,086, a product was returned as candidate $E_{3}=$140,008,811, was purchased $E_{4}=$13,138,944. The task here is given a query return the most probable product to be purchased. A natural formulation for this task is under the link prediction setting.

Next we explain the different parameters that define the various approaches considered in this work.

Encoders. We consider the following possible semantic encoders in this work. BERT is the pre-trained BERT model from [26]. graph-aware BERT is the pre-trained BERT model [26] that we further fine-tune it for graph structure prediction as in equation (2). BERT-PR is a BERT model that is pre-trained using the MLM objective in the proprietary data of the company.
Task encoders. We consider 2 candidate encoders for the experiments presented here. MLP is a single-layer MLP that projects the text embeding to an appropriate dimension for node classification or for link prediction and allows us to circumvent to directly compare with the language model. This is used as a baseline approach to directly use the semantic encoder model in the downstream tasks. GraphSAGE is the model presented in [12] and is used as our baseline GNN model.
Fine-tune. This parameter defines whether we will back-propagate the loss to the semantic encoder during learning or not. The training is orders of magnitude faster when the loss is not back-propagated.
Warm-start. This parameter defines whether we will warm-start the GNN model by keeping the semantic encoder parameters fixed for some iterations before end-to-end fine-tuning.
Model configuration. We optimize the parameters such that the validation set performance is optimized. We select the number of GNN layers from $1,2,3$, GNN hidden dimension from $128,256,512$ and learning rate from $10^{-3},10^{-4},10^{-5}$.

Table 1 collects the results for the public datasets for different training configurations and encoder models in node classification. Notice that the first two rows apply the node prediction loss directly on the node representation of the text embeddings after it is appropriately mapped by a single layer MLP. Fine-tuning in this context means that the gradient updates the parameters of the semantic encoder otherwise it is not besides the MLP parameters. The target node for the Yelp dataset does not have any text hence, we can not evaluate the first two settings for that. For this experiment the warm-start did not give significant improvement and hence was not included.

By comparing lines 2 and 3 we observe that fine-tuning the BERT directly for the downstream tasks and disregarding the graph structure achieves on-par performance as the one of keeping the BERT model fixed and using these representations as input to the GNN model. This suggests that the initial BERT embeddings are indeed not the most appropriate semantic embeddings. By comparing lines 3 and 4 we see a performance benefit of fine-tuning the BERT model through the GNN, since the multi-hop information is captured by the GNN. By comparing lines 3 and 5 we observe the clear advantage of the graph-aware BERT. The graph-aware pre-fine-tuning fuses the transformer with graph information and is the most suitable semantic encoder. Finally, line 6 coincides with the proposed LM-GNN framework. We observe that this configuration achieves the best overall performance and includes the proposed stage-wise fine-tuning approach. Hence, fine-tuning the BERT model for link prediction provides good performance in the node classification tasks. This result is very important since it allows to train a BERT model on link prediction and transfer the knowledge on different downstream tasks which may speed up the overall training.

Table 2 collects the link prediction performance of the various baselines measured using the MRR score. The first row applies the link prediction supervision in (2) directly on the node representation of the text encoding after mapped by a single layer MLP to an embedding and the whole architecture is fine-tuned end-to-end. This model corresponds to the graph-aware pre-fine-tuned model for link prediction and is the same as the used as a semantic encoder in lines 5 and 6.

By comparing lines 1 versus 2, 3, and 4 we observe that the original BERT model is indeed not appropriate as a semantic encoder used with the GNN. On the other hand, fine-tuning the BERT model for link prediction in row 1 achieves a very good MRR performance. Lines 5, 6, 7 relative to 2, 3, 4 showcase the advantage of using the graph-aware pre-fine-tuning as an essential step in our LM-GNN framework, where the former leads to a large performance boost. Furthermore, by comparing 5 with 6 and 7 we observe the necessity of warm-starting in certain cases of the GNN encoder to avoid non-desirable local minima. Since the GNN model is initialized at random and the graph-aware BERT is well-trained optimizing this model without warm-start is challenging. By comparing lines 3 and 4 we see a performance benefit of fine-tuning the BERT model through the GNN, since the multi-hop information is captured by the GNN.

The Table 3 contains the edge classification results for the various baselines in ESCI. Note that here we also report the performance for each individual class since we are interested in predicting well also for the rare classes in our application.

By comparing rows 2 and 4 that both fine-tune the BERT-PR model we observe a strong boost of 320 bps in performance when the GNN is used. This indicates that the GNN can indeed help boosting the performance probably for the rare classes (S-C-I) by a large extend. By comparing rows 3 and 4 we observe that it is very important to fine-tune the BERT-PR model during GNN training. By comparing rows 3 and 6 we observe that the graph-aware pre-fine-tuning is giving a significant boost when the BERT embeddings are fixed. This benefit diminishes when the BERT model is fine-tuned. Finally, the performance in lines 5-8 is quite similar, but interesting the performance in the rare classes is maximized in rows 5 and 6. We plan to dive deeper into these results and analyze the performance for different sample sizes besides the current split.

Table 4 collects the results for predicting if the search query leads to the purchase of a product and is treated as a link prediction task. The evaluation metric is the Macro recall at 100, which is the percentage of true products that exist in the top 100 retrieved products by each method. The results for fine-tuning the BERT model via the GNN model are ongoing. By comparing rows 1 and 3 that do not fine-tune the BERT-PR model we observe a very large performance boost by the GNN model that considers the graph structure. An even larger performance boost is observed when fine-tuning the BERT model via graph information in row 2. The best performance is observed by using the graph-aware BERT model as a fixed semantic-encoder for the GNN model. Future steps will focus on end-to-end fine-tuning.