State-Regularized Recurrent Networks

1 INTRODUCTION . Recurrent neural networks ( RNNs ) have found their way into numerous applications . Still , RNNs have two shortcomings . First , it is difficult to understand what concretely RNNs learn . Some applications require a close inspection of learned models before deployment and RNNs are more difficult to interpret than rule-based systems . There are a number of approaches for extracting finite state automata ( DFAs ) from trained RNNs ( Giles et al. , 1991 ; Wang et al. , 2018b ; Weiss et al. , 2018b ) as a means to analyze their behavior . These methods apply extraction algorithms after training and it remains challenging to determine whether the extracted DFA faithfully models the RNN ’ s state transition behavior . Most extraction methods are rather complex , depend crucially on hyperparameter choices , and tend to be computationally costly . Second , RNNs tend to work poorly on input sequences requiring long-term memorization , despite having this ability in principle . Indeed , there is a growing body of work providing evidence , both empirically ( Daniluk et al. , 2017a ; Bai et al. , 2018 ) and theoretically ( Miller & Hardt , 2018 ) , that recurrent networks offer no benefit on longer sequences , at least under certain conditions . Intuitively , RNNs tend to operate more like DFAs with a large number of states , attempting to memorize all the information about the input sequence solely with their hidden states , and less like automata with external memory . We propose state-regularized RNNs as a possible step towards addressing both of the aforementioned problems . State-regularized RNNs ( SR-RNNs ) are a class of recurrent networks that utilize a stochastic state transition mechanism between cell applications . The stochastic mechanism models a probabilistic state dynamics that lets the SR-RNN transition between a finite number of learnable states . The parameters of the stochastic mechanism are trained jointly with the parameters of the base RNN . SR-RNNs have several advantages over standard RNNs . Instead of having to apply post-training DFA extraction , it is possible to determine their ( probabilistic and deterministic ) state transition behavior more directly . We propose a method that extracts DFAs truly representing the state transition behavior of the underlying RNNs . We hypothesize that the frequently-observed poor extrapolation behavior of RNNs is caused by memorization with hidden states . It is known that RNNs – even those with cell states or external memory – tend to memorize mainly with their hidden states and in an unstructured manner ( Strobelt et al. , 2016 ; Hao et al. , 2018 ) . We show that the state-regularization mechanism shifts representational power to memory components such as the cell state , resulting in improved extrapolation performance . We support our hypotheses through experiments both on synthetic and real-world datasets . We explore the improvement of the extrapolation capabilities of SR-RNNs and closely investigate their memorization behavior . For state-regularized LSTMs , for instance , we observe that memorization can be shifted entirely from the hidden state to the cell state . For text and visual data , stateregularization provides more intuitive interpretations of the RNNs ’ behavior . 2 BACKGROUND . We provide some background on deterministic finite automata ( DFAs ) and deterministic pushdown automata ( DPDAs ) for two reasons . First , one contribution of our work is a method for extracting DFAs from RNNs . Second , the state regularization we propose is intended to make RNNs behave more like DPDAs and less like DFAs by limiting their ability to memorize with hidden states . A DFA is a state machine that accepts or rejects sequences of tokens and produces one unique computation path for each input . Let Σ∗ be the language over the alphabet Σ and let be the empty sequence . A DFA over an alphabet ( set of tokens ) Σ is a 5-tuple ( Q , Σ , δ , q0 , F ) consisting of finite set of states Q ; a finite set of input tokens Σ called the input alphabet ; a transition functions δ : Q × Σ → Q ; a start state q0 ; and a set of accept states F ⊆ Q . A sequence w is accepted by the DFA if the application of the transition function , starting with q0 , leads to an accepting state . Figure 1 ( center ) depicts a DFA for the language of balanced parentheses ( BP ) up to depth 4 . A language is regular if and only if it is accepted by a DFA . A pushdown automata ( PDA ) is defined as a 7-tuple ( Q , Σ , Γ , δ , q0 , ⊥ , F ) consisting of a finite set of states Q ; a finite set of input tokens Σ called the input alphabet ; a finite set of tokens Γ called the stack alphabet ; a transition function δ ⊆ Q × ( Σ ∪ ) × Γ → Q × Γ∗ ; a start state q0 ; the initial stack symbol ⊥ ; and a set of accepting states F ⊆ Q. Computations of the PDA are applications of the transition relations . The computation starts in q0 with the initial stack symbol⊥ on the stack and sequence w as input . The pushdown automaton accepts w if after reading w the automaton reaches an accepting state . Figure 1 ( right ) depicts a deterministic PDA for the language BP . 3 STATE-REGULARIZED RECURRENT NEURAL NETWORKS . The standard recurrence of an RNN is ht = f ( ht−1 , ct−1 , xt ) where ht is the hidden state vector at time t , ct is the cell state at time t , and xt is the input symbol at time t. We refer to RNNs whose unrolled cells are only connected through the hidden output states h and no additional vectors such as the cell state , as memory-less RNNs . For instance , the family of GRUs ( Chung et al. , 2014 ) does not have cell states and , therefore , is memory-less . LSTMs ( Hochreiter & Schmidhuber , 1997 ) , on the other hand , are not memory-less due to their cell state . A cell of a state-regularized RNN ( SR-RNN ) consist of two components . The first component , which we refer to as the recurrent component , applies the function of a standard RNN cell ut = f ( ht−1 , ct−1 , xt ) . For the sake of completeness , we include the cell state c here , which is absent in memory-less RNNs . We propose a second component which we refer to as stochastic component . The stochastic component is responsible for modeling the probabilistic state transitions mechanism that lets the RNN transition implicitly between a finite number of states . Let d be the size of the hidden state vectors of the recurrent cells . Moreover , let ∆D : = { λ ∈ RD+ : ‖ λ ‖= 1 } be the ( D − 1 ) probability simplex . The stochastic component maintains k learnable centroids s1 , ... , sk of size d which we often write as the column vectors of a matrix S ∈ Rd×k . The weights of these centroids are global parameters shared among all cells . The stochastic component computes , at each time step t , a discrete probability distribution from the output ut of the recurrent component and the centroids of the stochastic component α = g ( S , ut ) with α ∈ ∆k . ( 1 ) Crucially , instances of g should be differentiable to facilitate end-to-end training . Typical instances of the function g are based on the dot-product or the Euclidean distance , normalized into a probability distribution αi = exp ( ( ut · si ) /τ ) ∑k i=1 exp ( ( ut · si ) /τ ) ( 2 ) αi = exp ( − ‖ ut − si ‖ /τ ) ∑k i=1 exp ( − ‖ ut − si ‖ /τ ) ( 3 ) Here , τ is a temperature parameter that can be used to anneal the probabilistic state transition behavior . The lower τ the more α resembles the one-hot encoding of a centroid . The higher τ the more uniform is α . Equation 2 is reminiscent of the equations of attentive mechanisms ( Bahdanau et al. , 2015 ; Vaswani et al. , 2017 ) . Instead of attending to the hidden states , however , SR-RNNs attend to the k centroids to compute transition probabilities . Each αi is the probability of the RNN to transition to centroid ( state ) i given the vector ut for which we write put ( i ) = αi . The state transition dynamics of an SR-RNN is that of a probabilistic finite state machine . At each time step , when being in state ht−1 and reading input symbol xt , the probability for transitioning to state si is αi . Hence , in its second phase the stochastic component computes the hidden state ht at time step t from the distribution α and the matrix S with a ( possibly stochastic ) mapping h : ∆k × Rd×k → Rd . Hence , ht = h ( α , S ) . A simple instance of the mapping h is to sample j ∼ put and set ht = sj . ( 4 ) This renders the SR-RNN not end-to-end differentiable , however , and one has to use EM or reinforcement learning strategies which are often less stable and less efficient . A possible alternative is to set the hidden state ht to be the probabilistic mixture of the centroids ht = k∑ i=1 αisi . ( 5 ) Every internal state h of the SR-RNN , therefore , is computed as a weighted sum h = α1s1 + ... + αksk of the centroids s1 , ... , sk with α ∈ ∆k . Here , h is a smoothed variant of the function that computes a hard assignment to one of the centroids . One can show that for τ → 0 the state dynamics based on equations ( 4 ) and ( 5 ) are identical and correspond to those of a DFA . Figure 2 depicts two variants of the proposed SR-RNNs . Additional instances of h are conceivable . For instance , one could , for every input sequence and the given current parameters of the SR-RNN , compute the most probable state sequence and then backpropagate based on a structured loss . Since finding these most probable sequences is possible with Viterbi type algorithms , one can apply a form of differentiable dynamic programming ( Mensch & Blondel , 2018 ) . The probabilistic state transitions of SR-RNNs open up new possibilities for applying more complex differentiable functions . We leave these considerations to future work . The probabilistic state transition mechanism is also applicable when RNNs have more than one hidden layer . In RNNs with l > 1 hidden layers , every such layer can maintain its own centroids and stochastic component . In this case , a global state of the SR-RNN is an l-tuple , with the lth argument of the tuple corresponding to the centroids of the lth layer . Even though we have augmented the original RNN with additional learnable parameter vectors , we are actually constraining the SR-RNN to output hidden state vectors that are similar to the centroids . For lower temperatures and smaller values for k , the ability of the SR-RNN to memorize with its hidden states is increasingly impoverished . We argue that this behavior is beneficial for three reasons . First , it makes the extraction of interpretable DFAs from memory-less SR-RNNs straightforward . Instead of applying post-training DFA extraction as in previous work , we extract the true underlying DFA directly from the SR-RNN . Second , we hypothesize that overfitting in the context of RNNs is often caused by memorization via hidden states . Indeed , we show that regularizing the state space pushes representational power to memory components such as the cell state of an LSTM , resulting in improved extrapolation behavior . Third , the values of hidden states tend to increase in magnitude with the length of the input sequence , a behavior that has been termed drifting ( Zeng et al. , 1993 ) . The proposed state regularization stabilizes the hidden states for longer sequences . First , let us explore some of the theoretical properties of the proposed mechanism . We show that the addition of the stochastic component , when capturing the complete information flow between cells as , for instance , in the case of GRUs , makes the resulting RNN ’ s state transition behavior identical to that of a probabilistic finite state machine . The proofs of the theorems are part of the appendix . Theorem 3.1 . The state transition behavior of a memory-less SR-RNN using equation 4 is identical to that of a probabilistic finite automaton . We can show that the lower the temperature the more memory-less RNNs operate like DFAs . Theorem 3.2 . For τ → 0 the state transition behavior of a memory-less SR-RNN ( using equations 4 or 5 ) is equivalent to that of a deterministic finite automaton .

Summary: This paper is based on the observation that LSTMs use the hidden state to memorize information and the cell state (memory) is not fully utilized. To encourage the LSTM to utilize the cell state, authors constraint the hidden state to a set of centroid states and learn to transition between these centroids in a soft way. Authors demonstrate their model in learning simple regular and context-free languages and also in a couple of non-synthetic tasks. The proposed model also has some interpretability of internal state transitions.