Looped Transformers for Length Generalization

• Transformers still struggle with unseen input lengths, and fixed-depth computation may be insufficient for some tasks.Source: §1 · p. 1
• n-RASP-L tasks use repeated RASP-L operations whose required step count depends on problem complexity, including input length.Source: §1 · p. 2
• The approach trains Looped Transformers with end-to-end supervision and known step counts, then adapts the number of loops at inference.Source: §1 · p. 2
• The proposed method outperforms baseline approaches in length-generalization performance.Source: §1 · p. 2

2.1 RASP-L

• RASP-L programs map arbitrary-length input sequences to same-length outputs, using restricted operations and no branching or loops.Source: §2.1 · pp. 2–3
• Zhou et al. (2024a) report length-generalizable results for counting, mode, sorting, and unique-token copying, but difficulty with addition, parity, and copying under NTP.Source: §2.1 · p. 3

2.2 NEXT-TOKEN PREDICTION AND FULL-OUTPUT PREDICTION

• NTP trains on the full query-answer sequence and predicts shifted tokens, then continues token-by-token at inference until EOS.Source: §2.2 · p. 3, Figure 2 · p. 3
• FOP supplies only the query plus EOS padding and predicts all missing answer tokens after internal processing, rather than one token at a time.Source: §2.2 · p. 3, Figure 2 · p. 3

• Under NTP, addition, parity, and copy lack known single RASP-L programs for predicting each next token from any answer prefix.Source: §3 · p. 3
• An n-RASP-L program applies one RASP-L operation T(n) times, optionally with preprocessing and postprocessing.Source: §3 · p. 3
• Parity uses n iterations of shifting and XOR, while copy uses n iterations of shifting.Source: Figure 3 · p. 4, §3 · p. 4
• Addition uses n+1 iterations of shifted XOR and AND, with the inputs preprocessed.Source: Figure 3 · p. 4, §3 · p. 4

4.2.2 TRAINING ALGORITHM

• Without intermediate targets, final-answer training at sample-specific step counts lets the shared decoder receive supervision across different lengths and numbers of steps.The authors expect this to help the model learn intermediate steps that can generalize to more steps at test time.Source: §4 · p. 5, §4.2.2 · pp. 5–6

4.1 END-TO-END SUPERVISED DATA WITHOUT INTERMEDIATE STEP SUPERVISION

• Training assumes an n-RASP-L solution exists and supplies input-output pairs, required iteration counts, and varied lengths and step counts.Intermediate outputs are not provided.Source: §4.1 · p. 5, §4 · p. 5

4.2.1 ARCHITECTURE OF THE LOOPED TRANSFORMERS

• The architecture reuses one decoder block across adjustable loop steps, injects the original input at each step, and uses NoPE to match the RASP-L assumption.Input injection adds the original input embeddings to the previous step’s output embeddings.Source: §4.2.1 · p. 5

4.3 ADAPTIVE INFERENCE

• At inference, stopping can use the known required step count or maximum confidence; confidence may be assessed per length-batch or per sample, with Tmax as the cap.The batch size B may equal the test-set size for a given length or 1 for per-sample stopping.Source: §4.3 · p. 6

Positional embedding for length generalization

• Positional-embedding approaches have improved length generalization, but the paper treats them as orthogonal to its RASP-L focus, uses NoPE, and leaves other embeddings for future study.Other positional embeddings might combine with the approach, though they might not be expressible as RASP-L operations.Source: §5 · p. 6

RNNs and Chomsky Hierarchy

• A Chomsky-hierarchy study reports architecture-specific generalization limits, while this looped model combines recurrence and augmented memory with whole-sequence operations.That study reports failures by RNNs and Transformers on non-regular tasks, success by LSTMs on regular and counter-language tasks, and success on some context-free and context-sensitive tasks only with structured memory.Source: §5 · p. 6

Universal Transformers and other looped models

• Compared with UT and PonderNet, the proposed model uses decoder-only Transformers, FOP rather than NTP, no positional encoding, and simpler halting; input injection is also useful.The paper says its step-dependent supervision is a further difference from prior looped Transformer work.Source: §5 · p. 7, Table 1 · p. 7

Input representation

• Task-specific input reformatting can make addition easier to express as RASP-L, but such representations may not exist generally.The cited examples include index hints and reversed-order answers.Source: §5 · p. 7

COT

• Scratchpad or CoT can support length generalization by exposing intermediate results, but may hinder it when steps add complexity or require operations not easily expressed in RASP-L.The paper also contrasts this approach with fixed-depth models, noting that adaptive compute can be used at inference.Source: §5 · p. 7

6.1.1 TASKS

• The evaluation covers Parity, Copy, Addition, Binary Sum, Multiplication, and Unique Set, with task-specific definitions of problem length and required steps.For example, Parity uses the number of input digits as both problem length and step count; Multiplication sets steps to the product of operand lengths.Source: §6.1.1 · pp. 7–8

6.2 LENGTH GENERALIZATION RESULTS

• Trained up to 20 digits, the looped model generalizes on Parity to more than 40 digits near perfectly and nearly solves Addition and Copy at maximum training length +10, where NTP fails.The adaptive-depth model also outperforms NTP variants across the studied tasks.Source: Figure 4 · p. 8, §6.2 · p. 9

6.1.2 BASELINE METHODS

• The baselines are vanilla NTP with comparable effective training depth, NTP with pause tokens, and fixed-step weight-tied NTP called NTP-Loop.NTP-Loop uses 20 looped steps; NTP-Pause uses the same depth as vanilla NTP.Source: §6.1.2 · pp. 8–9

6.1.3 TRAINING AND EVALUATION SETUP

• All methods use a GPT-2 decoder and a curriculum that increases sequence length; evaluation measures exact-match accuracy with oracle or maximum-confidence stopping.The reported Sections 6.2 and 6.3 results use oracle steps, while Section 6.4 evaluates maximum-confidence stopping.Source: §6.1.3 · p. 9

6.3 ABLATION STUDIES

• Input injection generally improves generalization over its removal, with the effect most visible on Addition, Binary Sum, and Unique Set.The ablation compares the method without input injection against the full approach.Source: §6.3 · p. 9
• FOP pause tokens sometimes improve over vanilla FOP, but remain less effective than adaptive looped steps; deep FOP also has worse in-distribution Addition accuracy.The authors attribute the Addition accuracy issue mainly to training a model with 20× the decoder-block depth.Source: §6.3 · p. 9

6.4 THE STOPPING CRITERION AND VISUALIZATIONS

• At selected steps, maximum-confidence stopping gives near-perfect full-test accuracy across tasks; per-sample stopping performs worse on non-converging tasks.Figure 5 marks the step chosen using B = Ntest, the test-set size.Source: Figure 5 · p. 10, §6.4 · p. 10
• For Addition, Copy, Multiplication, and Unique Set, the looped model converges after solving without explicit convergence training, and its loss curves are smoother.Source: §6.4 · p. 10

7 LIMITATIONS AND CONCLUSION

• The current n-RASP-L definition excludes tasks requiring multiple loops in sequence, while direct looped training can be computationally demanding at large step counts.The authors report training only a limited number of steps and lengths because of limited computing resources; step-dependent supervision also requires ground-truth step counts.Source: §7 · p. 10
• The conclusion is that n-RASP-L problems can be learned with final-answer step-dependent supervision and handled with adaptive inference steps to improve generalization.Source: §7 · p. 10

A n-RASP-L PROGRAMS

• The appendix supplies n-RASP-L programs for Parity, Addition, and Copy, alongside library functions used to implement them.The provided code includes loop procedures and step functions for these tasks.Source: A n-RASP-L PROGRAMS · pp. 14–15

B RESULTS FROM THE ADAPTIVE STOPPING CRITERIA

• Per-sample stopping works almost as well as batch stopping on Addition, Copy, Multiplication, and Unique Set, but less well on Parity and Binary Sum.Batch stopping uses the test set for the stopping criterion; per-sample stopping chooses steps for each sample.Source: B RESULTS FROM THE ADAPTIVE STOPPING… · p. 16, Figure 6 · p. 17
• Variable depth generally outperforms fixing the depth of the model.The comparison includes FOP-Loop-Fixed.Source: B RESULTS FROM THE ADAPTIVE STOPPING… · p. 16
• The stopping variants use predefined steps, a batch criterion with B equal to test-set size, or a per-sample criterion with B = 1.For the per-sample variant, accuracy is still calculated across the dataset.Source: Figure 6 · p. 17

C ABLATION STUDY

• The adaptive-depth looped Transformer generalized better than FOP variants across the studied tasks.The comparison includes a version without input injection and FOP with pause tokens.Source: Figure 7 · p. 17
• The ablation includes the adaptive-depth method without input injection and FOP with pause tokens.The figure caption identifies these as comparison variants.Source: Figure 7 · p. 17
• The vertical dashed line marks the maximum training length.Source: Figure 7 · p. 17

D VISUALIZATION OF USING PAUSE TOKENS IN NTP AND FOP

• NTP-Pause and FOP-Pause are visualized with examples from the Copy task.The ellipsis denotes the pause token.Source: D VISUALIZATION OF USING PAUSE TOKENS… · p. 17
• The pause-token variants use 3 pause tokens in the figures and 20 in experiments, added before the final answer during training and inference.Source: D VISUALIZATION OF USING PAUSE TOKENS… · p. 17
• Table 2 defines the number of heads and block depth as looped-decoder-block size, and interval as the steps between sequence-length increases during training.Source: Table 2 · p. 18

E INFERENCE TIME COMPLEXITY

• With known a(n′), inference takes O(ka(n′)n′^2), compared with O(Cn′^2) for NTP with KV cache and FOP.Here n′ is test length, k is layers per step, and C is the fixed layer count for NTP and FOP.Source: E INFERENCE TIME COMPLEXITY · p. 18
• Confidence-based stopping takes O(kN′n′^2), where N′ is the maximum number of steps searched.Source: E INFERENCE TIME COMPLEXITY · p. 18
• In the experiments, NTP and FOP use C approximately kT(n), making their training compute similar; the looped method's inference is slightly longer because it uses more steps.Source: E INFERENCE TIME COMPLEXITY · p. 18

F EXPERIMENTAL DETAILS

• Looped computation uses a decoder-only GPT-2 block with NoPE and 256 embedding dimensions, operating in embedding space before decoding the final output.The original input embedding is added at each iteration.Source: F EXPERIMENTAL DETAILS · pp. 18–19
• Vanilla NTP uses autoregressive loss; NTP-Pause and FOP-Pause add 20 pause tokens, while FOP-Looped-Fixed uses the maximum depth of FOP-Looped-Adaptive.Input injection adds the original input embedding to each iteration, following a technique similar to Yang et al. (2024).Source: F EXPERIMENTAL DETAILS · pp. 18–19
• Training samples characters and lengths evenly online, gradually increases maximum training length, and runs 100k gradient steps with AdamW and cosine learning-rate decay.The learning rate decays from 10−4 to 0 after reaching the maximum training length; batch size is 64.Source: F EXPERIMENTAL DETAILS · p. 19
• Evaluation uses 6400 random samples and reports mean exact-match accuracy with standard error across five training runs with different random seeds.EMA of model parameters, with factor 0.9999, is used for Parity and Binary Sum.Source: F EXPERIMENTAL DETAILS · p. 19