CoTFormer: A Chain of Thought Driven Architecture with Budget-Adaptive Computation Cost at Inference
Authors: Amirkeivan Mohtashami
Matteo Pagliardini
Martin Jaggi
Venue: ICLR 2025 Poster (2025)
Keywords: language models
adaptive compute
chain of thought
Primary area: generative models

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.

Published as a conference paper at ICLR 2025
COTFORMER: A CHAIN-OF- DRIVEN AR- CHITECTURE WITH BUDGET-ADAPTIVE COMPUTATION COST AT INFERENCE
Amirkeivan Mohtashami EPFL
Matteo Pagliardini EPFL
Martin Jaggi EPFL
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-constrant environments, the race to continually develop ever larger and deeper foundational models is underway. At the same time-regardless of the model size-ttak-sseccfiic techniques continue to play a pivotal role in achieving op- timal downstream performance. One of these techniques, called Chain-of-Thougge (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, 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 com- pensate for it by leveraging CoTFormer's 's special compatibility with token-wing variable depth. Through a compute adaptive model-which automatically allocates the compute to tokens that need 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 competitive accuracy.
INTRODUCTION
Large foundational models have demonstrated remarkable performance across various tasks, predom- inantly employing the Transformer architecture (Vaswani et al., 2017). The ability to tackle new tasks in zero-shot or few-shot settings (Brown et al., 2020) has been attributed to emergent properties that become increasingly prominent with model size (Wei et al., 2022a). This observation has sparked race to build progressively larger models (Brown et al., 2020; OpenAL, 2023; Touvron al., 2023a;b).
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 (Cobbe et al., 2021). In response to these challenges, an alternative approach called Chain-of-Thhuulee (CoT)) (Wei et al., 2022b) has been proposed, requiring models to think step by step and articulate their thought processes, showing remarkable success (Kojima et al., 2022). In particular, using CoT can improve the general performance of even smaller models Ho et al. (2023); Li et al. (2024).
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 (Dehghani et al., 2019). However, in this work, we point out that there is 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.
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
"Equal contribution. order is alphabetical.

---

Published as a conference paper at ICLR 2025
token
TTandforme Bloker
Tranformer Modle
SooK todboss prdictin
N Transfonnee Blodis
Trandformmee Blockc
uongoo 1900
Model
Model
abared weichts
shared weights
(a) CoT
(b) Block Universal Transformer
(c) CoTFoormer
Figure I: Block universal transformer vs. CoTFormer vs. Chain-of- Chain-of-Thought (CoT) reasoning. In (a) we represent the shain-of-thoughee mechanism in which model is iteratively generating reasoning tokens to help solve downstream applications. Based on existing input red tokens, next token (blue) is generated and added the sequence, re-iterating this process yields the green and yellow tokens. Here we emphasize how (i) the last red tokens is "pushed" several times through the model-the yellow token being the red token after three successive applications of the model-and (ii), new (e.g. blue) tokens
attend to previous (e.g. red) tokens, this observation is the basis of CoTFormer. In (b) we represent the unive-uliverral transformer which recursively applies the same N transformer blocks to the input tokens. This approach is to be contrasted with the CoTFormer architecture (c) which interleaves old and new representations in between each block. In the figure this process is done times (Prepeat 2), but could be repeated many more times. As in COT, and unlike block universal transformers, later (c.g. blue) tokens can attend to carlier (e.g. red) tokens.
such as Universal Transformers Dehghani et al. (2019), and pushes the cerplexity-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 intermediary reasoning steps before reaching an answer. In
to make explicit greater number 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 long-standing goal, driving the exploration of architectures that can be recurrently applied controlling the compute cost through deciding the depth of the recursion Banino applied-controllingg
al. (2021); Dehghnn al. (2019); Elbayad et al. (2020); Graves (2017); Liu al. (2020); Tan et al. (2023). 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 ws and W2 at depth 3, given that W2 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 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 ws can simply access all the tokens which have been generated before through the attention mechanism. This makes the CoTFormer a much more natural 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 model that allows choosing the computation cost at inference time, and navigating the cewtrracc--oomuaatttiin trade-off without additional training. This is in contrast with most current models that have 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- ofoo-ghht-t-hhss and the recurrent application of a model with weight-tying.
Proposing CoTmerer which accounts for the aforementioned distinction, and demonstrating its superior performance over other weight-tied deep models (c.g. Universal Transformer (Dehghani al., 2019)).
2

---

Published as a conference paper at ICLR 2025
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 Bolukbasi et al. (2017); Graves (2017).
These methods usually rely applying the model multiple times, simulating deeper model with weight tying. In many aspects, this approach is similar the widely used technique of instructing the model generate intermediary thoughts before outputting the final answer, called Chain of Thought (CoT) Wei et al. (2022b). 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 Feng et al. (2023). Furthermore, while Transformers with fixed depth are not Turing complete on their own Merrill & Sabharwal (2023), combining them with the auto-regressive decoding used for generating the chain of thought can make them Turing complete Malach (2023); Merrill & Sabharwal (2024). 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. (2019) propose Universal Transformers which repeatedly applies single layer transformer model on the model. A predictor is trained using the ACT method Graves (2017) decide whether to stop or apply the model again. Due the instability of ACT and its sensitivity to hyperparameters, Banino et al. (2021) propose PonderNet which weights the predictions 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 Takase & Kiyono (2021).
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 Liu et al. (2021). 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 similar challenge when deciding whether to skip blocks of standard Transformer architecture, Raposo et al. (2024) propose Mixture-of-Deppe 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 Formrrress, Mol) 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 Formers which apply full prefix of blocks, MoDs decide separately whether to use each block or not, preventing early exiting.
(MoD) defining a
Recent work have explored and proposed a variety of alternative architectures which prove to be useful in different scenarios Pagliardini al. (2024); Wang et al. (2022). Most prominently Mixture- of-Expssss (MoEs) have been shown to improve performance of the model in many cases Jiang et al. (2024). For example Sparse Universal Transformers Tan et al. (2023) combine the idea of universal transformer with MoEs, allowing a router to choose possibly different model every time the input is processed again. In this work we mainly focus our experiments on the Pre-LayerNorm Transformer architecture Xiong et al. (2020), which is currently the most widely used architecture and is the backhone of the state of the art language models Jiang et al. (2023); Touvron et al. (2023c). 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 Nye al. (2021) 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
3

---

Published as a conference paper at ICLR 2025
work to generate artificial data Ho et al. (2023); Zelikman et al. (2024). Regardless, we believe this approach to be complimentary to CoTFormers. Intuitively, CoTFormers allow re-using basic modules such 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 higher level to make rationale arguments.
Finally, while Block Universal Transformers simulate a deeper model, and while alternative proposals such Pause Tokens simulate model with increased hidden dimension (i.e. width) Goyal et al. (2024), CoTFormers intuitively facilitates both. Still, width-increasing methods such as Pause Tokens can be combined with CoTFormers.
CHALN-OF-THUUGH AND MODEL DEPTH
Chain-of- Chain-of-Thought 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 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 Recurrent Neural Network, RNN Rumelhart et al. (1986); see Figure la). Consequently, one might be tempted to frame the chain-of-thought process as the utilization of a deeper model with tied weights (see Figure 1b). Indeed, such arrangement resembles a version of Universal Transformers (Dehghani et al., 2019) 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 CoTForrrer 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)
(0)
[
and a current input token Bio
1:maa
11 we describe the process of generating the next token for the Block Universal Transformer and our CoTFormer model. First, let (O c) be the i-th repeat of a set of Raaye transformer blocks taking input the token and being able to attend to the context through its attention mechanism. One can imagine x being the current token being processed and e being the key/value-cache, as often (i)
used during inference. For typical transformer, generating the output of ( !) for token 2 can be written as follows:
(++1) Bog == ((++1)) Ren
(I)
Repeating the above formula neepeat times with weight tying between the B(), S í Prepent yields the Block Universal Transformer:
(i+1) B
B(x
1:71
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 follows:
(xx
(0) Free
(4) 11msse
Figure Ic illustrates the above process. It can be seen that the sequence length grows linearly with the number of repeats. This does not have 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 inference. However, it may impact the computational cost which we discuss in Section 3.2. We use the notation MhyerrX/rreeatt to describe a CoTFormer or Block Universal Transformer with Mayer layers being repeated Nrepeat times. Furthermore, we use same position ids for the intermediary representations as the corresponding original token.
3.2 COMPARISON WITH BLOCK UNIVERSAL TRANSFORMER
Experimental setting. To establish the importance of attending previous intermediary states, we compare the performance of CoTFormer and Block Universal Transformer on the Open Web- Text2 (Gao al., 2020) dataset; dataset of websites linked from reddit between 2005 and 2020
4

---

Published as a conference paper at ICLR 2025
Table Performance of CoTFormer, Block Universal Transformer and Standard Transformers on Open WebText2. The mean perplexity of 3 runs is reported with the standard error of the mean in parenthesis. can be seen that CoTFormers clearly outperforms Block Universal Transformers. The best perplexity for given Peyeert combination is marked in bold.
Model
Base Layers (Mtayer)
Prepest
2
Standard
12
28.39 (0.01)
Block Universal Transformer CoTFormer
12
27.74(0.01)
27.47(0.01)
27.15(0.02)
12
27.55(0.02)
27.07(0.01)
26.64(0.04)
Standard
24
25.93 (0.02)
Block Universal Transformer CoTFormer
24
25.47(0.01)
19.19(0.03)
24.95(0.01)
24
25.28(0.00)
(0.85(0.04)
24.48(0.03)
Standard
48
24.17 (0.00)
initially released under MIT license. We train the models for 40k steps using the A (Loshchilov & Hutter, 2019) optimizer and apply linear warmup of the learning rate for the first 8000 steps. We use the Fre-rayerNorm Transformer Xiong et al. (2020) 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 ForFers 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.
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 Ganesh al. (2021); Tay et al. (2022). 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-accumulat accumulate operations needed to process different sequence lengths by block universal transformer with Prppeat = and a CoTFormer with nepeea = which obtains similar accuracy (on sequences of length 256). It can be seen that even for sequences as long as 8192, the 12x3 FoTFormer's is cost remains below that of 12x5 Block Universal Transformer.
The demonstrated or is further depicted in Figure which shows the perplexity against compute cost of processing sequence with 256 tokens for CoTFormers and Block Universal Transformers with Mayer = 12 and Faayer = 24. At both scales, it can be clearly seen that while Formrre with the same number of repeats
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 CorFrrmess. The above results clearly demonstrate the effectiveness of CoTFrrrers 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 exccute Neggia layers, followed by multiple passes through Redddlee layers. Finally the output is generated from the final pass of each token by passing it through the last Rend layers.
5

---

Published as a conference paper at ICLR 2025
12x2
Block Universal
Universal hanssorme (12x5)
27.50
Muloppy- Sapp-cccuuoooo Operatoons
1ga
ColiForner (12x3)
27.25
Perplexity
12x6
27.00
12x3
12x15
102
26.75
26.50
10
26.25
12x15
100
300
500
125
235
1034
2048
81R6
12246
-uttpllyAAccumultee Operations (x109)
aa Paaye
12 (x-axis is in log scale)
Figure 3: CoTFormer is less compute intensive than Block Universal Transformer of compa- rable performance. Comparing 12 layers CoT- Former with 3 repeats (12x3) and 12 layer Block Universal Transformer with 5 repeats (12x5) in terms of computation cost. The CoTFormer's s accu- racy is better than the Block Universal Transformer (see Figure 2). Despite the increase context length when processing the input with Cormorrer the computational cost of CoTFormer remains be- low the Block Universal Transformer for sequence lengths as high as 8192.
Block Universal Transformer CoTFormer
25.4
24x2
24x3
25.2
Perplexity
24.8
24.6
24x5
100
125
150
175
200
225
250
MultipyyAAcummlate Operations (x109)
Table 2: Ablation study for the architecture tweaks discussed in Section 3.3. The final archi- tecture with Neppeat 5 obtains lower perplexities than a 48 layers standard transformer which has double its size.
(b) Wayyer =
Figure 2: Comparison of Block Universal Transformer and CoTFormer in terms of accuracy-compitation tradeoff. It can be clearly seen that at both Maaer = 12 and Maaaer 24, CoTFormers are closer the Pareto frontier. The gap widens with larger number of repeats, suggesting better scaling properties of CoTFormers. The x-axis shows the number of operations for processing se- quence of length 256.
Model
Perplexity
Standard
48x1
24.17(0.00)
CoTFormer
24x5
24.48(0.03)
+ Reserved Layers
2-+21x5
24.51(0.01)
+ Layer Norm
2-+21x5-11
24.11(0.03)
Layer Norm After Each Repeat. Given the internal residual connections of the model, we conjecture that it is important to maintain consistent input's scale. Therefore 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 2. 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 24. Most noticeably, after applying these changes, the performance of a CoTFormer with 24 layers and 5 repeats, surpasses the performance
standard 48 layer Transformer. We present similar results for downstream tasks in Appendix B. We call the final resulting architecture LN-CoTFormer. We note that while LN-CoTFormer' 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, focus using FN-CoTFormmrs when building adaptive model in the next section.

---

Published as a conference paper at ICLR 2025
4 TOKEN- WISE ADAPTIVE REPEATS
The standard CoTFormer has the advantage of obtaining better performance with smaller models which is useful in conttro--nntt 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 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-Haiee Haltedd Tokens in Later Layers: If 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 '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.
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 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. (2024) 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 randomized capacities, detailed in the next subsection.
Efficient Batch Processing:
In addition to addressing the above challenges, we also aim to build 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 Mohtashami et al. (2022); Yu et al. (2018). 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 reeepet to refer to the maximum number of times a token can go through the model. For each of the passes 2 to nreepst we train a separate embedding vector, e.g. e( 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 (0) we compute the score s/ (+)) = (( (0(()) 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 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 & is chosen based on the capacity level for this repeat: 4. In particular, if we denote the sequence length by nheq. we will use A == [G x negg]]
Let us denote the output of the model on the input x by B(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 ()) This interpolation plays two roles. First of all, the gradient use (6+ (1 s,) I; + s; needed to train the embeddings e 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 repeat to 1,
7

---

Published as a conference paper at ICLR 2025
even tokens with very low scores will be selected. However, 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 q₁ we sample them at random for each batch. More particularly, we assign a capacity of 1 to the first repeat and sample Prepeat - 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 pass would affect the gradient. Therefore, the update for router weights would only take into account those high scoring tokens. As 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 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 depth embedding to the model.
Depth Embedding. In order to allow the model recognize how many number of passes it has done, we add depth embedding at the start of each pass. For this embedding we learn a single vector e (depph) and add (Nrepea i) e (depph) 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 4. 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 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 4 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 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-compe trade-off at inference time which 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 D and continue with reporting the results of training 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 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 sacrifice some performance when the number of repeats remains low.
8

---

Published as a conference paper at ICLR 2025
30
LN Formmer (Router)
CoTFormer (40k Training)
LN CoTFormer (Fxed Depth)
CoTFormer (60k Training)
Black Universal (Router)
Perplexity
28
0.5
50
100
150
200
0.2-
MultpplyAAcummulate Operations (x11")
Figure 4: Perplexity for different amount of compute budgets chosen at inference. The adap- tive CoTFormer can adapt to different budgets, re- ducing compute in exchange for reasonable loss in accuracy. Furthermore, using the router weights allocate the available compute (Router) much more effective than fixing the depth at inference time to a smaller value in order reduce compu- tation cost (Fixed Depth).
0.0
0.0
0.2
0.4
0.6
0.8
1.0
Repeat Weight Prediction
Figure 5: Distribution of router weights for the last repeat for different number of training steps. It can be seen that when training longer, the model learns more use the deepest repeat, leading to higher router weights.
5 DISCUSSION AND FUTURE WORK
Training of Deeper Layers. While the above performance is remarkable, between an adaptive CoTFormer and non-adaptive CoTFormer trained with exactly 5 repeats even when the adaptive variant is allowed the same amount of compute 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 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 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.5x longer.
observe gap
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, could not find method that performs better than randomly picking and sorting described in Section 4.1. Still, we expect better methods to exist and leave exploration of such methods as 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 48 layer model using a 24 layer model, 5 repeats needed whereas optimally we would need 2. As shown in Table 1, Using CoTForrrer 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 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 simple implementation (using non-causal version of Flash Attention Dao et al. (2022)) which leaves room for improvement. In particular, low-level kernels such as Pagliardini et al. (2023) 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 develop FoTForeer and show its superior performance to Block Universal Transformers. Most noticeably we propose additional small tweaks in the architecture, allowing CoTFormer obtain the same accuracy standard

---

Published as a conference paper at ICLR 2025
Transformer that has double its size. Moreover, propose 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.
ACKNOWLEDDEMENT
This project was supported by SNSF grant number (00020_200342
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Andrea Banino, Jan Balaguer, and Charles Blundell. PonderNet: Learning to Ponder, September 2021. URL http: /ttp///aavvvve.ee/eeeeeeeeaeaeva org. http://bb/221....5...... arXiv:2107.0548 07040007007 [es].
Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive Neural Networks for Efficient Inference, September 2017. URL
arXiv:1702.07811 [es, stat].
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voor Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter. Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray. Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners, July 2020. URL. http:://rxiv org/ abs 200020000.. 141 4166..6.. 65. arXiv:2005. 14165 [cs].
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training Verifiers to Solve Math Word Problems, November 2021. URL http:
14168. urXiv:2110, 14168 [cs].
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with 10-Awareness, June 2022. URL htt.htttp///rxi..orggg abs/22055 14135, arXiv:2205.1 14135 [es].
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In ICLR (Poster). OpenReview.net, 2019.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth Depth-Adaptive Transformer, February 2020. URL http: //tp://////iivii/ org/ /abs///////// 1910/991111111 0..00.10000 10073. arXiv:1910.10073 [cs].
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective, December 2023. URL http: // //arxiv..org./....a...a..
2305. . 1 5408. arXiv:2305. 15408 [cs, stat].
Prakhar Ganesh, Yao Chen, Xin Lou, Mohammad Ali Khan, Yin Yang, Hassan Sajjad, Preslav Nakov, Deming Chen, and Marianne Winslett. Compressing Large-Scale Transformer-Based Models: Case Study on BERT. Transactions of the Association for Computational Linguistics, 9:1061-1080, September 2021. ISSN 2307-387X. //arxiv. org/ abs 2002. 11985. arXiv:2002. 1995 [cs, stat].
doi: 10.1 atcpttallaa00Aae a_00413. URL http:
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. Open WebTex2 dataset, as part of 'the Pile: An 800gb dataset of diverse text for language modeling'. arXiv preprint arXiv:2101.0003 2020.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training Language Models With Pause Tokens, April 2024. org/ abs 2310. 02226, 6. arXiv:2310.022D. [cs].
URL wwwp::wwww.....//.pt/oooooo http://bbb//1...... 08983. 1603.066 [cs].
Alex Graves. Adaptive Computation Time for Recurrent Neural Networks, February 2017. URL http://arxiv.or
Namgyu Ho, Laura Schmid, and Se- Young Yun. Large Language Models Are Reasoning Teachers, 1007 71. arXiv:2212.10077 [cs].
June 2023. URL http:::/aaxxvv.oo////aaaaa
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud. Marie-Anne Lachaux. Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7B, October 2023. URL http: / //taxxi..rggaa/////////eo.ootata 66855. arXiv:2310.06825 [cs].
11

---

Published as a conference paper at ICLR 2025
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel. Guillaume Bour, Guillaume Lample. Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of Experts, January 2024. URL htttp://arxiv..../ arggabs///01..... 04088. arXiv:2401.04005 [es].
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35: 22199-22213, 2022.
Liunian Harold Li, Jack Hessel, Youngjae Yu, Xiang Ren, Kai- Wei Chang, and Yejin Choi. Symbolic Chain-of-Thought Distillation: Small Models Can Also "Think" Step-by-Step. April 2024. URL http:://axti orgo httg/ab//220 /306 14050. arXiv:2306. 14050 [cs].
Yijin Liu, Fandong Meng, Jie Zhou, Yufeng Chen, and Jinan Xu. Faster Adpth-veeiiieee Transformers, December 2020. URL http: arxiv. org/ / 004 13542. arXiv:2004.13542 [cs].
Yijin Liu, Fandong Meng, Jie Zhou, Yufeng Chen, and Jinan Xu. Faster depth-adaptive transformers. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 13424-13432. 2021.
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR (Poster). OpenReview.net, 2019.
Eran Malach, Auto-Regressive Next-Token Predictors are Universal Learners, Septemher 2023. URL
http::/arxie www./ogg////...oooo.oo.../....ee
http://b//23000 2309 06979. 2XIv:2339.06975 [cs].
William Merrill and Ashish Sabharwal. The Parallelism Tradcoff: Limitations Log-Precision
Transformers, April 2023. /Bt/stpp::/arr/...oor//a/b////// 00729. w0w.:0222.....220........ arXiv:2207.00729
[cs].
William Merrill and Ashish Sabharwal. The Expressive Power of Transformers with Chain of
Thought, April 2024. URL http::://axxvvoor//aaa/.. 079233 arXiv:2310.07923 [es].
Amirkeivan Mohtashami, Martin Jaggi, and Sebastian U. Stich. Masked Training of Neural Net- works with Partial Gradients, March 2022. URL http:://aax....ogg//aa///// 08895. arXiv:2106.08895 [es].
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show Your Work: Scratchpads for Intermediate Computation with Language Models, November 2021. URL /ttpo///..tppp/aaaa...o///////// wttp:///..g//////....///////// 00114. ar0vv211.22 [es].
OpenAL. GPT-4 Technical Report, March 2023, URL htt.://ttt////erarr.rrr//////22. arXiv:2303.08774 [es].
Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Faster Causal Attention Over Large Sequences Through Sparse Flash Attention, June 2023. URL htttttt//axxi..ooe abs/2306 01 01160. arXiv:2306.01 160 [cs].
Matteo Pagliardini, Amirkeivan Mohtashami, Francois Fleuret, and Martin Jaggi. DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging, March 2024. URL httt:://arxe orggabss 2402 02 622. arXiv:2402.02622 [es].
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mfxture-of-Deplle Dynamically allocating compute in transformer-batt based language models, April 2024. URL http: httt:://p//iiiiieee abs 4 02258. arXiv:2404.02258 [cs].
12

---

Published as a conference paper at ICLR 2025
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning internal representations by error propagation, parallel distributed processing, explorations in the microstructure of cognition, ed. de rumelhart and j. meclelland. vol. 1. 1986. Biometrika, 71:599-600 605599677,7 1986.
Sho Takase and Shun Kiyono. Lessons parameter sharing across layers in transformers. arXiv preprint arXiv:2 (6.2222 2021.
Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse Universal Trans- arXiv:2310.0709
former, October 2023. URL [es] version: I.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient Transformers: A Survey, March 2022. URL http: www.v.orggaas/aae..ooooot....aaa 2t0.//000000.2.....2 arXiv:2009.06733 [cs].
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan. Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux. Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu. Yuning Mao. Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open Foundation and Fine-Tuned Chat Models, July 2023c. URL BRLLhttp///axx org/ abs 2307 09288, arXiv:2307.00 .92288 [cs].
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, 0adp//0060077222 2017. URL http:://rxxti www.lv...rr//oooooeeoeeee
03762.
Hongyu Wang, Shuming Ma, Shaohan Huang. Li Dong, Wenhui Wang, Zhiliang Peng, Yu Wu, Payal Bajaj, Saksham Singhal, Alon Benhaim, Barun Patra, Zhun Liu, Vishrav Chaudhary, Xia Song, and Furu Wei. Foundation Transformers, October 2022. URL http: http:://axxiai /ttp:///axiv..v/ org/ hargg//bs/2211 06423. 2XXiv:22100064:4 [cs].
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent Abilities of Large Language Models, October 2022a. URL http: http:::://xxvvvv.a...o absrgg/bbs//// absa /200 07682. arXiv:2206.0768 [es].
13

---

Published as a conference paper at ICLR 2025
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. -hain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24833 2022b.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Ya Yan Liu. On Layer Normalization in the Transformer Archi- tecture, June 2020. URL /attp://araee stat].
owt.//ggg//a/ooooo 04745. arXiv:2002.04745 [es.
Jiahui Yu, Linjie Yang. Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable Neural Networks, December 2018. URL
1112. 08 928. arXiv:1812.08923 [cs].
Eric Zelikman, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, and Noah D. Goodman. Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking, March 2024. URL http:ttt///.. /tte://axxxii aawii. org/ /403 09629. arXiv:2403.09629 [cs].
14

---

Published as a conference paper at ICLR 2025
A CODE
Our implementations for all experiments is available https: sttpa:://ithub..com//gg cotformer.
B DOWNSTREAM TASKS
We evaluate the zero-shot performance of the models we trained on OpenWebText2 on several downstream tasks, and present the results in Table 3.
Model
MMLU
ARC
Hellaswag
PIQA
Average
Standard Transfromer (24)
25.83
29.34
27.41
58.54
35.28
Block Universal Transformer (24x2)
25.64
29.5
27.38
59.09
35.4
Block Universal Transformer (24x5)
26.07
30.25
27.73
58.27
35.58
CoTFormer (2->21x5>>1)
25.98
29.89
28.2
59.3
35.84
CoTFormer (24x5)
25.99
29.95
27.74
59.47
35.79
CoTFormer (24x2)
26.22
30.72
27.26
58.65
35.71
Standard Transformer (48)
26.11
29.31
27.93
60.28
35.91
Table Accuracy (Noramlized by Sequnece Length, Ignoring Space). The model's result with best performance between Block Universal Transformer and CoTFormers in each task is shown in bold.
As expected, we can observe patterns similar to the perplexity results presented Section 3.2, with CoTFormer outperforming Block Universal Transformers. We emphasize that these results should be interpreted only comparison with each other as obtaining good downstream tasks performance requires very long training which is not possible due to computational budge limitations.
C EFFECT OF DEPTH EMBEDDING
Table Effect of Depth Embedding on fixed and adaptive number of repeats. The performance is similar on fixed repeats but is noticeably improved in the adaptive case.
Model
Adaptive
neeepatt =
-o-mmer
25.08(0.03)
24.11(0.03)
+ Depth Embedding
24.94(0.01)
24.17(0.08)
15

---

Published as a conference paper at ICLR 2025
D ALTERNATIVE ADAPTIVE TRAINING METHODS
We experimented with Stick Breaking method proposed by Tan al. (2023) as well as the halting mechanism in PonderNet Banino et al. (2021).
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 5.
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 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 5: Comparison of Mixture of Repeats with Previous Mehtods.
Method
Perplexity Stick Breaking
33.82
PonderNet (Ap 0.4)
41.37
Mixture of Repeats
33.08
E COMPARISON WITH PAUSE TOKENS
In Goyal et al. (2024), the authors show adding a number of virtual tokens, called pause tokens, after each original token in the input, leads improvements in perplexity. While CoTFormer also adds additional tokens after each original token, these new tokens are not just placeholder tokens. Instead, they intermediate representations of the model and each subsequent token is created by passing the last token through the model again. To demonstrate that this is important for the performance of CoTFormer, in Table we compare model with 4 pause tokens and CoTForm with 5 repeats. can be clearly observed that CoTFormer significantly outperforms pause tokens.
Table 6: Comparing the performance of CoTFormer and pause tokens (Goyal et al., 2024). Adding pause tokens improves perplexity but is still heavily outperformed by CoT'orrmer that uses the output of previous repeat for the subsequent processing by the model.
Model
Base Layers (Maayer
Prepeat =
Standard
24
25.93 (0.02)
Pause tokens (Goyal et al., 2024)
24
25.05(0.03)
Block Universal Transformer
24
24.95(0.01)
CoTFormer
24
24.48(0..03)
Standard
48
24.17 (0.00)
F ATTENTION PATTERNS
In this section, we present some the attention patterns we observed in a 24 layer CoTFormer with 5 repeats. While, a thorough discussion around understanding how these models operate is outside the scope of the current work, we hope these results encourage such investigations.
In particular, we plot the average attention pattern the last token in the last repeat in Figure 6. We observe interesting patterns. In particular, current or recent tokens. Moreover, we observe heads that attend to tokens generated during a specific repeat. This is especially common for the first and last repeat though some heads also focus other
notice heads that attend to different repeats for the
16

---

Published as a conference paper at ICLR 2025
repeats. These patterns suggest that the model does not only rely having access to intermediate representations of the current token but also uses intermediate representations the previous tokens.
Head
0.03
Bappst
3-
0.02
4-
0.00
50
100
150
200
250
Layer Head
0.020
Repeat
0.015
0.010
0.005
0.000
50
100
150
200
250
Layer
Ropest
0.01
50
100
150
200
250
Layer 9 Head 10
0.020
0.015
Repeat
0.010
0.005
0.000
100
150
200
250
Layer 23 Head 2
0.06
Repeat
0.04
0.02
100
200
Layer 23 Head 6
0.02
Repoor
0.01
0.00
100
150
200
250
Figure 6: Attention Patterns of a 24 layer CoTFormer with Prepeat 5. Each figure shows the average of attention scores validation data from the last repeat of the last token in sequence all other token-repeat pairs. The x-axis shows the token index whereas the y-axis shows the repeat that the token belongs to. The intensity of the color shows how high the attention score specific token at a specific repeat has been for the head being considered (averaged over validation data).
17

---

Published as a conference paper at ICLR 2025
G EFFECT OF SEQUENCE LENGTH AND WIDTH
For the main experiments in the paper focus on models with 768 hidden dimension and 256 sequence length. Here, we also present results for 1024 hidden dimension Block Universal Trans- former and CoTFormer and also compare these models with 512 seqeunce length (with 768 hidden dimension). The results are shown in Table 7 and clearly demonstrate the benefits of CoTmorrer over Block Universal Transformers persist on different widths and sequence lengths. All models have 12 layers and use Prppett 5.
Table 7: Comparing CoTFormer and Block Universal Transformer different width and different sequence lengths.
Model
Base Layers (Haye)
Hidken Dimension
Sequence Lengih
Plopeet 5
Block Universal Transformer
256
27.15(0.02)
ColFormer
12
768
256
26.64(0.04)
Universal Transformer ColFormer
512
21.53
12
768
512
21.06
Block Universal Transformer CoTFormer
12
1024
256
24.89
12
1024
256
24.56
18


Reviews:
[{"review_id": "cEeV0gUERe", "reviewer": "Reviewer_YaZu", "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.", "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.\n* 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).\n* 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.= \n* 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?\n* 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.\n* 402: what does it mean to “activate a prefix of repeats”? is this the fixed depth baseline that is referenced in Figure 4?\nHow 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?\n\n\nBelow are some thoughts that might be helpful but are not critical to give insight into ways that might improve the paper.\n* Consider analyzing attention patterns and strengthening theoretical connections to transformer expressivity research (building on 3's architecture analysis).\n* Explore sparse variants to improve scaling for longer sequences beyond 8192 (extending the computational analysis in 3.2).\n* Focus on improving training efficiency, particularly for deeper layers and adaptive computation (addressing limitations discussed in 5).\n* Develop specialized attention implementations for better computational performance (following the implementation discussion in 5).\n* Expand evaluation to include longer sequences and broader comparisons with other adaptive approaches (extending the experimental work in 4.3).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "*The application of deploying LLMs to storage-constrained devices like mobile phones is relevant and timely.\n* 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).\n* 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.\n* Addition of depth embedding (Section 4.2) shows notable improvements in the adaptive setting.\n* While not outperforming a FLOP-matched non-repeated transformer, the authors improve upon existing parameter-matched weight sharing baselines.\n* The proposed architecture and adaptive repetition method are described clearly.", "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.\n* Could better connect to recent theoretical work on transformer expressivity discussed in Section 2.\n* 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.\n* Performance gap between adaptive and fixed-depth CoTFormers under same compute budget (Section 5)\n* 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jT1CC7wLU6", "reviewer": "Reviewer_UfJo", "summary": "This work proposes CoTFormer, a novel Transformer-based model architecture for generative language models.\nLike Universal Transformers, a CoTFormer re-applies the same Transformer blocks for an adaptive number of times for generating each token.\nThe major difference is, after each repeat, the output tokens interleaved with input tokens are used as the new input for the next repeat;\nthis is inspired by chain-of-thought where each generated \"thought token\" can attend directly to all previous thought tokens.\nOther details are handled, such as compatibility with KV cache and batch processing during inference.\nEmpirical results show that a CoTFormer achieves lower perplexity or higher scores in some common benchmarks \nthan Block Universal Transformer with the same configuration, \nor a standard Transformer of the same size.", "questions": "- Typo in Line 114, \"similar the\" --> \"similar to the\"\n\n\n- 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?\nI can't see an obvious way of doing this, since the behavior of a CoTFormer deviates significantly from that of a standard Transformer.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This work proposes a novel Transformer-based model architecture, and draws an interesting link to chain-of-thought.\nThe 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.\n\n- Overall the writing is good, and things are mostly well explained.\n\n- The source code is provided.", "weaknesses": "My major concern is that the empirical evidence for the efficacy of CoTFormer, or its advantages over the standard Transformer architecture, is insufficient.\n\n- 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. \nThe 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.\nThis 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).\nWithout 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.\n\n- The results of downstream performance in Appendix B have limited implications, as discussed by the authors in Line 725. \n    For example, all scores for MMLU are close to 25%, namely the accuracy of randomly picking option A/B/C/D.\n\n- The current work only contains end-to-end performance (perplexity or scores) on some common datasets and benchmarks.\n    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.\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XpKkHW0giX", "reviewer": "Reviewer_rjXu", "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.", "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?\n2. Could you add the standard Transformer to Figure 2?\n3. 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)? \n\n**After discussion period, questions were addressed by the authors:\n\nAnswer 1: Keep all the other factors constant, the scaling behavior with respect to the training FLOPs still holds.\n\nAnswer 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.\n\nAnswer 3: Theoretically, the memory footprint is the same for CoTFormer and Block Universal.\n \n\n[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).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The CoTFormer architecture and Mixture-of-Repeats approach effectively improve the performance and efficiency of the Universal Transformer.\n2. The evaluation of downstream tasks illustrates the CoTFormer's potential to surpass the standard 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?\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P9xPBk4dTl", "reviewer": "Reviewer_jCWB", "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.\n\n(Note: Thank the authors for their clarification; that addressed some of my concerns. I've adjusted the score accordingly)", "questions": "1. Is the compute budget a hyperparameter to tune to achieve an optimal balance between accuracy and computation?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The architecture is novel and the authors made smart observations with respect to the CoT can access previous tokens.\n2. The problem is very practical and of high interest to the community, especially with limited computation resources\n3. The paper is overall well-written and organized.", "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.\n2. The paper studies the performance of CoTFormer on a particular dataset; would be interesting to see the performance on other datasets\n3. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)