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

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