A Looped Transformer Router Shows Its First Replicated Gain

The useful result from this week is small, but real enough to change the research direction.

After many routed looped Transformer variants that were either unstable, collapsed, or only improved one metric while hurting another, we finally have a replicated router candidate that beats a matched fixed loop baseline on a small BPE language-model setup.

This follows an earlier negative result where a sequence-level router looked competitive but collapsed toward a low-entropy exit policy. The new result is different because the router is token-conditioned, sparse, and only active late in the recurrent computation.

The current best candidate is:

d512 / effective depth 16 / fixed_4x4 baseline
vs.
late-final-loop 4x4 token-feedback router

batch size = 8
steps = 2000
seeds = 0, 1, 2
tokenizer = GPT-2 BPE
training data = nanochat-style text
train cap = 8 MB
eval cap = 512 KB

This is not a solved architecture yet. The margins are still tiny, around 1e-3 in broad language-model loss. But the important change is that the router no longer wins only by sacrificing the general language-model objective. It is slightly better on validation loss, token accuracy, reasoning-slice loss, deterministic validation loss, easy-token loss, and fixed-defined hard-token loss at the same time.

The 8 MB train cap is intentional. This is a screening setup for architecture decisions, not evidence of pretraining-scale behavior.

The Core Idea

A normal Transformer spends depth in a straight line:

block 1 -> block 2 -> block 3 -> block 4

A looped Transformer reuses a smaller number of blocks multiple times. A 4x4 loop means:

4 unique recurrent blocks
4 recurrent passes
effective depth = 16 block applications

That already gives a useful parameter-sharing structure: fewer unique blocks, more effective computation. It does not automatically make inference cheaper, because recurrent visits still have to be executed.

The routed version asks for something stronger. Instead of forcing every token and every context through the same fixed recurrent path, can the model learn when later block state should feed back into earlier computation?

The research question is:

Can later block outputs feed back into earlier blocks, and can a router decide
when that feedback is useful based on token and context state?

Why Earlier Routers Failed

The negative results were the most useful part of the project.

The first sequence-level destination routers could sometimes beat fixed loops at medium training budgets, but the route policy often collapsed. Route entropy fell, final-exit probability rose toward one, and the router became more like a cheap exit policy than an adaptive computation policy.

Then token-level destination routers gave a better signal on hard tokens, but they often paid for it elsewhere. Some variants improved reasoning loss or fixed-defined hard-token loss while hurting generic validation loss. That made them interesting diagnostics, not architecture candidates.

The lesson was:

The router has signal, but the objective must protect generic likelihood while
letting the model spend a small amount of feedback on genuinely useful tokens.

Only winning on hard tokens is not enough. A candidate has to beat the fixed loop on the full bundle:

The Candidate That Worked

The best current router is not a free-form destination router. It is a late token-feedback router.

Instead of choosing an entirely different path at every step, it adds a small learned feedback signal late in the recurrent computation:

later block state
  -> token/context router
  -> learned source feedback
  -> add a small correction to the current hidden state

The strongest structure so far is:

fixed_4x4 backbone
4 unique recurrent blocks
4 recurrent passes
effective depth = 16

router active only in the final recurrent loop
router active only on late blocks 3 and 4
token_feedback_applications = 2

In plain language: let the fixed recurrent computation build a stable representation first, then allow a small amount of late feedback for tokens that appear to benefit from it.

That constraint matters. Earlier variants that applied feedback more broadly were noisier. The router needs to be useful, but also quiet.

Three-Seed Result

The exact promoted candidate is:

token_feedback_router_latefin4x4_learnedsrc_scale010_dynfloor000_gate005_utilbce_w005_hard030_negbce004_top010_protectbce002_top010_pos015_start300_biasm4

Matched baseline:

fixed_4x4

Three-seed aggregate:

Lower loss is better. Higher accuracy is better.

The router statistics are also important:

The router is mostly using final-ish source information, but it has not collapsed into a completely deterministic route. It learns a small source mixture and applies feedback to only a small share of token states.

That is exactly the behavior we wanted to see before scaling the experiment further.

This is a matched-backbone training-quality comparison, not a serving-cost claim. The fixed and routed models use the same 4x4 recurrent backbone and the same training setup, but the router adds extra routing machinery and auxiliary losses. A separate serving experiment would be needed to show an actual latency or FLOP advantage.

Why Batch Size Mattered

The same depth-16 direction was already slightly positive at batch size 4, but the margins were almost too small to trust:

At batch size 8, the same broad pattern became clearer:

The most interesting movement is hard-token loss: the advantage widened from about 0.0010 to about 0.0033.

My current hypothesis is that router training is more sensitive to gradient noise than the fixed baseline. The fixed model only has to learn the language-model objective. The router has to learn the language-model objective, identify useful token states, avoid damaging easy tokens, and keep the feedback policy sparse. A small batch can make those auxiliary signals too noisy.

What This Does Not Prove

This result should be framed carefully.

It does not prove that routed looped Transformers are generally better than fixed looped Transformers. It does not prove a scaling law. It does not yet prove that routing will help at larger natural-language pretraining scale. It also does not claim a compute-efficiency win at inference time.

What it does prove is narrower:

A late, sparse, token-feedback router can beat a matched fixed_4x4 looped
Transformer baseline in a replicated small-budget BPE language-model run.

That is enough to promote the design from “diagnostic experiment” to “current router candidate.”

Source Artifacts

The promoted result is tracked in the project registry:

docs/current_candidate.json

I also kept a compact project-local summary for this exact batch8 result:

docs/batch8_latefin4x4_evidence_20260706.md

The two Modal summaries behind the batch8 three-seed aggregate are:

runs/modal-downloads/nanochat_lm_bpe_feedback_latefin4x4_learnedsrc_negbce004_d512_eval4_b8_2000s_seed0_20260706/summary.json
runs/modal-downloads/nanochat_lm_bpe_feedback_latefin4x4_learnedsrc_negbce004_d512_eval4_b8_2000s_seeds12_20260706/summary.json

The earlier batch4 depth-16 aggregate is also kept in the same registry, so the batch4-to-batch8 comparison is not a separate hand calculation.

The Research Direction From Here

The next experiment should scale this exact structure before inventing a more complex router.

The decision gate should be:

1. Keep the fixed 4x4 matched baseline.
2. Scale model width or depth while preserving the late-final-loop feedback pattern.
3. Keep three-seed comparisons.
4. Require the router to win validation loss and reasoning loss, not only hard-token loss.
5. Track route statistics so a win caused by collapse does not get promoted.

The architectural lesson is also becoming clearer:

Do not let the router rewrite the computation path too early.
Let recurrence build stable state first.
Then use a small, late, token-conditioned feedback correction.

That may be the minimal structure that makes routed looped Transformers worth scaling.