CoTFormer: A Chain of Thought Driven Architecture with Budget-Adaptive Computation Cost at Inference

Dear Reviewer,

We provide the following answers to your questions:

1. We run our experiments on both 12 layer and 24 layer CoTFormers and compare with up to 48 layer Transformers to showcase consistency of our results in different scales. However, scaling up the model further requires much higher computation power. In the revision of our paper, we provide results for 512 sequence length in the appendix which shows the benefits of CoTFormer persists at these lengths as well.

2. In order to limit the number of experiments and compute costs, we focus on the OpenWebText2 dataset, we emphasize that this is a fairly generic and large dataset and is not task specific similar to the datasets used for state of the art pre-training.

3. We are not sure what kind of theoretical analysis is requested here. The goal of this paper is to clarify a clear advantage of CoTFormer over Block Universal due to access through attention which has not been pointed out in prior work. This is evidenced by empirical results. Any theoretical analysis remains far outside the scope of this work.

4. In the adaptive CoTFormer, the compute budget can be decided at inference and can be tuned depending on the cost of compute for the user. Of course, more compute usually leads to better accuracy (as shown in Fig. 4).

We hope the above answers adequately answer your concerns. However, looking at your review, we do not understand why you are suggesting a rejection of our work as most of the raised points are not major objections. If our answers have adequately answered your concerns, please consider raising your score or let us know any major objections that compel you to suggest a rejection so we can also answer them.

Dear Reviewer,

While we appreciate your consideration of our work, upon reading your review, we do not understand your justification behind rejecting our work with such high confidence. We hope the following comments clarify our contributions further and earnestly ask you to consider raising your score or provide reasons that have so strongly compelled you to reject our work.

1. Chain of thought is a prompting technique that relies on auto-regressive decoding. Figure 1 is demonstrating how each method works, comparing what happens in CoTFormer with what happens when a Chain of thought technique is used on a normal Transformer. The goal of this figure is to demonstrate the difference between generating CoT using auto-regressive decoding (Figure 1a) and repeatedly applying the same model (block universal, Figure 1b). In the former, a token is fed through the model to generate a new token, which is appended to the sequence, and itself fed again to the model, generating a new token, etc. This is conceptually similar to feeding the same token through multiple repetitions of the model. In the latter approach, each token is fed to multiple repetition of the model, yet—unlike in the former approach—intermediary representation of a given token cannot attend to past representations of that token. This observation led us to the design of CoTFormer. The colors are meant to distinguish tokens based on how many steps were used to generate them. For chain of thought (Figure 1a), the figure shows a process of generating an answer after 3 steps of decoding with the model whereas for block universal and CoTFormer, the figure only shows 2 repeats. That is why there are no yellow colors for CoTFormer and Block Universal. We understand that this can be slightly confusing and will remove the third step from CoT so all the figures correspond to 2 steps of using the model. That being said, given that this is the only weakness mentioned in your review, we strongly object that this would be grounds for rejection.

2. Regarding training vs inference FLOPs, we note that we are plotting against the FLOPs for a forward pass of sequence length 256. The backward pass’s FLOPs is usually within a constant factor of the forward pass. As such—given we are using a similar number of training iterations, batch sizes, etc. for all methods—training FLOPs and inference FLOPs are not that different. Note that we are not plotting against single token decoding FLOPs. We will add this remark to our next revision. Let us know if more clarifications are needed on this point.

3. The goal of this work is to point out the difference between doing CoT using auto-regressive decoding, and simply applying the model multiple times (block universal transformer). We do not claim that we are doing better than a larger standard Transformer and clearly discuss the benefits (such as adaptive compute) and limitations (needing 5 repeats to match 48 layer Transformer) in the paper (Section 5). However, compared with architectures that rely on repeatedly applying the same model, CoTFormer still improves significantly over block universal transformers due to attention’s access to intermediary representations. To draw a parallel with another recent line of research, it is known that dense transformer models are better than Mixture-of-Expert (MoE) models with similar number of total parameters, yet MoEs have other benefits that make them desirable leading to them being widely adapted by the community.

4. Please note that the referred sentence about the growth of computation cost of attention being less noticeable is about much larger models that have large feed forward layers. Training these models is not possible for us due to resource limitations. However, we point out that when using flash attention, the attention’s memory cost increases linearly (not quadratically) with length. Therefore, theoretically when factoring in the KV cache, the memory requirements of Block Universal and CoTFormer are the same.

We believe we have addressed (i) the misunderstanding regarding the link between our method and CoT, as well as (ii) the performance of our method in terms of training FLOPs, and finally (iii) compared the memory footprint of our method with our block universal transformer baseline. As those fully cover the concerns raised in your review, we hope you will either raise your score or provide further clarifications justifying your score.

Dear Reviewer,

We hope the following comments address your concerns:

1. The limitation you described is discussed in Section 5 of the main paper.. However, the goal of this paper is to demonstrate an inconsistency between block universal (i.e. applying the same model multiple times) and doing chain of thought (i.e. applying the same model multiple times and attending to past intermediary representations) which was otherwise missed and to show that it has significant effects. That is why our main comparison is with block universal, to show that allowing attention to previous intermediary states, pushes the pareto frontier forward.

2. While it is true that the results on downstream tasks are not very strong, they still show improvements of CoTFormer over Block Universal Transformer. These results are meant to complement our results in the main text. To obtain much better performance on these tasks, we need much larger models that are also trained for much longer which is not feasible in our academic settings.

3. We are evaluating the perplexity of language modeling which is the main target for training many large models. We are also showcasing performance on downstream tasks in the appendix. As such, we do not understand why training on a synthetic task would help. Additionally, understanding how Transformers or CoTFormers work is outside the scope of our work which is to point out the benefits of having attention access to intermediately representations and demonstrate new methods to achieve adaptive compute. Such analysis would of course be useful in future work.

4. We did not investigate fine-tuning a model and focused on training from scratch. We agree that it is not obvious to fine-tune a model as a CoTFormer and investigation into whether it is possible, for example by starting from a pre-trained model and using CoTFormer training with it, would be interesting in the future. However, we emphasize that new methods for training models from scratch are useful given that they enable additional benefits such as adaptive compute in this case.

We hope the above comments fully address your concerns. We ask that you kindly consider raising your score or otherwise please share the major objections that compel you to reject our work. We remain in your disposal if you have any additional comments or questions.

Thank you for your responses.

I think one important question here is: which model architecture is a more appropriate baseline for the current work, the standard Transformer or (Block) Universal Transformer? This question is related to what CoTFormer intends to achieve, and also to how evaluation of it should be done.

To explain this, let's first recall the original Universal Transformer paper (https://arxiv.org/pdf/1807.03819). The motivations and fundamental advantages of Universal Transformer was clearly stated in the abstract, for example:

• "Despite these successes, however, popular feed-forward sequence models like the Transformer fail to generalize in many simple tasks that recurrent models handle with ease, e.g. copying strings or even simple logical inference when the string or formula lengths exceed those observed at training time."
• "UTs combine the parallelizability and global receptive field of feed-forward sequence models like the Transformer with the recurrent inductive bias of RNNs."
• "In contrast to the standard Transformer, under certain assumptions UTs can be shown to be Turing-complete."

And to validate the claimed advantages, empirical results on both synthetic tasks and some common benchmarks were reported. Perplexity on text was reported only for LAMBADA, for which the accuracy was also reported.

Of course, the complexity of the Universal Transformer architecture limits its adoption in practice (and standard Transformer is still dominant), but that's fine as long as it has sufficient research value.

Now, let's get back to CoTFormer. If its goal is merely to improve language modeling perplexity or scores on common benchmarks (which is how most of the evaluation is done in the current work), then I think the standard Transformer (with equal inference FLOPs) would be a more appropriate baseline, and CoTFormer doesn't seem to outperform it. On the other hand, if its goal is to achieve fundamental advantages over the standard Transformer and Universal Transformer (via attention to previous intermediary states, which is the main innovation of CoTFormer), then it's totally fine to use Universal Transformer as the baseline, but language modeling perplexity or common benchmarks would not be sufficient in this case.

From my perspective, a nearly 100% accuracy on a carefully crafted synthetic task (no matter how toy it is) where standard / Universal Transformer totally fails would be much more exciting and insightful than a minor improvement in perplexity, and it is this kind of evidence that can truly demonstrate the fundamental benefits of scaling up inference FLOPs in the particular way that CoTFormer does (rather than simply using a standard Transformer with twice as many layers). The bottomline would be to show that CoTFormer retains the key advantages that Universal Transformer (with equal inference FLOPs) has, while improving perplexity or scores on common benchmarks.

Overall, I feel that the current work does not fully meet my expectations for a research work that introduces a novel model architecture, and I'm inclined to maintain my rating. However, I acknowledge that my evaluation might be somewhat subjective, so I'm not strongly opposed to acceptance.

Dear Reviewer,

Thank you for your careful consideration of our work and helpful suggestions. We provide the following answers to your comments:

1. The goal of this paper is to demonstrate the benefit of having access to past intermediate representations through attention and point out an often overlooked difference between block universal transformer and chain of thought. While understanding the inner workings of Transformers or CoTFormers is a very exciting and interesting direction it falls outside the scope of this work. Still, as you suggested we plotted the attention patterns observed in a CoTFormer to the appendix of our paper’s revised version. Interesting patterns similar to what you mentioned such as heads focusing on tokens generated during specific repeats can be observed.

2. While we agree that the used sequence length is shorter than state of the art models, we use it to save on computation costs. That being said, our intuition suggests that the method’s performance is disentangled from the context’s original length. To support this we provide two additional pieces of evidence in the appendix of the revised version of our paper. First, we demonstrate superiority of CoTFormer on longer sequence lengths of 512. Second, we also compare the performance of CoTFormer with using pause tokens suggested in [1]. While the latter also provides longer sequence lengths to the model during inference, it can be seen that this does not suffice to obtain the high accuracy obtained by CoTFormer.

3. As can be seen from Equation 3, the (i+1)-th pass of a token through the model is done by inputting the output of the i-th pass to the model. Therefore, similar to the figure, the last token representation is processed by the model to generate the next one. However, when using CoTFormer, earlier versions (e.g. the initial token representation) can also be accessed through attention.

4. This can be done by providing a special mask vector which allows tokens in the i-th pass to attend to tokens that are coming before it and are in earlier passes. For efficiency, in the implementation, we append the tokens in the new pass to the sequence instead of interleaving them. Thus the mask is no longer fully causal.

5. We use the same position id as the original tokens. Thank you for raising this question. We will include this clarification in the camera ready version of the paper.

6. As suggested, we also add the results for models with a larger width (1024) to the appendix in our revision. It can be seen that even for the larger widths the gap between CoTFormer and Block Universal Transformer persists. That being said, intuitively, going to extremely larger widths might allow Block Universal Transformers to obtain the same performance as CoTFormers since one can theoretically fit the same information as multiple tokens in one token with a much larger width. However, aside from possibly being impractical, such a setting diminishes the adaptivity benefits of CoTFormer as it needs to operate on the larger width for all tokens.

7. “Activating a prefix of repeats” is similar in inference to fixed depth, since we fix the depth in this case. However, note that in this case the model is not trained with fixed depth and therefore we can use the same model at different fixed depths. When using a mixture of depth, at training we fix the ratio of tokens reaching each depth. But at inference, e.g. for decoding, we usually use a fixed threshold to decide from the router’s weight whether to proceed to the next depth or not. Varying this threshold is how we vary the compute in Figure 4. When using this thresholding mechanism, the decision can be made auto-regressively during decoding. Thus there is no issue regardless of the batch size.

8. We appreciate the various suggestions you provided and have already implemented some of them as experimental results for longer sequence lengths and larger widths. We point out that in this work our aim is to clarify the benefits of having access to earlier representations of the model through attention which is clearly supported by our results. We agree that improvements in efficiency of the method are great directions for future work and as you mentioned we have already outlined them in Section 5. Similarly, theoretical analysis of this phenomenon would be interesting but falls outside of the scope of this work.

We hope the above comments fully address your concerns. We ask that you kindly consider raising your score and remain in your disposal if you have any additional comments or questions.