GNN-LM: Language Modeling based on Global Contexts via GNN

1 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 ) . 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 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 . 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 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 a strategy different from memorization or close-book examination – copy , or in other words , open-book examination . For example , given a prefix “ 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 1The code can be found at https : //github.com/ShannonAI/GNN-LM – 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 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 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. , 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 significantly easier for the model to predict upcoming words in the LM task . We further combine GNN-LM with kNN-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 . Specifically , applying the GNN-LM framework to a strong base LM leads to substantial performance boost ( -1.9 perplexity ) on WikiText-103 , and combining with kNN-LM achieves a new state-of-the-art perplexity of 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 1 . At each time step t , a neural language model ( LM ) f ( · ) first encodes a sequence of context tokens ct = ( w1 , w2 , ... , wt−1 ) to a high-dimensional representation ht = f ( ct ) ∈ Rd , where d is the dimension of hidden states . Then a transformation matrixW ∈ RV×d is used to estimate the probability of the t-th token p ( wt|ct ) = softmax ( Wht ) , 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 ht generated by the vanilla LM . The updated representation , which aggregates additional information from reference tokens , is then used to estimate pLM ( wt|ct ) . 2.2 GRAPH CONSTRUCTION . The first step of our proposed framework is to build a graph capturing the connections between the context tokens ct = ( w1 , w2 , ... , wt−1 ) and those similar to ct in the training set . To this end , we construct a directed heterogeneous graph , where the nodes are tokens from ct 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 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 define two types of nodes A = { ao , an } , where ao means that the node is within the input ct. an means the node is in N ( ct ) , the set of extracted contexts within the neighborhood of ct. We also define two types of edges R = { rinter , rintra } , where rinter means inter-context connection ( from an nodes to ao nodes ) and rintra means intra-context connection ( between two nodes of same type ) . Each token within the input is a node of type ao , and edges of type rintra are constructed from node wi to wj ( 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 ct , we retrieve k nearest neighbors N ( ct ) = { c ( 1 ) t1 , ... , c ( k ) tk } of ct from the training set as follows : we first use ht 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 similarity2 , 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 ht . 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 kNN-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 ) ti 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 find that using shallow ( i.e . 3 ) GNN layers and adding rintra edges between adjacent tokens can alleviate overfitting . Since a 3-layer GNN only aggregates information from 3-hop nodes in the graph , using larger l and r have no 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 , to accommodate the modeling of rintra from node wi to wj ( i ≤ j ) within ct , where Transformer with self-attention is usually adopted , we extend the self-attention mechanism to rinter , and construct a self-attention augmented GNN . Specifically , the l-th layer representation of node n is computed by ( here we use the superscript [ l ] to represent the l-th layer ) : h [ l ] n = Aggregate ∀s∈N ( n ) ( Attention ( s , e , n ) · Feature ( s , e , n ) ) + h [ l−1 ] n . ( 1 ) 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 . ( 2020 ) . 2In practice , we use FAISS ( Johnson et al. , 2019 ) for fast approximate kNN search . 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 final 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 : K ( s ) =W kτ ( s ) h [ l−1 ] s , Q ( n ) =W q τ ( n ) h [ l−1 ] n , Attention ( s , e , n ) = 1 Z exp ( K ( s ) W ATTφ ( e ) Q ( n ) > · µ〈τ ( s ) , φ ( e ) , τ ( n ) 〉√ d ) , Z = ∑ s′∈N ( n ) , e′∈φ ( e ) Attention ( s′ , e′ , n ) , ( 2 ) 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 defined by : Feature ( s , e , n ) =W vτ ( s ) h [ l−1 ] s W FEA φ ( e ) , ( 3 ) whereW vτ ( s ) ∈ R d×d andW 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 : Aggregate ( · ) =W oτ ( n ) ( ⊕ ∀s∈N ( n ) ( Attention ( s , e , n ) · Feature ( s , e , n ) ) ) ( 4 ) where ⊕ is element-wise addition and W oτ ( n ) ∈ R d×d is model parameter . The representation of token wt−1 from the last layer is used to compute the language model probability pLM ( wt|ct ) . 2.4 kNN BASED PROBABILITY FOR NEXT TOKEN We further incorporate the proposed model with kNN ( 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 ct = ( w1 , w2 , ... , wt−1 ) , we retrieve the k nearest neighbors N ( ct ) = { c ( 1 ) t1 , ... , c ( k ) tk } , and compute the kNN based probability for the next token by : p ( wt|ct ) = λpkNN ( wt|ct ) + ( 1− λ ) pLM ( wt|ct ) , pkNN ( wt|ct ) = 1 Z k∑ i=1 1 wt=w ( i ) ti exp ( cos ( f ( ct ) , f ( c ( i ) ti ) ) /T ) , ( 5 ) 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.3

Summary: 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.   Contributions: 1. a new retrieval-augmented language model implemented via GNN techniques. 2. improvements over state-of-the-art models on three benchmarks.