Learning Adaptive Control Flow in Transformers for Improved Systematic Generalization
Abstract
Despite successes 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 ﬁnd 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 ﬂow, we propose two modiﬁcations to the Transformer architecture, copy gate and geometric attention.
Our novel Neural Data Router (NDR) achieves 100% length generalization accuracy on the compositional table lookup task.
NDR’s attention and gating patterns tend to be interpretable as an intuitive form of neural routing .
Introduction
Systematic generalization [1] is a property of learning systems that enable them to generalize beyond their training distribution, by learning compositional rules that are generally applicable.
Examples include generalization to sequences much longer than those seen during training and algorithmic combinations of known rules.
Despite recent efforts [2, 3, 4, 5, 6, 7], systematic generalization generally remains unsolved [8, 9, 10, 11, 12].
On some datasets, the best performing models are neuro-symbolic hybrids [13, 14] which make use of task speciﬁc symbolic functions.
However, their applicability to other datasets remains limited [15, 16].
Typically many NNs solve the training set, but only few of them generalize to systematically different test sets.
A big question is: which type of external pressure or architectural inductive bias would encourage the training process to select a “real” solution which systematically generalizes?
While some approaches such as meta-optimization [17] address this issue, the gain is typically limited.
The popular Transformer [18] architecture also suffers from such problems [19, 7, 20].
Frustratingly,
Transformers fail to generalize in many algorithmic tasks (e.g. [10, 21]), even those that have intuitive solutions which can be simply expressed in terms of attention patterns in Transformers.
Indeed, given an input sequence of length N and a Transformer encoder of depth T , solving an algorithmic task is often all about routing the right information to the right node/operation at the right time in the T -by- N grid represented by Transformer columns.
Effectively the task is to learn to draw an adaptive control ﬂow on the canvas of Transformer columns.
In fact, a recent work by Weiss et al. [22] introduced a programming language called RASP which is speciﬁcally designed to express solutions for sequence processing problems, and which has a direct equivalent to the operations in
Transformers.
However, it is also shown that making Transformers learn a solution expressed in
RASP is only possible with intermediate supervision of attention patterns.
In some cases, even such supervision fails [22].
In general, Transformers fail to ﬁnd solutions of algorithmic tasks which are easily interpretable and/or symbolic.
We conversely hypothesize that attention-based NNs able to ﬁnd intuitive solutions (achieving interpretable attention patterns) can improve systematic generalization.
Here we argue that regular Transformers lack a few basic ingredients for learning such “intuitive” solutions to algorithmic problems.
We propose simple architectural modiﬁcations to help them learn data routing and easily interpretable solutions.
As a ﬁrst step towards validating our model, we focus on the popular length generalization task of compositional table lookup (CTL) [10, 23, 21].
Our model achieves 100% generalization accuracy on this task, which is previously unsolved by NNs [21], and we show that the resulting attention patterns are interpretable as plausible control ﬂows.
Improving Transformers for learning adaptive control ﬂow
We argue that the following components are needed to build Transformers capable of learning adaptive control ﬂow.
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 [24].
Second , the network should be sufﬁciently deep, at least as deep as the deepest data dependency in the computational graph (e.g., in case of a parse tree, this is the depth of the tree).
Otherwise, multiple operations would be fused into a single layer and hinder natural and elegant compositions.
Third , inputs in some columns should be kept unchanged until the right point in time when they can be processed in a later step/layer.
The regular Transformer lacks a mechanism for skipping the whole transformation step by simply copying the input to the next step/layer.
We propose a special gating function, copy gate , to implement such a mechanism (Sec. 2.1).
Finally , many algorithmic tasks require combining a number of local computations in the right order.
This typically implies that at a given time, attention should not focus on all possible matches, but only on the closest match.
We propose and investigate a new type of attention with the corresponding inductive bias (Sec. 2.2), which we call geometric attention .
We refer to the resulting new Transformer as the Neural Data Router (NDR).
2.1 Copy gate: learning to skip operations ( vertical ﬂow)
Each layer of the regular Transformer consists of one self-attention and one feedforward block.
The input to each of these blocks are directly connected to the corresponding output via a residual connection [25, 26].
However, such a connection does not allow for skipping the transformation of the entire layer and simply passing the unchanged input to the next layer.
Here we propose to add an explicit gate, which we call copy gate , to facilitate such a behavior.
We consider a T -layer Transformer encoder and an input sequence of length N .
Since each layer corresponds to one computational step , we often refer to a layer as a step t .
We denote the Transformer state of column i in layer t as h ( i,t ) =
H t,i ∈
R
d where d is the state size, and H t ∈ R N ×
d denotes the states of all N columns in layer t .
In the copy gate-augmented Transformer, each column i in layer ( t + 1) processes the input H t like in the regular Transformer:
[Equation]
but the output is gated as:
[Equation]
We use the basic two-layer feedforward block [18] for both FFN data and FFN gate which transforms input x ∈ R d to:
[Equation]
but with separate parameters and different dimensionalities: for FFN data
[Equation]
data [Equation] R d × d FF , while for FFN gate W
[Table]
, W
[Table]
When the gate is closed i.e. g ( i,t +1) = 0 , the entire transformation is skipped and the input is copied over to the next layer
h
( i,t +1)
=
h ( i,t ) .
Crucially, we parameterize the gate (Eq. 3) as a function of the output of the self-attention (Eq. 1), such that the decision to copy or to transform the input for each column depends on the states of all columns.
This is a crucial difference compared to previously proposed gating in Transformers which are solely motivated by the training stability [27] or a common practice from convolution based models [19].
2.2 Geometric attention: learning to attend to the closest match (horizontal ﬂow)
We propose geometric attention designed to attend to the closest matching element.
Like in regular self-attention, given an input sequence [ x (1) , x (2) , ..., x ( N ) ] with x ( i ) ∈ R d in , each input is projected to key k ( i ) ∈ R d key ,
value v ( i ) ∈ R d value , query q ( i ) ∈
R d key vectors, and the dot product is computed for each key/query combination.
In our geometric attention, the dot product is followed by a sigmoid function to obtain a score between 0 and 1:
[Equation]
which will be treated as a probability of the key at (source) position j matching the query at (target) position i .
These probabilities are ﬁnally converted to the attention scores A i,j as follows:
[Equation]
where S i,j denotes the set of all (source) indices which are closer to i than j to i , and when two indices have the same distance to i , we consider the one which is to the right of i (i.e., greater than i ) to be closer, i.e.,
[Equation]
In addition, we explicitly zero-out the diagonal by setting A i,i = 0 for all i = 1 , ..., N .
The ordering of source indices is illustrated in Figure 1/Right.
The resulting scores
A i,j are the attention scores used to compute the weighted average of the value vectors.
By using the term (1 − P i,k ) in Eq. 7, when there is a match, it downscales any other matches which are more distant.
Two recent works [28, 29] use such a parameterized geometric distribution in the form of Eq. 7.
In practice, Eq. 6 is augmented by an additional direction encoding (See Appendix B). [Figure]
E
E a b b c S
[Figure]
c c
E
E t [Figure]
t [Figure] t
[Figure] t [Figure] t
[Figure]
Experiments
Task.
Compositional table lookup task [10, 23, 21] is constructed from a set of symbols and unitary functions deﬁned over these symbols.
Each example in the task is deﬁned with one input symbol and a list of functions to be applied sequentially, i.e. the ﬁrst function is applied to the input symbol, and the resulting output becomes the input to the second function, and so forth.
There are 8 possible symbols.
Each symbol is traditionally represented by a 3-bit bitstring [10], however, in practice, they are simply processed as one token [21].
The functions are bijective and randomly generated.
Each function is represented by a letter.
An example input is ‘ 101 d a b ’, which corresponds to the expression b ( a ( d (101))) ; the model has to predict the correct output symbol.
We note that there exists a sequence-to-sequence variant of this task [21] where the model has to predict all intermediate steps (thus trained with intermediate supervision).
We directly predict the ﬁnal output.
An ideal model should be able to solve this task independently of the presentation order, i.e., whether the task is encoded as ‘ 101 d a b ’ or ‘ b a d 101 ’.
We thus study both forward (former) and backward (latter) variants of the task.
To evaluate systematic generalization, the train/valid/test sets have different numbers of compositions: samples of up to 5/6-8/9-10 operations respectively.
Results.
We consider four different baselines: an LSTM [30], DNC [31, 32], Universal Transformers [18, 24] and its variants with relative positional encodings [7].
For Transformers, the prediction is based on the last column in the ﬁnal layer.
Results are shown in Table 1.
The LSTM and DNC perform well in the forward variant, achieving perfect generalization for longer sequences, but fail on the backward variant.
In contrast, basic Transformers fail in both cases.
By introducing the copy gate (Sec. 2.1), the relative Transformer is able to solve the forward, but not the backward task.
Further 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 remediate this issue, we introduce a simple mechanism to let the model choose between absolute and relative positional encodings at each position (see Appendix A).
The resulting model effectively manages to make use of the absolute position for the prediction, and to perform well in both directions.
However, such a combination of absolute/relative positional encoding might be an overly speciﬁc bias.
A more generic solution, geometric attention (Sec. 2.2), also achieves perfect generalization, and was found easier to train.
[Table]
0.19 ± 0.04 + abs/rel + gate 1.00 ± 0.
1.00 ± 0.
0.98 ± 0.
0.98 ± 0.03 + geom.
att. 0.96 ± 0. 0.93 ± 0. 0.16 ± 0.
0.15 ± 0.02 + geom.
att.
+ gate (NDR) 1.00 ± 0.
1.00 ± 0.
1.00 ± 0.
1.00 ± 0.
Analysis.
Figure 2 shows the gating and attention patterns of our model with both copy gate and geometric attention for an example of the backward presentation task.
As shown in Fig.
Right, the gates of different columns open sequentially one after another, when the input is available for them.
Fig.
2/Left shows the attention map.
Each column attends to the neighboring one, waiting for its computation to be ﬁnished.
The behavior of the last column is different: it always attends to the second position of the sequence, which corresponds to the last operation to be performed.
[Figure]
Conclusion
We proposed a new view on the internal operations of Transformers, introducing a dynamic dataﬂow architecture between Transformer columns.
This overcomes two shortcomings of traditional Transformers: the problem of keeping and routing data in unaltered fashion, which we solve by an additional copy gate, and the problem of learning length-independent attention patterns, which we solve by geometric attention.
Our new model, NDR generalizes to longer lengths on the popular compositional lookup table task in both forward and backward directions.
The gates and the attention maps collectively make the architecture more interpretable compared to the baselines.
Acknowledgments and Disclosure of Funding
We thank Imanol Schlag and Sjoerd van Steenkiste for helpful discussions and suggestions on an earlier version of the manuscript.
This research was partially funded by ERC Advanced grant no: 742870, project AlgoRNN and by Swiss National Science Foundation grant no: 200021_192356, project NEUSYM.
We are thankful for hardware donations from NVIDIA & IBM.
The resources used for the project were partially provided by Swiss National Supercomputing Centre (CSCS) project s1023.
A Details of attention with absolute/relative positional encoding combination
The addition of copy gates enables Transformers to generalize to longer lengths in the forward presentation order, but not in the backward order.
Examining the attention maps reveals that the model uses position-based attention for reading out the result, as opposed to content-based attention.
In the backward presentation order, the last column of the transformer should focus on the second column, whose relative position changes dynamically with the length of the sequence.
We solve this issue by adding an option to the attention head to choose between absolute and relative positional encodings.
We use the relative positional embedding variant of self-attention Dai et al.
In what follows, we describe the operation within a single layer/step.
We thus omit the layer/step index t for better readability, and denote the state at column/position i as h i instead of h ( i,t ) .
Our attention matrix is decomposed as follows:
[Equation]
where the matrix W q ∈ R d head × d maps the states to queries, W k,E ∈ R
d head × d maps states to keys, while W k,P ∈ R d head × d maps positional embedding to keys.
d head is the size of the key, query and value vectors for each head, set as d head =
d n heads .
u , v ∈ R d head are learned vectors.
Softmax is applied to the second dimension of ˆ A to obtain the ﬁnal attention scores, A .
Component (a) corresponds to content-based addressing, (b, e) to content based positional addressing, (c) represents a global content bias, while (d, e) represents a global position bias.
We introduce the term (e) for the positional embedding which can switch between absolute and relative positional encodings using the gate r i
(Eq. 9; parameterized by W ar ∈ R d × 1 and b ar ∈ R ), which is the function of the state at target position i .
p i ∈
R d is a sinusoidal embedding for position i , which can be can be both positive and negative.
Following Vaswani et al.
[18], we scale A i,j by√ d head prior to applying the softmax, and we deﬁne p i,j as:
[Equation]
B Details of geometric attention
In Sec. 2.2, we introduce a basic form of geometric attention (Eqs.
In practice, Eq. 6 can be augmented by an additional directional encoding .
In fact, the only positional information available in the geometric attention presented above is the ordering used to deﬁne the product in Eqs.
In practice, we found it crucial to augment the score computation of Eq.
6 with an additional directional information .
The directional information we introduce is encoded as a scalar D i,j for each target/source position pair ( i, j ) computed as:
[Equation]
(12) where h i ∈ R d denotes the input/state at position i and W LR , W RL ∈ R 1 × d , b LR , b RL ∈ R are trainable parameters.
This directional information is integrated into the score computation of Eq.
6 as follows (akin to how Dai et al.
[33] introduce the relative positional encoding as an extra term in the computation of attention scores):
[Equation]
(13) where the matrix W q ∈ R d head × d maps the states to queries,
u ∈ R d head is a bias for queries,
W k,E ∈ R
d head × d maps states to keys (we note that d head is typically the size of the key, query and value vectors for each head, d head =
d n heads ), and α, β, γ ∈ R are learned scaling coefﬁcients and bias, initialized to α = 1 √ d head , β = 1 , γ = 0 .
Using this additional directional information, each query (position i ) can potentially learn to restrict its attention to either the left or right side.
The proposed attention function has a complexity of O ( N 2 ) , similar to the one of regular selfattention.
Eq.
7 can be implemented easily in a numerically stable way in log-space.
Products can then be calculated using cumulative sums, subtracting the elements for the correct indices in each position.
C Implementation details
A PyTorch implementation of our models together with the experimental setup is available under https://github.com/after_review .
The performance of all models are reported as mean and standard deviation over 5 different seeds.
C.1 Dataset details
Our implementation uses 8 symbols (all 3 bit patterns) as input arguments and 9 randomly sampled bijective functions denoted by small 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 number of samples for each depth.
There are different validation sets: an IID, which matches the distribution of the train set, and a depth validation, which includes samples of length 6, 7 and 8.
The test set consists of sequences of length 9 and 10.
C.2 Model details
We use the AdamW optimizer [34] for all of our models.
Standard hyperparameters are listed in
Tab. 2.
Additionally, models with gating use dropout [35, 36] of 0.1 applied to the content-based query and the position-query components.
In the case of geometric attention, since the channels of the positional encoding does not have any redundancy, dropout is applied just to the content-query.
In the case of Transformers with the copy gate but without geometric attention, we use tanh instead of LayerNorm in Eq. 2.
The hyperparameters of the gateless Transformers differ signiﬁcantly from the gated ones.
This is because they are very hard to train to achieve good performance even on the IID set, and they needed extensive hyperparameter optimisation.
One might argue that fewer layers (11 instead of 14) makes them less competitive on longer sequences, however, we were unable to train them to perform well even on the IID data with 14 layers.
Similarly, increasing the size of the baseline Transformer models made them fail to train.
All Transformer variants have a begin (B) and end (E) token included on sequence boundaries.
RNNs (LSTM and DNC) have no such addition.
All Transformers are encoder-only, and the result is read from the last column (corresponding to the end token).
A single feedforward layer maps it to theclasses corresponding to the 8 possible output symbols.
The DNC has 21 memory cells, 4 read heads and an LSTM controller, and contains recently introduced improvements [32].
Bias of the gates b gateare initialized to − 3 .
This ensures that initially no updates happen, and creates a better gradient ﬂow between layers.
We use gradient clipping with magnitude 5 for all of our models.
Hyperparameters were obtained by a Bayesian hyperparameter search of Weights & Biases 2 over the systematically different validation set for the +abs/rel + gate models and was reused for all other gated models.
For the non-gated models we used the +rel variant for tuning.
This failed to yield to any usable conﬁgurations, so we used the mixture of the IID and systematically different validation set to obtain the parameters in Tab. 2.
We train all models to a ﬁxed number of n iters iterations and measure its validation performance every 1000 iterations.
For each model, we select the best checkpoint according to the validation performance, and report its test accuracy.
[Table]
D Additional analysis
In the main text [Table] we only had space to show the gate and attention activity of a gated network with absolute/relative positional encodings in Figure 4 and Fig. 5.
Here we show the corresponding visualization for the geometric attention variant in Figures 6 and 7, as well as the attention map for baseline Transformer with relative positional encoding in Figure 3.
Please directly refer to the caption of the ﬁgures for the corresponding analysis.
In general, the visualization for our geometric model variant is interpretable unlike that of the baseline Transformer model.
One might wonder how geometric attention is able to focus next to its diagonal, if there is no positional embedding.
By combining direction sensitivity with the bias, it is possible to focus on such elements.
Even in cases where this is not possible, there might be alternative solutions relying solely on direction and content.
B b i b d g g a h e c 100 E
Bbibdggahec 100E t =
B b i b d g g a h e c 100 E
Bbibdggahec 100E t =
B b i b d g g a h e c 100 E
Bbibdggahec 100E t =
B b i b d g g a h e c 100 E
Bbibdggahec 100E t
[Figure]
B b i b d g g a h e c 100 E
Bbibdggahec 100E t
[Figure]
B b i b d g g a h e c 100 E
Bbibdggahec 100E t
[Figure]
B b i b d g g a h e c 100 E
Bbibdggahec 100E t
[Figure]
B b i b d g g a h e c 100 E
Bbibdggahec 100E
[Figure]
B b i b d g g a h e c 100 E
Bbibdggahec 100E
[Figure]
B b i b d g g a h e c
[Figure]
E
Bbibdggahec 100E
[Figure]
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h d h a d110 E
Bdbbibhdhad 110E t =
B d b b i b h
d h a d 1 10 E
Bdbbibhdhad 110E t =
B d b b i b h
d h a d 1 10 E
Bdbbibhdhad 110E t = 0.
B d b b i b h d h a d 11 0 E t = B d b b i b h d h a d 11 0 E t =
B d b b i b h d h a d 11 0 E t = B d b b i b h d h a d 11 0 E t =
B d b b i b h d h a d 110 E t = B d b b i b h d h a d 110 E t =
B d b b i b h d h a d 110 E t = B d b b i b h d h a d 110 E t =
B d b b i b h d h a d 110 E t = B d b b i b h d h a d 110 E t =
B d b b i b h d h a d 110 E t = B d b b i b h d h a d 110 E t =
B d b b i b h d h a d 110 E t =
B d b b i b h d h a d 110 E t = 0.
The attention maps for the same input can be seen on Figure 4.
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a001 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a 0 01 E
Bfdccfhaiba 001E t =
B f d c c f h a i b a 0 01 E
Bfdccfhaiba 001E t = 0.
B f d c c f h a i b a 00 1 E t =
B f d c c f h a i b a 00 1 E t =
B f d c c f h a i b a 00 1 E t =
B f d c c f h a i b a 00 1 E t =
B f d c c f h a i b a 001 E t =
B f d c c f h a i b a 001 E t = B f d c c f h a i b a 001 E t =
B f d c c f h a i b a 001 E t = B f d c c f h a i b a 001 E t =
B f d c c f h a i b a 001 E t = B f d c c f h a i b a 001 E t =
B f d c c f h a i b a 001 E t = B f d c c f h a i b a 001 E t =
B f d c c f h a i b a 001 E t = 0.