eac123/ptr-recurrent-depth

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

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

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