Self-Instantiated Recurrent Units with Dynamic Soft Recursion

1 Introduction . Models based on the notion of recurrence have enjoyed pervasive impact across various applications . In particular , most effective recurrent neural networks ( RNNs ) operate with gating functions . Such gating functions not only ameliorate vanishing gradient issues when modeling and capturing long-range dependencies , but also benefit from fine-grained control over temporal composition for sequences [ Hochreiter and Schmidhuber , 1997 , Cho et al. , 2014 ] . With diverse and even growing data modalities ( e.g. , logic , algorithmic input and output , music , code , images , and language ) that can be expressed in sequences and may benefit from more architectural flexibility , recurrent neural networks that only explicitly impose a chain structure on such data but lack an explicit bias towards recursive self-instantiation may be limiting . For example , their gating functions are typically static across the temporal dimension of sequences . In view of such , this paper aims at studying an inductive bias towards recursive self-instantiation where the extent of recursion is dynamic at different time steps . We propose a novel recurrent unit whose gating functions are repeatedly controlled by instances of the recurrent unit itself . Our proposed model is called the self-instantiated recurrent unit ( Self-IRU ) , where self-instantiation indicates modeling via the own instances of the model itself in a recursive fashion . Specifically , two gates of the Self-IRU are controlled by Self-IRU instances . Biologically , ∗Work was done at NTU . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . this design is motivated by the prefrontal cortex/basal ganglia working memory indirection [ Kriete et al. , 2013 ] . For example , a child Self-IRU instance drives the gating for outputting from its parent Self-IRU instance . Our proposed Self-IRU is also characterized by the dynamically controlled recursion depths . Specifically , we design a dynamic soft recursion mechanism , which softly learns the depth of recursive self-instantiation on a per-time-step basis . More concretely , certain gates are reserved to control the extent of the Self-IRU recursion . Since values of these gates are between 0 and 1 and may vary across the temporal dimension , they make dynamic soft recursion depth at each time step possible , which could lead to more architectural flexibility across diverse data modalities . This design of the Self-IRU is mainly inspired by the adaptive computation time ( ACT ) [ Graves , 2016 ] that learns the number of computational steps between an input and an output and recursive neural networks that operate on directed acyclic graphs . On one hand , the Self-IRU is reminiscent of the ACT , albeit operated at the parameter level . While seemingly similar , the Self-IRU and ACT are very different in the context of what the objective is . Specifically , the goal of the Self-IRU is to dynamically expand the parameters of the model , not dynamically decide how long to deliberate on input tokens in a sequence . On the other hand , the Self-IRU marries the benefit of recursive reasoning with recurrent models . However , in contrast to recursive neural networks , the Self-IRU is neither concerned with syntax-guided composition [ Tai et al. , 2015 , Socher et al. , 2013 , Dyer et al. , 2016 , Wang and Pan , 2020 ] nor unsupervised grammar induction [ Shen et al. , 2017 , Choi et al. , 2018 , Yogatama et al. , 2016 , Havrylov et al. , 2019 ] . Our Contributions All in all , sequences are fundamentally native to the world , so the design of effective inductive biases for data in this form has far-reaching benefits across a diverse range of real-world applications . Our main contributions are outlined below : • We propose the self-instantiated recurrent unit ( Self-IRU ) . It is distinctly characterized by a novel inductive bias towards modeling via the own instances of the unit itself in a recursive fashion , where the extent of recursion is dynamically learned across the temporal dimension of sequences . • We evaluate the Self-IRU on a wide spectrum of sequence modeling tasks across multiple modalities : logical inference , sorting , tree traversal , music modeling , semantic parsing , code generation , and pixel-wise sequential image classification . Overall , the empirical results demonstrate architectural flexibility and effectiveness of the Self-IRU . For example , the Self-IRU achieves state-of-the-art performance on the logical inference dataset [ Bowman et al. , 2014 ] even when comparing with competitive models that have access to ground-truth syntactic information . Notation For readability , all vectors and matrices are denoted by lowercase and uppercase bold letters such as x and X , respectively . When a scalar is added to a vector , the addition is applied element-wise [ Zhang et al. , 2021 ] . 2 Method . This section introduces the proposed Self-IRU . The Self-IRU is fundamentally a recurrent model , but distinguishes itself in that the gating functions that control compositions over time are recursively modeled by instances of the Self-IRU itself , where the extent of recursion is dynamic . In the following , we begin with the model architecture that can recursively self-instantiate . Then we detail its key components such as how dynamic soft recursion is enabled . 2.1 Self-Instantiation . Given an input sequence of tokens x1 , . . . , xT , the Self-IRU transforms them into hidden states throughout all the time steps : h1 , . . . , hT . Denoting by L the user-specified maximum recursion depth , the hidden state at time step t is ht = Self-IRU ( L ) ( xt , h ( L ) t−1 ) , where Self-IRU ( L ) is an instance of the Self-IRU model at recursion depth L and h ( L ) t−1 is a hidden state at time step t− 1 and recursion depth L. In general , a Self-IRU instance at any recursion depth 0 ≤ l ≤ L returns a hidden state for that depth : Self-IRU ( l ) ( xt , h ( l ) t−1 ) = h ( l ) t , which involves the following computation : f ( l ) t = σ ( α ( n ) t Self-IRU ( l−1 ) ( xt , h ( l−1 ) t−1 ) + ( 1− α ( n ) t ) F ( l ) f ( xt , h ( l ) t−1 ) ) ( 2.1 ) o ( l ) t = σ ( β ( n ) t Self-IRU ( l−1 ) ( xt , h ( l−1 ) t−1 ) + ( 1− β ( n ) t ) F ( l ) o ( xt , h ( l ) t−1 ) ) ( 2.2 ) z ( l ) t = tanh ( F ( l ) z ( xt , h ( l ) t−1 ) ) ( 2.3 ) c ( l ) t = f ( l ) t c ( l ) t−1 + ( 1− f ( l ) t ) z ( l ) t ( 2.4 ) h ( l ) t = o ( l ) t c ( l ) t + xt , ( 2.5 ) where denotes the element-wise multiplication , σ denotes the sigmoid function , scalars α ( n ) t and β ( n ) t are soft depth gates at time step t and recursion node n in the unrolled recursion paths , and F ( l ) f , F ( l ) o , and F ( l ) z are base transformations at recursion depth l. Without losing sight of the big picture , we will provide more details of such soft depth gates and base transformations later . On a high level , Figure 1 depicts the Self-IRU model architecture . We highlight that two gating functions of a Self-IRU , the forget gate f ( l ) t in ( 2.1 ) and the output gate o ( l ) t in ( 2.2 ) , are recursively controlled by instances of the Self-IRU itself . Therefore , we call both the forget and output gates the self-instantiation gates . The base case ( l = 0 ) for self-instantiation gates is f ( 0 ) t = σ ( F ( 0 ) f ( xt , h ( 0 ) t−1 ) ) and o ( 0 ) t = σ ( F ( 0 ) o ( xt , h ( 0 ) t−1 ) ) . At each recursion depth l , the candidate memory cell z ( l ) t at time step t is computed in ( 2.3 ) . Then in ( 2.4 ) , the forget gate f ( l ) t controls the information flow from z ( l ) t and the memory cell c ( l ) t−1 at the previous time step to produce the memory cell c ( l ) t at the current time step t. As illustrated by the bottom arrow starting from xt in Figure 1 , the output gating in ( 2.5 ) also adds a skip connection from residual networks to facilitate gradient flow throughout the recursive self-instantiation of the Self-IRU [ He et al. , 2016 ] . 2.2 Dynamic Soft Recursion . Now let us detail the soft depth gates α ( n ) t and β ( n ) t in ( 2.1 ) and ( 2.2 ) for time step t and recursion node n in the unrolled recursion paths . The index n is used to distinguish nodes at different positions in the recursion tree ( e.g. , in Figure 2 ) that is determined by the maximum recursion depth L. We propose learning them in a data-driven fashion . Specifically , we parameterize α ( n ) t and β ( n ) t with α ( n ) t = σ ( F ( n ) α ( xt ) ) and β ( n ) t = σ ( F ( n ) β ( xt ) ) , where F ( n ) ∗ ( xt ) = W ( n ) ∗ xt + b ( n ) ∗ ∈ R ( ∗ ∈ { α , β } ) with weight parameters W ( n ) ∗ and bias parameters b ( n ) ∗ both learned from data . Together with the sigmoid function σ , these simple linear transformations of the input token xt are applied dynamically at each time step t across the input sequence . Moreover , as shown in ( 2.1 ) and ( 2.2 ) , mathematically 0 < α ( n ) t , β ( n ) t < 1 control the extent of recursion at each recursion node n , enabling soft depth along any unrolled recursive self-instantiation path . Thus , α ( n ) t and β ( n ) t are called the soft depth gates . Putting these together , Figure 2 unrolls the recursive self-instantiation paths at two consecutive time steps to illustrate dynamic soft recursion depth . Specifically , the “ softness ” is indicated by greyscale of any node at the beginning of an arrow along an unrolled recursion path . In sharp contrast to multi-layer RNNs , Self-IRUs enable tree structures of self-instantiation , where the extent of recursion is dynamic ( to be visualized in Section 3.7 ) . 2.3 Base Transformations . At any recursion depth l , F ( l ) f in ( 2.1 ) , F ( l ) o in ( 2.2 ) , and F ( l ) z in ( 2.3 ) are base transformations of the input xt and the hidden state h ( l ) t−1 . For example , we can model base transformations using RNN units ( e.g. , LSTM ) : at recursion depth l , for ∗ ∈ { f , o , z } we have F ( l ) ∗ ( xt , h ( l ) t−1 ) = RNN ( l ) ∗ ( xt , h ( l ) t−1 ) . Alternatively , we may also model base transformations with linear layers that only transform the input xt using learnable weight parameters W ( l ) ∗ and bias parameters b ( l ) ∗ for ∗ ∈ { f , o , z } : F ( l ) ∗ ( xt ) = W ( l ) ∗ xt + b ( l ) ∗ . The Self-IRU is agnostic to the choice of base transformations and we will evaluate different choices in the experiments . We will discuss how the Self-IRU can be useful as a ( parallel ) non-autoregressive model and connects to other recurrent models in the supplementary material .

Summary: The paper proposes a new type of recurrent unit, analogous to GRU and LSTM, dubbed Self-IRU. The underlying update is similar to that used by LSTMs and GRUs. However, the authors suggest to additionally incorporate the output of a different Self-IRU in the forget and output gates, similar to stacking LSTMs or GRUs in multiple layers. The authors propose that this allows the model to learn a "dynamic soft recursion". The paper performs empirical experiments on pixel-wise image classification, a synthetic logical inference task, sorting, tree traversal, music modeling, and semantic parsing tasks.