GNN-LM:
L ANGUAGE
M ODELING
BASED
ON
G LOBAL C ONTEXTS VIA GNN
Anonymous authors
Paper under double-blind review
A BSTRACT
Inspired by the notion that “ to copy is easier than to memorize ”, in this work, we introduce GNN-LM, which extends vanilla neural language model (LM) by allowing to reference similar contexts in the entire training corpus.
We build a directed heterogeneous graph between an input context and its semantically related neighbors selected from the training corpus, where nodes are tokens in the input context and retrieved neighbor contexts, and edges represent connections between nodes.
Graph neural networks (GNNs) are constructed upon the graph to aggregate information from similar contexts to decode the token.
This learning paradigm provides direct access to the reference contexts and helps improve a model’s generalization ability.
We conduct comprehensive experiments to validate the effectiveness of the GNN-LM: GNN-LM achieves a new state-of-the-art perplexity of 14.8 on WikiText-103 (a 4.5 point improvement over its counterpart of the vanilla LM model), and shows substantial improvement on One Billion Word and
Enwiki8 datasets against strong baselines.
In-depth ablation studies are performed to understand the mechanics of GNN-LM.
I NTRODUCTION
Language modeling (LM) is a basic and long-standing task in natural language processing (Shannon, 2001; Bahl et al., 1983; Chen & Goodman, 1999; Mikolov et al., 2012;
Xie et al., 2017).
It aims at predicting the upcoming token given the sequence of previous context consisting of a sequence of tokens.
A common practice to train a language model is to enforce the model to maximize the probability of the upcoming ground-truth token at training time.
At test time, the next token to predict could be the one with the highest probability (via greedy search) or the one that maximizes a window of tokens through the beam search strategy.
This form of training-test procedure can be viewed as a process of memorization , or doing a close-book examination , if we compare the training data to a book and inference to doing an examination: The process of iterating N epochs over the training data is comparable to reviewing the book N times and the model needs to memorize what is the most likely to appear given speciﬁc context based on the training data.
At test time, the book needs to be closed, i.e., the model does not have means to refer to the training data at test time, and the model has to invoke related memory to predict the next token during inference.
There are two limitations to this close-book examination strategy: (1) the memorization-based language models are usually hard to memorize the knowledge of hard examples (e.g., long-tail cases in the training set); (2) memory required to memorize the whole training data is usually intensive.
The difﬁculty of resolving these two problems can be substantially alleviated if the model can be provided with related contexts from the training set so that the model can reference them for decisions.
This process can be viewed as a strategy different from memorization or close-book examination – copy , or in other words, open-book examination .
For example, given a preﬁx “ J. K. Rowling is best known for writing ” and we want to predict the upcoming token, a language model will more easily generate token “ Harry ” if it can refer to the context “ J. K. Rowling wrote the Harry Potter fantasy series”.
Motivated by the observation that “ to copy is easier than to memorize ”, or “ an open-book exam is easier than to a close-book exam ”, in this work, we introduce a new language modeling scheme – GNN-LM, which provides an LM model with the ability to reference similar contexts from the
Input Context 𝒄 𝒕
The movie is
Base LM 𝒉 𝑡 𝑐 (1) : This movie is great .
𝑐 (2) :
Those movies are bad .
𝑐 (3) :
This movie is what I like.
Neighbor Contexts
The movie is movies is what movie is great movies are bad
Graph Construction ( 𝒍 = 𝒓 = 𝟏 )
TrainingDatastore
Inter-context edge
Intra-context edge
Input query 𝒄 𝒕 𝑎 𝑜 𝑎 𝑜 𝑎 𝑜 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑛 𝑎 𝑜 𝑎 𝑛
Node from the original context
[Figure]
from the neighbor context entire training corpus as cues for prediction.
The similar contexts, deﬁned as the k neighbors of the input in the training corpus, are served as additional references for the model to predict the next token.
To integrate retrieved neighbors with the input, we build a directed heterogeneous graph on top of the input and the extracted contexts, where nodes are the tokens and edges represent the connections between them.
We deﬁne two types of nodes – the original node from the input context and the neighbor node from the extracted contexts, and two types of edges – the inter-context edge and the intra-context edge that respectively associate inter (i.e., context within the input) and intra (i.e., context in the retrieved sentences) contexts.
A graph neural network (GNN) is employed to aggregate information from both inter-context and intra-context, which is used to generate the target token.
We observe that the proposed scheme retrieves the related contexts as references, making it signiﬁcantly easier for the model to predict upcoming words in the LM task.
We further combine GNN-LM with k NN-LM (Khandelwal et al., 2019), an orthogonal technique enhancing language models, to improve the overall performance of our model.
We carry out experiments on three widely used language modeling benchmarks: WikiText-103, One Billion
Word and Enwik8.
Experimental results show that our proposed framework outperforms the strong baseline on all three benchmarks.
Speciﬁcally, applying the GNN-LM framework to a strong base
LM leads to substantial performance boost (-2.5 perplexity) on WikiText-103, and combining with k NN-LM achieves a new state-of-the-art perplexity of 14.2 – a 4.5 point improvement over the base
LM.
We perform comprehensive analyses including complexity analysis and the effects of different components to better understand the mechanics of GNN-LM.
GNN-LM 2. O VERALL PIPELINE
We present the overall pipeline of our model in Figure 1.
At each time step t , a neural language model (LM) f ( · ) ﬁrst encodes a sequence of context tokens c t = ( w 1 , w 2 , ..., w t − 1 ) to a high-dimensional representation h t = f ( c t ) ∈ R d , where d is the dimension of hidden states.
Then a transformation matrix W ∈ R V × d is used to estimate the probability of the t -th token p ( w t | c t ) =
softmax ( W h t ) , where V is the size of the vocabulary.
We augment the vanilla neural language model by allowing it to reference samples in the training set that are similar to the current decoded sequence.
Concretely, we leverage a novel self-attention augmented Graph Neural Network (GNN) on top of the vanilla LM to enable message passing between the context c and retrieved reference tokens from the training set, updating the representation h t generated by the vanilla LM.
The updated representation, which aggregates additional information from reference tokens, is then used to estimate p LM ( w t | c t ) .
2. G RAPH C ONSTRUCTION
The ﬁrst step of our proposed framework is to build a graph capturing the connections between the context tokens c t = ( w 1 , w 2 , ..., w t − 1 ) and those similar to c t in the training set.
To this end, we construct a directed heterogeneous graph, where the nodes are tokens from c t or the tokens from the neighbor contexts retrieved from the training set, and the edges represent different relationships between the nodes to be discussed below.
Formally, we deﬁne a graph as G =
( V , E , A , R , τ, φ ) , where V is a collection of nodes v and E is a collection of edges e .
We deﬁne two types of nodes
A = { a o , a n } , where a o means that the node is within the input c t .
a n means the node is in N ( c t ) , the set of extracted contexts within the neighborhood of c t .
We also deﬁne two types of edges R = { r inter , r intra } , where r inter means inter-context connection (from a n nodes to a o nodes) and r intra means intra-context connection (between two nodes of same type).
Each token within the input is a node of type a o , and edges of type r intra are constructed from node w i to w j ( i ≤ j ), which can be viewed as a graph interpretation of the transformer structure.
Both nodes and edges are associated with their respective type mapping functions τ ( v ) :
V → A and φ ( e ) : E → R .
For an input context c t , we retrieve k nearest neighbors N ( c t ) =
{ c (1) t 1 , ..., c ( k ) t k } of c t from the training set as follows: we ﬁrst use h t to query the cached representations of all tokens for training samples, where the cached representations are obtained by a pretrained LM.
The distance is measured by the cosine similarity 2 , and we retrieve the top K tokens denoted by { w ( i ) j } .
The superscript ( i ) denotes the i -th training sample and the subscript j denotes the j -th time step.
w ( i ) j thus means that the j -th time step of the i -th training sample is retrieved as one of the nearest neighbors to h t .
w ( i ) j is expanded to c ( i ) j by adding both left and right contexts, where c ( i )
j = { w ( i ) j + p } rp = − l , where l and r respectively denote the left and right window size.
The corresponding representations { h ( i ) j + p } rp = − l are used as the initialized node embeddings
Different from k NN-LM (Khandelwal et al., 2019) that uses w ( i ) j +1 , which is the token right after the retrieved token w ( i ) j , to directly augment the output probability, we explicitly take advantage of all contextual tokens near w ( i )
t
i as additional information in the form of graph nodes.
In this way, the model is able to reference similar contexts in the training set and leverage the corresponding ground-truth target tokens via the heterogeneous graph built on both the original input tokens and the context reference tokens.
For the neighbor context window size l and r , we set l = r = 1 in all experiments.
During experiments, we ﬁnd that using shallow (i.e. 3) GNN layers and adding r intra edges between adjacent tokens can alleviate overﬁtting.
Since a 3-layer GNN only aggregates information from 3-hop nodes in the graph, using larger l and r have no inﬂuence on GNN representations.
2. GNN ON THE C ONSTRUCTED G RAPH
We now use graph neural networks (GNNs) to aggregate and percolate the token information based on the graph constructed in Section 2.2.
In this work, to accommodate the modeling of r intra from node w i to w j ( i ≤ j ) within c t , where Transformer with self-attention is usually adopted, we extend the self-attention mechanism to r inter , and construct a self-attention augmented GNN.
Speciﬁcally, the l -th layer representation of node n is computed by (here we use the superscript [ l ] to represent the l -th layer):
[Equation]
Attention ( s, e, n ) estimates the importance of the source node s on target node n with relationship e ,
Feature ( s, e, n ) is the information feature that s should pass to n , and Aggregate ( · ) aggregates the neighborhood message with the attention weights.
To draw on the information in the heterogeneous graph, we use different sets of parameters for different node types τ ( · ) and different edge types φ ( · ) akin to Hu et al.
Attention Similar to the multi-head attention mechanism of Transformer (Vaswani et al., 2017), the Attention ( · , · , · ) operator in our model consists of h heads, which compute attention weights independently, followed by concatenation to get the ﬁnal output.
For simplicity, we only describe the single-head situation below.
For each edge ( s, e, n ) , the representation of target node n is mapped to a query vector Q ( n ) , and the representation of source node s is mapped to a key vector K ( s ) .
The scaled inner-production is then used to compute the attention weight between Q ( n ) and K ( s ) , which is further normalized over all edges that have the same edge type:
[Equation]
where d is the hidden dimensionality, and W qτ ( s ) ∈ R
d × d , W kτ ( n ) ∈ R
d × d , W ATT φ ( e ) ∈ R d × d , µ ∈ R |A|×|R|×|A| are learnable model parameters.
Feature Parallel to the calculation of attention weights, we propagate information from source node s to target node n .
The single-head feature is deﬁned by:
[Equation]
where W vτ ( s ) ∈ R
d × d and W FEA φ ( e )
∈ R
d × d are learnable model parameters.
Aggregate Aggregate ( · ) weight-sums the feature Message ( s, e, n ) within the vicinity using
Attention ( s, e, n ) , and the result is then linearly projected into a d -dimensional representation:
[Equation]
where ⊕ is element-wise addition and W oτ ( n
) ∈ R d × d is model parameter.
The representation of token w t − 1 from the last layer is used to compute the language model probability p LM ( w t | c t ) .
2. k NN B ASED P ROBABILITY FOR N EXT T OKEN
We further incorporate the proposed model with k NN (Khandelwal et al., 2019; 2020; Meng et al., 2021), a related but orthogonal technique, to improve the performance of our model.
It extends a vanilla LM by linearly interpolating it with a k -nearest neighbors (kNN) model.
Concretely, for each input context c t = ( w 1 , w 2 , ..., w t − 1 ) , we retrieve the k nearest neighbors N ( c t ) =
{ c (1) t 1 , ..., c ( k ) t k } , and compute the k NN based probability for the next token by:
[Equation]
with Z being the normalization factor, f ( · ) is the neural language model encoding contexts to high dimensional representations, cos( · , · ) is cosine similarity, and λ and T are hyperparameters.
E XPERIMENTS
We conduct experiments on three widely-used language modeling datasets:
WikiText-
(Merity et al., 2016), One Billion Word (Chelba et al., 2013) and Enwik8 (Mahoney, 2011).
For all experiments, we add a 3-layer self-attention augmented GNN on top of the pretrained base LM, and use the same hidden dimension and number of heads as our base LM.
We retrieve k = 1 , 024 nearest neighbors for each source token, among them the top 128 neighbors are used in graph, and all of them are used in computing the k NN-based probability p kNN ( w t | c t ) .
For the neighbor context window size l and r in Section 2.2, we set l = 1 and r = 1 .
3. T RAINING D ETAILS
KNN Retrieval
In order to reduce memory usage and time complexity, in practice we use FAISS (Johnson et al., 2019) for fast approximate k NN search.
Concretely, we quantized each dense vector to q bytes, followed with a clustering of all vectors to C clusters.
During retrieval, we only search in 32 clusters whose centroids are nearest to query vector.
For WikiText-103 and Enwik8 datasets, which contain approximately 100M tokens, we set q = 128 and c = 4 , 096 .
For One Billion Word dataset, we set q = 64 and c = 1 , 048 , 576 (2 20 ) for faster search.
Data Leakage Prevention When searching for the k nearest neighbors of c t = ( w 1 , w 2 , ..., w t − 1 ) , we need to make sure each reference neighbor token does not leak information for w t .
Speciﬁcally, we should not retrieve c t +1 =
( w 1 , w 2 , ..., w t ) as reference, otherwise the model prediction is trivial to optimize since the information of target token is already included in the graph.
Let T be the maximum sequence length and L be the number of layers.
Practically, the representation of each token is dependent on previous T and
T ×
L tokens for Transformer and Transformer-XL, respectively.
Therefore we ignore all the neighboring nodes within this interval in graph construction during training.
During inference, we do not impose this constraint.
Feature Quantization
The input node representations of the graph neural network H [0] are generated by a pretrained neural language model.
To accelerate training and inference, we wish to cache all token representations of the entire training set.
However, frequently accessing Terabytes of data is prohibitively slow.
To address this issue, we followed Meng et al.
(2021) to use product quantization (PQ) (Jegou et al., 2010; Ge et al., 2013) to compress the high-dimensional representation of each token.
In our experiments, quantizing representations from 1,024-dimension ﬂoating-point dense vectors to 128 bytes reduces the memory consumption from 2.3TB to 96GB for the One Billion Word dataset, thus making the end-to-end model training feasible.
3. M AIN R ESULTS
WikiText-103 WikiText-103 is the largest available word-level language modeling benchmark with long-term dependency.
It contains 103M training tokens from 28K articles, and has a vocabulary of around 260K. We use the base version of deep Transformer language model with adaptive embeddings (Baevski & Auli, 2018) as our base LM.
This model has 16 decoder layers.
The dimensionality of word representations is 1,024, the number of multi-attention heads is 16, and the inner dimensionality of feedforward layers is 4,096.
During training, data is partitioned into blocks of 3,072 contiguous tokens.
During evaluation, blocks are complete sentences totaling up to 3,072 tokens of which the ﬁrst 2,560 tokens serve as context to predict the last 512 tokens.
As shown in Table 1, GNN-LM reduces the base LM perplexity from 18.7 to 16.8, which demonstrates the effectiveness of the GNN-LM architecture.
The combination of GNN and k NN further boosts the performance to 14.8, a new state-of-the-art result on WikiText-103.
One Billion Word One Billion Word is a large-scale word-level language modeling dataset of short-term dependency.
It does not preserve the order of sentences, contains around 768M training tokens and has a vocabulary of around 800k.
We adopt the very large version of Transformer model in Baevski & Auli (2018) as our base LM.
Results in Table 2 show that GNN-LM helps base LM reduce 0.5 perplexity with only 27M additional parameters.
For comparison, Baevski & Auli (2018) use 560M additional parameters to reduce perplexity from 23.9 to 23.0.
Enwik
Enwik8 is a character-level language modeling benchmark that consists of 100M characters from English Wikipedia articles, and has a vocabulary of 208.
For base LM, we use TransformerXL
(Dai et al., 2019) with 12 layers, 8 heads, 512 dimensional embedding and 2,048 dimensional inner feed forward layer.
Table 3 shows that GNN- k NN-LM outperforms base LM by 0.03 Bit per
[Table] [Table] Character (BPC), achieving 1.03 BPC with only 48M parameters, comparable to 18L Transformer-XL with 88M parameters.
[Table]
A NALYSIS 4. C OMPLEXITY A NALYSIS
Space Complexity
In our model, we consider k nearest neighbors for each token c i in context; the number of nodes in the graph is k times larger than vanilla LM during training.
Accordingly, training GNN requires approximately k times larger memory than vanilla LM, since we have to maintain hidden representations of each node for backward propagation.
We propose two strategies to alleviate the space issue: (1) For all datasets, we ﬁrst train with a smaller k = 32 , then further ﬁnetune the model with a larger k = 128 ; and (2) For datasets with extremely long dependency (e.g., WikiText-103), we truncate the context to a smaller length (e.g., 128) instead of the original longer context (e.g., 3,072) used by vanilla Transformer (Baevski & Auli, 2018).
Note that we build
GNN model on top of the vanilla Transformer, and the parameters of Transformer are ﬁxed when
GNN parameters are being trained.
Hence, the GNN could exploit long dependency information learned by Transformer without having to build a large graph with long context.
Figure 2(a) shows the comparison of base LM and GNN-LM on GPU memory usage with variant k in WikiText-103.
[Figure]
Time Complexity Both GNN and Transformer consist of two basic modules: the feed forward layer and the attention layer.
Let | V | be the number of nodes and | E | be the number of edges in the graph.
Then the time complexity of the feed forward layer is O ( | V | ) and the time complexity of the attention layer is O ( | E | ) .
The GNN model increases | V | by ( l + r + 1) k times in the graph, and adds ( l + r + 1)
k | V | edges to the graph.
Note that | E | = | V | 2 in Transformer, and thus the increased time complexity is acceptable if k (cid:28) | V | holds.
Figure 2(b) shows the comparison between base LM and
GNN-LM in speed in WikiText-103.
We observe that the speed of GNN-LM is approximately 8 to 20 times slower than the base LM (Baevski & Auli, 2018) with respect to different k .
It is worth noting that the overhead of the proposed model comes from k NN retrieval, which can be done in advance and thus does not result in time overhead when running the model.
Speciﬁcally, the time overhead for retrieval comes from two processes: 1) building data indexes with token representations in the train set; 2) collecting nearest neighbors by querying the data indexes.
For
WikiText-103, building data indexes takes approximately 24 hours on a CPU machine with 64 cores.
And querying data indexes for all tokens in train set takes approximately 30 hours.
4. A BLATION S TUDY
Number of Neighbors per Token
The number of neighbors per source token (i.e., k ) signiﬁcantly inﬂuences how much information could be retrieved from the training set.
Figure 2(c) shows that test perplexity monotonically decreases when k increases from 8 to 128.
This trend implies that even larger improvements can be achieved with a larger value of k .
Neighbor Quality We evaluate the quality of k NN retrieval by examining whether the target token to predict (i.e., w t ) is the same as the token that comes right after the retrieved nearest sequence using the recall metric.
Given a sample c t = ( w 1 , w 2 , ..., w t − 1 ) and its k NN N ( c t ) = { c (1) t 1 , ..., c ( k ) t k } , the quality of k NN retrieval is deﬁned by
[Equation]
[Equation]
where w t is the target token to predict at time step t , and w ( i ) t
i is the token that comes right after the i -th retrieved neighbor.
We calculate and then divide all samples in the WikiText-103 test set by the [Table] recall value into 5 buckets, with each bucket containing around 50k tokens.
Results are reported in Table 4.
We observe that GNN- k NN-LM gains more relative improvements to base LM when the quality of k NN retrieval reaches a relatively high level.
Representation in k NN We ﬁnally study the effect of using different representations in the k NN scoring function in Section 2.4.
We experiment with two types of representations: (1) from the last layer of Transformer, which is the default setting, and (2) from the last layer of GNN.
The model performances with different choices for query and key are reported in Table 5.
Results show that using GNN representations for both query and key leads to the best performance.
It suggests that
GNN learns better representations for context similarity.
We also observe that the performance is marginally worse when both query and key are using Transformer representations.
Considering that building an additional datastore for GNN representations is computationally intensive, in practice we can directly use Transformer representations (the default setting).
[Table]
query and key.
4. E XAMPLES
Table 6 presents two examples showing the input and the corresponding extracted three neighbor contexts.
The two examples demonstrate that the extracted contexts have a strong connection in semantics to the input, and thus leveraging the neighboring information will beneﬁt model predictions.
[Table]
bold token is the gold token to predict, and the underlined are the extracted context tokens.
R ELATED W ORK
Language Modeling Traditional methods for language modeling use n -gram statistics to compute the probability of the next token given the ( n − 1) -gram context (Bahl et al., 1983; Nadas, 1984; Chen & Goodman, 1999).
With the development of neural language models (NLMs) (Mikolov et al., 2012), deep learning based methods begin to dominate the learning paradigm of language modeling.
For example, Jozefowicz et al.
(2016) built a strong language model by combining the LSTM (Schuster & Paliwal, 1997) model and the CNN structure; Melis et al.
; Merity et al.
(2017) applied a variety of regularizations to LSTMs and achieved state-of-the-art results; Baevski & Auli (2018) proposed adaptive input embeddings, which can improve performance while drastically reducing the number of model parameters.
On top of Transformer (Vaswani et al., 2017), considerable efforts have been devoted to building stronger and more efﬁcient language models (Shazeer et al., 2018; Dai et al., 2019; Beltagy et al., 2020; Press et al., 2020b;a).
BERT (Devlin et al., 2018) proposed the Masked Language Modeling (MLM) pretraining paradigm to train a deep bidirectional Transformer model;
RoBERTa (Liu et al., 2019) removed the Next Sentence Prediction (NSP) task in BERT;
XLNet (Yang et al., 2019) generalized BERT pretraining to the autoregressive manner; Span-level BERTs (Lewis et al., 2019;
Song et al., 2019; Joshi et al., 2020) introduced span-level masks rather than just relying on token-level masks. ELECTRA (Clark et al., 2020) proposed to detect token replacement as opposed to token generation, improving both the efﬁciency and effectiveness of pretraining.
(Sun et al., 2021) extends BERT to accommodate glyph information.
Graph neural networks (GNNs) capture the dependencies and relations between nodes connected with edges, which propagate features across nodes layer by layer (Scarselli et al., 2008; Kipf & Welling, 2016; Hamilton et al., 2017).
GNNs have demonstrated effectiveness in a wide variety of tasks in natural language processing such as text classiﬁcation (Yao et al., 2019; Lin et al., 2021), machine translation (Bastings et al., 2017), question answering (Song et al., 2018; De Cao et al., 2018), recommendation (Wu et al., 2019) and information extraction (Li et al., 2020a).
For example, Guo et al.
(2019) proposed Star Transformer, a Transformer backbone but replaces the fully-connected structure in self-attention with a star-like topology.
Ye et al.
(2019) adopted a ﬁne-to-coarse attention mechanism on multi-scale spans via binary partitioning (BP).
Li et al.
(2020b) proposed to learn word connections speciﬁc to the input via reinforcement learning.
Retrieval-augmented Models Retrieving contexts from another corpus as additional information improves the model’s robustness towards infrequent data points.
A typical application of retrievalaugmented models is open-domain question answering, which solicits related passages from a large open-domain database to answer a given question.
The dominant approach is to cache dense representations of the passages and retrieve the closest ones to the input during inference (Lewis et al., 2020b; Karpukhin et al., 2020;
Xiong et al., 2020; Lee et al., 2020; Li et al., 2020b).
Lewis et al.
(2020a) proposed to ﬁrst extract a set of related texts and condition on them to generate the target text. allowing for strong zero-shot performance.
Besides open-domain QA, other tasks such as language modeling (Khandelwal et al., 2019; Guu et al., 2020), machine translation (Zhang et al., 2018; Tu et al., 2018; Jitao et al., 2020), text classiﬁcation (Lin et al., 2021), and task-oriented dialog generation (Fan et al., 2020; Thulke et al., 2021) also beneﬁt from the additionally retrieved information.
For example, Khandelwal et al.
(2019) retrieved k nearest neighbors from a large-scale unannotated corpus and interpolates with the decoded sentence for language modeling.
Khandelwal et al.
Meng et al.
(2021) retrieved k NNs from the parallel translation corpus to augment the machine translation outputs.
However, these methods retrieve related texts independently .
C ONCLUSION AND F UTURE W ORK
In this work, we propose GNN-LM, a new paradigm for language modeling that extends vanilla neural language model by allowing to reference similar contexts in the entire training corpus.
High dimensional token representations are used to retrieve k nearest neighbors of the input context as reference.
We build a directed heterogeneous graph for each input context, where nodes are tokens from either the input context or the retrieved neighbor contexts, and edges represent connections between tokens.
Graph neural networks are then leveraged to aggregate information from the retrieved contexts to decode the next token.
Experimental results show that our proposed method outperforms strong baselines in standard benchmark datasets, and by combining with k NN LM, we are able to achieve state-of-the-art results on WikiText-103.
In future work, we will consider improving efﬁciency for building the graph and retrieving nearest neighbors.