Technical Summary: Dense Supervision Is Not Enough: The Readout Blind Spot in Looped Language Models

1. Problem Statement

Looped language models (LLMs) reuse the same transformer block across recurrent depth, allowing test-time depth to function as a dynamic compute knob (e.g., stopping early for easy predictions). In this architecture, a hidden state $H_k$ serves a dual role: it is decoded for prediction (the "readout" interface) and fed back as input for future computation (the "recurrent" state).

The central problem identified is a visibility–activity mismatch. While dense supervision (applying cross-entropy loss at every loop) ensures that intermediate states are usable for prediction, it does not necessarily control all variables active within the recurrent transition. Specifically, the paper investigates hidden-state scale (the radial magnitude of the hidden vector).

In standard pre-norm residual loops with scale-invariant readouts (e.g., RMSNorm or LayerNorm), the readout normalizes the hidden state, making the immediate cross-entropy loss insensitive to the radial scale of the input. However, the recurrent path (the skip connection and residual update) continues to carry and update this scale. Consequently, the model can be trained at every loop while the hidden-state norms drift uncontrollably to thousands or tens of thousands, rendering the "readout blind spot" a failure mode where dense supervision fails to constrain active recurrent variables.

2. Methodology and Mechanism

Theoretical Mechanism

The authors formalize the failure mode through two lemmas:

1. Visibility Lemma: If a readout $r$ is scale-invariant (i.e., $r(\alpha H) = r(H)$ for $\alpha > 0$), the immediate cross-entropy loss $L$ has a radial derivative of approximately zero ($\langle \nabla_H L, H \rangle \approx 0$). The loss cannot "see" or correct the scale of the hidden state.
2. Activity Lemma: In pre-norm residual updates, the scale is active. The update function $F(su) = su + b(u)$ (where $s$ is scale and $u$ is direction) allows the residual branch $b(u)$ to update the scale at the leading order. Even if the readout hides the scale, the recurrence preserves and potentially amplifies it.

This creates a local mismatch: the loss signal cannot oppose the learned radial velocity of the recurrent update, leading to scale drift.

Experimental Setup

The authors trained autoregressive looped transformers on the WikiText-103 dataset with $K=4$ recurrent applications of a shared decoder stack. They utilized two model sizes: 44M and 129M parameters.

• Core Ablation: A $2 \times 2$ design crossing Loss Placement (Terminal-only CE vs. Per-loop CE) and Readout Type (Normalized RMSNorm vs. Raw/Unnormalized).
• Controls:
  • Norm Penalty: Adding an explicit auxiliary loss term $\lambda \sum \|H_k\|^2$ to make scale visible without changing the readout.
  • Final-only Norm: Using raw readouts for intermediate loops but a normalized readout for the final output.
  • Recurrent Scale Clamp: Manually rescaling hidden states to their initial loop-1 magnitude to test if accumulated scale growth is predictive.
• Evaluation: Measured final-loop hidden-state norms, perplexity (PPL), and the "compute-quality frontier" (PPL vs. throughput) at variable inference depths ($K=1$ to $K=4$).

3. Key Contributions

1. Identification of the Readout Blind Spot: The paper defines a specific failure mode where dense cross-entropy supervision trains prediction interfaces but fails to control active recurrent variables (specifically scale) if those variables are hidden by scale-invariant readouts.
2. Mechanism Validation: Through ablations and diagnostics, the authors demonstrate that:
   • Per-loop CE with normalized readouts leads to massive norm drift (e.g., norms reaching 39,000+ in the 44M model).
   • Raw readouts or explicit norm penalties collapse these norms to the tens, even with per-loop CE.
   • The drift is token-level and heavy-tailed, not just a mean shift.
3. Design Rule: The authors propose a simple design rule: Recurrent scale control requires either making hidden-state scale visible to a loss (via raw readouts or norm penalties) or removing it from the recurrent path (via inter-loop normalization). Dense CE alone is insufficient.
4. Variable-Depth Performance: Scale-controlled models (using raw readouts or norm penalties) maintain usable early exits while achieving lower perplexity at matched inference depths compared to scale-invariant baselines. The baselines often exhibit "K-invariance" (performance does not improve with more loops) because the scale drift shrinks angular updates ($1/s_k$), effectively freezing the normalized state.

4. Results

• Norm Drift: In 44M and 129M models with per-loop CE and RMSNorm readouts, final hidden-state norms drifted to 39,207 and 56,051, respectively. In contrast, models with raw readouts or norm penalties kept norms in the range of 17–83.
• Gradient Diagnostics: Radial gradient fractions for normalized readouts were near zero ($\sim 10^{-9}$), while raw readouts showed significant radial sensitivity ($\sim 10^{-2}$).
• Perplexity and Depth:
  • RMSNorm Baseline: Models were nearly K-invariant; stopping at $K=1$ yielded similar perplexity to $K=4$. This resulted in high apparent speedup but indicated a lack of useful depth utilization.
  • Scale-Controlled Variants: These models showed significant perplexity improvement as depth increased ($K=1$ to $K=4$). They occupied a superior PPL-throughput frontier, offering lower perplexity at the same throughput compared to the RMSNorm baseline.
• Scale Clamp: Artificially clamping the scale of hidden states in trained models resulted in negligible changes to cross-entropy, suggesting that the accumulated scale growth itself is not the primary source of predictive information, but rather a byproduct of the uncontrolled recurrence.

5. Significance and Claims

The paper claims that the "readout blind spot" is a fundamental architectural consideration for looped and adaptive-depth models. The significance lies in shifting the focus from merely applying dense supervision to ensuring that the supervision signal is visible to the specific variables being updated in the recurrence.

• Practical Implication: For practitioners using normalized readouts, the paper suggests adding an explicit norm penalty as a "drop-in" baseline to control scale without altering the output interface.
• Theoretical Insight: The work clarifies that dense supervision trains the interface (the readout) but does not automatically regulate the state dynamics (the recurrence) if the interface masks the state's magnitude.
• Scope Limitation: The authors modestly note that their conclusions are based on specific settings (WikiText-103, 44M/129M models, teacher-forced evaluation) and that the mechanism is local. They do not claim to solve all stability issues in looped models but specifically address the under-constrained nature of scale in scale-invariant readout loops.

The paper concludes that adaptive-depth looped models require both trained exits (via per-loop loss) and a mechanism to control active recurrent scale (via visibility or removal), and that these are orthogonal design requirements.