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 |

Seeds: 3 per family (2 for fixed-4-long). All checkpoints in `ckpts/`.

## 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 |

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) means
   single-seed comparisons of length generalisation are unreliable here.

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).