Rectified Looped Transformer (RLT)

Summary (Overview)

• Novel feedback mechanism: The paper introduces Rectified Looped Transformer (RLT), which replaces direct feedback in looped Transformers with a flow-inspired linear interpolation schedule between the current state and the network's endpoint prediction.
• Theoretical guarantees: RLT provides depth-uniform bounds on state magnitudes (constant in loop depth $K$) and reduces worst-case perturbation amplification from exponential ($L^K$) to polynomial ($K^{L-1}$) under a uniform Lipschitz condition.
• Empirical improvements: RLT consistently achieves lower perplexity than Standard Looped Transformer (SLT) across multiple depths (K=2, 4, 8) and evaluation corpora (The Pile, FineWeb-Edu, OpenWebText), with gains of up to 5.5% in perplexity reduction.
• Enhanced robustness: RLT demonstrates greater tolerance to input perturbations, inference-depth mismatches (both fewer and more loops), and quantization, while maintaining more stable hidden-state norms across loops.
• Zero additional cost: These gains require no extra trainable parameters or recurrent block evaluations—only elementwise interpolation between states.

Introduction and Theoretical Foundation

Background and Motivation

Looped Transformers repeatedly apply a shared stack of layers to increase effective computational depth without increasing physical parameters. This makes recurrence depth an additional axis for scaling language models. However, the standard direct-feedback formulation (SLT) applies the entire change from current state to network prediction at every loop, which can lead to:

• Hidden-state scale growth as recurrence deepens
• Perturbation amplification
• Performance degradation at larger recurrent depths (as shown by STARS)

Key Insight

The paper draws inspiration from rectified flow (Liu et al., 2023), which uses linear interpolation between source and target samples as training paths. RLT treats each loop's network output as an endpoint prediction and combines it with the current input state through a remaining-time interpolation schedule.

Core Formulation

The looped computation is associated with a uniform time grid on $[0, 1]$:

$$\Delta t = \frac{1}{K}, \qquad t_k = \frac{k}{K}, \qquad k = 0, \ldots, K-1$$

The shared Transformer stack $\mathbf{F}_\theta$ performs one loop of computation:

$$\mathbf{h}_{k+1} = \mathbf{F}_\theta(\mathbf{x}_k, t_k), \quad \mathbf{x}_{k+1} = (1 - \eta_k)\mathbf{x}_k + \eta_k \mathbf{h}_{k+1}, \quad 0 \leq \eta_k \leq 1$$

where $\mathbf{x}_k$ is the input state, $\mathbf{h}_{k+1}$ is the hidden output, and $\eta_k$ is the feedback coefficient.

SLT uses direct feedback ($\eta_k = 1$), while RLT uses the remaining-time schedule:

$$\eta_k = \frac{\Delta t}{1 - t_k} = \frac{1}{K - k}, \qquad k = 0, \ldots, K-1$$

The weight increases from $\eta_0 = 1/K$ to $\eta_{K-1} = 1$, forming a smoothly interpolated trajectory from $\mathbf{x}_0$ to the final prediction $\mathbf{x}_K = \mathbf{h}_K$.

Methodology

Architecture

• Base model: GPT-2-style looped Transformer with time and step-size conditioning via AdaLN-style modulation
• Default configuration: 3 shared layers × 8 loops (3 ⊗ 8), 278M physical parameters
• Supervision: Uniform all-step next-token prediction with loss $K^{-1}\sum_{k=1}^{K}\ell_k$
• Readout: $\mathcal{D}(\mathbf{h}) = \mathbf{W}\,\text{RMSNorm}(\mathbf{h})$ with cross-entropy loss

Training Setup

• Dataset: Deduplicated The Pile (6.88B tokens)
• Optimization: 14,000 steps, peak learning rate $6 \times 10^{-4}$, global batch size 480 sequences of 1,024 tokens
• Evaluation: Perplexity on The Pile, FineWeb-Edu, OpenWebText; zero-shot accuracy on 10 benchmarks (COPA, HellaSwag, LAMBADA, OpenBookQA, PIQA, RACE, SciQ, ARC, Social IQA, WinoGrande)

Theoretical Analysis Framework

The analysis establishes:

1. State norm bounds under time-dependent residual control
2. Perturbation amplification bounds under Lipschitz conditions
3. Fixed-point residual decay for nonexpansive networks

Empirical Validation / Results

Language Modeling Performance

At matched training and inference depths ($K_{train} = K_{test} = K$):

RLT reduces perplexity by 4.2–5.5% at K=4 across corpora. LAMBADA accuracy improves by 5 percentage points at both K=2 and K=4.

Recurrent Stability Results

Hidden-state scale: SLT exhibits pronounced scale accumulation (mean ratio relative to $\mathbf{h}_1$ reaches 5.08 by loop 8), while RLT maintains 0.65—consistent with the theoretical bound in Corollary 1.2.

Perturbation robustness: RLT shows smaller output displacements under input noise and better preserves predictive performance under per-loop perturbations. At $\rho = 0.5$, output perturbations increase Pile cross-entropy by 0.61 nats/token for RLT vs. 2.86 for SLT.

Quantization robustness: RLT retains its perplexity advantage across W8A16, W4A16, and W8A8 quantization settings at all inference depths.

Inference-Depth Adaptation

For models trained with $K_{train} = 8$:

At $K_{test} = 32$, RLT achieves PPL 17.72 vs. 57.72 for SLT, and average accuracy 42.88% vs. 38.05%.

Theoretical and Practical Implications

Theoretical Results

Theorem 1 (State norm bounds): Under the residual control condition $\sup_{\mathbf{x}} \|\mathbf{F}_\theta(\mathbf{x}, t) - \mathbf{x}\| \leq M(1-t)^\alpha$:

• SLT bound: $\|\mathbf{x}_K\| \leq \|\mathbf{x}_0\| + M \underbrace{K^{-\alpha}\sum_{j=1}^{K} j^\alpha}_{\longrightarrow \infty (K \to \infty)}$ — unbounded growth
• RLT bound: $\max\{\max_k \|\mathbf{x}_k\|, \max_k \|\mathbf{h}_k\|\} \leq \underbrace{\|\mathbf{x}_0\| + \frac{M}{\alpha}}_{\text{constant in } K}$ — uniformly bounded

Theorem 2 (Perturbation amplification): Under uniform Lipschitz bound $L > 1$:

• SLT: $\|\mathbf{x}_K - \tilde{\mathbf{x}}_K\| \lesssim L^K \|\mathbf{x}_0 - \tilde{\mathbf{x}}_0\|$ — exponential
• RLT: $\|\mathbf{x}_K - \tilde{\mathbf{x}}_K\| \lesssim K^{L-1} \|\mathbf{x}_0 - \tilde{\mathbf{x}}_0\|$ — polynomial

Proposition 3 (Fixed-point convergence): For nonexpansive networks, RLT achieves:

$$\|\mathbf{F}_\theta(\mathbf{x}_K) - \mathbf{x}_K\| \lesssim \frac{\text{dist}(\mathbf{x}_0, \text{Fix}(\mathbf{F}_\theta))}{\sqrt{\log K}}$$

Practical Implications

1. Stable scale enables decoding techniques: RLT's stable representation scales support effective contrastive decoding (accuracy gains up to 44.45% from 42.73% baseline) and guidance during recurrence (0.76 percentage point gain at $\lambda = 0.3$).

2. Flexible inference budgets: RLT supports a broader range of inference depths, making it practical for adaptive computation scenarios where the loop count may vary.

3. Quantization-friendly: The controlled state norms make RLT more amenable to low-precision inference.

Conclusion

RLT introduces a simple yet effective feedback mechanism for looped Transformers: interpolate between the current state and each loop's endpoint prediction according to remaining time. This design:

• Separates prediction from state advancement, preventing uncontrolled accumulation of changes
• Provides theoretical guarantees of bounded state norms and polynomial (rather than exponential) perturbation amplification
• Achieves consistent empirical gains in perplexity, robustness, and latent stability without additional parameters or compute

The work highlights state feedback as a crucial design choice for jointly improving language modeling quality and the stability of latent recurrent computation. Future directions may include exploring adaptive interpolation schedules, extending to other recurrent architectures, and scaling to larger models and datasets.