Looped Transformers for Length Generalization
Authors: Ying Fan
Yilun Du
Kannan Ramchandran
Kangwook Lee
Venue: ICLR 2025 Poster (2025)
Keywords: Transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Recent work has shown that Transformers trained from scratch can successfully solve various arithmetic and algorithmic tasks, such as adding numbers and computing parity. While these Transformers generalize well on unseen inputs of the same length, they struggle with length generalization, i.e., handling inputs of unseen lengths. In this work, we demonstrate that looped Transformers with an adaptive number of steps significantly improve length generalization. We focus on tasks with a known iterative solution, involving multiple iterations of a RASP-L operation—a length-generalizable operation that can be expressed by a finite-sized Transformer. We train looped Transformers using our proposed learning algorithm and observe that they learn highly length-generalizable solutions for various tasks.

Published as a conference paper at ICLR 2025
LOOPED TRANSFORMERS FOR LENGTH GENERALIZA- TION
Ying Fan', Yilun Du², Kannan Ramchandran', Kangwook Lee'
University of Wisconsin-Madison Massachusetts Institute of Technology UCC Berkeley
ABSTRACT
Recent work has shown that Transformers trained from scratch can successfully solve various arithmetic and algorithmic tasks, such as adding numbers and comput- ing parity. While these Transformers generalize well on unseen inputs the same length, they struggle with length generalization, i.c., handling inputs of unseen lengths. In this work, we demonstrate that looped Transformers with an adaptive number of steps significantly improve length generalization. We focus on tasks with a known iterative solution, involving multiple iterations of a RASP-L operation-a gength-generalizable operation that can be expressed by finite-sized Transformer. We train looped Transformers using our proposed learning algorithm and observe that they learn highly fngth-ggnneralizable solutions for various tasks.
INTRODUCTION
Most algorithmic tasks such as coding, writing mathematical proofs, and reasoning are defined with inputs of variable length. The length of an input often correlates with the difficulty of the problem instance. For example, the longer the input, the more difficult the problem tends to be. We say a model perfectly length-generalig
if it can solve an algorithmic task on inputs of any length, even if it was only trained on data with inputs up to finite length (Anil et al. 2022). Generally, it is hard to expect models to be trained on inputs with all possible lengths, and we need to rely on length generalization. Also, if model can length- getthheenrralzze, it means the model has truly learned the correct algorithmic solution to the task, not just spurious solution that works only for certain lengths.
Recently, many works on Large Language Models (LLMM)) have shown that we can get more powerful Al models by scaling both compute and data training time. This scaling approach has indeed succeeded in improving accuracies on various benchmarks. However, even the largest and latest LLMS like Achiam et al. (2023) trained on much of the existing text on the Internet, still struggle with length generalization (Wu et al.. 2023 Anil et al. 2022. Lee et al. 2024). One possible cause is the particular computing model. LLMs are built based mostly on the Transformer architecture (Vaswani et al. 2017). While Transformers can accept variable length of inputs (that can be processed in paralle parallel), they usually have a fixed depth. This might be sufficient for certain tasks, but not always.
To learn model that effectively generalize to longer problems, it is important to consider archi- tectures that can adaptively adjust the computational budget to the difficulty of the tasks (Anil et al. 2022 Du et al. 2022 2024). One approach achieve this is to explicitly generate intermediate output tokens, similar to writing down a scratchpad, which improves LLMs' capability for solving harder problems (yye et al. 202T). In theory, LLMs may generate more scratchpad tokens representing intermediate computation when solving more difficult task, indicating that they can allocate elastic computation according to the length and difficulty of the given instance. This approach can be learned by explicitly training model on data with intermediate computation steps (Ling et al., 2017 Cobbe jet al. 202T). Alternatively, it can be achieved via Chain-of-Though (CoT) reasoning with few-shot examples (Wei et al. 2022) even in a zero-shot manner (Kojima et al., 2022). Notice that these approaches still use fixed-depth models. While these approaches help solve more complex reasoning tasks, they still far from achieving near-perfect length generalization for simple algorithmic tasks. For instance, Lee et al. applied CoT for arithmetic tasks but observed that Transformers cannot length generalize even for simple addition tasks (Lee et al. 2024).
Recently, there has been growing interest in using recurrent architectures for reasoning (Dehghani et al. 2018 Bai et al. 2019 Bansal et al. 2022 Yang al. 2024). Unlike standard RNN-type

---

Published as a conference paper at ICLR 2025
architectures that process the input sequence incrementally, one consider recurrent architecture that processes the entire input sequence multiple times, passing the intermediate output to the next iteration's 's input, possibly along with the original input. In particular, if the base model in each iteration is a Transformer, this model is called a Looped Transformer (Yang et al. 2024).
naturally break the limita- tion of the fixed depth in the standard Transformer architecture: One can adjust the number of looped steps based on the computational complexity of the underlying algorithmic solution. Consider a prob- lem set where 1) The problems can be solved by a loop of one RASP-I. (Zhou et al., 2024a) progran[] i.e., each step in the loop can be performed by a decoder-only Transformer with a fixed depth; 2) The number of steps needed in the loop depends on the problem's ' complexity, i.e., more difficult problems could potentially require more steps to solve. Under the length generalization scheme, we consider the number of steps depending on the problem length, and define this problem set as n- n-RASP-L. problems. For r--ASP-L problems, if we can learn these length- independent steps, we can utilize an adaptive number of steps achieve length generalization.
Looped Transformer
Output
*
*
L
In
#
Decoder Block
X n
E
L n
#
#
Input
Figure 1: Method Overview. During training, we supervise the output of the model to match the target data only after certain number of steps of applying the same decoder block, helping the model learn intermediate steps that can be reused and can handle input of arbitrary lengths. All grey blocks share the same parameters. Examples are from the Copy task with symbols. "#" indicates EOS, "e* indicates ignored output, and ">" indi- cates the end of the query (EOQ).
Inspired by this observation, we study training Looped Transformers models for length generaliza- tion. Specifically, we consider a training setup where we not require any intermediate supervision data (such as reasoning steps or scratchpad). We only assume access to end-to-end supervision (input and output) and the number of steps needed. Depending on the number of steps, iteratively apply the same decoder block and then decode the final answer; See Figure I for illustration. At inference time, the model could either decide when to stop with predefined stopping criteria stop when reaching the ground-truth number of steps. Empirically, we show that looped Transformers with an adaptive number of steps can successfully length-generalize to longer lengths simply by appropriately adapting the number of loops at inference time, indicating that our approach encourages the model to implicitly learn the necessary steps solve a task.
Our contributions can be summarized as follows: (1) We first formally define n n-RASP-L problems, and provide examples of n-RASP-L solutions the Copy, Parity, and Addition tasks (Section B: (2)
We propose to learn n--ASP-L problems with Looped Transformers where we supervise the final answer step-dependent way, which enables us to use an adaptive number of steps depending on the problem complexity (Section 4 (3) Empirically, we show that our proposed method outperforms the baseline approaches in terms length generalization performance (Section 6)
2 BACKGROUND
2.1 RASP-L
A decoder-only Transformer is a type of Transformer architecture that consists of only the decoder part of the original Transformer model introduced by Vaswani et al. (2017). where a causal mask is applied to the attention weights to prevent the model from attending to future tokens.
RASP (Restricted Access Sequence Processing) (Weiss et al. 2021 is a computational model for the Transformer architecture in the form of a programming language. RASP-L (Zhou et al. 2024a). is learnable subset of the RASP language. Some key points about RASP-L are:
RASP-L programs accept an input sequence and return an output sequence of the same length for arbitrary length, like decoder-only Transformers.
'eere consider more general way to loop, i.e., predicting all missing tokens at the end of the loop, not necessarily in the way of predicting the single next token time. See more discussions in Section 2.2
2

---

Published as a conference paper at ICLR 2025
The core operations in RASP-L include element-wist operations sequences and a specific type of non-elementwiss operation called kqv, which simulates a causal attention layer.
RASP-L has restrictions the allowed operations to ensure learnability: It does not allow arbitrary index arithmetic, and restricts operations on token indices to order comparisons and computing successoo/predecessorr
RASP-L does not allow control flow statements like branching or loops. Programs must be straight-line code, with each line being call to core function or another RASP-L program.
In Zhou et al. (2024a), they show algorithmic tasks that can be written as RASP-L program can be easily learned by a Transformer in a length -intth-gnneraliabbl way with next-token prediction. The eength-generalizable tasks include counting, finding the mode, copying the input sequence (consisting of unique tokens), and sorting. However, they also showed that for algorithmic tasks whose RASP-L program representation is not known exist, such as addition, parity, and copying the input sequence, way. In other words, once the Transformer is trained on
is hard to learn a in-distribution data up to particular length, it fails generalize unseen lengths.
2.2 NEXT-TOKEN PREDICTION AND FULL-OUTPUL PREDICTION
Decoder-only Transformers naturally con- venient for next-token prediction (NTP) which could be efficiently trained in parallel. In /Zoou et al. (2024a). their setup and RASP-L solutions are both constrained predicting the single next token: During training, the full sequence (both the query and the answer) provided as input and the output is expected to be the shifted se- quence. During inference, only the query part is provided, and the model continues to output the next token and append the token to the current sequence until the output token is EOS. The out- put locations before the end of query (EOQ) sign ignored. See (a) in Figure for illustration.
Output
Output
model
TF model
0110111
Input
Input
(a) NTP
(b) FOP
Figure Visualization the next-token prediction (NTP) and full-output prediction (FOP) schemes. "#" indicates EOS, "*" indicates ignored output, and ">" indicates the end of the query (EOQ).
On the other hand, we can also consider more general way of predicting the answer: full- output prediction (FOP). During both training and inference time, the input given is just the query part, and the rest of the locations filled with multiple EOS tokens to keep the imput and the output to be the same length. The model is supposed to output the answer with shifted location, and the output locations before the EOQ sign are ignored; see (b) in Figure 2 Notice that in FOP, the model is not forced to predict token-by-toker as NTP. Instead, the model is expected to predict all missing tokens after all internal processing steps.
D-RASP-L
Recall that RASP-L programs do not allow loops. If we consider the next-token prediction (NTP) scheme, it means that we need find the same RASP-L program (which can be represented with a fixed-depth decoder-only Transformer) to predict the next token given any possible prefix in the answer sequence. Such solutions might not always exist for all problems: there no known RASP-L program for addition, parity, and copy under the NTP scheme (Zhou et al. 2024a).
On the other hand, architectures such as the Looped Transformer have external loops embedded in the architecture which naturally provides adaptive depth. Thus, natural question is: what kind of algorithmic tasks can we represent with a decoder-only Transformer in a loop? Specifically, what also allow the number of iterations to explicitly depend the input length, say n2? Moreover, what we are not constrained by the NTP scheme, but more general FOP scheme?
Inspired by these questions, we define the following class of algorithmic tasks:
Definition 3.1 (n-RASP-L). program P is called n-RASP- program if (1) there exist T + N, and (2) P can be decomposed to a sequential application of P' for T(n) steps with a possible pre-pre peepprccssiing step Ppm and post-processing step Ppost: P = Ppr o ( ((r) Ppes where
P', P P Pasr E RASP-I.
3

---

Published as a conference paper at ICLR 2025
ans
110
carry-on
addition
ans
0 0
partial ans carry-on
0 0 0
ans
0
shifted imput
1
cappy
P parity
P
additiion
partial ans shifted input 0 1
0 1
partial ans cerry-on
0 1
0 1
0
1
P
P parity
P
addition
init ans input 1
0
summand 1 summand 2
1
input 0 >
0
1.0.0.00
dpa to ople
spirsoos craced
summard Addition
(a) Copy
(b) Parity
Figure 3: Visualization of the n-RASP-L solutions for Copy, Parity, and Addition with 2. Copy is implemented by iterations of shifting; Parity implemented by iterations of shifting and XOR; Addition implemented by n + iterations of shifted XOR and AND; inputs are preprocessed. details Section
We show that n-ditt addition, n-bit parity, copying symbols indeed have n-RASP-L solutions. Proposition 3.2. (Parity.) There exists n n-RASP-L program with T(n) = n. that solves the n-bit parity check task:
tokens
n'' tolens, n
n tokens
nt' tokens
where y is the parity check result for the arbitrary binary input sequence {xi}.
Proof. See Listing [ in Appendix A where the number of steps required in parity_ is T(n) = n for the input query with n bits.
Proposition 3.3. (Copy.) There exists n-RASP-L prrogrr with T(n) v. that solves the n-symbol copy task:
tokens
vx! tokens, 2
rokens
where {x:} denotes arbitrary binary input symbols.
Proof. See Listing in Appendix[A) where the number of steps required in copy_loop is T(n) for the input query with n symbols.
Proposition 3.4. (Addition.) There exists n--RAP-L progran with T(n) + that solves the n-digit addition task:
11
Z
2/
21
Zm++
tokens
tokens
2n
rokens
tokens
re'
where {xa), {ui} are arbitrary binary summands and { denotes the result of adding (xx) and {ys)
Proof. See Listing in Appendix A where the number of steps required in addition is T(n) = n + for the input summands with n digits each.
We present visualizations of the intermediate steps in the loops of our n- n-RASP-L solutions Figure 3 For the parity task, P Piurity is to shift the input sequence to the right by and calculate XOR of the answer sequence and the input sequence; For the copy task, P Propy is shift the input sequence the right by 1; For the addition task Pr Fddttionn to calculate the XOR of two sequences and shift the results to the right by I position as the partial answer, and calculate the AND of two sequences as the carry-on sequence
2 For simplicity, we include leading 0's to keep the same length the output all possible inputs.
Herr omit the pre-processing and post-processing steps like handling EOS ("u4") and FOQ (">"") tokens which be done by fixed-depth attention layers outside of the loop (see Listings
4

---

Published as a conference paper at ICLR 2025
4 LEARNING 72- 1--AAP-LL PROBLEMS WITH LOOPED TRANSFORMERS
Consider task solvable by n-RASP-L program. It is straightforward to learn the looped Trans- former model with the supervision of the ground truth intermediate outputs: One can use a fixed-depth TF block and simply supervise the input and output for each step. However, such intermediate super- vision can be difficult to get, just like collecting helping CoT steps could be difficult. Therefore, a more interesting setup consider is to learn looped Transformers in end-to-end manner without intermediate-step supervision.
Here we present novel framework for length generalization: In the absence of ground truth CoT data/iitermediite output, we propose to leverage the inherent structure of the problem with the help of "knowing when to stop". We present the setup for training data in Section 4.1 the model architecture and training algorithm in Section 4.2 and the inference algorithm in Section 4.3
4.1 END-TO-END SUPERVISED DATA WITHOUT INTERMEDIATE STEP SUPERVISION
We consider the following settings for the training data and the tasks:
There exists an D-RASP-L program that solves the given task.
Training data consists only of (I, 3) pairs, but not intermediate steps. That is, we do not have access P (x)),, P'(((())
T(n). i.e., the pre-defined number of iterations solve the problem (with some P') is available in the training date
The length is diversely distributed in the dataset, e.g., n E {1,.... (1max) where Tlmax is the maximum number of lengths in the dataset; The pre-defined number of steps needed T(n) is also diversely distributed in the dataset, e.g., T'n) E (T(I), (T where T TTmmaax the maximum number of steps in the datase
4.2
LOOPED TRAINING WITH STEP SUPERVISION
4.2. 1
ARCHITECTURE OF THE LOOPED TRANSFORMERS
We present the architecture for Looped Transformer model in Figure The key characteristics are:
Recurrence: Instead of having a simple stack of blocks, the Looped Transformer is recurrent (like Giannou et al. (2023) but with decoder-only structure) in the sense that we reuse the same decoder block (which consists of a certain number of layers) for number of looped steps, and we adjust the number of looped steps at will.
Input injection:
For each step, the original input sequence is injected together with the output from the previous decoder block, i.e. the input embeddings are added the output embeddings of the previous step as the input of the current step. With input injection, the model can maintain a strong connection the original input, preventing information loss with improved performance (Bai et al.,
2019
Yang et al. 2024).
Notice that there is no positional encoding in the RASP-L operations (Zhou al., (2024a). To follow our n-RASP-L assumption and test the effect of the looped training only, we use NoPE (Kazemnejad et al. 2024) in decoder-only Transformers to avoid the impact from different positional embedding."
Positional embeddings:
4.2.2 TRAINING ALGORITHM
Given a dataset D (((2)) h. tokens, (2n) J. is the output with L; tokens, and T; is pre-defined number of steps of sample i. We aim to learn the transformer model by minimizing the following loss:
L.) where {xxi) J: is the input with L;
'This assumption is to provide supervision for when stop during training: for inference. the pre-defined steps or leverage the confidence the output as stopping criterion (see Section for details.)
either
'The length of the problem is not necessarily the same the actual length of the input due to EOS and EOQ tokens; see Section 1. for the definition of the length of the specific tasks.
'NoPE is shown to inherently learn to use relative positional embeddings in practice Kazemnejad et al.
2024p.
My only handles the embedding space and we use greedy decoding to get the decoded output.
5

---

Published as a conference paper at ICLR 2025
ED|c (fr, (Moo ((((m)) Rd) ( ((0))
where C is the cross entropy loss and (Mo. {(x)) Mt) = Mo(Mo( - Ma Mo(((
Ti iterations
How would our training help to learn more generraiane steps without intermediate supervi- sion? Consider that for some input, want to match the target after T steps, T 2 1. Although we do not have intermediate supervision for this input, there could exist some other input where we match the output after T - 1 steps as shown in Figure which means the model also receives supervision from T - steps with a different input length. The same argument applies to other T values covered in the training set, so the same decoder block receives supervision from not only variable lengths, but also a variable number of steps. As a result, we expect the model to learn length-generrlizabl intermediate steps that could potentially generalize to more steps at test time.
4.3 ADAPTIVE INFERENCE
Recall that looped Transformers can use adaptive depth at inference time, so we need certain rules to decide when stop. In this paper, we consider two rules: 1) Oracle: We can assume that the number of steps needed is given; 2) Maximum confidence: We can use confidence base rules to decide when to stop, i.e., stop when we are confident about the output at the current step. More specifically, for 2), and trained model Me, get the number
given B test sequences with length L: {<x)//1))
of steps T from Equation
T arg m (fo ((2))))))) (OLL)
(2)
where ((ftt ) is the decoded sequences from (M,, (fa)) at step Tmax is the maximum number of steps. Here could be Niet, which is the size of the test set with some specific length. Also, we can choose B = 1 to calculate the per-sample stopping criterion (see Section
5 RELATED WORK
Positional embedding for length generalization. Positional embeddings have been shown to greatly affect Transformers' ability to generalize longer lengths (Press et al. [021 Kazemnejad et al. 2024
Ruoss et al. 2023; Su et al. 2024 Li et al. 2023; Cho et al. [0024] Sabbaghi et al., 2024; McLeish
al. 2024 Golovneva et al. 2024). By designing positional embedding schemes that better capture relative positional information with techniques such as randomization and functional representations, researchers have made significant progress in improving length generalization. Especially, [Coo (t al. (2024) and McLLiish et al.) (2024) use tailor-made positional embeddings for some arithmetic problems without potential generality. (3) This direction is orthogonal to our work since there is no positional encoding in RASP-L operations. We choose no positional embedding in our experiments, but other positional embeddings could further synergistically applied with our approach. However, they might not be expressed as RASP-L operations. We leave further investigation with different positional embeddings to future work.
RNNs and Chomsky Hierarchy. Delétang et al. 2022) conduct extensive empirical study to investigate the limits of the generalization performance of different neural network structures, demonstrate that grouping tasks according to the Chomsky hierarchy allows forecasting whether certain architectures will be able to generalize to out-of of-oiistiibution inputs. Their results show that RNNs and Transformers fail to generalize on non-regular tasks, LSTMs can solve regular and counter-language tasks, and only networks augmented with structured memory (such as stack or memory tape) can successfully generalize on some context-free and context-sensitiee tasks. In our paper, the Looped Transformer architecture also has augmented memory and the recurrent structure but is potentially more powerful since each itcration contains an operation of the whole sequence.
2024), they show that models with weight- layers (but with a fixed depth) can improve the generalization ability when comparing with the variants of the same positional embedding. but they do not find adaptive depths to be helpful since they do not perform the step- specific training as our method, while the key to our method is to use models with adaptive depths. To also compare with this baseline, we add NTP-Loop in 6.1.2
*nn McLeish

---

Published as a conference paper at ICLR 2025
Method
Encoder/Decoder
Prediction Type
PE
Input Injection
Halting Mechanism
UT
Both
NTP
Yes
No
ACT Bolukbast al. 2017
PonderNet
Both
NTP
Yes
No
Halting node
Ours
Decoder-nnly
FOP
No
Yes
Confidence hased or predefined
Table I: Comparison between UT, PonderNet, and ours. PE is short for "Positional Embeddings".
Universal Transformers and other looped models. Our method is highly inspired by Universal Transformers (UT) Dehghaan et al. 2018), but we introduce several novel modifications to design looped Transformers that are compatible with our n-RASP-L assumption. One major architectural innovation is the use of FOP, while all the other prior works are based on NTP. We also only use decoder-only Transformers, which is different from UT and the follow-up work PonderNet (Banino ct al. 202T), which use both encoder and decoder Transformers.
addition to these two critical differences, we do not any positional encoding, and use simpler halting mechanism. Moreover, we find input injection useful to further improve the performance (see details in Section Table summarizes the differences between ours and the previous approaches. Besides architectural differ- ences, we are also the first to show the benefit of using step-dependent supervision for training looped Transformers. Apart from Transformers, [Bansal al. (2022) study learning recurrent networks to generalize to harder maze problems than seen during training, but with a focus on CNNs.
Input representation. Recall that adding two numbers length n could not be solved by a RASP-L program where the difficulty mainly comes from indexing operations (Zhou et al. (0024)). It could be solved by reformatting the input so that each digit is presented to the model with "index hints" in Zhou al. (0224)). Such reformatting enables simple RASP-L program for addition. Similarly, representing the answer in reversed order also helps because the corresponding RASP-L program gets much shorter, providing a concrete justification of the empirical observation made in Lee al. (2024). However, such input representations are highly dependent on the specific problems and might not necessarily exist general.
COT. Seratchpad CoT reasoning Nye al. 2021 Ling al. 2017 Cobbe al. 2021 Wei
et al. 2022) Hoo et al. 2024) is also useful for length generalization as it could simplity the next- token prediction task with intermediate results presented the input layer. There also potential drawbacks and limitations to CoT reasoning. First, CoT training data could be hard to collect. Training and inference with pause tokens (Goyal al. 2023) has been proposed learn implicit CoT steps without CoT data, but pause tokens only increase horizontal compute, not sequential compute. Second, require operations not easily expressible in RASP-L, then CoT may hinder length generalization, as shown in Zhou et al. (2024a). Moreover, CoT steps that could convert the next token prediction task to RASP-L programs might not always exist. Besides, CoT is normally constrained to fixed-depth models, while we study a more general and powerful way to use adaptive compute at inference time.
all CoT steps are helpful. CoT steps introduce additional complexity or
EXPERIMENTS
We evaluate the efficacy of looped Transformers in solving tasks that require length generalization. We introduce the experimental setup in Section 6.I present length generalization results in Section 6.2 and ablation studies in Section 6.3 and visualize the stopping criterion in Section 6.4 Code available at /
www.c//w//wwoeeeeeoo Madi son Lee- Lab/ haobb/loped-tF
6.1
EXPERIMENTAL SETUP
6.1.1 TASKS
Here we consider tasks with 1-RASP-L solutions presented in Section Parity, Copy, and Addition, together with more tasks like calculating the sum, multiplication, and calculating the unique set.
Parity. Checking the parity of the binary string. Example input: output:
example
We define the length of the problem be the number of the digits, set T (the number of steps needed) be the same as the length, and train with length [1, 20).
Copy (with repeated tokens). Copying the binary string. Example input: example output:
We define the length of the problem to be the number of
7

---

Published as a conference paper at ICLR 2025
Parity
Addition
Copy
1.0
1.0
1.0
Accuracy
Accuracy
TOP-Loop Adaptive (Ours)
FOP Loo-Loop-ee Adaptive (Durs)
FOP Adaptive (Durs)
0.4
0.4
0.4
NIP
NTP
NTP-Pause
0.2
NTP-Pause
0.2
NTT.Rause
NTP-Loop
NTP NTP-Loop
NTP-Loop
20
25
30
35
40
20
22
24
26
28
30
20
22
24
26
28
30
32
34
Test Length
Test Length
Test Length
Multiplication
Binary Sum
Unique Set
1.0
1.0
Accuracy
FOP-Loop Loop Adaptiee (Durs)
FOP
FOP FOP- Loop-Adaptive (Ours)
FOP-Locp App-Locp-Adaative (Durs)
0.4
0.4
NTP Pause
NTP Pouse
NTP.Pauuse
NTP
NTP-Locp
11
12
24 Test Length
13
15
16
19
20
21 Test Length
22
23
24
20
22
24
26
28
30
32
34
Test Length
Figure Length Generalization Performance. Our looped Transformer model with adaptive depth generalized better than NTP methods across studied tasks, including the variants with pause tokens and weight-tied layers. The vertical dashed line indicates the maximum training length.
the binary digits to copy, set I" to be the same as the problem length, and train with length [1, 20). has been shown that copy with unique tokens could be easily solved by inductive head (Olsson et al. 2022). but copy with repeated tokens (e.g., binary) does not geigh--gnerrlize with vanilla NTP training (Zhou et al. 2024a).
Binary Addition. Performing binary addition of binary numbers with the same number of digits, and the output has one more digit (without removing leading O if it appears). Example input: example output: -
# We highlight that we do 2024 McLeish al. 2024 Zhou al., (2024b). We define the length of the problem to be the number of digits to be added, set 7' to be the same as the problem length, and train with length [1. 20). It has been shown that binary addition without index hint is hard to generalize in vanilla NTP (Zhou et al. 2024a).
reverse output like recent works (Lee
Binary Sum. Calculating the sum the binary string in the binary form (in reversed order). Example input:
We define the length of the problem to be the number of binary digits to be added, set T be the same the problem length, and train with length (1,20).
# example output
Binary Multiplication. Multiplying two binary numbers, while the first number has up to two digits. The output is in reversed order and the length is the sum of the lengths of two numbers, without removing leading
0. Example input: 1
x
example output:
# We define the problem length be the length of the second number, and set to the product of the lengths of two numbers, and train with length [1, 12).
o
Unique Set. Calculating the unique set with the first occurrence order with alphabet of 50 tokens. example output: We define the length of the problem to be the number of digits to be calculated, set to be the same as problem length, and train with length 20).
Example input:
6.2.2 BASELINE METHODS
Vanilla NTP. We use vanilla next-token prediction as a baseline, referred to as "NTP" in Figure To ensure that the baseline method uses maximum effective depth comparable ours during training, we train the transformer model with a depth 20 times the depth of the looped block in our approach.
NTP with pause tokens. Training and inference with pause tokens (Goyal et al. 2023) is a way to implicitly leam implicit CoT steps without CoT data by enabling extra compute pathways before outputting the answer in NTP. We use it as baseline with the same depth as in vanilla NTP, referred to "NTPP-Puse" in Figure We include visual illustration of NTP- in Figure Appendix [
8

---

Published as a conference paper at ICLR 2025
NTP with weight-tied layers. Using weight-tied layers but with fixed number of overall depths in NTP is also shown to improve the performance in McLLeih et al. (2024). Here we fix the number of looped steps 20, use the same depth as the decoder block of our looped model, and train the model with NTP as another baseline which is referred to as "NTP-LLopp" in Figure 4
6.1.3 TRAINING AND EVALUATION SETUP
For training, we use a decoder-only GPT-2 architecture (Radford et al. 2019). We adopt a curriculum learning strategy for all methods that starts from the smallest length and incrementally increases the length during training till it reaches the maximum length in Garg et al. (2022).
For evaluation, we measure the exact match accuracy for the whole output sequence. For our looped
inference, test two possible stopping criteria discussed in Section 4.3 1) Oracle: Adopt the same
rule when generating the dataset as the number of steps to perform. 2) Maximum confidence: Run a
maximum number of steps, and choose the step using Equation 22 We report test results from 1) in
Section 6.2 and 6.3 and we also find 2) to be an effective stopping criterion in Section 6.4
Full details of training and evaluation are in Appendix F
6.2
LENGTH GENERALIZATION RESULTS
We present the generalization performance on various reasoning tasks in Figure 4
Looped Transformers help with length generalization. Our looped training significantly improves the length generalization performance. For example, for Parity, it can generalize to more than 40 digits near perfectly when trained with up to 20 digits. Moreover, for tasks like addition and copy, where the next token prediction failed when tested maximum training length +10, our looped model can still perform almost perfectly.
Variants of NTP could improve generalization but not as effectively as our adaptive-depth model. Compared with vanilla NTP, we observe that NTP-Loop could lead to improved generalization in tasks like Addition, Copy and Multiplication. Similarly, NTP-pause could introduce slight improvement in Binary Sum and Unique Set. However, they all fall behind compared with our method.
6.3 ABLATION STUDIES
In Section 6.2 we compare with NTP baselines while the efficacy of components in our architecture design remains unclear. In this section, compare with FOP variants of model in Figure 7 (Appendix -od-LLop-Adaptive-- WO" indicates our method but without input injection; "FOP- Pause" indicates FOP with pause token "FOP" indicates vanilla FOP without weight-tied layers and adaptive depths.
Effect of input injection. We observe the generalization performance with input injection is generally better than without it, which aligns with the findings in (Bai et al. (2019) and Yang et al. (2024). The effect of input injection more visible tasks like Addition, Binary Sum, and Unique Set.
Comparison with pause tokens and vanilla FOP. Training with pause tokens in FOP could sometimes boost the generalization performance compared to vanilla FOP, but not effective as our method with looped steps and adaptive depth. As discussed in Goyal et al. (2023), pause tokens mainly introduce parallel but not sequential compute, which is less powerful than adaptive depth. Besides, we find worse in in-distribution accuracy for both FOP and FOP-Pause in Addition, which mainly comes from the difficulty in training deep model x20x the depths of the decoder block used in the looped model) in FOP. It further highlights the importance of supervision with variant depths used in our training.
6.4
THE STOPPING CRITERION AND VISUALIZATIONS
In this section, we visualize the accuracy and the cross-entropy loss with respect to the decoded output in each iterative step across tasks in Figure We also provide more visualizations from other test lengths in Appendix B The vertical lines in Figure 5 are chosen based on the full test set using
Another option is to set a threshold for cross-cntropy loss and stop when the threshold first met. This will also succeed if the maximum confidence rule works.
Visual illustration of FOP-Pause is in Figure in Appendix

---

Published as a conference paper at ICLR 2025
Accuracy
Accuracy
Accuracy
Step
Step
Step
Log Loss
Log Loss
Log Loss
-2.5
-00
-5.0
20
-75
20
60
20
Step
Step
Step
(a) Parity (test length=40)
(b) Addition (test longth=30)
(c) Copy (test length=35)
Accuracy
Accuracy
Accuracy
0
Stop
Log Loss
Log Loss
Log Loss
-10
-5
-10
60
30
Step
Step
Step
(d) Multiplication (test 1ength=16)
(e) Binary Sum (test length=24)
(f) Unique Set (test 3ength=35)
Figure 5: Stopping criterion visualizations. Plot of the stopping criterion on the test set. The vertical line indicates the step chosen from Equation ( (where B = Nus equals the size of the test set) within the step range shown in the plots. The accuracy is on the full test set. The chosen steps have accuracy z across tasks.
Equation 2 with B Nie given the specific length. We also include the full test results with B 1 and B= Mest from all test lengths in Figure [ Appendix B
Convergence in Addition, Copy, Multiplication, and Unique Set. We notice that for Addition, Copy, Multiplication, and Unique Set, the looped model somehow learns to converge for a certain number of steps after solving the task, even though we do not explicitly train the model to converge. The loss curves for these tasks are also smoother than those without convergence behaviors.
The maximum confidence stopping criterion. In Figure 5 the cross entropy loss (with B = Ness) reaches the lowest when the generalization performance is near perfect at the given test length, which indicates the maximum confidence rule chooses the right time to exit. We also show the full results in Figure (Appendix BB. The convergence nature of the tasks also affects the performance of the per-sample-baas based stopping criterion (B- = 1) in Figure (Appendix B): the per-sample stopping criterion performs almost as well as when using B = Niest to decide for converging tasks, but the performance is worse in non-converging tasks.
[
7 LIMITATIONS AND CONCLUSION
Our current definition of n-RASP-L does not support tasks that require multiple loops followed by each other. Thus, one important future direction would be extending our current definition to support multiple loops and identifying larger class of tasks that can be implemented only with multiple loops but not with a single loop. For training, direct looped training could be computationally demanding when the number of looped steps is too large. A possible workaround for more efffccinnt training could be stopping the gradient tracking for earlier steps like Clark et al. (2023). but there might be trade-off in performance and computation. We only train the looped Transtormers for limited number of steps and lengths due to a lack of computing resources. With more diverse training data, the looped model has the potential to generalize to even longer test lengths. We use NoPE for simplicity, and an orthogonal direction is to use more delicate positional embedding to further improve length generalization performance. Moreover, our step-dependent supervision requires the ground-truth truth number of steps in the training data, which is an additional requirement compared with normal end-to-end training. However, we still require fewer assumptions than CoT training.
In conclusion, we show that R--RSP-L problems can be learned by looped Transformer with step- dependent supervision on the final answer, and can be applied with an adaptive number of steps during inference time to improve generalization. Note that ---PSPLL problem set, could cover more challenging reasoning problems than presented in the paper, and we believe that our method has the potential to generalize to more challenging tasks.
challenging algorithmic
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The works is supported by NSF Award DMS-2023239, NSF CAREER Award CCF-2339978, Amazon Research Award, and a grant from FuriosaAI.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 35:38546-38555 2022.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances neural information processing systems, 32, 2019.
Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. In 8th ICML Workshop on Automated Machine Learning (AutoML), 2021.
Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances Neural Information Processing Systems, 35:20232-200 20242, 2022.
Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In International Conference on Machine Learning, pp. 527-536. PMLR, 2017.
Hanseul Cho, Jaeyoung Cha, Pranjal Awasthi, Srinadh Bhojanapalli, Anupam Gupta, and Chul- hee Yun. Position coupling: Leveraging task structure for improved length generalization of transformers. arXiv preprint arXiv:2405. 20671, 2024.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309. /7400, 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems, 2021. URL hupp::/arrsv. org/ubw21110.// 1416,, 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint
2018.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, et al. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.02098, 2022.
Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference on Machine Learning, pp. 5570-5582. PMLR, 2022.
Yilun Du, Jiayuan Mao, and Joshua B Tenenbaum. Learning iterative reasoning through energy diffusion. In Forty-first International Conference on Machine Learning. 2024.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers. In International Conference on Machine Learning, pp. 1 11398-114 11422 PMLR, 2023.
Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Contextual position encoding: Learning to count what's important, arXiv preprint arXiv:2405.188 187.1 2024.
11

---

Published as a conference paper at ICLR 2025
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint vrX1v:2310.02226, 2023.
Kaiying Hou, David Brandfonbrener, Sham Kakade, Samy Jelassi, and Eran Malach. Universal length generalization with turing programs, arXiv preprint urXiv:2407.03310, 2024.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35: 22199-22211 2022.
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024. URL
Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang. Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. arXiv preprint arXiv:2310.04418 2023.
Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale genera- tion: Learning to solve and explain algebraic word problems. arXiv preprinl arXiv:1705.0414 2017.
Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, et al. Transformers can do arithmetic with the right embeddings. arXiv preprint arXiv:2405.11739,, 2024.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114 14, 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint 2rXiv:2209.1189) 2022.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2 /2409, 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint arXiv:2305. 16843, 2023.
Mahdi Sabbaghi, George Pappas, Hamed Hassani, and Surbhi Goel. Explicitly encoding structural symmetry is key to length generalization in arithmetic tasks. arXiv preprint a:Xiv:2406.01895, 2024.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-2444 2022.
12

---

Published as a conference paper at ICLR 2025
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pp. 11080-11090. PMLR, 2021.
Zhaofeng Wu, Linlu Qiu, Alexis Ross, Ekin Akyürek, Boyuan Chen, Bailin Wang, Najoung Kim, Jacob Andreas, and Yoon Kim. Reasoning reciting? exploring the capabilities and limitations of language models through counterfactual tasks. arXiv preprini arXiv:2307.02477 2023.
Liu Yang, Kangwook Lee, Robert D Nowak, and Dimitris Papailiopoulos, Looped transformers are better at learning learning algorithms. In The Twelfth International Conference on Learning Representations, 2024.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M. Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. In The Twelfth International Conference Learning Representations, 2024a. URL
htttps://openreview.. net Aow.umidddAssuuuiiiii...
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou, Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024b.
13

---

Published as a conference paper at ICLR 2025
A n-RASP-L PROGRAMS
Here we provide the n-RASP-L programs for our Parity, Addition and Copy tasks in Listings | We also present the RASP-L library functions we Listing which is partially taken from Zhou et al. (2024a).
Input example: 1
4 #
output example:
indicates ignored token, is E0Q, and # BOS.
def (partiiil_
seq):
align the last digit with the answer location seq shift_right (seg, 1)
calculate XOR
partial_ans_see partial_ans_sent seq) & (>(partiaal_ns__ee) seq)
return partial ans seq, seq
def parity_loopp (seq, num :
get the question in the prompt
prompt_mask seq mask prompt_mask)
has_scen (seq, full (seg, B
# init answer seq with 0
parttal_ans_sst = full (seq, generate EOS seq after BOQ
seq where (promppttasssseeeeeeeeaeppppeeee perform parity steps
full, 0), full(seg, BOS)
for i in range cum_steep)=
partial_ans_sed get answer with 8OS
seq (pratt_st seq)
ans_seq partial_ans_eet
seq shift_rights (end_seq, seq, 1)
ans_sed _ss_eeq where (end_ end end_seq, seq, ans_seq) return ans_seq
Listing I: Parity.
Input example: 0101100 Output example: * 0 # indicates ignored token, 1s ROQ, and K08.
def copy_step =nd_seq): seq):
seq = shift shift_right (seg, 1)
end_seq shift_right send_seq, 1)
return seq, end_seq
def copy_ [cy_lccpssee, num_step):
generate EOS seq after B0Q
end_mask has_seen (seg, full (seq, E0Q)
enc_seq where (end
(seq, 0), full (seq, ECS) )
perform copy steps
for i in range (num_step)
seq, end seq copy_step (seq, end seg)
get answer with 80S
seq where (end_ return seq
seq, seq)
Listing 2: Copy.
Input example: 0.0
Output example:
indicates ignored token, > is E0Q, and # is BOS.
def itt_tiionnsepsseeeeeeee seq?, end_seq)
14

---

Published as a conference paper at ICLR 2025
end seq shift right (end seq,
seel np. (seql, atype=bccl)
see2 = np.array (seq2, dtype=bool)
carry_on = seql & seq2
$ A XOR B = (A OR B) AND (NOT (A AND B))
in_place ((seq] seq2)
(seq] seq2)))
in_place shift_right (in_place, 1)
seel = ar. (in_pace, atype=int)
see2 = ap.array (carry_on) atype=int)
return seql, seq2, end_seq
def adcition__reproeess (seq):
generate EOS seq after EQQ
end mask = sas_seen full (seq, E0Q1
seq where (end mask==0, sull(seq,
full (seg, ECS)
generate masks for the first and second summands secn_tckk has_scon (seg, full (scq, ADD_SIGN)
secn_tok1 = has_scen (seq, full (seq, BOQ))
maskl "seen_tokO
mask2 seen_tokō (~seen_tckl)
mask2 mask2 shift_right (mask2, 1)
get the first and second summands
seel = mask (seq, maskl)
seq2 mask (seq, mask2)
align the first summand with the second induct_numl cumsum (maskl)
induct_num2 = cumsum (mask2)
target_index firsts (induct_numl, induct_num2, default-0)
seel index_seles select (seql, target_index)
seel mask (seql, mask2) return seql, seq2, end_seq
def additicn (seq, num_step);
seql, seq2, end_seq = additicn__reerroesss (seq)
perform addition steps
for in range (num_step):
seql, , seq?, end_seq addition (seq), seq?, end_seq)
get answer with 8OS
ans seql
ans = where -eeedddsee=-O,,, end seq, ans)
return ans
Listing 3: Addition (in forward order).
import numpy as np
def full (x, const):
return ap.full_like(s const, dtype-int)
def indi cesicesii (x):
return np.arange (Len (x), dtype=int)
def select (k, q, pred, causal=True):
# compute attention matrix
len (k)
np.zeros (ss, s), dtype-becl) for gi in range
for kj in (range (qi+1) if causal else range (s)
# i_ide <= indde if causal
Alqi, kj]
pred pred (k[kj), q(qi))
return
def sel (B)
return no. (A, np.ones (1en (A) (intt
15

---

Published as a conference paper at ICLR 2025
def aggr_ mean (A, v, default=0): out np. (A, v)
norm = sel sel_width
out np. np.divide
out, norm, out=np. out=np.full where=(norm 0)
default, dtype=float)),
return out, astype (int)
def kgy (k, q, v, pred, default=0, reduction='mean'
return aggr_mean(
select (k, 9 pred), v, default=default, reductionnredustion
def sec map Yr fund):
tokenwise map over two sequences
return np.array [func (xi, yi) for xi, yi in (x,y),)
def shift_right (x, n, default=0):
# shifts sequence x to the right by n positions (other positions filled with default)
return kqv (indices (x) +n, (xdices(x), x, equals, default=defawit)
def where (cendition, x if, y_else):
equivalent to np. wherh (condition, x_if, y_else)
x_masked - seq_map (x_if, condition, lambda m: if n. else 0) y_masked seq_map (y_else, condition, lambda y, y if not else 0) return seq map_map (x masked, y m_mased, lambda x, y: if 0 else y)
def has_seen (x, queries):
return (x, queries, full 1), equals, default=0)
def firsts (x, queries, default=-1);
find the index of the first occurrence of each puery[i] in x return kqv(
queries, indices (x), equals, default-default, reduction='min!
def mask (x, bool mask, mask v equivalent to x=bool_mask + default* (-bool_mask)
return where (bool_mask, x, full (x, mask_val))
Listing 4: Library functions from Zhou et al. 2024a).
B RESULTS FROM THE ADAPTIVE STOPPING CRITERIA
Here we present the length generalization results in Fugire 6 with different stopping criteria: Oracle 2) Maximum confidence based on the full test set using Equation 2 with B = Nus given the specific test length and 3) Maxumum confidence using Equation ? where sample (we still calculate the accuracy across the full test set). For tasks with converging behaviors (Addition, Copy, Multiplication, Unique Sct), per sampler stopping criterion works almost as well using B Nisa while for tasks without the converging behavior (Parity and Binary Sum), per sampler stopping criterion does not work as well as using
calculated per
N
Besides, also visualize another baseline which is using FOP with looped structure but with fixed number of steps, which is denoted as "FOP-Loop-Fine "ixed"" Figure 6, showing that using variable depth generally outperforms fixing the depth of the model.
16

---

Published as a conference paper at ICLR 2025
Parity
Addition
Copy
1.0
1.0
1.0
Accurayy
Accuracy
Loo-Loo--ar Adspiive
FOP Loop Adaptive
FOPLocp- Adsptive
0.4
0.4
0.4
FDP- AoM-LogpAdepiiv--iee Adaptive-Batch
hep-LaoopAAapttvee-nne nsttanee
FDP wew.ooooeodpptll--nnttnn
Loo-Loop- Rixdd
FOP Loop- Rxed
LOP- Fxnd
20
25
30
35
40
20
22
24
26
28
30
20
22
24
26
28
30
32 34
Test Length
Test Length
Test Length
Multiplication
Binary Sum
Unique
1.0
1.0
1.0
Accursy
FOP FOP-Loop Loop-Adaptive
FOP Laop- hopp-ddative
FOPLoc--Adsptoog
0.4
Saap-AAdativv-Bttht
0.4
FoP-Laop-Adapee Saap-Adaptiv--Batte Batch
0.4
AiPLLog-Adiitiv--att
LoP-Loop- Adaptive instance
FDP Laop Adaptive Instance
FOP Adaptive Instance
FOP Loop fixed
Loop Fixed
Lopp Fixed
11
12
13
14
15
16 19
20
21
22
23
24
20
22 24
26
30
32
34
Test Length
Test Length
Test Length
Figure 6: Effect of different stopping criteria. Loop--oope-AAaatiiee indicates stopping with the pre-defined
number of steps: 1o0P-LLoped-Aapppe Adaptive-Batch" Batch" indicates stopping with Equation (2) where B is the size of the test set; "FOP-l.opeedAAaaple -ati--ddpptiienttaacee indicates the number of steps is chosen per sample, with Equation
where B 1 (the accuracy is still calculated across the dataset).
C ABLATION STUDY
We present the generalization performances of our method compared with FOP variants in Figure as mentioned in Section 3
Parity
Addition
Copy
0.8
0.8
0.8
Accuraty
Accuracy
0.6
0.6
0.6
Loop-Adaptive (Ours)
FOP LoP-Loop-Adappie (Durs)
FOP-Leop Adeptive (Durs)
FOP-1 oop--Aaaptive--N
hop--Loop-aaptiiveeee
FDP- MoP-Locp-daative-WWW
0.2
FOP-Fause
0.2
FOP-Pause
0.2
FOP-Pause
0.0
0.0
0.0
32
Test Langth
Test Langth
Test Length
Multiplication
Binary Sum
Unique Set
1.0
0.8
0.8
Acorody
Acoracy
0.9
Accuncy
PoP.Loop-Adaptvve (Durs)
0.8
FOF-Loop-Adaapiivn (Durg)
FOPLech-Adaptve
WO
FOP FOP-Loop-Aee Loop-Adaptive-WW
LOP Wo
0.2
Pause
0.7
Pouse
0.2
FOP-Peuse
FOP
FOP
0.0
0.6
0.0
21
Test Length
TASt Length
Test Langth
Figure 7: Ablation study. Our looped Transformer model with adaptive depth generalized better than FOP variants across studied tasks, including the variant of method without input injection, and FOP with pause tokens. The vertical dashed line indicates the maximum training length.
D VISUALIZATION OF USING PAUSE TOKENS IN NTP AND FOP
We visualize NTP-Pause in Figure and FOP-Pause in Figure respectively, where add fixed number of pause tokens (3 in the figures, 20 in our experiments) before outputting the final answer during both training and inference.
17

---

Published as a conference paper at ICLR 2025
Table Task-specific experimental hyperparameters. "Number of Heads" and "Block Depth" define the size of the looped decoder block. "Interval" denotes the number training steps between successive increases in input sequence length.
Task
Number of Heads
Block Depth
Interval
Parity
64
500
Copy
2
1000
Addition
3
1600
Multiplication
8
500 4
Binary Sum
16
2
500
Unique Set
3
1000
INFERENCE TIME COMPLEXITY
Here we present the inference time complexity for our method, vanilla NTP and vanilla FOP.
Assume that the maximum length of the training set is the number of steps needed is T(2). and the number of layers each step Assume that NTP and FOP are using fixed number of layers C. And we test on length n!.
For the first stopping criterion where we know a(n'), our inference time would be O(a((n' ((ka()")n"?) and NTP (with KV cache) and FOP will be OCOn?). For the second criterion, we need to specify the maximum number of steps in order to find the step with maximum confidence. So our inference time would be OKkN'n where N' is the maximum number of steps.
In NTP and FOP, we use some Ra (n) in our experiments such that they use similar compute during training. Our inference time is then slightly longer than NTP with KV cache and FOP since we use more steps than the fixed-depth models.
Output
Output
Output
TF model
model
TF model
O
0 #
0.1.1P
1
Input
Input
Input
Figure NTP-Pause visualization. Examples are from the Copy task. ".. indicates the pause token.
Output
Output
Output
model
model
TF model
O 1
Input
Input
Input
Figure 9: FOP-Pause visualization. Examples from the Copy task. indicates the pause token.
F EXPERIMENTAL DETAILS
Model architecture. We use the decoder-only GPT-2 model with NoPE and 256 embedding dimensions the basic block for the looped iterations with task specific configuration in Table We first convert the input to the embedding space, perform the looped computation in the embedding space, and then decode the final output. For input injection,
similar technique Yang et al.
18

---

Published as a conference paper at ICLR 2025
(2024) that adds the original imput embedding each iteration as part of the input. For vanilla NTP, we adopt the same training scheme, but trained with autoregressive loss instead. For WTP--Puue and FOP-Pause add 20 pause tokens before outputting the final answer. For FOP-Looped-Fixe -ixeed fixed the model depth to be the maximum depth used in -OP-Looppe-AAaapiiee
Training and evaluation details. For the training distribution, we adopt the online training scheme following Zhou et al. (2024a) where each batch is i.i.d. sampled. Given any length, the probability of each possible character is evenly distributed instead of from finite train set to avoid over wver-fitting, and the length is also evenly distributed. We also
curriculum gradually increase the maximum training length (see Table 2 for the specific setup for each task). We use Adam W optimizer and decay the learning rate from 10-4 to 0 with cosine decay schedulers after reaching the maximum training length with batch size 64, and train for a total number of 100k gradient steps. Additionally, non-converging tasks are less tolerant when choosing which step stop, and we find using the exponential moving average (EMA) of model parameters helpful for Parity and Binary Sum with factor 0.9999. For evaluation, we test with 6400 random samples in Figure
6 and report the mean exact match accuracy and standard error from five training runs with different random seeds.
19


Reviews:
[{"review_id": "zdcgZb1Jue", "reviewer": "Reviewer_gBs4", "summary": "- This work studies the efficacy of Looped Transformers for Length Generalization of several algorithmic tasks whose computation complexity is known (as a function of the query length).\n- The paper proposes the definition of $n$-RASP-L, a generalization of the RASP-L computation model allowing the loop of RASP-L programs. It is shown, under a general framework called full-answer prediction (FAP), that some tasks (Copying binary sequence (allowing duplicates), Parity, and Binary Addition) have their own $n$-RASP-L program with a linear number of steps in problem length.\n- The authors propose training Looped Transformers (with input injection and curriculum learning) to learn $n$-RASP-L-programmable tasks, where the ground-truth number of steps is known for each task during training. They also propose two variants of inference methods: either we retain the knowledge about the number of steps at inference time (*Oracle*), or we adaptively decide the number of iterations based on the confidence of FAP (*Maximum confidence*).\n- The proposed method is tested on several algorithmic tasks.", "questions": "**Q1. Question on the visualization in Figure 3**\n\n- Why don’t the illustrations in the figure contain any “#” (EOS) tokens? Is it due to the pre-processing?\n\n**Q2. Do the trained Looped Transformers simulate the $n$-RASP-L program?**\n\n- Although it might be difficult to reverse-engineer a trained transformer model to figure out what algorithm it actually simulates or implements, it might be interesting if we can observe any kind of similarity between it and the $n$-RASP-L program.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1. The paper is written and organized well. Overall, the presentation of the methodology and empirical results is clear and easy to follow.\n\nS2. The idea behind the proposed method is neat and plausible. It is natural to think about adaptively scaling the depth of the model according to the problem length or the problem complexity. This paper successfully implements this idea to solve various interesting algorithmic tasks with the power of Looped Transformers. Also, $n$-RASP-L is an interesting but intuitive generalization of the RASP-L framework by allowing the loops. \n\nS3. The proposed answer-generation framework called FAP is also an interesting component of this work. It might be of separate interest to study.\n\nS4. The paper presents extensive ablation studies on several components of the proposed method. Also, the empirical results (length generalization performances) are impressive enough to convince the readers about the proposed method’s efficacy.", "weaknesses": "**W1. The definition of $n$-RASP-L (Definition 3.1) can be improved.**\n\n- I think the equation “$T(n): \\mathbb{N} \\rightarrow \\mathbb{N}$” should be corrected to “$T: \\mathbb{N} \\rightarrow \\mathbb{N}$” because $T$ (instead of $T(n)$) is a function of input length $n$ representing the number of steps inside a task-solving $n$-RASP-L program.\n- In (2), I guess $P’$ should be a RASP-L program, which is unspecified in the definition.\n- Should $P$ be decomposed to a sequential application of $P’$, i.e., $P = (P’)^{T(n)}$? I don’t think this is exactly true because there are pre-/post-processing parts inside the proposed $n$-RASP-L programs (in Appendix A). Can the same RASP-L program $P’$ handle such parts? (It might be true because of the experimental results, but I cannot fully understand this part.) If not, I guess the definition should be modified to include the pre-/post-processing parts. For example, $P = P_{\\tt pre} \\circ (P’)^{T(n)} \\circ P_{\\tt post}$.\n\n**W2. “Ground truth” number of steps?**\n\n- According to Definition 3.1, a program $P$ suffices to be an $n$-RASP-L if a corresponding $T(n)$ exists. Indeed, Propositions 3.2, 3.3, and 3.4 claim and prove the existence of $T(n)$ for the Parity, Copy (with duplicates), and Binary Addition tasks, respectively.\n- My question is about the uniqueness or optimality of such $T(n)$’s. There might be a clever way to construct another RASP-L program $\\tilde{P}$ so that $P$ can be implemented with $\\tilde{T}(n)$ steps of applying $\\tilde{P}$, where $\\tilde{T}(n)$ is much smaller than the previously known $T(n)$ (e.g., $\\tilde{T}(n) \\in o(T(n))$). It may happen since there is no uniqueness guarantee or lower bound result on $T(n)$.\n    - If I venture a guess, I would say it might be possible to implement an $O(\\log n)$-step $n$-RASP-L solution for the Parity task by using the parallelism of the transformer architecture. Please correct me if I am wrong. Also, I understand if it is impossible to show whether this bold guess is true. If you are interested, there are some (probably) useful references about logarithmic-depth transformers [1,2].\n- However, the authors keep using the phrase “ground truth number of steps” throughout the paper, which may lead to misunderstanding that the only way to implement the given $n$-RASP-L program is by using a loop of length $T(n)$.\n- If two different $T(n)$’s can be applied to a single $n$-RASP-L-programmable task, it might be interesting to observe whether the model’s performance changes depending on the choice of $T(n)$.\n- Furthermore, if multiple choices of $T(n)$’s exist for a given task, does knowing only one of them suffice to train reasonably performant Looped Transformers? If we know more than one, how should we choose $T(n)$ when we train the model?\n\n**W3. Shouldn’t we consider the input injection when implementing an $n$-RASP-L program for the given task?**\n\n- The input injection seems to be an important component of their experiments. Since it changes the input vectors of each layer, I guess the task-solving algorithm under input injection might be different from that without it.\n- However, I can’t see that the $n$-RASP-L programs provided in Appendix A reflect the input injection. As I inspect inside the loop of each program, every iteration only reuses the calculation(s) from the previous iteration right before the current one.\n- Shouldn’t we consider the very first input sequence and the result from the previous iteration when implementing the loops? Or is it a valid implementation of input injection? Getting even further, Is there any way to embed the input injection into the $n$-RASP-L programs?\n\n**W4. The proposed training method requires prior knowledge of the task’s structure.**\n\n- The proposed method is limited in that it requires a prior understanding of the structure (e.g., $T(n)$) of the task where we want to train a model. This is because it hinders fully end-to-end training.\n- Are Looped Transformers still useful for achieving length generalization even when we don’t (or cannot) know the exact expression of $T(n)$?\n- Besides, it seems that the depth of the decoder block is determined based on the complexity/difficulty of the subroutine $P’$ at each step inside the loop (Appendix F). How are they actually chosen? Or, how should we decide the size of the repeating decoder block?\n\n**W5. Some experimental details seem missing or wrong.**\n\n- I guess Equation (2) has a typo: shouldn’t it be arg-main instead of arg-max?\n- In Binary Addition, it seems that $T$ is chosen to be $n$ (the length of each operand). However, Proposition 3.4 claims that $T(n)=n+1$ for the same task. Why is there a discrepancy between theory and experiment?\n- In Binary Multiplication, I guess some words are used in a wrong way. In Lines 417-418, I think it should be: “We define the problem length to be the **length** of the second **number**, and set $T$ to be the product of the lengths of two **numbers**.”\n- In Section 6.1.2, are input injections also applied to NTP-based methods? Also, I’m not sure why it is fair to compare their method (based on FAP) to NTP methods with the architectural setting “…with a depth 20 times the depth of the looped block” because such depth might be suboptimal for NTP-based methods.\n- Although the paper indirectly showcases that their adaptive decision of the number of steps works quite well via Figure 5, it would be better to display similar performance plots to Figure 4 (plots based on the “Oracle” inference) but using the adaptive confidence-based method instead, at least in their appendix.\n\n**W6. Minor writing issues**\n\n- Section 4.1, fourth bullet point: I guess $T(n) \\in \\\\{T(1), \\ldots, T(n_{\\rm max})\\\\}$ is correct ($T(1)$ instead of $1$).\n- Equations (1) and (2) have several weird-looking brackets (too many open brackets etc.)\n- Line 510: Use *less* abbreviations like “w.r.t.”\n\n---\n\n**References**\n\n[1] Sanford, Clayton, et al. \"Transformers, parallel computation, and logarithmic depth.\" ICML 2024.\n\n[2] Sanford, Clayton, et al. \"Understanding transformer reasoning capabilities via graph algorithms.\" NeurIPS 2024.", "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": "ACHZBffFKW", "reviewer": "Reviewer_3KX9", "summary": "This paper investigates the length generalization problem of Transformer models, which refers to the inability of the model to deal with longer samples than encountered during the training phase. While recent literature has focused on modifying the positional embeddings and the input formats, this paper proposes to use Looped Transformers, which can dynamically adjust their computation steps according to the problem length. The authors define n-RASP-L problems to figure out which problems can be solved by Looped Transformers. Then, they train the models on these tasks (parity, copy, binary addition, binary sum, binary multiplication, unique set) under a full-answer prediction setup. Empirically, the trained models could successfully length-generalize to longer lengths by appropriately adapting the number of loops at inference time.", "questions": "Q1. In Figure 5, why do some tasks perform well even when exceeding the step count, while others degrade immediately? For instance, the performance of the parity task and the binary sum task immediately drops when executed with additional steps, whereas the addition, multiplication, and copy tasks retain accuracy to some extent.\n- Particularly for the copy task, the selected step count is significantly higher than the actual number of steps required, which seems unusual to me.\n\nQ2. Are there any tasks whose T(n) is nonlinear (e.g. sqrt(n), n^2) to the length of the input sequence? It would be interesting to see experimental results for such tasks.\n\nQ3. Why is the output reversed for binary multiplication (but not for binary addition)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-structured and clearly written.\n- The introduction of Looped Transformers is well-motivated and effectively argued.\n- The results are strong and solid. They do not require the use of a scratchpad. Also, the prediction is conducted using an end-to-end, full-answer prediction setup, which is a more general way than the conventional next-token prediction setup.\n- The paper clearly illustrates that the model can determine the number of steps to take on its own and does not require T(n) in the test time.", "weaknesses": "Weakness 1: Applicability Limited to n-RASP-L Tasks\n\n- The approach is limited to tasks that belong to n-RASP-L categories, as it requires the ground-truth number of steps in the training data.\n\nWeakness 2: Insufficient Experimentation.\n\n- ***Effect of Curriculum Learning.*** How does the model perform without curriculum learning? Is the use of curriculum learning necessary?\n\n- ***Tolerance to Step Counts.*** I am curious whether this method will still perform well with different choices of T(n). For example, for tasks like parity, would the model maintain its performance if T(n) were set to n+1 rather than n? What about 2n instead of n? This question stems from the possibility that there might be more efficient solutions to n-RASP-L problems than human-designed ones, which could work with fewer steps. Testing whether the model is robust under overestimated T(n) values could help verify the robustness of this approach.\n\n- Overall, the paper requires more ablation studies.", "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": "7H6Li5UrrK", "reviewer": "Reviewer_imEP", "summary": "This paper examines how looped transformers perform in terms of length generalization. The focus is on n-RASP-L problems, which are problems that can be tackled using a loop of a single RASP-L program. The concept is that Transformers can learn steps that are independent of length, employing a flexible number of iterations in the looped transformer to achieve length generalization. The authors first demonstrate that n-digit addition, n-bit parity, and copying n symbols can be addressed with n-RASP-L solutions. They then reveal that when utilizing the looped transformer with adaptive stopping time, the results exhibit significantly stronger length generalization compared to next token prediction (NTP) and other methods like using pause tokens or NTP-loop with a fixed stopping time.", "questions": "I listed my questions in the weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, I really liked the paper, I think that using a looped transformer to achieve length generalization is an interesting idea that was not studied in the past to my knowledge. This paper complements all the other techniques (universal transformers, different types of position emebedding, etc.) that were used in the past for length generalization The paper is well-written and well-explained. This is why I advocate for acceptance of this paper.", "weaknesses": "I would like to raise the following weaknesses/questions regarding this paper: \n\n- **Lack of other baselines**: What would happen if you have a very deep universal transformer? Universal transformers also have shared parameters and looks equivalent to the loop transformer. The depth may play the role of the number of loops. Would this be equivalent to the fixed loop NTP? It would be interesting to run the same experiments with a universal transformer.\n\n- **Comparison with other methods**: Where would you position the looped transformers in the list of all the tricks for length generalization? Are the effects similar or complementary to change of the input (index hinting, reverse order of operands, etc.) ? Changes of positional encoding? Chain of Thought? It would be interesting to understand this by making combinations of the tricks with looped transformers with other tricks and analyze the performance differences.\n\n-  What is the depth of the encoder block in the loop transformer? I think this information is important to put in the main paper. \n\n- **Adaptive inference time**: I think one weak point of the method is actually coming up with an adaptive inference time. The methods that are proposed are nice but may look a bit hacky. Do you think one could learn this adaptive inference time?\n\n- In Figure 2, which adaptive inference time method is used for FAP-Loop-Adaptive?\n\n- Lastly, this is a wild question: have you tried your method on problems where there is no n-RASP-L solutions?  Would it still work better than just doing NTP?", "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": "czmRbrFet6", "reviewer": "Reviewer_Fhkc", "summary": "Empirically explores the ability of looped Transformers, i.e. Transformers that repeatedly apply the same block of layers, to length-generalize on several algorithmic tasks, including copy, parity, and addition. First, the authors manually derive length-generalizing solutions to the considered tasks through a variant of the RASP language, which they term n-RASP-L. Then, based on these ground truth solutions, they show that looped Transformers length-generalize well when trained with access to the true number of steps required to compute the output for a given input.", "questions": "1. Are the quantities reported in Figure 5 indeed for a single training example? When using the maximum confidence criterion, how do the results compare to the ones reported in Figure 4 with access to the ground truth number of steps?\n\n2. In Bansal et al. 2022, they avoid the need for knowing the exact number of steps during training and inference. Have you tried using similar heuristics?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. The paper is mostly well-written and easy to follow.\n\n2. Demonstrates that, given knowledge of the number of steps required to perform a given task, a certain looped Transformer, which jointly predicts the full output sequence, tends to learn a length-generalizing solution. The length-generalizing capabilities of this looped Transformer are shown to surpass baselines that use next-token prediction.", "weaknesses": "1. The main weakness of the current paper is that the significance of the results is somewhat limited. In particular, I find that it falls somewhere in between works that are practically relevant and those that may not be practically relevant, but improve our understanding of certain phenomena. On the one hand, this work shows empirically that, in some algorithmic tasks for which we already know how to write explicitly a length-generalizing solution (in terms of looped Transformer weights), looped Transformers generalize well to longer lengths, if they have access during training to the number of steps required for solving the task for a given input. Consequently, the practical relevance is limited since the proposed method requires that we already know how to manually write a length-generalizing solution, in which case there is arguably no point in learning. On the other hand, this work does not provide much in terms of understanding why or how looped Transformers are able to length-generalize.\n\n    Note that one may consider the demonstration of such length generalization to be possible as a main contribution. Yet, the ability to extrapolate through recurrence of layers has been demonstrated in the past, albeit for other architectures (see Bansal et al. 2022 [1], which notably do not require knowing the ground truth number of steps in training).\n\n2. A related issue is the usage of ground truth stopping time during inference. The quantities reported in Figure 5 seem to be for a single training example, yet it is not entirely clear. If so, then how does the maximum confidence stopping criterion fair across the dataset? It would be useful to report results similar to those of Figure 4 but when using the proposed stopping criterion as opposed to the ground truth stopping time, which should be unknown.\n\nOverall, my assessment of the paper tends towards the positive side, yet it is not a clear accept due to the substantial limitations mentioned above. Specifically, the significance of the contributions can be greatly improved if it would be possible to remove the dependence on knowing the ground truth number of steps required to solve the task for a given input during training (and by how it seems from the current results, during test time as well).\n\n\nAdditional (more minor) comments:\n- In Definition 3.1, it seems that the intention is for $P’$ to be some RASP-L program, as opposed to just a program. Otherwise, trivially any program $P$ is an n-RASP-L program by choosing $P’ = P$ and $T(n) = 1$.\n- In Equation (2), I believe that the criterion should be an argmin over the cross entropy loss instead of an argmax.\n\n\n[1] Bansal, A., Schwarzschild, A., Borgnia, E., Emam, Z., Huang, F., Goldblum, M., & Goldstein, T. (2022). End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, 20232-20242.", "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)