README.md

---
tags:
- ml-intern
- research
- recurrent-depth
- length-generalization
- pointer-chasing
library_name: vanilla-pytorch
---

[![Built with ML Intern](https://huggingface.co/datasets/huggingface/badges/resolve/main/built-with-ml-intern-sm.svg)](https://hf.co/chat?mode=ml-intern)

# Do recurrent transformers learn length-generalising algorithms for pointer chasing?

**Status: COMPLETE.** 17 trained models, full evaluation sweeps, figures and
the technical report are in this repo. See **[REPORT.md](REPORT.md)** for the
full write-up, `results/results_table.csv` for per-seed numbers, and
`plots/` for figures. Live training dashboards:
[eac123/ptr-recurrent-depth-trackio](https://huggingface.co/spaces/eac123/ptr-recurrent-depth-trackio).

## Research question

Can recurrent depth (a weight-tied block applied R times) give better
out-of-distribution generalisation than a conventional fixed-depth transformer,
at matched parameter and training-FLOP budgets, on synthetic multi-hop
reasoning (predict f^k(x) given a random mapping f)?

## Task

Each example: a fresh uniform-random mapping f: S → S, a start state x, and a
hop count k. The answer is f^k(x). Training: k ~ U{1..8}, S = 64. Evaluation:
fresh mappings, k up to 64 (OOD), smaller state spaces S' ∈ {16, 32}, and
inference-time unroll counts R ∈ 1..16 (more or fewer than trained).
Nothing from training is reused at eval; the deterministic LookupSolver
reference scores 1.0 by construction and validates the pipeline every run.

Shortcuts are controlled: f is uniform random per example (no counter-free
structure, cf. arXiv:2210.10749); answer marginals are ~uniform per k
(verified, TV < 0.02); k is the only length-correlated token.

## Models

One implementation (`ptr/models.py`): `depth` pre-LN blocks applied `n_iters`
times with per-iteration step embeddings and zero-initialised output
projections (stability for long unrolls, cf. arXiv:2604.07822).
d_model=128, 4 heads, d_ff=512, AdamW lr=1e-3 (chosen by a 6-run sweep),
cosine decay, batch 256, 25k steps (100k for fixed4-long).

| family | blocks × unrolls | params | total training FLOPs |
|---|---|---|---|
| fixed-4 | 4 × 1 | 829,696 | 2.19e15 |
| rec-2×2 | 2 × 2 | 433,152 | 2.19e15 |
| rec-1×4 | 1 × 4 | 234,880 | 2.19e15 |
| rec-1×dynamic | 1 × R~U{2..8} | 234,880 | 2.73e15 |
| rec-4×4 | 4 × 4 | 829,696 | 8.68e15 |
| fixed-4-long | 4 × 1, 100k steps | 829,696 | 8.74e15 |
| rec-1×aligned | 1 × k (R = k) | 240,256 | 2.46e15 |
| rec-1×m4 | 1 × (k+4) | 242,048 | 4.62e15 |
| rec-1×m4+dense | 1 × (k+4), stem, dense loss | 440,320 | 4.62e15* |

Seeds: 3 per family (2 for fixed-4-long). All checkpoints in `ckpts/`.
(*The m4+dense row's stored FLOPs under-count its stem by ~12%.)

## Headline results (mean over seeds, exact-match accuracy)

| family | acc k=1..8 | acc k=64 (OOD) | best k=64 over inference R |
|---|---|---|---|
| fixed-4 | 0.487±0.003 | 0.102±0.012 | 0.123 |
| rec-2×2 | 0.580±0.052 | 0.177±0.037 | 0.195 |
| rec-1×4 | 0.761±0.147 | 0.107±0.062 | 0.115 |
| rec-1×dynamic | 0.899±0.008 | 0.103±0.036 | 0.195 |
| rec-4×4 | **1.000** | 0.251±0.099 | 0.269 |
| fixed-4-long (iso-FLOP) | 0.716±0.009 | 0.171±0.018 | 0.172 |
| rec-1×aligned (R=k) | 0.106±0.003 | 0.017±0.003 | 0.139 |
| rec-1×m4 (R=k+4) | 0.262±0.003 | 0.029±0.017 | 0.196 |
| rec-1×m4+dense | **1.000±0.000** | 0.075±0.032 | **0.437** |

Key observations (empirical):

1. **At matched parameters**, rec-4×4 reaches 1.000 in-distribution while
   fixed-4 plateaus at 0.487 — and fixed-4-long still only reaches 0.716
   after 4× the total training FLOPs. Serial block-applications, not
   parameters, are the binding resource.
2. **Length generalisation beyond k=8 is weak everywhere** (max ~0.27 mean
   at k=64): no architecture in this budget class learned a clean
   length-generalising algorithm. OOD accuracy is nonetheless 6–16× chance.
3. **More inference-time unrolls than trained sometimes help a lot** —
   rec-4×4 seed2: 0.394→0.401 at R=5; rec-1×dynamic seed1: 0.088→0.319 at
   R=5 — but the effect is seed-dependent, not systematic at this scale.
4. **Large seed variance in OOD** (rec-4×4 k=64: 0.18/0.18/0.39;
   m4+dense k=64: 0.116/0.038/0.070) means single-seed comparisons of
   length generalisation are unreliable here.
5. **The v2 ablation (§4.7-4.8) decomposes what makes aligned training
   work**: R=k with end-of-chain loss fails (0.106 — credit assignment
   across depths); a +4 unroll margin alone only partially helps (0.262);
   adding an untied stem block and dense per-iteration supervision with
   hold-after-k targets f^min(t,k)(x) fully rescues it (1.000), and its
   fixed-R sweep shows a staircase (+0.634) — the model genuinely does
   one hop per unroll. **But even this model does not extend past k=8**
   (best k=64 over R: 0.437) — the frontier tracks the training-data
   limit, making curriculum/data coverage the most promising next lever.

See REPORT.md for the full analysis, figures, and limitations.

## Prior work

- Merrill & Sabharwal, "The Parallelism Tradeoff" — https://arxiv.org/abs/2305.15401
- Li et al., "Chain of Thought Empowers Transformers to Solve Inherently Serial Problems" — https://arxiv.org/abs/2402.12875
- Liu et al., "Transformers Learn Shortcuts to Automata" — https://arxiv.org/abs/2210.10749
- Anil et al., "Exploring Length Generalization in Large Language Models" — https://arxiv.org/abs/2207.04901
- Dehghani et al., "Universal Transformers" — https://arxiv.org/abs/1807.03819
- Yang et al., "Looped Transformers are Better at Learning Learning Algorithms" — https://arxiv.org/abs/2310.00197
- Fan et al., "Looped Transformers for Length Generalization" — https://arxiv.org/abs/2409.15647
- Bae et al., "Relaxed Recursive Transformers" — https://arxiv.org/abs/2410.20672
- Li et al., "Loop, Think, & Generalize: Implicit Reasoning in Recurrent-Depth Transformers" — https://arxiv.org/abs/2604.07822
- Ye et al., "Transformers Struggle to Learn to Search" — https://arxiv.org/abs/2412.04603

## Reproduce

```bash
python validate.py                       # architecture validation suite
python run_all.py lrsweep --hub-repo eac123/ptr-recurrent-depth
python run_all.py stage1 --hub-repo eac123/ptr-recurrent-depth
python run_all.py stage2 --hub-repo eac123/ptr-recurrent-depth
python run_all.py eval   --hub-repo eac123/ptr-recurrent-depth   # recovery
python make_report.py                                            # figures+report
```

Dependencies: torch==2.14.1, numpy==2.5.2, huggingface_hub==1.33.0,
trackio==0.42.0, matplotlib==3.11.2 (report only).
NOTE: `run_all.py lrsweep/stage1/stage2` each run their own training;
stage1/stage2 fetch `results/chosen_lr.json` from the Hub, so the lrsweep
must complete first (or run them with `--hub-repo` unset for a fully local
pipeline with `chosen_lr.json` in `out/`).

## Compute

~$9 of HF Jobs compute on t4-small ($0.40/hr): GPU smoke test, lr sweep
(6 runs), stage1 (12 runs, 4h19m), stage2 (5 runs, 7h10m), two recovery
eval passes, report generation. Two duplicate jobs created by interrupted
tool calls were cancelled (~$0.30 wasted, reported not hidden).