T-LoopFormer: Token-Level Elastic-Depth Looped Transformers for Latent Reasoning with Dynamic Routing

1 INTRODUCTION

• Existing looped approaches apply a fixed number of loops uniformly across tokens, despite their differing hidden states.The introduction identifies this as a source of reduced generated-token accuracy, unnecessary memory overhead, and lower inference throughput.Source: §1 · p. 2
• A dynamic token-choice router uses each token’s hidden state to decide whether it exits early or continues recursing.The paper says this token-level elastic depth can improve token generation accuracy.Source: §1 · p. 2
• Recursion-wise KV caching retains entries for active tokens and discards those belonging to tokens that have exited, accelerating autoregressive decoding.Attention is restricted to locally cached entries, and cache size shrinks as tokens exit deeper recursions.Source: §1 · p. 2
• Experiments report robust performance on language modeling and reasoning benchmarks, alongside decoding efficiency compared with baselines.The results are presented as validating both the token-choice router and recursion-wise KV cache.Source: §1 · p. 2

2 RELATED WORK

• Looped Transformers reuse shared parameters to provide effective depth; prior timestep-conditioned approaches enable elastic-depth inference but apply the same recursion steps to all tokens [Jeddi et al., 2026].The Universal Transformer, ALBERT, and DEQ are also presented as parameter-sharing approaches [Dehghani et al., 2018; Lan et al., 2019; Bai et al., 2019].Source: §2 · p. 3
• Latent reasoning uses looped Transformers’ effective depth without verbalized reasoning steps, reducing context length and token consumption [Saunshi et al., 2024; 2025].The paper frames its method as extending this token-efficient reasoning to individual hidden states.Source: §2 · p. 3
• Prior dynamic routing allocates computation across tokens, layers, or experts, whereas T-LoopFormer routes token-level loop unrolling within a parameter-tied recurrent block [Raposo et al., 2024].The passage notes that MoD’s non-causal routing and discontinuous KV cache hinder parameter-tied looped architectures.Source: §2 · p. 3
• KV-cache methods include eviction and quantization, while T-LoopFormer’s recursion-wise cache is exact, loop-indexed, and retains entries only for active tokens.It is distinguished from eviction heuristics and lossy-approximation cache methods.Source: §2 · p. 4

3 T-LOOPFORMER

• T-LoopFormer is a looped Transformer with a shared stack, dynamic token-choice routing for token-level depth, and recursion-wise KV caching during inference.The stated goals are reducing memory use and accelerating autoregressive decoding.Source: §3 · p. 4

3.1 LONG AND SHORT TRAJECTORY ALIGNMENT

• Training aligns sampled shorter trajectories with a full trajectory, using next-token losses for both and a stop-gradient per-token logit-alignment loss.The method samples a shortcut length and step schedule; the full trajectory uses the maximum number of loops.Source: §3.1 · p. 4

3.2 DYNAMIC TOKEN-CHOICE ROUTING

• A two-layer MLP router maps each token’s initial hidden state to a recursion-depth distribution, determining whether it remains active at each loop.Active token sets are nested, so a token dropped from recursion does not re-enter; the router uses a straight-through estimator for backpropagation.Source: §3.2 · pp. 4–5
• Gather-scatter training compacts active-token states so the shared block processes survivors, then restores states to their original sequence positions.The computation at each loop scales with the number of active tokens; Figure 2 depicts gather-scatter during training and prefill and decode during inference.Source: §3.2 · pp. 5–6, Figure 2 · p. 5

3.3 RECURSION-WISE KV CACHE

• Inference keeps a separate KV cache for each recursion loop and attends only to valid historical entries from that same loop.Gather-scatter preserves original positions for cache writes; tokens that exited before a loop are masked from its attention.Source: §3.3 · p. 6

4 EXPERIMENTS

• Experiments compare a 24-layer, approximately 1B-parameter non-looped Transformer with FLOP-matched looped variants trained on FineWeb-Edu for 25B tokens.Evaluation uses perplexity on FineWeb-Edu and OpenWebText and zero-shot accuracy on ten reasoning benchmarks; baselines include fixed-depth and depth-elastic models.Source: §4 · pp. 7–8

4.1 MAIN RESULTS

• T-LoopFormer achieves the best performance among same-parameter models across tested FLOP settings and outperforms the 24×-parameter non-looped base at 24× FLOPs.It also has the lowest decoding latency among baselines, runs 2.24× faster than LoopFormer, and has the lowest training memory; most tokens use at most 4 loops (average depth 3.69).Source: §4.1 · pp. 8–9, Table 3

4.2 ABLATION STUDY

• At 24× FLOPs on FineWeb-Edu validation, removing the router increases PPL by 2.12 and TPOT by 0.70 ms, while removing recursion-wise KV cache increases PPL by 0.58 and TPOT by 0.196 ms.Source: Figure 3 · p. 8, §4.2 · p. 9
• Without gather–scatter, training memory increases 21.9% and inference latency increases 26.4% on FineWeb-Edu validation.Figure 4 examines its effects on training memory and inference latency.Source: Figure 4 · p. 8, §4.2 · p. 9
• At 24× FLOPs on FineWeb-Edu validation, router threshold 0.5 is best, with PPL following a U-shaped curve as the threshold varies.Source: Figure 4 · p. 8, §4.2 · p. 9
• As input length grows from 128 to 2048 tokens, T-LoopFormer's TTFT increases less than the base model's and remains the lowest across tested lengths.The comparison uses 1,000 FineWeb-Edu validation documents truncated to the target prompt lengths.Source: Figure 4 · p. 8, §4.2 · p. 9

• T-LoopFormer combines a dynamic token-choice router with gather–scatter and recursion-wise KV caching for token-level elastic-depth latent reasoning.Source: §5 · p. 9
• The authors report robust language-modeling and zero-shot reasoning performance under the same parameters, including surpassing the base Transformer at 24× FLOPs, and the lowest decoding latency.Source: §5 · p. 9
• Future work will extend T-LoopFormer to larger-scale pretraining and broader reasoning tasks.Source: §5 · p. 9

• DeepSeek assisted with language polishing and data-preprocessing tasks, including drafting or refactoring preprocessing code and checking data-cleaning scripts.Source: §6 · p. 10
• The authors state that DeepSeek was not used to generate scientific ideas, formulate methods, conduct experiments, interpret results, or write core scientific claims.Source: §6 · p. 10
• Authors reviewed, verified, and edited AI-assisted outputs; final data preprocessing used deterministic scripts, and authors take responsibility for the work.Source: §6 · p. 10

• The method, experimental setup, and evaluation protocol are described in Section 4 and Appendix A.Source: §7 · p. 10
• The anonymous supplementary repository provides source code, hyperparameters, random seeds, and step-by-step reproduction instructions.Source: §7 · p. 10
• The authors state that these resources are sufficient to reproduce the reported results.Source: §7 · p. 10

A IMPLEMENTATION DETAILS

• Models are trained on 8×A100 (80 GB) GPUs using the NanoGPT training stack, with 250,000 optimizer steps, batch size 48 sequences, and block size 2048.This corresponds to approximately 25B training tokens.Source: A IMPLEMENTATION DETAILS · p. 14
• Training uses AdamW with weight decay 2 × 10−1, cosine learning-rate decay, a peak learning rate of 6 × 10−4, a minimum of 6 × 10−5, and 4,000 warmup steps.Warmup and cosine decay mitigate observed depth-related training instabilities.Source: A IMPLEMENTATION DETAILS · p. 14
• All configurations use hidden size 2048, 32 attention heads, a 5120-dimensional two-layer GELU MLP, RMSNorm, and learned positional embeddings.Source: A IMPLEMENTATION DETAILS · p. 14
• Separate embedding modules map normalized loop time and step size to d-dimensional conditioning vectors using Fourier features followed by a two-layer SiLU MLP.The Fourier features use width 256 and maximum period 10000.Source: A IMPLEMENTATION DETAILS · p. 14
• For each active token, the time and step-size embeddings are summed and passed to a shared, token-wise AdaLN-style modulator that produces scaling and residual-gating parameters.The modulator is zero-initialized so initial behavior matches the unmodulated backbone and conditioning can be learned stably.Source: A IMPLEMENTATION DETAILS · p. 14

B THE DIFFERENCE BETWEEN GATHER–SCATTER DURING TRAINING AND INFERENCE.

• Training scatter restores compactly processed token states to their original sequence positions, allowing later loops to gather changing active sets and keeping states aligned with the loss.Indexed assignment also ensures gradients flow to the corresponding compact-tensor rows.Source: B THE DIFFERENCE BETWEEN… · pp. 14–15
• At inference, scatter updates token states and writes new keys and values to their original positions in the loop-wise cache, preserving the causal alignment future tokens use.Source: B THE DIFFERENCE BETWEEN… · p. 15
• Inference must propagate ρ(i) = sort(Ai) from gathering so compact rows write to original sequence positions in the loop-i cache; training uses Ai directly for scatter addresses.Source: B THE DIFFERENCE BETWEEN… · p. 15
• At inference, the loop-i cache needs validity masking because slots for tokens with mτ < i are empty and otherwise can receive non-negligible attention probability.A block-causal mask is also used among currently gathered tokens during batched prefill; training has no cache reads or corresponding mask.Source: B THE DIFFERENCE BETWEEN… · p. 15
• During single-token autoregressive decoding, gathering and scattering reduce to an identity operation and in-place write, so recursion ends when the token leaves the active set.Cache updates become ordinary tail concatenation, and ρ(i) reduces to identity.Source: B THE DIFFERENCE BETWEEN… · p. 15
• The training and inference versions of scatter preserve different invariants: the computation graph during training and the KV cache during inference.Source: B THE DIFFERENCE BETWEEN… · pp. 14–15

C ROUTING DESIGN CONSIDERATIONS

• Hard thresholding and discrete gather–scatter indexing make the routing computation graph discrete, but the router receives gradients through complementary paths.Source: C ROUTING DESIGN CONSIDERATIONS · p. 15
• The selection rule is a differentiable function of router logits, so gradients flow into the softmax distribution at every loop despite the discrete binary mask.Source: C ROUTING DESIGN CONSIDERATIONS · p. 15
• For selected tokens, task loss backpropagates through the shared block to the gathered representation and router weights; a straight-through estimator can pass task-loss gradients through routing when desired.Source: C ROUTING DESIGN CONSIDERATIONS · p. 15

D DISCUSSION FOR PERFORMANCE DEGENERATION

• When a token’s initial-state loop count exceeds the maximum loop limit, decoding degenerates; because loop counts are not recomputed, later context understanding may be constrained and reasoning performance may degrade.The fixed initial-state count supports static tensor compaction and avoids sequential halting checks at each loop.Source: D DISCUSSION FOR PERFORMANCE… · p. 16