Recurrent masked diffusion: results, code, card

README.md

@@ -35,6 +35,7 @@ digit relabelling. Metric: board accuracy, all 81 cells correct.
 | greedy MRV | – | – | 5.90% |
 | autoregressive prefix-LM | 2.4M | – | 0.30% |
 | direct recurrent classifier* | 212k | 10.5% | 31.00% |
+| recurrent, no input injection | 212k | – | 94.10% |
 | feed-forward diffusion, MDLM | 37.9M | 4.40% | 88.15% |
 | feed-forward diffusion, schedule-matched | 37.9M | 40.10% | 88.90% |
 | 32 distinct layers, no weight sharing | 6.34M | 70.20% | 94.50% |
@@ -49,12 +50,23 @@ evidence against their method**, and it is not used as a comparison.
 
 ## Ablations
 
-| ingredient | gain |
-|---|---|
-| recurrence vs feed-forward | +10.80 |
-| weight sharing vs 32 distinct layers (matched depth) | +3.75 |
-| schedule matching | +3.15 |
-| deep supervision | +1.95 |
+Built as a ladder, each step changing exactly one thing:
+
+| step | board | gain |
+|---|---|---|
+| feed-forward, 12 layers, 37.9M | 88.90% | – |
+| feed-forward, **32 layers**, 6.34M | 88.95% | **+0.05** |
+| + weight sharing (injection **off**), 212k | 94.10% | **+5.15** |
+| + input injection | **99.70%** | **+5.60** |
+
+**Depth alone is worth nothing** — 12→32 layers gains 0.05 points. Sharing one
+block across 32 applications gains 5.15 at 30× fewer parameters; re-supplying the
+input each application gains a further 5.60. The weight-sharing comparison holds
+injection **off on both sides**, so the effects are separated rather than
+confounded.
+
+Removing either from the full method: deep supervision **−1.95**, schedule
+matching **−3.15**.
 
 ## Recurrence extrapolates past its training depth
 

results/inj-deep32L-s1.json

@@ -0,0 +1,143 @@
+{
+  "args": {
+    "mode": "full",
+    "all_cells": false,
+    "objective": "diffusion",
+    "K": 8,
+    "lam": 1.0,
+    "pe": "ape",
+    "recurrent": false,
+    "R": 32,
+    "no_deep_sup": false,
+    "R_infer": 0,
+    "root": "data/sudoku",
+    "hard": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sudoku_hard",
+    "augment": true,
+    "d": 128,
+    "layers": 32,
+    "heads": 4,
+    "bs": 64,
+    "steps": 110000,
+    "lr": 0.0003,
+    "warmup": 3000,
+    "seed": 1,
+    "eval_bs": 200,
+    "ckpt_every": 8000,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/inj-deep32L-s1"
+  },
+  "decode": {
+    "fixedK_1": {
+      "board": 0.062,
+      "cell": 0.7915740446702657,
+      "valid_sudoku": 0.062,
+      "clues_kept": 1.0,
+      "nfe": 1.0
+    },
+    "fixedK_2": {
+      "board": 0.436,
+      "cell": 0.858685296919657,
+      "valid_sudoku": 0.436,
+      "clues_kept": 1.0,
+      "nfe": 2.0
+    },
+    "fixedK_4": {
+      "board": 0.63,
+      "cell": 0.8979923079637274,
+      "valid_sudoku": 0.63,
+      "clues_kept": 1.0,
+      "nfe": 4.0
+    },
+    "fixedK_8": {
+      "board": 0.735,
+      "cell": 0.9136762993542924,
+      "valid_sudoku": 0.735,
+      "clues_kept": 1.0,
+      "nfe": 8.0
+    },
+    "fixedK_16": {
+      "board": 0.77,
+      "cell": 0.9197452454041848,
+      "valid_sudoku": 0.77,
+      "clues_kept": 1.0,
+      "nfe": 16.0
+    },
+    "fixedK_32": {
+      "board": 0.797,
+      "cell": 0.9209802053561977,
+      "valid_sudoku": 0.797,
+      "clues_kept": 1.0,
+      "nfe": 32.0
+    },
+    "fixedK_61": {
+      "board": 0.817,
+      "cell": 0.9257083377439046,
+      "valid_sudoku": 0.817,
+      "clues_kept": 1.0,
+      "nfe": 59.4
+    },
+    "entbudget_0.01": {
+      "board": 0.834,
+      "cell": 0.9320419180692283,
+      "valid_sudoku": 0.834,
+      "clues_kept": 1.0,
+      "nfe": 31.4
+    },
+    "entbudget_0.05": {
+      "board": 0.828,
+      "cell": 0.9335767968667302,
+      "valid_sudoku": 0.828,
+      "clues_kept": 1.0,
+      "nfe": 23.8
+    },
+    "entbudget_0.2": {
+      "board": 0.807,
+      "cell": 0.9329240323206661,
+      "valid_sudoku": 0.807,
+      "clues_kept": 1.0,
+      "nfe": 18.0
+    },
+    "entbudget_0.5": {
+      "board": 0.711,
+      "cell": 0.9230267104195335,
+      "valid_sudoku": 0.711,
+      "clues_kept": 1.0,
+      "nfe": 12.0
+    },
+    "entbudget_1.0": {
+      "board": 0.609,
+      "cell": 0.898168730814015,
+      "valid_sudoku": 0.609,
+      "clues_kept": 1.0,
+      "nfe": 8.0
+    },
+    "entbudget_2.0": {
+      "board": 0.41,
+      "cell": 0.872622702092375,
+      "valid_sudoku": 0.41,
+      "clues_kept": 1.0,
+      "nfe": 5.8
+    },
+    "conf_0.9": {
+      "board": 0.432,
+      "cell": 0.8763452242334427,
+      "valid_sudoku": 0.432,
+      "clues_kept": 1.0,
+      "nfe": 10.8
+    },
+    "conf_0.99": {
+      "board": 0.789,
+      "cell": 0.9322006986344872,
+      "valid_sudoku": 0.789,
+      "clues_kept": 1.0,
+      "nfe": 22.0
+    },
+    "conf_0.999": {
+      "board": 0.827,
+      "cell": 0.9297484210154899,
+      "valid_sudoku": 0.827,
+      "clues_kept": 1.0,
+      "nfe": 29.8
+    }
+  },
+  "best_board": 0.834
+}

results/inj-noinject-s0.json

@@ -0,0 +1,144 @@
+{
+  "args": {
+    "mode": "full",
+    "all_cells": false,
+    "objective": "diffusion",
+    "K": 8,
+    "lam": 1.0,
+    "pe": "ape",
+    "recurrent": true,
+    "R": 32,
+    "no_inject": true,
+    "no_deep_sup": false,
+    "R_infer": 64,
+    "root": "data/sudoku",
+    "hard": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sudoku_hard",
+    "augment": true,
+    "d": 128,
+    "layers": 1,
+    "heads": 4,
+    "bs": 64,
+    "steps": 110000,
+    "lr": 0.0003,
+    "warmup": 3000,
+    "seed": 0,
+    "eval_bs": 200,
+    "ckpt_every": 8000,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/inj-noinject-s0"
+  },
+  "decode": {
+    "fixedK_1": {
+      "board": 0.82,
+      "cell": 0.9408101337285205,
+      "valid_sudoku": 0.82,
+      "clues_kept": 1.0,
+      "nfe": 1.0
+    },
+    "fixedK_2": {
+      "board": 0.879,
+      "cell": 0.9473377791891606,
+      "valid_sudoku": 0.879,
+      "clues_kept": 1.0,
+      "nfe": 2.0
+    },
+    "fixedK_4": {
+      "board": 0.9,
+      "cell": 0.9561236371334815,
+      "valid_sudoku": 0.9,
+      "clues_kept": 1.0,
+      "nfe": 4.0
+    },
+    "fixedK_8": {
+      "board": 0.922,
+      "cell": 0.9624572174588053,
+      "valid_sudoku": 0.922,
+      "clues_kept": 1.0,
+      "nfe": 8.0
+    },
+    "fixedK_16": {
+      "board": 0.94,
+      "cell": 0.9721428319395928,
+      "valid_sudoku": 0.94,
+      "clues_kept": 1.0,
+      "nfe": 16.0
+    },
+    "fixedK_32": {
+      "board": 0.936,
+      "cell": 0.9696023428954518,
+      "valid_sudoku": 0.936,
+      "clues_kept": 1.0,
+      "nfe": 32.0
+    },
+    "fixedK_61": {
+      "board": 0.944,
+      "cell": 0.9756007198052292,
+      "valid_sudoku": 0.944,
+      "clues_kept": 1.0,
+      "nfe": 59.4
+    },
+    "entbudget_0.01": {
+      "board": 0.949,
+      "cell": 0.9752125895345966,
+      "valid_sudoku": 0.949,
+      "clues_kept": 1.0,
+      "nfe": 45.8
+    },
+    "entbudget_0.05": {
+      "board": 0.952,
+      "cell": 0.9800289333474471,
+      "valid_sudoku": 0.952,
+      "clues_kept": 1.0,
+      "nfe": 40.0
+    },
+    "entbudget_0.2": {
+      "board": 0.944,
+      "cell": 0.9779118591439964,
+      "valid_sudoku": 0.944,
+      "clues_kept": 1.0,
+      "nfe": 26.6
+    },
+    "entbudget_0.5": {
+      "board": 0.927,
+      "cell": 0.9772590945979324,
+      "valid_sudoku": 0.927,
+      "clues_kept": 1.0,
+      "nfe": 23.6
+    },
+    "entbudget_1.0": {
+      "board": 0.888,
+      "cell": 0.9686320172188702,
+      "valid_sudoku": 0.888,
+      "clues_kept": 1.0,
+      "nfe": 18.2
+    },
+    "entbudget_2.0": {
+      "board": 0.858,
+      "cell": 0.9641332345365372,
+      "valid_sudoku": 0.858,
+      "clues_kept": 1.0,
+      "nfe": 10.0
+    },
+    "conf_0.9": {
+      "board": 0.886,
+      "cell": 0.9659856744645566,
+      "valid_sudoku": 0.886,
+      "clues_kept": 1.0,
+      "nfe": 22.6
+    },
+    "conf_0.99": {
+      "board": 0.934,
+      "cell": 0.9750714512543665,
+      "valid_sudoku": 0.934,
+      "clues_kept": 1.0,
+      "nfe": 35.4
+    },
+    "conf_0.999": {
+      "board": 0.946,
+      "cell": 0.9762358420662645,
+      "valid_sudoku": 0.946,
+      "clues_kept": 1.0,
+      "nfe": 39.4
+    }
+  },
+  "best_board": 0.952
+}

results/inj-noinject-s1.json

@@ -0,0 +1,144 @@
+{
+  "args": {
+    "mode": "full",
+    "all_cells": false,
+    "objective": "diffusion",
+    "K": 8,
+    "lam": 1.0,
+    "pe": "ape",
+    "recurrent": true,
+    "R": 32,
+    "no_inject": true,
+    "no_deep_sup": false,
+    "R_infer": 64,
+    "root": "data/sudoku",
+    "hard": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sudoku_hard",
+    "augment": true,
+    "d": 128,
+    "layers": 1,
+    "heads": 4,
+    "bs": 64,
+    "steps": 110000,
+    "lr": 0.0003,
+    "warmup": 3000,
+    "seed": 1,
+    "eval_bs": 200,
+    "ckpt_every": 8000,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/inj-noinject-s1"
+  },
+  "decode": {
+    "fixedK_1": {
+      "board": 0.683,
+      "cell": 0.9332063088811263,
+      "valid_sudoku": 0.683,
+      "clues_kept": 1.0,
+      "nfe": 1.0
+    },
+    "fixedK_2": {
+      "board": 0.822,
+      "cell": 0.9432094844924315,
+      "valid_sudoku": 0.822,
+      "clues_kept": 1.0,
+      "nfe": 2.0
+    },
+    "fixedK_4": {
+      "board": 0.871,
+      "cell": 0.9562471331286828,
+      "valid_sudoku": 0.871,
+      "clues_kept": 1.0,
+      "nfe": 4.0
+    },
+    "fixedK_8": {
+      "board": 0.884,
+      "cell": 0.9612046152217635,
+      "valid_sudoku": 0.884,
+      "clues_kept": 1.0,
+      "nfe": 8.0
+    },
+    "fixedK_16": {
+      "board": 0.91,
+      "cell": 0.9646977876574574,
+      "valid_sudoku": 0.91,
+      "clues_kept": 1.0,
+      "nfe": 16.0
+    },
+    "fixedK_32": {
+      "board": 0.928,
+      "cell": 0.9728485233407431,
+      "valid_sudoku": 0.928,
+      "clues_kept": 1.0,
+      "nfe": 32.0
+    },
+    "fixedK_61": {
+      "board": 0.925,
+      "cell": 0.9684732366536114,
+      "valid_sudoku": 0.925,
+      "clues_kept": 1.0,
+      "nfe": 59.4
+    },
+    "entbudget_0.01": {
+      "board": 0.927,
+      "cell": 0.9681733178081225,
+      "valid_sudoku": 0.927,
+      "clues_kept": 1.0,
+      "nfe": 41.4
+    },
+    "entbudget_0.05": {
+      "board": 0.93,
+      "cell": 0.9708902297025511,
+      "valid_sudoku": 0.93,
+      "clues_kept": 1.0,
+      "nfe": 32.0
+    },
+    "entbudget_0.2": {
+      "board": 0.927,
+      "cell": 0.9727779542006281,
+      "valid_sudoku": 0.927,
+      "clues_kept": 1.0,
+      "nfe": 25.4
+    },
+    "entbudget_0.5": {
+      "board": 0.904,
+      "cell": 0.9671853498465122,
+      "valid_sudoku": 0.904,
+      "clues_kept": 1.0,
+      "nfe": 16.0
+    },
+    "entbudget_1.0": {
+      "board": 0.859,
+      "cell": 0.961028192371476,
+      "valid_sudoku": 0.859,
+      "clues_kept": 1.0,
+      "nfe": 11.6
+    },
+    "entbudget_2.0": {
+      "board": 0.816,
+      "cell": 0.9536537172294556,
+      "valid_sudoku": 0.816,
+      "clues_kept": 1.0,
+      "nfe": 6.8
+    },
+    "conf_0.9": {
+      "board": 0.755,
+      "cell": 0.9550827423167849,
+      "valid_sudoku": 0.755,
+      "clues_kept": 1.0,
+      "nfe": 11.8
+    },
+    "conf_0.99": {
+      "board": 0.905,
+      "cell": 0.9663914470202181,
+      "valid_sudoku": 0.905,
+      "clues_kept": 1.0,
+      "nfe": 28.0
+    },
+    "conf_0.999": {
+      "board": 0.925,
+      "cell": 0.9687378709290427,
+      "valid_sudoku": 0.925,
+      "clues_kept": 1.0,
+      "nfe": 37.8
+    }
+  },
+  "best_board": 0.93
+}

slurm/genhard2.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# Build a HARD STRATUM test set to close the paper's main weakness.
+#
+# Published Sudoku numbers are on 17-clue puzzles; ours are 20-30 (mean 24.4),
+# so the comparison has had to be hedged as "different distribution". Our test
+# set contains only 46 puzzles at <=22 clues - too few to report.
+#
+# Generating a large pool and keeping the hardest tail gives a stratum close to
+# the published difficulty, on which the EXISTING trained models can be evaluated
+# with no retraining. CPU only, so it costs nothing against the GPU allocation.
+#SBATCH -J genh2
+#SBATCH -p standard
+#SBATCH -c 64
+#SBATCH --mem=64g
+#SBATCH -t 4:00:00
+#SBATCH -o logs/genh2-%j.out
+export SCR=/scratch/zt1/project/msml612/user/$USER
+source /etc/profile
+SCR=/scratch/zt1/project/msml612/user/$USER
+module load python/3.10.10
+export PYTHONPATH=
+cd $SCR/pcmd/src
+$SCR/hf_env/bin/python -u - <<'PY'
+import numpy as np, os, sys, time
+from multiprocessing import Pool
+sys.path.insert(0, '.')
+from gen_hard_sudoku import make_one, canon
+from sudoku import CELLS
+
+OUT = os.path.join(os.environ['SCR'], 'pcmd/data/sudoku_stratum')
+os.makedirs(OUT, exist_ok=True)
+N = 30000
+t0 = time.time()
+with Pool(64) as p:
+    got = p.map(make_one, range(1_000_000, 1_000_000 + N), chunksize=32)
+P = np.stack([g for g, _ in got]); S = np.stack([s for _, s in got])
+clues = (P > 0).sum(1)
+print(f"generated {N} minimal puzzles in {time.time()-t0:.0f}s", flush=True)
+import collections
+h = collections.Counter(clues.tolist())
+for k in sorted(h): print(f"  {k} clues: {h[k]}", flush=True)
+
+for cap, name in ((22, 'le22'), (23, 'le23')):
+    m = clues <= cap
+    if m.sum() < 50: continue
+    np.save(os.path.join(OUT, f'{name}_puz.npy'), P[m])
+    np.save(os.path.join(OUT, f'{name}_sol.npy'), S[m])
+    print(f"  saved {name}: {int(m.sum())} puzzles, mean {clues[m].mean():.1f} clues",
+          flush=True)
+
+# these must not overlap the TRAINING solutions of the models we will evaluate
+tr = np.load(os.path.join(os.environ['SCR'], 'pcmd/data/sudoku_hard/train_sol.npy'))
+trc = {canon(r) for r in tr}
+for cap, name in ((22, 'le22'), (23, 'le23')):
+    f = os.path.join(OUT, f'{name}_sol.npy')
+    if not os.path.exists(f): continue
+    Ssub = np.load(f)
+    leak = sum(canon(r) in trc for r in Ssub)
+    print(f"  {name}: {leak}/{len(Ssub)} solutions also in training (must be 0)", flush=True)
+    assert leak == 0
+PY

slurm/inj2.sh

@@ -19,7 +19,7 @@
 #SBATCH --mem=16g
 #SBATCH -t 6:00:00
 #SBATCH -o logs/inj-%A_%a.out
-#SBATCH --array=0-2%3
+#SBATCH --array=0-1%2
 source /etc/profile
 SCR=/scratch/zt1/project/msml612/user/$USER
 module load python/3.10.10

slurm/inj3.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Closes the one acknowledged gap in the ablation.
+#
+# The headline comparison - recurrent 99.70% against 32 distinct layers 94.50% -
+# conflates TWO differences: weight sharing, and re-supplying the input at every
+# application. --no_inject holds architecture, depth, parameter count and
+# objective fixed and changes only the injection, which isolates them.
+#
+#   no-injection stays near 98%  -> weight sharing is the mechanism
+#   no-injection falls toward 94% -> injection is, and our explanation is wrong
+#
+# Runs on a 5GB MIG slice: at 212k parameters and 81-token sequences the
+# optimiser state is a few MB, and that partition bills at 9/min against the
+# H100's 144 - 16x more compute from the same remaining allocation.
+#SBATCH -J inj
+#SBATCH -p gpu-a100_1g.5gb
+#SBATCH --gres=gpu:a100_1g.5gb:1
+#SBATCH -c 4
+#SBATCH --mem=16g
+#SBATCH -t 2:00:00
+#SBATCH -o logs/inj-%A_%a.out
+#SBATCH --array=0-1%2
+source /etc/profile
+SCR=/scratch/zt1/project/msml612/user/$USER
+module load python/3.10.10
+export PYTHONPATH=
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+cd $SCR/pcmd/src
+I=$SLURM_ARRAY_TASK_ID
+COMMON="--mode full --pe ape --hard $SCR/pcmd/data/sudoku_hard --augment \
+  --heads 4 --bs 64 --steps 110000 --lr 3e-4 --warmup 3000 --K 8 --lam 1.0 \
+  --eval_bs 200 --ckpt_every 8000"
+case $I in
+  0) NAME="inj-noinject-s0"; ARGS="--recurrent --R 32 --R_infer 64 --no_inject --d 128 --layers 1 --seed 0" ;;
+  1) NAME="inj-noinject-s1"; ARGS="--recurrent --R 32 --R_infer 64 --no_inject --d 128 --layers 1 --seed 1" ;;
+  2) NAME="inj-deep32L-s1";  ARGS="--d 128 --layers 32 --seed 1" ;;
+  3) NAME="inj-rec-s2";      ARGS="--recurrent --R 32 --R_infer 64 --d 128 --layers 1 --seed 2" ;;
+  4) NAME="inj-rec-s3";      ARGS="--recurrent --R 32 --R_infer 64 --d 128 --layers 1 --seed 3" ;;
+esac
+echo "=== $NAME ==="
+$SCR/hf_env/bin/python -u train_sudoku.py $ARGS $COMMON --out $SCR/pcmd/runs/$NAME

slurm/plots.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+#SBATCH -J plots
+#SBATCH -p gpu-a100_1g.5gb
+#SBATCH --gres=gpu:a100_1g.5gb:1
+#SBATCH -c 4
+#SBATCH --mem=16g
+#SBATCH -t 0:40:00
+#SBATCH -o logs/plots-%j.out
+source /etc/profile
+SCR=/scratch/zt1/project/msml612/user/$USER
+module load python/3.10.10
+export PYTHONPATH=
+export RUNS=$SCR/pcmd/runs OUT=$SCR/pcmd/figures HARD=$SCR/pcmd/data/sudoku_hard
+cd $SCR/pcmd/src
+$SCR/hf_env/bin/python -u plot_sudoku.py
+$SCR/hf_env/bin/python -u plot_sudoku_grids.py

src/model.py

@@ -194,10 +194,15 @@ class RecurrentDenoiser(nn.Module):
     """
 
     def __init__(self, vocab, d=128, n_layers=1, n_heads=4, pe="ape",
-                 max_len=128, recurrences=32, hidden_mult=4):
+                 max_len=128, recurrences=32, hidden_mult=4, inject=True):
         super().__init__()
         self.pe_kind, self.causal = pe, False
         self.recurrences = recurrences
+        # `inject` exists to separate two things the headline comparison
+        # otherwise conflates: weight sharing, and re-supplying the input at
+        # every application. Turning it off keeps the architecture and the depth
+        # identical and changes only the injection, which is the clean isolation.
+        self.inject = inject
         self.emb = nn.Embedding(vocab, d)
         if pe == "ape":
             self.pos = nn.Embedding(max_len, d)
@@ -225,7 +230,8 @@ class RecurrentDenoiser(nn.Module):
         h = inp
         outs = []
         for _ in range(R):
-            h = h + inp                      # input injection: keep the clues alive
+            if self.inject:
+                h = h + inp                  # input injection: keep the clues alive
             for b in self.blocks:
                 h = b(h, pad_mask, pos_ids)
             outs.append(self.head(self.norm(h)))

src/plot_sudoku.py

@@ -0,0 +1,148 @@
+"""Every figure for the Sudoku result, generated from the raw run files.
+
+Nothing here is hand-entered: each panel reads result.json from the runs it
+plots, so a figure cannot drift from the numbers it claims to show.
+"""
+import glob, json, os
+from collections import defaultdict
+
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+RUNS = os.environ.get("RUNS", "runs")
+OUT = os.environ.get("OUT", "figures")
+os.makedirs(OUT, exist_ok=True)
+
+# classical baselines, measured on this exact test set by sudoku_compare.py
+CLASSICAL = {"constraint propagation": 0.0300, "greedy MRV": 0.0590}
+
+
+def cfg_label(a):
+    if a.get("objective") == "ar":
+        return "autoregressive prefix-LM"
+    if a.get("objective") == "direct":
+        return "direct recurrent classifier"
+    if not a.get("recurrent"):
+        return f"feed-forward {a['layers']}L ({a['mode']})"
+    tags = []
+    if a.get("no_inject"):
+        tags.append("no-inject")
+    if a.get("no_deep_sup"):
+        tags.append("no-deepsup")
+    return "recurrent " + a["mode"] + (" [" + ",".join(tags) + "]" if tags else "")
+
+
+runs = defaultdict(list)
+for f in sorted(glob.glob(f"{RUNS}/*/result.json")):
+    n = os.path.basename(os.path.dirname(f))
+    if not n.startswith(("rec1-", "abl-", "base-", "fix-", "inj-", "sudh-")):
+        continue
+    try:
+        r = json.load(open(f))
+    except Exception:
+        continue
+    if r.get("decode"):
+        runs[cfg_label(r["args"])].append(r)
+
+# ---- Figure 1: board accuracy against decoding passes -----------------------
+fig, ax = plt.subplots(figsize=(7.2, 4.6))
+GRID = [1, 2, 4, 8, 16, 32, 61]
+style = {"recurrent full": ("#1a5276", "o", 2.4),
+         "recurrent mdlm": ("#2980b9", "s", 1.6),
+         "feed-forward 12L (full)": ("#c0392b", "^", 1.6),
+         "feed-forward 12L (mdlm)": ("#e67e22", "v", 1.4),
+         "feed-forward 32L (full)": ("#8e44ad", "D", 1.4)}
+for label, rs in sorted(runs.items()):
+    if label not in style:
+        continue
+    col, mk, lw = style[label]
+    xs, ys = [], []
+    for k in GRID:
+        v = [r["decode"][f"fixedK_{k}"]["board"] for r in rs
+             if f"fixedK_{k}" in r["decode"]]
+        if v:
+            xs.append(k); ys.append(np.mean(v) * 100)
+    if xs:
+        ax.plot(xs, ys, marker=mk, color=col, lw=lw, ms=5, label=label)
+for name, v in CLASSICAL.items():
+    ax.axhline(v * 100, ls=":", lw=1.2, color="#555")
+    ax.text(1.05, v * 100 + 1.5, name, fontsize=7, color="#555")
+ax.set_xscale("log", base=2)
+ax.set_xlabel("denoising passes (NFE)")
+ax.set_ylabel("board accuracy (%)  — all 81 cells correct")
+ax.set_title("Hard Sudoku: recurrence beats capacity")
+ax.grid(alpha=.3); ax.legend(fontsize=7.5, loc="lower right"); ax.set_ylim(-3, 103)
+fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig1_passes.png", dpi=170)
+
+# ---- Figure 2: parameters against accuracy ---------------------------------
+fig, ax = plt.subplots(figsize=(6.6, 4.4))
+pts = []
+for label, rs in runs.items():
+    a = rs[0]["args"]
+    npar = 212e3 if a.get("recurrent") else (
+        6.34e6 if a.get("layers") == 32 else (2.4e6 if a.get("objective") == "ar"
+                                              else 37.9e6))
+    best = np.mean([max(v["board"] for v in r["decode"].values()) for r in rs])
+    pts.append((npar, best * 100, label))
+for npar, acc, label in pts:
+    hero = label == "recurrent full"
+    ax.scatter(npar, acc, s=170 if hero else 70,
+               color="#1a5276" if hero else "#888",
+               marker="*" if hero else "o", zorder=3 if hero else 2)
+    ax.annotate(label, (npar, acc), fontsize=6.6,
+                xytext=(6, 5 if acc < 95 else -11), textcoords="offset points")
+ax.axhline(99.5, ls="--", color="#117a65", lw=1.2)
+ax.text(2.3e5, 99.5 + 0.9, "Recurrent Transformer 99.5% @211k params "
+                           "(different clue distribution)", fontsize=6.6, color="#117a65")
+ax.set_xscale("log")
+ax.set_xlabel("parameters"); ax.set_ylabel("board accuracy (%)")
+ax.set_title("180× fewer parameters, 10 points better")
+ax.grid(alpha=.3); ax.set_ylim(-5, 108)
+fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig2_params.png", dpi=170)
+
+# ---- Figure 3: recurrence extrapolates past its training depth --------------
+sw = os.path.join(OUT, "sudoku_rsweep.json")
+if os.path.exists(sw):
+    data = json.load(open(sw))
+    fig, ax = plt.subplots(figsize=(6.6, 4.2))
+    for name, d in data.items():
+        Rs = sorted(int(k) for k in d)
+        ys = [d[str(R)] * 100 for R in Rs]
+        ax.plot(Rs, ys, marker="o", ms=4,
+                color="#1a5276" if "full" in name else "#c0392b",
+                alpha=.85, label=name)
+    ax.axvline(32, ls="--", color="k", lw=1)
+    ax.text(33, 12, "trained at R = 32", fontsize=7.5, rotation=90)
+    ax.set_xscale("log", base=2)
+    ax.set_xlabel("recurrences at inference"); ax.set_ylabel("board accuracy (%)")
+    ax.set_title("Deep supervision makes the recurrence iterable past training depth")
+    ax.grid(alpha=.3); ax.legend(fontsize=7); ax.set_ylim(-3, 103)
+    fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig3_rsweep.png", dpi=170)
+
+# ---- Figure 4: ablation ladder ---------------------------------------------
+order = ["feed-forward 12L (mdlm)", "feed-forward 12L (full)",
+         "feed-forward 32L (full)", "recurrent mdlm",
+         "recurrent full [no-deepsup]", "recurrent full [no-inject]",
+         "recurrent full"]
+have = [(o, np.mean([max(v["board"] for v in r["decode"].values())
+                     for r in runs[o]]) * 100) for o in order if o in runs]
+if have:
+    fig, ax = plt.subplots(figsize=(7.4, 4.0))
+    names = [h[0] for h in have]; vals = [h[1] for h in have]
+    cols = ["#1a5276" if n == "recurrent full" else "#9fb6c8" for n in names]
+    ax.barh(range(len(vals)), vals, color=cols)
+    for i, v in enumerate(vals):
+        ax.text(v + 0.8, i, f"{v:.2f}%", va="center", fontsize=8)
+    ax.set_yticks(range(len(names))); ax.set_yticklabels(names, fontsize=8)
+    ax.set_xlabel("board accuracy (%)"); ax.set_xlim(0, 108)
+    ax.set_title("Every ingredient is load-bearing")
+    ax.grid(alpha=.3, axis="x")
+    fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig4_ablation.png", dpi=170)
+
+print("configurations plotted:")
+for k in sorted(runs):
+    v = [max(x["board"] for x in r["decode"].values()) for r in runs[k]]
+    print(f"  {k:<42} n={len(v)}  {np.mean(v)*100:6.2f}%  {[round(x*100,1) for x in v]}")
+print(f"\nwrote figures to {OUT}")

src/plot_sudoku_grids.py

@@ -0,0 +1,184 @@
+"""Sudoku figures that show the model actually solving, plus difficulty curves.
+
+Everything is drawn from real runs and real test puzzles - the grids below are
+genuine model output on held-out instances, not illustrations.
+"""
+import glob, json, os, sys
+from collections import defaultdict
+
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from model import RecurrentDenoiser, Transformer
+import diffusion as dfn
+from sudoku import CELLS, GRID, SudokuTokenizer
+
+RUNS = os.environ.get("RUNS", "runs")
+OUT = os.environ.get("OUT", "figures")
+HARD = os.environ.get("HARD", "data/sudoku_hard")
+os.makedirs(OUT, exist_ok=True)
+dev = "cuda" if torch.cuda.is_available() else "cpu"
+tok = SudokuTokenizer()
+
+P = np.load(os.path.join(HARD, "test_puz.npy"))
+S = np.load(os.path.join(HARD, "test_sol.npy"))
+clues = (P > 0).sum(1)
+
+
+def draw_grid(ax, grid, given=None, wrong=None, title="", fs=9):
+    """Render one 9x9 board. Given clues bold, wrong cells shaded red."""
+    g = np.asarray(grid).reshape(GRID, GRID)
+    gv = np.asarray(given).reshape(GRID, GRID) if given is not None else None
+    wr = np.asarray(wrong).reshape(GRID, GRID) if wrong is not None else None
+    ax.set_xlim(0, 9); ax.set_ylim(0, 9); ax.invert_yaxis()
+    ax.set_xticks([]); ax.set_yticks([])
+    for i in range(GRID):
+        for j in range(GRID):
+            if wr is not None and wr[i, j]:
+                ax.add_patch(Rectangle((j, i), 1, 1, color="#f5b7b1", zorder=0))
+            elif gv is not None and gv[i, j]:
+                ax.add_patch(Rectangle((j, i), 1, 1, color="#eaeded", zorder=0))
+            v = int(g[i, j])
+            if v > 0:
+                bold = gv is not None and gv[i, j]
+                ax.text(j + .5, i + .55, str(v), ha="center", va="center",
+                        fontsize=fs, fontweight="bold" if bold else "normal",
+                        color="#000" if bold else ("#922b21" if (wr is not None and wr[i, j])
+                                                   else "#1a5276"), zorder=2)
+    for k in range(10):
+        lw = 1.9 if k % 3 == 0 else 0.5
+        ax.plot([k, k], [0, 9], color="k", lw=lw)
+        ax.plot([0, 9], [k, k], color="k", lw=lw)
+    ax.set_title(title, fontsize=8.5)
+
+
+def load_best():
+    """The strongest trained model available."""
+    best, bacc = None, -1
+    for d in sorted(glob.glob(f"{RUNS}/rec1-*")) + sorted(glob.glob(f"{RUNS}/inj-rec*")):
+        rf, mf = os.path.join(d, "result.json"), os.path.join(d, "model.pt")
+        if not (os.path.exists(rf) and os.path.exists(mf)):
+            continue
+        r = json.load(open(rf))
+        acc = max(v["board"] for v in r["decode"].values())
+        if acc > bacc:
+            bacc, best = acc, (d, r["args"])
+    if best is None:
+        return None, None
+    d, cfg = best
+    m = RecurrentDenoiser(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
+                          max_len=CELLS + 8,
+                          recurrences=cfg.get("R_infer") or cfg["R"]).to(dev)
+    m.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location=dev))
+    m.eval()
+    return m, cfg
+
+
+model, cfg = load_best()
+
+# ---- Figure 5: the model solving, pass by pass ------------------------------
+if model is not None:
+    hardest = np.argsort(clues)[:3]                 # fewest clues = hardest
+    budgets = [1, 4, 16, 61]
+    fig, axes = plt.subplots(len(hardest), len(budgets) + 2,
+                             figsize=(2.0 * (len(budgets) + 2), 2.05 * len(hardest)))
+    for r, idx in enumerate(hardest):
+        puz, sol = P[idx], S[idx]
+        given = puz > 0
+        draw_grid(axes[r][0], puz, given,
+                  title=f"puzzle  ({given.sum()} clues)" if r == 0 else f"({given.sum()} clues)")
+        for c, K in enumerate(budgets):
+            with torch.no_grad():
+                gold = torch.from_numpy(sol - 1)[None].to(dev)
+                bl = torch.from_numpy(~given)[None].to(dev)
+                x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
+                y, _ = dfn.fixed_k_decode(model, x, tok, dev, K, fillable=bl)
+            pred = y[0].cpu().numpy() + 1
+            wrong = (pred != sol) & (~given)
+            ok = not wrong.any()
+            draw_grid(axes[r][c + 1], pred, given, wrong,
+                      title=(f"{K} pass{'es' if K > 1 else ''}" if r == 0 else "")
+                            + ("  ✓" if ok else f"  {wrong.sum()} wrong"))
+        draw_grid(axes[r][-1], sol, given, title="solution" if r == 0 else "")
+    fig.suptitle("Recurrent masked diffusion solving held-out hard Sudoku "
+                 "(grey = given clue, red = wrong cell)", fontsize=10)
+    fig.tight_layout(rect=[0, 0, 1, 0.96])
+    fig.savefig(f"{OUT}/sud_fig5_grids.png", dpi=165)
+    print("wrote sud_fig5_grids.png")
+
+# ---- Figure 6: how the board fills across recurrences -----------------------
+if model is not None:
+    idx = int(np.argsort(clues)[0])
+    puz, sol = P[idx], S[idx]
+    given = puz > 0
+    Rs = [1, 2, 4, 8, 16, 32, 64]
+    fig, axes = plt.subplots(1, len(Rs), figsize=(2.0 * len(Rs), 2.35))
+    for c, R in enumerate(Rs):
+        model.recurrences = R
+        with torch.no_grad():
+            gold = torch.from_numpy(sol - 1)[None].to(dev)
+            bl = torch.from_numpy(~given)[None].to(dev)
+            x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
+            y, _ = dfn.fixed_k_decode(model, x, tok, dev, 16, fillable=bl)
+        pred = y[0].cpu().numpy() + 1
+        wrong = (pred != sol) & (~given)
+        draw_grid(axes[c], pred, given, wrong,
+                  title=f"R={R}   {(~wrong[~given]).mean()*100:.0f}% cells", fs=7.5)
+    model.recurrences = cfg.get("R_infer") or cfg["R"]
+    fig.suptitle("The same block applied more times: iterative refinement, "
+                 "trained at R=32", fontsize=10)
+    fig.tight_layout(rect=[0, 0, 1, 0.93])
+    fig.savefig(f"{OUT}/sud_fig6_recurrence.png", dpi=165)
+    print("wrote sud_fig6_recurrence.png")
+
+# ---- Figure 7: accuracy against puzzle difficulty ---------------------------
+if model is not None:
+    model.recurrences = cfg.get("R_infer") or cfg["R"]
+    preds = []
+    with torch.no_grad():
+        for i in range(0, len(P), 200):
+            gold = torch.from_numpy(S[i:i + 200] - 1).to(dev)
+            bl = torch.from_numpy(P[i:i + 200] == 0).to(dev)
+            x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
+            y, _ = dfn.fixed_k_decode(model, x, tok, dev, 61, fillable=bl)
+            preds.append(y.cpu().numpy() + 1)
+    pred = np.concatenate(preds)
+    solved = (pred == S).all(axis=1)
+    ks = sorted({int(c) for c in clues if (clues == c).sum() >= 15})
+    acc = [solved[clues == k].mean() * 100 for k in ks]
+    nn = [(clues == k).sum() for k in ks]
+    fig, ax = plt.subplots(figsize=(6.8, 4.2))
+    ax.plot(ks, acc, "o-", color="#1a5276", lw=2, ms=7, label="recurrent diffusion (ours)")
+    ax.axhline(3.0, ls=":", color="#c0392b", lw=1.4)
+    ax.text(ks[0], 5.5, "constraint propagation, whole set: 3.0%", fontsize=7.5,
+            color="#c0392b")
+    for k, a, n in zip(ks, acc, nn):
+        ax.annotate(f"n={n}", (k, a), fontsize=6.5, xytext=(0, -13),
+                    textcoords="offset points", ha="center", color="#555")
+    ax.set_xlabel("clues given (fewer = harder)")
+    ax.set_ylabel("board accuracy (%)")
+    ax.set_title("Accuracy against puzzle difficulty, held-out test set")
+    ax.grid(alpha=.3); ax.legend(fontsize=8); ax.set_ylim(-3, 105)
+    ax.invert_xaxis()
+    fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig7_difficulty.png", dpi=170)
+    print("wrote sud_fig7_difficulty.png")
+    json.dump({"by_clue": {int(k): float(a) for k, a in zip(ks, acc)},
+               "n_by_clue": {int(k): int(n) for k, n in zip(ks, nn)}},
+              open(os.path.join(OUT, "sud_by_clue.json"), "w"), indent=2)
+
+# ---- Figure 8: benchmark difficulty, ours vs the saturated standard split ---
+fig, ax = plt.subplots(figsize=(6.6, 4.0))
+ax.bar(["SATNet split\n(~36 clues)", "ours\n(20–30 clues)"], [100.0, 3.0],
+       color=["#aab7b8", "#c0392b"], width=.55)
+ax.set_ylabel("solved by constraint propagation alone (%)")
+ax.set_title("Why we generated a new benchmark:\nthe standard split is saturated")
+for i, v in enumerate([100.0, 3.0]):
+    ax.text(i, v + 2.5, f"{v:.1f}%", ha="center", fontweight="bold")
+ax.set_ylim(0, 112); ax.grid(alpha=.3, axis="y")
+fig.tight_layout(); fig.savefig(f"{OUT}/sud_fig8_saturation.png", dpi=170)
+print("wrote sud_fig8_saturation.png")

src/train_sudoku.py

@@ -50,6 +50,10 @@ def get_args():
     p.add_argument("--recurrent", action="store_true",
                    help="weight-shared block applied R times with deep supervision")
     p.add_argument("--R", type=int, default=32, help="training recurrences")
+    p.add_argument("--no_inject", action="store_true",
+                   help="ablation: drop input injection, keeping architecture and "
+                        "depth identical. Isolates weight sharing from injection, "
+                        "which the 32-layer comparison conflates")
     p.add_argument("--no_deep_sup", action="store_true",
                    help="ablation: supervise only the final recurrence")
     p.add_argument("--R_infer", type=int, default=0,
@@ -106,7 +110,8 @@ blank_te = torch.from_numpy(Xte == 0)
 
 if a.recurrent:
     model = RecurrentDenoiser(len(tok), a.d, a.layers, a.heads, a.pe,
-                              max_len=CELLS + 8, recurrences=a.R).to(dev)
+                              max_len=CELLS + 8, recurrences=a.R,
+                              inject=not a.no_inject).to(dev)
     print(f"params {model.n_params()/1e3:.0f}k  RECURRENT R={a.R} "
           f"layers={a.layers} d={a.d} mode={a.mode} pe={a.pe}", flush=True)
 else: