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

## Abstract

Abstract Scaling language models to larger and deeper sizes has led to significant boosts in performance. Even though the size of these models limits their application in compute-constrained environments, the race to continually develop ever larger and deeper foundational models is underway. At the same time—regardless of the model size—task-specific techniques continue to play a pivotal role in achieving optimal downstream performance. One of these techniques, called Chain-of-Thought (CoT), is particularly interesting since, as we point out in this work, it resembles employing a deeper transformer through re-applying the model multiple times. However, a key subtlety in computing the attention of past tokens differentiates CoT from simply applying the model several times. Based on this insight, we propose CoTFormer, a novel architecture which closely mimics CoT at the token level, allowing us to obtain significantly improved accuracies close to much larger models. While applying CoT introduces additional computation costs, we compensate for it by leveraging CoTFormer’s special compatibility with token-wise variable depth.
Through a compute adaptive model—which automatically allocates the compute to tokens that need it most—we show that it is possible to reduce the computation cost significantly without any reduction in accuracy, and with further compute cost reductions possible while maintaining a competitive accuracy.

## 1 Introduction

Large foundational models have demonstrated remarkable performance across various tasks, predominantly employing the Transformer architecture []. The ability to tackle new tasks in zero-shot or few-shot settings [] has been attributed to emergent properties that become increasingly prominent with model size []. This observation has sparked a race to build progressively larger models [, , , ].

However, despite the evident improvement in performance with size, certain challenges persist even in very large and deep models. One example is their proficiency in mathematical tasks []. In response to these challenges, an alternative approach called Chain-of-Thought (CoT) [] has been proposed, requiring models to think step by step and articulate their thought processes, showing remarkable success []. In particular, using CoT can improve the general performance of even smaller models [, ].

In this work, we draw attention to the intrinsic connection between constructing deeper Transformers and employing CoT. At a first glance, applying CoT with $n$ thought tokens can resemble an $n$-times deeper Transformer with weight tying implemented on every $n$-th layer (see Figure 1). Such weight tying schemes have been explored in the past []. However, in this work, we point out that there is a distinction between CoT and conventional weight tying. More particularly, when applying CoT, the attention mechanism can access previous intermediary tokens whereas in the weight tying such access is not granted.

Figure: (a) CoT
Refer to caption: x1.png

Based on this observation, we propose CoTFormer, a Transformer that implicitly applies a similar mechanism to CoT. We empirically show that using CoTFormer allows us to obtain much better performances than deeper baseline models. Especially, CoTFormers surpasses existing methods such as Universal Transformers [], and pushes the perplexity-compute Pareto frontier forward.

Through asking the model to think step by step, CoT generates a variable amount of intermediary tokens. More complex next token predictions tasks (e.g. an advanced math question) might require to make explicit a greater number of intermediary reasoning steps before reaching an answer. In contrast, tokens which are simpler to predict might not require any intermediary reasoning step at all.
This adaptability of CoT to the difficulty of the prediction is remarkable. Indeed, building compute adaptive models has been a long-standing goal, driving the exploration of architectures that can be recurrently applied—controlling the compute cost through deciding the depth of the recursion [, , , , , ].
However, one challenge those prior methods face is how to allow deeper layers to attend to tokens that have been assigned less depth—e.g. what is the expected interactions between tokens $w_{5}$ and $w_{2}$ at depth $3$, given that $w_{2}$ stopped at depth $1$? Existing works have proposed possible solutions such as copying the output of the last layer onward. However, these solutions require the model to be able to process the output of different layers at each layer. In contrast, by treating each new application of the model as creating a new token, CoTFormers can completely bypass this problem. Token $w_{5}$ can simply access all the tokens which have been generated before through the attention mechanism. This makes the CoTFormer a much more natural fit to use in a computation adaptive setting.

In this work, we also propose a new adaptive training method for CoTFormer.
We show that using this method yields a model that allows choosing the computation cost at inference time, and navigating the accuracy-computation trade-off without additional training. This is in contrast with most current models that have a fixed computation requirement, preventing them from functioning in more constrained settings. Our model automatically allocates more compute to the tokens that need it while cutting back on others to remain within budget. We observe that, as expected based on our conjecture, the computation cost can be reduced to a certain level with a negligible impact on the accuracy.
We also show that, as expected, reducing the computation cost beyond a certain level inevitably reduces the accuracy of the model.

Our main contributions can be summarized as follows:

- •
Pointing out an important distinction between Chain-of-Thoughts and the recurrent application of a model with weight-tying.
- •
Proposing CoTFormer which accounts for the aforementioned distinction, and demonstrating its superior performance over other weight-tied deep models (e.g. Universal Transformer []).
- •
Proposing a training method that allows adjusting the per-token depth distribution at inference, controlling the computation costs while trading compute for accuracy.

## 2 Related Works

A model usually receives a mix of easy and hard examples which encourages the idea of adapting computation cost based on the input’s difficulty. Prior work proposed different approaches to achieve this adaptiveness for various models [, ].

These methods usually rely on applying the model multiple times, simulating a deeper model with weight tying. In many aspects, this approach is similar the widely used technique of instructing the model to generate intermediary thoughts before outputting the final answer, called Chain of Thought (CoT) []. Previous work has observed that applying CoT significantly improves performance on various tasks such as mathematical reasoning and its effect on increasing depth has been studied from a theoretical perspective []. Furthermore, while Transformers with fixed depth are not Turing complete on their own [], combining them with the auto-regressive decoding used for generating the chain of thought can make them Turing complete [, ]. In this work, while we acknowledge the similarity between CoT and recurrently applying the model, we point out an important difference between these two approaches. Taking this difference into account to mimick the CoT process leads to the development of CoTFormer.

For Transformers, Dehghani et al. [] propose Universal Transformers which repeatedly applies a single layer transformer model on the model. A predictor is trained using the ACT method [] to decide whether to stop or apply the model again. Due to the instability of ACT and its sensitivity to hyperparameters, Banino et al. [] propose PonderNet which weights the predictions at each depth using a probability distribution close to a geometric distribution. This architecture has been extended to cases where the base model has more than one layer. The Block Universal Transformer architecture we explored in this work as a baseline is an example of such architecture while other weight tying arrangements are possible and are explored in [].

In these methods, the artificial depth is determined separately for each token. The varying depth between tokens introduces the problem of missing information from tokens that terminated in an early layer when using the attention mechanism in deeper layers. Various approaches have been proposed to address this issue such as copying the latest layer representation forward []. In contrast, no such problem exists in CoTFormer since going deeper is equivalent to producing a new implicit thinking token.

Furthermore, the token-based variability of depth makes it challenging to implement batching for these models efficiently. To address a similar challenge when deciding whether to skip blocks of a standard Transformer architecture, Raposo et al. [] propose Mixture-of-Depth (MoD) defining a fixed capacity for each block which determines the number of tokens that will go through that block. We use a similar method to allow efficient implementation of our depth adaptive CoTFormer models. However, unlike CoTFormers, MoD uses different weights for each block and therefore does not benefit from the smaller size induced by weight-tying as in CoTFormers. Furthermore, in contrast with CoTFormers which apply a full prefix of blocks, MoDs decide separately whether to use each block or not, preventing early exiting.

Recent work have explored and proposed a variety of alternative architectures which prove to be useful in different scenarios [, ]. Most prominently Mixture-of-Experts (MoEs) have been shown to improve performance of the model in many cases []. For example Sparse Universal Transformers [] combine the idea of universal transformer with MoEs, allowing a router to choose a possibly different model every time the input is processed again. In this work we mainly focus our experiments on the Pre-LayerNorm Transformer architecture [], which is currently the most widely used architecture and is the backbone of the state of the art language models [, ]. However, we emphasize that our method uses the Pre-LayerNorm Transformer architecture as a black box and therefore could be directly applied to any of its other variants.

Recent work have also studied explicitly teaching the model to reason by training it on a corpus containing step by step reasoning [] and have shown it to be useful. The main obstacle with this approach is the lack of abundant volumes of high quality reasoning data, encouraging recent work to generate artificial data [, ]. Regardless, we believe this approach to be complimentary to CoTFormers. Intuitively, CoTFormers allow re-using basic modules such as extracting information from the context and applying them multiple times. On top of that, the explicit CoT training teaches the model how to reason in a higher level to make rationale arguments.

Finally, while Block Universal Transformers simulate a deeper model, and while alternative proposals such as Pause Tokens simulate a model with increased hidden dimension (i.e. width) [], CoTFormers intuitively facilitates both. Still, width-increasing methods such as Pause Tokens can be combined with CoTFormers.

## 3 Chain-of-Thought and Model Depth

Chain-of-Thought involves asking the model to output the solution step-by-step (a process similar to thinking) before outputting the final answer. This process results in the generation of thought tokens in addition to the normal tokens. These thought tokens are generated using auto-regressive decoding.
Notably, the whole process of generating thought tokens and finally generating the next normal token is similar to recursively applying the same model multiple times (similar to a Recurrent Neural Network, RNN []; see Figure 1(a)). Consequently, one might be tempted to frame the chain-of-thought process as the utilization of a deeper model with tied weights (see Figure 1(b)). Indeed, such arrangement resembles a version of Universal Transformers [] generalized to allow multi-layer base blocks (instead of only single layer). However, in this work, we point out one critical distinction between the described generalization of universal transformers (which we call Block Universal Transformer), and Chain-of-Thought: When applying CoT, the generated thought tokens can directly attend to previous thought tokens.

Having emphasized the above distinction, we propose CoTFormer to closely resemble the CoT process at the token level, taking the highlighted distinction into account.

### 3.1 CoTFormer

Given a context input sequence at depth 0 0: ${\bm{x}}_{1:{n_{\text{seq}}}}^{(0)}=[{\bm{x}}_{1}^{(0)},\ldots,{\bm{x}}_{n_{
\text{seq}}}^{(0)}]$ and a current input token ${\bm{x}}_{{n_{\text{seq}}}+1}^{(0)}$, we describe the process of generating the next token for the Block Universal Transformer and our CoTFormer model. First, let $B^{(i)}({\bm{x}},{\bm{c}})$ be the i-th repeat of a set of ${n_{\text{layer}}}$ transformer blocks taking as input the token ${\bm{x}}$ and being able to attend to the context ${\bm{c}}$ through its attention mechanism. One can imagine ${\bm{x}}$ being the current token being processed and ${\bm{c}}$ being the key/value-cache, as often used during inference. For a typical transformer, generating the output of $B^{(i+1)}$ for token ${\bm{x}}_{n_{\text{seq}}}^{(i)}$ can be written as follows:

$$ $\displaystyle{\bm{x}}_{{n_{\text{seq}}}+1}^{(i+1)}:=B^{(i+1)}({\bm{x}}_{{n_{ \text{seq}}}+1}^{(i)},{\bm{x}}_{1:{n_{\text{seq}}}}^{(i)})\,.$ (1) $$

Repeating the above formula ${n_{\text{repeat}}}$ times with weight tying between the $B^{(i)},1\leq i\leq{n_{\text{repeat}}}$, yields the Block Universal Transformer:

$$ $\displaystyle{\bm{x}}_{{n_{\text{seq}}}+1}^{(i+1)}:=B({\bm{x}}_{{n_{\text{seq} }}+1}^{(i)},{\bm{x}}_{1:{n_{\text{seq}}}}^{(i)})\,.$ (2) $$

CoTFormer also use weight tying, but in contrast with the Block Universal Transformer, it provides intermediary representations from previous repeats in the attention. The CoTFormer can be specified as follow:

$$ $\displaystyle{\bm{x}}_{{n_{\text{seq}}}+1}^{(i+1)}:=B({\bm{x}}_{{n_{\text{seq} }}+1}^{(i)},[{\bm{x}}_{1:{n_{\text{seq}}}}^{(i)},\ldots,{\bm{x}}_{1:{n_{\text{ seq}}}}^{(0)}])\,.$ (3) $$

Figure 1(c) illustrates the above process. It can be seen that the sequence length grows linearly with the number of repeats. This does not have an effect on memory since the intermediate representations need to be stored in any case, either for the backpropagation during training, or for the KV-cache at inference. However, it may impact the computational cost which we discuss in Section 3.2. We use the notation ${n_{\text{layer}}}\textsf{x}{n_{\text{repeat}}}$ to describe a CoTFormer or Block Universal Transformer with $n_{\text{layer}}$ layers being repeated ${n_{\text{repeat}}}$ times.

### 3.2 Comparison with Block Universal Transformer

Experimental setting. To establish the importance of attending to previous intermediary states, we compare the performance of CoTFormer and Block Universal Transformer on the OpenWebText2 [] dataset; a dataset of websites linked from reddit between 2005 and 2020 initially released under MIT license. We train the models for $40$k steps using the AdamW [] optimizer and apply linear warmup of the learning rate for the first 8000 steps. We use a the Pre-LayerNorm Transformer [] with $12$ heads, hidden dimension $768$, sequence length $256$, and maximum learning rate $0.001$ and feed the data in batches of $128$ sequences. We run all our experiments on Nvidia A100 80GB GPUs.

Perplexity comparison. The results are reported in Table 1. It can be clearly seen that CoTFormers significantly outperform Block Universal Transformers with the same size and the same number of repeated applications of the model. We emphasize that using CoTFormers does not introduce an overhead in terms of memory. This is because the storage of intermediary tokens is needed given the need for the KV cache even when using Block Universal Transformers.

**Table 1: Performance of CoTFormer, Block Universal Transformer and Standard Transformers on OpenWebText2. The mean perplexity of 3 runs is reported with the standard error of the mean in parenthesis. It can be seen that CoTFormers clearly outperforms Block Universal Transformers. The best perplexity for a given ${n_{\text{layer}}}\textsf{x}{n_{\text{repeat}}}$ combination is marked in bold.**
| Model | Base Layers (${n_{\text{layer}}}$) | ${n_{\text{repeat}}}$ |  |  |
| --- | --- | --- | --- | --- |
| 2 | 3 | 5 |  |  |
| Standard | 12 | 28.39 (0.01) |  |  |
| Block Universal Transformer | 12 | 27.74(0.01) | 27.47(0.01) | 27.15(0.02) |
| CoTFormer | 12 | 27.55(0.02) | 27.07(0.01) | 26.64(0.04) |
| Standard | 24 | 25.93 (0.02) |  |  |
| Block Universal Transformer | 24 | 25.47(0.01) | 25.19(0.03) | 24.95(0.01) |
| CoTFormer | 24 | 25.28(0.00) | 24.85(0.04) | 24.48(0.03) |
| Standard | 48 | 24.17 (0.00) |  |  |

Compute cost comparison. As such, the only downside of using CoTFormers instead of a Block Universal Transformer is the growth in the computation cost of the attention layer. This growth occurs because when using CoTFormers, the outputs of all previous repeats are accessible. Therefore, given the quadratic cost of the attention, one might expect the cost of CoTFormer to grow quadratically with number of repeats. However, for current models, the main bottleneck in computation when processing average length sequences is the feed-forward network in each block, not the attention layer [, ]. It is only for very long sequences that the attention layer becomes a bottleneck. Therefore, the growth in computation cost is actually much less noticeable. At the same time, using CoTFormer comes with significant boost in accuracy. The same pattern can be observed in Figure 3 which shows the number of multiply-accumulate operations needed to process different sequence lengths by a block universal transformer with ${n_{\text{repeat}}}=5$ and a CoTFormer with ${n_{\text{repeat}}}=3$ which obtains a similar accuracy (on sequences of length $256$). It can be seen that even for sequences as long as $8192$, the $12\textsf{x}3$ CoTFormer’s cost remains below that of $12\textsf{x}5$ Block Universal Transformer.

Figure: (a) ${n_{\text{layer}}}=12$ (x-axis is in log scale)
Refer to caption: x4.png

The demonstrated trade-off is further depicted in Figure 3 which shows the perplexity against compute cost of processing a sequence with $256$ tokens for CoTFormers and Block Universal Transformers with ${n_{\text{layer}}}=12$ and ${n_{\text{layer}}}=24$. At both scales, it can be clearly seen that while CoTFormers with the same number of repeats are more costly, they come with significant improvement in accuracy which overall puts them in the front of the Pareto frontier. Furthermore, the performance gap widens as the number of repeats increases, suggesting better scaling properties of CoTFormers. The above results clearly demonstrate the effectiveness of CoTFormers over Block Universal Transformers.

### 3.3 Architecture Tweaks & LN-CoTFormer

The previous section introduced as little innovations as possible to clearly demonstrate that the improved performances are due to the built-in CoT mechanism, and not to some other tweak. Having established the better performance of CoTFormer, we now introduce several modifications which we found further improved the results.

#### Reserving Beginning and Final Layers.

In order to allow the model to operate on an intermediary space which is not necessarily the same as the word embedding space, we propose separating the first and last few layers from repeats. In this scenario, the model will first execute $n_{\text{begin}}$ layers, followed by multiple passes through $n_{\text{middle}}$ layers. Finally the output is generated from the final pass of each token by passing it through the last $n_{\text{end}}$ layers.

#### Layer Norm After Each Repeat.

Given the internal residual connections of the model, we conjecture that it is important to maintain a consistent input’s scale. Therefore we additionally inject a layer norm at the end of each repeated pass, similar to the final layer norm applied in the standard architecture before predicting the next token.

The clear positive effect of the above tweaks on performance can be seen in the ablation study in Table 3. In the case of reserved beginning and final layers, note that while the accuracy does not improve, the computation cost decreases since the total number of layers is kept fixed at $24$. Most noticeably, after applying these changes, the performance of a CoTFormer with $24$ layers and $5$ repeats, surpasses the performance of a standard $48$ layer Transformer.
We call the final resulting architecture LN-CoTFormer. We note that while LN-CoTFormer’s final performance is better than CoTFormer, we observed some benign spikes in the loss during training. Though the model quickly recovers without intervention from these spikes, this might suggest CoTFormer are more stable to train than LN-CoTFormers. Still, we focus on using LN-CoTFormers when building our adaptive model in the next section.

## 4 Token-Wise Adaptive Repeats

The standard CoTFormer has the advantage of obtaining better performance with smaller models which is useful in memory-constrained environments such as mobile phones. Moreover, the recurrent application of the small model also opens up the direction of varying the number of times the model is applied, i.e. the number of repeats, on a more granular level. In particular, the intuition that the difficulty of predicting the next token varies over the sequence, encourages that dynamically varying the number of repeats on a token level based on the context can yield computational savings.

Prior work, in particular universal transformers, also aim to create such adaptive models that use a different number of repeats based on the difficulty of the current token. This is done through a halting module which is called at the end of each repeat to decide whether the current state should be used as the output (i.e. halt) or to continue with another pass through the small model. However, two challenges remain persistent:

- •
Attending to Previously-Halted Tokens in Later Layers: If a token decides to halt early, subsequent repeats after the one where the token halts still need to have a representation of the token available in the attention layer to allow tokens that are still processing access the already halted token. Prior work have suggested approximating this representation by copying the last output for the token forward. In contrast CoTFormer does not face this challenge. Since the model can attend to each token’s representation after each of previous repeats, a halted token is already represented when invoking the attention mechanism. As such, CoTFormer adapts much more naturally to the adaptive depth setting.
- •
Efficient Batch Processing: Since the decision of halting or continuing happens on a token level, sequences of the same length may end up with different number of tokens in the subsequent repeats. As a result, efficient processing of batches of multiple sequences becomes challenging. Therefore, in this work, we propose a different approach where a certain capacity is assigned to each iteration of processing the sequence using the small model, and the most eager tokens are assigned to pass through that model again. Our approach is similar to the method proposed by Raposo et al. [] to train Mixture of Depth (MoD) models in some aspects, namely assigning capacities for each iteration instead of for tokens, but deviates from MoD’s training method in other aspects, such as randomized capacities, as detailed in the next subsection.

In addition to addressing the above challenges, we also aim to build a model that can work under different constraints. In particular, we aim to offer the flexibility to choose the compute budget during inference, with more compute yielding a better accuracy. Therefore, we randomly pick the compute budget at each iteration instead of fixing it in advance, allowing the model to adapt to different constraints. Similar approaches have been used to build models with varying width or rank [, ]. We now explain the details of our method.

### 4.1 Mixture of Repeats

We assume all tokens go at least one time through the model. We now use ${n_{\text{repeat}}}$ to refer to the maximum number of times a token can go through the model. For each of the passes $2$ to ${n_{\text{repeat}}}$ we train a separate embedding vector, e.g. ${\bm{e}}^{(i)}$ for the $i+1$-th pass, and use the dot product between this vector and the current representation of the token. In particular, if the output after the $i$-th pass for $j$-th token is denoted by ${\bm{x}}_{j}^{(i)}$, we compute the score $s_{j}^{(i)}:=\sigma({{\bm{e}}^{(i)}}^{\top}{\bm{x}}_{j}^{(i)})$ to determine whether the $j$-th token should pass for the $i+1$-th time through the model. We interchangeably use the terms router to refer to this mechanism. The router weights (the vectors ${\bm{e}}^{(i)}$) are trained alongside other parameters of the model.

To determine which tokens pass through the next repeat of the model, we sort the tokens based on their score as defined above and pick the top $k$ where $k$ is chosen based on the capacity level for this repeat: $c_{i}$. In particular, if we denote the sequence length by ${n_{\text{seq}}}$, we will use $k:=\lfloor c_{i}\times{n_{\text{seq}}}\rfloor$.

Let us denote the output of the model on the input $x$ by $B({\bm{x}})$. We use an interpolation between the previous pass’s output and the new output of model to get the output of this pass. In particular, we use ${\bm{x}}^{(i+1)}:=(1-s_{i})\cdot x_{i}+s_{i}\cdot B({\bm{x}}^{(i)})$. This interpolation plays two roles. First of all, the gradient needed to train the embeddings ${\bm{e}}^{(i)}$ is obtained only based on this interpolation since the process of token selection is not differentiable. More importantly, it provides a way to the model to ensure that increasing capacity will not hurt the performance. For example, if we set the capacity of a repeat to $1$, even tokens with very low scores will be selected. However, a low score indicates that such additional processing of these tokens might adversely affect the accuracy of the prediction. As a result of this interpolation, the representations of such tokens remain unchanged.

Finally, instead of using fixed capacities $c_{i}$, we sample them at random for each batch. More particularly, we assign a capacity of 1 to the first repeat and sample ${n_{\text{repeat}}}-1$ numbers and sort them in decreasing order to get the capacity for other repeats. This random sampling has two key effects. First of all, intuitively, sampling allows tokens to explore being passed through deeper layers. Otherwise, only tokens with a high score that were selected for a pass would affect the gradient. Therefore, the update for router weights would only take into account those high scoring tokens. As a result low scoring tokens will continue to be excluded. This challenge of exploration vs exploitation arises from simultaneous training of the router weights and the model parameters. The second effect of the sampling is ensuring the model functions with different capacity factors. This allows adjusting the capacity factors at inference, which in turn allows customizing the computation budget without unreasonable losses in accuracy.

### 4.2 Adaptive Architecture

We mainly use a LN-CoTFormer to build the adaptive model which has the same architecture as Section 3.2 except for the architecture tweaks discussed in Section 3.3. In particular, we fix (meaning we do not repeat) the first 2 layers and the last layer. Additionally, we train our models for 60k steps instead of 40k steps. Finally, we introduce a depth embedding to the model.

Depth Embedding. In order to allow the model recognize how many number of passes it has done, we add a depth embedding at the start of each pass. For this embedding we learn a single vector ${\bm{e}}^{(\text{depth})}$ and add $({n_{\text{repeat}}}-i)\cdot{\bm{e}}^{(\text{depth})}$ as the embedding for the $i$-th repeat. Intuitively, this should condition the model based on the maximum number of repeats it has left. We investigate the effect of this embedding in Appendix Table 3. While the performance is similar on fixed number of repeats, a noticeable boost is observed in the adaptive case.

### 4.3 Results

By design, an adaptive model’s performance depends on the amount of compute it is allocated. In order to measure the performance of the model for different computes, we compare two methods. The first is to activate a prefix of repeats, and the second is to rely on the router weights learned by the model. For the second approach, we compute the ratio of tokens that enter each repeat on a subset of the training set and use the obtained ratios as the capacity factors. Alternatively, one could directly threshold the router weight without needing such statistic measurement. However, we chose the former approach for simplicity and maintaining the ability of batch inference.

In Figure 5 we plot the accuracy against the number of multiply accumulate operations. We vary the router threshold to move between compute budgets. In order to show the effectiveness of the router in allocating compute to tokens, we compare the results with the alternative method of running all tokens at inference time on a smaller number of repeats to reduce cost.
The results clearly show that relying on the learned router weights provides a far more effective way of allocating compute and manages the accuracy-compute trade-off more efficiently. As a result, the computation cost can be significantly reduced without noticeable loss of accuracy. Further reduction of computation cost is also possible in exchange for reasonable accuracy losses, allowing us to traverse the accuracy-compute trade-off at inference time which is not possible with the standard models.

We also report the results for an adaptive Block Universal Transformer. While alternative methods for adaptive training of Universal Transformers have been proposed in the literature, we could not obtain a better performance using those methods in small scale experiments. We provide additional details in Appendix C and continue with reporting the results of training a Block Universal Transformer using Mixture of Repeats. Following previous work, for already halted tokens, we copy their last representation forward.

Similar to our previous results, we observe that CoTFormer outperforms Block Universal Transformer when enough compute is available. However, when moving to the lowest computation budget, in particular when no repeat is allowed, the adaptive Block Universal Transformer outperforms the adaptive CoTFormer. Intuitively, this can be because CoTFormer learns to better utilize the additional number of repeats to obtain better performance but has to sacrifice some performance when the number of repeats remains low.

Figure: Figure 4: Perplexity for different amount of compute budgets chosen at inference. The adaptive CoTFormer can adapt to different budgets, reducing compute in exchange for reasonable loss in accuracy. Furthermore, using the router weights to allocate the available compute (Router) is much more effective than fixing the depth at inference time to a smaller value in order to reduce computation cost (Fixed Depth).
Refer to caption: x7.png

## 5 Discussion and Future Work

#### Training of Deeper Layers.

While the above performance is remarkable, we can observe a gap between an adaptive CoTFormer and a non-adaptive CoTFormer trained with exactly 5 repeats even when the adaptive variant is allowed the same amount of compute at inference. For example, after 60k steps, the former reaches perplexity 23.83 while the latter achieves 23.19. One possible reason is the reduced amount of gradient information reaching higher number of repeats in the adaptive training since a good portion of tokens will halt before reaching those repeats. As such, we conjecture that the adaptive CoTFormer would benefit more from longer training. We verify this in Figure 5 where we plot the ratio of different values of router weights for the final repeat when the model is trained for 40k steps and compare it with training for 60k steps. We can clearly see that the model starts to favor using the final repeat more when the model is trained for longer. We note that training time of an adaptive model is *significantly lower* than training directly at the maximum number of repeats. For example, when training the model with fixed $5$ repeats, the training takes roughly around $1.5$x longer.

Alternative Sampling Methods. In addition to longer training, we conjecture that the sampling plays an important role in the quality of the final model. While we tried some alternative sampling methods, we could not find a method that performs better than randomly picking and sorting as described in Section 4.1. Still, we expect better methods to exist and leave exploration of such methods as a direction for future work.

Comparison with Standard Transformer. Currently, getting the same performance as a deeper standard transformer model, requires more number of repeats than the depth difference factor. For example, to get better performance than a $48$ layer model using a $24$ layer model, $5$ repeats is needed whereas optimally we would need $2$. As shown in Table 1, Using CoTFormer instead of Block Universal Transformer, significantly reduces this gap while at the same time maintaining alternative advantages such as smaller model size and the adaptivity of compute at inference. Still, reducing this gap further remains an important direction for future work.

Efficient Implementation. Since CoTFormer introduces additional tokens to the sequence, the attention module’s implementation, in particular the causal mask, needs to be adapted. In this work, we rely on a simple implementation (using non-causal version of Flash Attention []) which leaves room for improvement. In particular, low-level kernels such as [] can be directly used to improve the speed of CoTFormer’s implementation.

## 6 Conclusion

In this work, we point out an often overlooked distinction between chain of thought and iteratively applying the model. By taking this distinction into account we develop CoTFormer and show its superior performance to Block Universal Transformers. Most noticeably we propose additional small tweaks in the architecture, allowing a CoTFormer to obtain the same accuracy as a standard Transformer that has double its size. Moreover, we propose an adaptive training method and showed it enables adjusting the compute budget at inference in exchange with reasonable impact on accuracy. Unlike prior methods, our method does not introduce sensitive additional hyperparameters and allows for stable training. Finally, we discuss different avenues to improve the results, particularly in adaptive setting, in future work.

## Appendix A Code

Our implementations for all expeirments is available at [https://anonymous.4open.science/r/cotformer_neurips2024_anonym-BF66](https://anonymous.4open.science/r/cotformer_neurips2024_anonym-BF66).

## Appendix B Effect of Depth Embedding

**Table 3: Effect of Depth Embedding on fixed and adaptive number of repeats. The performance is similar on fixed repeats but is noticably improved in the adpative case.**
| Model | Adaptive | ${n_{\text{repeat}}}=5$ |
| --- | --- | --- |
| LN-CoTFormer | 25.08(0.03) | 24.11(0.03) |
| + Depth Embedding | 24.94(0.01) | 24.17(0.08) |

## Appendix C Alternative Adaptive Training Methods

We experimented with Stick Breaking method proposed by Tan et al. [] as well as the halting mechanism in PonderNet [].

As acknowledged in the same work, we found training with PonderNet mechanism to be challenging and sensitive to the choice of hyperparameter, specifically the weight of the KL divergence. We tried tuning this parameter and report the best result in Table 4.

When using Stick Breaking Halting, we observed that the model tends to be very conservative and as a result the average depth remains too low.

We compare the results of training block universal transformer for 10k iterations in Table 4. We observed better final perplexity with our method (Mixture of Repeats) than the other two methods. Due to computational limits, we could not perform longer experiments but decided to use our method given the more stable and less sensitive training dynamics as well as the better performance.

**Table 4: Comparison of Mixture of Repeats with Previous Mehtods.**
| Method | Perplexity |
| --- | --- |
| Stick Breaking | 33.82 |
| PonderNet ($\lambda_{p}=0.4$) | 41.37 |
| Mixture of Repeats | 33.08 |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: * I’m confused by Figure 1(c): it seems to indicate that the earlier token representations (i.e. the red rectangles) are reprocessed by the model to make new token representations. But Section 3.1 seems to contradict this and instead describes that these earlier representations are only used as context in attention.
* It would be helpful to state the total parameter counts of the models used in each experiment, as well as the total number of training tokens in each experiment (either in a table or in the prose describing experimental setup).
* 471: It would be helpful if the authors provided more details on their “efficient implementation”, and specifically how the authors are using a non-causal FlashAttention kernel to implement their  proposed method.= 
* How do position embeddings work with the added interleaved tokens? Are the interleaved tokens given the same position id as the original tokens they came from, do the position ids change between repetitions, or something else?
* Do the authors have any intuitions as to how their method behaves as the width of the model changes? It appears to be held constant across all experiments.
* 402: what does it mean to “activate a prefix of repeats”? is this the fixed depth baseline that is referenced in Figure 4?
How does mixture of repeats work during, for example, batch size 1 transformer decoding, where there is only a single token being processed through the model?


Below are some thoughts that might be helpful but are not critical to give insight into ways that might improve the paper.
* Consider analyzing attention patterns and strengthening theoretical connections to transformer expressivity research (building on 3's architecture analysis).
* Explore sparse variants to improve scaling for longer sequences beyond 8192 (extending the computational analysis in 3.2).
* Focus on improving training efficiency, particularly for deeper layers and adaptive computation (addressing limitations discussed in 5).
* Develop specialized attention implementations for better computational performance (following the implementation discussion in 5).
* Expand evaluation to include longer sequences and broader comparisons with other adaptive approaches (extending the experimental work in 4.3).
soundness: 3
strengths: *The application of deploying LLMs to storage-constrained devices like mobile phones is relevant and timely.
* The proposed CoTFormer architecture (Figure 1(c)) effectively translates the CoT principle into architectural design, showing clear improvements over baseline approaches while maintaining parameter efficiency through weight sharing (Section 3.1).
* The architectural tweaks introduced in Section 3.3, particularly reserved layers and layer normalization after each repeat (LN-CoTFormer), prove crucial for achieving state-of-the-art results.
* Addition of depth embedding (Section 4.2) shows notable improvements in the adaptive setting.
* While not outperforming a FLOP-matched non-repeated transformer, the authors improve upon existing parameter-matched weight sharing baselines.
* The proposed architecture and adaptive repetition method are described clearly.
summary: This paper introduces CoTFormer, a novel transformer architecture that draws inspiration from chain-of-thought (CoT) reasoning. The key insight is recognizing that CoT differs from simple weight-tying in how attention operates across intermediary reasoning steps. The authors leverage this insight to develop an architecture that allows tokens to attend to representations from all previous "thought" steps, leading to improved performance compared to baseline approaches like Block Universal Transformers. Additionally, they propose an adaptive computation mechanism that allows dynamic allocation of computational resources at inference time.
weaknesses: * While Section 3.2 provides empirical evidence, the theoretical understanding of why CoTFormer works better could be deeper. Through analysis of attention patterns, we observe that tokens in later repeats tend to focus heavily on earlier representations that capture key contextual information, suggesting the model learns to leverage complementary features detected at different processing stages. This selective attention to informative past representations may help explain why CoTFormer outperforms the baseline Block Universal Transformer, where such cross-repeat attention patterns are not possible.
* Could better connect to recent theoretical work on transformer expressivity discussed in Section 2.
* The sequence lengths that are used for training (256) are quite short relative to the lengths that are used for training modern language models and are shorter relative to common LLM evals and typical chatbot conversations.
* Performance gap between adaptive and fixed-depth CoTFormers under same compute budget (Section 5)
* Training efficiency of deeper layers could be improved (e.g. increasing the gradient information during adaptive training), as shown by the analysis of router weights distribution (Figure 5)

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - Typo in Line 114, "similar the" --> "similar to the"


- Is it possible to convert a standard pre-trained Transformer to a CoTFormer via a post-training or fine-tuning phase, which can be much more efficient than pre-training a CoTFormer from scratch?
I can't see an obvious way of doing this, since the behavior of a CoTFormer deviates significantly from that of a standard Transformer.
soundness: 2
strengths: - This work proposes a novel Transformer-based model architecture, and draws an interesting link to chain-of-thought.
The proposed CoTFormer is compatible with KV cache and batch processing, which is not the case for many other adaptive-computation architectures tailored to autoregressive generation.

- Overall the writing is good, and things are mostly well explained.

- The source code is provided.
summary: This work proposes CoTFormer, a novel Transformer-based model architecture for generative language models.
Like Universal Transformers, a CoTFormer re-applies the same Transformer blocks for an adaptive number of times for generating each token.
The major difference is, after each repeat, the output tokens interleaved with input tokens are used as the new input for the next repeat;
this is inspired by chain-of-thought where each generated "thought token" can attend directly to all previous thought tokens.
Other details are handled, such as compatibility with KV cache and batch processing during inference.
Empirical results show that a CoTFormer achieves lower perplexity or higher scores in some common benchmarks 
than Block Universal Transformer with the same configuration, 
or a standard Transformer of the same size.
weaknesses: My major concern is that the empirical evidence for the efficacy of CoTFormer, or its advantages over the standard Transformer architecture, is insufficient.

- Generally speaking, the empirical results from the main text suggest that a CoTFormer with $n_{\text{repeat}} \ge 2$ slightly outperforms a standard Transformer of the same size in terms of perplexity, but underperforms a standard Transformer with twice as many layers, except in Table 2 where a CoTFormer with $n_{\text{repeat}}$ as large as 5 (and other tweaks) achieves a perplexity that is lower by only 0.06. 
The issue is that the inference cost (in terms of time or memory, or both) of a CoTFormer, with the total number of tokens growing linearly with $n_{\text{repeat}}$, can possibly be larger than that of a standard Transformer with twice as many layers.
This raises the question of whether CoTFormer actually pushes forward the Pareto frontier of accuracy and cost; to support such a claim, it is necessary to compare CoTFormer's accuracy-cost curves with those of standard Transformers (not just Block Universal Transformer).
Without clear evidence of its advantages over standard Transformers, the additional complexity overhead to code and infrastructure might further hinders the adoption of CoTFormer in future research or applications.

- The results of downstream performance in Appendix B have limited implications, as discussed by the authors in Line 725. 
    For example, all scores for MMLU are close to 25%, namely the accuracy of randomly picking option A/B/C/D.

- The current work only contains end-to-end performance (perplexity or scores) on some common datasets and benchmarks.
    There is no intermediate empirical result (except for Figure 5) or synthetic task, like those in the original paper of Universal Transformers (Dehghani et al., 2019), for truly understanding when, why and how CoTFormer works or fails.
The authors might consider visualizing the attention patterns of CoTFormer, or designing synthetic tasks that highlight CoTFormer's fundamental advantages over standard Transformers or Universal Transformers.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Figure 2 shows the inference FLOPs vs. Perplexity. However, it cannot suggest better "scaling properties of CoTFormers" (quote Line 257 of the paper) because scaling properties should be suggested by the training FLOPs vs Perplexity following Kaplan et al.[1]. Could you provide the training FLOPs vs. Perplexity plot for Figure 2?
2. Could you add the standard Transformer to Figure 2?
3. The paper claims that "The growth in computation cost is actually much less noticeable". Could you provide the real measurement of computation cost in terms of memory footprint (Figures 2 and 3 only show FLOPs)? 

**After discussion period, questions were addressed by the authors:

Answer 1: Keep all the other factors constant, the scaling behavior with respect to the training FLOPs still holds.

Answer 2: The accuracy of the standard Transformer in Table 1 can indicate the distance between the CoTFormer and the standard Transformer. Therefore, it is necessary to add the standard Transformer to Figure 2.

Answer 3: Theoretically, the memory footprint is the same for CoTFormer and Block Universal.
 

[1] Kaplan, Jared, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. "Scaling laws for neural language models." *arXiv preprint arXiv:2001.08361* (2020).
soundness: 2
strengths: 1. The CoTFormer architecture and Mixture-of-Repeats approach effectively improve the performance and efficiency of the Universal Transformer.
2. The evaluation of downstream tasks illustrates the CoTFormer's potential to surpass the standard Transformer.
summary: The paper proposes a new model architecture called CoTFormer that improves the Block Universal Transformer by providing intermediary representations from previous repeats in the attention. Besides CoTFormer architecture, the paper also proposes a training approach called Mixture of Repeats that varies the number of model passes for individual tokens based on their difficulty. Results show that CoTFormer substantially improves accuracy and inference computation efficiency over Block Universal Transformer.
weaknesses: 1. Possible misuse of technical terms: Chain-of-thought is a prompting technique. The process illustrated by Figure 1 (a) is called auto-regressive, which is orthogonal to CoT. Could the authors clarify how the CoTFormer model architecture model relates to the CoT prompting? Could CoT prompt be applied to CoTFormer model?
2. The model architecture is not clearly explained. Specifically, the meaning of different colors in Figure 1 is vague. Why are there no yellow tokens in Figure 1 (b) and (c)? The figure can be more clear if the caption explains the reason for the absence of yellow tokens and the meaning of different numbers of tokens.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Is the compute budget a hyperparameter to tune to achieve an optimal balance between accuracy and computation?
soundness: 2
strengths: 1. The architecture is novel and the authors made smart observations with respect to the CoT can access previous tokens.
2. The problem is very practical and of high interest to the community, especially with limited computation resources
3. The paper is overall well-written and organized.
summary: The paper presents CoTFormer, a novel Transformer architecture that leverages the Chain-of-Thought mechanism to enhance model performance while allowing for budget-adaptive computation at inference. CoTFormer enables intermediate tokens to be accessible, improving accuracy without significantly increasing computational costs. The authors further propose an adaptive training method that dynamically allocates computational resources based on the needs of individual tokens. Empirical results demonstrate that CoTFormer outperforms existing models, such as the Block Universal Transformer, while maintaining a smaller model size.

(Note: Thank the authors for their clarification; that addressed some of my concerns. I've adjusted the score accordingly)
weaknesses: 1. The paper could have a more detailed discussion on the scalability of the architecture, with respect to larger models and higher sequence lengths, since the paper discusses that the attention computation is not the bottleneck.
2. The paper studies the performance of CoTFormer on a particular dataset; would be interesting to see the performance on other datasets
3. The paper could have benefited from a more thorough theoretical analysis of COTFormer, especially with the number of repeats compared to the block universal transformer

Meta-review: Summary:
The paper introduces CoTFormer, a novel transformer architecture that draws inspiration from chain-of-thought reasoning. The key innovation is allowing tokens to attend to representations from all previous "thought" steps, which differs from both standard transformers and Block Universal Transformers. The authors claim this leads to improved accuracy while maintaining parameter efficiency through weight sharing. They also propose an adaptive computation mechanism that dynamically allocates computational resources during inference based on token difficulty.

Strengths:
- Novel architectural insight connecting chain-of-thought reasoning to transformer design
- Practical application for compute-constrained environments
- Clear empirical improvements over baseline approaches
- Effective adaptive computation mechanism
- Thorough ablation studies and analysis
- Well-documented implementation details

Weaknesses:
- Limited theoretical understanding of why the architecture works better
- Experiments focused on relatively short sequence lengths (256) compared to modern standards
- Some gaps in performance between adaptive and fixed-depth variants
- Training efficiency challenges for deeper layers
- Limited exploration of performance at larger scales
- Could benefit from more diverse dataset evaluation

Reasons for Acceptance:
- The paper presents a novel and practical architectural innovation that shows clear improvements over existing approaches
- The work addresses an important practical challenge (compute-efficient language models)
- The empirical results, while not revolutionary, demonstrate consistent improvements
- The authors have been responsive to reviewer concerns and provided additional experiments
- The adaptive computation mechanism adds practical value
- The work opens up interesting directions for future research

Decision: Accept (Poster)