Adaptive Control Flow in Transformers Improves Systematic Generalization

1 INTRODUCTION . Neural networks ( NNs ) may easily learn certain training sets , but typically they do not generalize on systematically different test sets . Examples of systematic generalization ( Fodor et al. , 1988 ) include generalization to sequences longer than those seen during training—productivity , and algorithmic combinations of previously learned rules—systematicity . Despite recent efforts ( Bahdanau et al. , 2019 ; Korrel et al. , 2019 ; Lake , 2019 ; Li et al. , 2019 ; Russin et al. , 2019 ; Csordás et al. , 2021 ) , systematic generalization generally remains unsolved ( Fodor & McLaughlin , 1990 ; Lake & Baroni , 2018 ; Liska et al. , 2018 ; Greff et al. , 2020 ; Hupkes et al. , 2020 ) . On some datasets , the best performing models are neuro-symbolic hybrids ( Chen et al. , 2020 ; Liu et al. , 2020 ) using task-specific symbolic functions . However , their applicability to other datasets remains limited ( Furrer et al. , 2020 ; Shaw et al. , 2020 ) . A big question is : which type of architectural inductive bias encourages the training process to select “ good ” solutions which generalize systematically ? The popular Transformers ( Vaswani et al. , 2017 ) also often fail to generalize on algorithmic tasks ( e.g . Liska et al . ( 2018 ) ; Dubois et al . ( 2020 ) ; Chaabouni et al . ( 2021 ) ; Csordás et al . ( 2021 ) ; Ontañón et al . ( 2021 ) ) , even on tasks with intuitive solutions that can be simply expressed in terms of Transformer attention patterns . Given an input sequence of length N and a Transformer encoder of depth T , solving an algorithmic task is often all about routing the relevant information to the right node/operation at the right time in the T -by-N grid represented by Transformer columns . Effectively the task is to learn to draw an adaptive control flow on the canvas of Transformer columns . In fact , recent work by Weiss et al . ( 2021 ) introduced a programming language called RASP , which is specifically designed to express solutions to sequence processing problems , and which has a direct equivalent to the operations in Transformer encoders . However , it is shown that Transformers learn solutions expressed in RASP only through intermediate supervision of attention patterns , and sometimes , even such supervision fails . Generally speaking , Transformers fail to find easily interpretable and/or symbolic solutions to algorithmic tasks . We conversely hypothesize that attention-based NNs that are able to find intuitive solutions ( achieving interpretable attention patterns ) could improve systematic generalization . 1https : //github.com/robertcsordas/ndr Here we point out that regular Transformers lack some basic ingredients for learning such “ intuitive ” solutions to algorithmic problems . As a remedy , we propose simple architectural modifications to help them learn data routing . As a first step towards validating our model , we focus on the popular length generalization task of compositional table lookup ( CTL ; Liska et al . ( 2018 ) ; Hupkes et al . ( 2019 ) ; Dubois et al . ( 2020 ) ) , as well as two more complex tasks : a simple arithmetic task and a variant of ListOps ( Nangia & Bowman , 2018 ) designed to test the compositional generalization ability of NNs . Our novel Neural Data Router ( NDR ) achieves 100 % generalization accuracy ( never reported before ; Dubois et al . ( 2020 ) ) on the CTL task , and obtains nearly perfect accuracy on both the proposed simple arithmetic and ListOps tasks . We show that the attention and gating patterns of NDR tend to be interpretable as plausible control flows . 2 IMPROVING TRANSFORMERS FOR LEARNING ADAPTIVE CONTROL FLOW . We argue that the following components are needed to build Transformers capable of learning adaptive control flow . First , composing known operations in an arbitrary order requires that all operations are available at every computational step . This can be easily achieved by sharing the weights of the layers , as is done in Universal Transformers ( Dehghani et al. , 2019 ) . Second , the network should be sufficiently deep , at least as deep as the deepest data dependency in the computational graph built from elementary operations ( e.g. , in the case of a parse tree , this is the depth of the tree ) . Otherwise , multiple operations must be fused into a single layer and hinder natural and elegant compositions . Third , inputs in some columns should be kept unchanged until it is their turn to be processed . The regular Transformer lacks a mechanism for skipping the whole transformation step by simply copying the input to the next step/layer . We propose a special gating function , copy gate , to implement such a mechanism ( Sec . 2.1 ) . Finally , many algorithmic tasks require combining several local computations in the right order . This typically implies that attention should not focus on all possible matches at a given time but only on the closest match . We propose and investigate a new type of attention with a corresponding inductive bias called geometric attention ( Sec . 2.2 ) . Using both the geometric attention and copy gate , our model implements a “ neural data routing mechanism ” , which can adaptively serialize the input problem . We refer to the resulting new Transformer as Neural Data Router ( NDR ) . In the experimental section ( Sec . 3 ) , we evaluate this model on three algorithmic tasks requiring length generalization and demonstrate its effectiveness . 2.1 COPY GATE : LEARNING TO SKIP OPERATIONS ( VERTICAL FLOW ) . Each layer of the regular Transformer consists of one self-attention and one feedforward block . The input to each of these blocks is directly connected to the corresponding output via a residual connection ( Srivastava et al. , 2015 ; He et al. , 2016 ) . However , such a connection does not allow for skipping the transformation of the entire layer and simply passing the unchanged input to the next layer . Here we propose to add an explicit gate , which we call copy gate , to facilitate such a behavior . We consider a T -layer Transformer encoder and an input sequence of length N . Since each layer corresponds to one computational step , we often refer to a layer as a step t. We denote the Transformer state of column i in layer t as h ( i , t ) = Ht , i ∈ Rd where d is the state size , and Ht ∈ RN×d denotes the states of all N columns in layer t. In the copy gate-augmented Transformer , each column i in layer ( t+ 1 ) processes the input Ht similarly to regular Transformers : a ( i , t+1 ) = LayerNorm ( MultiHeadAttention ( h ( i , t ) , Ht , Ht ) + h ( i , t ) ) ( 1 ) ĥ ( i , t+1 ) = LayerNorm ( FFNdata ( a ( i , t+1 ) ) ) ( 2 ) using the standard multi-head attention operation ( Vaswani et al. , 2017 ) MultiHeadAttention with a query obtained from h ( i , t ) and keys/values from Ht , but the output is gated ( using g ( i , t+1 ) ∈ Rd ) as : g ( i , t+1 ) = σ ( FFNgate ( a ( i , t+1 ) ) ) ( 3 ) h ( i , t+1 ) = g ( i , t+1 ) ĥ ( i , t+1 ) + ( 1− g ( i , t+1 ) ) h ( i , t ) ( 4 ) We use the basic two-layer feedforward block ( Vaswani et al. , 2017 ) for both FFNdata and FFNgate which transforms input x ∈ Rd to : FFN ( x ) = W2 max ( W1x+ b1 , 0 ) + b2 ( 5 ) but with separate parameters and different dimensionalities : for FFNdata W data1 ∈ RdFF×d , W data2 ∈ Rd×dFF , while for FFNgate W gate1 , W gate 2 ∈ Rd×d , with biases bdata1 ∈ RdFF and bdata2 , b gate 1 , b gate 2 ∈ Rd . When the gate is closed i.e . g ( i , t+1 ) = 0 in Eq . 4 , the entire transformation is skipped and the input is copied over to the next layer h ( i , t+1 ) = h ( i , t ) . Crucially , we parameterize the gate ( Eq . 3 ) as a function of the output of the self-attention ( Eq . 1 ) , such that the decision to copy or transform the input for each column depends on the states of all columns . This is a crucial difference compared to previously proposed gatings in Transformers , which are solely motivated by training stability ( Parisotto et al. , 2020 ) or by a common practice from convolution-based models ( Chaabouni et al. , 2021 ) . None of the previous approaches can implement the behavior of our copy gate ( see Sec . 6 on related work ) . The bias of the gate bgate2 is initialized to −3 ( Hochreiter & Schmidhuber , 1997 ) . This ensures that no update happens initially to create a better gradient flow between layers . It also encourages the model to skip layers unless they have an important contribution in the corresponding step . 2.2 GEOMETRIC ATTENTION : LEARNING TO ATTEND TO THE CLOSEST MATCH ( HORIZONTAL FLOW ) . We propose geometric attention designed to attend to the closest matching element . Like in regular self-attention , given an input sequence [ x ( 1 ) , x ( 2 ) , ... , x ( N ) ] with x ( i ) ∈ Rdin , each input is projected to key k ( i ) ∈ Rdkey , value v ( i ) ∈ Rdvalue , query q ( i ) ∈ Rdkey vectors , and the dot product is computed for each key/query combination . In our geometric attention , the dot product is followed by a sigmoid function to obtain a score between 0 and 1 : Pi , j = σ ( k ( j ) > q ( i ) ) ( 6 ) which will be treated as a probability of the key at ( source ) position j matching the query at ( target ) position i . These probabilities are finally converted to the attention scores Ai , j as follows : Ai , j = Pi , j ∏ k∈Si , j ( 1− Pi , k ) ( 7 ) where Si , j denotes the set of all ( source ) indices which are closer to i than j is to i , and when two indices have the same distance to i , we consider the one which is to the right of i ( i.e. , greater than i ) to be closer , i.e. , Si , j = { k ∈ { 1 , ... , N } \ { i , j } : |i− k| ≤ |i− j| , if i < j k ∈ { 1 , ... , N } \ { i , j } : |i− k| < |i− j| , if j < i ( 8 ) In addition , we explicitly zero out the diagonal by setting Ai , i = 0 for all i = 1 , ... , N . The ordering of source indices is illustrated in Figure 1/Right . The resulting scores Ai , j are the attention scores used to compute the weighted averages of the value vectors . By using the terms ( 1− Pi , k ) in Eq . 7 , when there is a match , it downscales any other more distant matches . Two recent works ( Brooks et al. , 2021 ; Banino et al. , 2021 ) use such a parameterized geometric distribution in the form of Eq . 7 ( see Sec . 6 on related work ) . The resulting attention function has a complexity of O ( N2 ) , similar to the regular self-attention used in Transformers ( Vaswani et al. , 2017 ) . Eq . 7 can be implemented in a numerically stable way in log space . The products can then be calculated using cumulative sums , subtracting the elements for the correct indices in each position . Directional encoding . In practice , we augment Eq . 6 with an additional directional encoding . In fact , the only positional information available in the geometric attention presented above is the ordering used to define the product in Eqs . 7-8 . In practice , we found it crucial to augment the score computation of Eq . 6 with additional directional information , encoded as a scalar Di , j ∈ R for each target/source position pair ( i , j ) : Di , j = { WLRh ( i ) + bLR , if i ≤ j WRLh ( i ) + bRL , if i > j ( 9 ) where h ( i ) ∈ Rd denotes the input/state at position i and WLR , WRL ∈ R1×d , bLR , bRL ∈ R are trainable parameters . This directional information is integrated into the score computation of Eq . 6 as follows ( akin to how Dai et al . ( 2019 ) introduce the relative positional encoding ( Schmidhuber , 1992 ) as an extra term in the computation of attention scores ) : Pi , j = σ ( α ( Wqh ( i ) + u ) > Wk , Eh ( j ) + βDi , j + γ ) ( 10 ) where the matrix Wq ∈ Rdhead×d maps the states to queries , u ∈ Rdhead is a bias for queries , Wk , E ∈ Rdhead×d maps states to keys ( we note that dhead is typically the size of the key , query and value vectors for each head , dhead = dnheads ) , and α , β , γ ∈ R are learned scaling coefficients and bias , initialized to α = 1√ dhead , β = 1 , γ = 0 . Using this additional directional information , each query ( position i ) can potentially learn to restrict its attention to either the left or right side .

Summary:  The authors propose Transformer Control Flow (TCF), a set of improvements to the Universal Transformer (Dehghani et al, ICLR 2019). They show that, for three compositional problems, TCF allows trained models to generalize to longer sequences, a common problem of many transformer implementations.  As in the Universal Transformer (UT), the encoder consists of one shared transformer layer (self attention + fully connected network) which is iterated through a fixed number of times, by feeding the output of each iteration back into the input of the shared layer. However, whereas the UT uses a sequence to sequence model, TCF is an encoder-only architecture, which decodes the last element in the output sequence as the final result.   Two new features are introduced :  - a gating mechanism that allows the model to "skip a layer" (the input is then copied to the output), on the basis of the self-attention output, - a weighting system for the outputs of attention heads, which favors short-range attention (i.e. tokens close to the one currently considered), and can be trained to be biased towards a certain direction (before or after the current token).  Experiments are conducted over three tasks: - predicting the output of sequences of permutations of 8 elements, in prefix or postfix notation, - predicting the result of of additions and multiplications modulo 10, in infix notation, - predicting the result of operations on lists of small integers, in prefix notation.  For each task, TCF is shown to be capable of extrapolation to larger problems (i.e. longer sequences) than those seen at training.