RecursiveThink

A weekend research experiment on a laptop GPU

Current transformers spend identical compute on every token regardless of difficulty. A period at the end of a sentence gets the same forward pass as a critical step in a mathematical proof. Humans don't work this way — you can feel yourself "thinking harder" on difficult problems, allocating more mental effort where it's needed. What if a transformer could learn to do the same?

This project started as a conversation about transformer limitations and turned into a hands-on experiment: can a small transformer learn to iteratively refine its internal representations, spending more compute on harder problems and less on easy ones?

The short answer: yes, but only under specific conditions — and the results reveal something interesting about when and why adaptive computation emerges.

Take a standard small decoder-only transformer and add a Recurrent Refinement Block — a single transformer layer that can loop. Instead of passing through once and moving on, this block iterates, refining its hidden state with each pass. A learned halting mechanism (inspired by Adaptive Computation Time) decides when to stop: a small network reads the hidden state and outputs a halting probability. When cumulative probability crosses a threshold, the block stops iterating and passes its output to the remaining layers.

A ponder cost penalty in the loss function discourages unnecessary iteration — the model has to earn each additional thinking step by improving its output.

We tested two variants against a baseline:

• baseline: Standard transformer with fixed layers — same compute for every token
• recursive_replace: One base layer is replaced by the recurrent block (same unique layers, one can loop)
• recursive_insert: The recurrent block is added as an extra layer (one more layer than baseline, and it can loop)

All experiments ran on a single NVIDIA RTX A2000 8GB Laptop GPU. Training runs took 30–90 minutes each. The models are tiny — 4 layers, 128-dimensional hidden states, ~5M parameters. The tasks are synthetic: multi-digit addition (1–8 digits), multiplication (1–3 digits), and bracket matching with variable nesting depth.

We used PyTorch with mixed-precision training, MLflow for experiment tracking, and generated all training data synthetically with stratified sampling across difficulty levels.

The most consistent finding across all experiments: the model learns to iterate more on harder problems, but only when it's capacity-constrained.

Both recursive variants learn difficulty-correlated iteration counts when the base model is small (4 layers, 128 hidden dim). Easy 1-digit problems get 1–2 iterations; hard 6-digit problems get 8–12 iterations.

This finding was remarkably sensitive to model size. When we ran the same experiment with a larger base model (6 layers, 256 hidden dim), the recursive_insert variant collapsed to 1 iteration on everything — it had enough fixed-depth capacity to handle the task without needing recurrence. The ponder cost penalty won, and the recurrent block became a dead layer.

Same experiment, bigger model. recursive_insert (green) flatlines at 1 iteration — the model has enough depth to brute-force the task. recursive_replace still iterates somewhat on multiplication, but far less than in the capacity-constrained setting.

With a larger model, the baseline matches or beats both recursive variants. When the model has enough fixed depth, recurrence adds overhead without benefit.

The key insight: adaptive computation is a substitute for depth, and models only learn to use it when they're short on depth. Give a model enough fixed layers and it will brute-force the task through memorized patterns rather than learning to deliberate. Starve it of layers and recurrence becomes the escape valve.

When the model does learn to iterate adaptively, the accuracy gains on hard problems are substantial.

Addition accuracy by digit count. At 6 digits, recursive_insert achieves ~95% sequence accuracy vs ~8% for the baseline. Both recursive variants dramatically outperform a same-parameter-count baseline on hard in-distribution problems.

The training curves tell the same story: both recursive models converge faster and to higher final validation accuracy than the baseline.

Both recursive variants reach higher validation accuracy while the baseline plateaus earlier.

The iteration heatmaps are perhaps the most visually striking result. They show how many iterations the recurrent block uses at each token position for different problems.

recursive_insert heatmap. Simple problems like 0+6=6 barely register. Complex problems like 1700970+8052126=9753096 light up intensely, with iteration counts reaching 20+ at critical positions. The model concentrates computation on the positions where the actual arithmetic happens.

recursive_replace shows similar patterns with even higher iteration counts, reflecting its greater need for recurrence due to fewer base layers.

The model isn't just applying blanket "more compute" to hard problems. It's learning position-specific deliberation — spending more iterations on positions where carries propagate or where the output digits are being computed, while largely ignoring padding and input tokens.

One unexpected finding: both recursive models show a performance dip at intermediate difficulty levels, specifically around 5-digit addition, even though they excel at both easier and harder problems.

Looking at the iteration counts explains why. Around difficulty 4–5, the models are transitioning between a low-iteration strategy (pattern matching, works for 1–4 digits) and a high-iteration strategy (deliberation, works for 6+ digits). At the transition point, the model allocates extra iterations but hasn't yet learned to use them effectively at that specific difficulty. It's like a student who recognizes a problem is hard and stares at it longer, but hasn't developed the right strategy for the extra time yet.

recursive_replace hits this wall earlier (difficulty 4–5) because it has fewer base layers and needs to switch to high-iteration mode sooner. recursive_insert hits it later (difficulty 5–6). By the next difficulty level, the high-iteration strategy has clicked and accuracy surges.

This suggests that learning when to think harder and learning how to think harder effectively are somewhat separate skills that develop at different rates during training.

We tested whether the adaptive computation mechanism could generalize to unseen difficulty levels — training on 1–6 digit addition and evaluating on 7–10. The result was unambiguous: all models, including the recursive variants, dropped to near-zero accuracy on unseen lengths.

The iteration-vs-difficulty plots showed something interesting though: when encountering OOD problems, the recursive_insert model didn't try to iterate more — it actually reduced iterations, as if recognizing "I can't productively use more thinking on this." A form of learned uncertainty, perhaps, but not the OOD generalization we hoped for.

The model learned to detect difficulty and allocate compute within its training distribution, but didn't learn a generalizable algorithm (like "iterate once per carry") that could extend to novel lengths. This is consistent with the broader finding in the field that length generalization remains a hard problem for transformers.

The central finding is that learned adaptive computation is real and useful, but conditional. It emerges when models are capacity-constrained, it improves accuracy on hard in-distribution problems substantially, and it learns fine-grained position-specific deliberation patterns. But it doesn't solve generalization, and the recurrent block's shared weights across iterations limit what each additional pass can contribute.

The deeper insight is about the relationship between recurrence and depth. In current practice, we make models deeper (more layers) to handle harder tasks. This project suggests an alternative: make models shallower but give them the ability to dynamically loop, trading fixed depth for adaptive depth. A 4-layer model with a recurrent block that can iterate 20 times on hard problems achieves accuracy that a 4-layer fixed model can't match — and it does so while breezing through easy problems in 1–2 iterations, saving compute where it isn't needed.

However, there's a gap between "the model learns to iterate" and "each iteration is maximally productive." The performance dip at intermediate difficulties, and the failure to generalize OOD, both suggest that the shared-weight recurrent block isn't learning rich, step-by-step algorithms. It's more like a refinement process — each pass polishes the representation slightly — than true multi-step reasoning where each step builds qualitatively on the last.

Curriculum learning. We observed that adaptive computation emerged more readily with skewed data distributions (mostly easy examples, some hard ones) than with balanced data. This suggests the model needs to "master the basics first, then discover it can iterate to handle harder cases" — a natural curriculum.

Non-shared weights across iterations. The biggest limitation may be that each iteration uses the same weights. If different iterations could specialize — first pass for coarse understanding, second for carries, third for verification — the quality of deliberation could improve dramatically.

Multiple recurrent blocks at different layers. Our experiments used a single recurrent block at one position. Multiple blocks could learn different types of deliberation at different levels of abstraction.

Larger models on harder tasks. Testing on natural language tasks where difficulty varies (simple factual recall vs. multi-step reasoning) would be the real test of whether adaptive computation can scale.

# Install dependencies (requires uv)
uv sync

# Run a quick smoke test (baseline + 2 recursive, ~5 min)
uv run python sweep.py sweeps/smoke_test.yaml

# Or train a single model
uv run python train.py --model_type recursive_replace --task arithmetic

# Evaluate
uv run python evaluate.py --task arithmetic

# Generate plots
uv run python visualize.py --task arithmetic

# MLflow dashboard
uv run mlflow ui --backend-store-uri sqlite:///mlruns.db

uv run python train.py --model_type baseline --task arithmetic --epochs 50

All config fields are CLI flags — see config.py for the full list, or run:

uv run python train.py --help

./run_experiment.sh arithmetic 50

This trains baseline, recursive_insert, and recursive_replace in sequence, running evaluation and plot generation after each.

# Preview what will run
uv run python sweep.py sweeps/weekend_arithmetic.yaml --dry-run

# Run the sweep (sequential, with resume support)
uv run python sweep.py sweeps/weekend_arithmetic.yaml

Sweep configs are YAML files defining a cartesian grid:

name: my_sweep
fixed:
  task: arithmetic
  epochs: 50
extra_runs:
  - model_type: baseline
grid:
  model_type: [recursive_replace, recursive_insert]
  recursive_layer_idx: [1, 2, 3, 4]
  act_lambda: [0.001, 0.005, 0.01]

Resume support: if a sweep is interrupted, re-running the same command skips already-finished runs (checked via MLflow).

# Default: evaluate the 3 standard model types
uv run python evaluate.py --task arithmetic

# Sweep mode: discover and evaluate all *_best.pt checkpoints
uv run python evaluate.py --sweep --task arithmetic

# Single checkpoint
uv run python evaluate.py --checkpoint checkpoints/my_run_best.pt --task arithmetic

config.py                  # Frozen dataclass with all hyperparameters
train.py                   # Training loop with MLflow logging
evaluate.py                # Per-difficulty evaluation with OOD breakdown
visualize.py               # Accuracy, iteration, and sweep plots
sweep.py                   # YAML-driven hyperparameter sweep runner
run_experiment.sh          # Run all 3 models end-to-end

models/
  baseline.py              # Standard decoder-only transformer
  recursive_think.py       # RecursiveThinkTransformer (replace/insert modes)
  recursive_block.py       # RecursiveBlock with ACT halting

data/
  arithmetic.py            # Addition & multiplication dataset
  brackets.py              # Bracket depth prediction dataset

sweeps/                    # Sweep YAML configs
docs/                      # Architecture diagram and experiment plots

All experiments log to MLflow (SQLite backend):

uv run mlflow ui --backend-store-uri sqlite:///mlruns.db

Logged metrics include training loss, validation accuracy, ponder cost, per-difficulty accuracy, and per-difficulty iteration statistics for recursive models.

All experiments were conducted on a single NVIDIA RTX A2000 8GB Laptop GPU with CUDA 12.2. Training runs ranged from 30–90 minutes

This project was a collaboration between a human researcher and Claude (Anthropic), from initial architecture discussion through experimental design, debugging, and analysis.