🤖 AI

Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

This paper demonstrates that depth-recurrent transformers exhibit token-specific convergence to fixed points, enabling a training-free, per-token early-stopping mechanism that achieves uniform quality at significantly reduced average computational depth compared to both fixed-depth inference and learned routing strategies.

Original authors: Joe Logan

Published 2026-07-17

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Joe Logan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The World of Thinking Machines

Imagine a giant library where every book is written by a super-smart robot. This robot doesn't just read words; it tries to understand the meaning behind them. In the world of artificial intelligence, these robots are called "transformers." They work by looking at a sentence and asking, "What word comes next?" To do this, they pass the sentence through a series of mental "layers," like a ball of yarn being pulled through a series of hoops. Each hoop adds a little more understanding.

Usually, these robots are built to pull the yarn through a fixed number of hoops—say, eight. If the sentence is simple, like "The cat sat," the robot might only need three hoops to get it right. But if the sentence is a complex riddle, it might need all eight. The problem is that the robot is built the same way every time. It pulls the yarn through eight hoops for every sentence, even the easy ones. This is like using a sledgehammer to crack a nut: it works, but it wastes a lot of energy and time. Scientists have been trying to figure out how to let the robot stop early when it's done thinking, saving energy without losing its smarts. This paper dives into a clever new way to do exactly that.

The Paper: When to Stop Thinking

This study looks at a special kind of AI model called a "depth-recurrent transformer." Think of this model not as a factory assembly line with a fixed number of stations, but as a single, super-smart station that can be visited over and over again. The robot takes a word, runs it through the station, gets a new version of the word, and then runs that new version through the same station again. It can do this loop as many times as needed.

The researchers wanted to know two things: First, does this robot actually get smarter the more times it loops, or does it just spin its wheels? Second, can we tell exactly when the robot has "figured out" a specific word so we can stop the loop for that word and move on?

The Discovery: Finding the "Fixed Point"

The team trained a model on a huge collection of educational text (about 12.9 billion words) and watched what happened inside its brain. They found something amazing: for every single word, the robot eventually stops changing its mind.

Imagine you are trying to guess the ending of a story. At first, you might think, "Maybe it's a dog?" Then you think, "No, maybe a cat?" Then, "Wait, it's definitely a dog." Eventually, you stop guessing and just say, "It's a dog." That moment of certainty is called a "fixed point."

The researchers measured how much the robot's guess changed from one loop to the next. They found that the "thinking" settles down incredibly fast. By the 16th loop, the change in the robot's guess was so tiny it was almost zero (specifically, the difference dropped from 0.39 to 0.0000085). But here is the kicker: not all words take the same amount of time to settle.

The "Easy" Words: Simple words like spaces or punctuation marks usually figure themselves out by the 6th loop.
The "Hard" Words: Complex content words keep thinking longer. About 10% of the words are still changing their minds even at the 8th loop, which was the average number of loops the robot was trained to do.

The Big Test: Reading the Mind vs. Guessing the Mind

The most exciting part of the paper is how they used this discovery. They asked: "Can we just watch the robot's brain and stop it the moment it stops changing?" Or, do we need to train a separate, complicated system to predict when to stop?

They tried two approaches:

The "Free" Rule (Reading the Mind): They set up a simple rule: "If the robot's guess for a word doesn't change much between loops, stop processing that word and move to the next." This rule required no extra training and no extra brainpower.
The "Trained" Rule (Guessing the Mind): They tried to teach a separate, simple computer program to look at the robot's state and predict, "Ah, this word is done!" This is like hiring a manager to tell the workers when to clock out.

The Result: The "Free" rule was a huge winner. It achieved the same high quality as running the robot for a full 8 loops, but it only used an average of 4.94 loops. That's a 38% reduction in work! The "Trained" rule, however, failed to save any time. It needed to run almost the full 8 loops to get the same result.

The paper suggests that the robot's own internal signals are so clear that trying to learn a new way to predict them is actually a waste of time. The robot tells you exactly when it's done; you just have to listen.

What This Means (and What It Doesn't)

The researchers are careful to say this isn't a magic speed-up button for every computer yet. They measured the "average depth" (how many loops on average), which is a good guess for how much work was saved, but they didn't build a special system to actually run the loops faster on a real computer chip yet. So, while the math says they saved 38% of the work, the actual time saved on a computer screen might be less until engineers build better systems to handle these "early exits."

Also, they found that the robot stops getting smarter once it hits the number of loops it was trained on. If you trained it to average 8 loops, it stops improving after 8 loops. It doesn't magically get better if you force it to run 32 loops. This suggests that the "smartness" of the robot is set during its training, not just by how long you let it think.

The Takeaway

This paper shows that in these looping AI models, every word has its own "thinking speed." Some words are solved quickly, while others need more time. By simply watching the robot's brain and stopping the loop the moment a word settles, we can save a massive amount of energy without losing any intelligence. And the best part? We don't need to teach the robot how to stop; it already knows when it's done. We just need to let it tell us.

Technical Summary: Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

Problem Statement
Standard fixed-depth transformers bind model quality to a single compute operating point. Serving multiple quality/cost tiers requires training and storing separate models, while adapting compute to input difficulty at inference forfeits the prompt cache. Depth-recurrent transformers (applying a weight-tied core a variable number of times) offer a theoretical solution, allowing a single parameter set to operate across a range of depths. However, prior work often conflates two distinct properties: depth robustness (producing useful output at untrained depths) and depth allocation (determining how much depth each specific token requires). Existing approaches to per-token allocation suffer from significant limitations: early-exit methods on non-shared stacks fail due to key-value cache inconsistencies; mixture-of-recursions methods fix the compute budget in the weights during pretraining, preventing post-hoc retargeting; and jointly trained halting policies often collapse to degenerate depths or require delicate tuning. Crucially, no prior work has directly measured the per-token computation of a trained elastic model to determine if a simple, training-free policy could outperform learned ones.

Methodology
The authors propose a two-stage approach separating depth robustness from depth allocation:

Depth Robustness (Pretraining): A randomized-depth backbone is trained using a minimal recipe. At each optimizer step, a recursion count $r$ is sampled from a log-normal distribution and applied to the whole step. The model uses standard cross-entropy loss with no halting units, ponder loss, or FLOP penalties. The architecture consists of a unique prelude, a weight-tied core (with input re-injection and sandwich normalization to prevent state collapse), and a unique coda.
Depth Allocation (Post-hoc Policy): Once the backbone is frozen, three policies are evaluated on the same model to determine per-token depth:
- A0 (Uniform Depth): Applies a fixed loop count $r$ to all tokens (baseline).
- A1 (Convergence Exit): A training-free rule that halts a token once the KL divergence between successive output distributions ( $KL(p_i || p_{i-1})$ ) falls below a threshold $\epsilon$ .
- A2 (Oracle-Labeled Router): A linear probe trained on the model's own hidden states to predict whether a token's argmax prediction will change in future loops.

Key Contributions

Direct Per-Token Measurement: The study provides the first characterization of per-token convergence in a depth-recurrent language model. It demonstrates that the recurrent state converges to a per-token fixed point. While the median token converges by loop six, a tail of approximately 10% of tokens continues to update at the training-mean depth of eight. Convergence depth is ordered by token type (whitespace converges shallowest, content words deepest).
Reading vs. Learning Convergence: The paper demonstrates that directly reading the model's state dynamics outperforms learning to predict them. The training-free convergence exit (A1) achieves uniform depth-8 quality at an average of 4.94 loops (a 38% reduction in average depth). In contrast, the linear router (A2), trained on convergence labels harvested from the same model, requires nearly full depth (7.99 loops) and yields no reduction.
Training-Mean Saturation Hypothesis: Supported by internal data and external consistency checks, the authors hypothesize that a recursive model's useful depth saturates at its training-time mean loop distribution. Models trained with a mean of 4 saturate at depth 4; those with a mean of 8 saturate at depth 8.

Results

Elasticity: The randomized-depth backbone exhibits monotonic quality improvement from depth 1 to 8, after which validation loss stabilizes (3.20 nats) even up to 32 loops. This confirms the model is depth-robust without needing specific tuning for each depth.
Convergence Dynamics: Mean successive-output KL divergence drops from $3.9 \times 10^{-1}$ at loop 2 to $8.5 \times 10^{-6}$ at loop 16. The state change decays in step, confirming the system reaches a fixed point rather than spinning uselessly.
Allocation Efficiency: The training-free A1 policy matches the quality of a uniform depth-8 model at 4.94 average loops. It outperforms the uniform baseline at shallow depths and matches it across the average-depth range. The learned router (A2) fails to provide savings, requiring nearly full depth to match the baseline.
Token-Level Behavior: Structural tokens (whitespace) exit earliest, while content tokens compute longer. The relative ordering of numerals and words flips between datasets (TinyStories vs. FineWeb-Edu), attributed to distribution shifts rather than architectural failure.
Efficiency Bounds: While the study reports a 38% reduction in average depth (a FLOP proxy), it explicitly notes that realized wall-clock speedup depends on depth-wise continuous batching, which was not implemented. The best-case wall-clock speedup is estimated at 1.48–1.51x, while the current harness shows 0% realized speedup due to masking rather than skipping work.

Significance and Claims
The paper claims that the "elasticity" of depth-recurrent transformers is a stable property that can be leveraged via simple, training-free policies. The central finding is that per-token variation in convergence depth is directly readable and exploitable, rendering complex learned halting policies unnecessary and less effective. The authors position this work as a measurement and allocation study rather than a contribution to the underlying elastic backbone (which they attribute to prior work like Huginn and LoopFormer).

The study is modest in its scope:

Results are established at a single scale (135M-class parameters) and a single seed (for the main model).
No FLOP-matched parity is claimed against standard transformers; comparisons are made against a ceiling reference (SmolLM2-135M) which has a 150x token count advantage.
Realized speedup is not reported as a final metric but as a bracket between FLOP reduction and best-case batching potential.
The hypothesis that saturation depth equals training-mean depth is presented as an open hypothesis supported by limited data points, not a proven law.

The complete study was reproducible on a single RTX 4090 in approximately 100 GPU-hours.