The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization
Authors: Róbert Csordás
Kazuki Irie
Jürgen Schmidhuber
Venue: ICLR 2022 Poster (2022)
Keywords: transformer
compositionality
systematic generalization
algorithmic reasoning
arithmetic

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

Published as a conference paper at ICLR 2022
THE NEURAL DATA ROUTER: ADAPTIVE CONTROL FLOW IN TRANSFORMERS
IMPROVES SYSTEMATIC GENERALIZATION
Róbert Csordás Kazuki Iric' Jürgen Schmidhuber
'The Swiss AI Lab, IDSIA, University of Lugano (USI) & SUPSI, Lugano. Switzerland 2 Kinn Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia frobert, kazuki, juergen e idsia.
ABSTRACT
Despite progress across broad range of applications, Transformers have limited success in systematic generalization. The situation 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 repre- sented 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 generaliza- tion accuracy on the classic compositional table lookup task, as well as near-perfect accuracy on the simple arithmetic task and new variant of ListOps testing for generalization across computational depths. NDR's 's attention and gating patterns tend be interpretable as an intuitive form of neural routing. Our code is public
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 algo-
rithmic combinations of previously learned rules systematicity. Despite recent efforts (Bahdanau
al.
2019: Korrel al. 20199] Lake, 2019 Li al.
2019 Russin et al. [019 Csordás al.
2021
systematic generalization generally remains unsolved (Fodor & Mcl aughhin 1990; Lake &
Baroni 2018 Liska et al.) 2018 Greff et al.) 2020 Hupkes al. 2020)) On some datasets, the best
performing models are neuro-symbolic hybrids (Chen et al. 2020 Liu et al. (0200) using task-specific symbolic functions. However, their applicability 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. (020); Chaabouni et al. (202T); [sordás et al. (202T): Ontañó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 solving 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 (illustrated in Figure Left). Effectively the task is to learn draw 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 express solutions sequence processing problems, and which has a direct equivalent to the operations 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 cetentioo-bbaeed NNs that are able to find intuitive solutions (achieving interpretable attention patterns) could improve systematic generalization.

---

Published as a conference paper at ICLR 2022
Here 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 model, we focus the popular length generalization task of compositional table lookup (CTI.; Liska et al. (2018): Hupkes et al. [Duboi[Dubois Dubois al.] (2020)), as well as two more complex tasks: simple arithmetic task and 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.
help them learn data routing. As a first step towards validating
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 parse tree, this is the depth of the tree). Otherwise, multiple operations must be fused into 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 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.I). 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 "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 the corresponding output via 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 step t. We denote the Transformer state of column i in layer as (f.1) = H E Rd where d is the state size, and H, E TR-N Nxd denotes the states of all N columns in layer In the copy auttenee Transformer (Figure 5 in the appendix), each column i in layer (t + 1) processes the input H, similarly to regular Transformers:
a (1,++1) LayaaNemm((hhhilllcdAAttttttot (f.1)) (f.0), He, H.) + (t.00)
u (1++1) = LayerNorm (a(+t())))
(2)
using the standard multi-head attention operation Vaswani et al. 2017) MultiHeadAAtention with query obtained from (6.4) and keys/values from H,, but the output is gated (using g (++++1 E Rd) as:
g (8,6+1) oFFFN (alitt)))
(3)
h (1+41) = (1+++1) (12++) + (1 - g (1++1)) C (.00
(4)
We use the basic two-layer feedforward block Vaswani et al. 2017 for both FFN and FFN which transforms input æ E d to:
FFN(#:) We max max(Wax + b 0) b2
(5)
2

---

Published as a conference paper at ICLR 2022
but with separate parameters and different dimensionaliites: for FFN Willaa Rdmxd W Rdxder, while for FFNEaee W W E paxd with biases bltta E Re and B b b Rd.
When the gate is closed i.e. g ((x++)) 0 in Eq. 4 the entire transformation is skipped and the input is copied over to the next layer h (1.++1) = (2.11) Crucially, we parameterize the gate (Eq. 3 as func- tion of the output of the self-attention (Eq.[) such that the decision to copy or transform the input for each column depends on the states of all columns. This is crucial difference compared to previously proposed gatings in Transformers, which are solely motivated by training stability (Parisotto et al. 2020) by common practice from convolution- the previous approaches can implement the behavior of our copy gate (see Sec. 16 on related work).
models (Chaabouni al. 2021). None of
The bias of the gate by is initialized to -3 (Hochreieer 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(11, x!2) [M)] with (x() E R each input is projected to key ks()) E Rd value (f) E RR query q0)) E R vectors, and the dot product is computed
for each key/query combination. In our geometric attention, the dot product is followed by sigmoid function to obtain a score between () and 1:
P.j or (k(
which will be treated probability of the key at (source) position j matching the query at (target) position i. These probabilities are finally converted to the attention scores A., as follows:
Aij = P.ù II (1.P).
AES,,,
where S denotes the set of all (source) indices which are closer to than j is to i, and when two indices have the same distance i, we consider the one which is to the right of (i.e., greater than i) to be closer, i.e., S. = ( E N3 (i -
E (1,-.,, NJ \ (1 : < 1 ifi
(8)
if
In addition, explicitly of source indices illustrated in Figure Right. The resulting scores A.,j are the attention scores used compute the weighted averages of the value vectors.
the diagonal by setting A., 0 for all i = 1, .... N. The ordering
By using the terms (1 P..k) ) in Eq. when there is a match, it downscales any other more distant matches. Two recent works (Brooks al.)
2021 Banino et al.) (02T)) use such a parameterized geometric distribution in the form of Eq. (see Sec. 6n related work).
The resulting attention function has complexity of O(N2), similar the regular self-attention used in Transformers (Vaswani et al.] (017). Eq. [q be implemented 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. 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. 8 In practice we found erucial to augment the score computation of Eq. 6 with additional directional information, encoded as a scalar Dr.j E R for each target/source position pair (i,j):
(L bLR, ifi (R) + bRL, if
D.J
where (0) E Rd denotes the input/state at position and WIR WRL E 1xdd byR. bRL. E RR are trainable parameters. This directional information is integrated into the score computation of Eq. aa
3

---

Published as a conference paper at ICLR 2022
Target pos.
1=2
33
3
30
3
2
[-0
Source pos.
Figure 1: Left: an ideal sequence of computations in a Transformer for an arithmetic expression.
Right: ordering (numbers in the grid) of source positions used in geometric attention (Eq. 8 N = 5).
follows (akin to how Dai et al. (2019) introduce the relative positional encoding (Schmidhuber extra term in the computation of attention scores):
1992)
P., (a + bb)) Wi, + BDi +7)
(10)
where the matrix W₂ Rddeaadd maps the states to queries, by E TRc is a bias for queries, W Rilaadd maps states to keys (we note that chaad is typically the size of the key, query and value vectors for each head, dhead
Theade )M and ax 3, 7 E IR are learned scaling coefficients and bias,
initialized to Vihyos - 1, O. Using this additional directional information, each query (position i) can potentially learn to restrict its attention to either the left or right side.
EXPERIMENTS
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 com- positional way (productivity). Further experimental details for each task can be found in Appendix C
3.1
COMPOSITIONAL TABLE LOOKUP
Task. The compositional table lookup task (Liska et al. 2018 Hupkes et al., 2019 Dubois et al.. 2020) is constructed based 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 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 represented by 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 sequence- to- sequence variant of this task (Dubois al. 2020) where the model has 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 forward (former) and backward (latter) variants of the task. To evaluate systematic generalization, the train/valid//est sets reflect different numbers of compositions: samples with 1-5/6-8/9. 1016-889-10 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 Sec. for further details on the previous work.
b' or 'b d 101'. We thus study both
Results. We consider five different baselines: an LSTM (Hochreiter & Schmidhuber 1997), bidi- rectional LSTM (Schuster & Paliwal 1997). DNC GGrves et al.) 2016) Csordás & Schmidhuber 2019), Universal Transformers (Vaswani et al. 2017 Dehghani et al. 2019), and its relative position variants (Csoráás et al.) 202T). ). For Transformers, the prediction is based on the last column in the final layer (we conduct ablation study on this choice in Appendix A). The syper-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 claim on the necessity
4

---

Published as a conference paper at ICLR 2022
for "sufficiently" deep architecture. The main results this task are shown in Table I 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 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 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 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. 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, 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 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 (See. 2.2). also achieved perfect generalization and was found easier to train. We present the corresponding visualization of our model in Sec.
Table 1: Accuracy compositional table lookup dataset.
IID
Longer
Model
Forward
Backward
Forward
Backward
LSTM
1.00 0.00
0.59 0.03
1.00 + 0.00
0.22 + 0.03
Bidirectional LSTM
1.00 0.00
1.00 + 0.00
1.00 + 0.00
1.00 1 0.00
DNC
1.00 + 0.00
0.57 + 0.06
1.00 + 0.00
0.18 + 0.02
Transformer
1.00 0.00
0.82 1 0.39
0.13 1 0.01
0.12 1 0.01
rel
1.00 + 0.00
1.00 1 0.00
0.23 1: 0.05
0.13 1 0.01
rel gate
1.00 + 0.00
1.00 0.00
0.99 + 0.01
0.19 0.04
abs/rel+ gate
1.00 0.00
1.00
0.98
0.98
+ geom, att.
0.96 0.04
0.06
0.06
0.15 0.02
+ geom. att. + gate (NDR)
1.00 + 0.00
1.00 0.00
1.00 0.00
1.00 + 0.00
3.2 SIMPLE ARITHMETIC
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 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) generated randomly. The tree depth is up The depth is measured as the number of operations, ignoring the leaves, the example above has depth of 2. The sequence length is limited to at most 50 tokens.
should evaluate to '0' (30 modulo 10). The expressions
to 5 for the training set, 6 for the validation set, and 7-8 for the test set.
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 converge, the NDR achieves near-perfect accuracy after 50 K steps of training.
Results.
5

---

Published as a conference paper at ICLR 2022
Table 2: Performance of different models on the simple arithmetic dataset. All models are trained for 200 K iterations, except the NDR which we stop training at 100 K. We also report the performance after 50 K iterations, where it can be seen that NDR converges significantly faster than the others.
IID (1..5)
Test (7..8)
200 K
200 K
50 K
LSTM Bidirectional LSTM
0.99 + 0.00
0.74 + 0.02
0.72 = 0.01
0.98 1 0.01
0.82 + 0.06
0.80 = 0.04
Transformer
0.98 + 0.01
0.47 = 0.01
0.29 = 0.01
rel
1.00 0.00
0.77 = 0.04
0.40 0.05
abs/rel gate
1.00 0.01
0.80 + 0.16
0.73 0.15
geom. att. gate (NDR)
1.00 0.00
0.98
0.01
0.98 0.01
3.3 LISTOPS
We also evaluate our model on a variant of the ListOps task (Nangia & Bowman 2018) which is a popular task commonly used evaluate parsing abilities of NNs (Havrylov et al. [019 [Shen] et al. 2019 Xiong et al., 2021, Tay et al. [2021 Irie et al., 2021). Some special architectures such as Chowchury& Caragea (2021) can almost perfectly generalize to longer sequences on this task. However, as far we know, 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 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 should return 6. There are two knoln 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 for validation and depths and 8 for test. Importantly, we make sure that a depth- K sample effectively requires computation until depth-A 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 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.
9. For example: [MED 4 8 5 [MAX 8 4 9
Table 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 near-perfect accuracy on IID data. In contrast, our model achieves near-perfect generalization.
Results.
Table 3:
Performance of different models on balanced ListOps dataset. All models are trained for 200 K iterations, except all +gate variants which converge after 100 K steps. The numbers in the parentheses indicate the problem depths (1-5 for the IID, and 7-8 for the test set).
IID (1..5)
Test (7..8)
LSTM Bidirectional LSTM
0.99 1: 0.00
0.71 4: 0.03
1.00 0.00
0.57 + 0.04
Transformer
0.98 + 0.00
0.74 + 0.03
rel
0.98 0.01
0.79 + 0.04
abs/rel + gate
1.00 + 0.01
0.90
geom. att. gate (NDR)
1.00 + 0.00
0.99

---

Published as a conference paper at ICLR 2022
4 ANALYSIS
In this section, we provide some visualizations of attention and gating patterns of the NDR and the corresponding analyses. For more visualizations, refer the readers to Appendix D
Compositional Table Lookup. Figure shows the gating and attention patterns of the NDR model for 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. 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 be performed.
06
[-0
(-8
0.4
Figure 2: Example visualization of NDR. For other models, see Appendix D Top: Attention map for different steps. The x/y-axis corresponds to source/target positions, respectively. Each position focuses on the column to the right, except the last one where the result is read from, which focuses on the last operation. The focus becomes clear only once the result available. Bottom: gate activations for different steps/layers. The gates remain closed until the data dependencies are satisfied.
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 ex- ample is: [SM [MED [MIN
[MAX
9
[MED 8 5
First of all, we find that there is 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 at- tend to the operator. In step 0 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-tor titiistttopperttion attention patterns remain for few steps). We note that at this stage, the last argument of MIN is not ready yet ( ([MIN 1 7 4 [MAX 2 4
O 8 9 1). We can see that only arguments which are already ready 7 4) attend the operator (see the column of MIN). In step (1 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 is always the result of MAX when 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. /st row), we recognize that the result of MAX marked as an argument for MIN in head 13 which is responsible for communication between operators and their arguments. This shown by the new attention which appears at = 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, f = 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 (Ist row), all arguments for the first MED (the result of MIN which is 1. and 7). Finally in = 7 (2nd row), two heads, head 3 and head seem to copy/gather two inputs needed to compute the corresponding median, 1 and 7, and store them in
also recognized
7

---

Published as a conference paper at ICLR 2022
Figure 3: Example visualization of NDR on ListOps. The top row shows head 13 in different steps, which controls which arguments are used in which step. The bottom row shows different heads in different key steps. Please refer to Sec. 4 for the step-by-step description. More visualizations are provided in the appendix: Fig. 12 shows the max of attention over all heads for all steps, Fig. I3 shows all steps of head 13, and Fig. 14 shows the corresponding gates.
the column of the operator MED. A complete visualization of further steps can be found in Appendix D.2 We noticed that some 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, attention patterns described above in the baseline Transformer with relative positional encoding, at least some inspected examples, which also explains its rather high accuracy.
also partially find the
5 DISCUSSION
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 a single step makes the processing faster (fewer steps), but also more specialized. Learning the most granular solutions thus preferable for generalization. At the same time, not all processing should happen serially: be processed
branches of the computational graph that do not have common data dependencies independently in parallel, which we empirically observe in our NDR in the ListOps example (Sec. 4 This enables the architecture 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 ct al. 202I). 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 sub-optimal choice of architectureoptiiizztioo 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 bottom-up approach: building on toy tasks that focus on individual aspects of generalization and using them for improving models.
6 RELATED WORK
Gating inside Transformers. Several prior works have proposed to use some sort of gating within Transformer architectures (Parisotto al. 2020 Chaabouni et al. 2021). Our proposed copy gate
8

---

Published as a conference paper at ICLR 2022
is different from those 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 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 improve over regular Transformers in our preliminary experiments on compositional table lookup. Recent work by CChabooni et al. (202T) also makes use of "gating" in Transformers through gated linear unit (GLU) activation function com- monly used in convolutional NNs (Dauphin et al. 2017). Transformer models with such an activation function were reported to outperform RNN baselines on 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.c. a closed gate zeros out the input) which allows the model to skip transformation. In more general context, benefits of the GLU activation in Transformers vary across tasks (Irie et al. 2019 Shazeerr 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 al.
2021) have used a form of parameterized geometric distribution (PGD; in the form of Eq. 7 Brooks al. (202T) have used such a distribution to parameterize the movement of pointer on a sequence of instructions. Banino et al. (2021) have used it to implement adaptive computation time (Schmidhuber 2012 Graves 20166). 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 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] [992] 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 (Coordás et al., 2021). Here we also present two new approaches related to positional encoding. One is the gated combination of absolute and relative posi- tional encoding (Sec.B.I] details in Appendix We show that absolute positional encoding can com- plement relative positional encoding. The former enables the model always attend to specific po- sition, 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.
CONCLUSION
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 Trans- formers: the problem of routing and retaining data in an unaltered fashion, which we solve by an additional copy gate, and the problem of learning solve by geometric attention. Our new model, the Neural Data Router (NDR), generalizes to compo- sitions 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 full sequence wwweenne-to-sequeeee to-sequence model and evaluate it on other standard tasks in systematic generalization requiring generation of variable-length output sequences.
iigth-inndeeeddent attention patterns, which we

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
We thank Imanol Schlag and Sjoerd van Steenkiste for helpful discussions and suggestions on an carlier 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: 100021_192358 project NEUSYM. We thankful for hardware donations from NVIDIA & IBM. The resources used for the project were partially provided by Swiss National Supercomputing Centre (CSCS) project s0023.
REFERENCES
Dzmitry Bahdanau, Harm de Vries, Timothy J O' O'Donnel, Shikhar Murty, Philippe Beaudoin, Yoshua Bengio, and Aaron Courville. CLOSURE: Assessing systematic generalization of CLEVR models. In ViGIL workshop, NewrlPS, Vancouver, Canada, December 2019.
Andrea Banino, Jan Balaguer, and Charles Blundell. PonderNet: Learning to ponder. Preprint wrXiv:2107.05407, 2021.
Ethan A. Brooks, Janarthanan Rajendran, Richard L. Lewis, and Satinder Singh. Reinforcement learning of implicit and explicit control flow instructions. Proc. Int. Conf. on Machine Learning (ICML), pp. 1082-1091, Virtual only, July 2021.
Tom B Brown et al. Language models are few-shot learners. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
Rahma Chaabouni, Roberto Dessi, and Eugene Kharitonov. Can transformers, jump around right in natural language? Assessing performance transfer from SCAN. Preprint arXiv:2107.01366. 2021.
Xinyun Chen, Chen Liang, Adams Wei Yu, Dawn Song, and Denny Zhou. Compositional generaliza- tion via neural-symbolic stack machines. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
Jishnu Ray Chowdhury and Cornelia Caragea. Modeling hicrarchical structures with continuous recursive neural networks. In Proc. Int. Conf. on Machine Learning (ICML), pp. 1975-1988, Virtual only, July 2021.
Róbert Csordás and Jürgen Schmidhuber. Improving differentiable neural computers through mem- ory masking, de-allocation, and link distribution sharpness control. In Int. Conf. on Learning Representations (ICLR), New Orleans, LA,, USA, May 2019.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of Transformers. In Proc. Conf. Empirical Methods in Natural Language Processing (EMNLP), Punta Cana, Dominican Republic, November 2021.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. T::::former-XI:: Attentive language models beyond fixed-length context. In Proc. Association for Computational Linguistics (ACL), pp. 2978-2988, Florence, Italy, 2019.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proc. Ini. Conf. on Machine Learning (ICML), pp. 933-941, Sydney, Australia, August 2017.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal Transformers. In Int. Conf. on Learning Representations (ICLR), New Orleans, LA, USA, May 2019.
Roberto Dessi and Marco Baroni. CNNs found to jump around more skillfully than RNNs: Composi- tional generalization in seq2seq convolutional networks. In Proc. Association for Computational Linguistics (ACL), pp. 3919-3922 Florence, Italy, July 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional Transformers for language understanding. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL-HLT), pp. Minneapolis, MN, USA, June 2019.
10

---

Published as a conference paper at ICLR 2022
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
Yann Dubois, Gautier Dagan, Dieuwke Hupkes, and Elia Bruni. Location attention for extrapolation to longer sequences. In Proc. Association for Computational Linguisties (ACL), pp. 403-413, Virtual only, July 2020.
Jerry Fodor and Brian McLaughlin. Connectionism and the problem of systematicity: Why Smolensky's s solution doesn't 'tt work. Cognition, 35(2):183-204, 1990.
Jerry A Fodor, Zenon W Pylyshyn, et al. Connectionism and cognitive architecture: A critical
analysis. Cognition, 28 28(1-)):3-71, 1988.
Daniel Furrer, Marc van Zee, Nathan Scales, and Nathanael Schärli. Compositional generalization in semantic parsing: Pre-training vs. specialized architectures. Preprint arXiv:2007.08970, 2020.
Alex Graves. Adaptive computation time for recurrent neural networks. In Int. Conf. on Learning Representations (ICLR) Workshop Track, Vancouver, Canada, April 2016.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinska, Sergio Gomez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John P. Agapiou, Adrià Puigdomènech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukcuoglu, and Demis Hassabis. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. Preprint arXiv:2012.05208, 2020.
Stephen José Hanson. A stochastic version of the delta rule. Physica D: Nonlinear Phenomena, 42
(1-3):265-272, 1990.
Serhii Havrylov, Germán Kruszewski, and Armand Joulin. Cooperative learning of disjoint syntax and semanties. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACI-HLT)) pp. 1118-1128, Minneapolis, USA, June 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Proc. The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, Las Vegas, NV, USA, June 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, pp. 1735 1997.
Dieuwke Hupkes, Anand Singh, Kris Korrel, German Kruszewski, and Elia Bruni. Learning com- positionally through attentive guidance. In Proc. Int. Conf. on Computational Linguistics and Intelligent Text Processing, La Rochelle, France, April 2019.
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, pp. 757-795, 2020.
Kazuki Irie. Advancing Neural Language Modeling in Automatic Speech Recognition. PhD thesis, Computer Science Department, RWTH Aachen University, Aachen, Germany, May 2020.
Kazuki Irie, Albert Zeyer, Ralf Schlüter, and Hermann Ney. Language modeling with deep Trans- formers. In Proc. Interspeech, pp. 3905-3909, Graz, Austria, September 2019.
Kazuki Irie, lmmool Schlag, Róbert Csordás, and Jürgen Schmidhuber. Going beyond linear Trans- formers with recurrent fast weight programmers. Preprint 1rXiv:2106.0633 2021.
11

---

Published as a conference paper at ICLR 2022
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang. Mare van Zee, and Olivier Bousquet. Measuring compositional generalization: A comprehensive method on realistic data. In Int. Conf. on Learning Representations (ICLR), Addis Ababa, Ethiopia, 2020.
Yoon Kim and Yacine Jernite David Sontag Alexander Rush. Character-aware neural language models. In Proc. AAAI Conference on Artificial Intelligence, Phoenix, A7, USA, February 2016.
Kris Korrel, Dieuwke Hupkes, Verna Dankers, and Elia Bruni. Transcoding compositionally: Using attention to find more generalizable solutions. In Proc. BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, ACL, pp. 1-11, Florence, Italy, 2019.
Brenden M Lake. Compositional generalization through meta sequencce-t-sequeene learning. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pp. 9788-9798, Vancouver, Canada, December 2019.
Brenden M. Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-to-sequenee recurrent networks. In Proc. Int. Conf. on Muchine Learning (ICML), pp. 2873-2883 Stockholm, Sweden, July 2018.
Yuanpeng Li, Liang Zhao, Jianyu Wang, and Joel Hestness. Compositional generalization for primitive substitutions. In Proc. Conf. on Empirical Methods in Natural Language Processing and Iut.Joint Conf. on Natural Language Processing (EMNLP-IJCNLA) pp. 4292-422244 4301, Hong Kong. China, November 2019.
Adam Liska, Germán Kruszewski, and Marco Baroni. Memorize generalize? Searching for compositional RNN a haystack. In AEGAP Workshop ICML, Stockholm, Sweden, July 2018.
Qian Liu, Shengnan An, Jian-Guang Lou, Bei Chen, Zeqi Lin, Yan Gao, Bin Zhou, Nanning Zheng, and Dongmei Zhang. Compositional generalization by learning analytical expressions. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Int. Conf. on Learning Representations (ICLR), New Orleans, LA, USA, May 2019.
Nikita Nangia and Samuel R. Bowman. ListOps: A diagnostic dataset for latent tree learning, In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL-HLT), pp. 92 99, New Orleans, USA, June 2018.
Santiago Ontañón, Joshua Ainslie, Vaclav Cvicek, and Zachary Fisher. Making Transformers solve compositional tasks. Preprint arXiv: 200Xivv2108... 04378, 2021.
Emilio Parisotto, H. Francis Song. Jack W. Rae, Razvan Pascanu, Çaglar Gülçehre, Siddhant M. Jayakumar, Max Jaderberg, Raphaël Lopez Kaufman, Aidan Clark, Seb Noury, Matthew Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing Transformers for reinforcement learning. In Proc. Int. Conf. on Machine Learning (ICML), volume 119, pp. 7487-7498, Virtual only, July 2020.
Jake Russin, Jason Jo, Randall C O'RRilly, and Yoshua Bengio. Compositional generalization in deep seq2seq model by separating syntax and semantics. Preprint arXiv: 1904::190400978: 2019.
Jürgen Schmidhuber. Self-delimiting neural networks. Technical
Report IDSIA-08-12 12,
arXiv:1210.0118v1, The Swiss AI Lab IDSIA, 2012.
Jürgen Schmidhuber. Learning complex, extended sequences using the principle of history compres-
sion. Neural Computation, 4(2):234-242. 1992.
Mike Schuster and Kuldip K Paliwal. Bidirectional recurrent neural networks. IEEE Transactions on 1997.
Signal Processing.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACI-HILT) pp. 464 468, New Orleans, Louisiana, USA, June 2018.
12

---

Published as a conference paper at ICLR 2022
Peter Shaw, Ming- Wei Chang, Panupong Pasupat, and Kristina Toutanova. Compositional general- ization and natural language variation: Can a semantic parsing approach handle both? Preprint arXiv:2010.12725. 2020.
Ming-Wei
Noam Shazeer. GLU variants improve transformer. Preprint arXiv:2002.05200,, 2020.
Yikang Shen, Shawn Tan, Seyed Arian Hosseini, Zhouhan Lin, Alessandro Sordoni, and Aaron C. Courville. Ordered memory. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pp. 5038-50800 5049, Vancouver, Canada, December 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1):1929-1958 2014.
Rupesh K Srivastava, Klaus Greff, and Jürgen Schmidhuber. Training very deep networks. In Proc. Advances in Neural Information Processing Systems (NIPS), pp. 2368-2376, Montreal, Canada, December 2015.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long Range Arena A benchmark for efficient transformers. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Advances in Neural Information Processing Systems (NIPS), pp. 5998-69088 Long Beach, CA, USA, December 2017.
Gail Weiss, Yoav Goldberg. and Eran Yahav. Thinking like Transformers. In Proc. Int. Conf. on Machine Learning (ICML), pp. 11080-11090, Virtual only, July 2021.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A Nyström-based algorithm for approximating self-attention. In Proc. AAAI Conf. Artificial Intelligence, pp. 14138 14148. Virtual only, February 2021.
13

---

Published as a conference paper at ICLR 2022
A ABLATIONS
Number of layers. In Sec. 2 we hypothesized that decomposition of the problem into its elementary operations is necessary property of a model which generalizes. This motivated us to configure our models to have at least as many layers as the depth of the computation involved, plus few additional layers for writing the output and for gathering an overview of the problem the beginning. We assumed that in such model with sufficient number of layers, each layer learns the underlying "elementary" operation. The resulting models are thus deeper than those typically used in the literature for similar tasks (Keysers et al. 2020; Tay et al. (021). Here provide an ablation study to demonstrate that such depths are effectively necessary for generalization. We measure the IID and generalization performance with various numbers of layers on the compositional table lookup dataset. Since our test set on the CTL task consists of up to 10 function applications, it should require about 12 layers according our hypothesis. Table 4 shows the results. We clearly observe that, while the shallow models can also solve the IID split, only the deep models generalize to the longer problems (here the 12-layer model generalizes almost perfectly, but the 10-layer one does not).
Table 4: The performance of NDR on the compositional table lookup dataset, with different number of layers.
IID
Test
Flayers
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
1.00 + 0.00
1.00 0.00
0.23 + 0.02
0.24 + 0.03
1.00 + 0.00
0.96 0.03
0.22 +
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 the sequence which might increase the difficulty of the problem, in particular for the models using absolute positional embeddings. Table 5 ablation study. We observe that this choice has only marginal impact on the model performance. As a side note, 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.
shows the corresponding
Table Accuracy on compositional table lookup dataset with the results read from the first or last column (Readout).
IID
Longer
Model
Readout
Forward
Backward
Forward
Backward
Transformer
First
1.00
0.00
0.82 + 0.39
0.12 0.01
0.13 + 0.01
Last
1.00 0.00
0.82 + 0.39
0.13 0.01
0.12 + 0.01
rel
First
1.00
0.00
1.00 + 0.00
0.12 0.01
0.22 + 0.05
Last
1.00 + 0.00
1.00 0.00
0.23 0.05
0.13 + 0.01
+ rel + gate
First
1.00 0.00
1.00 0.00
0.17 0.02
1,00 + 0.00
Last
1.00 + 0.00
1.00 + 0.00
0.99 0.01
0.19 0.04
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.I We could also consider using Adaptive Computation Time (ACT) to dynamically determine the number of steps. Furthermore, ACT introduces 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 certain time and copy the result to the output) is fundamentally different
14

---

Published as a conference paper at ICLR 2022
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 provide some experimental results models with ACT which confirm that the proposed copy gate is a crucial component for generalization which not be replaced by ACT.
from copy gate (Sec. 2.1
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 (0166) to Transformers, and another one used by Dehghami et al. (2019). We start with the description of the former.
An extra sigmoidal unit p(6.0) is computed for each column in each timestep t as:
pl. o WHh (6.1) + bH)
(11)
where WH E Rixd and by E IK are trainable parameters. By comparing the cumulative sum of g(4.0) over time steps certain threshold value (1 - <) with a pappmmrrr (0.01 in our experiment), we determine the termination step I" for column i as:
T (Tixx: mix: min(l's
Σ 216021-0
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.tt
ift < T"
(13)
Re
if T
R
R (6.00
(14)
which is used to compute the final output of column as:
T
(t.0 (a.))
(15)
In Dehghani al. (2019)3 's variant, different equation is used in lieu of Eq. 15 above and the
computation of the reminder term It' in Eq. a4 is not properly handled in case where Eq. 12 terminates because of the first condition on 7 Tax. For further details, we refer the readers to Listing
and 2 in Dehghani et al. (2019) and/or our public code.
One subtlety introduced by Dehghani et al. (2019) which we note here that the computation of the final output o² of column i effectively "halts" after T (since o' only depends on h (1.4) for 0 < t < T but column itself still continues transforming the hidden states (f.t) for steps t T 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. TV T³). 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. [5p. We follow this decision in our implementation of both variants.
In addition, new regularizer term, LACT = a E Ri is added to the loss function, where N is the length of the input sequence. This makes the network prefer short computations. We ran hyper-parameter search for a 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 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 many 30k used for the models without ACT. Table
15

---

Published as a conference paper at ICLR 2022
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, 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 a 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 which is orthogonal to our study.
Table 6: Accuracy on compositional table lookup dataset with adaptive computation time (ACT). Two variants of ACT are shown: "U" corresponds Dehghani al. (2019), while "A" is the variant described in Appendix [ We also include baselines without ACT from Table as a reference. Generalization performance after 30k and 60k training steps are shown.
IID
Longer, 60k
Longer, 30k
Model
ACT
Forward
Backward
Forward
Backward
Forward
Backward
Transformer
1.00 0.00
0.82 + 0.39
0.13 + 0.01
0.12 + 0.01
A
1.00 + 0.00
1.00 = 0.00
0.12 + 0.02
0.12 + 0.01
0.13 + 0.01
0.13 + 0.01
U
1.00 0.00
1.00 = 0.00
0.12 + 0.01
0.11 + 0.01
0.13 1 0.01
0.12 + 0.01
rel
1.00 + 0.00
1.00 = 0.00
0.23 0
0.13 = 0.01
A
1.00 0.00
1.00 0.00
0.99
0.02
0.13 1 0.00
0.13 0.01
1.00 0.00
1.00 +
0.92 0.14
0.12 1 0.02
0.67 0.41
+ 0.00
+00
0.96 0.04
0.93 = 0.06
0.16 0.02
0.15 = 0.02
1.00 0.00
1.00 = 0.00
0.97
0.05
0.45
+ 0.21
0.58 0.16
0.30 + 0.17
0.96 0.10 U
1.00 + 0.00
0.72 0.35
0.44 + 0.19
0.31 + 0.22
0.21 + 0.07
+ gale
1.00 0.00
1.00 1 0.00 0.00
0.99 0.01
0.19 1 0.04
abs/rell gate
1.00 0.00
1.00 = 0.00
0.98
0.02
0.98 = 0.03
geo + gate (NDR)
1.00 + 0.00
1.00 = 0.00
1.00 = 0.00
1.00 = 0.00
15
Num. of steps
10
3
4
5
10
11
12
13
Sequence length
Figure 4: Average number of steps/layers for different sequence lengths on the compositional table lookup task for the Transformer with relative positional encodings and the ACT variant described in Appendix A The red line shows Tmax = 14. Note that the sequence length shown here includes the begin and end tokens. Thus, the sequence length of 4 corresponds to one function application (3 for the identity function i.e. no function is applied).
B DETAILS OF ATTENTION WITH COMBINED ABSOLUTE/RELATIVE POSITIONAL ENCODING
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-bes based attention to read out the result instead content- based attention. presentation order, the last column of the transformer should focus on the second column, whose relative position changes dynamically with the length the sequence. We solve this issue by adding option choose between absolute and relative positional encodings to the attention head.
the backward
16

---

Published as a conference paper at ICLR 2022
In what follows, we describe the operation within single layer/step. This allows to omit the layer/step-inde for better readability, and thus denote the state of column/position as h, instead of h,.0) 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 follows:
rj - o War + bar)
(16)
A.j h W. Wa,n + b W WR.r.h +
(h. W W + b (K.)) (pi-)ri + Pp( ri)
(a)
(6)
(d)
(c)
(17)
where the matrix W E Rdaaadd maps the states to queries, WA, Rddeatdd maps states to keys, while WK, E Rduas xd maps positional embeddings to keys. dhead the size the key, query and value vectors for each head, set as daead =
Thead E, by. E [R are learned vectors. p; E Rd the standard sinusoidal embedding for position (Vaswani et al. 2017). Softmax is applied to the second dimension of Á to obtain the final attention scores, A. Component (a) corresponds to content-based addressing, (b, e) content-basee positional addressing, (c) represents 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 r; (Eq. 16 parameterized by War E Rd and bør E R), which is the function of the state target position
C IMPLEMENTATION DETAILS
PyTorch implementation of our models together with the experimental setup is available under https: : /github. com/ robert csoretssssdsss ndz The performance of all models is reported as mean and standard deviations over different seeds.
C. CHOOSING THE NUMBER OF LAYERS
In Sec. 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 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) x (steps per operation) + a few more layers.
Determining the number of steps needed by the elementary operation not straightforward but can be done empirically. For example, for ListOps, as is shown in Sec. it requires two steps per operation: one step in which the operands attend 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 safe, and can be used 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 used read out the final result, and one or a few more can be needed for communication between columns (c.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 to train models using a depth which is "sufficient" to solve the training set, but increase at test time to generalize to 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 performance penalty, while it speeds up training.
17

---

Published as a conference paper at ICLR 2022
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 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. 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 sum- modulo (SM), which is 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. M samples are used for training, 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 or 2. SUM, all operands are needed perform the operation. If we construct 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 the tree that are too deep and still getting the correct result with a high chance. We also ensure that the number of samples the same for all possible dependency depths in each split.
C.3 MODEL DETAILS
We the Adam W optimizer (I Loshchilov & Hutter 2019) for all of our models. Standard hyperpa- rameters are listed in Tab. 7 and 9 Additionally, models with gating use dropout (Hanson 1990 Srivastava al. 2014) applied to the content-based query and the position-query components 0.1 for most models, except for non-gaten 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 Narı in Eq. 2 The Transformer/NDR layer with copy gate is illustrated in Figure 5
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, were unable train them to perform well even on IID data with comparable sizes.
All Transformer variants have 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 (Csordás & Schmidhuber] 2019).
We use gradient clipping with magnitude 5 (for CTL) or (for simple arithmetic and ListOps) for all of our models.
18

---

Published as a conference paper at ICLR 2022
1-X
Layernorm
Sigmoid
Update
FF: Gate
Layernor
Attention
Figure 5: Structure of Transformer/NDR layer with a copy gate (Sec. 2.1 The blue part corresponds to the standard Transformer, except for the missing residual connection around the feedforward block ("FF: Update"). The gray part is the copy gate. The feedforward part corresponding to the gate is usually significantly smaller than the one used for the update.
Hyperparameters were obtained by Bayesian hyperparameter search Weights & Biase over the systematically different (OOD) validation set for the + +abs/ / rel + gate models and were reused for all other gated models. For the non- gategated models, we used the +rel variant for tuning. It was possible to tune the baselines using only the OOD validation set because their performance was too bad on that set. We thus used mixture of IID and OOD validation sets to tune the hyperparameters for the baselines. Table shows used to calculate dFF from dmnddl--
shows the range of hyperparameters used for tuning. "FF multiplier" is
We train all models for fixed number of Plters 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.
Table 7: Hyperparameters used for diffeernnt models on the compositional table lookup task. We denote the feedforward size as drr, weight decay as "wd.", dropout as "do.". The model is trained for Titers iterations.
dmaaal
dyF
Talaees batch
learning rate
wd.
do.
L.TM
200
256
10 10
0.5
20k
Bidirectional LSTM
400
256
10 10
0.5
20k
DNC
200
256
10 10-4
0.5
20k
Transformer
128
256
4
11
512
1.5 * 10
0.0025
0.1
30k
+ rel
128
256
512
1.5 10-4
0.0025
30k
+ rel + gate
256
512
1
14
512
1
0.01
0.5
30k
abs/rel gate
256
512
512
2
0.01
0.5
30k
+ geom, att.
128
256
4
11
512
1.5 10
0.0025
0.1
30k
+ geom. att. gate (NDR)
256
512
14
512
1.5 10
0.01
0.5
30k
D ADDITIONAL ANALYSIS
D.1 COMPOSITIONAL TABLE LOOKUP
An idealized sequence computations in Transformer for an example from CTL task is shown in
Fig. Each column waits for its input from the left side, then performs update. Finally, the last
hhttps:/waant
19

---

Published as a conference paper at ICLR 2022
Table 8: Hyperparameters used for different models on the simple arithmetic task. We denote the feedforward size as drr, weight decay as "wd.", dropout as "do.". The model is trained for fittrs iterations.
dmadel
dFF
7headd
Niaerrs
batch
learning rate
wd.
do.
Yitees
L.TTM
200
2
2.66
10 10
0.5
200k
Bidirectional LSTM
400
2
256
10 + 10
0.5
200k
Transformer
128
256
11
512
1.5 x 10
0.0025
0.5
200k
rel
128
256
4
11
512
1,5 10
0.0025
0.5
200k
+ abs/rel + gate
256
1024
15
512
1.5
10
0.01
0.5
100k
geom. att. gate (NDR)
256
1024
4
15
512
1.5
10-4
0.01
0.5
100k
Table 9: Hyperparameters used for different models on the ListOps task. We denote the feedforward size as drr, weight decay as "wd.", dropout as "do.". The model trained for fiiers iterations.
dmnddl
dFF
7hennd
Taayers
batch s.
learning rate
wd.
Witers
LSTM
512
512
10
0.08
200k
Bidirectional LSTM
1024
4
512
10 + 10
0.08
0.1
200k
Transformer
256
1024
16
6
512
4 10
0.05
0.015
200k
+ rel
256
1024
512
10-4
0.05
0.015
200k
+ abs/ree + gate
512
1024
16
20
512
2 10-4
0.09
0.1
100k
geom. gate (NDR)
512
1024
16
20
512
2 10-1
0.09
0.1
100k
column copies the result. So far, in the main text, we only had space to show the gate and attention activity of the NDR for few timesteps. Here show the corresponding visualization of all steps in Figures 10 and 11 as well as the attention map for the baseline Transformer with relative positional abs /rel gate variant in Fig.
encoding Figure
We also show the Trans former and Fig. / Please directly refer to the caption of the figures for the corresponding analysis. In general, the visualization for our NDR and the abs / rre gate variant is easily interpretable, unlike that of the baseline Transformer model.
000
Lm3
D60
D-2
010
110
Figure 6: An ideal sequence of computations in a Transformer for an example CTL task.
D.2 LISTOPS
Figures 12 and 14 shows the attention and gate patterns of our NDR architecture example from the ListOps dataset. We highlighted notable attention patterns in Sec.
Different heads seem specialize in different functions. As already mentioned in Sec. head 13 of the NDR architecture, shown in Figure 13 seems to specialize selecting which arguments belong to which operator.
The gating patterns are also very interesting. the early stages, the deepest parts of the input are 8), which independent branches of the parse
updated: [MAX 2 4 0 8 9] and [MED 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.
20

---

Published as a conference paper at ICLR 2022
Table 10: Parameter ranges for hyperparameter tuning
Parameter
Range
learning rate
0,00005 0.001
Tlayers
4. ... 20
dmodel
128, 256,, 512
Thead
2, 4, 8, 16
weight decay
0.0
dropout
0.0 0.5
attention dropout
0.0 0.5
FF multiplier
1,2,4
1.5
1-5
1-8
1-9
1.0
0.8
0.6
0.4
Figure 7: Attention map for every computational step for a baseline Transformer with relative positional encoding on CTL. The attention pattern gets blurry very quickly, and the model does not generalize to longer sequences.
21

---

Published as a conference paper at ICLR 2022
1=0
1=1
1+2
1-5
1-8
1-8
1x18
0.8
0.6
0.4
0.2
Figure Attention map for every computational step for a Transformer with gating and rela- tive/absolute positional encoding (presented in Figure CTL. The attention pattern relatively stable over time, and it gets blurrier only after the given column is processed and updated. The gate sequence for the same input can be seen in Figure 9
22

---

Published as a conference paper at ICLR 2022
1-5
0.4
0.2
0.0
Figure 9: Gates for every computational step for Transformer with gating and relative/absolute positional encoding on CTL. The gates are closed until all arguments of the given operation become available. The attention maps for the same input can be seen in Figure
23

---

Published as a conference paper at ICLR 2022
1=0
1s1
2+2
10
1-5
1-8
1-8
1-8
1.18
0.8
0.6
0.4
0.2
Figure 10: Attention map for every computational step of the NDR on CTL. The network correctly and clearly focuses the last element of the sequence, and the last sharp read happens in step 10- corresponding to the 10 function calls in the example. The gate sequence for the same input can be seen in Figure II
24

---

Published as a conference paper at ICLR 2022
0.4
0.2
0.0
Figure 11: Gates every computational step of the NDR on CTL. The gates remain closed until all arguments of the given operations become available. The attention maps for the same input can be seen in Figure 10
25

---

Published as a conference paper at ICLR 2022
Attention maps for every computational step of the NDR on ListOps. The network has 16 heads; the max of them shown. The input has only depth 4, which explains the early stopping of the computation, roughly after 8-9 steps, after which the attention barely changes. The corresponding gate maps for the same input can be seen in Figure 14
Figure 12:
26

---

Published as a conference paper at ICLR 2022
Figure 13: Attention maps for head 13 of the NDR in every computational step on ListOps. This head shows the operands for each operation. Following it, we observe the hierarchy and the order in which the operations are performed.
27

---

Published as a conference paper at ICLR 2022
Figure 14: Gates for every computational step of the NDR on ListOps. Gates open for the deepest operations in the tree, processing proceeds upwards in the computational tree. The input has only depth 4, which explains the early stopping of the computation, roughly after 8-9 steps. The attention maps for the same input can be seen in Figure 12
28


Reviews:
[{"review_id": "mUpvvUH7uRj", "reviewer": "Reviewer_zmLb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe paper is very clearly written, and proposes an interesting solution to an important question. The tasks chosen are meaningful, and the experimental results suggest that the proposed architecture can solve the extrapolation problem. The technical aspects are precisely documented, which makes the research easy to reproduce. \n\nMy main concerns are related to the experimental comparisons, and the impact of certain design choices, such as the use of a fixed model depth at training and test time, the use of the last word in the encoder representation as the basis for model output, and the absence of the Adaptive Computation Time (ACT) in the Universal Transformer implementation that serves as the main baseline. This makes it difficult to judge the impact of the two improvements suggested (copy-gate and geometric attention), and the benefits of the new architecture, compared to an encoder-only state-of-the-art version of the Universal Transformer (with relative positional embedding, and ACT). I believe improving this part of the experimental design and discussion would greatly reinforce the paper. Below are my concerns and questions, split into four themes. \n\n*Computational, and model, depth*\n\nAt the beginning of section two, the authors argue that four properties are needed for network to extrapolate to larger problems: \n- shared layers\n- depth of the computational graph\n- step skipping\n- short-range attention\n\nI would disagree with the second point, for two reasons. First, computational depth is a relative notion. In an arithmetic task, I can choose to represent modular addition as one operation, or two (addition and modulo), or even three (digit addition, carry propagation, modulo). On the other hand, some linear algebra packages define \"add and mul\" as a single operation. There is no doubt that network depth should somehow increase with complexity, but defining it from computational depth seems unpractical. \n\nSecond, since you use shared layers, model depth can be varied without having to retrain. Specifically, model depth could be adjusted to the complexity of the training examples, and then increased at inference to fit the complexity of the test set. Using the maximum depth in both train and test sets (provided it can be defined from computational trees) is not necessary and might not even be beneficial. \n\nIn a recent paper (https://arxiv.org/abs/2106.04537), Schwarzschild et al. have shown (using different architectures, and testing on different tasks) that adding iterations during inference could help models extrapolate to larger problems. It would be interesting to test this on TCF (and baselines).\n\n*Copy-gating and variable depth*\n\nThe original Universal Transformers paper proposes a copy-gating mechanism, which uses the Adaptive Computation Time (ACT) mechanism (Graves 2016) at the token level. The gating works differently than in TCF: all gates begin closed, and once opened remain so. However, I believe this (universal transformer plus token-level ACT gating) is the correct baseline for TCF. Can such a comparison be provided? This is all the more important as gating has a large impact on performance, for the three problems considered. \n\nACT-gating has another merit: it adaptively controls the depth of the Universal Transformer, which goes on iterating until all gates are open. This means that the model can adjust to longer sequences by iterating for more \"ponder time\". This would provide an adaptive solution to the depth adjustment problem discussed in the previous section. Do you think an adaptive control for the number of iterations/layers could be implemented? \n\n*Encoder-only output, and geometric attention* \n\nThe original Universal Transformer is a sequence to sequence model. When decoding a solution, all the output sequence of the encoder is attended to. In your implementation, only the last element in the output is taken into account. As you observe in the results discussion of section 3.1, this creates problems at test time because the position of the last word changes as sequence length increases. It also complicates training, since the output position depends on the sequence length (which varies in all the problems you propose). As you show, this can be alleviated by relative positional embeddings and directional encodings, which can be used to force the result to \"move right\" as the computation proceeds. It also seems to be the main justification of geometric attention (which seems to bring very little when used alone, cf table 1).\n\nBut could the original problem, variable output position, be eliminated? What would happen if the output position had a fixed positional embedding? This could be done in many ways: reading the output from the first position instead of the last (since the transformer is bidirectional, this should have no adverse effect), or from some other fixed value (e.g. the fifth output word), or enumerating positions so that the last token has a fixed embedding (e.g. counting backward, or from both ends to the center). \n\nAnother question is the use of a single output element to decode the solution. Would the model be improved by using a longer part of the output sequence (e.g. the N first output words, with N the minimal size of output sequences, or shorter output padded to this size)?\n\nAn alternative (and in my opinion much better) solution would be to use a special attention-based layer for the decoder (an attention plus a linear layer working from the output sequence of the encoder). This amounts to a minimal seq2seq model, with one non-shared, cross-attention-only layer in the decoder. This would eliminate the variable output position problem, and allow the full output sequence to be taken into account while decoding. \n\nI believe these alternative decoders need to be tested. Without them, it is hard to assess the importance of relative positions and geometric  attention.\n\n*Compositional table lookup: the backward case*\n\nUnless their architectures are bidirectional, the backward task is very unfair on LSTM and DNC, which are causal models. To solve the backward task, they would need to memorize all the tables before seeing the value to be operated on, an impossible task given their capacity. \n\nTransformers, on the other hand, are bidirectional. Their bad performance on the IID backward case comes as a surprise, and the unusually large error on this observation suggests an experimental problem. Do you have explanations about this high standard deviation of experimental results? \n\nOn the test data, the fact that relative positional embeddings seem to improve the forward, but not the backward case, might be due to the choice of the last term in the output as the result to be decoded. Would, for instance, the results be inverted if the first output were chosen instead (or some fixed middle position)? \n\nOverall, I am not certain this backward case helps the argument about the efficiency of TCF (what it demonstrates, I think, is that causal models like the LSTM need their inputs to be presented in a particular order, which is no new news...).\n", "summary_of_the_paper": "\nThe 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.\n\nAs 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. \n\nTwo new features are introduced : \n- 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,\n- 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).\n\nExperiments are conducted over three tasks:\n- predicting the output of sequences of permutations of 8 elements, in prefix or postfix notation,\n- predicting the result of of additions and multiplications modulo 10, in infix notation,\n- predicting the result of operations on lists of small integers, in prefix notation.\n\nFor each task, TCF is shown to be capable of extrapolation to larger problems (i.e. longer sequences) than those seen at training.", "summary_of_the_review": "\nTCF appears to be a promising architecture, and the experimental results seem good. However, some design and experimental choices make them difficult to compare with the original Universal Transformers. In particular, the choice of a fixed depth at train and test time might be an unnecessary constraint, the proposed gated-copy operation should be compared with the ACT-based copy described in the UT paper, and the use of the last element in the output  sequence as the basis for decoding might unduly increase the necessity of relative positions and short-range attention (ie the geometric attention and directional encoding recommended by the authors). \n\nHence my note of 6, which I would gladly increase if additional experimental results are provided. \n\nEdit: Thank you very much for the detailed response. I have raised my note to 8.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "lQFp13tKreq", "reviewer": "Reviewer_rpDa", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Main strengths:\n1) The paper is well written and does a great job in introducing the problem and revealing the flaws of the universal transformer in achieving good performance in the described tasks, as well the authors' intuition of the properties of a good solution.\n2) The main components of the proposed method are explained in sufficient details to help reproducing the proposed method.\n3) The proposed benchmarks and datasets, the empirical approach, and chosen hyperparameters are provided and discussed in details.\n4) The paper is well positioned with regards to the related work.\n\n## Main Weaknesses:\n5) It is not clear if the the considered benchmarks cover all required aspects of task generalization, or the generalization is only valid for tasks that are to some extent similar to the considered experiments. The authors should further explain which aspects, if any, are missing and are not addressed in this work.\n6) It is not clear if the considered assumptions are always necessarily and correct. The authors should address the following questions in the paper either in form of justified explanations or if required with ablation studies:\n\n\t6.1) Is there any task which would benefit or require settings that are not covered by the considered settings described in section 2?\n\n\t6.2) Regarding point 2 in section 2: What if the data dependency graph was too long that memory complexity would not practically allow to use such a depth? In other words, to what extent the proposed depth is necessary?\n\n\t6.3) Regarding point 3 in section 2: Could the gating function result in a shortcut/collapse in optimization? (Considering a far more complex task that is generally addressed by transformers could reveal such issues.)\n\n\t6.4) Regarding the final point in section 2: Could a task would prefer non-local operations to local ones? Does the performance of the proposed method degrade in that situation?\n\n7) Some previous works, for example on the ListOps task, consider sequences that are orders of magnitude longer than the ones considered in this paper (A couple of examples are [1], [2]). It is not clear if the claim that previous results did not achieve the perfect accuracy is well-supported? It seems like that to be fair, the authors should have considered some of the SOTA methods and adapt their hyperparameters for these tasks with limited sequence length before testing how they would perform.\n\n### Minor points:\n8) In section B.2, the set of values or ranges over which the hyperparameters are searched should also be mentioned.\n9) First line of page 14. \"sample\" -> \"sampled\"\n\n### References:\n\n[1] \"Modeling Hierarchical Structures with Continuous Recursive Neural Networks\" by Chowdhury, J.R. and Caragea, C., (arXiv:2106.06038v1)\n\n[2] \"Nystr ̈omformer: A Nystr  ̈om-based Algorithm for Approximating Self-Attention\" by Xiong et al., (arXiv:2102.03902v3)", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is well organized and covers the background knowledge required to follow the discussions. It motivates the goal, provides justifications for the choices made in proposed methods, follows a thorough empirical approach, and achieves near perfect results in the considered benchmarks. There are a few points of concerns that need to be cleared before I can fully support the submission. Regardless, I see many of the properties of a good research project and so I lean towards accepting the paper at this stage. I look forward to authors' feedback on my concerns before I finalize my decision.\n\nEdit: I thank the authors for responding to my comments in details. After reviewing their response and the changes made in the paper, many of my concerns are resolved. Therefore, I change my recommendation to accept this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "8PAcvFxXA2_", "reviewer": "Reviewer_TUV2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I enjoyed reading this work. The writing is clear and to the point, and the approach itself is very well motivated (see questions for more), and simple to implement without too many tunable hyperparameters. And as such, the experiments are cleanly setup and do suggest improved improved generalization. From the analysis of the attention maps, we can see that the method is doing exactly what it is supposed to as well (that is, copying previous values until other intermediates have been computed, paying more attention to local hidden states etc). Based on these strengths, I recommend that this paper be accepted to the conference. \n\nSo now, let me focus on some weaknesses / suggestions / questions: \n\nOverall positioning: Firstly, I think the paper should probably make it more clear that it’s only focusing on a very specific notion of systematicity that has to do with length / depth generalization, and not other more traditional notions like generalizing to new compositions (which isn’t really something that is evaluated) like SQOOP from Bahdanau 2019 etc. \n\nEvaluation: Secondly, while not a strict requirement, there is no evaluation on language tasks / pseudo language tasks like SCAN - there is a length generalization benchmark within SCAN itself and it would be good to know how this method does on that. \n\nAnalysis: In Figure-2 bottom, what does It is unclear what the y axis is. Isn’t the copy gate just a single number for each time step, for each layer? If so, i would’ve expected the figure to just be a single number for each time step for the various layers, so I don’t understand what the grid signifies.\n", "summary_of_the_paper": "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.\n", "summary_of_the_review": "Overall, I think based on the results I recommend that the paper be accepted into the main conference. The problem is well motivated, the comparisons are fair and the results compelling though there is some scope for improvement that i highlight in the weakness section of the main review", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work proposes a novel Transformer Control Flow model and achieves near-perfect accuracy on length generalization, simple arithmetic tasks, and computational depth generalization.  All reviewers give positive scores. AE agrees that this work is very interesting and has many potentials. It would be exciting if the author could extend this framework to more challenging tasks (e.g. visual reasoning [1. 2]).  Given the novelty of the proposed model, AC recommends accepting this paper!

[1]  CLEVR: A Diagnostic Dataset for Compositional Language and Elementary Visual Reasoning. ICCV 2017.

[2] PTR: A Benchmark for Part-based Conceptual, Relational, and Physical Reasoning, NeurIPS 21
Decision: Accept (Poster)