THE NEURAL DATA ROUTER: ADAPTIVE CONTROL FLOW IN TRANSFORMERS IMPROVES SYSTEMATIC GENERALIZATION
Year: 2022

Abstract:
Despite progress across a broad range of applications, Transformers have limited success in systematic generalization. The situation is especially frustrating in the case of algorithmic tasks, where they often fail to find intuitive solutions that route relevant information to the right node/operation at the right time in the grid represented by Transformer columns. To facilitate the learning of useful control flow, we propose two modifications to the Transformer architecture, copy gate and geometric attention. Our novel Neural Data Router (NDR) achieves 100% length generalization accuracy on the classic compositional table lookup task, as well as near-perfect accuracy on the simple arithmetic task and a new variant of ListOps testing for generalization across computational depths. NDR’s attention and gating patterns tend to be interpretable as an intuitive form of neural routing. Our code is public.1

heading: 1 INTRODUCTION
text: 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.

heading: 2 IMPROVING TRANSFORMERS FOR LEARNING ADAPTIVE CONTROL FLOW
text: 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.

heading: 2.1 COPY GATE: LEARNING TO SKIP OPERATIONS (VERTICAL FLOW)
text: 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.

heading: 2.2 GEOMETRIC ATTENTION: LEARNING TO ATTEND TO THE CLOSEST MATCH (HORIZONTAL FLOW)
text: 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.

heading: 3 EXPERIMENTS
text: We evaluate the proposed methods on three tasks: the compositional table lookup (Liska et al., 2018; Hupkes et al., 2019), a custom variant of ListOps (Nangia & Bowman, 2018), and a simple arithmetic task which we propose. In all cases, the task is designed to test the compositional generalization ability of NNs: the model has to learn to apply operations seen during training in a longer/deeper compositional way (productivity). Further experimental details for each task can be found in Appendix C.

heading: 3.1 COMPOSITIONAL TABLE LOOKUP
text: Task. The compositional table lookup task (Liska et al., 2018; Hupkes et al., 2019; Dubois et al., 2020) is constructed based on a set of symbols and unary functions defined over these symbols. Each example in the task is defined by one input symbol and a list of functions to be applied sequentially, i.e., the first function is applied to the input symbol and the resulting output becomes the input to the second function, and so forth. There are eight possible symbols. Each symbol is traditionally represented by a 3-bit bitstring (Liska et al., 2018). However, in practice, they are simply processed as one token (Dubois et al., 2020). The functions are bijective and randomly generated. Each function is represented by a letter. An example input is ‘101 d a b’, which corresponds to the expression b(a(d(101))); the model has to predict the correct output symbol. We note that there exists a sequenceto-sequence variant of this task (Dubois et al., 2020) where the model has to predict all intermediate steps (thus trained with intermediate supervision). We directly predict the final output. An ideal model should be able to solve this task independently of the presentation order, that is, it should not matter whether the task is encoded as ‘101 d a b’ or ‘b a d 101’. We thus study both forward (former) and backward (latter) variants of the task. To evaluate systematic generalization, the train/valid/test sets reflect different numbers of compositions: samples with 1-5/6-8/9-10 operations, respectively. To best of our knowledge, no previous work has reported perfect accuracy on this task through an NN. We refer the readers to Sec. 6 for further details on the previous work.
Results. We consider five different baselines: an LSTM (Hochreiter & Schmidhuber, 1997), bidirectional LSTM (Schuster & Paliwal, 1997), DNC (Graves et al., 2016; Csordás & Schmidhuber, 2019), Universal Transformers (Vaswani et al., 2017; Dehghani et al., 2019), and its relative position variants (Csordás et al., 2021). For Transformers, the prediction is based on the last column in the final layer (we conduct an ablation study on this choice in Appendix A). The hyper-parameters used for each model can be found in Table 7 in the appendix. We also provide an ablation study on the number of layers needed for generalization in Appendix A, which supports our claim on the necessity for a “sufficiently” deep architecture. The main results on this task are shown in Table 1. The LSTM
and DNC perform well in the forward variant, achieving perfect generalization for longer sequences, but fail on the backward variant. This is not surprising since in the forward case, input symbols are presented in the “right” processing order to the LSTM. As expected, the bidirectional LSTM performs well in both presentation orders, since one of its processing directions is always aligned with the order of computation. However, for an arbitrary task, the order of processing is not given. For example, for ListOps (Sec. 3.3), the processing should start from the deepest point in the parse tree, which is probably somewhere in the middle of the sequence. The experiments on other tasks (Sec. 3.2 and 3.3) requiring arbitrary processing orders show that bidirectional LSTMs do not generalize well in such tasks. This is not satisfactory since our goal is to create a generic architecture which can solve arbitrary problems with an arbitrary underlying input processing order. While the Transformer seems to be a good candidate for learning problem dependent processing orders, the baseline Transformer variants fail to generalize in this task in both directions.
By introducing the copy gate (Sec. 2.1), the relative Transformer can solve the forward task, but not the backward one. Our analysis showed that the network learns to attend to the last operation based on the relative position information. Since the result is read from the last column, this position changes with the sequence length. The model thus fails to generalize to such arbitrary offsets. To address this issue, we introduce a simple mechanism to let the model choose between absolute and relative positional encodings at each position (see Appendix B). The resulting model effectively manages to use the absolute position for the prediction and perform well in both directions. However, such a combination of absolute/relative positional encoding might be an overly specific bias. A more generic solution, geometric attention (Sec. 2.2), also achieved perfect generalization and was found easier to train. We present the corresponding visualization of our model in Sec. 4.

heading: 3.2 SIMPLE ARITHMETIC
text: In order to validate the success of the proposed model on a task that involves more complex data flows and operations, we propose the simple arithmetic task.
Task. The task is to execute an arithmetic expression consisting of nested modulo 10 additions and multiplications. This requires the model to process tree-structured data flows, which is presumably more difficult than the sequential processing required for the CTL task. Each operation is surrounded by brackets, such that the boundaries of operations are easy to determine. For example ‘((4*7)+2)’ should evaluate to ‘0’ (30 modulo 10). The expressions are generated randomly. The tree depth is up to 5 for the training set, 6 for the validation set, and 7-8 for the test set. The depth is measured as the number of operations, ignoring the leaves, so the example above has a depth of 2. The sequence length is limited to at most 50 tokens.
Results. Table 2 shows the results. All considered models perform well on the IID validation data, but none except the NDR performs well on the generalization test set, which achieves near-perfect accuracy of 98%. We also note that the NDR learns very quickly: while all other models require about 200 K steps to converge, the NDR achieves near-perfect accuracy after 50 K steps of training.

heading: 3.3 LISTOPS
text: We also evaluate our model on a variant of the ListOps task (Nangia & Bowman, 2018) which is a popular task commonly used to evaluate parsing abilities of NNs (Havrylov et al., 2019; Shen et al., 2019; Xiong et al., 2021; Tay et al., 2021; Irie et al., 2021). Some special architectures such as Chowdhury & Caragea (2021) can almost perfectly generalize to longer sequences on this task. However, as far as we know, no Transformer variant has been reported to be fully successful.
Task. The task consists of executing nested list operations written in prefix notation. All operations have a list of arguments that can be either a digit (from 0 to 9) or recursively another operation with its own list of arguments. The operations are min, max, median and sum. The sum is modulo 10, and the median is followed by the floor function such that the output of any operation lies between 0 and 9. For example: [MED 4 8 5 [MAX 8 4 9 ] ] should return 6. There are two well-known variants: the original one by Nangia & Bowman (2018) and the “Long Range Arena” variant by Tay et al. (2021) which have different maximum numbers of arguments in each function and maximum sequence lengths. In both variants, there is no strict control of the depth of data samples: there is simply a certain pre-defined probability that each argument in the list is expanded into another list (which may increase the tree depth). This is not suitable for evaluating systematic generalization in terms of compositionality (over the problem depth). We propose instead to generate clean train, valid, and test splits with disjoint depths: up to depth 5 for training, depth 6 for validation and depths 7 and 8 for test. Importantly, we make sure that a depth-K sample effectively requires computation until depth-K (otherwise min, max, and med operations could potentially find the output without executing all of its arguments). By dissociating the splits by the depth, we can clearly identify models which fail to generalize compositionally. Apart from the depth specifications, all train/valid/test sets share the same settings as follows: the maximum sequence length is 50 (tokens), the probability of recursively sampling another function inside a list is 30% at each position, and the maximum number of arguments for a function is 5. The train set consists of 1M, the valid and test sets of 1K sequences.
Results. Table 3 shows the results. Like on the other tasks, the baseline LSTM and Transformers do not generalize well on the test set consisting of deeper problems, while they achieve a near-perfect accuracy on IID data. In contrast, our model achieves near-perfect generalization.

heading: 4 ANALYSIS
text: In this section, we provide some visualizations of attention and gating patterns of the NDR and the corresponding analyses. For more visualizations, we refer the readers to Appendix D.
Compositional Table Lookup. Figure 2 shows the gating and attention patterns of the NDR model for an example of the backward presentation task. As shown in Fig. 2/Bottom, the gates of different columns open sequentially one after another when the input is available for them. Fig. 2/Top shows the corresponding attention maps. Each column attends to the neighbouring one, waiting for its computation to be finished. The behavior of the last column is different: it always attends to the second position of the sequence, which corresponds to the last operation to be performed.
ListOps. We can also identify how the NDR processes the data in ListOps. Different attention heads play different roles. We highlight the core observations in Figure 3. The input for this example is: [SM [MED [MIN 1 7 4 [MAX 2 4 0 8 9 ] ] 7 ] 5 [MED 8 5 8 ] 0 7 ]. First of all, we find that there is a head (head 13 in Figure 3, first row) which seems to be responsible for connecting operators and their arguments: the operands/arguments of an operation attend to the operator. In step 0 (t = 0 in the figure), we can recognize that the operations at the deepest level, namely MAX and the second MED have all the arguments ready (as is shown by vertical lines on the columns corresponding to MAX and MED). The model indeed identifies that these two operations are ready to be executed and that they can be processed in parallel (these arguments-to-operation attention patterns remain for a few steps). We note that at this stage, the last argument of MIN is not ready yet ([MIN 1 7 4 [MAX 2 4 0 8 9 ] ]). We can see that only arguments which are already ready (1 7 4) attend to the operator (see the column of MIN). In step 1 (t = 1, 2nd row), we can see that head 5 copies the expected result of MAX, 9 to the column of the operator (we note that this only requires one step as 9 is always the result of MAX when it is one of the arguments of MAX). Similarly in step 2, head 7 (2nd row) seems to copy the result of the second MED, 8 to the operator column. In step 3 (t = 3, 1st row), we recognize that the result of MAX is marked as an argument for MIN in head 13 which is responsible for communication between operators and their arguments. This is shown by the new attention which appears at t = 3 in head 13 from the source position MAX to the target position MIN (a pattern which is not visible at t = 2). In head 3, t = 6 (2nd row), the expected result of MIN, which is 1, is copied to the operator, similarly to the patterns we observed above for MAX and MED. In head 13, t = 6 (1st row), all arguments for the first MED are now also recognized (the result of MIN which is 1, and 7). Finally in t = 7 (2nd row), two heads, head 3 and head 5 seem to copy/gather two inputs needed to compute the corresponding median, 1 and 7, and store them in
the column of the operator MED. A complete visualization of further steps can be found in Appendix D.2. We noticed that some of the heads do not seem to play a key role; we focused on interpreting those which seem to participate in the main computation. For ListOps, we also partially find the attention patterns described above in the baseline Transformer with relative positional encoding, at least on some inspected examples, which also explains its rather high accuracy.

heading: 5 DISCUSSION
text: Learning adaptive serialization. The NDR architecture can be understood as performing adaptive serialization of the problem. A key requirement for reusable computation is decomposing the problem into reusable building blocks, typically applied in sequential steps. The granularity of the decomposition determines the degree of reusability: fusing operations in a single step makes the processing faster (fewer steps), but also more specialized. Learning the most granular solutions is thus preferable for generalization. At the same time, not all processing should happen serially: branches of the computational graph that do not have common data dependencies can be processed independently in parallel, which we empirically observe in our NDR in the ListOps example (Sec. 4). This enables the architecture to get away with a number of computational steps reflecting the depth of the computational graph rather than the length of the input.
Bottom up approach for improving model architectures. Transformers have seen tremendous successes across various application domains (Devlin et al., 2019; Brown et al., 2020; Dosovitskiy et al., 2021). Impressive results have been reported when they are scaled up with a large amount of data (Brown et al., 2020). On the other hand, simple tasks like those highlighted in the present work demonstrate that the Transformer architecture still struggles with basic reasoning. Particularly in algorithmic tasks, it is often the case that a sub-optimal choice of architecture/optimization method makes the model fall back to simple memorization. We argue that it is crucial to look at isolated problems which test specific generalization capability. This calls for a bottom-up approach: building on toy tasks that focus on individual aspects of generalization and using them for improving models.

heading: 6 RELATED WORK
text: Gating inside Transformers. Several prior works have proposed to use some sort of gating within Transformer architectures (Parisotto et al., 2020; Chaabouni et al., 2021). Our proposed copy gate
is different from those as it satisfies two important properties. First, our copy gate allows the model to skip the entire Transformer layer (i.e., both the self-attention and the feedforward blocks) when the gate is closed. Second, the gate function is conditioned on the attention output such that the decision of opening or closing depends on information from all columns. While multiple gating variants have been proposed by Parisotto et al. (2020) to stabilize Transformers for reinforcement learning, none of them can produce this behavior. Empirically, we also tried out a few other gating variants which do not satisfy the two properties above; we found them not to improve over regular Transformers in our preliminary experiments on compositional table lookup. Recent work by Chaabouni et al. (2021) also makes use of “gating” in Transformers through a gated linear unit (GLU) activation function commonly used in convolutional NNs (Dauphin et al., 2017). Transformer models with such an activation function were reported to outperform RNN baselines on a systematic generalization task (Dessı̀ & Baroni, 2019). Unlike our copy gate or Parisotto et al. (2020)’s gating, such a gating activation does not have the “residual” term (i.e. a closed gate zeros out the input) which allows the model to skip a transformation. In a more general context, benefits of the GLU activation in Transformers vary across tasks (Irie et al., 2019; Shazeer, 2020). In language modeling, no improvement is typically obtained by using the standard highway gate instead of the residual connection in Transformers (Irie, 2020), while it yields improvements when combined with convolutional layers (Kim & Rush, 2016).
Parameterized geometric distributions. Two recent works (Brooks et al., 2021; Banino et al., 2021) have used a form of parameterized geometric distribution (PGD; in the form of Eq. 7). Brooks et al. (2021) have used such a distribution to parameterize the movement of a pointer on a sequence of instructions. Banino et al. (2021) have used it to implement adaptive computation time (Schmidhuber, 2012; Graves, 2016). We use the PGD to obtain a generic attention mechanism as a replacement of the standard self-attention used in Transformers (Vaswani et al., 2017).
Compositional table lookup. CTL task was proposed for evaluating the compositional ability of NNs (Liska et al., 2018). Previous works evaluated RNNs, RNNs with attention, and Transformers on this task with limited success (Hupkes et al., 2019; Dubois et al., 2020). Dubois et al. (2020) have proposed a special attention mechanism to augment the recurrent architecture. While they obtained good performance for the forward presentation order, the proposed model failed in the backward one. In contrast, two of our approaches (Sec. 3.1) achieve 100% generalization accuracy for both orders.
Positional encodings. Many previous works have focused on improving positional encoding (Schmidhuber, 1992; Vaswani et al., 2017) for self-attention. Most notably, the relative positional encoding (Schmidhuber, 1992; Shaw et al., 2018; Dai et al., 2019) was found useful for improving systematic generalization of Transformers (Csordás et al., 2021). Here we also present two new approaches related to positional encoding. One is the gated combination of absolute and relative positional encoding (Sec. 3.1; details in Appendix B). We show that absolute positional encoding can complement relative positional encoding. The former enables the model to always attend to a specific position, as is needed for the CTL task in the last step, while the gating allows it to use relative positional encoding for other positions/steps. Second, we introduce directional encoding to augment geometric attention. Unlike positional encoding which can overfit to a range of positions seen during training, the direction information is found to be robust and to be a crucial augmentation of the geometric attention.

heading: 7 CONCLUSION
text: We proposed a new view on the internal operations of Transformer encoders as a dynamic dataflow architecture between Transformer columns. This overcomes two shortcomings of traditional Transformers: the problem of routing and retaining data in an unaltered fashion, which we solve by an additional copy gate, and the problem of learning length-independent attention patterns, which we solve by geometric attention. Our new model, the Neural Data Router (NDR), generalizes to compositions longer than those seen during training on the popular compositional lookup table task in both forward and backward directions. NDR also achieves near perfect performance on simple arithmetic and ListOps tasks in settings that test systematic generalization in terms of computational depth. In general, the gates and the attention maps collectively make the architecture more interpretable than the baselines. Future work will extend this encoder-only architecture to a full sequence-to-sequence model and evaluate it on other standard tasks in systematic generalization requiring generation of variable-length output sequences.

heading: ACKNOWLEDGMENTS
text: We thank Imanol Schlag and Sjoerd van Steenkiste for helpful discussions and suggestions on an earlier version of the manuscript. This research was partially funded by ERC Advanced grant no: 742870, project AlgoRNN, and by Swiss National Science Foundation grant no: 200021 192356, project NEUSYM. We are thankful for hardware donations from NVIDIA & IBM. The resources used for the project were partially provided by Swiss National Supercomputing Centre (CSCS) project s1023.

heading: A ABLATIONS
text: IID Test
nlayers Forward Backward Forward Backward
14 1.00 ± 0.00 1.00 ± 0.00 1.00 ± 0.00 1.00 ± 0.00 12 1.00 ± 0.00 1.00 ± 0.00 1.00 ± 0.00 0.99 ± 0.02 10 1.00 ± 0.00 1.00 ± 0.00 0.75 ± 0.04 0.62 ± 0.05
8 1.00 ± 0.00 1.00 ± 0.00 0.23 ± 0.02 0.24 ± 0.03 6 1.00 ± 0.00 0.96 ± 0.03 0.22 ± 0.05 0.15 ± 0.01 4 0.96 ± 0.04 0.68 ± 0.11 0.14 ± 0.01 0.13 ± 0.01
Readout from the first instead of the last column. In our experiments with the Transformer models, the last column was used for the readout of the result. Under this configuration, the readout position depends on the length of the sequence which might increase the difficulty of the problem, in particular for the models using absolute positional embeddings. Table 5 shows the corresponding ablation study. We observe that this choice has only marginal impact on the model performance. As a side note, we also tried the variant where an additional cross-attention layer is used for the readout. Again, the generalization performance was not better. In fact, these results are not surprising since none of these changes fundamentally addresses the problem of length generalization.
Does Adaptive Computation Time (ACT) help? In this work, we determined the number of layers/steps to be used in the model based on heuristics (see Appendix C.1). We could also consider using Adaptive Computation Time (ACT) to dynamically determine the number of steps. Furthermore, ACT introduces a form of gating which creates shortcuts in the credit assignment path between the output and a result of an intermediate layer. This “copying” mechanism resulting from the ACT (i.e. stop computation at a certain time and copy the result to the output) is fundamentally different
from our copy gate (Sec. 2.1). Our copy gate allows Transformer columns to keep the input unchanged until it’s their turn to be processed (a crucial property to implement control flow like behavior). This behavior can not be simulated by the ACT. Here we provide some experimental results on models with ACT which confirm that the proposed copy gate is a crucial component for generalization which can not be replaced by ACT.
We note that there are various versions of ACT in the literature, e.g., the variant used by Dehghani et al. (2019) in Universal Transformers is different from the one used by Graves (2016). Here we focus on two variants: one in which we directly apply Graves (2016) to Transformers, and another one used by Dehghani et al. (2019). We start with the description of the former.
An extra sigmoidal unit p̂(i,t) is computed for each column i in each timestep t as:
p̂(i,t) = σ(WHh (i,t) + bH) (11)
where WH ∈ R1×d and bH ∈ R are trainable parameters. By comparing the cumulative sum of p̂(i,t) over time steps to a certain threshold value (1− ) with a hyper-parameter (0.01 in our experiment), we determine the termination step T i for column i as:
T i = min{Tmax,min{t′ : t′∑
t=1
p̂(i,t) ≥ 1− }} (12)
where Tmax is the pre-defined maximum number of steps.
The corresponding halting probability p(i,t) is then computed as:
p(i,t) =
{ p̂(i,t) if t < T i
Ri if t = T i (13)
Ri = 1− T i−1∑ t=1 p̂(i,t) (14)
which is used to compute the final output of column i as:
oi = T i∑ t=1 p(i,t)h(i,t) (15)
In Dehghani et al. (2019)’s variant, a different equation is used in lieu of Eq. 15 above and the computation of the reminder term Ri in Eq. 14 above is not properly handled in case where Eq. 12 terminates because of the first condition on Tmax. For further details, we refer the readers to Listing 1 and 2 in Dehghani et al. (2019) and/or our public code.
One subtlety introduced by Dehghani et al. (2019) which we note here is that the computation of the final output oi of column i effectively “halts” after T i (since oi only depends on h(i,t) for 0 < t < T i), but column i itself still continues transforming the hidden states h(i,t) for steps t > T i until all columns reach the termination step, and its updated states can be attended/read by another column j which has not halted yet (i.e. T j > T i). In this sense, computation is never stopped independently for each column. The mechanism described above instead finds the readout steps for each column (as is used in Eq. 15). We follow this decision in our implementation of both variants.
In addition, a new regularizer term, LACT = α 1N ∑N i=1R i is added to the loss function, where N is the length of the input sequence. This makes the network prefer short computations. We ran a hyper-parameter search for α from the following values: 0.001, 0.003, 0.01, 0.03, 0.1. We found α = 0.03 to work the best.
We conducted experiments on the compositional table lookup task. We first noted that ACT helps training our baseline Transformer models with a maximum step of 14 layers which was not possible without ACT (our baseline Transformer had only 11 layers for this reason; see Table 7). The shortcut in the credit assignment path introduced by ACT certainly helps training of this 14 layer model. As we noticed that the models with ACT learn slower than those with gating, we increased the number of training steps to 60k steps which is twice as many as 30k used for the models without ACT. Table
6 shows the results. We observe that, interestingly, ACT enables generalization for longer lengths in the forward direction of the Transformer with relative positional encoding and the one with geometric attention. However, we were not able to find any configuration that generalizes in the backward case. This demonstrates that the copy gate is effectively a crucial component for generalization which can not be replaced by ACT. Furthermore, the convergence of models with ACT is significantly slower than those of models with our gating, and they are more unstable and very sensitive to the value of α on the regularization term, even in the successful forward case. Overall, the only benefit of ACT is thus the adaptive depth, as is illustrated in Figure 4, which is orthogonal to our study.

heading: B DETAILS OF ATTENTION WITH COMBINED ABSOLUTE/RELATIVE POSITIONAL ENCODING
text: The use of copy gates enables Transformers to generalize to longer lengths in the forward presentation order of the CTL task (Sec. 3.1), but that alone was not enough to make the model generalize in the backward order variant of the task. Examining the attention maps reveals that the model uses position-based attention to read out the result instead of content-based attention. In the backward presentation order, the last column of the transformer should focus on the second column, whose relative position changes dynamically with the length of the sequence. We solve this issue by adding an option to choose between absolute and relative positional encodings to the attention head.
In what follows, we describe the operation within a single layer/step. This allows us to omit the layer/step-index t for better readability, and thus denote the state of column/position i as hi instead of h(i,t). We use the relative positional embedding variant of self-attention by Dai et al. (2019). Our attention matrix with the gated absolute/relative positional encodings can be decomposed as follows:
ri = σ(hiWar + bar) (16)
Âi,j = h > i W > q Wk,Ehj︸ ︷︷ ︸ (a) +u>Wk,Ehj︸ ︷︷ ︸ (c)
+ ( h>i W
> q Wk,P︸ ︷︷ ︸ (b) +v>Wk,P︸ ︷︷ ︸ (d)
) ( pi−jri + pj(1− ri)︸ ︷︷ ︸
(e)
) (17)
where the matrix Wq ∈ Rdhead×d maps the states to queries, Wk,E ∈ Rdhead×d maps states to keys, while Wk,P ∈ Rdhead×d maps positional embeddings to keys. dhead is the size of the key, query and value vectors for each head, set as dhead = dnhead . u,v ∈ R
dhead are learned vectors. pi ∈ Rd is the standard sinusoidal embedding for position i (Vaswani et al., 2017). Softmax is applied to the second dimension of Â to obtain the final attention scores, A. Component (a) corresponds to content-based addressing, (b, e) to content-based positional addressing, (c) represents a global content bias, while (d, e) represent a global position bias.
We introduce term (e) for the positional embedding which can switch between absolute and relative positional encodings using the scalar gate ri (Eq. 16; parameterized by War ∈ Rd×1 and bar ∈ R), which is the function of the state at target position i.
C IMPLEMENTATION DETAILS
A PyTorch implementation of our models together with the experimental setup is available under https://github.com/robertcsordas/ndr. The performance of all models is reported as mean and standard deviations over 5 different seeds.
C.1 CHOOSING THE NUMBER OF LAYERS
In Sec. 2, we hypothesized that one of the conditions for our model to generalize is to be “sufficiently” deep such that elementary operations are learned in separate layers which would then become composable. In practice, a “sufficient” depth can be determined by the basic units of compositions implicitly defined by the dataset. The depth of the model must be at least as deep as the deepest path in the computation graph defined by these basic operations. This hypothesis was empirically validated in the ablation study presented above (Appendix A). In general, we used the following heuristics to choose the depth of the Transformers:
(length of the deepest path in the graph) × (steps per operation) + a few more layers.
Determining the number of steps needed by the elementary operation is not straightforward but it can be done empirically. For example, for ListOps, as is shown in Sec. 4, it requires two steps per operation: one step in which the operands attend to the operation, followed by another one where the result is written back to the operation. For other tasks, we found that a single step per operation was enough. Choosing more layers than needed is safe, and it can be used to determine the required number of layers, for example by looking at the gate activity. Finally, “+ a few more layers” are needed because one additional layer should be used to read out the final result, and one or a few more can be needed for communication between columns (e.g., to determine operator precedence).
Since parameters are shared across layers, we can optionally train models with a certain number of layers and increase the number of computational steps at test time. This allows us to train models using a depth which is “sufficient” to solve the training set, but increase it at test time to generalize to a test set requiring more computational steps. We did this for the ListOps experiment (Sec. 3.3): the model was trained with 20 layers and tested with 24. Our preliminary experiments confirmed that this practice has no performance penalty, while it speeds up training.
C.2 DATASET DETAILS
Compositional table lookup. Our implementation uses 8 symbols as input arguments and 9 randomly sampled bijective functions denoted by lower case letters of the English alphabet. All functions are included in the train set in combination with all possible input symbols. The rest of the training set consists of random combinations of functions applied to a random symbol as an argument, up to length 5. The total size of the train set is 53,704 samples. The samples are roughly balanced such that there are similar numbers of samples for each depth. There are different validation sets: an IID set, which matches the distribution of the train set, and a depth validation, which includes samples of lengths 6, 7 and 8. The test set consists of sequences of lengths 9 and 10.
Simple arithmetic. The dataset is constructed by sampling random digits (0-9) and operations + (add) and ∗ (multiply). The operations are performed modulo 10. Parentheses surround the arguments of the operations. The depth of the resulting tree is computed, and rejection sampling is used to ensure that the same number of samples from each depth is present in the given split. The maximum length of samples is 50 tokens, sub-operations are sampled with probability 0.2. 100 K samples are used for training, 1 K for both test and validation sets. The train set consists of 0-5 operations, the validation set of 6 and the test set of 7 operations.
ListOps. Random digits are sampled from range 0-9. Operations are sample from the set summodulo (SM), which is a sum modulo 10, min (MIN), max (MAX) and median followed by the floor function (MED). The maximum number of arguments for each operation is 5. A sub-operation is sampled with probability 0.3. 1 M samples are used for training, 1 K for test and validation. The train set consists of 0-5 operations, 6 for the validation set, and 7 for the test set.
For each sample, we calculate a number which we call dependency depth. To understand it, note that MIN and MAX operations only select one of their operands, MED selects 1 or 2. In SUM, all operands are needed to perform the operation. If we construct a parse tree and prune away the branches which were not selected by any operation and measure the depth of such a tree, the resulting number is the dependency depth. This ensures that the deeper parts of the tree contribute to the result calculation, preventing shallow heuristics, like ignoring all branches of the tree that are too deep and still getting the correct result with a high chance. We also ensure that the number of samples is the same for all possible dependency depths in each split.
C.3 MODEL DETAILS
We use the AdamW optimizer (Loshchilov & Hutter, 2019) for all of our models. Standard hyperparameters are listed in Tab. 7, 8 and 9. Additionally, models with gating use dropout (Hanson, 1990; Srivastava et al., 2014) applied to the content-based query and the position-query components of 0.1 for most models, except for non-gated Transformers on ListOps, where this value is 0.05. In the case of geometric attention, since the channels of the directional encoding does not have any redundancy, dropout is applied just to the content-query.
In the case of Transformers with the copy gate but without geometric attention, we use tanh instead of LayerNorm in Eq. 2.
The hyperparameters of the gateless Transformers differ significantly from the gated ones. This is because they were very hard to train to achieve good performance even on the IID set, requiring extensive hyperparameter tuning. One might argue that fewer layers make them less competitive on longer sequences. However, we were unable to train them to perform well even on IID data with comparable sizes.
All Transformer variants have a begin (B) and end (E) token included in the sequence. RNNs (LSTM and DNC) have no such tokens. All Transformers are encoders only, and the results are read from the last column (corresponding to the end token).
The DNC has 21 memory cells, 4 read heads, and an LSTM controller. It contains recently introduced improvements (Csordás & Schmidhuber, 2019).
We use gradient clipping with magnitude 5 (for CTL) or 1 (for simple arithmetic and ListOps) for all of our models.
Hyperparameters were obtained by a Bayesian hyperparameter search of Weights & Biases2 over the systematically different (OOD) validation set for the +abs/rel + gate models and were reused for all other gated models. For the non-gated models, we used the +rel variant for tuning. It was not possible to tune the baselines using only the OOD validation set because their performance was too bad on that set. We thus used a mixture of IID and OOD validation sets to tune the hyperparameters for the baselines. Table 10 shows the range of hyperparameters used for tuning. “FF multiplier” is used to calculate dFF from dmodel.
We train all models for a fixed number of niters iterations and measure their validation performance every 1000 iterations. For each model, we select the best checkpoint according to the validation performance, and report its test accuracy.

heading: D ADDITIONAL ANALYSIS
text: D.1 COMPOSITIONAL TABLE LOOKUP
In the main text, we only had space to show the gate and attention activity of the NDR for a few timesteps. Here we show the corresponding visualization of all steps in Figures 8 and 9, as well as the attention map for the baseline Transformer with relative positional encoding in Figure 5. We also show the Transformer + abs/rel + gate variant in Fig. 6 and Fig. 7. Please directly refer to the caption of the figures for the corresponding analysis. In general, the visualization for our NDR and the abs/rel + gate variant is easily interpretable, unlike that of the baseline Transformer model.
D.2 LISTOPS
Figures 10 and 12 shows the attention and gate patterns of our NDR architecture on an example from the ListOps dataset. We highlighted notable attention patterns in Sec. 4.
Different heads seem to specialize in different functions. As already mentioned in Sec. 4, head 13 of the NDR architecture, shown in Figure 11, seems to specialize in selecting which arguments belong to which operator.
The gating patterns are also very interesting. In the early stages, the deepest parts of the input are updated: [MAX 2 4 0 8 9] and [MED 8 5 8], which are independent branches of the parse tree that can be processed in parallel. In the following steps, the update patterns spread up in the parse tree, updating the operations that have their arguments available. In this task, the input is read from the first column, which is written in a very late stage.

Review summaries:
This paper addresses an issue of transformers that sometimes they fail to find solutions that are easily expressible by attention patterns. The issue is justified to be the same as the problem of learning useful control flow. The authors propose two modifications, namely adding a copy gate functionality and a geometric attention module which facilitates focusing on local useful operations. The resulting method achieves near perfect accuracy on the considered benchmarks for length generalization, simple arithmetic tasks, and computational depth generalization.

 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.

This paper proposes two modifications to provide additional inductive bias to the attention mechanism in the transformer architecture. The first modification adds a copy mechanism to simulate a “no-op” at a given transformer layer, and the second modification is an attention mechanism that is biased towards attending to local context. Both of these modifications are motivated as being useful for algorithmic tasks like compositional table lookup and arithmetic. From experiments that are *mostly* concerned with some kind of length/depth generalization, we see very significant improvements.