GNN-LM: Language Modeling based on Global Contexts via GNN
Authors: Yuxian Meng
Shi Zong
Xiaoya Li
Xiaofei Sun
Tianwei Zhang
Fei Wu
Jiwei Li
Venue: ICLR 2022 Spotlight (2022)

Abstract: Inspired by the notion that "it 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 3.9 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. The code can be found at https://github.com/ShannonAI/GNN-LM.

Published as a conference paper at ICLR 2022
GNN-LM: LANGUAGE MODELING BASED ON GLOBAL CONTEXTS VIA GNN
Yuxian Meng', Shi Zong?, Xiaoya Li', Xiaofei Sun Tianwei Zhang', Fei Jiwei Li "Shannon.AI, "Nanjing University," Nanyang Technological University. "Zhejiang University [yuxian meng. xiaoya_ li, xiaofei sun, jiwei_ li] @ shannonai.com, szong nju.edu.cn zianwei.zhang @ eeeetteedd sdg. wufei @zj..duu..n
ABSTRACT
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 s generalization ability. We conduct comprehensive experiments to validate the effectiveness of the GNN-LM: GNN-LM achieves a new state-of-the-ard perplexity of 14.8 WikiText-103 (a 3.9 point improvement over its counterpart of the vanilla LM model), and shows substantial improvement on One Billion Word and Enwiki8 datasets against strong basclines. In-depth ablation studies are performed to understand the mechanics of GNN-LM
INTRODUCTION
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). 7/. It aims at predicting the upcoming token given the sequence of previous context consisting of a sequence of tokens. A common practice to train 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- process of memorization, or doing a close-book examination, if we compare the training data to 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 specific 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.
procedure can be viewed as a
There are two limitations to this close-book examination strategy: (1) the memorization- language models are usually hard to memorize the knowledge of hard examples (c.g., long-tail cases in the training set); (2) memory required to memorize the whole training data is usually intensive. The difficulty 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 strategy different from memorization or close-book examination copy, or in other words, open-book examination. For example, given a prefix "J. Rowling is best known for writing" and we want predict the upcoming token, 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 a close-book exam", , in this work, we introduce a new language modeling scheme-
The code can be found at htt.://////////ttttooooooo/////a
/GNN-1

---

Published as a conference paper at ICLR 2022
Graph Construction (t r- 1)
Neighbor Contexts
This movie is great.
On
Gg
Intrr-onntect idge
())
; Those movies are bad
This movie is what like.
movie
great
ag Node from original contert
Intra cort--ontectt edge
Training Datastooe
ht
query
Input
The
movie
is
0g Node from the neighbor conlest
Input Context
Gg
do
Base LM
C2
movies
are
bad
movies
is
what
The
movie
is
On Bn
Gg
An An
On
Figure 1: An overview of the proposed GNN-LM model pipeline. Left: Given input context Ca = (Wi, ... w(-1) (here the context "The movie is"), a base LM model encodes it into a high-dimensional representation h,, which is then used to query the training datastore retrieve the nearest contexts along with the visited tokens (marked in red). Right: The tokens in the input context and the retrieved tokens comprise a graph and are viewed as two types of nodes: nodes from the original text and nodes from the neighbor text. Intra-context edges link tokens within the same input, and inter-context edges link tokens from the retrieved contexts to the original context. After modeling the graph as a whole with GNNs, we use the updated representation of w (token "is" in this example) to compute the likelihood of the next token.
GNN-LM, which provides an LM model with the ability to reference similar contexts from the entire training corpus as cues for prediction. The similar contexts, defined as the 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 define 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., between retrieved contexts and input context) and intra (i.e., context within the input) contexts. 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 significantly easier for the model to predict upcoming words in the LM task.
We further combine GNN-LM with ANN-LM (Khandelwal 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. Specifically, applying the GNN-LM framework to a strong base LM leads to substantial performance boost (-1.9 perplexity) on WikiText-103 103, and combining with ANN-LM achieves a new state-of-the-art perplexity 14.8 - a 3.9 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.
2
GNN-LM
2.1
OVERALL PIPELINE
We present the overall pipeline of our model in Figure ! At each time step t, neural language model (LM) f()) first encodes a sequence of context tokens C, = (W,, W.,.... W.-1) to high- high-dimensional representation h, = (c:) Rdd where is the dimension of hidden states. Then a transformation matrix W E R. Vxdd is used to estimate the probability of the t-th token p( (w, ic.) softmax (Wh Wh.). where V is the size of the vocabulary. We augment the vanilla neural language model by allowing 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 e and retrieved reference tokens from the training set, updating the representation h, generated by the vamilla LM. The updated representation, which aggregates additional information from reference tokens, then used to estimate PLM (we))))).
2

---

Published as a conference paper at ICLR 2022
2.2 GRAPH CONSTRUCTION
The first step of our proposed framework is to build a graph capturing the connections between the context tokens C₁ (Wi, W2,,..., W;-1) and those similar to c, in the training set. To this end, construct a directed heterogeneous graph, where the nodes are tokens from e 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 define graph as (,, Ø), where V is a collection of nodes v and & is collection of edges e. We define two types of nodes A = {ao. )n where d, means that the node is within the input C₁. an means the node is in ((ee), the set of extracted contexts within the neighborhood of C.. We also define two types of edges R {Tinter: Ti where Tinter means inter-context connection (from an nodes to a, nodes) and rinta means intra-context connection (between two nodes of same type). Each token within the input is node of type ag and edges of type Tintra are constructed from node w, to wy (i S 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 r (w)) V + A and Ø(e)= : 3 -> R.
(k) of c from the For input context Ct. we retrieve k nearest neighbors M(c.) [c ..... Ct.
training set as follows: we first use h, to query the cached representations of all tokens for training samples, where the cached representations
obtained by pretrained L.M. The distance is measured
by the cosine similarity P
3.). denotes the i-th training sample and the subscript, denotes the j-th time step. w/ thus means that the j-th time step of the i-th training sample is retrieved as one of the nearest neighbors to he. w/" is expanded to cf by adding both left and right contexts, where cf = {w }ppp)pp--.. where I and r respectively denote the left and right window size. The corresponding representations ( used as the initialized node embeddings.
and we retrieve the top K tokens denoted by {w)) J. The superscript (i)
Different from &NN-LM (Khandelwal et al. 2019) that uses w0 j++1 which is the token right after the retrieved token w to directly augment the output probability, we explicitly take advantage of all contextual tokens near w.(.) 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 and r. we set = = 1 in all experiments. During experiments, find that using shallow (i.e., 3) GNN layers and adding Tim edges between adjacent tokens can alleviate overfitting, Since a 3-layer GNN only aggregates information from 3-hop nodes in the graph, using larger / and have influence on GNN representations.
2.3 GNN ON THE CONSTRUCTED GRAPH
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, accommodate the modeling of rinta from node wi to wy (i < j)) within C, where Transformer with self-attention is usually adopted, we extend the self-attention mechanism to Timeer. and construct a self-attention augmented GNN.
the superscript []] to Specifically, the 1-th layer representation of node is computed by (here represent the I-th layer):
h
AggrrggateAAttenionnnsss Attentionns, e, Feature e, n)) + h! 1-11 (sE
(1)
Attention(s,, e, n) estimates the importance of the source node s on target node with relationship e, Feature(s, n)) is the information feature that should pass to n, and Aggregate(-) aggregates the neighborhood message with the attention weights. To draw on the information in the heterogeneous graph, use different sets of parameters for different node types 7(-) and different edge types Φ(-) akin to Hu et al. (2020).
II practice, we FAISS Johnson al. 2019 for fast approximate ANN search.
3

---

Published as a conference paper at ICLR 2022
Attention Similar the multi-head attention mechanism of Transformer (Vaswani al. (017).
the Attentionn'...> operator in our model consists of h heads, which compute attention weights independently, followed by concatenation get the final output. For simplicity, only describe the single-head situation below. For each edge (s, e, n), the representation of target node n is mapped to query vector (n), and the representation of source node mapped to key vector ((s). The scaled inner-production is then used to compute the attention weight between Q(n) and K (s), which further normalized over all edges that have the same edge type:
(s) W h T(s)
1
Q(n) W h!- r(n)
1
K (l) Mirls) vd
Attention(s, (s,
7 exp
eex
(2)
Z Σ
Autention( ('',,,
('E((e),
where the hidden dimensionality, and W (()) E Rdxd W r(n) E Rddd μ E [|4]x|R|x||] are learnable model parameters.
WATT d(e)
E
R
Feature Parallel the calculation of attention weights, propagate information from source node 8 to target node n. The single-head feature is defined by:
Feature(s, n) W h 1- p(e))
(3)
where We r(s) Rdd and F-FEA Rdxd are learnable model parameters. ((e)
Aggregate Aggregate(.) weight-sums the feature Message(s, e, n) within the vicinity using Attention s, n), and the result is then linearly projected into d-dimensional representation:
Aggregate()) W 7(n)
+ YeEN(()
(Attention(s, es e,
(4)
where @ is element- wise addition and Wi(m)
Rdxd is model parameter. The representation of token W from the last layer is used to compute the language model probability (L,(w)))))
r(n)
2.4 ANN BASED PROBABILITY FOR NEXT TOKEN
We further incorporate the proposed model with ANN (Khandelwal et al. 2019 2020 Meng et al.
2021), related but orthogonal technique, to improve the performance of our model. It extends vanilla LM by linearly interpolating it with k-nearest neighbors (kNN) model. Concretely, for each
input context (wi, W2, ..., w;-1), we retrieve the & nearest neighbors M(ec)) {ct, Ci
(1) (k)
and compute the KNN based probability for the next token by:
p(we(ce) \PaNy(w,le.) + (1 A)PLM (w)l),
1 (INN(WelC.) = Σ w exp (cos(f(a)..((()))))
(5)
with 7 being the normalization factor, f(.) is the neural language model encoding contexts high dimensional representations, cos(:,) :) is cosine similarity, and À and are hyperparameters
3 EXPERIMENTS
We conduct experiments on three widely- widely-used used language modeling datasets: WikiText- 103 (Merity al. 2016), One Billion Word (Chelba et al.) 2013) and Enwik8 (Mahoney 2011). For all experi- ments, we add a 3-layer self-attention augmented GNN on top of the pretrained base LM, and use
'The original version of ANN-LM (Khandelwal et al. 2019 uses negative Le distance as vector similarity, and does not have hyperparameter T. We followed Khandelwal et al. (2020) add hyperparameter T and followed MenooweeMenggeeee et al. 2021
cosine similarity.
4

---

Published as a conference paper at ICLR 2022
the same hidden dimension and number of heads neighbors for each source token, among them the top 128 neighbors are used in graph, and all of them used in computing the kNN-based probability PENx PeNs(w, (c)lc,) For the neighbor context window size and r in Section 2.2
base LM. We retrieve k = 1, 024 nearest
we set = and r
3.1 TRAINING DETAILS
KNN Retrieval In order reduce memory usage and time complexity, in practice we use FAISS (Johnson et al. 2019) for fast approximate kNN search. Concretely, we quantized each dense vector to q bytes, followed with clustering of all vectors to 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 e = 4, 096. For One Billion Word dataset, we set q 64 and c 1, 048, 576 (220) for faster search.
Data Leakage Prevention When searching for the l nearest neighbors c; (Wi, W2, W/-1 1), we need to make sure each reference neighbor token does not leak information for wi Specifically, we should not retrieve Ct+1 (W1, W2,,.., Wx) 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 x L tokens for Transformer and Transformer-Xe respectively. Therefore we ignore all the neighboring nodes within this interval 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 gener- ated by pretrained neural language model. To accelerate training and inference, we wish to cache all token representations the entire training set. However, frequently accessing Terabytes of data is prohibitively slow. To address this issue, we followed Meng et al. (202T) to product quantization (PQ) (Jegou et al., 2010 Ge et al. 2013 3 to compress the high- high-dimensional representation of each token. In our experiments, quantizing representations from 1, 1,024ddimension floating-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.2 MAIN RESULTS
Wiki'eex- 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 decoder layers. The dimensionality of word representations 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 complete sentences totaling up 3,072 tokens of which the first 2,560 tokens serve as context to predict the last 512 tokens. As shown in Table 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 ANN further boosts the performance to 14.8, new state-of-the-arr result on WikiText-103.
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) our base LM. Results in Table that GNN-ANN-LM helps base LM reduce 0.5 perplexity with only 27M additional parameters. For comparison, Baevski & Auli (0018) use 560M additional parameters to reduce perplexity from 23.9 to 23.0.
Enwik8 Enwik8 is a character-level language modeling benchmark that consists 100M characters from English Wikipedia articles, and has vocabulary of 208. For base LM, use Transformer- XL (Dai et al., 2019) with 12 layers, 8 heads, 512 dimensional embedding and 2,048 dimensional inner feed forward layer. Table shows that GNN-ENNNLLM outperforms base LM by 0.03 Bit per
5

---

Published as a conference paper at ICLR 2022
Model
Param
Test ppl (-)
Hebbian + Cache (Rae el al.
2018
151M
29.9
Transformer- XL (Dai
2019
257M
18.3
Transformer-X XL Dynamic Eval (Krasse
2019
257M
16.4
Compressive Transformer (Rac et al. 2019
17.1
KNN-LM + Cache (Khandēlwal ct al. 2019
257M
15.8
Sandwich Transformer (Press et al.
2020a
247M
18.0
Shortformer (Press et al. 20220
247M
18.2
PegaTTransoomee--I.. Bai
2021
257M
17.1
Routing Transformer
[ooy
2021
15.8
base LM Baevski & Auli 2018
247M
18.7
+GNN
274M
16.8
+GNN+&NN
274M
14.8
Table Test perplexity on WikiText- 103 dataset.
Model
# Param
Test ppl (L)
LSTM+CNN (Jozefowicz et al.
2016
1.04B
30.0
High-Buut Bigh-Budget MoE (Shazeer
2016
SB
28.0
DynamicConv Wu al.
0.34B
26.7
Mesh-Tensorflow (Shazecr et al. 2018
4.9B
24.0
Evolved Transformer (Saazeer el al. 2018
28.6
Transformer-a XL (Dai èt al.
2019)
0.8B
21.8
Adaptive inputs (base) Bacvski & Auli
2018
0.36B
25.2
Adaptive inputs (large) [Baevski Auli
2018
0.46B
23.9
base LM (Bevsski & Auli 2018
1,03B
23.0
+ANN
1,02B
22.8
+GNN
1.05B
22.7
+GNN+ËNN
1.05B
22.5
Table 2: Test perplexity on One Billion Word dataset.
Character (BPC), achieving 1.03 BPC with only 48M parameters, comparable to 18I. Transformer-Xi with 88M parameters.
Model
# Param
BPC ({)
Transformer (Al-Rfou et al.
2019
235M
1.06
18L Transformer-Xi XL Dai ct al.
2019
88M
1.03
24L Transformer-XL. DDi cl al.
2019
277M
0.99
24L
Transformer-XI. Dynamic Evall
Krause et al. 2019
277M
0.94
Longformer (Beltagy al. 2020
102M
0.99
Adaptive Transtormer (Sukhbaatar et
2019
209M
0.98
Compressive Transformer (Rae et al. 2019
277M
0.97
Sandwich Transformer (Press el al.
020a
209M
0.97
Transformer-XX Dai
2019
41M
1.06
+NNN
41M
1.04
+GNN
48M
1.04
+GNN+ANN
48M
1.03
Table 3: Bit per Character the Enwik8 dataset.
ANALYSIS
COMPLEXITY ANALYSIS
Space Complexity In model, consider & nearest neighbors for each token e; in context; the number of nodes in the graph is times larger than vanilla LM during training. Accordingly, training GNN requires approximately & times larger memory than vanilla LM, since have to

---

Published as a conference paper at ICLR 2022
maintain hidden representations of each node for backward propagation. We propose two strategies to alleviate the space issue: (1) For all datasets, we first train with a smaller = 32, then further finetune the model with larger & 128; and (2) For datasets with extremely long dependency (e.g., Wiki 10),, 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 build GNN model on top of the vanilla Transformer, and the parameters of Transformer are fixed when GNN parameters are being trained. Hence, the GNN could exploit long dependency information learned by Transformer without having to build large graph with long context. Figure 2[a) shows the comparison of base LM and GNN-LM on GPU memory usage with variant in WikiText- 103
(a)
(b)
(c)
17500
400
GNN-LLM
15000
CPU memory (GB)
1000
Base LM
wonE pet sacora
12500
18-
10000
GNN-LM Base LM
17
2.00
7500-
5000
1.00
2500
SANLM
15
1.6
32
64
128
16
32
128
32
64
128
k
Figure 2: Comparisons between base LM and GNN-LM WikiText- 103 with respect to different k. (a) GPU memory usage. (b) Speed (word per second). (c) Test perplexity.
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 ((V)) and the time complexity of the attention layer is O(E)). ). The GNN model increases [ by (1+r+ 1)kk times in the graph, and adds r1+r+1 1) edges the graph. Note that EE| V in Transformer, and thus the increased time complexity is acceptable if l < [V] holds. Figure 2 b) shows the comparison between base LM and GNN-LM in speed in WikiTextt 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 ANN retrieval, which can be done in advance and thus does not result in time overhead when running the model. Specifically, 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 Wiki 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.2 ABLATION STUDY
Number of Neighbors per Token The number of neighbors per source token (i.e., k) significantly influences how much information could be retrieved from the training set. Figure shows that test perplexity monotonically decreases when A increases from 8 to 128. This trend implies that even larger improvements can be achieved with larger value of k.
Neighbor Quality We evaluate the quality of kNN retrieval by examining whether the target token to predict (i.e., we) is the same as the token that comes right after the retrieved nearest sequence using the recall metric. Given a sample c = (W1, W2,... ..., w and its ANN (ca)) = ((c/)) ,..., Ctty the quality of ANN retrieval is defined by
(8)
R(C,)) = Σ 1
Wi:
(
(6)
(4)
where We is the target token to predict at time step t, and wy 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
WWe note hase LM uses a context length of 3,072, while the context length of GNN-LM is 128. We scale up the value of GNN-LM 24 times for fair comparison.
7

---

Published as a conference paper at ICLR 2022
KNN recall range
[0. 4)
27)
[27, 137)
[137,463
[463, 1024]
base LM
-7.14
-3.84
-2.21
-1.19
-0.30
+GNN+KNN
-7.15
-3.46
-1.71
-0.80
-0.21
absolute improvement
-0.01
0.38
0.50
0.39
0,09
relative improvement
-0.0%
10%
23%
33%
32%
Table 4: Comparison of base LM and GNN-LM in different kNN recall buckets. We report average log probabilities within each bucket, and compute the absolute and relative improvement.
recall value into 5 buckets, with each bucket containing around 50k tokens. Results are reported in Table 4 We observe that GNN-ANN-LM gains relative improvements to base LM when the quality ANN retrieval reaches relatively high level.
Representation in ANN We finally study the effect of using different representations in the ANN scoring function in Section 2.4. We experiment with two types of representations: (1) from the last layer of Transformer, which 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. 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 additional datastore for GNN representations is computationally intensive, in practice can directly use Transformer representations (the default setting).
Query Repres.
Key Repres.
Test ppl (4)
Transformer
Transformer
14.82
Transformer
GNN
15.16
GNN
Transformer
14.97
GNN
GNN
14.76
Table 5: Test perplexity WikiText- 103 with different representations as query and key.
4.3 EXAMPLES
Table 6 presents two examples showing the input and the corresponding extracted three neighbor contexts. The two examples demonstrate that the extracted contexts have strong connection in semantics to the input, and thus leveraging the neighboring information will benefit model predictions.
Input: In 2000 Boulter had guest @-@ starring
Extracted 1: In 2009, Beghe had guest @-@ starring role on the television show Califoriiattin.. Extracted 2: had previously worked on Hack, for guest @-@ starring episode are on the show.
Extracted 3: and because of Patrick Stewart 's hilarious guest @-@ starring role Number One
Input:
Tourism vital industry in Manila and Extracted 1: large audience in Mogadishu, and was widely sold prior to the civil war Extracted 2: industry is well established, with Mumbai Port being one of the oldest and most
Extracted 3: transportation has become large business in Newark. accounting for than 17
Table Two examples showing the input context and the corresponding extracted three neighbors. The bold token is the gold token predict, and the underlined are the extracted context tokens.
RELATED WORK
Language Modeling Traditional methods for language modeling 2-gram statistics 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
8

---

Published as a conference paper at ICLR 2022
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 [997) model and the CNN structure; MMlis al. 2017); Merity et al. (2017) applied
variety of regularizations to LSTMs and achieved state-of-the-an results; Baevski & Auli (2018)
proposed adaptive imput embeddings, which can improve performance while drastically reducing the
number of model parameters, On top of Transformer (Vaswani et al., (017). considerable efforts have been devoted to building stronger and more efficient language models (Shazeer et al. 2018 Dai et al.
(019)Beltagge et al. 2020; Press et al. 2020b a). BERT (Devlin et 2018)) proposed the Masked
Language Modeling (MLM) pretraining paradigm 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.) [019]Joshh 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 effecincc and effectiveness of pretraining. (unn
et al. 2021 extends BERT to accommodate glyph information.
Graph Neural Networks Graph neural networks (GNNs) capture the dependencies and relations
between nodes connected with edges, which propagate features across nodes layer by layer (Scarselli
jt al. 2008 Kipf & Welling, 2016 Haamiloon et al. 2017). GNNs have demonstrated effectiveness in
wide variety of tasks in natural language processing such as text classification (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 cully-connected structure in self-attention with a star-like topology. Ye et al. (2019) adopted
tine-toccarse attention mechanism on multi-scale spans via binary partitioning (BP). Li et al. (2020b)
proposed to learn word connections specific to the input via reinforcement learning.
Retrieval-augmented Models Retrieving contexts from another corpus additional information improves the model's robustness towards infrequent data points. typical application of retrieval- augmented models is open-domain question answering, which solicits related passages from 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 ct al. 20206. Karpukhin et al. 20200 Xiong et al. 2020 Lee et al. 2020: Li et al., 2020b). Lewis et al. (2020a) proposed to first extract 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 al. 2020), machine translation (Zhang et al. 2018; Tu et al.) [018:JJiaa et al. 2020). text classification (Lin et al., [02I],, and task-oriented dialog generation (Fan et al. 22000) Thulke et al. (202)) also benefit from the additionally retrieved information. For example, Khandelwal et al. (2019) retrieved & nearest neighbors from a large-scale unannotated corpus and interpolates with the decoded sentence for language modeling. Khandelwal et al.] (2020); Meng et al. (2021) retrieved kNNs from the parallel translation corpus to augment the machine translation outputs. However, these methods retrieve related texts independently.
CONCLUSION AND FUTURE WORK
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 i nearest neighbors of the input context as reference. We build 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 ANN LM, we are able to achieve state-of-ther the-art results on WikiText- 103. In future work, we will consider improving efficiency for building the graph and retrieving nearest neighbors.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
This work is supported by the Science and Technology Innovation 2030 "New Generation Artificial Intelligence" Major Project (No. 122100110201) and the Key & D Projects of the Ministry of Science and Technology (2020YFC0832500), We would like to thank anonymous reviewers for their comments and suggestions.
REFERENCES
Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-lever language modeling with deeper self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3169-3166, 2019.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference Learning Representations, 2018.
Lalit R Bahl, Frederick Jelinek, and Robert L. Mercer. A maximum likelihood approach continuous speech recognition. IEEE transactions on pattern analysis and machine intelligence, (2):179-190, 1983.
He Bai, Peng Shi, Jimmy Lin, Yuqing Xie, Luchen Tan, Kun Xiong, Wen Gao, and Ming Li. Segatron: Segment-aware transformer for language modeling and understanding. 2021.
Jasmijn Bastings, Ivan Titov, Wilker Aziz, Diego Marcheggiani, and Khalil Sima' an. Graph con- volutional encoders for syntax-aware neural machine translation. Proceedings of the 2017 Conference Empirical Methods in Natural Language Processing, pp. 1957-1967, Copenhagen, Denmark, September 2017. Association for Computational Linguistics.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05 /50, 2020.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005 2013.
Stanley F Chen and Joshua Goodman. An empirical study of smoothing techniques for language modeling. Computer Speech & Language, 13(4):359-394, 1999.
Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning. Electra: Pre-training text encoders as discriminators rather than generators. International Conference on Learning Representations, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978-2222882 2988, 2019.
Nicola De Cao, Wilker Aziz, and Ivan Titov. Question answering by reasoning across documents with graph convolutional networks. arXiv preprint 1XXx:1808.099200 2018.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1300044805 2018.
Angela Fan, Claire Gardent, Chloe Braud, and Antoine Bordes. Augmenting transformers with knn-based composite memory for dialogue. arXiv preprint arXiv:2004.12744 2020.
Tiezheng Ge, Kaiming He, Qifa Ke, and Jian Sun. Optimized product quantization. IEEE transactions on pattern analysis and machine intelligence, 36(4):744 2013.
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang, Star- transformer. arXiv preprint arXiv: 09113, 2019.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming Wei Chang. Realm: Retrieval- augmented language model pre-training. arXiv preprint arXiv:2002.08990 2020.
10

---

Published as a conference paper at ICLR 2022
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in neural information processing systems, pp. 1024-1034, 2017.
Ziniu Hu, Yuxiao Dong, Kuansan Wang, and Yizhou Sun. Heterogeneous graph transformer. In Proceedings of The Web Conference 2020, pp. 2704-2710, 2020.
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-228, 2010.
XU Jitao, Josep M Crego, and Jean Senellart. Boosting neural machine translation with similar translations. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 1580-1590, 2020.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 2019.
Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S Weld, Luke Zettlemoyer, and Omer Levy. Spanbert: Improving pre-training by representing and predicting spans. Transactions of the Association for Computational Linguistics, 8:64-77, 2020.
Rafal Jozefowiez, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv:1602.0244 2016.
Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04900 2020.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00111 2019.
Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Nearest neighbor machine translation. arXiv preprint arXiv:2010.007/00, 2020.
Thomas Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv: wxXvv:6600.099000 2016.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of trans- former language models. arXiv preprint 1XX4:0994008778 2019.
Jinhyuk Lee, Mujeen Sung, Jaewoo Kang, and Danqi Chen. Learning dense representations of phrases scale. arXiv preprint 4rX1::0012.12624, 2020.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising sequence-to-sequent pre-training for natural language generation, translation, and comprehension. arXiv preprint 1rX1v:1910.1346/ 2019.
Mike Lewis, Marjan Ghazvininejad, Gargi Ghosh, Armen Aghajanyan, Sida Wang, and Luke Zettlemoyer. Pre-training via paraphrasing. arXiv preprint arXiv:2006. /5020, 2020a.
Patrick Lewis, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, al. Retrieval-augmented genera- tion for Knowledge-intensiee nlp tasks. arXiv preprini arXiv:2005. 11401, 2020b.
Bo Li, Wei Ye, Zhonghao Sheng, Rui Xie, Xiangyu Xi, and Shikun Zhang. Graph enhanced dual attention network for document-level relation extraction. In Proceedings the 28th International Conference on Computational Linguistics, pp. 1551-1560, 2020a.
Xiaoya Li, Yuxian Meng, Mingxin Zhou, Qinghong Han, Fei Wu, and Jiwei Li. Sac: Accelerating and structuring self-attention via sparse adaptive connection. arXiv preprint arXiv: 09833, 2020b.
11

---

Published as a conference paper at ICLR 2022
Yuxiao Lin, Yuxian Meng, Xiaofei Sun, Qinghong Han, Kun Kuang, Jiwei Li, and Fei Wu. Bertgen: Transductive text classification by combining gen and bert. arXiv preprint arXiv:2105.05727, 2021.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeitlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprini arXiv: 1907.1 1 1692, 2019.
Matt Mahoney, Large text compression benchmark, 2011.
Gábor Melis, Chris Dyer, and Phil Blunsom. On the state of the art of evaluation in neural language models. arXiv preprint arXiv:1007.05589, 2017.
Yuxian Meng, Xiaoya Li, Xiayu Zheng, Fei Wu, Xiaofei Sun, Tianwei Zhang, and Jiwei Li. Fast nearest neighbor machine translation. arXiv preprint 1tfb::205..15288, 2021.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: (00.:v660000844. 2016.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing Istm language models. arXiv preprint arXiv: /70::070..220 0888002/82, /82, 2017.
Tomáš Mikolov al. Statistical language models based on neural networks. Presentation at Google, Mountain View, 2nd April, 80:26, 2012.
Arthur Nadas. Estimation of probabilities in the language model of the ibm speech recognition system. IEEE Transactions on Acoustics, Speech, and Signal Processing, 32(4):859-811, 1984.
Ofir Press, Noah A Smith, and Omer Levy. Improving transformer models by reordering their sublayers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2996-3005, 2020a.
Ofir Press, Noah A Smith, and Mike Lewis. Shortfommer: Better language modeling using shorter inputs. arXiv preprint arXiv:2012. /5832, 2020b.
Jack Rae, Chris Dyer, Peter Dayan, and Timothy Lillicrap. Fast parametric learning with activation memorization. In International Conference on Machine Learning, pp. 4228-4237. PMLR, 2018.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations, 2019.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. JEEE Transactions on Neural Networks, 20(1):61-80, 2008
Mike Schuster and Kuldip K Paliwal. Bidirectional recurrent neural networks. IEEE transactions on Signal Processing, 45(1 1):2673-2681 1997.
Claude Elwood Shannon. A mathematical theory of communication. ACM SIGMOBILE mobile computing and communications review, 5(1):3-55, 2001.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outragcously large neural networks: The sparsely-gated mixture-of-experts layer. 2016.
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al. Mesh-tensorfloww: Deep learning for supercomputers. In NewrIPS, 2018.
Kaitao Song. Xu Tan, Tao Qin, Jianfeng Lu, and Tie- Yan Liu. Mass: Masked sequence sequence pre-trainnt te-trainnng for language generation. In International Conference on Machine Learning, pp. 5926-5936, 2019.
12

---

Published as a conference paper at ICLR 2022
Linfeng Song, Zhiguo Wang, Mo Yu, Yue Zhang, Radu Florian, and Daniel Gildea. Exploring graph-structured passage representation for multi-hop reading comprehension with graph neural networks. arXiv preprint arXivv1809.02000 2018.
Sainbayar Sukhbaatar, Édouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics,, pp. 331-335, 2019.
Zijun Sun, Xiaoya Li, Xiaofei Sun, Yuxian Meng, Xiang Ao, Qing He, Fei Wu, and Jiwei Li. Chinesebert: Chinese pretraining enhanced by glyph and pinyin information. arXiv preprint 1rX1::2106.6608,, 2021.
David Thulke, Nico Daheim, Christian Dugast, and Hermann Ney. Efficient retrieval augmented generation from unstructured knowledge for task-oriented dialog. arXiv preprint vcXiv:2102.046433 2021.
Zhaopeng Tu, Yang Liu, Shuming Shi, and Tong Zhang. Learning to remember translation history with continuous cache. Transactions of the Association for Computational Linguistics, 6:407-420, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Felix Wu, Angela Fan, Alexei Bacvski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. In International Conference Learning Representations, 2018.
Shu Wu, Yuyuan Tang, Yanqiao Zhu, Liang Wang, Xing Xie, and Tieniu Tan. bassion-based recommendation with graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 346-353 2019.
Ziang Xie, Sida I Wang, Jiwei Li, Daniel Lévy, Aiming Nie, Dan Jurafsky, and Andrew Y Ng. Data noising as smoothing in neural network language models. arXiv preprint 1003v:1703.0220 2017.
Lee Xiong. Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. arXiv preprint arXiv:2007.000 000000 2020.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding, Advances in neural information processing systems, 32, 2019.
Liang Yao, Chengsheng Mao, and Yuan Luo. Graph convolutional networks for text classification. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 7370-77707 7377, 2019.
Zihao Ye, Qipeng Guo, Quan Gan, Xipeng Qiu, and Zheng Zhang. Bp-transformer: Modelling long-range context via binary partitioning. arXiv preprint arXiv: 191.0000, 2019.
Jingyi Zhang, Masao Utiyama, Eicchro Sumita, Graham Neubig, and Satoshi Nakamura. Guiding neural machine translation with retrieved translation pieces. arXiv preprint 1804.v:804.025555 2018.
13


Reviews:
[{"review_id": "_Mq1Bokw4_f", "reviewer": "Reviewer_eytL", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I would recommend  accepting the paper based on the novel idea to build heterogeneous graph to do LM and the impressive model performance on PPL.  This work further extends KNN-LM to utilize not next tokens but all neighborhood information to get global context. This good combination of GNN and LM can be valuable to the community. The improvement on PPL also shows the importance to use global knowledge.  \n\nThe heterogeneous GNN is standard. I would like to know if the authors have thought about designing specific graph structure or avoid inter-context edge when building the graph or considering \"is\" and \"are\" are the same node when building the graph in Fig 1. \n\nI would consider KNN-LM is a special version of GNN-LM. Then why adding KNN can greatly further improve the model performance?", "summary_of_the_paper": "This work build a novel GNN-LM to do language modeling by using global context information. The proposed model is novel and quite different from previous LM structures. This work in my view draws the connection between traditional n-gram language model and neural language model. The overall performance is quite impressive in all standard LM datasets. Extensive ablation study is conducted to understand the model. ", "summary_of_the_review": "Overall, base on the novel idea of creating global context graph, GNN-LM and show the significant improvements on all LM datasets, I would like to recommend accept this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "P5baAzR3F_", "reviewer": "Reviewer_BmDz", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The work is among the studies of enhancing language generation by context retrieval, and the new idea is modeling the retrieved neighbor contexts through a graph neural network.  My concerns lie in the following aspects:\n\n1. Comparison: it seems that the proposed model is compared with a retrieval-augmented baseline  only on WikiText-103. Why the comparison is not conducted on the other two benchmarks? Is it because there are no available results on the two other dataset? In this case, can you implement the model and make the comparison? Moreover, there have been many retrieval-augmented models, though some of them are not applied to LM tasks in the initial papers, then can you make adaptations and compare with them as well? Since the major contribution of the work is the GNN module, it would be important to demonstrate how useful the component is. \n\n2. Clarity: I feel confused at several places during the review process: (1) Eq.(5), the text below explains that f(\\cdot) is a neural language model, but afterwards in Table 5, it seems that GNN is a better choice. What is the exact setting for this equation? (2) The first paragraph of Section 3, the authors say \"we retrieve k=1024 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 kNN-baed probability\", then can I understand as the graph only involves 128 tokens? If it is, given GNN is better than Transformer in Table 5, how to calculate f(\\cdot) in Eq.(5) for other tokens? (3) how do you train the transformer and the GNN？There are some comments above Figure 2, but do you train the two components iteratively? If yes, how do you set up the training procedure? (e.g., how many iterations do you do? how to warm up the parameters of the GNN?) (4) Eq. (6) is really confusing. You only compare w_t with the i-th token of w_{t_i}? Why don't you use soft similarity (e.g., Cosine)? ", "summary_of_the_paper": "The paper presents a GNN based language model where neighbor contexts are retrieved, encoded via a graph neural network, and used to enhance generation. Evaluation on three benchmarks indicates that the proposed approach can outperform a bunch of baseline models. \n\nContributions:\n1. a new retrieval-augmented language model implemented via GNN techniques.\n2. improvements over state-of-the-art models on three benchmarks. ", "summary_of_the_review": "1. Relatively incremental technical contribution to the community. \n2. Relatively weak comparison with baselines.\n3. Many vague points that impede us from understanding (reproducing) the work. \n\n——————————————————————————————————————————————————————————————\n\nThe authors' response answered most of my questions. though I still feel that the technical contribution is big enough. I slightly raise my score accordingly. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9gde-Q3a1xf", "reviewer": "Reviewer_54po", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n- The authors propose an interesting hypothesis that referring to the training data could be helpful for language modeling, and they showed that the method is able to make considerable improvements over the vanilla LM.\n- The method achieves the new SOTA on Wiki103, which is impressive\n\nWeaknesses:\n- The paper lacks a discussion part about the actual overhead for retrieval and the time overhead for running the model seems to be significant (8-20X slower). \n- Some hyperparameters seem to be chosen quite arbitrarily like l and r, maybe the authors could provide more insights as in why they chose such a small number. Does it affect the efficiency or the performance much?", "summary_of_the_paper": "This paper proposes a method to model original texts and similar texts in a graph structure for language modeling with graph neural networks. In the graph, the nodes are texts or similar contexts and the edges are connections between the nodes. The new model achieves the new state-of-the-art on WikiText-103 and shows substantial improvements over other language modeling datasets such as One Billion Word and Enwiki8 datasets .", "summary_of_the_review": "The paper proposes a retrieval augmentation to the language models and use GNNs on top of a graph structure of the original input and its neighbors and shows that it consistently improves over the vanilla lm model. The empirical results are good and the authors provide examples showing that the retrieved examples indeed help prediction in language models. The only concern is that through modeling additional neighboring contexts, the method introduces significant overhead in running time. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new type of language model, the GNN-LM, which uses a graph neural network to allow a language model to reference similar contexts in the training corpus in addition to the input context.  The empirical results are good, and the model sets a new SOTA on the benchmark Wikitext-103 corpus, as well as improving over strong baselines on two other language modeling datasets (enwiki8 and Billion Word Benchmark).  The main drawback, as noted by one reviewer, is the computational expense of the method with significant slowdowns compared to the baseline.

Two reviewers voted strong accept, with a third raising several concerns.  The largest concern was the lack of comparison to prior work, especially prior retrieval based methods on two datasets.  The authors responded with an ablation study comparing their method to KNN-LM and showed their proposed GNN-LM performs better.  Other concerns raised by the reviewer were the paper's lack of clarity (the authors should address the reviewers questions during the next revision) and incremental technical contribution.  Another reviewer highlighted the paper's novelty, and this AC agrees it is sufficient for publication.

Overall, the method is an interesting, if expensive, extension of retrieval based language models, and the empirical results support its effectiveness.
Decision: Accept (Spotlight)