Learning Hierarchical Structures with Differentiable Nondeterministic Stacks

1 Introduction . Manymachine learning problems involve sequential data with hierarchical structures , such as modeling context-free languages ( Grefenstette et al. , 2015 ; DuSell & Chiang , 2020 ) , evaluating mathematical expressions ( Nangia & Bowman , 2018 ; Hao et al. , 2018 ) , logical inference ( Bowman et al. , 2015 ) , and modeling syntax in natural language ( Dyer et al. , 2016 ; Shen et al. , 2019b ; Kim et al. , 2019 ) . However , recurrent neural networks ( RNNs ) have difficulty learning to solve these tasks , or generalizing to held-out sequences , unless they have supervision or a hierarchical inductive bias ( van Schijndel et al. , 2019 ; Wilcox et al. , 2019 ; McCoy et al. , 2020 ) . A limiting factor of RNNs is their reliance on memory whose size is constant across time . For example , to predict the second half of a string of the form F # FR , a simple RNN would need to store all of F in its hidden state before predicting its reversal FR ; a memory of finite size will inevitably fail to do this for inputs exceeding a certain length . To remedy this , some previous work has investigated the addition of differentiable stack data structures to RNNs ( Sun et al. , 1995 ; Grefenstette et al. , 2015 ; Joulin & Mikolov , 2015 ; DuSell & Chiang , 2020 ) , which is closely related to work on neural networks that model shift-reduce parsers ( Bowman et al. , 2016 ; Dyer et al. , 2016 ; Shen et al. , 2019a ) . Just as adding a stack to a finite state machine , which makes it a pushdown automaton ( PDA ) , enables it to recognize context-free languages ( CFLs ) , the hope is that adding stacks to RNNs will increase the range of problems on which they can be used effectively . We also expect stacks to aid training by introducing an inductive bias for learning hierarchical patterns , and to increase generalization power by structuring the model ’ s memory in a way that better predicts held-out hierarchical data . DuSell & Chiang ( 2020 ) recently proposed a stack-based RNN called the Nondeterministic Stack RNN ( NS-RNN ) that outperformed other stack RNNs on a range of CFL language modeling tasks . This model ’ s defining feature is that its external data structure is a nondeterministic PDA , allowing it to simulate an exponential number of sequences of stack operations in parallel . This is in contrast to prior stack RNNs ( Grefenstette et al. , 2015 ; Joulin & Mikolov , 2015 ) which model deterministic stacks , being designed to learn one correct stack operation at each time step . One reason nondeterminism is important is that deterministic CFLs are a proper subset of CFLs . If the analogy with PDAs holds true , then equipping an RNN with a deterministic stack would only enable it to model deterministic CFLs , whereas a nondeterministic stack should enable it to model all CFLs . This is important for natural language processing , as human language is known to be full of syntactic ambiguity . Another benefit of nondeterminism , even on deterministic CFLs , applies to training . In order for a model to receive a reward for an action , it must try the action ( that is , give it nonzero probability so that it receives gradient during backpropagation ) . For example , in the digit-recognition task , a classifier tries all ten digits , and is rewarded for the correct one . But in a stack-augmented model , the space of possible action sequences is very large . Whereas a deterministic stack can only try one of them , a nondeterministic stack can try all of them and always receives a reward for the correct one . But as explained in §3.1 , because the NS-RNN ’ s probability for an action sequence is the product of many probabilities , it can be extremely small , so the NS-RNN sometimes learns very slowly . In this paper we present a new model , the Renormalizing NS-RNN ( RNS-RNN ) , which is based on the NS-RNN , but improves its performance on all of the CFL tasks it was originally tested on , thanks to two key changes . The first is that stack actions have weights that do not necessarily form a probability distribution ( §3.1 ) . They define an unnormalized distribution over stacks that is renormalized whenever the model queries it . Second , the RNS-RNN includes not only top stack symbols but also PDA states in this query ( §3.2 ) . These changes allow the RNS-RNN to attain lower cross-entropy on CFL tasks ( in fact , very close to the information-theoretic lower bound ) and to surpass deterministic stack RNNs on a task on which the NS-RNN fails to do so ( “ padded reversal ” ) . Finally , as a third modification , we present a memory-restricted version of the RNS-RNN that requires only $ ( = ) time and space ( §5 ) . This restricted RNS-RNN can be run incrementally on arbitrarily long sequences , which is a necessity for language modeling on natural language , for which we provide experimental results . We will release our code publicly . 2 Previous stack RNNs . We begin by discussing three previously proposed stack RNNs , each of which uses a different style of differentiable stack : stratification ( Das et al. , 1992 ; Sun et al. , 1995 ; Grefenstette et al. , 2015 ) , superposition ( Joulin & Mikolov , 2015 ) , and nondeterminism ( DuSell & Chiang , 2020 ) . 2.1 Controller-stack interface . Each type of stack RNN consists of a simple RNN ( or variant such as LSTM ) , called the controller , connected to a differentiable stack . The stack has no parameters of its own ; its role is to accept actions from the controller to push and pop elements at each time step , simulate those actions , and return a reading to the controller as an extra input to the next time step that serves as a representation of the updated top element of the stack . The stack actions and stack reading take on continuous values so that they may be differentiable ; their form and interpretation vary with architecture . Following DuSell & Chiang ( 2020 ) , we make minor changes to the original model definitions given by Grefenstette et al . ( 2015 ) and Joulin & Mikolov ( 2015 ) to ensure that all three of these stack RNN models conform to the same controller-stack interface . This allows us to isolate differences in the style of stack data structure employed while keeping other parts of the network the same . We assume the input F = F1 · · ·F= is encoded as a sequence of vectors x1 , · · · , x= . In all of our experiments , we use an LSTM ( Hochreiter & Schmidhuber , 1997 ) as the controller , whose memory consists of a hidden state hC and memory cell cC . The controller computes the next state ( hC , cC ) given the previous state ( hC−1 , cC−1 ) , input vector xC , and stack reading rC−1 : ( hC , cC ) = LSTM ( ( hC−1 , cC−1 ) , [ xC rC−1 ] ) . We set h0 = c0 = 0 . The hidden state is used to compute the stack actions 0C and predict the logits yC for the next word FC+1 . The previous stack and new actions are used to compute a new stack BC , which in turn is used to produce a new reading rC : 0C = Actions ( hC ) yC = WhyhC + bhy BC = Stack ( BC−1 , 0C−1 ) rC = Reading ( BC ) In order to change the stack data structure , we need only change the definitions of Actions , Stack , Reading , and B0 , which may depend on parameters of the model ; for our changes to the NS-RNN , we will only need to change Actions and Reading . 2.2 Stratification . Based on work by Das et al . ( 1992 ) and Sun et al . ( 1995 ) , the stack of Grefenstette et al . ( 2015 ) relies on a strategy dubbed “ stratification ” by DuSell & Chiang ( 2020 ) . The elements of the stack are vectors , each of which is associated with a “ thickness ” between 0 and 1 , which represents the degree to which the vector element is present on the stack . A helpful analogy is that of layers of a cake ; the stack elements are like cake layers of varying thickness . In this model , 0C = ( DC , 3C , vC ) , where the pop signal DC ∈ ( 0 , 1 ) indicates the amount to be removed from the top of the stack , vC is a learned vector to be pushed as a new element onto the stack , and the push signal 3C ∈ ( 0 , 1 ) is the thickness of that newly pushed vector . This model has quadratic time and space complexity with respect to input length . We refer the reader to Appendix A.1 for full details . 2.3 Superposition . The stack of Joulin & Mikolov ( 2015 ) simulates a combination of partial stack actions by computing three new , separate stacks : one with all cells shifted down ( push ) , kept the same ( no-op ) , and shifted up ( pop ) . The new stack is then an element-wise interpolation ( “ superposition ” ) of these three stacks . In this model , stack elements are again vectors , and 0C = ( aC , vC ) , where the vector aC is a probability distribution over three stack operations : push a new vector , no-op , and pop the top vector ; vC is the vector to be pushed . The vector vC can be learned or can be set to hC ( Yogatama et al. , 2018 ) . The stack reading is the top cell of the stack . This model has quadratic time and space complexity with respect to input length . We refer the reader to Appendix A.2 for full details . 2.4 Nondeterminism . The stack module in the Nondeterministic Stack RNN ( NS-RNN ) model of DuSell & Chiang ( 2020 ) maintains a probability distribution over whole stacks by simulating a weighted PDA . It has cubic time complexity and quadratic space complexity with respect to input length , leading to higher wall-clock run time than other stack RNNs , but often better task performance . The simulated weighted PDA maintains a state drawn from a finite set & , which includes an initial state @ 0 , and a stack with symbols drawn from an alphabet Γ , which includes an initial symbol ⊥ . At each time step , the PDA executes a weighted transition that changes its state and manipulates the stack . Stack operations are drawn from the set Op ( Γ ) = •Γ ∪ Γ ∪ { n } , where for any H ∈ Γ , •H means “ push H , ” H means “ replace the top element with H , ” and n means “ pop the top element. ” A valid sequence of transitions is called a run , and the weight of a run is the product of the weights of its transitions . The RNN controller emits transition weights to the stack module . Note that the stack module , not the controller , keeps track of PDA states and stack configurations , so the controller emits distributions over transitions conditioned on the PDA ’ s current state and top stack symbol . More precisely , 0C = Δ [ C ] is a tensor where the meaning of element Δ [ C ] [ @ , G → A , h ] is : if the PDA is in state @ and the top stack symbol is G , then , with weight Δ [ C ] [ @ , G → A , h ] , go to state A and perform stack operation h ∈ Op ( Γ ) . The original NS-RNN definition requires that for all C , @ , and G , the weights form a probability distribution . Accordingly , they are computed from the hidden state using a softmax layer : Δ [ C ] = softmax @ , G ( WhahC + bha ) . ( 1 ) The stack module marginalizes over all runs ending at time step C and returns the distribution over top stack symbols at C to the controller . It may appear that computing this probability distribution is intractable because the number of possible runs is exponential in C , but Lang ( 1974 ) gives a dynamic programming algorithm that simulates all runs of a nondeterministic PDA in cubic time and quadratic space . Lang ’ s algorithm exploits structural similarities in PDA runs . First , multiple runs can result in the same stack . Second , for : > 0 , a stack of height : must have been derived from a stack of height : − 1 , so in principle representing a stack of height : requires only storing its top symbol and a pointer to a stack of height : − 1 . The resulting data structure is a weighted graph where edges represent individual stack symbols , and paths ( of which there are exponentially many ) represent stacks . We may equivalently view this graph as a weighted finite automaton ( WFA ) that encodes a distribution over stacks , and accordingly the NS-RNN ’ s stack module is called the stack WFA . Indeed , the language of stacks at a given time step C is always regular ( Autebert et al. , 1997 ) , and Lang ’ s algorithm gives an explicit construction for the WFA encoding this language . Its states are PDA configurations of the form ( 8 , @ , G ) , where 0 ≤ 8 ≤ = , @ ∈ & , and G ∈ Γ is the stack top . A stack WFA transition from ( 8 , @ , G ) to ( C , A , H ) means the PDA went from configuration ( 8 , @ , G ) to ( C , A , H ) ( possibly via multiple time steps ) , where the only difference in the stack is that a single H was pushed , and the G was never modified in between . The weights of these transitions are stored in a tensor W of shape = × = × | & | × |Γ| × | & | × |Γ | , where elements are written W [ 8 −→ C ] [ @ , G −→ A , H ] . For 0 ≤ 8 < C ≤ = , W [ 8 −→ C ] [ @ , G −→ A , H ] = I [ 8 = C−1 ] Δ [ C ] [ @ , G → A , •H ] push + ∑ B , I W [ 8 −→ C−1 ] [ @ , G −→ B , I ] Δ [ C ] [ B , I → A , H ] repl . + C−2∑ : =8+1 ∑ D ∑ B , I W [ 8 −→ : ] [ @ , G −→ D , H ] W [ : −→ C−1 ] [ D , H −→ B , I ] Δ [ C ] [ B , I → A , n ] pop ( 2 ) The NS-RNNmarginalizes over all stacks/accepting paths in the stackWFA using a tensor U of forward weights of shape = × | & | × |Γ | . The weight U [ C ] [ A , H ] is the total weight of reaching configuration ( C , A , H ) . These weights are normalized to compute the distribution over top stack symbols at C : U [ 0 ] [ A , H ] = I [ A = @ 0 ∧ H = ⊥ ] ( 3 ) U [ C ] [ A , H ] = C−1∑ 8=1 ∑ @ , G U [ 8 ] [ @ , G ] W [ 8 −→ C ] [ @ , G −→ A , H ] ( 1 ≤ C ≤ = ) ( 4 ) rC [ H ] = ∑ A U [ C ] [ A , H ] ∑ H′ ∑ A U [ C ] [ A , H′ ] . ( 5 ) We refer the reader to DuSell & Chiang ( 2020 ) for details regarding the derivation of these equations from Lang ’ s algorithm . To avoid underflow and overflow , in practice , Δ , W , and U are computed in log space . The model ’ s time complexity is $ ( | & |4 |Γ |3=3 ) , and its space complexity is $ ( | & |2 |Γ |2=2 ) .

Summary: This work continues a recent work on nondeterministic stack RNNs (in which an RNN controls a nondeterministic pushdown automaton, which I will refer to here as an N-PDA), proposing two changes to the architecture which create an RNS-RNN (renormalised NS-RNN) as follows: 1. Where in the NS-RNN the RNN-controller outputs a set of distributions over the next possible PDA actions (one distribution for every current state and stack-top pair), the RNS-RNN outputs scores over these actions. This is achieved simply by not softmaxing the scores that are anyway computed in the NS-RNN. The goal/effect is that the score of each (current state,current stack) pair being maintained in the N-PDA does not tend to zero as the input sequence length increases (as it would when probabilities are used), and in fact can even be increased at certain transitions. 2. Where in the NS-RNN the RNN-controller receives as part of its input at each time step a distribution over the current stack top token, it now also receives a distribution over the current PDA state. (These are normalised back to distributions here, normalised from the scores kept in the N-PDA.)  The authors evaluate the RNS-RNN against the NS-RNN and several other architectures, showing its success on various formal languages and even a greater success on a "hardest" context free grammar.   The NS-RNN requires quadratic space and cubic time in the input sequence length*, which is impractical for natural language processing. The authors propose a memory-limited version which reduces the complexity to linear time and space in the input sequence length, and discuss truncated BPTT for this version. They evaluate a memory limited RNS-RNN on the Penn Treebank.   *Because it is maintaining the probabilities of all possible configurations of non-deterministic PDA on the current input---in fact we are lucky this is not exponential! A dynamic algorithm by Lang is used to maintain the PDA, I am not sure of the details