# Towards Looped Models Done Right, Part II: Rethinking at Fixed Points > A learned, entropy-regularized depth prior and orthogonal input injection improve fixed-point convergence in looped language models, enabling 2x faster RL training and 1.8x faster prefill. - **Source:** [arXiv](https://arxiv.org/abs/2610.06833) - **Published:** 2026-10-12 - **Permalink:** https://picx.dev/p/85xVun - **Whiteboard:** https://picx.dev/p/85xVun/image ## Summary # 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_{