Multi-timescale Representation Learning in LSTM Language Models

1 INTRODUCTION . Autoregressive language models are functions that estimate a probability distribution over the next word in a sequence from past words , p ( wt|wt−1 , . . . , w1 ) . This requires capturing statistical dependencies between words over short timescales , where syntactic information likely dominates ( Adi et al. , 2017 ; Linzen et al. , 2016 ) , as well as long timescales , where semantic and narrative information likely dominate ( Zhu et al. , 2018 ; Conneau et al. , 2018 ; Gulordava et al. , 2018 ) . Because this probability distribution grows exponentially with sequence length , some approaches simplify the problem by ignoring long-range dependencies . Classical n-gram models , for example , assume word wt is independent of all but the last n− 1 words , with typical n = 5 ( Heafield , 2011 ) . Hidden Markov models ( HMMs ) assume that the influence of previous words decays exponentially with distance from the current word ( Lin & Tegmark , 2016 ) . In contrast , neural network language models such as recurrent ( Hochreiter & Schmidhuber , 1997 ; Merity et al. , 2018 ; Melis et al. , 2018 ) and transformer networks ( Melis et al. , 2019 ; Krause et al. , 2019 ; Dai et al. , 2019 ) include longer-range interactions , but simplify the problem by working in lower-dimensional representational spaces . Attention-based networks combine position and ∗Current affiliation : Apple Inc. content-based information in a small number of attention heads to flexibly capture different types of dependencies within a sequence ( Vaswani et al. , 2017 ; Cordonnier et al. , 2019 ) . Gated recurrent neural networks ( RNNs ) compress information about past words into a fixed-length state vector ( Hochreiter & Schmidhuber , 1997 ) . The influence each word has on this state vector tends to decay exponentially over time . However , each element of the state vector can have a different exponential time constant , or “ timescale ” ( Tallec & Ollivier , 2018 ) , enabling gated RNNs like the long short-term memory ( LSTM ) network to flexibly learn many different types of temporal relationships ( Hochreiter & Schmidhuber , 1997 ) . Stacked LSTM networks reduce to a single layer ( Turek et al. , 2020 ) , showing that network depth has an insignificant influence on how the LSTM captures temporal relationships . Yet in all these networks the shape of the temporal dependencies must be learned directly from the data . This seems particularly problematic for very long-range dependencies , which are only sparsely informative ( Lin & Tegmark , 2016 ) . This raises two related questions : what should the temporal dependencies in a language model look like ? And how can that information be incorporated into a neural network language model ? To answer the first question , we look to empirical and theoretical work that has explored the dependency statistics of natural language . Lin & Tegmark ( 2016 ) quantified temporal dependencies in English and French language corpora by measuring the mutual information between tokens as a function of the distance between them . They observed that mutual information decays as a power law , i.e . MI ( wk , wk+t ) ∝ t−d for constant d. This behavior is common to hierarchically structured natural languages ( Lin & Tegmark , 2016 ; Sainburg et al. , 2019 ) as well as sequences generated from probabilistic context-free grammars ( PCFGs ) ( Lin & Tegmark , 2016 ) . Now to the second question : if temporal dependencies in natural language follow a power law , how can this information be incorporated into neural network language models ? To our knowledge , little work has explored how to control the temporal dependencies learned in attention-based models . However , many approaches have been proposed for controlling gated RNNs , including updating different groups of units at different intervals ( El Hihi & Bengio , 1996 ; Koutnik et al. , 2014 ; Liu et al. , 2015 ; Chung et al. , 2017 ) , gating units across layers ( Chung et al. , 2015 ) , and explicitly controlling the input and forget gates that determine how information is stored and removed from memory ( Xu et al. , 2016 ; Shen et al. , 2018 ; Tallec & Ollivier , 2018 ) . Yet none of these proposals incorporate a specific shape of temporal dependencies based on the known statistics of natural language . In this work , we build on the framework of Tallec & Ollivier ( 2018 ) to develop a theory for how the memory mechanism in LSTM language models can capture temporal dependencies that follow a power law . This relies on defining the timescale of an individual LSTM unit based on how the unit retains and forgets information . We show that this theory predicts the distribution of unit timescales for LSTM models trained on both natural English ( Merity et al. , 2018 ) and formal languages ( Suzgun et al. , 2019 ) . Further , we show that forcing models to follow this theoretical distribution improves language modeling performance . These results highlight the importance of combining theoretical modeling with an understanding of how language models capture temporal dependencies over multiple scales . 2 MULTI-TIMESCALE LANGUAGE MODELS . 2.1 TIMESCALE OF INFORMATION . We are interested in understanding how LSTM language models capture dependencies across time . Tallec & Ollivier ( 2018 ) elegantly argued that memory in individual LSTM units tends to decay exponentially with a time constant determined by weights within the network . We refer to the time constant of that exponential decay as the unit ’ s representational timescale . Timescale is directly related to the LSTM memory mechanism ( Hochreiter & Schmidhuber , 1997 ) , which involves the LSTM cell state ct , input gate it and forget gate ft , it = σ ( Wixxt +Wihht−1 + bi ) ft = σ ( Wfxxt +Wfhht−1 + bf ) c̃t = tanh ( Wcxxt +Wchht−1 + bc ) ct = ft ct−1 + it c̃t , where xt is the input at time t , ht−1 is the hidden state , Wih , Wix , Wfh , Wfx , Wch , Wcx are the different weights and bi , bf , bc the respective biases . σ ( · ) and tanh ( · ) represent element-wise sigmoid and hyperbolic tangent functions . Input and forget gates control the flow of information in and out of memory . The forget gate ft controls how much memory from the last time step ct−1 is carried forward to the current state ct . The input gate it controls how much information from the input xt and hidden state ht−1 at the current timestep is stored in memory for subsequent timesteps . To find the representational timescale , we consider a “ free input ” regime with zero input to the LSTM after timestep t0 , i.e. , xt = 0 for t > t0 . Ignoring information leakage through the hidden state ( i.e. , assuming Wch = 0 , bc = 0 , and Wfh = 0 ) the cell state update becomes ct = ft ct−1 . For t > t0 , it can be further simplified as ct = f t−t0 0 c0 = e ( log f0 ) ( t−t0 ) c0 , ( 1 ) where c0 = ct0 is the cell state at t0 , and f0 = σ ( bf ) is the value of the forget gate , which depends only on the forget gate bias bf . Equation 1 shows that LSTM memory exhibits exponential decay with characteristic forgetting time T = − 1 log f0 = 1 log ( 1 + e−bf ) . ( 2 ) That is , values in the cell state tend to shrink by a factor of e every T timesteps . We refer to the forgetting time in Equation 2 as the representational timescale of an LSTM unit . Beyond the “ free input ” regime , we can estimate the timescale for a LSTM unit by measuring the average forget gate value over a set of test sequences , Test = − 1 log f̄ , ( 3 ) where f̄ = 1KN ∑N j=1 ∑K t=1 f j t , in which f j t is the forget gate value of the unit at t-th timestep for j-th test sequence , N is the number of test sequences , and K is the test sequence length . 2.2 COMBINING EXPONENTIAL TIMESCALES TO YIELD A POWER LAW . From earlier work , we know that temporal dependencies in natural language tend to decay following a power law ( Lin & Tegmark , 2016 ; Sainburg et al. , 2019 ) . Yet from Equation 1 we see that LSTM memory tends to decay exponentially . These two decay regimes are fundamentally different–the ratio of a power law divided by an exponential always tends towards infinity . However , LSTM language models contain many units , each of which can have a different timescale . Thus LSTM language models might approximate power law decay through a combination of exponential functions . Here we derive a theory for how timescales should be distributed within an LSTM in order to yield overall power law decay . Let us assume that the timescale T for each LSTM unit is drawn from a distribution P ( T ) . We want to find a P ( T ) such that the expected value over T of the function e −t T approximates a power law decay t−d for some constant d , t−d ∝ ET [ e−t/T ] = ∫ ∞ 0 P ( T ) e−t/T dT . ( 4 ) Solving this problem reveals that P ( T ) is an Inverse Gamma distribution with shape parameter α = d and scale parameter β = 1 ( see Section A.1 for derivation ) . The probability density function of the Inverse Gamma distribution is given as P ( T ; α , β ) = β α Γ ( α ) ( 1/T ) α+1e ( −β/T ) . This theoretical result suggests that in order to approximate the power law decay of information in natural language , unit timescales in LSTM language models should follow an Inverse Gamma distribution . We next perform experiments to test whether this prediction holds true for models trained on natural language and models trained on samples from a formal language with known temporal statistics . We then test whether enforcing an Inverse Gamma timescale distribution at training time improves model performance . 2.3 CONTROLLING LSTM UNIT TIMESCALES . To enforce a specific distribution of timescales in an LSTM and thus create an explicit multi-timescale model , we drew again upon the methods developed in Tallec & Ollivier ( 2018 ) . Following the analysis in Section 2.1 , the desired timescale Tdesired for an LSTM unit can be controlled by setting the forget gate bias to the value bf = − log ( e 1 Tdesired − 1 ) . ( 5 ) The balance between forgetting information from the previous timestep and adding new information from the current timestep is controlled by the relationship between forget and input gates . To maintain this balance we set the input gate bias bi to the opposite value of the forget gate , i.e. , bi = −bf . This ensures that the relation it ≈ 1− ft holds true . Importantly , these bias values remain fixed ( i.e . are not learned ) during training , in order to keep the desired timescale distribution across the network .

Summary: This work investigates representational power of LSTM to model natural language, in particular how well it models temporal dependencies within text. They define a notion of timescale of each LSTM unit and analitycally show that LSTM memory exhibits exponential decay, while natural language tends (based on prior work) to decay following the power law. Based on this, they figure that LSTM memory may decay following the power law *if the timescales approximate samples from the particular Inverse Gamma Distribution*. To achieve that they propose the multi-timescale LSTM unit, where the desired timescale is explicitly controlled via the forget gate bias.