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: The paper proposes a new stack-augmented RNN, RNS-RNN that includes two modifications to the Nondeterministic Stack RNN: 1. RNS-RNN uses unnormalized transition weights to avoid probability vanishing problems, such that gradient can be easier backpropagated to certain decisions at a previous time step. 2. RNS-RNN feeds the hidden state y into the controller while reading the stack, this allows the controller has more complete information to make decisions.