Mapping the Timescale Organization of Neural Language Models

1 INTRODUCTION . Language processing requires tracking information over multiple timescales . To be able to predict the final word “ timescales ” in the previous sentence , one must consider both the short-range context ( e.g . the adjective “ multiple ” ) and the long-range context ( e.g . the subject “ language processing ” ) . How do humans and neural language models encode such multi-scale context information ? Neuroscientists have developed methods to study how the human brain encodes information over multiple timescales during sequence processing . By parametrically varying the timescale of intact context , and measuring the resultant changes in the neural response , a series of studies ( Lerner et al. , 2011 ; Xu et al. , 2005 ; Honey et al. , 2012 ) showed that higher-order regions are more sensitive to longrange context change than lower-order sensory regions . These studies indicate the existence of a “ hierarchy of processing timescales ” in the human brain . More recently , Chien & Honey ( 2020 ) used a time-resolved method to investigate how the brain builds a shared representation , when two groups of people processed the same narrative segment preceded by different contexts . By directly mapping the time required for individual brain regions to converge on a shared representation in response to shared input , we confirmed that higher-order regions take longer to build a shared representation . Altogether , these and other lines of investigation suggest that sequence processing in the 1The code and dataset to reproduce the experiment can be found at https : //github.com/ sherrychien/LSTM_timescales brain is supported by a distributed and hierarchical structure : sensory regions have short processing timescales and are primarily influenced by the current input and its short-range context , while higher-order cortical regions have longer timescales and track longer-range dependencies ( Hasson et al. , 2015 ; Honey et al. , 2012 ; Chien & Honey , 2020 ; Lerner et al. , 2011 ; Baldassano et al. , 2017 ; Runyan et al. , 2017 ; Fuster , 1997 ) . How are processing timescales organized within recurrent neural networks ( RNNs ) trained to perform natural language processing ? Long short-term memory networks ( LSTMs ) ( Hochreiter & Schmidhuber , 1997 ) have been widely investigated in terms of their ability to successfully solve sequential prediction tasks . However , long-range dependencies have usually been studied with respect to a particular linguistic function ( e.g . subject-verb number agreement , Linzen et al . 2016 ; Gulordava et al . 2018 ; Lakretz et al . 2019 ) , and there has been less attention on the broader question of how sensitivity to prior context – broadly construed – is functionally organized within these RNNs . Therefore , drawing on prior work in the neuroscience literature , here we demonstrate a model-free approach to mapping processing timescale in RNNs . We focused on existing language models that were trained to predict upcoming tokens at the word level ( Gulordava et al. , 2018 ) and at the character level ( Hahn & Baroni , 2019 ) . The timescale organization of these two models both revealed that the higher layers of LSTM language models contained a small subset of units which exhibit long-range sequence dependencies ; this subset includes previously reported units ( e.g . a “ syntax ” unit , Lakretz et al. , 2019 ) as well as previously unreported units . After mapping the timescales of individual units , we asked : does the processing timescales of each unit in the network relate to its functional role , as measured by its connectivity ? The question is motivated by neuroscience studies which have shown that in the human brain , higher-degree nodes tend to exhibit slower dynamics and longer context dependence than lower-degree nodes ( Baria et al. , 2013 ) . More generally , the primate brain exhibits a core periphery structure in which a relatively small number of “ higher order ” and high-degree regions ( in the prefrontal cortex , in default-mode regions and in so-called “ limbic ” zones ) maintain a large number of connections with one another , and exert a powerful influence over large-scale cortical dynamics ( Hagmann et al. , 2008 ; Mesulam , 1998 ; Gu et al. , 2015 ) . Inspired by the relationships between timescales and network structure in the brain , we set out to test corresponding hypotheses in RNNs : ( 1 ) Do units with longer-timescales tend to have higher degree in neural language models ? and ( 2 ) Do neural language models also exhibit a “ core network ” composed of functionally influential high-degree units ? Using an exploratory network-theoretic approach , we found that units with longer timescales tend to have more projections to other units . Furthermore , we identified a set of medium-to-long timescale “ controller ” units which exhibit distinct and strong projections to control the state of other units , and a set of longtimescale “ integrator units ” which showed influence on predicting words where the long context is relevant . In summary , these findings advance our understanding of the timescale distribution and functional organization of LSTM language models , and provide a method for identifying important units representing long-range contextual information in RNNs . 2 RELATED WORK . Linguistic Context in LSTMs . How do LSTMs encode linguistic context at multiple timescales ? Prior work suggested that the units sensitive to information that requires long-range dependencies are sparse . By ablating one unit at a time , Lakretz et al . ( 2019 ) found two units that encode information required for processing long-range subject-verb number agreement ( one for singular and one for plural information encoding ) . They further identified several long-range “ syntax units ” whose activation was associated with syntactic tree-depth . Overall , Lakretz et al . ( 2019 ) suggests that a sparse subset of units tracks long-range dependencies related to subject-verb agreement and syntax . If this pattern is general – i.e . if there are very few nodes tracking long-range dependencies in general – this may limit the capacity of the models to process long sentences with high complexity , for reasons similar to those that may limit human sentence processing ( Lakretz et al. , 2020 ) . To test whether long-range nodes are sparse in general , we require a model-free approach for mapping the context dependencies of every unit in the language network . Whole-network context dependence . Previous work by Khandelwal et al . ( 2018 ) investigated the duration of prior context that LSTM language models use to support word prediction . Contextdependence was measured by permuting the order of words preceding the preserved context , and observing the increase in model perplexity when the preserved context gets shorter . Khandelwal et al . ( 2018 ) found that up to 200 word-tokens of prior context were relevant to the model perplexity , but that the precise ordering of words only mattered within the most recent 50 tokens . The tokenbased context-permutation method employed in this study was analogous to the approach used to measure context-dependence in human brain responses to movies ( Hasson et al. , 2008 ) and to auditory narratives ( Lerner et al. , 2011 ) . Inspired by the findings of Khandelwal et al . ( 2018 ) and Lakretz et al . ( 2019 ) , in the present study we set out to map the context-dependence across all of the individual units in the LSTM model . This enabled us to relate the timescales to the effects of node-specific ablation and the network architecture itself . In addition , our context manipulations included both context-swapping ( substituting alternative meaningful contexts ) and context-shuffling ( permuting the words in the prior context to disrupt inter-word structure ) , which allowed us to better understand how individual words and syntactically structured word-sequences contribute to the the context representation of individual hidden units . 3 METHODS . 3.1 LANGUAGE MODELS AND CORPUS . We evaluated the internal representations generated by a pre-trained word-level LSTM language model ( WLSTM , Gulordava et al. , 2018 ) as well as a pre-trained character-level LSTM model ( CLSTM , Hahn & Baroni , 2019 ) as they processed sentences sampled from the 427804-word ( 1965719-character ) novel corpus : Anna Karenina by Leo Tolstoy ( Tolstoy , 2016 ) , translated from Russian to English by Constance Garnett . For the WLSTM , we used the model made available by Gulordava et al . ( 2018 ) . The WLSTM has a 650-dimensional embedding layer , two 650-dimensional hidden layers and an output layer with vocabulary size 50,000 . The model was trained and tested on Wikipedia sentences and was not fine-tuned to the novel corpus . Therefore , we only used sentences with low perplexity from the novel in our main timescale analysis . We performed the same analysis using the Wikipedia test set from Gulordava et al . ( 2018 ) and obtained similar results ( See Section 5.3 , Figure A.4A , Appendix A.2.1 ) . For the CLSTM , we used the model made available by Hahn & Baroni ( 2019 ) . The CLSTM has a 200-dimensional embedding layer , three 1024-dimensional hidden layers and an output layer with vocabulary size 63 . The model was trained on Wikipedia data with all characters lower-cased and whitespace removed . We tested the model with sentences sampled from Anna Karenina as the WLSTM model , and we obtained bits-per-character ( BPC ) similar to what Hahn & Baroni ( 2019 ) reported in their original work . 3.2 TEMPORAL CONTEXT CONSTRUCTION PARADIGM . In order to determine the processing timescales of cell state vectors and individual units , we modified the “ temporal context construction ” method developed by Chien & Honey ( 2020 ) . Thus , the internal representations of the model were compared across two conditions : ( 1 ) the Intact Context condition and ( 2 ) the Random Context condition . In both conditions , the model was processing the same shared sequence of words ( for example , segment B ) , but the preceding sentence differed across the two conditions . In the Intact Context condition , the model processed segment B ( the shared segment ) preceded by segment A , which was the actual preceding context from the original text . In the current study , for example , segment A and B are connected by “ , and ” within long sentences from the novel corpus ( Figure 1A ) , to ensure the temporal dependencies between A and B . In the Random Context condition , however , the model processed the same shared input ( segment B ) , but the context was replaced by segment X , which was a randomly sampled context segment from the rest of the corpus . Segment X was therefore not usually coherently related to segment B . For the WLSTM timescale analysis , we chose long sentences in the Intact Context condition that satisfied the following constraints : ( 1 ) mean perplexity across all words in the sentence < 200 , ( 2 ) the shared segment was longer than 25 words , and ( 3 ) the context segment was longer than 10 words . 77 sentences are included as trials in our analyses . In the Random Context condition , we preserved the same shared segments and randomly sampled 30 context segments ( each longer than 10 words ) from other parts of the novel . For the CLSTM timescale analysis , we used the same 77 long sentences in the Intact Context condition , and randomly sampled 25 context segments ( with length > 33 characters ) for the Random Context condition . In brief , the model is processing the same input ( the shared segment ) with different preceding context ( the intact vs. random context ) . We can now measure the context dependence of individual units by examining how the cell state activations differ between the two conditions , while the network is processing the shared segments with identical input . Any difference in internal representations must arise from the context manipulation , since the current input is the same . A decrease in activation difference over time implies that the units exposed in the Intact context and Random context start to build a similar representation as they process the shared input . For a long-timescale unit , whose current state is dependent on information in the far-preceding context , we will see that the activation difference is preserved across contexts ( Figure 1B , green curve ) , even while the unit is processing the shared input . On the other hand , for a short-timescale unit whose activation is driven largely by the current input , we will see that the activation difference drops quickly ( Figure 1B , red curve ) as the unit processes the shared input .

Summary: This paper applies tools from neuroscience to understand how language models integrate across time. The basic approach is to present a phrase, preceded by two different context phrases: one that is natural (i.e. the phrase that actually preceded it in the corpus) and one that is randomly selected. The authors then measure how long it takes for the unit activations to become similar for the two different contexts, which provides a measure for how long the context impacts the representation. They find that (1) timescales increase at later layers of the language model (2) that only a small fraction of units exhibit long timescales (3) that long/medium-timescale units appear to come in two forms which they try and characterize using graph-style analyses.