Towards Looped Models Done Right: Part II — Rethinking at Fixed Points

Summary (Overview)

• Fixed points enable efficiency shortcuts: When looped language model states converge to fixed points, the endpoint can replace the full recurrent trajectory, enabling truncated backpropagation in training, terminal KV sharing in decoding, distilled prefill, and rollout-state reuse in RL.
• Two training components shape fixed points: The depth prior (distribution of training depths) and input injection (how the input enters each recurrence) determine whether models actually converge to useful fixed points.
• Learned depth prior outperforms fixed priors: A categorical prior learned from prediction feedback with entropy regularization lowers perplexity at every scale (100M–1.6B) compared to Huginn's fixed Poisson-lognormal prior, while keeping terminal KV sharing intact.
• Orthogonal input injection (OrthoInj) improves convergence: Projecting out the state's component along the injected input ensures every recurrence receives the input at consistent strength, outperforming Parcae's injection scheme.
• Two new shortcuts enabled: Distilled prefill achieves up to 1.79× faster prefill on 8K prompts, and rollout-state reuse in RL halves scoring/backward time while maintaining quality.

Introduction and Theoretical Foundation

Looped language models repeatedly apply a shared block of layers, with two competing design philosophies: one camp keeps loop counts small and fixed (forgoing fixed points), while the other shapes fixed points and loops much deeper. The paper argues that fixed points hold the key to cutting costs at every stage: pretraining (activation storage), serving (KV cache), prefill (compute), and post-training (rollout replay in RL).

The core theoretical foundation rests on the fixed-point equation:

$$z^{\star} = F_{\theta}(z^{\star}; x)$$

where $z^{\star}$ is the fixed-point state, $F_{\theta}$ is the shared update map, and $x$ is the input. The sensitivity of the fixed point to parameters is given by implicit differentiation:

$$\frac{\mathrm{d}z^{\star}}{\mathrm{d}\theta} = (I - \partial_{z}F_{\theta})^{-1}\partial_{\theta}F_{\theta}$$

This depends only on the endpoint, not the path to it — the key insight that enables all the efficiency gains.

Convergence measurement: Tokens converge at heterogeneous depths, measured by the tokenwise relative change:

$$\rho_{r,t} = \frac{\|h_t^r - h_t^{r-1}\|_2}{\|h_t^{r-1}\|_2 + \epsilon}, \quad \epsilon = 10^{-8}$$

Token $t$ is converged from depth $r$ if $\rho_{r',t} < \tau$ (threshold $\tau = 2\%$) for all $r \leq r' \leq R_{\text{obs}}$.

Methodology

1. Truncated Backpropagation (TBPTT)

The equilibrium gradient can be expressed via Neumann series:

$$D_{\star} = (I - J_{\star})^{-1}B_{\star} = \sum_{k=0}^{\infty}J_{\star}^{k}B_{\star}, \quad D_{b} = \sum_{k=0}^{b-1}J_{\star}^{k}B_{\star}$$

where $J_{\star} = \partial_z F_{\theta}(z^{\star};x)$ and $B_{\star} = \partial_{\theta}F_{\theta}(z^{\star};x)$. TBPTT differentiates through only the last $b$ recurrences, with gradient:

$$\tilde{D}_b = \sum_{k=0}^{b-1}\left(\prod_{j=1}^{k}J_{R-j}\right)B_{R-1-k}$$

This is stable from initialization (unlike pure implicit gradients which require warm-up) and recovers the fixed-point gradient as states converge.

2. Rollout-State Reuse for RL

Full BPTT gradient $\tilde{D}_R$ can be approximated using only the endpoint $z^R$:

$$\tilde{D}_R = \sum_{k=0}^{R-1}\left(\prod_{j=1}^{k}J_{R-j}\right)B_{R-1-k} \approx \sum_{k=0}^{b-1}J_R^k B_R$$

The reuse algorithm (Neumann-(b−1)): applies the recurrent block once at $z^R$, then loops the backward pass $b-1$ times over that single recorded recurrence, iterating $v \leftarrow g + J_R^{\top}v$.

3. Learned Depth Prior

Replaces the fixed PLN prior with a categorical prior $p_{\phi} = \text{softmax}(\phi)$ over depths $\mathcal{R} = \{1, \ldots, R_{\max}\}$, learned by minimizing:

$$\mathcal{J}_{\text{prior}}(\phi) = -\hat{\mathbb{E}}_n[\text{stopgrad}(A_n)\log p_{\phi}(r)] - \lambda_H H(p_{\phi}) + \lambda_m(\mathbb{E}_{p_{\phi}}[R] - \bar{R})^2$$

• Prediction quality: Advantage $A_n = s_{n-1}(u - \hat{u}_{n-1})$ where $u = \exp(-\ell)$ (inverse perplexity) and EMA statistics center/rescale the reward
• Entropy regularization: $\lambda_H H(p_{\phi})$ keeps the distribution broad for KV sharing
• Budget control: $\lambda_m(\mathbb{E}[R] - \bar{R})^2$ maintains the target mean depth

4. Orthogonal Input Injection (OrthoInj)

Standard additive injection (Parcae): $\zeta_t^r = Mh_t^r + q_t$ with diagonal carryover $\Lambda = \text{diag}(\exp(-\Delta \odot \alpha))$.

OrthoInj projects out the state's component along the input:

$$\zeta_t^r = Q_{q_t}\Lambda h_t^r + q_t, \quad Q_{q_t} = I - \frac{q_t q_t^{\top}}{\|q_t\|_2^2 + \epsilon}$$

This ensures the component along $q_t$ is exactly $q_t$ at every recurrence, with contractive projected carryover: $\|Q_{q_t}\Lambda\|_2 \leq \|\Lambda\|_2 < 1$.

5. Distilled Prefill

A lightweight student $S_{\psi}$ replaces the iterated recurrent core during prefill, predicting $\hat{z} = S_{\psi}(e) \approx z_{\theta}^R(x)$ from the prelude output $e$, trained with:

$$\mathcal{L}_{\text{distill}}(\psi) = \mathbb{E}_x\left[\frac{\|\hat{z} - z_{\theta}^R(x)\|_F^2}{Nd\sigma_T^2} + \frac{1}{N}\sum_{t=1}^{N}\text{KL}(p_t \| \hat{p}_t)\right]$$

6. Theoretical Guarantee (Lemma 1)

Under contraction conditions ($\kappa < 1$) and Lipschitz continuity ($\beta$-Lipschitz on context), the joint and frozen-prefix iterations:

$$h^{r+1} = f_{\theta}(h^r; C_{<t}^r), \quad \tilde{h}^{r+1} = f_{\theta}(\tilde{h}^r; C_{<t}^{\star})$$

both converge to the same fixed point $h^{\star} = f_{\theta}(h^{\star}; C_{<t}^{\star})$, proving terminal KV sharing transfers from training to decoding.

Empirical Validation / Results

Key Results at 1.6B Scale (L)

Table: Learned Prior vs. Baselines (L scale, R=5, terminal KV sharing)

Key findings:

• The learned prior with entropy regularization matches the full-cache downstream average of fixed-depth training while using a 3× smaller KV cache
• Fixed-depth training breaks KV sharing: GSM8K drops from 50.6 to 21.2 under sharing
• The learned prior's lead over Untied 4 grows with scale: +2.2 (S), +5.0 (M), +8.9 (L) downstream points

OrthoInj Results

Rollout-State Reuse (RL)

Reuse achieves 1.99–2.02× faster scoring/backward time with quality within run-to-run variation (whole-model gradient cosines 0.997–0.999 vs. recomputation).

Distilled Prefill

At 8K prompts, distilled prefill is 1.5–1.8× faster than full teacher prefill:

• S: 125.9ms vs 195.0ms (1.55×), trails downstream AVG by 1.0 points
• M: 299.3ms vs 495.8ms (1.66×), trails by 2.3 points
• L: 830.3ms vs 1482.1ms (1.79×), trails by 4.5 points

At matched latency, distilled prefill scores 0.5–0.9 points above teacher prefill at R=2.

Theoretical and Practical Implications

Theoretical implications:

• Provides formal justification (Lemma 1) for why terminal KV sharing works when training supervises near fixed points
• Unifies TBPTT and implicit differentiation: TBPTT is stable early (exact gradient of truncated objective) and recovers equilibrium gradients once fixed points form, with no warm-up or switch needed
• Establishes that truncation alone shapes fixed points (unlike fixed-depth training with full BPTT)

Practical implications:

• Pretraining: TBPTT with $b \geq 4$ matches full BPTT within 1% perplexity at 400M scale
• Decoding: Terminal KV sharing cuts cache 3× at five recurrences with minimal quality loss
• RL post-training: 2× faster updates from saved rollout states, enabling cheaper iterative refinement
• Serving: Distilled prefill trades controllable accuracy for latency, with gains growing at scale

Conclusion

The paper demonstrates that fixed points in looped language models enable decoupling training and inference costs from recurrence depth. The two components that shape fixed points — depth prior and input injection — are both improved: a learned, entropy-regularized depth prior outperforms fixed priors while preserving KV-sharing compatibility, and OrthoInj removes the state's component along the injected input for cleaner convergence.

Future directions identified by the authors:

• Scaling beyond 1.6B parameters and testing other model configurations
• Better distillation objectives and student architectures (e.g., mixture-of-experts)
• Student as standalone language model or speculative decoding drafter
• Connections between recurrent state refinement and flow-based language modeling
• Wider application of fixed-point reuse (e.g., warm-starting solvers on related inputs)

Limitations: Models trained up to 1.6B with one seed per configuration; ablations do not exhaust the design space; real-world feasibility needs further study.