Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
main_review: ## Main strengths:
1) 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.
2) The main components of the proposed method are explained in sufficient details to help reproducing the proposed method.
3) The proposed benchmarks and datasets, the empirical approach, and chosen hyperparameters are provided and discussed in details.
4) The paper is well positioned with regards to the related work.

## Main Weaknesses:
5) 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.
6) 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:

	6.1) Is there any task which would benefit or require settings that are not covered by the considered settings described in section 2?

	6.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?

	6.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.)

	6.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?

7) 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.

### Minor points:
8) In section B.2, the set of values or ranges over which the hyperparameters are searched should also be mentioned.
9) First line of page 14. "sample" -> "sampled"

### References:

[1] "Modeling Hierarchical Structures with Continuous Recursive Neural Networks" by Chowdhury, J.R. and Caragea, C., (arXiv:2106.06038v1)

[2] "Nystr ̈omformer: A Nystr  ̈om-based Algorithm for Approximating Self-Attention" by Xiong et al., (arXiv:2102.03902v3)
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.

Edit: 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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.

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. 

So now, let me focus on some weaknesses / suggestions / questions: 

Overall 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. 

Evaluation: 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. 

Analysis: 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.

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
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: 
The authors propose Transformer Control Flow (TCF), a set of improvements to the Universal Transformer (Dehghani et al, ICLR 2019). They show that, for three compositional problems, TCF allows trained models to generalize to longer sequences, a common problem of many transformer implementations.

As in the Universal Transformer (UT), the encoder consists of one shared transformer layer (self attention + fully connected network) which is iterated through a fixed number of times, by feeding the output of each iteration back into the input of the shared layer. However, whereas the UT uses a sequence to sequence model, TCF is an encoder-only architecture, which decodes the last element in the output sequence as the final result. 

Two new features are introduced : 
- a gating mechanism that allows the model to "skip a layer" (the input is then copied to the output), on the basis of the self-attention output,
- a weighting system for the outputs of attention heads, which favors short-range attention (i.e. tokens close to the one currently considered), and can be trained to be biased towards a certain direction (before or after the current token).

Experiments are conducted over three tasks:
- predicting the output of sequences of permutations of 8 elements, in prefix or postfix notation,
- predicting the result of of additions and multiplications modulo 10, in infix notation,
- predicting the result of operations on lists of small integers, in prefix notation.

For each task, TCF is shown to be capable of extrapolation to larger problems (i.e. longer sequences) than those seen at training.
main_review: 
The 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. 

My 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. 

*Computational, and model, depth*

At the beginning of section two, the authors argue that four properties are needed for network to extrapolate to larger problems: 
- shared layers
- depth of the computational graph
- step skipping
- short-range attention

I 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. 

Second, 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. 

In 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).

*Copy-gating and variable depth*

The 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. 

ACT-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? 

*Encoder-only output, and geometric attention* 

The 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).

But 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). 

Another 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)?

An 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. 

I believe these alternative decoders need to be tested. Without them, it is hard to assess the importance of relative positions and geometric  attention.

*Compositional table lookup: the backward case*

Unless 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. 

Transformers, 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? 

On 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)? 

Overall, 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...).

summary_of_the_review: 
TCF 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). 

Hence my note of 6, which I would gladly increase if additional experimental results are provided. 

Edit: Thank you very much for the detailed response. I have raised my note to 8.
review_title: null

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)