Reasoning with Latent Thoughts: On the Power of Looped Transformers
Authors: Nikunj Saunshi
Nishanth Dikkala
Zhiyuan Li
Sanjiv Kumar
Sashank J. Reddi
Venue: ICLR 2025 Poster (2025)
Keywords: looped models
reasoning
language model
iterative algorithm
inductive bias
scaling
Primary area: foundation or frontier models, including LLMs
TL;DR: Looped models can solve many reasoning problems and have an inductive bias towards improving reasoning of language models

Abstract: Large language models have shown remarkable reasoning abilities and scaling laws suggest that large parameter count, especially along the depth axis, is the primary driver. In this work, we make a stronger claim --- many reasoning problems require a large depth but not necessarily many parameters. This unlocks a novel application of looped models for reasoning. Firstly, we show that for many synthetic reasoning problems like addition, $p$-hop induction, and math problems, a $k$-layer transformer looped $L$ times nearly matches the performance of a $kL$-layer non-looped model, and is significantly better than a $k$-layer model. This is further corroborated by theoretical results showing that many such reasoning problems can be solved via iterative algorithms, and thus, can be solved effectively using looped models with nearly optimal depth. Perhaps surprisingly, these benefits also translate to practical settings of language modeling --- on many downstream reasoning tasks, a language model with $k$-layers looped $L$ times can be competitive to, if not better than, a $kL$-layer language model. In fact, our empirical analysis reveals an intriguing phenomenon: looped and non-looped models exhibit scaling behavior that depends on their effective depth, akin to the inference-time scaling of chain-of-thought (CoT) reasoning. We further elucidate the connection to CoT reasoning by proving that looped models implicitly generate latent thoughts and can simulate $T$ steps of CoT with $T$ loops. Inspired by these findings, we also present an interesting dichotomy between reasoning and memorization, and design a looping-based regularization that is effective on both fronts.

Published as a conference paper at ICLR 2025
REASONING WITH LATENT THOUGHTS: ON THE POWER OF LOOPED TRANSFORMERS
Nikunj Saunshi', Nishanth Dikkala' Zhiyuan Li Sanjiv Kumar', Sashank J. Reddi
{nsaunshi, nishanthd, lizhi yuan, sanjivk, sashank jaahank)@@ooglle Jégoogle. com "Google Research, Toyota Technological Institute at Chicago
ABSTRACT
Large language models have shown remarkable reasoning abilitics and scaling laws suggest that large parameter count, especially along the depth axis, is the pri- mary driver. In this work, we make a stronger claim - many reasoning problems require large depth but not necessarily many parameters. This unlocks a novel application of looped models for reasoning. Firstly, we show that for many syn- thetic reasoning problems like addition, y-hop induction, and math problems, a k-layer transformer looped L. times nearly matches the performance of a k L--layer non-looped model, and significantly better than a k-layer model. This is further corroborated by theoretical results showing that many such reasoning problems can be solved via iterative algorithms, and thus, can be solved effectively using looped models with nearly optimal depth. Perhaps surprisingly, these benefits also translate to practical settings of language modeling - on many downstream rea- soning tasks, language model with k-layers looped L times can be competitive to, not better than, a L--layer language model. In fact, our empirical analysis reveals intriguing phenomenon: looped and non-looped models exhibit scaling behavior that depends on their effective depth, akin to the inference time seal- ing of chain-of-thought (CoT) reasoning. We further elucidate the connection CoT reasoning by proving that looped models implicitly generate latent thoughts and can simulate steps of CoT with T loops. Inspired by these findings, we also present an interesting dichotomy between reasoning and memorization, and design a looping- based regularization that is effective on both fronts.
INTRODUCTION
Language models have shown a lot of promise in solving problems that require strong reasoning abil- ities like math, coding, common sense reasoning and logical puzzles (Brown et al.. 2020 Team et al.,
2023 This has sparked interest in developing techniques improve reasoning on harder problems WWei et al.] 2022b) and has inspired theoretical studies on how Transformers are able to perform reasoning Feng et al. 2024 Sanford et al. 2024a). Reasoning abilities are often emergent in larger language models (Wei et al. 2022a) this aligns with various scaling laws (Kaplan et al. 2020 Hoff- mann et al.) 2022 Allen-Zhu Li 2024) that show that the performance language models is very the model size (i.c., number of parameters) and much lesser other archi- tectural design choices. However, recent works have started to question this view. Ye et al. (2024) argue that scaling laws for reasoning parameter count at the same parameter count, deeper but shallower models are better. This is a de- viation from the conventional scaling law wisdom, but intuitively makes sense because reasoning problems often requires multi-step compositional thinking, and thus, depth can play a crucial role.
strongly dependent
more subtle, and depth is very important in addition to
In this work, we make stronger claim - while depth important, many reasoning problems do not necessarily require a lot of parameters. How does solve reasoning problems with large depth but few parameters? We argue that looped models are perfectly suited, where a function, parameterized with few parameters, is iteratively applied on the input. This leads
first important claim:
Claim 1: Many reasoning problems require depth but not necessarily parameters. That is, they can be solved via looped models

---

Published as a conference paper at ICLR 2025
Looped models have been studied in the literature for parameter efficiency Lan et al. 2020 adaptive compute (Dehghani et al., 2018), equilibrium models (Bai et al. 2019) and for in-context learning (Yang et al. [023] Gatmiry et al. 2024a). In this work, we initiate the study of looped models in the context of reasoning. Admittedly, reasoning is not very well-defined and can be of various forms (Sun al., 2023). Acknowledging this hurdle,
focus a non-exhaustive list of problems that intuitively require reasoning and that are inspired by reasoning benchmarks. Throughout the paper, we use the notation (k @ L) to denote a k-layer model looped L times (precise definition in Section 2) which has the same number of parameters as a (k @ 1) model and same flops as a (A:L @ 1) non-looped model (see Figure 5). As first step towards connecting looped models and reasoning. we empirically evaluate looped models on several simple reasoning tasks in the literature (e.g. Section 2) Perhaps surprisingly, we find that a (A: @ L) looped models does almost as well as, if not better than, a non-looped model (I: 60 1) that has the same effective depth but L times more parameters on these reasoning tasks. The looped model is also significantly better than (k @ 1) model which has the same number of parameters. Our theoretical results on the expressiveness of looped models in representing iterative algorithms with short description further corroborate these empirical findings and provide strong support for our claim. This naturally raises an important question: do looped models benefit language modeling in a similar manner?
Claim 2: For language modeling, looped models have inductive bias towards good reasoning despite having worse perplexity and memorization to an iso-flop non-looped model
For the above claim, we again train a (A 60 I.) looped model on causal language modeling and com- pare it the iso-param (A @ 1) and iso-flop (k& @ 1) non-looped baselines. While the looped model improves over the iso-param baseline, perhaps unsurprisingly, it ends up with worse perplexity than iso-flop baseline, since perplexity depends strongly on number of parameters. However, the down- stream evaluations reveal an intriguing trend: looped models have tendency to improve tasks that require reasoning lot more than memorization tasks. Specifically, the looped model has reasoning performance much closer to the iso-flop baseline, sometimes even exceeding it despite having L times fewer parameters and worse perplexity. This contrasting behavior between the pretraining and downstream metrics has been a subject of study lately (Saunshi et al. 2022 Liu et al. 2023) and is attributed to the inductive biases introduced due to different architectures and training algorithms. Our empirical analysis also uncovers an interesting phenomenon: accuracy downstream tasks scale as logarithm of the effective depth. In particular, more loops enhances performance, and the relative benefit of loops is higher for tasks that require more reasoning. This is conceptually similar to inference time scaling discovered for CoT, but with looping as a central component. To further elucidate this interesting relationship of looped models with CoT, we present the following claim.
Claim 3: Looped models generate latent thoughts and can, in theory, simulate CoT reasoning
Note that COT reasoning gives the model more time and compute by generating multiple thought tokens before the answer, and it has powered the recent paradigm of inference-time scaling for "thinking" models like O1 and DeepSeek's 's RI (Guo et al. 2025). We make an observation about CoT reasoning is essentially a looped model that generates thought token in each iteration. However, looped models seem to be much more powerful, since they can generate multiple latent thoughts in each iteration. We translate this intuition into theoretical result about how looped mod- els can simulate CoT reasoning. Motivated by these findings, we propose regularization scheme that aims to tap into the inductive bias of looped models towards reasoning.
Claim 4: Looping- regularization can leverage this inductive bias towards better reasoning
With the backdrop of these claims, we concretely present the contributions of the paper below:
In this paper we study looped models - multilayer models with weight sharing and their role in reasoning. In particular, we compare a k-layer model looped times, denoted by (k @ L). with an iso-param (k & 1) non-looped model with & layers and an iso-flop (kL @ 1) model with kL layers and L times more parameters.
We conduct experiments on synthetic reasoning tasks like addition, p-hop induction and GSM- style math word problems in Section [ For these tasks, we surprisingly find that iso-flop looped models, despite having way fewer parameters, can nearly match outperform non-looped
2

---

Published as a conference paper at ICLR 2025
model. Supporting these experiments, in Section 5 we present theoretical results for why looped
models can solve such problems with almost optimal depth.
In Section []] we train looped models causal language modeling at IB parameter scale. Here, we show that looped models have an inductive bias towards doing well on reasoning benchmarks, despite having much worse perplexity. This finding is novel, since most prior work on looping focused more on perplexity metrics rather than downstream reasoning tasks. We validate this in- ductive bias by visualizing the perplexity vs downstream performance plots training process. Additionally, we show that looped models demonstrate good scaling behavior on various bench- marks as the number of loops are increased, akin to CoT reasoning. Finally, we show that looped models, along with scratchpad, can simulate chain-of- tfthouhh reasoning.
Inspired by this inductive bias, in Section 4 we propose a regularization that encourages layers to be more similar to each other. We find that training with such regularization inherits the inductive bias of looped models towards reasoning without affecting perplexity.
LOOPED MODELS ON SIMPLE REASONING TASKS
We first explore our hypothesis of looped models helping reasoning tasks on of tasks con- structed in a procedural manner. The illustrative reasoning tasks we consider are: n-ary addition, p-hop induction head that tests the model's ability to track back for p steps, and i-GSM which con- sists of synthetically constructed grade-school math problems. While these obviously do not cover the whole spectrum of reasoning problems, they provide useful insights into looped models and provide a basis for the theoretical results in Section
Looped models. While many variants of looped model have been proposed (Lan et al. 2020 De-
hghann al. 2018 Giannou et al., 2023: Yang et al., 2023 Mohtashami et 2023 use the vanilla version for simplicity. For any sequence-to-sequenee I to be the function that is f looped I, times. In general, the looping mechanism is independent of architectural choices for f. For the rest of the paper, we typically use f to denote k-layer Trans- former backbone of model. Thus, f looped L times is the same as kL layer model with weight sharing between all L blocks of & consecutive layers. We denote such looped models with the nota- tion (R+8 L). Please refer to Figure Figure5)for for a succinct illustration of the looping mechanism. Section 5.1 provides a more formal definition of looped transformers that is used for theoretical analysis.
function f. we denote f(
fof.....
2.1 EXPERIMENTS WITH SIMPLE REASONING PROBLEMS
n-ary addition. We consider the problem of adding numbers with 3 digits each. Addition is pop- ular in the literature to study aspects of reasoning with Transformers, such as use of scratchpad (Nye) al. 202 1) chain of thought reasoning (Lee et al.) 2024: Li al., (024) and length generalization (Cho et al. 2024). One reason for its popularity is that addition can have algorithmic solutions, which is a feature of many reasoning problems. For our experiments, we train on a uniform mixture on numbers of operands n E {2, 4, 8, 16, 32} and sample each 3-digit operand uniformly at random between [0, 999). We train all models directly on the input-output pair, without any chain-of-thought steps. For instance, for n = 4, we have Input: "315 + 120 + 045 + 824 ="; Output "1304".
We train a standard Transformer-based baseline (12 6 1) model with 12 layers. Please refer to Appendix B.I for details on the training setup. We also train (k & 12/k) looped model and an iso-param (A & 1) baseline models for comparison, and vary k E {2, 3, 4, 6). All trained models are finally evaluated separately on each of E {8, 16, 24, 32) measure accuracy on increasingly difficult problems. Results are presented in Table I We find that, while the shallower baselines (A & 1) degrade with lower it, the looped model (A @ 12/k) performs very well, and nearly matches the iso-flop (12 @ 1) baseline. In fact, even 1-layer network looped 12 times is able to solve this, despite using merely 1 1/12th of the parameters of the baseline. This suggests the addition problem primarily requires depth, but not necessarily more parameters.
p-hop induction. The p-hop problem is a synthetic induction task studied in Sanford et al. 2024b), who were inspired by the analysis of induction heads from Elhage et al. (2021). Specif- ically, given a sequence of letters v = (v) ...... v.) from an alphabet Σ, an induction head tries to find the penultimate occurrence of ", (from left right) in the sequence and output the character
3

---

Published as a conference paper at ICLR 2025
Table 1: Accuracy of looped and non-looped models on the addition problem (left) and p-hop induc-
tion (middle), and i-GSM (right) as described in Section 2 Left. For addition, we report accuracies for different number of operands (n). For all budgets, a (k @ 12/k) looped model is significantly better than the iso-param (A: @ 1) model and also nearly as good as the non-looped iso-flop (12 & 1) baseline model. Middle. The findings are very similar for the p-hop problem for different values of Note that a random guessing baseline would get at least 25% accuracy (since only 4 choices for an- swer). This suggests that depth via looping and small number of parameters is very effective for these problems. Right. Accuracy on i-GSM task from Section 2 (k @ L) looped model is significantly better than the iso-param (A:1 1) model and performs as well as non-looped iso-flop (L 1) model.
Addition of n numbers
p-uup with
i-GSM
Params/ / /LOPs
Accuracy
Params/ FLOPs
n
24
n
52
Params / HL.Ps
p
p - 32
Base (8 1)
Sx
73,2
256
256
layer mode)
Base (12 1)
12x 12x
100.0
100.0
Base (6 @ 1)
6x
99.9
99.6
Base (1
Loop (1 ()))
1x/ 1x
1x/2x
52.3
24.5
layer model
layer model
Loop (180)
69.9
Base @1 1)
1x/
0.1
0.0
Base (1 @
Ix/Ix
48.9
29.0
Loop (1 8
1x/Sx 8x
73.2
Loop (1 6: 12)
12x
99.9
99.6
Loop (1 6)
1x/6x
99.9
99.5
layer modell
layer
layer model
Base (2
2x. 2x
54.0
Rase (2:01) 1)
/ 2x
49.3
38.8
Hase (2 1)
/x
68.8
59.4
Loop 0 2) Loog (2 10 4)
2x 4x
66.9
Loop (2 8: 6)
2x /
99.7
99.5
Loop (2 2:33
2x/6x
99.9
99.8
2x 8s.
73.6
layer model
layer model
layer model
Basc (3
69.2
60.7
00
73.0
Base (4&11 Loop (1
4x /4x
71.3
Loop (3 8 4)
/x/12x 12x
97.0
96.6
Loop (3:88 2)
/x
99.9
99.5
4x/So Sx.
71.6
immediately succeeding it. The p-hop problem generalizes this idea to sequentially hop p times. Intuitively, the p-hop problem tests model's 's ability to recursively backtrack and retrieve the answer for a given query. This is reminiscent of the reasoning abilities required to solve reading compre- hension kind of problems. We present the formal definition of the p-hop problem in Definition B. We perform experiments with looped and non-looped models on the p-hop problem, with alphahet size set 4 and sequences of length 256, and vary p between 16 and 32 to control the problem difficulty. Our observations are presented in Table [ Similarly to our findings on the addition task, reasonably deep looped models perform as well as the baseline using much fewer parameters.
i-GSM (Synthetic Grade School Math Problems). Inspired by et al. (2024), we built our own version of grade-school level math word problems. While we follow many of the design guidelines Ye et al. 2024), we make few simplifying changes. We generate the math problem as a DAG of arithmetic computations modulo 7, and restrict the depth of the graph to 4. For simplicity. we retain the problems in the symbolic form and do not map them to English (e.g.. see Table 8) We train models of depths I, 2, 4 and and compare them with different looped variants in Table] The answer is computed modulo 7. Hence, random guessing baseline would get at least 14% accuracy. Appendix B.I has more details. Remarkably, we again observe that a depth model looped I, times matches or outperforms a depth KL model and far outperforms a depth & non-looped model. This suggests that even a more complex and realistic looking reasoning problem does not need too many parameters and can benefit from depth via looping.
LANGUAGE MODELING WITH LOOPED MODELS
In this section, we pretrain and evaluate looped models for causal language models. We train models on 250B tokens of the Pile dataset (Gao et al. 2020) and use a 24-layer 1B parameter model for most experiments, motivated by the setting in Tay et al. (2022) (refer to Appendix B.2 for more details).
3.1 EXPERIMENTS WITH LANGUAGE MODELING
For causal language modeling, we pretrain various looped models on the standard GPT2-style next token prediction objective (Radford et al. 2019). We train models with different parameter budgets to make sure that the findings are robust. We remind the reader that the notation (A & L) corresponds to k layer model looped L times. For each setting, we compare 3 models: (a) (24 @ 1): 24-layer 1B model, (b) (k @ 1): k-layer model with the same configuration as the 24- model for other dimensions, (c) (k @ 24/k): k-layer model looped 2 times to match the parameter count of (b)
Similar to Ye et al. 2024 the simplified setting still allows for over 7 billion unique solution templates.
4

---

Published as a conference paper at ICLR 2025
Table 2: Downstream evaluations for language models trained on the Pile dataset. Comparisons in- clude a 24-layer IB 1B-parameter baseline model, iso-flop looped models (k@24/k) /k) for various param- eter budgets k, and the corresponding iso-param baselines (k@1). Downstream evaluations are aver- aged over tasks within 4 task groups. We also include the % Gap metric for each k to measure the gap between the iso-param and iso-flop baselines that is covered by the looped model (see Equation ) Overall the looped models are worse on perplexity and closed book QA (memorization benchmarks), but the % Gap is much higher for task groups that require reasoning (open book QA, math word problems). In fact for reasoning primitives, which are purely testing for reasoning skills, the looped models are much better than the 1B baseline for all k; despite having 24/hx fewer parameters.
Puras FLOPs
Perplexity (1)
Closed Book QA (1)
Open Book QA (f)
Math Word
All Tasks Average (?)
Reasoning Primitives (1)
Problems (1)
Hayers
Baseline
24k 24x
7.40
11.2
33.9
29.3
26.0
47.5
13 layers
12x/
26.9
21.8
Loop (12:82 8 2)
12x /24x
7,90
9.3
30.8
34,3
26.5
51.2
Gap
56
Middle Loop
12x 24x
7.81
11.0
32.3
28.3
25.0
56.5
8 1,4.1)
%: Gap
46 %
94 56
28 9
62 %
95 56
176 %
8 layers
Base (8 Loop 3)
8x Xx
8.35
22.7
17.1
16.1
33.0
8x 24x
8.19
8.5
30.8
28.4
23.9
55.3
41
44 9
72
92
78 56
56
layers
Base 16 0:
6x 6x
9:25
4.0
19.3
17.7
14.6
24.1
Loop (6 00-00 4)
6x 24x
8.42
8.2
28.7
29.8
23.7
56.1
Gop
58 96
64 5
104 56
80 56
136 %
layers.
Base (4
4x/4k
10:12
1.8
13.8
9.7
9.0
19.4
Loop (4:00 6)
/xxx
8.79
6.7
26.2
24.8
20.4
56.9
% Gap
48 %
52 %
61 %
27 %
67 %
163 %
and match the effective depth/FLOPs, of (a). We run experiments for & E {4, 6, 8, 12] to ensure that the findings are robust. After pretraining on Pile, we evaluate the models on validation perplexity and on downstream benchmarks using k-shot evaluations. Results are summarized in Table 2
Evaluation metrics. We evaluate the models on perplexity metric after training is completed. Since there is growing evidence that perplexity, although very useful for training, is narrow measure of model quality, we also track more holistic downstream evaluations (Liang et al., 2023). Thus, we evaluate the model on 4 important slices: closed book QA, open book QA, math word problems and reasoning primitives. These comprise of 19 different tasks in total. We defer details for the evaluation benchmarks in Appendix B.3
For each task group G' from above, in Table 2 we report the average accuracy for that task group, denoted by Avgg. Furthermore, for each layer budget k; we report the % gap between the iso-param and -so-flop models that is covered by the looped model. More specifically
% Gap = Avga(k @ 24 24/k) Avga (k: 09 1)
(I)
Avg. (24 @ 1) Avga @ 1)
This measures how effectively looping can bridge the gap between iso-param and iso-flops baselines. Implicitly, it measures how different task groups behave when a model only has a few parameters but is given depth through looping. % Gap being closer to 0% means that providing depth via looping does not benefit the task group, and number of parameters is the most important factor. On the other hand, % Gap closer 100% means parameter count matters much less for the task group, and that depth via looping is more essential.
Perplexity results. Firstly we notice that all (k @ 24/k) looped models have better perplexity compared to the iso-param (k @ 1) baseline, but worse perplexity compared to the non-looped 24- layer baseline. The looped models only covers up roughly 34 50% of the perplexity gap between the -so-param and iso-floo flop baselines for various values of k. This perplexity gap is not too surprising since the looped model has 24/k times fewer parameters, and thus, lower capacity than the 24-le baseline. This was also been observed in prior works (Lan et al. 2020) Mohtashami et al. 2023)
5

---

Published as a conference paper at ICLR 2025
Beseline
56.
Looged (12
52.5
Closed Book
2.-
Open Book
-0-
Looged 21
Baseline (12 e 1)
Math Word Probilerrs
20
30
501
415
45.1
Baseline (12 9 21
4/3
Baseline (12 11
45-
Looped (12
36.5
21s 2:2
Perplexity
Log Perplexity
Perplexity
Perplexity
Figure I: Downstream evaluation for various task groups on the x-axis, vs validation log perplexity on the y-axis (reversed), training proceeds. We compare 12-layer baseline model (12 @ 1) with a looped model (12 @ 2). For closed book QA (memorization) tasks looping has very similar trends baseline. For open book QA tasks and math word problems, looping has much better downstream performance at an equivalent log perplexity.
and is the primary reason looped models have been ignored for language modeling. However, as we shall see shortly, the downstream metrics paint a more interesting and favorable picture.
Results on QA tasks. We first consider closed book and open book QA categories in Table 2 Closed book QA tasks are primarily testing the model's memorization abilities. Open book QA on the other hand tests the model's ability to infer the answer from the additional context that is provided. Thus, intuitively, open book QA tasks require more reasoning. Firstly we notice that the % Gap for closed book QA (memorization) is very similar to % Gap for perplexity. The % Gap for open book QA, on the other hand, is much higher for all parameter budgets. This suggests that looped models relatively improve contextual QA much more than memorization based QA.
Math problems and reasoning primitves. We also present the % Gap for the math word problem in Table 2 Surprisingly, we find that (k 24/k) looped model can almost match the baseline (24 & 1) model for & 2 6, despite having & times fewer parameters. In fact, the 12 layer model looped twice even signficantly better (34.3) than the 24 layer baseline (29.3), despite having 50% of parameters and worse perplexity; suggesting that looping disproportionately improves mathematical reasoning.
To better understand the effect on reasoning, we direct the readers attention to the evaluations for reasoning primitives in Table The results quite remarkable: (k @ 24/k:) looped models are better than the iso-flop baseline (24 @ 1) at reasoning primitives, for all values of k. This is priori very surprising, since these are synthetic generated tasks and have nothing to do with the pretraining data or the model architecture. Thus, solving these tasks necessarily requires reasoning from context, and memorization abilities will not help here. These results clearly suggest that looped models have a bias towards improving reasoning, despite having worse perplexity and memorization.
3.2 INDUCTIVE BIAS TOWARDS REASONING
In this section, we formalize the inductive bias by plotting the perplexity vs downstream metric iso-plots, as introduced in Saunshi et al. (002)). Section showed that looped models have higher than expected performance on reasoming problems. However, since looped models are worse on perplexity, it is hard to make a direct comparison between various models. One way to bring parity between models is to look at their downstream performances at the same validation pretraining loss (Liu et al. 2023). [Sausshi et al. (2024) proposed plotting pretraining loss vs downstream metrics as training proceeds, as a way to study the inductive bias of various methods. For each model, we evaluate the log perplexity and downstream metrics at every 20k steps, starting from 120k steps. We plot these values in a scatter plot and fit a linear function with log perplexity and the corresponding downstream metric being input and output respectively. Please refer to Figure T for two sets of isoplots. For all values of k, we observe the following:
The isoplots for (k & L) looped model and (k& 1) baseline are very aligned for closed book QA tasks (if extrapolated). This suggests that log perplexity is a very strong indicator of downstream performance on memorization based tasks.
For open book QA and math word problems, the isoplot line for the looped model is always higher than the baseline model. This suggests that at the same log perplexity, looped models will tend have higher evaluation on these tasks that require more reasoning.

---

Published as a conference paper at ICLR 2025
Closed Mose
Open (Magglhtae
Reasoning Primitives
1.2)
Baseline a
Baseline ID 1)
Baseline s0 1)
Looped
Accuracc 25
Looped (4 014)
50
20
24
Efeccivve Model Depth (2)
Effective Depth IDI
Effective Model Depth (0)
Figure 2: Scaling behavior for various task groups as the effective depth increases (also see Fig- ure 7). The blue curve shows how performance scales for (4 @ D/4) looped models as the number
of loops increases, without increasing parameters. The orange curve visualizes the scaling behavior of (D & 1) which increases the depth by adding fresh parameters. For reasoning primitives, the looped model scales as well, or even better, than the baseline despite having D/4 4 fewer parameters.
For reasoning primitives, there is stark difference between looped and baseline models. The looped model seems to have good performance at most points training.
This suggests a strong inductive bias of looped models towards improving reasoning. We discuss connections to gradual stacking (Appendix B.4) approaches that have similar inductive bias.
3.3 SCALING BEHAVIOR LOOPING
In this section, discuss an intriguing scaling behavior of looping, specifically the impact of the number of loops on various evaluation metrics. In particular, we are interested in scaling behavior of: (a) accuracy as function of number of loops and (b) comparison of looped and non-looped baseline of the same effective depth. To this end, we pretrain various looped language models of the form (4 s0 L), i.e., 4-layer model looped times, for T. E {1, 3, 6, 9, 12). To enable iso-FL.OPs comparison, we also train baseline models of the form (4L 1). which has L times more parameters. For each task group, we plot the average accuracy as a function of the effective depth, i.c. D 4L. From the results presented in Figure we observe the following.
1. In both cases, we find that the accuracies for all task groups continue to increase with more loops/depth, although, unsurprisingly, the returns are diminishing with depth for both looped and non-looped models. Interestingly, for both looped and non-looped models, we found that one can fit a simple scaling law of the following form:
Acc = log (D) + 3,
where D is the effective depth and a measures the impact of depth on downstream performance.
2. Furthermore, for each task group, we compute (Koop (loop/Obasse /Obase to assess the relative impact of "depth via looping" compared to "depth via additional parameters". We find that more loops continue to help, and the relative benefit of loops is higher for reasoning tasks like open book QA and math problems. Remarkably, the impact of loops is even higher (1.19x) than impact of depth for reasoning primitives, which further consolidates the benefit of looped models for reasoning.
Latent thoughts and connections to CoT reasoning. We end this discussion with an important question: why should we expect this interesting scaling behavior of looped models? We argue that looped models have a strong connection to chain-of-though (CoT) reasoning. Such a connection is insightful because recent works on thinking have shown CoT can demonstrate inference-ting scaling behavior, where the accuracy on reasoning datasets continues to improve as the length of the models chain of thought increases; i.e., generating more thoughts leads to better reasoning.
To establish this connection, we make simple observation about CoT reasoning - it is essentially a looped model that generates a single thought token in each iteration. However, looped models can be more powerful since they can generate multiple "latent thoughts" in each iteration. This can be visualized in Figure g We further translate this intuition into a theoretical result (see Section 5.4 about how looped models can in fact simulate CoT reasoning. Given this connection to CoT rea- soning, it is believable that looping can scale well for harder reasoning. We believe that leveraging looping explicitly for inference-time scaling is very promising future direction.
7

---

Published as a conference paper at ICLR 2025
Table 3: Results for the 24-layer IB model with and without the regularization introduced in
Section 4 We try various block sizes k motivated by the looped model settings from Table 2 Overall, regularization helps retain the inductive bias towards reasoning. with notable improvements on math word problems and reasoning primitives, without almost neutral perplexity.
Perplexity (1)
Closed Book QA (?)
Open Book QA (t)
Math Word Problems (1)
All Tasks Average (f)
Reasoning Primitives (1)
Raseline
7.40
11.2
33.9
29.3
26.0
47.5
Regularized
Auy
7.41
11.2
34,8
31.6
27.2
42.5
Regularized (k Aug
-00)
7.38
12.5
36.2
36.4
30.0
57.2
Regularized Aag 10)
7.40
12.0
35.8
31.0
27.5
55.8
Regularized 8, Aug
10)
7.43
11.3
34.4
32.8
27.6
56.3
Regularized 12,
10)
7.51
10.1
34.1
32.3
27.0
50.7
4 LOOPING-INSPIRED REGULARIZATION
In the previous section, we observed the looped models can improve reasoning with worse perplex- ity. Can we leverage this observation to improve reasoning without affecting perplexity? Here, we propose simple approach: regularize the weights of the model to encourage them to be close to looped model. This could have two advantages, (a) the model still has free parameters to improve perplexity, (b) the closeness to looped model can inherit the desirable inductive bias. In particular, an L-layer model denoted as fo fi.... oftk/-1- where each fi a block of k layers, we add regularization term that makes all the weights of the block close fi+1 in terms of cosine similarity. For parameter group G (e.g. first feed-forward layer, or query matrix in Transformer), use ((0) 0(0 "6 to denotes the weights in all layers. Then the regularization term is
(L.--))
$
2-2k-1
glik-j) g (((4DR4t))
Ro(k)
Cosine
(3)
L
2=0
The final loss function a sum of the standard cross-entropy loss and the regularization term aver- aged over all groups, multiplied by a scalar hyperparameter. Let G denote the set of all parameter groups; G = [Attn-Q, Attn-K,...,,FFE FFN-W2]
L- L + 1r
Σ (R)(A) GEO
In the above formulation, Aveg = would recover standard training and Aveg + oo would converge fully looped model. Intermediate values of Aveg will lead to "approximately" looped models. For instance, to emulate the (4 @ 6) looped model setting, we use pick k = 4, L = 24 and a large regularization strength like Are 10. All other hyperparameters are kept the same as baseline training for fair comparison. We tried options other than cosine similarity, like E norm, to bring the weights closer but found that cosine was more robust and effective.
Cosine similarities. Firstly, check if the regularization had the right effect by measuring the cosine similarities between the successive blocks of & layers at the end training. We, indeed, find that for all parameter groups, the cosine similarity around 0.98 higher (see Figure 3)
Inductive bias. To confirm the inductive bias of the regularizer, we visualize the log perplexity vs downstream isoplots for Areg = 10 and baseline models in Figure 6 While the plots are similar for closed book QA, a strong inductive bias shows up for open book QA and reasoning problems. Crucially, the regularized model does well on reasoning without hurting perplexity (see Table 3).
THEORETICAL ANALYSIS FOR LOOPED MODELS
In this section, we present theoretical results to understand the phenomenon from the previous sections why can looped model with few parameters match an iso-flops non-looped baseline reasoning problems? While a complete theory is challenging, since "reasoning" is a very broad concept. the goal is to provide some intuition and formalization for the expressive power of looped models. First, we show that looped Transformers can effectively solve group composition (a generalization of the addition problem). Then we show very general result how non-looped
8

---

Published as a conference paper at ICLR 2025
model with very few distinct layers can be simulated by looped model with small blowup in model size. This result is then used to solve the p-hop problem using a one-layer looped transformer. Our construction for group composition and p-hop problem nearly optimal in terms of depth and much more efficient in terms of parameters compared to non-looped models.
5.1 PRELIMINARIES AND NOTATIONS
We first define the standard transformer architecture. Throughout the paper we will fix the dimension of the embedding to be d E N+, the vocabulary to V and maximum sequence length to 7mmx-- We will use id to denote the identity mapping. Here, refer to Appendix C.I for detailed notations. We use FF and MHA to denote the feed-forward and attention layers respectively, and OμHA and OFF denote the parameters in these layers.
describe the high-level notation. Please
Definition 5.1 (Transformer Block). Given number of layers E N+ and parameter θTB = (1)) 1=0 L-layer transformer block TBθTn (RR)" (R for any E N+ is defined MHA
TB
A (id + FF #F
()( o (id + MHA (L- ... (id + FF ((0)) (id + MHA (())
)
(5)
MHA
We also denote EMBED and OUTPUT to be the input embedding and output softmax layers respec-
tively. Please refer to Appendix C.T for precise definitions. Finally, we define the entire transformer model that maps sequence of tokens a distribution over tokens: pø Un y A 1
OUTPUTTUTOOO SourPUr o TB8T EMBED
(6)
where 0 = (OTB, QTE, OPE, OOUTPUT) denote all the transformer parameter. In particular, we use TFe TFe( ) - arg m v Pe (elee..... Un) to denote the deterministic version of the trans- former model. We now define looped Transformer model that also subsumes non-looped model.
Definition 5.2 (( T) Looped Transformer). Given the number of loops T T N+, parameters 0 (OTB, OTE, OPE, (OUTPUT), where OTF (OHHA 0 0-1 we define a (L & T)) looped Transformer po,r OUTPUT COUTPTT (TBoTa) EMBED 0oe
5.2 GROUP COMPOSITION PROBLEM
We consider the problem of composing n elements from a group, and prove that a standard 1-layer transformer looped (log(n)) times can solve this problem. This is a generalization of the modular addition problem and has a long history (see Appendix C.2.2). Recently, Liu et al. (2022) show that transformers with log, n depth can compute composition over group elements, regardless of whether the group is solvable or not. However, the construction in Liu et al. (022) uses different attention parameter for each layer. Here, we provide a more parameter efficient construction where solve this problem by looping one-layer transformer log(n) times. This result matches the depth upper bound shown for non-looped models by Liu et al.] (2022). and is very close to their super constant lower bound. Thus looped models can solve the problem with best known depth.
N+, there exists n elements from G with a 1-layer transformer block, T = flog, n loops, G U (#} being the vocabulary. d = 3 [loga |G|l + [log, n + 11) embedding size, OFF = |G|2 + 6 log, |G|] hidden dimension in MLP, dATTN = [log, n hid- den attention dimension, and 2 attention heads. More specifically, for any I1..... E G.
Theorem 5.1. For any finite group G and every
constant-precision looped
transformer F computing the composition
TF (Fo(#191.. 9t,.99,,.................... = o .... In-
5.3 LOOPED MODELS CAN SIMULATE NON-LOOPED MODELS
Our second theoretical result shows that non-looped transformer with repeated layers can be sim- ulated by a looped transformer with fewer parameters and same depth.
Theorem 5.2. For any transformer Pa with L layers, embedding size, dFF hidden dimension for ML.. H attention heads with dATTN hidden dimension, most R distinct transformer layers and

---

Published as a conference paper at ICLR 2025
bounded activation value, there is looped transformer pe L working on the same vocabulary V plus a dummy token #, which loops a 1-layer transformer block for L times, with d R + embedding size, RhFF + ((L) hidden dimension for MLP, R.H attention heads with dATTN hidden dimension, such that for any string 01, Un E V, po(u1,-... ;.) = po por. (L pa).
We can also use the above theorem to convert the OOlog2 n) depth transformer that simulates group composition into a 1-layer looped transformer, although less parameter efficient than the result from the previous section. Please refer to Appendix [C.2 for the full proof.
Theory for p-hop. The experiments on p-hop induction from Section 2.1 surprisingly show that a small model looped multiple times can solve it very effectively. We establish theoretical basis for this finding. More specifically, we show that a constant layer transformer with log(p) loops suffices to solve p-hop induction problem. This result, in fact, matches the lower bound for layers required for non-looped models proved in Sanford et al. (2024b). The result follows from Theorem [ and the construction for non-looped models from (aaford et al. (2024b) (restated in Theorem C.
Corollary 5.3. p-hop problem (Definition B.I can be solved by looping a I-layer transformer [loge p] + 2 times, which has Ollogn) bits of precision, d = dFF dATTN = O(1) embedding size, hidden dimension for MI.P and attention, and at most attention heads.
5.4 LOOPED MODELS CAN SIMULATE COT REASONING
we discussed how looped models can be viewed as generating multiple latent thoughts in each iteration. The following theorem shows that one can use looped transformer with loops simulate steps CoT of another transformer with similar sizes.
In Section 3.3
Theorem 5.4 (Looped transformer simulates CoT), For any L-layer non-looped transformer TF. with fixed input length n and the number of CoT steps m, there exists looped transformer TFe with + (1) layers, (noog(n + m)). more embedding dimension and constantly many more attention heads, such that for any input v (u:)² the output of non-looped transformer after m steps of COT, TF2 (v), is the same as that of the looped transformer on input x concatenated by m dummy tokens with m loops, TF (v, )
We sketch the high-level idea behind Theorem 5.4 the full proof can be found in Appendix
(Masking all- all-but-one) We maintain counter t (Lemma C.10) at each position for the current number of loop and for the ith position, we will only "update" the embedding if i-n 2 and reset the input the looped transformer layer otherwise. This can be implemented by MLP. So similar to CoT, the first n + i - 1 embedding won't 't't be changed in the ith loop.
the embedding
(Shift by one) We can use attention to obtain the output of the current loop at the previous position and use that as the input of the next loop at current position. (Lemma C.9)
(Token decoding) We show that we can MLP with ReLU activation to simulate the encoding
and decoding process of COT. (Lemma C.7
CONCLUSIONS, LIMITATIONS AND FUTURE WORK
This work explores a new direction of "looped models for reasoning". Not only are looped models able to solve many reasoning problems with very fewer parameters, they also have an inductive bias towards disproportionately improving the reasoning performance of language models, over memorization. The theoretical results on the expressivity of looped models start to provide some hints into their deptth-opiimality. While we test looped models on a subset of reasoning problems, a natural question is whether the results hold for many other forms of reasoning (e.g. multimodal and common-sense reasoning). particular, succinct formalization of reasoning problems itself is interesting future direction. Furthermore, the inductive bias towards improved reasoning performance at the same perplexity is very intriguing and deserves further exploration. We find the scaling behavior of looped models very fascinating, and the connections to latent thoughts and CoT reasoning start to provide hints into this behavior. We hope this inspires future exploration on using looped models for more efficient inference-time scaling to aid with deeper reasoning.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws. arXiv preprint M,XY::2404.05405,, 2024.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances neural information processing systems, 2019.
Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Extrapolation with- out overthinking. Advances in Neural Information Processing Systems, 2022.
David A. Barrington. Bounded-width polynomial-sie branching programs recognize exactly those languages nc. pp. 1-5, 1986.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 2020.
Xingwu Chen and Difan Zou. What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404. 01601, 2024.
Hanseul Cho, Jaeyoung Cha, Pranjal Awasthi, Srinadh Bhojanapalli, Anupam Gupta, and Chulhee Yun. Position coupling: Leveraging task structure for improved length generalization of trans- formers. arXiv preprint wrXiv:2405.20611 2024.
Eunsol Choi, He He, Mohit lyyer. Mark Yatskar, Wen-tau Yih, Yejin Choi, Percy Liang, and Luke Zettlemoyer. QuAC: Question answering in context. In Proceedings of the 2018 Conference Empirical Methods Natural Language Processing, 2018.
Jonathan H. Clark, Eunsol Choi, Michael Collins, Dan Garrette, Tom Kwiatkowski, Vitaly Nikolaev, and Jennimaria Palomaki. TyDi QA: A benchmark for information-ssekingg question answering in typologically diverse languages. Transactions of the Association for Computational Linguistics, 2020.
Artur Back de Luca and Kimon Fountoulakis. Simulation of graph algorithms with looped trans- formers. arXiv preprint 17Xiv:2402.01107, 2024.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representutions, 2018.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proc. of NAACL, 2019.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, pp- 5793-5831, PMLR, 2022.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1:1, 2021.
Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. arXiv preprint arXiv:2305. 15408, 2023.
Guhao Feng. Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems, 36, 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0mm 2020.

---

Published as a conference paper at ICLR 2025
Yihang Gao, Chuanyang Zheng, Enze Xie, Han Shi, Tianyang Hu, Yu Li, Michael K Ng. Zhenguo Li, and Zhaoqiang Liu. On the expressive power of a variant of the looped transformer. arXiv preprint 4rXiv:2402.13557, 2024.
Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn implement multi-step gradient descent for in-context learning? In Forty-first International Conference on Machine Learning. 2024a.
Khashayar Gatmiry, Nikunj Saunshi, Sashank J Reddi, Stefanie Jegelka, and Sanjiv Kumar. On the role of depth and looping for in-context learning with task diversity. arXiv preprint arXiv:2410.21696 2024b.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning. PMLR, 2023.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang. and Tieyan Liu. Efficient training of BERT by progressively stacking. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Daya Guo, Dejian Yang. Haowei Zhang. Junxiao Song. Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501. 12948, 2025.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meet- ing of the Association for Computational Linguistics (Volume I: Long Papers), 2017.
Jared Kaplan, 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.0883 2020.
Rik Kenzes-Kedziorski, Subhro Roy, Aida Amini, Nate Kushman, and Hannaneh Hajishirzi. Mawps: A math word problem repository. In Proceedings of the 2016 conference of the north american chapter of the association for computational linguistics: human language technologies, 2016.
Kenneth Krohn and John Rhodes. Algebraic theory of machines. i. prime decomposition theorem for finite semigroups and machines. Transactions of the American Mathematical Society, 116, 1965.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wes Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 2019.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020.
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learn ing Representations, 2024.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In The Tweifth International Conference on Learning Represen- fations, 2024. RRK.htttps::/pperevvieee
net/ htt.:eeneeeeooomommm/a/a.rr
12

---

Published as a conference paper at ICLR 2025
Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. Transactions on Machine Learning Research, 2023.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprini arXiv:2210. 10749, 2022.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre- rre-training loss, better down- stream: Implicit bias matters for language models. In International Conference on Machine Learning. PMLR, 2023.
Zechun Liu, Changsheng Zhao, Forrest landola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong, Ernie Chang, Yangyang Shi, Raghuraman Krishnamoorthi, et al. Mobilellm: Optimizing sub-billion parameter language models for on-device use cases. arXiv preprint 242.:224402.14055 2024.
William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923, 2023a.
William Merrill and Ashish Sabharwal. The parallelism tradeoff:: Limitations of log- precision trans- formers. Transactions of the Association for Computational Linguistics, 11:531-5aa 1555,,55,,505 2023b.
Shen- Yun Miao, Chao-Chun Liang, and Keh- Yih Su. A diverse corpus for evaluating and developing english math word problem solvers. Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 975-984, 2020.
Amirkeivan Mohtashami, Matteo Pagliardini, and Martin Jaggi. Cotformer: More tokens with tention make up for less depth. arXiv preprint arXiv:2310. 10845, 2023.
Rodrigo Nogueira. Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2102.13013 2021.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez. Maarten Bosma, David Luan, your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2/12.00/1 14, 2021.
Show
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are nlp models really able to solve simple math word problems? In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Linguistics, 2021.
Jackson Petty, Sjoerd van Steenkiste, Ishita Dasgupta, Fei Sha, Dan Garrette, and Tal Linzen. The impact of depth and width transformer language model generalization. arXiv preprint arXiv:2310.19956, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Pranav Rajpurkar, Robin Jia and Percy Liang. Know what you don't 't know: Unanswerable questions for SQuAD. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), Melbourne, Australia, 2018
Sashank Reddi, Sobhan Miryoosefi, Stefani Karp, Shankar Krishnan, Satyen Kale, Seungyeon Kim, and Sanjiv Kumar. Efficient training of language models using few-shot learning. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Siva Reddy, Danqi Chen, and Christopher D. Manning. CoQA: conversational question answering challenge. Transactions of the Association for Computational Linguistics, 2019.
Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Halcrow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. arXiv preprint arXiv:2405. 185 12, 2024a.
13

---

Published as a conference paper at ICLR 2025
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and loga- rithmic depth. arXiv preprint arXiv:2402.09268, 2024b.
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. In Proceedings of the 39th International Conference on Machine Learning, 2022.
Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryooseñi, Sashank J. Reddi, and San- jiv Kumar. On the inductive bias of stacking towards improving reasoning. arXiv preprint arXiv:2409. 19044, 2024.
Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn algorithm? generalizing from easy to hard problems with recurrent networks. Advances in Neural Information Processing Systems, 2021.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604, PMLR, 2018.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. Journal of Machine Learning Research, 2018.
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. Transformers as rec- ognizers of formal languages: A survey on expressivity. arXiv preprint urXiv:231 1.00208, 2023.
Jiankai Sun, Chuanyang Zheng, Enze Xie, Zhengying Liu, Ruihang Chu, Jianing Qiu, Jiaqi Xu, Mingyu Ding, Hongyang Li, Mengzhe Geng, et al. A survey of reasoning with foundation models. arXiv preprint arXiv:2312.11500 2023.
Alon Talmor and Jonathan Berant. The web as a knowledge-base for answering complex questions. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long Papers), 2018.
Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung. Dara Bahri, Tal Schuster, Steven Zheng, et al. U12: Unifying language learning paradigms. In The Eleventh International Conference on Learning Representations, 2022.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- Cean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv: / 2023.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.0766 2022a.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 2022b.
Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. arXiv preprint arXiv:2311.12423 2023.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.1, grade-school math and the hidden reasoning process. arXiv preprint 2rX1v.2407.03311, 2024.
A RELATED WORK
Reasoning is recognized significant focus over the last few years. The synthetic reasoning problems we consider in this work have all been used in the prior works of Sanford et al.] (2024b); Ye et al. (2024); Sanford] (et al.] (2024a): Nogueira et al. (202T) to theoretically analyze the strengths and limitations of Transformers. There is also interest in the representation power of Transformers for computational
core ability for intelligent and robustly model and has thus seen
14

---

Published as a conference paper at ICLR 2025
problems (Liu et al., 2022 Strobl et al., 2023) and for chain-of-thought reasoning (Merril & Sabharwal 2023a Feng et al. 2023 Li et al. (004). The necessity of model depth for performance
has been remarked upon. for small models (Liu et al. 2024) and reasoning (Chen & Zou 2024 Ye et al. 2024 Petty et al. 2023). In this work, we make a finer observation that albeit larger depth is necessary, this can be achieved with limited parameter budget via looping.
Looping in transformer models has been studied since the works (Dehghani et al., 2018 Lan et al., 2020) where they showed the benefits of looping for supervised learning tasks and BERT pretrain-
ing respectively. Looping also appears in (Schwarzschild et al., 2021: Bansal et al. 2022) that
study the extrapolation properties of looping for certain algorithmie tasks. More recently Gian- nou al. )023); de Luca & Fountoulakis (2024) have studied the theoretical properties of looped decoder models and show that looping can simulate arbitrary Turing machines. In addition Yang al. (2023); Gao al. (2024); Gatmiry et al. (2024a b) study looped models
way to sim- ulate iterative algorithms for in-context learning. Recently, MMhhtasamii et al.) (2023) introduced Formorrer, which tries improve the perplexity of looped language models and (Liu et al. 2024) explore latency-efficiene parameter sharing for on-device LLMs. In contrast, our work focuses on the surprising inductive bias of looping to improve downstream reasoning tasks, and goes beyond algorithmic and in-context learning tasks.
Different training algorithms (e.g. gradient descent (Soudry et al. 2018) and architectural choices (e.g. attention (Edelman et al.) 2022)) have been shown to have certain implicit biases. There is increasing interest in such inductive biases during pretraining (Saunshi et al.) 2022] Liu et al., 2023 More recently, (Saunshi et al.] (202)) showed an inductive bias of stacking (Reddi et al. 2023) towards improving reasoning and hypothesize that a connection of stacking to looped models
could be responsible for this. Our results provide further verification for this hypothesis.
B EXPERIMENTS
B. SIMPLE REASONING SETUP DETAILS
n-ary addition. All experiments are run on standard Transformer architecture with input di- mension of 256, 8 heads and 1024 hidden dimension in the feed-forward layers. We train using Adafactor (Shazeer & Stern 2018) employing linear warmup coupled with a cosine decay sched- ule for the learning rate. All runs use batch size of 1024, learning rate of 0.005 and run for 200k steps. This corresponds to 200M examples, which is insignificant compared to the total possible ex- amples (> 10320). Thus memorization of the answer is not an issue. Since training is a bit noisy, for each setting, we run 3 different random seeds and pick the run with the maximum average accuracy. We pick maximum instead of average because we care about expressivity power of these models.
p-hop induction. We formally define the p-hop problem below:
Definition B.1 (p-hop, Sanford et al. (2 For finite alphabet Σ, define the map hopp Σn EU(L) as hn),,,)) (w.nd)(w.n) if find, (v, n) o 0 else 1, where
wH}) find,( i) find, (v, 1in i)) for p 2 2.
find; (v, i) max ( S -,,
For the p-hop problem we sample instances randomly while enforcing that there always exists p- hops present the input sequence. We do this by first picking the sequence of p-hops randomly and then shuffling them around in sequence with filler tokens to be filled by the remaining characters. After the shuffle, we sample the remaining characters to occur in place of the filler tokens while respecting the p-hop order. Our train set consists of 4M examples and our test and validation sets consist of 262k examples each. For all models we train on this dataset, the model dimension is 128, hidden dimension is 512 and 8 attention heads are used. Rotary positional encodings well. We train using Adafactor for 200,000 steps with a batch size of 256 using a base learning rate of 10-3 and use linear warmup coupled with cosine decay schedule for the learning rate.
used as
15

---

Published as a conference paper at ICLR 2025
8:915
35045
0.881
0.333
0.740
0.9033
C.0000
2.0074
1,0000
5.5011
03213
0.379
0.740
0.915
1,000
35906
0.554
0.333
0.771
3.5003
1.2003
0.9917
C.006
25000
35000
D3915
1,0000
03254
03313
09771
03303
19500
1,000
0.503
0.773
16
3.5003
0,337
1,000
C.91T
15001
M5
55500
0.345
09703
CA
04
12
02
D31/1
25155
0.303
10000
25803
0.383)
COPIU
LECGC
25111
03/91 03953 01359 1000
05
(a) Attn:
(b) FFN:WI
(c) Post:oo
Figure 3: Cosine similarities for different layers in the model trained with the regularization strength
Aeg 10 for block size k (see Section 4 for details). The 24 layer model will have such blocks of size 4. The heatmap above shows the cosine similarities between weights for all pairs of blocks. Overall we find the final cosine similarity very high, thus suggesting a strong connection to looped models.
3,000
0000
0
0.3322
00000
L0001
03201
0.302
3.000
50000
30000
5C02
DCOS
0
03322
D0000
50000
00000
1.000
0/000
S1000
10000
03231
2,00
2901
600011
0.997
1/00
4:000
BJ007
10333
(10333
00007
3006
0.0000
0.662
1.303
0,002
3.00C3
0.0000
0.300
00000
0000
02000
0.007
0.0000
0023
5000
0.0053
4.007
.2.000
2.000
4.9335
6.000
1,000
(a) Attn:Q
(b) FFN:W1
(c) Attn: Aott:oos
Figure 4: Cosine similarities for different layers in the baseline model trained without any regular- ization strength. Overall the cosine similarities are very low for large matrices, as expected for high dimensions matrices.
16

---

Published as a conference paper at ICLR 2025
Ise-param Boseline
Looped Model
Iso- OP Baseline (ML-6
Midil Looping
(4
koops
bypers
Eeyoes
Figure 5: Illustration of the simple and architecture agnostic looping mechanism that we consider. A k-layer block looped L times (middle) is denoted by (A: @ L), which can essentially be viewed as a weighted shared model. The iso-param baseline. (A: 1). number of distinct parameters. The iso-FLOP baseline, (k& @ 1), is a kL-layer model with the same depth but times more parameters. Middle looping is strategy that is inspired from prior works on model stacking (e.g. (Saunshi et al. 2024)).
k-layer model with the same
i-GSM. We describe how the i-GSM dataset is generated in more detail here. We start with hierarchy of entities of depth 4 from which we build randomly sampled structure graph where directed edges connect entities in level defines a instance parameter which is an integer (for e.g. an edge between city center and car parks denotes the number car parks in city center). We then construct randomly sampled mathematical dependency graph which is DAG over all the instance parameters by relating each to the others. Finally we pick one of the nodes of the dependency graph to query and the goal is to compute the numerical value of this node modulo some prime number P. For more details process for the structure and dependency graphs, we refer the reader to Ye et al. (2024). We make 3 simplifications compared Ye et al. 2024): without the English construction of sentences (see Table 8); we do not allow abstract parameters in our problems; we perform arithmetic modulo as opposed 23. Our train dataset consists of around 4 million examples and we test on around 50k examples. Given the significantly larger number of unique solution templates possible, train-test overlap the problem template space going to be limited with high probability. For all models we train this dataset, the model dimension is 128, hidden dimension is 512 and attention heads are used. Rotary positional encodings are used as well. We train using Adafactor for 200,000 steps with batch size of 256 using base learning rate of 10 3 and use a linear warmup coupled with a cosine decay schedule for the learning rate.
those in level i 1. Each edge the structure graph
the sampling
phrase our problems in a symbolic language
B.2 LANGUAGE MODELING SETUP
We train on the Pile data using causal language modeling objective. The dataset is pre- rre-processed and cached to ensure that all models are trained on exactly the same tokens in the same order. For all experiments, we use batch size of 512 and sequence length of 1280. We use cosine learning rate schedule decaying over 400k steps with a peak learning rate of 0.01, tuned based on the baseline model. The base model is 1.5B parameter decoder only Transformer model, with 24 layers, model dimensions of 2048, hidden dimension 5120 and 32 heads. For the shallower baselines and looped models, we only change the number of layers and keep all other hyperparameters the same.
B.3 RESULTS FOR FACH TASK GROUP
In Section 3 we discussed results for various task groups like closed book QA, math word problems etc. Here we present details on the task groups and results for each individual task for completeness.
Closed book QA: This includes tasks like TriviaQA (Joshi et al. 2017), TydiQA-NoContext (Clark et al.] 2020). Natural Questions (Kwiatkowski et al. 2019) and Web Questions (Talmor) & Berant 2018) that test the model's 's ability to answer questions without any context, and thus, primarily measure the memorization abilities of language models.
Open book QA This includes tasks like TydiQA-Golde (Clark et al. 2020), SquadV2 (Ra-
jjuurkar al. 2018). Drop (Dua et al. 2019). QuAC (Choi et al. 2018 CoQA (Reddy al.
17

---

Published as a conference paper at ICLR 2025
Table 4: Downstream evaluations for the models in Table for closed book QA tasks.
Params/ FLOPs
TrivinQA
TydiQA- NoContext
Natural Questions
Average Average
Questions
Baseline
24x 24x
24.5
10.6
4.3
5.6
11.2
Base (12 09 1)
12x 12x
15.9
9.3
2.4
5.2
8.2
Loop (12 2)
12x 24x
18.6
9.6
3.4
5.8
9.3
Regularized (k 12)
24x 24x
20.9
10.9
3.6
10.1
Base [8 1)
8x 8x788
12.3
7.4
188
3,9
6.3
Loop
8x 24x
17.3
8.8
2.5
8.5
Regularized (k
24x
23.6
10.6
6.7
11.3
Base [6 1)
6x 6x
7.4
4.8
1.1
2.7
4.0
Loop (6 4)
6x 24x
16.0
9.3
2.7
4,9
8.2
Regularized (k
24x 24x
25.8
12.2
12.0
Base (4 @ 1)
4x
3.3
1.9
0.6
1.5
1.8
Loop (1 6)
4x 24x
11.8
8.8
2.4
3.5
6.7
Regulurized (A 4)
24x
26.1
13.3
4.5
Table 5: Downstream evaluations for the models in Table for open book QA tasks.
Parums/ / FLOPs
TydiQA-NoContext
SquadV2
DROP
QuAC
CoQA
Average
Baseline
24x / 24x
33.4
41.3
23.4
18.6
52.7
33.9
Base (12 @
12x
34.6
20.0
41.1
26.9
Loop [12
/ 24x
27.7
39.5
22.3
17.6
30.8
Regularized [k 12)
24x / 24x
33.0
44.3
23.8
17.7
51.9
34.1
Base (8 @ 11
8x 8x
12.7
33.8
16.0
15.3
35.9
22.7
Loop (8 3)
/
29.8
38.1
21.6
30.8
Regularized (A: 8)
24x 24x
33.0
41.7
25.2
19.3
52.9
34.4
Base (6 @ 1]
6x 6x
8.2
26.8
14.8
14.4
32.3
19.3
Loop (6 4)
20.8
18.1
Regularized (A- 6)
24x 24x
34.5
46.1
26.2
18.6
53.4
35.8
Base (4 1]
4x 4x
3.4
19.0
11.1
13.1
22.3
13.8
Loop (4 6)
24x
22.0
32.4
20.1
15.8
Regularized
24x 24x
33.6
49.4
24.1
19.8
54.0
36.2
2019) that evaluate the model's ability to infer the answer to a question from the extra context that provided, akin to reading comprehension.
Math word problems: To evaluate the model's ability to reason, we test them on math word
problems considered (Wei al. (0226). This includes tasks like SVAMP (Patel al. (02I),
ASDiv (Miao et al.) (020). the MAWPS benchmark Kooncel-Kezziosse et al. (0166). We report
5-shot evaluation for the pretrained model on these tasks.
Reasoning primitives: (Saunshi et al. (2024) introduced these datasets to study the inductive bias of stacking towards improving reasoning, by isolating simple reasoning abilities. One such primitive is depth-k variable assignment that requires the model to resolve chain of assignments of length k. An example depth-0 var-assign is a=1, b=2, c=6, b=?, and example for depth-1 var-assign a=1, b=2, c=a, d=b, d=?. We evaluate the math and coding variants of the depth-0 and depth-1 problems using 5-shot evaluation.
Detailed results for closed book QA are in Table 4 open book QA in Table math word problems in Table reasoning primitives in Table These tables include, both, looped models from Table and the models trained with regularization from Table 3
B.4 MIDDLE LOOPING VARIANT AND RELATIONSHIP WITH GRADUAL STACKING
Recently Saunshi et al. (024) introduced gradual stacking (Gong et al. 2019 Reddi et al. 2023 approach for training language models called MidAS. This approach gradually grows the model depth as training proceeds by duplicating certain layers of the model in each stacking operation. Surprisingly, they found that MidAS not only speeds up pretraining, but also improves reasoning in the same sense Figure I better reasoning at the same perplexity. Furthermore, the paper established a strong connection between stacking via MidAS and looped models, owing to the layer duplication operation. and conjectured that this is the reason for such inductive bias. Our results from the previous section provides a compelling evidence for this conjecture by showing that looped models also show very similar inductive bias, thus, further strengthening the connection
18

---

Published as a conference paper at ICLR 2025
Table 6: Downstream evaluations for the models in Table 2 for math word problems.
Purams FLOPs
ASDiv
MAWPS- AddSuh
MAWPS- MultiArith
MAWPS- SingleF4
MAWPS- SingleOp
SVAMP
Average
Baseline
24x 24x
26.9
49.9
2.7
38.6
40.2
17.9
29.3
Base (12 1)
12x 12x
23.1
44.8
2.0
36.8
37.9
15.3
26.7
Loop (12 @
12x 24x
31.5
47.4
50.2
18.5
34.3
Regularized (8 12)
24x
27.4
54.4
43.7
45.9
19.8
32.3
Base (8 1)
8x 8x
12.9
32.7
2.0
19.9
21.9
13.4
17.1
Loop (8
8x 24x
23.2
36.2
38.3
28.4
Regularized (A:
24x 24x
31.0
40.7
32.8
Base (6 1)
6x 6x
15.4
31.1
1.3
21.3
26.3
10.9
17.7
Loop (6 4)
6x 24x
24.5
51.6
0.7
38.0
47.7
16.3
29.8
Regularized (A:
24x 24x
32.0
31.0
Base (4 1)
4x74
7.9
10.4
11.0
19.0
8.3
9.7
Loop (1 6)
4x 24x
19.9
49.1
1.7
29.3
35.9
12.6
24.8
Regularized (k:
24x 24x
35.0
52.2
23.0
36.4
Table 7: Downstream evaluations for the models in Table for reasoning primitives.
Params FL.PPs
Code Depth-O
Code Depth-1
Average
Depth-0
Depth-1
Baseline
24x 24x
71.1
72.5
24.2
22.3
47.5
Base 0522811 1)
12x 12x
52.2
51.6
20.7
18.3
35.7
Loop (12
24x
73.5
86.5
Regularized (k 12)
24x 24x
75.1
74.9
27.0
25.7
50.7
Base (8& 1)
8x 8x
48.7
42.0
21.4
19.8
33.0
1oop
Regularized (k
24x 24x
92.2
86.7
23.1
23.3
56.3
Base (6 @6:811 1)
6x 6x
26.3
29.7
20.2
20.1
24.1
Loop 4)
90.6
88.1
24.1
21.7
56.1
Regularized (k:
24x 24x
84.0
79.7
31.4
28.3
55.8
Base (4 6 1)
4x 4x
19.3
23.3
17.3
17.6
19.4
Loop (4
4x
87.9
90.0
24.8
24.8
Regulurized (k
24x /4x
88.0
86.9
27.9
25.8
between stacking and looped models. Why such an inductive bias occurs is still an open question, and we believe that understanding this an important future direction.
Furthermore, inspired by their findings, we explore middle looping (see Figure for an illustration) variant of looping which maintains independent layers at the start and the end of the network, and perform looping on the middle block of layers. The high-level intuition from Saunshi et al. (2024) is that the first and last layers play special role in the model and thus, should be treated differently from the middle layers. In Table we report results for a version of middle looping that iso-param with (12 @ 1) bascline and iso-flop with (24 % 1) baseline, just like the (12 & 2) model. Overall, we find that middle looping has better perplexity and more uniform improvements than the default looping of (12 @ 2) (except for math word problems), and thus, might be more practical looping approach. We leave the exploration of the best looping strategies for future work.
C THEORETICAL RESULTS
C.I DETAILED NOTATIONS
Definition C.1 (Embedding Layer). Given finite vocabulary V, embedding dimension d E N+, token embedding parameter ÔTE E Id and position embedding parameter ÓPE E pxd we define the embedding layer as a scquence-to-eqquene map, denoted by EMBED V (Rd)" for any 1 < n S Mmaxx where
EMBED STED@Te.0or (U1,...,y,) = (fTE(v1) +
+ θpE(n))
Multi-Head Self-Attention Mechanism: Given attention parameters fATTN (Wo, WK. Wy, Wo). where each W W M. W W E Rdarinxd we define the Self- Attention layer with causal mask for decoder-only transformer in Algorithm [ We also
19

---

Published as a conference paper at ICLR 2025
Table 8: Left. Symbolic i-GSM problem and its solution. Right. Accuracy of looped and non- looped models on the i-GSM task from Section 2 (k @ L.) looped model is significantly better than the iso-param (k @ 1) model and performs as well as non-looped iso-flop (k- @ 1) model.
Parans FLOPs
Accuracy
Question. EBI EHJ ESI. K&N I8N J#0 F#K. FRK
8x
73.2
== E&J. INO == FAKK + KB0 + EHJ. HBJ := EBI + FAK. I8P := L8M I#N + K&O. 1&M
layer model
JNP
H8J- FAK.
18 Loop (1 2)
24.5
IAN JMP F#K, IAV := 2 J&P + HBI ERI. K#0 JAP + IAN
1x/2s
52.3
EAI. ISP?
Loop (1 (04
1x/4x
69.9
Loop (1
layer model
Answer with COT. = E81 LHI
ESJ
Base (2 0
2x 2x
54.0
= 4. FAK ERI. = FOR 4. HHJ ERJ+F#K. = HHI
Loop (2 (2 2)
2x 4x
66.9
HRI-FHK. = 4.
2INP+2HMM++F0L
Loop 4)
2x
4. LAM IANAJAP-FEK. = LAM 5, K#O JMP+IAN+ERI
4 layer model 4k
K#O = 5. IMP LAM = H8P
(4 8
71.3
Loop (1 @ 2)
4k 8%.
71.6
define Multi-Head Attention layer as a collection of self-attention layer with non-shared pa- (h)
rameters OMHA 10
ATTN h=11 and its output is the sum of the outputs from each head. That is, h==1
H
MHAm = Σ ATTN
ATTN
Algorithm Causal Self-Attention ATTN
Input: Parameter GTTNN (Wq. WK. Wv, Wo), Input embedding 21...., xn) (R!)".
Output: Output embedding x' = (af,..., r/1) å ATTNeATT (T).
1: q & WQz;, K; & WKF,, 22 a Wyz, Vi [m]
2: 8, softmax (g. k1) ..... (g, (O),
...,, 3: M4 A W Ey=1( (8.),
Feed-Forward Network: Given the parameters of the fully-connected feedforward network layer OFF = (W1, b, W2, 62) E R xd x gdsr RR iRder we define the feedforward layer FFo : Rd - d as F (h) A W2, relu h + bi) + bb2.
Definition C.2 (Output Layer). Given parameter OOUTPUT E Rdx we denote the output layer as OUTPUT OUTPUTTOUTRUU BOUTPUT (Rdd) + APP-- where
OUTPUT BOUTPUR
a softmax(1,, (I, T (OUTPUT)
(8)
Finally, we define the entire transformer model pa Unstmmar v A VI-1
P8 OUTPUT BOUTPUT TBOTe EMBED FEE 9PE
(9)
for any 0 = (0TB, ÔTE, OPE, (OUTPUT) ..r For convenience, we also write [pa (pa(v)..... Un (v) as pa (v U....., v. In particular, we use (Fe("1....,, U.)) arg max pe(u|ui...., 2,) to denote the deterministic version of the transformer model.
Finite-precistie Modeling: In this paper assume the transformer is of finite precision. More specifically, we follow the setting in Li et al. (2024) and use the shorthand F, {c 2- C (-1,1),0 < k < 22 1, k E NJ to denote fixed-point numbers of constant precision & and rounding operation Hl. : IR F. to denote the correcting rounding, namely the mapping from I to the closest representable number in F.. (We break the tie by picking the number with smaller absolute value). We assume that (1). all the parameters of the transformer are in F, and (2). correct rounding is performed after every binary operation in the forward pass of the transformer. We will refer the readers to Li et al. (2024) for detailed discussion on such finite-pne preci-poe modeling and only list important notations and lemmas that will be used in this paper below.
Thoough in this paper focus attention with casual mask, our definition of looped transformer gener- alizes to the cases with other attention masks.
20

---

Published as a conference paper at ICLR 2025
Beseline
36.
Closed Book OR.
2.
-0-
Looged 212
Baseline (12 1)
Math Word Probilerss
20
30
Looged (12
52.5
501
415
45.1
Baseline (12 9 92)
4/
Baseline (12 11
20.
455
Looped (12 2))
36.5
21s 2:2
Perplexity
Log Perplexity
Perplexity
Perplexity
Math Word Probleems (calc)
Closed Book OA-
12
Regularized (k
Baseline (24 11
Beseline (24 e 1) Regularized (X
-4-
Regui arguuerriedd
Baseline (24 e 1]
ng Primitives
-2-
Baseline
12
50
Reqularizd (K
66
2:0 Log Perplexity
2.50
2.25 2.05
2.04 2.50
2.50 Log Perplexityy 233
2014 2.50
2300
2.00
2.00
2.04
Log Perplexity
Log Perplesity
Figure 6: Downstream evaluation for various task groups on the x-axis, vs validation log perplexity on the y-axis (reversed), as training proceeds. The top plots compare a 12-layer baseline model (12 @ 1) and the looped model (12 @ 2). The second row compares the baseline 24-layer model and the 24-layer model trained with regularization using block size k: = 4 and A = 10 (See Equa- tion (4)) For both comparisons we have similar observations. For closed book QA (memorization) tasks looping has very similar trends to baseline. For open book QA tasks and math word problems, looping has much better downstream performance an equivalent log perplexity. This verifies the inductive bias looping and regularization towards better reasoning abilities.
We use 1 to denote all-one vectors of length s. Similarly define (: .. xs, and softmax,- We recall that for any E N+ and integer 0 < x S 2 1, we use (in E {0, 1]. to denote the usual binary encoding of integer æ using binary bits in the sense that i -.= 22(nin, (x)); and szin,(2) E 1 to denote the signed binary encoding, which is 2bin, (x) (l.....1) 1). Finally
we define B max F,, 2" --
Lemma C.1. [Lemma E.I, ( al.
2024 For any E N+ it holds that (exp(-B,)) O.
Lemma C.2. [Lemma E.2, (Li et al.
2024 ] For any s E N+, it holds that (exp(B.)l. B..
Lemma C.3. [Lemma E.5, (Li et al.
2024)] Unlimited-fanin AND, OR (resp. MAJORITY) :
{0, 13" {0, }} can be simulated by some 2-layer feedforward ReLU network with constant (resp. log n) bits of precision constant hidden dimension and additional n constant inputs of value
Mathematically, let FF[s(n)] be the set of functions C {0, 10, + {0. 1} which can be two- layer feedforward ReLU network with at most s(n) bits of precision and constant hidden dimension FF. : {0, 1)2n > (0, 1), FF)()") = W2 x, relu([W; x, x' + bı),), where (W2, W1, b1), such that for any x E {0, 13"
In, , 1) = C(x). FFo 12 I,...,,n,,
(10)
We have unlimited-fami AND, OR FF[1] and MAJORITY E FF[log n].
Given two vectors y of the same length s, we use xyy to denote their interleaving, that is, (x y)z-1 = Dz, (x y)2 = Ye for all E [e].
Lemma C.4. [Lemma E.3, Li et al.
2024 For any s E N+, let qi (i) 1, and k; B, (sbin,(i) (-1)) for all E 2 [2-1]. holds that [exp((1; (1) [ = for all i, [2 -1}
21

---

Published as a conference paper at ICLR 2025
Evall Porplexty (0uepleew 0.5)
Closec Book (ffop/
Open Book
10.0
Baseline (D @
Baseline 10 11
Loopeci a DM)
Looped (4 s D/4)
Jccuracy
Looged (4 @ DVA)
Effective Mocel
Effective Kodel Death
Effective Model Depth
Problems lQlcellteaee
Rossoning (chooe Soare
Eval Token Acc (0 Mouse 0.5)
45
Boseline 1D
Baseline 1D 11
Baseline I0
Looped Di4)
Accuracy
Looped DV4)
Looged' (4 e D14)
15
13 24 76
12 24
Effeclive Model Depth (D)
Effective Mocel Depth (D)
Effective Model Depth (D)
Figure 7: Scaling behavior for various task group as the effective depth increases. The blue curve
shows how performance scales the number of loops increases, without increasing parameters, ing models the form (4g) D/4) for various values of D. The orange curve visualizes the scaling be- havior of (D&1) which increases the depth by adding fresh parameters. For reasoning primitives, the looped model scales as well, or even better, than the baseline despite having D/4 fewer parameters.
C.2 PROOFS
C.2.1 LOOPED MODELS CAN SIMULATE NON-LOOPED MODELS
Proof of Theorem 5.2 We start by introduce some more notations. We will proof the theorem for
any fixed sequence of (UI...., vn). We use (!)) (1) ( to denote the intermediate embedding of pe in the Ith layer. More specifically, we define
x' = (id FF a-)) o (id + MHA (1-1)) o.... (id FF ) (id MHA (03
EMBED RMBEDOEWBED
DEMBED (v).
Off
WHA
(11)
We also use (1+.55) = (x) (l+0.5) after the attention layer.
1=1 to denote the intermediate embedding of pa in the 7th layer
x! 1+0.5
(id MHA pit (xx)..
(12)
Similarly, for the constructed looped transformer Po,r, use V' (#1 U1..... va) to denote its input. For simplicity, we use the convention that # is at position O. The proof still works if # starts position because we can just transfer the later tokens by position. We define (1) = (1) 2/1 ....,,,,,
(() In (()) as the intermediate embedding of po in the Ith layer and (1+00.5)
(140.5) (140.5) Fn ((+00.5) the intermediate embedding of po in the ith layer.
Below we first state the key properties that our construction will satisfy, which imply the correctness of the theorem and then we state our construction of por, and show the properties are actually satisfied:
I (x) 1R Cr(), 4, 1 = =]).
(1+0.5)
(x. (1+0.5) 1 Ex(), 1, = 0)).
(1) (l+0.)) 0
not defined in the original non-looped transformer. The key point here is that they are 0 vectors throughout the forward pass.
Here we abuse the notation for simplicity of presentation. zo x (1+055)
(0)
22

---

Published as a conference paper at ICLR 2025
Algorithm 2 Group Composition
Input: Group elements gu. 91....... In E G, where Po = e. Output: 90 91 o ... In-
1: 9: = 9., v0 S. i S n.
2: for
+ 1log n do
3: a 9 922-n-11 6/ = (1-1) n2+
(!) (1-))
40 siiss sism
(1) (1) bb 44 H n.
4: 9 a;
5: end for (flog, ml)
6:
return On
To get the last coordinate 1, which is a depth counter, we just need to add more hidden dimension in MLP.
Next we show we can use two-layer with L + 2 MLP to get the mapping from l + 1 Er(1). Let (ofF) o / M be the parameters of the distinct layers in We assume in the 1th layer, r())s
parameters are used. This is because er(z) = Σ C.5 * ( [1-i+t1 1|+ - 2[1 - +] +
1]+)-
Now we explain how to deactivate the undesired MLP neurons. In other words, our construction of OF is essentially concatenation of off for E [r] in the hidden dimension of FF, with the additional control that FF 0/F ((x) l Er()), 1), 1 [i 0)) = - FF ata (x: ( ) [r(I] = i, * 0] at /th layer.
This control can be done by subtracting (r(1) + 1 0] by a constant which is larger than the maximum pre-activation in the hidden layer.
Finally we explain how to deactivate the undesired attention. We will only use attention to update the first part of the embedding, which is x ((+0.5) A erucial step here is that set the token embedding of # as O We construct keys and queries as follows:
1. wrinr) (z (0))
(W (2) ( 11 = r()] for for E [R] and = .,....
.
2. winr))
(EF 11)) = (w(r)) (0) BB1 [ = 0)) for r' E [|]] and = O,..., where K 24 B some constant larger than the maximum previous inner product in attention, max (WWxz," (Woz!") (!)
K
3. "org""""& (241)) (1)) (W! W 20.22,00000000 0.0,0)
This construction works because only the 'desired' attention head r = =I()) will be activated and behave in the non-looped case, because otherwise all position in that attention head will be com- pletely attended to position and returns zero vector. (We can choose / to be large enough and distribution calculated by the attention score is delta distribution) position 0, which yields zero vector as its value. This completes the proof.
C.2.2 GROUP COMPOSITION.
The landmark result in automata theory, Krohn-Rhodes Decompotision Theorem (Krohn & Rhodes 1965), shows that all semi-automaton with solvable transformation group (which includes compo- sition problem of solvable groups) can be simulated by a cascade of permutation-reset automata, which can be simulated by TC circuits. (Liu et al. 2022) further showed that such automaton with solvable transformation group can be continuously simulated by constant-depth transformers. However, it is also shown (Barrington) 1986) that the composition problem of unsolvable groups are NC! compleee,,, for example, the composition of permutation group over 5 elements, Ss. Under the common hardness assumption that NC² + TC constant depth transformer cannot solve composi- tion of S5 using single forward pass (Merrill & Sabharwal 2023b) Liu et al.) 2022 Li al. 2024). But with CoT, very shallow transformers (depth equal to one or two) can simulate the composition problem of any group(Li et al. 2024 Merrill & Sabharwal 2023a).
23

---

Published as a conference paper at ICLR 2025
Proof of Theorem 5.7 We will set the token embedding of # the same that of e, which is the identity of G. In the following proof, we will just treat # as e. We will construct the transformer simulating group composition following the following algorithm AlgorithmB which gives the high- level idea of the construction. The correctness of Algorithm@folloo follows from the associativity of group composition. More concretely,
verify by induction that gó of p..-9h In is the same for all and in the final round, i.e., when / = [log2 n1,, 92 = e for all i <
1=
[log2 ni
Below we show how to construct a transformer of the given sizes to simulate the above Al- gorithm 2 We will embed each g E G' as a different vector g E {-1, no [G]] and each position 0 S I n as its binary representation in i E {-1,1} (o02 =+11, which is a shorthand for (i) with = [log2 n2 + 17. We concatenate them to get {-; (0) J= that is, 2; (9:, 7, [2 - n +]+ 22 n 1]+, 0 oflog, |Gll (olog (Gll). For convenience, we will drop the O's in the end (also in the other proofs of the paper) and write it as 1;
(g., i, [2i n 1|+, [2i n 1]+). Below we show we can construct I-layer transformer block with parameter (OMHA, OFF) satisfying that
1. MHAGMHA
( 0,.i,,2 2i n 1|+ 2i n 11 1 = = 1ogg GI] ++1], 912k 1- n 1|, 912k-n]+
for all 90 = 0, E GVi E [ni, k= O......
2. FF J. g. gg) = ( 3.0% 03 Tome n-11 - -g")"" for all i,j, = 0.
9.9'. g" E EG.
The first claim is because we can two attention heads to retrieve g 122k-n-2 n- 1] and g(2k-n]
spectively, where both of them use as the key and use -12k n 1 and - 12k - n]+ as queries respectively. This is possible because all the required information already in 2;, We further make attention temperature low enough so the probability returned by attention is a one-hot distribution at the position whose key is equal to the negative query after rounding.
Now turn to the second claim about MLP. We will use |G|2 neurons with ReLU activation and bias to simulate the product of g' and g". We can index each neuron by (h, //) for ht, h' E G and set its incoming weight [Wi [Wila,,'), = (h, I7) and set bias (bz) = 2 flog2 |G|] +1, which ensures that the activation of neuron (h, h²) will only be when g' = h.g" h and otherwise. Then setting the outgoing weight of neuron (h, h²) as h h and the bias in the second layer to be 0 finishes the construction for simulating the group composition. Finally we use the remaining 6 6llog, G+] to simulate negative identity mapping æ + -r for the remaining 3 logg |G]] embedding dimension. This completes the proof.
C.3 CONNECTION TO CHAIN-OF THOUGHT REASONING
In this section, we establish connection betwee looped models and CoT reasoning. We first define the recursion for CoT reasoning as follows:
TF (11..... v.)) A TF (v) (f1....,, ..... Un TF,(
for inn E N+ satisfying i + n < Mmnx 1 along with the base case of TF (U),....... v.) A (U)....... v.). For all < S Mmax n 1, the output with steps of CoT Un4i41
TF (v.,..., mn) = TF (V1. Pay Pat1, Un+)).
We first give the formal statement below.
Theorem C.5 (Looped transformer simulates CoT). For any L-layer non-looped transformer TF there exists a looped transformer TFş with + O(1) layers, constantly many more dimensions in embedding. MLP and attention layers and constantly many more attention heads, such that for any input e = (#=) and integer m, the output of non-looped transformer after m steps of CoT, TFI (v), is the same as that of the looped transformer on input concatenated by m dummy tokens with m loops, TF (w, #
Below some helping lemmas towards showing Theorem
C.5 is least powerful COT.
24

---

Published as a conference paper at ICLR 2025
0
0
0
Figure 8: Left. Chain-of-thought reasoning can viewed as a looped model, where each iteration produces one new thoughts token. The new tokens are highlighted in red. Right. A looped model can instead generate multiple latent thoughts in parallel and, in theory, can simulate CoT reasoning my masking the updates appropriately (see Theorem 5.4)
Lemma C.6 (Simulating arg max using MLP). For every d N and precision E N+, there exists a 3-layer network with relu activation and d² width I with precision, such that for any x E Ry there is [d], such that 2 maxjzk.je|d] Zyz (xx ek.
Proof of Lemma C.6 Define 9i = 2 relu(2- -jzt relu(a, xx)) for each i [n]. We claim
that if there is k [d]. such that wh max zij 2 2- 9: 1 iff k for all [d]. Next for i * k, it clearly holds that Σj relu(z, x;) 2-a and thus 9s = 0. This construction clearly be implemented by 3-layer relu network with precision.
First Is = 2 relu(2-*)
Lemma C.7 (Simulating Decoding and Embedding using MLP). Given any s-precision ÔTE E Rdx: and fouTPUT. there is 5-layer network f R - Rd with relu activation and max([E]²) width with s-precision, such that for all s-precision a E Red which admits unique arg max for v & arg max,ex(1* (O))PUT) it holds that
f(x) = OTE(U).
Proof of Lemma
C.7 This is simple application of Lemma C.6
Lemma C.8 (Control Gate). A 2-layer relu network with precision can implement F: x F3 (0, F(x,y,
Mx + M y.
Proof of Lemma C.8 Note that F(x, y, M) = relu(z - 2 (1- M)) - relu(- 2 (M) + 2 M) relu(-11 2 M). The proof is completed.
relu(y-
Definition C.3 (Transformer Block with Mask). Given number of layers L N+ parameter θTE (0
L-1 and mask function M .N {0, 1}, define the L-layer transformer block with maskTB ThakkBBore.. / (Rd)) + (IRd)" for any n E N+ as
MHA OF FF)1=0
(13) Definition C.4 (Looped Transformer with Mask), Given number of loops T E N+, pa- rameters A (0Tg, ÔTE. 8pE. (OOUPUU), and mask functions {M') where BTF define the looped transformer with mask as Р8,T, & OUTPUTROUTPUT ... TBo EMBED and the corresponding deterministic version as
TB M (x)), (1 M())zz + M(i) (z)],
(o
Oll' L-1
(MHHA FFF
TB
TF M.M. M(1 (U1...., Un) & arg max.ey P0,T,M (M/ (v/01,--.,U5,
25

---

Published as a conference paper at ICLR 2025
Definition C.5 (Shifting Layer). We define the shifting layer SHIFT following for any d, n E N+ and 21,.... In E Rd:
(Rd) - (I")" as the
(HI 22. 23. ,3, (#1, 23,....... " 5n..1.
(14)
Lemma C.9. For input sequence length up some integer n, SHIFT could be implemented by a attention layer by concatenating each embedding with sbin, where n [log2 n 17.
Proof of Lemma C.9 It is equivalent show we can construct attention heads which computes f(i) at each position i.
To do so, we just need to invoke Lemma C.6 and use that to set key and query, so position i attends to position f()). We set value position to be T₁. This completes the proof.
Lemma C.10. For any positive integer > O, there is a constant-depth MLP F with O(s) hidden neurons per layer and parameters in F., such that for any input bin(z) (-1,+ where <
S 2 - 1, it holds that
F(x) bin(x: 1)-
Proof of Lemma 1.10 By Lemma C.3 suffices to show that simulate bin(x:) bin using (s)) wide, constant-depth boolean circuits with AND, OR, NOT gates with unlimited fan-ins. This is immediate by noting that
[bin(x + 1)1, = [bin(= + 1) + bin(# x + 1)1;
(15)
Lemma C.11. For any positive integer s > O, there a constant-depth MLP with O(s) hidden neurons per layer and parameters in F., such that for any input (bin(x), bin(y)) E {-1, +13* where 0 S x. S 2 1, holds that
F(x, v) = 1 [z > yl
Proof of Lemma C.11 By Lemma C.3 it suffices to show that we can simulate bin(z), bin(y) + Ja > y) using O(s) wide, constant-depth boolean circuits with AND, OR, NOT gates with unlim- ited fan-ins. This is immediate by noting that
1 1 > v) V
([bin(y)): = = 0) ([bin(z)]; A
A 15jei
([bin())], [bin())l))
(16)
1E(x]
Proof of Theorem C.5
We consider mask 1 n). which we call it CoT masking. Let [log(n + m 1J. we use i to denote (b)n, for 1 n m for convenience. The embedding size of our newly constructed transformer is larger than the target transformer to be simulated by an additive constant of 3s. We denote the new embedding at position after ith loop (t) where x will be the original embedding of the transformer to be simulated, and (e) of dimension 3s and only depends on i and
by (x)" .P:
p;
particular, we can show that we can set (e) à (7, ( 2]_ + I, f) to save information about position and the number of loops - ,, is from the positional embedding and the update is due to Lemma C.10, The size of hidden dimension of MLP and attention (key, query, value) will also be increased by O(s).
P;
The proof contains two steps:
1. To show that there is a transformer with CoT masks simulating the target transformer with m steps of CoT and L layers by looping its own L + O(1) layer block times.
26

---

Published as a conference paper at ICLR 2025
2. To show that the above looped transformer with CoT masks can be simulated by standard looped transformer without mask and with constantly many more layers.
For the first claim, starting from the same parameter of the transformers with CoT 8, we build new looped model with parameter 0' with constantly many more layers in each transformer block and most constantly many more heads per attention layer. First, we can add constantly many more layers to use MLP to simulate the decoding- secoodingenncding process using Lemma C.7 Next, we can add one more transformer layer in each block and use the attention layer to simulate the shifting layer by Lemma C.9 since we have the additional position embedding 1) In particular, the embedding we get at position n + after loops, (4))) now simulates the token embedding of n + L of the CoT transformer. By the way we define CoT mask M, for every t 2 -n 1, the embedding a will keep the same for all 2 max(1, 0). In /th loop, the only embedding update that matters happens at n. + /th position, because no updates happen at earlier positions, and updates at later positions t' for some t' > t will be overwritten eventually in the future loops t', by some value which independent of their value at the current loop In the end, we know the embedding x(²) in Definition C.4 is exactly equal to that in CoT transformer, and so does the final output.
For the second claim, because CoT mask can be computed by a O(log(n + +m) wide, constant- depth MLP [Lemma[C.11) together with Lemma know it suffices to increase the number of layers per transformer block and embedding size and hidden dimensions by constant to simulate the transformer with mask by a transformer without mask.
Theorem C.12 (Restatement of Theorem 4.2, (Sanford et al. 2024b)). p-hop problem (Defini- tion B.1 can be solved by |log2P| P] + 2-layer non-looped transtormer with log n bits of precision, at most different layers, d = dFF dATTN O(1) embedding size, hidden dimension for MLP and attention, attention head.
27


Reviews:
[{"review_id": "2A198YklTR", "reviewer": "Reviewer_Rw4n", "summary": "This paper studies performance of looped transformers on reasoning tasks. They show that looped models result in worse memorization of facts reflected by lower perplexity but can outperform non-looped models on some tasks that require compositional generalization. They highlight that perplexity may be unsuitable metric because looped and non-looped models with same perplexity can result in substantially different performance on downstream tasks requiring reasoning. Lastly, they introduce regularization of non-looped models which forces the weights to be closer to each other in terms of cosine similarity. This leads to a same perplexity but better downstream performance on reasoning tasks.", "questions": "1. An obvious goal for the looped models would be to figure out how to make the looping adaptive, i.e. that the model would be able to decide when to stop the looping. Did you considered experiments which would test this ability? This could provide clear benefits over the non-looped baseline and I believe it would be substantial contribution. From Table 3, it is visible that different tasks require different number of loops.\n2. Why do the plots in Figure 1 contain datapoints for different perplexities for the two compared models?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "In terms of originality and contribution, the paper demonstrates that a simple regularization can improve reasoning abilities of transformers and provides a basic theoretical analysis of looped models. The writing is clear and the experiments seem sound.", "weaknesses": "Personally, I do not find the results very surprising. The fact that looped models can improve performance on algorithmic tasks is known and the extension to the tasks tested in the paper does not seem substantial. It seems to me that that the authors are not aware of the work done in the group of Tom Goldstein related to reasoning abilities of recurrent models (different term for looped models):\nhttps://arxiv.org/abs/2202.05826\nhttps://arxiv.org/abs/2106.04537\nAnd also Daniel Selsam's paper which also showed the importance of recurrence: https://openreview.net/forum?id=HJMC_iA5tm&referrer=%5Bthe%20profile%20of%20Daniel%20Selsam%5D(%2Fprofile%3Fid%3D~Daniel_Selsam1)\n\nMinor:\nThe first sentence in related work contains a typo: intelligent and robustly model", "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": "zSGNDS9Kjn", "reviewer": "Reviewer_2Phq", "summary": "This paper investigates using looped language models for reasoning tasks. The authors demonstrate that reasoning tasks can benefit from deeper models achieved through looping without significantly increasing the parameter count. They show that on synthetic reasoning tasks like multi-digit addition and p-hop induction, looped models can match or even surpass the performance of deeper models with more parameters. Additionally, they identify an inductive bias in looped models toward reasoning tasks and propose a looping-inspired regularization technique to capture these advantages without substantially compromising language modeling efficiency.", "questions": "1. Given that looped models may have higher perplexity but better reasoning performance, how do you suggest balancing these aspects in practical applications where both are important?\n\n2. Can the looping approach be extended to other domains, such as computer vision models like vision transformers, or models employing different attention mechanisms?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The experiments on synthetic tasks like multi-digit addition and p-hop induction provide a valuable starting point for understanding the strengths of looped models, particularly in recursive reasoning scenarios.\n\n2. Proposing a looping-inspired regularization technique is a noteworthy contribution. It offers a practical method to capture the inductive bias of looped models while maintaining competitive performance in perplexity metrics.", "weaknesses": "1. Some explanations, particularly regarding the looping mechanism and training setup, are not as clear as they could be. The paper lacks detailed explanations of how the transformer architecture implements the looping mechanism. For instance, it’s unclear how the KV cache is managed when layers are looped multiple times. Additionally, how does looping interact with the residual connections and layer normalization typically present in transformers? More detailed descriptions and possibly visual aids would enhance understanding.\n\n2. There’s limited discussion on the computational efficiency and scalability of looped models, including potential overhead during training and inference. Specifically, it’s important to know how the training time and inference speed of looped models compare to standard models. Does looping result in longer training times due to repeated computations, or does it impact memory consumption during training and inference?", "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": "qcPV95oxE5", "reviewer": "Reviewer_Ribu", "summary": "The work explores the looped models in reasoning tasks, demonstrating that they can achieve deeper reasoning depth with fewer parameters compared to the traditional models. The authors also conducted comprehensive experiments on reasoning tasks like addition and math word problems, and looped models showed improved performance in reasoning, despite having worse performance on memorization tasks. The findings may offer a potential new direction for language model development to improve its reasoning capacity.", "questions": "please see the weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This work provides extensive empirical studies showing that looped models perform well on reasoning tasks, even outperforming traditional models in some cases.\n2. This work offers theoretical results that support the expressiveness and efficiency of looped models in solving reasoning problems.\n3. The paper is well-presented, making it easy for readers to follow, despite some complex theoretical analysis.", "weaknesses": "1. There is limited evidence on the performance of looped models in real-world applications or datasets.\n2. It would be much better if the authors could scale up the experiments, and validate the effectiveness of the looped model in real-world tasks.", "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": "7qFvESnaJy", "reviewer": "Reviewer_iCpz", "summary": "The paper investigates the performance of looped Transformers on arithmetic and reasoning tasks and compares it with feedforward models with the same amount of flops. The authors show that looped models can equal and sometimes outperform models with a comparably larger amount of parameters by simply looping over the network, although they demonstrate a worse perplexity. The authors also propose a regularization method to induce feedforward models to behave as looped models and show that it yields better performance on reasoning without reducing perplexity.", "questions": "1. In Figure 1, can you explain how each of the data points are generated?\n\n2. In Table 4, do you have an explanation as of why performance increases when setting $k=4$ but then is reduced when $k$ increases?\n\n3. I am curious to know how looping compares with chain-of-thought as the latter effectively performs looping via autoregression but introduces a bottleneck between each loop with the next-token selection. Have you investigated the similarities and differences between the two approaches in your experments?\n\n4. Have you further studied the connection between the number of loops and performance in your experiments, particularly with a high number of loops? Does accuracy keep increasing as the model loops or where is the upper bound? Does it differ depending on the model size?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The experimental setup is sound and the findings are interesting. Although the good performance of looped models is not novel, the paper provide clear insights into the reasons, and makes a disctinction between reasoning and memorization performance. The claims are properly backed-up by the experiments.\n\nThe looping-inspired regularization is a particularly interesting and original idea to improve reasoning in Transformers.\n\nThe paper is also well-written and easy to follow.", "weaknesses": "Experiments lack a comparison with other categories of models that use looping in the computation (e.g. recurrent Transformers [1,2] or Transformers autoregressively returning an output via chain-of-thought or other prompting strategies). Adding comparisons on the arithmetic and reasoning tasks could help to further discriminate what makes looping efficient.\n\nThe results in Figure 1 and Table 4 are not thoroughly explained and are hard to interpret. In Figure 1, it is not exactly clear what the y-axis represents and therefore what each data point represents. In Table 4, the choices for the values of $k$ and $\\lambda_{\\text{reg}}$ are not well justified and a deeper analysis of the results could be helpful.\n\n[1] Bulatov, A., Kuratov, Y., & Burtsev, M. (2022). Recurrent memory transformer. Advances in Neural Information Processing Systems, 35, 11079-11091.\n\n[2] Hutchins, D., Schlag, I., Wu, Y., Dyer, E., & Neyshabur, B. (2022). Block-recurrent transformers. Advances in neural information processing systems, 35, 33248-33261.", "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)