FPSA — Spatial-Reasoning Experiments

Iterated, weight-tied self-attention (FPSA, Fixed-Point Self-Attention) applied to hard spatial-reasoning benchmarks, and benchmarked against tiny recursive reasoners (TRM).

Instead of stacking L independent self-attention layers, a single attention block is iterated to refine a hidden state. This repo applies that idea to grid puzzles: Sudoku-Extreme, Maze-Hard, and ARC.

Architecture (high level)

A single weight-tied attention block is applied k_max times over a running state z:

z_{k+1} = (1 - α) · z_k  +  α · W_O · Attn(Q(z_k), K(z_k), V)

• Weight-tied iteration — num_layers blocks unrolled k_max times give a large effective depth with O(1) parameter cost (default L=4, k_max=32 → effective depth 128).
• Per-head learnable α — each head learns its own update / mixing rate.
• 2D-RoPE — rotary position embeddings over (row, col) so spatial structure survives iteration.
• Value stream — how V is fed into the loop is the main design axis (see variants below).
• Stability controls — optional spectral norm / bounded Lipschitz on the Q/K/O and V paths (--v_constraint {none,spectral,bounded}) to keep the iteration well-behaved.
• Solvers — standard finite unroll, or a DEQ fixed-point path (--deq --deq_solver {picard,anderson}) with implicit (IFT/adjoint) gradients.
• Training recipe — optional TRM-style deep supervision (--deep_sup S: readout + loss after each of S segments, state detached between them) and DDP multi-GPU training.

Value-stream variants

vanilla (single-pass baseline) · fixed (V computed once outside the loop) · evolving (V recomputed each step) · scalar_blend · blended (learned per-token gate between the input and evolving streams — highest ceiling) · fixed_ffn.

Repository layout

• src/ — reference FPSA components (attention, iterated block, RoPE, implicit-gradient autograd).
• experiments/ — the spatial FPSA model, per-task trainers, the DDP TRM-comparison trainer, and dynamics diagnostics.
• scripts/ — end-to-end runners.

Experiments

• Value-stream ablation — which variant best refines a grid state under iteration; blended is the strongest but also the most sensitive to the contraction/expressivity trade-off.
• Scaling — widening the model (hidden_dim) toward TRM's parameter class on full-augmented Sudoku-Extreme, to separate a capacity bottleneck from a recipe bottleneck.
• TRM recipe — does TRM-style deep supervision lift FPSA beyond plain training?
• Seed robustness — multi-seed reproducibility of the scaled result.
• Dynamics diagnostics — is the model a true fixed-point solver or finite-depth iterative refinement? (residuals, churn, margin, test-time k-extrapolation, implicit-gradient fidelity.)
• DEQ — Picard vs Anderson fixed-point forward solve with implicit differentiation.

Benchmarks follow the TRM data format; TRM numbers are cited from the paper (no on-hardware TRM run).

Running

# single-task train (see each train_*.py for flags)
python experiments/sudoku_ablation/train_trm_compare.py --task sudoku_extreme \
    --variant blended --hidden_dim 320 --num_heads 10 --num_layers 4 --k_max 32
# multi-GPU (DDP)
torchrun --standalone --nproc_per_node=4 \
    experiments/sudoku_ablation/train_trm_compare.py --task sudoku_extreme --variant blended ...