Paper: How Capable Can a Transformer Become? A Study on Synthetic, Interpretable Tasks (ICLR-2024)

Reviewer comment (Experiments):
The experimental setup lacks sufficient clarity, making it challenging to understand the specific procedures.

Author response:
Thank you for raising this concern. We have significantly expanded Section 3.1 and Appendix A to include more details. We also commit to releasing the code.

Reviewer comment (Evaluation):
The paper's conclusion regarding the significance of later attention layers is not adequately substantiated by the experimental data; it lacks insight.

Author response:
We have added two new experiments to further substantiate this claim. 

In the first experiment, we visualize the attention maps of a Transformer on the step-by-step prompt. In Figure 7(right), we observe that the attention map selects two tokens, 1 task token (which picks the function) and 1 input token which is the intermediate result of a composition. Hence, attention seems to play an important role in selecting the function and the input that the function should be applied to.

In the second experiment, presented in Appendix B.4 (figure 14), we study the attention layers of transformers of different depths (1 to 60 layers). Across most Transformers, we find that compositions occur in the later layers of a Transformer. We also observe that a sharp upward inflection in accuracy in all the plots occurs after an MLP layer. The attention layers seem to select the input and the function to apply and the computation seems to occur in the MLP. 

Despite the strong evidence, these claims are still speculative and we have made sure to convey the same in the manuscript.

Reviewer comment (Writing):
Section 3.2 lacks a clear definition and differentiation between permutations and bijections. It should be explicitly stated that the results in Section 4.1 pertain only to bijections.

Author response:
We have added details to the section (we note it is now section 3.1) and to appendix A.1 where we describe the set of permutations and bijections. We have also clarified in sub-section 4.1 that the results only pertain to bijections.

Reviewer comment (Presentation):
Figure 3 needs more explanation about the notations. For instance, what $f_{i-j}$ means?

Author response:
We have tried to explain the notation further in the updated manuscript and have also simplified it for ease of understanding. Specifically, Figure 2(a) now explains the role of $i$ and $j$ for a function $F^{(i)}_j$. The superscript $i$ determines the position in the composition for an in-order composition. For example in a composition of 3 functions, the function $F^{(2)}_3$, must be in the second position for an in-order composition, i.e., an in-order composition of 3 functions should look like $g \circ F^{(2)}_3 \circ h$ since $i=2$.

The index $j$ is used to iterate over the set of all possible functions that are allowed to be at position $i$. For example $F^{(2)}_1$, $F^{(2)}_2$ and $F^{(2)}_3$ (j=1,2,3) are all functions that can appear in the second position in an in-order compositions. All of these functions have the same color in Figure 3.

We have improved the writing in Section 3.1 to complement Figure 3. We would be glad to further clarify if the reviewer believes that the details are still not entirely clear.

Reviewer comment (Experiments):
The experimental setup for Figure 6(b) in Section 4.3 is not elucidated. What is the rationale for having 25 bijections and 25 permutations?

Author response:
We would be grateful if the reviewer could clarify which specific details are unclear in Figure 6(b). The experimental setup (architecture, training methodology) is identical to all other experiments in the paper, as discussed in the beginning of the section. Specifically, Figure 6(b) considers the composition of two functions — one of which is a bijection and the other is a permutation.

Since we consider the composition of only two functions, we increase the number of function choices at each position to 25. This results in a total of 625 compositions of functions. If we instead considered only 5 bijections and 5 permutations, then the total number of compositions is only 25, which would make our evidence for compositional generalization weaker.

Reviewer comment (Theory):
The paper claims that adding permutations aids direct composition; this is a counterintuitive finding that warrants further elaboration.

Author response:
We believe there is a misunderstanding and would like to clarify that this is not the exact claim that we have made in the paper. We restate our result below.

The direct prompt format is successful if we train on compositions of two functions: a permutation function and a bijection function. Hence the function composition is an element from the set $\mathcal{F}_b \circ \mathcal{F}_p$. It is surprising that direct prompt format fails when we compose 2 functions which are both bijections, i.e. Transformers fail to generalize to compositions in the set $\mathcal{F}_b \circ \mathcal{F}_b$.

We do not have a precise answer for why this is the case and find this observation surprising. But we speculate why this could occur. Permutations operate on the position of a token and bijections operate on the value of a token; the position and value are (nearly) orthogonal features of the input. Okawa et al., Lewis et al. (https://arxiv.org/abs/2310.09336, https://arxiv.org/abs/2212.10537) make similar observations for vision models and show compositional generalization on an orthogonal set of attributes. We believe a similar phenomenon is at play when a Transformer is forced to compose in a single step.

Reviewer comment (Writing):
In page 2, remove one of the "considered" and the quotation masks are not correct.

Author response:
Thank you for bringing this to our attention.

Reviewer comment (Writing):
In Figure 11: "direction composition" should be corrected to "direct composition."

Author response:
Thank you for bringing this to our attention.

Reviewer comment (Novelty):
I don’t think that this paper makes a meaningful contribution to the field of deep learning on top of the already available work. The general question of whether Transformers or other neural architectures can generalize compositionally has already been discussed in numerous papers, including similar function compositional setups ([1, 2]). The general consensus in the literature is that with enough diversity, any neural architecture can learn any compositional behaviors.

Author response:
**>>> The general consensus in the literature is that with enough diversity, any neural architecture can learn any compositional behaviors.**

We disagree with the reviewer’s claim that any neural architecture can learn compositional behaviors with enough diversity. We present a new result in the manuscript where we train an LSTM on the same task as used for our Transformer models in both the direct and step-by-step prompt formats. LSTMs perform significantly worse compared to Transformers (~30% vs ~100%), highlighting that **the architecture has a large impact on compositional generalization for the same training data**.  

**We would also like to point out that “enough diversity” is ill-defined and our work also attempts to characterize the same.** We study how precisely defined properties of the data affect downstream generalization: we show how factors like spurious correlations in the training data (in-order compositions), details of the prompt format (step-by-step and direct), and the number of compositions in the training data have vastly different effects on compositional generalization. Importantly, we emphasize that diversity in the data does not necessarily guarantee compositional behavior. In our work, we find that training on 2000 of the 3125 bijections (Figure 6 (left)) does not result in compositional behavior if we use the direct prompt format. However we can compositionally generalize using the direct prompt format if we use permutations and bijections. **Hence, our results indicate that there are different types of compositional data and Transformers can fail to learn some of them even with sufficient diversity.** 

Finally, we would like to add that **Transformers trained with the step-by-step data format are able to compose with very little diversity** in the training data. In Figure 5, we can train on as few as 125 functions (**0.003% of the total number of functions**) but we can still generalize to a combinatorial set of functions not present in the training data (4 million of them). This is a surprising result and points to how data diversity, while important, can be rather minimal when the appropriate architecture and data format are used. 

**>>> I don’t think that this paper makes a meaningful contribution to the field of deep learning on top of the already available work.**

**>>> While the paper acknowledges that there is ample prior work on the topics, the paper fails to explain what it adds on top**

**>>> including similar function compositional setups ([1, 2])**

Thank you for pointing us to these papers. We have included them in our discussion in the Related work section. However, we hope to convince the reviewer of the importance of our work in the context of these prior works. 

Our work differs from prior work like [1, 2] (and several others that are reviewed in Section 2) in the following ways:


1. The goal of works like [1, 2] is to design new or improve existing architectures such that they compose better. While this is an important problem, we re-iterate that **our goal is different**: we would like to demonstrate that Transformers can learn to compose functions with very little training data and minimal data diversity, and we intend to understand what drives the existence of compositionality in autoregressive Transformers trained using the standard training pipeline. 
2. [1, 2] observe that Transformers and RNNs do not learn to compose on datasets like CTL and CTL++. However, **we present results that, on the contrary**, show that the **standard Transformer** architecture is capable of almost perfect compositional generalization. It would be interesting to understand why these differences exist and this warrants further exploration. We hypothesize that a reliable training pipeline with all the bells and whistles is largely responsible for this difference. No prior work has shown minimal demonstrations, that show this is possible for Transformers that perform stepwise computation–which is very natural in autoregressive models. 
3. Our **work evaluates functions outside of table lookups** and identifies that compositionality depends on the nature of the composed functions (permutations vs. bijections), which isn’t considered in prior works like [1, 2].
4. **The scale of our experiments is larger.** We evaluate these models systematically on as many as 4 million functions (even if the model was trained on as few as 100 of these functions) and take steps towards identifying precise conditions under which compositionality occurs and fails. In comparison some table-lookup tasks like [1, 2] are often limited to evaluation on just 128–256 functions

We also note that we have significantly expanded the Related work section to emphasize some of these points.

Reviewer comment (Experiments):
Further interesting questions can be asked: what architectures can learn compositional behavior from less diverse training data, can we get diverse enough training data to achieve compositionality in real-world tasks (GPT-4 seems to partially answer that), how compositional generalization abilities of neural models compare to those of humans. This paper, however, does not go deeper into one of these or any other direction.

Author response:
**>>> Further interesting questions can be asked ………..**

All these questions are indeed very interesting. However, we emphasize that our work has a very different goal. We do not intend to develop novel methods for improving Transformers’ ability to compositionally generalize. Our goals are twofold: 1) to demonstrate that compositionality can occur in autoregressively trained Transformers without changes to standard training pipelines, and if the underlying data domain itself is compositional, and 2) to understand what drives the existence of compositionality in large models. Importantly, we emphasize the motivation for our goals, as discussed in detail in the introduction, is grounded in recent literature on eliciting or predicting capabilities in pretrained models [1, 2]. Specifically, we aim to argue that if via training on a compositional domain (e.g., language), a model learns to compose its capabilities, then empirical benchmarking will be insufficient to characterize what capabilities the model possesses, i.e., what tasks it can perform.

1. Percy Liang, et al. Holistic evaluation of language models. arXiv preprint arXiv:2211.09110, 2022
2. Jordan Hoffmann, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 202

Reviewer comment (Novelty):
The paper discusses compositional generalization at the highest most abstract level, at which the answer is: it depends. While the paper acknowledges that there is ample prior work on the topics, it fails to explain what it adds on top.

Author response:
**>>>  it discusses compositional generalization and the highest most abstract level, at which the answer is: it depends.**

A strong claim requires strong evidence. We disagree that our experiments and discussion is only at the abstract level and we request the reviewer to justify why they think this is the case. Our work demonstrates that Transformers can easily learn to compose under a precisely defined data generating process.

To the best of our knowledge, we are the first to demonstrate how drastically small the training data can be, such that a Transformer learns to compositionally generalize. Our work also characterizes some properties of the training data (in-order compositions, step-by-step format, base21 vs. random) which contribute to compositional generalization. There is little to no other work that precisely characterizes when Transformers learn to compose, let alone demonstrate that it occurs with so little data. Our work makes progress on both these fronts.

Reviewer comment (Novelty):
The finding that step-wise composition is easier than direct is rather unsurprising, especially in the view of chain-of-thought prompting of LLM that has been getting popular lately.

Author response:
**>>> finding that step-wise composition is easier than direct is rather unsurprising,**

We arguably agree that this result is unsurprising, but emphasize that the step-by-step prompt format does better with extremely little data: we find that training on as few as 50 functions (1% of total) is enough to perfectly generalize to all other compositions not present in the training data. While prior works on the CTL dataset have identified that Transformers are poor at compositional generalization, our experiments in the step-by-step format highlight why the same need not be true for an autoregressive model, and how it significantly changes such a conclusion.

The compositional abilities of chain-of-thought in language models are still not well understood, For example works like Dziri et al. ([https://arxiv.org/abs/2305.18654](https://arxiv.org/abs/2305.18654)) argue that  language models have limited compositional abilities. Our work is a demonstration that it doesn’t seem to be the case in the synthetic setup even with limited data diversity. We hypothesize that the community may be severely underestimating the capabilities of pretrained language models, since we do not know the ideal prompt (unlike in the synthetic setup) to elicit a capability, even though it may exist in the model.

Reviewer comment (Writing):
Section 3.2 is a bit unclear on what are “bijection” and “permutation” mappings in the paper’s context. My understanding is that bijection here means per-token bijection, whereas permutation means shuffling tokens of the word.

Author response:
**>>>  My understanding is that bijection here means per-token bijection, whereas permutation means shuffling tokens of the word.**

That is indeed correct. We apologize if this wasn’t clear. We have added Figure 10 to Appendix A.1 that hopefully clarifies this detail.

Reviewer comment (Reproducibility):
The details of what the vocabulary size and what the size of the word is were difficult to find.

Author response:
**>>> The details of what the vocabulary size and what the size of the word is were difficult to find**

The details are present in Appendix A. We have also included them in Section 3.2 in the updated version of the manuscript.

Reviewer comment (Experiments):
I think an ablation over the number of layers in the Transformer is a key missing study here. Many studies (e.g. [Weiss et al](https://arxiv.org/abs/2106.06981), [von Oswald et al](https://arxiv.org/abs/2212.07677)) show that the depth of a Transfomer is key to what it can compute, so I'd like to see a sweep over the number of layers in the Transfomer and how that affects compositional generalization.

Author response:
Thank you for this suggestion! We have two new results that vary the number of layers in a Transformer (Appendix B.1). We find that Transformers **compositionally generalize even with 1 layer and up to 28 layers** in the step-by-step format. Models also compositionally generalizes in the **direct prompt format** but **requires at least 2 to 3 layers** and works only for a composition of a permutation and a bijection.

Reviewer comment (Theory):
I am not fully convinced that step-by-step computation implies compositional generalization. If the Transfomer is supervised with the outputs of the intermediate function results, isn't it then just learning 1. the individual functions and 2. that it should apply them sequentially? Maybe I'm missing something, but I'd definitely appreciate some sort of discussion around this in the paper.

Author response:
We emphasize that the step-by-step prompt format does not necessarily imply compositional generalization and it was not our intent to indicate that this is the case. We have updated the manuscript to address this (for example, see Section 4.1, green text) and have also highlighted two results (one of them new) which show that the ability to compose also depends on the architecture and training data, not just the use of step-by-step prompting format.


* On Architecture: We have conducted new experiments (see Appendix B.2)  with **LSTMs and observe that they do not compositionally generalize in the step-by-step and direct prompt formats**. Transformers seem to have an inductive bias to capture compositional structures in the training data. As a result, they can have a significantly larger set of capabilities compared to what’s present in the training data.
* On training data: The training data also influences if the model learns to compose. Figure 4 shows that training on just the individual functions (21 base, which acts like a null model for the choice of the training data) does not result in compositional generalization. In addition, the model must see compositions of some of these functions in the training data—albeit very few of them—for it to generalize to unseen compositions.

Reviewer comment (Writing):
(nitpick, did not affect review score) Suggestion: I think the title of the paper could be a lot better to reflect what's in the paper. Perhaps "Studying compositional generalization in transfomers via synthetic bijective functions" or something along those lines.

Author response:
Thanks for this suggestion! We will try to select a better title and are currently workshopping a few options. Our choice of current title was motivated by recent works that try to estimate the existence of a capability in a pretrained Language model, or predict its existence in a future version of the models [1,2]. We argue that if a model learns to compose its capabilities, such capability predictions will face challenges and likely underestimate the set of tasks that a pretrained Transformer can learn to perform, i.e., it is difficult to judge “how capable it can become”.

1. Percy Liang, et al. Holistic evaluation of language models. arXiv preprint arXiv:2211.09110, 2022

2. Jordan Hoffmann, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 202

Reviewer comment (Experiments):
To test if the Transfomer learned a certain composition of functions, how many examples (i.e. permutations of tokens) are used, and how are they constructed? Are experiments restricted to 6 non-unique tokens per example like in Figure 3? Unless I missed something, I think it'd be useful to clarify this.

Author response:
These details are present in Appendix A, but to emphasize them further, we have also added them to Section 3.2 (highlighted in Green) in an updated version of the manuscript. In brief: The vocabulary is of size 10 and we sample 6 tokens from this vocabulary to create the input string in all our experiments (like in Figure 3). To test if a Transformer has learnt a certain composition, we compute the accuracy of prompt completions on 1000 samples (10% of the entire input domain).

Reviewer comment (Reproducibility):
Will the authors release code/data to reproduce this paper? I imagine the code is not hard to write, but one benefit of such papers that use synthetic tasks and small models is that they should be easy to reproduce.

Author response:
We commit to releasing all code and data.