How Do Learned Depth Priors and Orthogonal Injection Improve Looped Language Models?

Fixed points let the endpoint replace the trajectory

A looped language model applies the same shared block of layers again and again. Every extra recurrence adds cost: activations to store during training, keys and values to cache during decoding, compute to prefill each prompt, and a replay to score each rollout during reinforcement learning. The paper's starting claim is that these costs shrink when recurrent states sit near a fixed point, because the fixed-point equation involves only the endpoint, not the path taken to reach it.

The paper tests that claim with terminal KV sharing, which keeps only each prefix token's terminal key-value pair instead of one bank per visit. Reusing terminal KV caches barely changes the zero-shot GSM8K accuracy of Huginn, but collapses that of Ouro at four loops. The paper attributes the split to convergence: Huginn's states approach fixed points, Ouro's do not. This is the paper's explanation of an observed contrast, not a separate controlled experiment.

The same fixed point gives each remaining cost a shortcut. Truncated backpropagation through only the last few loops approximates the equilibrium gradient without the warm-up implicit gradients require. A student distilled to predict the fixed point directly replaces the loops during prefill. In RL, scoring rollouts from saved fixed-point states avoids replaying the loops. Each shortcut removes a depth-dependent cost, so its savings grow with every added loop.

The depth prior decides what supervision sees

The depth prior is the distribution of training depths. It matters because tokens converge at different depths, so the sampled depth decides whether training sees the converged prefix context that terminal KV sharing supplies at inference. A model trained at a single depth predicts best there but breaks KV sharing. Huginn's fixed broad prior supports sharing but dilutes supervision at the target depth more than sharing requires.

The paper learns the prior from prediction feedback, which pulls supervision toward depths that predict well, and adds an entropy term that keeps it broad enough for sharing. Prediction feedback under a budget constraint favors depths that already predict well, so left alone it concentrates supervision on a narrow range of depths and can weaken fixed-point shaping. The entropy term counters that tendency, and because depths start at 1, the wider allocation extends mainly above the mean depth toward the maximum.

The entropy coefficient is therefore a direct knob for balancing predictive performance against memory efficiency. In the matched M-scale sweep, raising the coefficient from 0 to 0.1 keeps the mean depth near five while the probability mass at depth 5 falls from 67.4% to 14.0%. WikiText perplexity at the training depth of 5 rises from 12.89 to 13.20, while at depth 32 the coefficients 0.01 and 0.05 tie for the lowest perplexity at 12.93.

The learned prior beats fixed depth and PLN sampling

The paper compares three ways of choosing training depth. Fixed-depth training is best at the depth it was trained on, R = 5, but degrades sharply at other depths. PLN sampling performs well across all depths but falls behind at R = 5. The learned prior lies between them, narrowing the gap at R = 5 while staying stable across depths.

The clearest consequence appears in KV sharing. Scoring each checkpoint with per-visit and shared KV shows that sharing raises the continuation perplexity on WikiText, after a 1,024-token prefill, of fixed-depth training from 12.19 to 16.71. PLN-5 and all learned priors retain or slightly improve it. Higher entropy does not monotonically improve sharing on WikiText and LAMBADA, so the entropy term is not a simple more-is-better setting.

Without KV reuse, the coefficient 0 gives slightly lower validation perplexity than 0.01 at M and L scale (4.55 vs. 4.57 and 3.53 vs. 3.54) and ties it at S. With terminal KV sharing it falls behind at L scale on GSM8K, 36.47 versus 47.92. The same checkpoints can therefore rank differently depending on whether the cache is shared, which is why the paper reports both layouts.

Orthogonal injection removes the state component along the input

Input injection feeds the same embedding into every recurrence, keeping the fixed points input conditioned. In the additive form used by Huginn and Parcae, the carried-over state is added to the injected input. Parcae uses a diagonal carryover with entries in (0,1), so the carryover can have a component along the injected input that amplifies or cancels it. The input a fixed point retains then depends on the state, which the paper reports can hurt performance.

Orthogonal injection removes that component by projecting the carried-over state off the injected input. For a nonzero input and zero stabilizer, the component of the injected quantity along the input is exactly the input at every recurrence, including at a fixed point, while the orthogonal part evolves freely. A positive stabilizer makes this approximate. Because the projection has norm at most one, the projected carryover stays contractive.

The practical effect is that every recurrence receives the input at the same strength. The paper reports that orthogonal injection gives lower perplexity than the injection schemes of Huginn and Parcae, and that this holds from 100M to 1.6B parameters. Orthogonal injection also drops the RMSNorm that Parcae applies to the prelude output before injection; the paper places that ablation in its appendix.

Scale results and the 3x smaller cache

From 100M to 1.6B parameters, the learned prior and orthogonal injection lower perplexity at every scale relative to Huginn's prior and existing injection schemes, respectively. The headline comparison is at 1.6B parameters: the learned prior with a 3x smaller KV cache matches the downstream average of fixed-depth training with the full cache.

The cache reduction comes from terminal KV sharing at five recurrences, which cuts the KV cache 3x. The paper reports that sharing does not lower the average task accuracy of any learned prior, with average changes from +0.25 to +0.70 points, whereas fixed training at R = 5 loses 0.61 points. Those paired numbers come from the M-scale entropy sweep at prefill and decode depth 5, so they describe that setting rather than every scale.

The paper also reports two other shortcuts. A distilled student prefills up to 1.79x faster, and RL updates compute gradients from saved rollout states 2x faster than backpropagating through the replayed trajectory. The paper describes distilled prefill as trading downstream accuracy for prefill latency, with a gap that grows with scale, so the prefill speedup is not a free improvement.

Where these results stop transferring

The evidence covers models up to 1.6B parameters with one seed per configuration, and the ablations do not exhaust the design space. Larger scales and other model configurations remain untested, and the paper states that feasibility in real-world applications needs further study. The scale comparison therefore supports claims about the tested range, not about looped models in general.

Several reported numbers are tied to specific protocols. The sharing penalties are within-checkpoint ratios of terminal to full-KV perplexity, not differences from native MixVal. The paired task accuracies use 2,000 paired bootstrap draws per task with equal task weights, and the intervals quantify evaluation uncertainty rather than training-seed variability or differences between coefficients. Document-level pretraining decontamination was not performed.

The entropy sweep is also a single-scale result. The M-scale runs differ only in the entropy coefficient and share model, data, seed, learning-rate schedule, and controller, which makes that comparison clean but narrow. The claim that higher entropy spreads the prior and lowers the sharing penalty should be read as an M-scale finding, not as a scaling law.

The knob to tune first is the entropy coefficient

The paper's most decision-relevant result is a trade, not a single best setting. Increasing the entropy coefficient spreads the learned prior toward later depths, which lowers the KV sharing penalty, from +0.96% at coefficient 0 to -0.22% at coefficient 0.1, but raises perplexity at the training depth of 5 from 12.89 to 13.20. The coefficient is the direct control over that balance.

The absolute shared MixVal perplexity is lowest at coefficient 0.01, at 4.48 and within 0.01 of coefficients 0 and 0.05. That makes 0.01 the natural starting point when the goal is shared-cache quality rather than maximum depth spread. The choice between 0.01 and 0.1 depends on whether the priority is the training-depth perplexity or the sharing penalty at later depths.

The second priority is the injection scheme. Orthogonal injection is the change that lowers perplexity relative to existing schemes across the tested scales, and it is independent of the depth-prior choice. For a reader deciding what to adopt first, the paper's evidence points to orthogonal injection as the change with the broadest scale support, and to the entropy coefficient as the setting that must be matched to the intended cache regime.

Sources used in this answer