This repository contains the code, raw predictions, logs, figures, and summary
statistics from a set of experiments that test whether latent chain-of-thought
tokens propagate logical constraints across curriculum stages.
It is a companion artifact to the main code repo
Avra98/curriculum-cot-code
and the three checkpoint repos referenced below.
The headline question:
When a Sudoku cell-policy model is trained with a 3-stage curriculum
(S1 → enumerate legal values; S2 → narrow with one-step look-ahead;
S3 → commit to the unique solution), do latent CoT tokens allow the
Stage-3 model to retain the constraint information from Stages 1 and 2,
in a way that a vanilla data-curriculum baseline cannot?
Spoiler: yes, and the effect grows with cell difficulty. The latent ATC
model keeps its Stage-3 commitment inside its own nested earlier-stage
candidate sets — the full chain Ŝ3 ⊆ Ŝ2 ⊆ Ŝ1 is preserved on 97.0% of
cells, versus 92.3% for the data-curriculum baseline (DC). When DC makes
mistakes at Stage 3, ~10% of those mistakes are catastrophic (a digit that
is not even in the cell's true candidate set); for ATC the catastrophic rate is
~1% — an ~8× reduction. The gap is larger on harder cells.
Two figures carry the story:
figs/fig_constraint_propagation.png — the flagship. The Stage-3
commitment stays inside the model's own S1/S2 candidate sets (constraint
propagates across stages) for ATC, and increasingly does not for DC as
cells get harder.
figs_overnight/fig_latent_budget_sweep.png — the causal control. Knock
the latent budget k = num_cot_tokens past its trained value and accuracy
collapses off a cliff (0.94 → 0.04). The latent tokens are load-bearing, not
inert padding.
TL;DR — headline numbers
All numbers are over a 100-puzzle eval set × ~20 empty cells / puzzle =
2000 cells evaluated per checkpoint.
metric
ATC (latent CoT + curriculum)
DC (vanilla data curriculum)
Full nested chain Ŝ3 ⊆ Ŝ2 ⊆ Ŝ1 (all cells)
0.970
0.923
↳ on medium cells (`
S1
=2`)
↳ on hard cells (`
S1
=3`)
Containment Ŝ3 ⊆ Ŝ2 (averaged over all cells)
0.971
0.924
Stage-3 solve rate, `
target S1
= 2` (medium)
Stage-3 solve rate, `
target S1
= 3` (hard)
Stage-3 errors (out of 2000 cells)
76
107
Among wrong S3 cells: pred ⊆ Ŝ1
0.99
0.90
Among wrong S3 cells: catastrophic (pred ∩ Ŝ1 = ∅)
0.01
0.10
Latent-budget cliff (exact-set-match at k_train → k_train+1..2)
All adapters are LoRA on top of Qwen/Qwen2.5-1.5B-Instruct with
r=32, alpha=64, dropout=0.05. The latent checkpoints additionally use
num_cot_tokens ∈ {1, 3} recurrent-hidden tokens (see
code/predict_one.py).
Evaluation set
sudoku_t3_20empty_value_qwen_text_stage1_eval.jsonl from the training
pipeline (100 puzzles). Each puzzle has ~20 empty cells expanded into a
per-cell prediction problem, giving 2000 (puzzle, cell) instances per
checkpoint.
Ground-truth targets
For every empty cell we precompute three "stage-i consistent" candidate
sets directly from the grid:
target_S1 = digits that satisfy row/column/box uniqueness (legal moves).
target_S2 = target_S1 further pruned by one-step look-ahead
consistency.
target_S3 = { unique solution } (always a singleton).
These three sets are written into every per-cell record in preds/ so the
analysis is self-contained.
Experiments
Experiment 0 — flagship: constraint propagation across stages
This is the figure the whole repo is built around. Ground truth is a nested
chain target_S3 ⊆ target_S2 ⊆ target_S1 (commit ⊆ look-ahead-legal ⊆
row/col/box-legal). We test whether each model's own predictions form the
same chain Ŝ3 ⊆ Ŝ2 ⊆ Ŝ1 — i.e. whether the constraint information survives
being carried from one curriculum stage to the next.
constraint propagation
figs/fig_constraint_propagation.{pdf,png} (produced by
code/analyze_propagation.py). Three panels:
(a) Candidate-set funnel — mean predicted set size at S1 → S2 → S3,
overlaid on the ground-truth funnel. Both methods funnel from ~1.27
candidates to a singleton, but ATC tracks the ground-truth narrowing at S2
(1.00 vs GT 1.02) whereas DC over-retains (1.07).
(b) Nested-chain links preserved — each link of the chain as a
probability. Ŝ2⊆Ŝ1 ≈ 1.0 for both, but the Ŝ3⊆Ŝ2 link is where DC
leaks (0.924 vs ATC 0.971), and the joint chain Ŝ3⊆Ŝ2⊆Ŝ1 is
0.970 (ATC) vs 0.923 (DC).
(c) Full chain preserved vs difficulty — the divergence is concentrated
on harder cells: on |S1|=2 cells ATC preserves the chain 89.6% of the time
versus DC's 71.8% (an +17.8 pp gap).
Every number is in figs/propagation_summary.json. Standalone versions of the
funnel and the difficulty panel are also saved as
figs/fig_propagation_funnel.* and figs/fig_chain_preservation.*.
Each of the 6 adapters is evaluated on its own training stage (so the
output should match its target_S<train>). This gives the "marginal"
solve rates per stage:
S1
S2
S3
ATC exact-set-match
0.998
0.944
0.961
DC exact-set-match
0.990
0.873
0.945
(Computed by code/analyze.py from preds/.)
Experiment 2 — cross-stage containment
For every cell that has predictions from all 6 (method, stage) adapters,
we check the nested-set properties of the same model's predictions
across the three curriculum stages:
Ŝ3 ⊆ Ŝ1 — "Stage-3 commits to a value that was already enumerated
by Stage 1 of the same model."
Ŝ3 ⊆ Ŝ2 — "Stage-3 commits to a value the Stage-2 model also flagged
as look-ahead-consistent."
Ŝ3 ∩ Ŝ1 = ∅ — catastrophic rewrite: Stage-3 picks a digit that
isn't even in Stage-1's candidate set.
figs/fig_containment_basic.{pdf,png}. Both methods preserve broad
legality (Ŝ3 ⊆ Ŝ1 ≈ 99% for both), but ATC is 2.8× less likely to
violate Stage-2 consistency at Stage-3, and 6.5× less likely to do so
catastrophically against Stage-1.
Experiment 3 — difficulty-stratified solve rate
The 2000 cells are bucketed by |target_S1| (the size of the true
legal-moves set — our proxy for cell difficulty):
solve by difficulty
figs/fig_solve_by_diff.{pdf,png}. The ATC advantage roughly doubles
as cells get harder: +1 pp on easy singletons, +3.6 pp on medium cells
(|S1|=2), +8.7 pp on hard cells (|S1|=3).
Companion plots: fig_c13_by_diff, fig_c23_by_diff, fig_d23_by_diff
show the same breakdown for the containment metrics.
Experiment 4 — failure-mode taxonomy
Among the cells where the Stage-3 model is WRONG, where did it go wrong?
failure modes
figs/fig_correctness_breakdown.{pdf,png}. ATC has fewer wrong cells in
absolute terms (76 vs 107). Conditional on being wrong:
ATC stays inside Ŝ1 99% of the time (its mistakes are "narrow
failures" within the legal set).
DC stays inside Ŝ1 only 90% of the time and picks an illegal value
10% of the time when it errs.
Experiment 5 — set-size trajectory
set sizes
figs/fig_set_size_trajectory.{pdf,png}. Average predicted set size at
S1 / S2 / S3 per method. Both methods commit (size → 1) by Stage 3, but
ATC commits earlier (smaller average at S2: 1.00 vs 1.07).
Experiment 6 — per-cell Sankey visualisation
sankey
figs/fig_sankey_example.{pdf,png}. For a single example puzzle, shows
the S1 → S2 → S3 trajectory of predicted candidates per cell, side by
side for ATC and DC. Red rows are cells where Ŝ3 ⊄ Ŝ1 (containment
violation).
Experiment 7 — method agreement
method agreement
figs/fig_method_agreement.{pdf,png}. For each difficulty bucket: how
often do ATC.S3 and DC.S3 produce the same commit, and within those
disagreements, which method is correct? Agreement falls off sharply with
difficulty (99% on |S1|=1, 78% on |S1|=2, 54% on |S1|=3), and ATC
is the more accurate method on every bucket.
For each of the 6 adapters we additionally evaluated with an
off-diagonal stage_i prompt — e.g. the ATC S3 adapter asked to
enumerate target_S1. This produces a 2-methods × 3-train-stages ×
3-prompt-stages = 18-cell matrix.
These off-diagonal evaluations live in preds_xprompt/. The set is
filled in by code/run_cross_prompt.sh and
code/run_cross_prompt_phase2.sh. The matrix and the "did the prompt change
anything?" heatmap are produced by code/analyze_cross_prompt.py into
figs_xprompt/.
All 12 off-diagonals are now in (figs_xprompt/xprompt_summary.json), and
the result is a clean negative that sharpens the story. Look at the drift
table: when an adapter trained to commit at S3 is prompted to enumerate S1,
it ignores the prompt 99.75% of the time (frac_ignored_prompt ≈ 0.99–1.00
for every off-diagonal cell, for both methods). So the differentiator is not
"ATC retains the ability to enumerate when asked" — a trained-to-commit model
commits regardless of prompt. The differentiator is what value it commits
to: ATC.S3 asked for S1 lands inside the legal set 99.95% of the time
(subset_of_prompt_target), DC.S3 only 99.40%. The constraint structure lives
in the latent state carried across stages, not in a re-promptable surface
behaviour.
Experiment 9 — latent-budget (k) sweep: are the latent tokens load-bearing?
The propagation result is correlational on its own: maybe ATC just happens to
commit to legal values. Experiment 9 is the causal control. Each latent
adapter was trained with a fixed number of recurrent-hidden CoT tokens
(k_train: S1 → 1, S2/S3 → 3). At inference we sweep
k = num_cot_tokens ∈ {0,1,2,3,4,5,6,8,10,12} and re-measure exact-set-match.
latent budget sweep
figs_overnight/fig_latent_budget_sweep.{pdf,png} (produced by
code/analyze_ksweep.py; raw numbers in
figs_overnight/ksweep_summary.json). Accuracy is flat at near-peak up to the
trained budget and then falls off a cliff within one or two extra tokens:
adapter
k_train
exact @ k_train
exact one/two steps past
ATC S1
1
0.998
0.333 @ k=3 → 0.000 @ k≥4
ATC S2
3
0.944
0.037 @ k=4 → 0.000 @ k≥5
If the latent tokens were inert padding you would expect a flat line or a
graceful slope. A discontinuous collapse keyed to k_train is the signature of
the latent tokens carrying load-bearing computation. (The S3 sweep and the
off-diagonal × k perturbations are still draining from the overnight queue;
this figure auto-refreshes as preds_overnight/ fills in.)
Each JSONL line in preds/ and preds_xprompt/ has the same schema:
jsonc
1{
2 "method_tag": "atc_s2", // free-form id
3 "puzzle_id": 42, // 0-based row index in eval set
4 "target_cell": [3, 7], // 0-based (row, col)
5 "target_solution": 5, // unique correct digit
6 "stage_prompted": 2, // the stage_i passed to the prompt
7 "predicted_values":[5], // model's parsed output (sorted ints)
8 "predicted_text": "5\n", // raw decoded text
9 "parse_ok": true,
10 "exact_set_match": true, // pred == target_S{stage_prompted}
11 "target_S1": [3, 5, 8], // gt legal candidate set
12 "target_S2": [5, 8], // gt after 1-step lookahead
13 "target_S3": [5] // gt solution set
14}
The presence of target_S1/S2/S3 on every record is what makes the
cross-stage analysis trivial to run offline — no need to recompute legal
moves from the grid in any analysis script.
Reproduce
bash
1# 1. clone this repo + the curriculum-cot-code repo (for shared utilities)2git clone https://huggingface.co/Avra98/sudoku-cross-stage-experiments
3git clone https://huggingface.co/Avra98/curriculum-cot-code
4cd curriculum-cot-code
5ln -s ../sudoku-cross-stage-experiments/code _experiments/cross_stage_repo
67# 2. fetch the 6 LoRA adapters (selective LFS pull from the model repos)8git clone https://huggingface.co/Avra98/sudoku-latent-recurrent-hidden-20empty-stages
9git clone https://huggingface.co/Avra98/sudoku-9x9-20empty-baseline-1p5b-sweep
10# git lfs pull --include "<adapter subdir>" inside each as needed1112# 3. run the 6 diagonal evals (one per GPU)13bash _experiments/cross_stage_repo/run_all.sh
1415# 4. (optional) fill in the 12-cell off-diagonal matrix16bash _experiments/cross_stage_repo/run_cross_prompt.sh
17bash _experiments/cross_stage_repo/run_cross_prompt_phase2.sh
1819# 5. produce all figures20python _experiments/cross_stage_repo/analyze_v2.py \21 --preds_dir ./preds --out_dir ./figs
22python _experiments/cross_stage_repo/analyze_propagation.py \23 --preds_dir ./preds --out_dir ./figs # flagship Exp 024python _experiments/cross_stage_repo/analyze_cross_prompt.py \25 --diag_dir ./preds --xprompt_dir ./preds_xprompt --out_dir ./figs_xprompt
26python _experiments/cross_stage_repo/analyze_ksweep.py \27 --preds_dir ./preds_overnight --out_dir ./figs_overnight # Exp 9 latent-budget sweep
A single L4 / A10 / H100 finishes one diagonal eval in 12–45 min,
depending on whether it's the vanilla baseline (~12 min) or the latent
recurrent-hidden adapter (~45 min). The 18 cells can be run in parallel
across 8 GPUs in ~1 hour wall-clock.
Why this matters for the paper
The "ATC vs Data Curriculum" comparison would otherwise just be a single
solve-rate number on the final test set. These cross-stage experiments
let us tell a mechanistic story:
Both methods learn to commit to a single digit at Stage 3.
But ATC's commits are systematically constraint-aware — they live
inside the same model's earlier-stage candidate sets. The full nested
chain Ŝ3 ⊆ Ŝ2 ⊆ Ŝ1 holds on 97.0% of cells for ATC vs 92.3% for DC,
and the gap widens to +17.8 pp on medium-difficulty cells (Exp 0).
The DC baseline's Stage-3 model has effectively "overwritten" the
earlier-stage candidate structure, and pays for it with a 10×
catastrophic-error rate.
The cross-prompt experiment (Exp 8) rules out the boring explanation:
the structure is not a re-promptable surface behaviour (models ignore
the stage prompt ~99.75% of the time); it lives in the latent state.
The latent-budget sweep (Exp 9) shows the latent tokens are causally
load-bearing — accuracy collapses off a cliff the moment you perturb
k = num_cot_tokens past its trained value.
Together this is exactly the picture you would predict if the latent CoT
tokens were carrying logical constraint information forward across
curriculum stages.
