| \documentclass[11pt]{article} |
| \usepackage[a4paper,margin=1in]{geometry} |
| \usepackage[utf8]{inputenc} |
| \usepackage[T1]{fontenc} |
| \usepackage{lmodern} |
| \usepackage{microtype} |
| \usepackage{amsmath,amssymb} |
| \usepackage{booktabs} |
| \usepackage{array} |
| \usepackage{longtable} |
| \usepackage{tabularx} |
| \usepackage{xcolor} |
| \usepackage{listings} |
| \usepackage{hyperref} |
| \usepackage{enumitem} |
| \usepackage{parskip} |
|
|
| \hypersetup{colorlinks=true,linkcolor=blue,urlcolor=blue,citecolor=blue} |
|
|
| \lstset{ |
| basicstyle=\ttfamily\small, |
| breaklines=true, |
| columns=fullflexible, |
| keepspaces=true, |
| frame=single, |
| framerule=0.4pt, |
| xleftmargin=0.5em, |
| xrightmargin=0.5em, |
| showstringspaces=false, |
| } |
|
|
| \newcommand{\code}[1]{\texttt{#1}} |
| \newcommand{\indic}{\mathbf{1}} |
|
|
| \title{Curriculum CoT for $9{\times}9$ Sudoku\\[2pt] |
| \large Rebuttal / Paper-Section Material} |
| \author{} |
| \date{Last updated: 2026--05--24} |
|
|
| \begin{document} |
| \maketitle |
|
|
| \noindent |
| This document is a comprehensive, paper-ready reference of (a) the data |
| pipeline, (b) the instruction-tuning prompt format, (c) the curriculum |
| and reward design, (d) the latent thought-token architecture, (e) the |
| multi-stage SFT-then-GRPO training recipe, and (f) the headline numerical |
| results --- so a rebuttal section can be assembled directly from this |
| document. |
|
|
| \bigskip |
| \hrule |
| \bigskip |
|
|
| \section{Task} |
|
|
| We use the model as a \textbf{per-cell value policy} for $9\times 9$ |
| Sudoku. For a fixed target empty cell, the model emits a JSON set of |
| candidate digits that are ``i-consistent'' with the current grid |
| (definition in \S 4). We evaluate two metrics: |
|
|
| \begin{itemize}[leftmargin=*] |
| \item \textbf{per-cell exact set match} (\code{exact\_set\_match}) --- |
| predicted set equals the ground-truth i-consistent set; |
| \item \textbf{whole-puzzle solve rate} (\code{solve}) --- every empty |
| cell on a 20-empty puzzle produces an exact set match. |
| \end{itemize} |
|
|
| Because $\text{solve} = \prod \text{exact\_set\_match}$ across the $\sim 20$ |
| empty cells of a puzzle, the two metrics are non-linearly coupled: |
| \[ |
| \text{solve} \approx \text{exact\_set\_match}^{N_{\text{empty}}} |
| \] |
| so $0.95^{20} \approx 0.358$ and $0.97^{20} \approx 0.544$ --- every |
| percentage point of per-cell exact maps to a much larger swing in solve. |
|
|
| \section{Data pipeline} |
|
|
| \subsection{Puzzle generation} |
|
|
| Generated by \code{simple\_9x9\_curriculum/build\_dataset.py}: |
|
|
| \begin{itemize}[leftmargin=*] |
| \item Start from a base Latin-square grid; randomly relabel digits, |
| permute rows and columns within bands, and transpose. |
| \item Sample \code{empties=20} cell positions uniformly at random and |
| erase them. |
| \item Save 10\,000 train + 1\,000 eval puzzles (seed 0, seed 1). |
| \item Output JSONL files \code{data/sudoku\_t3\_20empty\_value\_qwen\_text\_stage1\_\{train,eval\}.jsonl}. |
| \end{itemize} |
|
|
| A single record contains: |
|
|
| \begin{lstlisting} |
| { |
| "prompt": "<full Qwen chat-templated prompt for one (puzzle, target_cell) pair>", |
| "completion": "[7,3,8,2,6,9,4,5,...]", |
| "metadata": { |
| "grid_size": 9, "box_size": 3, "empties": 20, |
| "empty_locs_1based": [[1,4],[1,9],...], |
| "target_triples_1based": [[1,4,7],[1,9,3],...] |
| } |
| } |
| \end{lstlisting} |
|
|
| The 20 \code{target\_triples} give the \textbf{solved} value at each of |
| the 20 empty positions, so per-cell training targets are always |
| available. At training time we expand each puzzle into 20 (puzzle, |
| target\_cell) examples. |
|
|
| \subsection{Cell-policy framing} |
|
|
| The model is never asked to solve a whole puzzle in one shot. Each |
| example is one (current\_grid, target\_cell) pair, and the supervised |
| target is the set of digits that are ``i-consistent'' with the current |
| grid (see \S 4). This turns Sudoku into a |
| \textbf{classification-into-a-set} problem and lets us share parameters |
| across cells, stages, and puzzle sizes. |
|
|
| \subsection{Multi-value oversampling (data-side trick)} |
|
|
| Implemented in \code{multi\_output\_cell\_policy/sft\_multi\_output\_train.py} |
| via \code{tokenizer.\_multi\_value\_oversample\_factor} and the CLI flags |
|
|
| \begin{lstlisting} |
| --multi_value_oversample_factor INT (default 1) |
| --train_target_size_min INT (default 0) |
| --train_target_size_max INT (default 0) |
| \end{lstlisting} |
|
|
| Inside the dataset builder, examples whose target set has more than one |
| digit are repeated \code{multi\_value\_oversample\_factor} times in the |
| training mix. This biases gradient steps toward exactly the cells the |
| model gets wrong (multi-value cells). Empirically, this is the single |
| biggest data-side lever --- see \S 10. |
|
|
| \subsection{Where the bottleneck lives} |
|
|
| For 20-empty puzzles in stage 3, only $\sim 25\%$ of empty cells have a |
| multi-value target set (the rest collapse to one i-consistent value). |
| Yet those multi-value cells are responsible for the entire solve-rate |
| gap: they are the cells where the model under-predicts (returns a |
| singleton when the target is a 2- or 3-element set), and a single |
| failed cell kills the whole-puzzle solve. The reward shaping in \S 6 |
| and the oversample in 2.3 both attack this single failure mode. |
|
|
| \section{Instruction format} |
|
|
| \subsection{System prompt} |
|
|
| (verbatim from \code{multi\_output\_cell\_policy/prompt\_builder.py}) |
|
|
| \begin{lstlisting} |
| You are a Sudoku value policy. |
| This setup uses puzzles with about 20 empty cells. |
| You will be given one target empty cell. |
| Return ONLY one JSON object of the form {"values":[...]}. |
| The JSON object must contain exactly one key named "values". |
| The "values" field must be a JSON array of unique integers in [1,9]. |
| You may return as many candidate values as you want, including one, |
| several, or many values. |
| Choose the number of returned values yourself based on which values seem |
| i-consistent. |
| The order of the values does not matter. |
| Do not output any explanation, markdown, punctuation outside JSON, or |
| extra text. |
| Current stage objective: i={i} consistency. |
| \end{lstlisting} |
|
|
| \subsection{User message} |
|
|
| \begin{lstlisting} |
| Sudoku grid (0 means empty): |
| <grid_to_text(grid)> |
| Empty cells in row-major order (20 total): (1,4), (1,9), (2,8), ... |
| Target cell to fill now: (R,C). |
| Turn: t/T. |
| Return only JSON with candidate values for this target cell: {"values":[...]} |
| \end{lstlisting} |
|
|
| We use the Qwen2.5-Instruct chat template |
| (\code{tokenizer.apply\_chat\_template}, \code{add\_generation\_prompt=True}) |
| to wrap system + user into the actual prompt ids. |
| \code{max\_prompt\_length = 768}. |
|
|
| \subsection{Output format} |
|
|
| \begin{lstlisting} |
| {"values":[3,7]} |
| \end{lstlisting} |
|
|
| Strictly canonical JSON (single key \code{values}, sorted unique digit |
| list, no whitespace). Outputs are scored by \code{parse\_values\_json} |
| (\code{shared\_multi\_output\_policy.py}); any deviation collapses the |
| whole prediction to \code{parse\_ok=0} and a hard-coded malformed |
| penalty. |
|
|
| \code{max\_completion\_length = 24} tokens --- enough to emit any |
| 9-digit set. |
|
|
| \section{Curriculum: stage-i consistency} |
|
|
| The curriculum lives in \code{\_stage\_i\_consistent\_values\_for\_grid}: |
|
|
| \begin{itemize}[leftmargin=*] |
| \item \textbf{Stage 1 --- $i=1$ (legal moves).} A value $v$ is $i=1$ |
| consistent at cell $c$ iff placing $v$ at $c$ violates no Sudoku |
| constraint (row, column, $3\times 3$ box). This is just ``legal |
| candidates''. |
|
|
| \item \textbf{Stage 2 --- $i=2$.} $v$ is $i=2$ consistent at $c$ iff |
| (a) it is $i=1$ consistent AND (b) after placing $v$, every other |
| empty cell in the grid still has at least one $i=1$-consistent value |
| (i.e.\ placing $v$ does not immediately make the puzzle unsolvable |
| by 1-step propagation). |
|
|
| \item \textbf{Stage 3 --- $i=3$.} Same recursion one more level deep: |
| $v$ is $i=3$ consistent iff after placing $v$, every other empty cell |
| still has at least one $i=2$ consistent value. |
| \end{itemize} |
|
|
| This is bounded look-ahead constraint propagation. Stage-3 sets are |
| tighter than stage-2 sets which are tighter than stage-1 sets. The |
| curriculum goal at deployment time is stage-3. |
|
|
| In data, we use the same source records and just change \code{--stage\_i}; |
| the target set is regenerated on the fly by |
| \code{stage\_i\_consistent\_values}. |
|
|
| \section{Latent thought-token architecture} |
|
|
| Base model: \textbf{Qwen/Qwen2.5-1.5B-Instruct} + LoRA |
| ($r=32$, $\alpha=64$, dropout $=0.05$) on |
| \code{q,k,v,o,gate,up,down}. The latent variant adds \textbf{$k$ |
| thought-token slots} between the prompt and the next-token logits. |
|
|
| Four modes are implemented (\code{latent\_multi\_output\_cell\_policy/}); |
| the winning mode for the final number is \textbf{\code{recurrent\_hidden}}: |
|
|
| \begin{quote} |
| \code{build\_recurrent\_hidden\_latent\_hidden(model, ids, mask, k)} |
| \begin{enumerate}[leftmargin=*,nosep] |
| \item Run the backbone once on the prompt. Keep |
| \code{base\_hidden = h[:,-1,:]}. |
| \item Set \code{latent\_token = base\_hidden}. |
| \item Repeat $k$ times: append \code{latent\_token} (as an embedding) |
| to the running sequence, run the backbone again on the extended |
| sequence, and replace \code{latent\_token} with the new last hidden |
| state. |
| \item After $k$ recursions, \code{latent\_hidden} is fed through the LM |
| head to produce the next-token distribution. |
| \end{enumerate} |
| \end{quote} |
|
|
| In equations, with $E$ the input embedding lookup, $f_\theta$ the |
| LoRA-decorated backbone, $U$ the LM head: |
| \begin{align*} |
| z_0 &= f_\theta\bigl(E([x_1,\dots,x_T])\bigr)_T \\ |
| z_{j+1} &= f_\theta\bigl([E(x_1),\dots,E(x_T), z_0, z_1, \dots, z_j]\bigr)_{T+j+1},\quad j=0,\dots,k-1 \\ |
| p(\cdot \mid x_{1:T}) &= \mathrm{softmax}(U z_k) |
| \end{align*} |
|
|
| The model can therefore ``iterate'' $k$ extra forward passes on the |
| same prompt before committing to a token, with the $k$ extra hidden |
| states carrying intermediate computation. Setting $k=0$ recovers the |
| vanilla baseline. |
|
|
| The other three latent modes are alternatives that we ablated: |
| \code{fixed\_slots} (concatenate $k$ trainable seed embeddings --- |
| Option-2), \code{latent\_seeds} (similar to \code{fixed\_slots}), and |
| \code{residual} (project $k$ extra hidden states back onto the base |
| hidden state via a learned residual). All modes share the SFT and GRPO |
| trainers; only the next-token logit function changes. |
|
|
| For the curriculum, we grow $k$ stage by stage: |
|
|
| \begin{center} |
| \begin{tabular}{ccl} |
| \toprule |
| \textbf{stage} & \textbf{num\_cot\_tokens} & \textbf{comment} \\ |
| \midrule |
| 1 & 1 & one extra recursion as soon as the model has the surface form \\ |
| 2 & 2 & two --- needed for 1-step propagation reasoning \\ |
| 3 & 3 & three --- needed for 2-step propagation reasoning \\ |
| \bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| \section{The reward function} |
|
|
| Defined in \code{multi\_output\_cell\_policy/rewards.py}. |
|
|
| Given target set $T$, predicted set $P$ (after JSON parse), let |
| \begin{itemize}[leftmargin=*,nosep] |
| \item \code{num\_good} $= |P \cap T|$ |
| \item \code{num\_bad} $= |P \setminus T|$ |
| \item \code{num\_missing} $= \max(0, |T| - \text{num\_good})$ |
| \item \code{is\_exact} $= (P \neq \varnothing) \land (P = T)$ |
| \item $\mathrm{tri}(n) = n(n+1)/2$ (rewards larger correct sets superlinearly) |
| \end{itemize} |
|
|
| Then |
| \begin{align*} |
| r &= \mathrm{tri}(\text{num\_good}) \cdot R_g \;-\; \text{num\_bad} \cdot P_b \\ |
| &\quad - \indic[P=\varnothing]\, P_e \;-\; \indic[|P|=1, |T|>1, i<2]\, P_s \\ |
| &\quad - \text{num\_missing}\cdot P_m \;+\; \indic[\text{is\_exact}]\, B_x \\ |
| &\quad - \indic[|P|<|T|, |T|>1]\, P_c |
| \end{align*} |
|
|
| with parameters (this is the recipe that produced the 0.58/0.68 latent |
| solve): |
|
|
| \begin{center} |
| \begin{tabular}{cllr} |
| \toprule |
| \textbf{symbol} & \textbf{flag} & \textbf{role} & \textbf{value} \\ |
| \midrule |
| $R_g$ & \code{--reward\_good\_value} & per-correct-value reward (triangular shape) & 1.25 \\ |
| $P_b$ & \code{--penalty\_bad\_value} & per-extra-wrong-value penalty & 1.0 \\ |
| $P_{\text{mal}}$ & \code{--penalty\_malformed} & flat penalty if JSON parse fails & 4.0 \\ |
| $P_e$ & \code{--penalty\_empty} & flat penalty if predicted set is empty & 0.5 \\ |
| $P_s$ & \code{--penalty\_singleton} & only at stage$<$2: punishes singleton on multi-value targets & 1.5 \\ |
| $P_m$ & \code{--penalty\_missing} & per-missing-value (recall pressure) --- \textbf{NEW} & \textbf{0.75} \\ |
| $B_x$ & \code{--exact\_match\_bonus} & only when $P = T$ --- \textbf{NEW} & \textbf{2.0} \\ |
| $P_c$ & \code{--cardinality\_mismatch\_penalty} & when $|P| < |T|$ and $|T|>1$ --- \textbf{NEW} & \textbf{1.0} \\ |
| \bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| Parse failures short-circuit to $r = -P_{\text{mal}}$ and zero per-cell |
| metrics. |
|
|
| \subsection{Why those three new terms exist (the breakthrough)} |
|
|
| Diagnosis: at the v3/v4 plateau, eval reported |
|
|
| \begin{lstlisting} |
| exact=0.95 precision=0.95 recall=0.95 solve=0.30 avg_set_size=1.000 |
| \end{lstlisting} |
|
|
| across all checkpoints. Per-cell exact and precision/recall were all |
| near 0.95 but the model \textbf{always predicted a single digit} |
| (\code{avg\_set\_size=1.000}). On a multi-value target $T=\{8,9\}$, |
| predicting $\{8\}$ keeps precision $=1.0$, recall $=0.5$ and yet |
| \code{exact\_set\_match}$=0$. Solve $= \text{exact\_set\_match}^N$ is |
| catastrophic in $N$ ($=20$), so even a small fraction of multi-value |
| cells killed it. |
|
|
| Without any of the new terms the optimum of $r$ on a multi-value cell |
| is trivially ``predict the singleton you are most confident about'' --- |
| there is no upside to enumerate the second value. The three new terms |
| close exactly that hole: |
|
|
| \begin{itemize}[leftmargin=*,nosep] |
| \item $P_m$ (\code{penalty\_missing}) directly penalises recall; |
| \item $B_x$ (\code{exact\_match\_bonus}) makes $P=T$ strictly dominate any singleton; |
| \item $P_c$ (\code{cardinality\_mismatch\_penalty}) is a flat hammer whenever $|P|<|T|$. |
| \end{itemize} |
|
|
| After these terms were added, GRPO on the latent variant moved solve |
| from $\sim 0.30$ to $\sim 0.58$ (100-puzzle eval) over $\sim 200$ |
| steps. The same fix is what we ported back into the baseline pipeline |
| this evening (see \S 10). |
|
|
| \section{Multi-stage warm-baseline pipeline (the recipe that worked)} |
|
|
| Master script: |
| \code{hard\_9x9\_stage1\_consistency\_queue/launch\_20empty\_warm\_baseline\_all\_latent\_modes\_stages123.sh}. |
|
|
| For each curriculum stage we run \textbf{three sub-phases in order}: |
|
|
| \begin{lstlisting} |
| [stage i] |
| (1) baseline warm SFT (no latent tokens, k=0, vanilla LM) |
| (2) latent SFT (k = i, latent mode = recurrent_hidden) |
| (3) latent GRPO (k = i) |
| \end{lstlisting} |
|
|
| \textbf{The warm baseline phase (1) is the trick that makes the |
| curriculum work.} At every stage transition the data distribution |
| changes ($i$ increases $\Rightarrow$ target sets shrink) and a new |
| latent slot appears. Doing a vanilla SFT on the new distribution first |
| lets the LM relearn the surface form on familiar parameters; THEN the |
| latent SFT adds the extra thought slot on top of an already-good policy. |
| When we tried to add a new latent slot directly on top of the previous |
| stage's GRPO checkpoint, training loss did NOT decrease. |
|
|
| Concrete LR schedule used for the champion run: |
|
|
| \begin{center} |
| \begin{tabular}{lllc} |
| \toprule |
| \textbf{phase} & \textbf{init from} & \textbf{LR} & \textbf{k} \\ |
| \midrule |
| S1 baseline SFT & base Qwen & 2e-4 & 0 \\ |
| S1 latent SFT & S1 baseline & 2e-4 & 1 \\ |
| S1 latent GRPO & S1 latent SFT & 1e-6 & 1 \\ |
| S2 baseline warm SFT & S1 GRPO & 5e-5 & 0 \\ |
| S2 latent SFT & S2 baseline & 5e-5 & 2 \\ |
| S2 latent GRPO & S2 latent SFT & 1e-6 & 2 \\ |
| S3 baseline warm SFT & S2 GRPO & 5e-5 & 0 \\ |
| S3 latent SFT & S3 baseline & 5e-5 $\rightarrow$ 1e-5 (champion) & 3 \\ |
| S3 latent GRPO & S3 latent SFT & 5e-6 ($\beta=0$) & 3 \\ |
| \bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| Other shared knobs: |
|
|
| \begin{lstlisting} |
| LoRA: r=32 a=64 dropout=0.05 on q,k,v,o,gate,up,down |
| SFT: per_device_bs=8 grad_accum=2 nproc=8 -> eff_bs=128 |
| GRPO: per_device_bs=4 grad_accum=2 nproc=8 -> eff_bs=64 |
| num_generations=4 beta=0.0 max_prompt_length=1024 |
| max_completion_length=24 |
| multi_value_oversample_factor=5, exact_match_bonus=2.0, |
| penalty_missing=0.75, cardinality_mismatch_penalty=1.0 |
| \end{lstlisting} |
|
|
| \section{GRPO settings that mattered} |
|
|
| \begin{itemize}[leftmargin=*] |
| \item \textbf{$\beta = 0$.} The KL anchor was harmful in every sweep |
| where we tried $\beta>0$. \code{s3\_grpo\_kl04} ($\beta=0.04$) peaked |
| at solve $=0.625$ (40p) at step 100 and regressed to $0.525$ by step |
| 500. |
|
|
| \item \textbf{\code{num\_generations} $= 4$.} With \code{num\_generations}$=2$ |
| we routinely saw \code{reward\_std}$=0$ (all sampled completions |
| identical $\Rightarrow$ no gradient). Bumping to 4 fixed it. |
|
|
| \item \textbf{Low LR.} \code{lr=5e-6} was the steadiest. \code{lr=1e-5} |
| peaked at step 200 (solve $0.65$) then collapsed back to $0.54$ --- |
| classic mode collapse. |
|
|
| \item \textbf{Effective bs $\geq 64$.} TRL's GRPOConfig requires |
| \code{eff\_bs * grad\_accum \% num\_generations == 0}; with 8 GPUs we |
| hit this trivially, but we caution single-GPU rerunners to set |
| \code{per\_device\_bs=4 grad\_accum=2 num\_generations=4}. |
|
|
| \item \textbf{\code{enable\_input\_require\_grads()} on the wrapped backbone.} |
| Required for TRL 0.15.x + PEFT LoRA + gradient checkpointing --- |
| otherwise the loss tensor produced by GRPOTrainer has |
| \code{requires\_grad=False} and \code{.backward()} raises. Also |
| \code{unwrapped.config.use\_cache = False}. |
| \end{itemize} |
|
|
| \section{Final hyperparameters table --- champion latent run} |
|
|
| \begin{center} |
| \begin{longtable}{lll} |
| \toprule |
| \textbf{group} & \textbf{hyperparameter} & \textbf{value} \\ |
| \midrule |
| \endfirsthead |
| \toprule |
| \textbf{group} & \textbf{hyperparameter} & \textbf{value} \\ |
| \midrule |
| \endhead |
| Backbone & model & Qwen/Qwen2.5-1.5B-Instruct \\ |
| Backbone & dtype & bf16 \\ |
| Backbone & LoRA target modules & q,k,v,o,gate,up,down \\ |
| Backbone & LoRA $r$ / $\alpha$ / dropout & 32 / 64 / 0.05 \\ |
| Latent & mode & \code{recurrent\_hidden} \\ |
| Latent & \code{num\_cot\_tokens} (S1/S2/S3) & 1 / 2 / 3 \\ |
| Latent & \code{max\_latent\_slots} / seeds & 8 / 8 \\ |
| Data & total empties & 20 \\ |
| Data & train rows / eval rows & 10\,000 / 100 \\ |
| Data & \code{multi\_value\_oversample\_factor} & 5 \\ |
| Data & \code{mixed\_stage1\_ratio} (S1) & 1 \\ |
| Data & \code{mixed\_stage2\_ratio} (S$\geq 2$) & 1 \\ |
| SFT & per\_device\_bs / grad\_accum & 8 / 2 \\ |
| SFT & \code{num\_epochs} (cap) & 64 \\ |
| SFT & LR (S1 latent) & 2e-4 \\ |
| SFT & LR (S2/S3 baseline warm + latent) & 5e-5 \\ |
| SFT & LR (S3 latent champion \code{s3b\_lr1e5\_o5}) & 1e-5 \\ |
| SFT & weight\_decay & 0.0 \\ |
| SFT & gradient checkpointing & on \\ |
| GRPO & per\_device\_bs / grad\_accum & 4 / 2 \\ |
| GRPO & \code{num\_generations} & 4 \\ |
| GRPO & LR & 5e-6 (S3); 1e-6 (S1, S2) \\ |
| GRPO & $\beta$ (KL) & 0.0 \\ |
| GRPO & \code{max\_prompt\_length} & 1024 \\ |
| GRPO & \code{max\_completion\_length} & 24 \\ |
| Reward & \code{reward\_good\_value} & 1.25 \\ |
| Reward & \code{penalty\_bad\_value} & 1.0 \\ |
| Reward & \code{penalty\_malformed} & 4.0 \\ |
| Reward & \code{penalty\_empty} & 0.5 \\ |
| Reward & \code{penalty\_singleton} & 1.5 \\ |
| Reward & \code{penalty\_missing} & 0.75 \\ |
| Reward & \code{exact\_match\_bonus} & 2.0 \\ |
| Reward & \code{cardinality\_mismatch\_penalty} & 1.0 \\ |
| Eval & early-stop on prec/recall & 0.98 \\ |
| \bottomrule |
| \end{longtable} |
| \end{center} |
|
|
| \section{Headline results} |
|
|
| \subsection{Latent (with thought tokens, \code{recurrent\_hidden})} |
|
|
| \begin{center} |
| \begin{tabular}{llrrrrr} |
| \toprule |
| \textbf{eval} & \textbf{model / phase} & \textbf{step} & \textbf{exact} & \textbf{prec} & \textbf{recall} & \textbf{solve} \\ |
| \midrule |
| \textbf{100p (auth.)} & \code{s3\_grpo\_baseline} (S3 GRPO, $\beta=0$, lr=5e-6) & 200 & 0.9665 & 0.9673 & 0.9680 & \textbf{0.580 (58/100)} \\ |
| 40p & \code{s3\_grpo\_sharp\_rwd} ($B_x{=}4$, $P_c{=}3$) & 300 & --- & --- & --- & \textbf{0.675 (27/40)} \\ |
| 40p & \code{s3\_grpo\_lr1e5} & 200 & 0.978 & 0.978 & 0.979 & 0.650 \\ |
| 40p & \code{s3b\_lr1e5\_o5} (S3 SFT champion) & 2400 & 0.974 & 0.974 & 0.975 & 0.600 \\ |
| \bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| \subsection{Vanilla baseline (no thought tokens, same Qwen2.5-1.5B + LoRA)} |
|
|
| \begin{center} |
| \begin{tabular}{llrrr} |
| \toprule |
| \textbf{sweep} & \textbf{best variant} & \textbf{best step} & \textbf{exact} & \textbf{solve (100p)} \\ |
| \midrule |
| v3 (single-GPU, no oversample, no new reward) & \code{baseline\_3stage\_20260522} & --- & 0.730 & \textbf{0.000} \\ |
| v4 (LR sweep, multi-GPU, original reward) & \code{pipe\_v\_sft\_extend} (S3 SFT extended) & 4000 & 0.948 & \textbf{0.400} \\ |
| \textbf{v6 (this evening; ports latent reward + oversample)} & \code{v6\_i\_sft\_v\_oversample10} & running & 0.952$+$ & \textbf{0.440 (best so far)} \\ |
| \bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| The v6 sweep is still running --- \code{v6\_e/f/i} are in S3 SFT |
| continuation, GRPO follow-on phases queued. The \code{v6\_i} variant |
| has hit \textbf{solve $=0.44$} at SFT eval (new baseline best, |
| $+0.04$ over v4) and is still climbing. |
|
|
| \subsection{Stage-by-stage trajectory (latent, 40-puzzle eval)} |
|
|
| \begin{lstlisting} |
| S1 SFT : exact ~ 0.85, solve ~ 0.20 |
| S1 GRPO : exact ~ 0.90, solve ~ 0.20 |
| S2 SFT (no oversample) : exact ~ 0.94, solve ~ 0.20-0.25 <- the wall |
| S2 SFT + multi_value_oversample=5 : exact ~ 0.96, solve ~ 0.30-0.35 |
| S2 GRPO + new reward terms : exact ~ 0.96, solve ~ 0.35-0.40 |
| S3 SFT (s3b_lr1e5_o5 step 2400) : exact 0.974, solve 0.600 <- SFT champion |
| S3 GRPO (s3_grpo_baseline step 200,100p): exact 0.967, solve 0.580 <- 100p champion |
| S3 GRPO (s3_grpo_sharp_rwd step 300,40p): solve 0.675 <- 40p peak |
| \end{lstlisting} |
|
|
| \subsection{Latent vs baseline gap (head-to-head, same 100p eval, same prompts)} |
|
|
| \begin{center} |
| \begin{tabular}{lrrrrr} |
| \toprule |
| \textbf{model} & \textbf{exact} & \textbf{prec} & \textbf{recall} & \textbf{solve} & \textbf{solved/100} \\ |
| \midrule |
| Latent \code{recurrent\_hidden}, S3 GRPO & 0.9665 & 0.9673 & 0.9680 & \textbf{0.580} & 58 \\ |
| Vanilla baseline, \code{v6\_i} (best at time of writing) & 0.952 & 0.952 & 0.952 & \textbf{0.440} & 44 \\ |
| \bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| Gap on 100-puzzle solve: $\approx$ \textbf{$+0.14$ absolute / $+32\%$ |
| relative} for latent over the strongest baseline we have. |
|
|
| \section{Why the latent works (interpretation hypotheses)} |
|
|
| These are the working hypotheses the experiments are consistent with; |
| none is fully proven and ablations are still WIP. |
|
|
| \begin{enumerate}[leftmargin=*] |
| \item \textbf{Constraint-propagation depth.} Stage-3 i-consistency is |
| essentially 2-ply lookahead. With $k=3$ recurrent hidden tokens the |
| model gets exactly three extra forward passes between prompt and |
| output --- one for the legality check, one for 1-step propagation, |
| one for the second step of propagation. Empirically the gap to the |
| no-thought-token baseline appears at stages where multi-step |
| propagation matters (stage 2 onward; stage 1 numbers are essentially |
| identical). |
|
|
| \item \textbf{Multi-value cells require enumeration, which a singleton |
| softmax can't do in one forward pass.} A vanilla LM at 1.5B |
| parameters predicts essentially deterministically once temperature is |
| low; for a target set $\{8, 9\}$ the LM picks one of the two and |
| stops. The latent model can use one of the recurrent hidden steps to |
| ``consider'' each option without committing yet, which is exactly |
| the failure mode in the data (\code{avg\_set\_size} $= 1.000$ for the |
| baseline, $\approx 1.05$ for the latent S3 model on the same eval). |
|
|
| \item \textbf{Stable curriculum capacity growth.} Adding a new latent |
| slot at every stage gives the model a ``fresh slate'' of |
| representational capacity at the exact transition where the task |
| gets harder. The warm-baseline SFT between stages prevents the new |
| slot from corrupting the previously learned policy. Without warm |
| baseline, training loss did not decrease at all (we observed this |
| directly when we tried to skip the warm baseline). |
|
|
| \item \textbf{GRPO without latent slots is starved of variance.} With |
| \code{max\_completion\_length} 24 and the model essentially |
| deterministic, GRPO's 4 sampled completions per prompt collapse to a |
| single answer --- \code{reward\_std}$=0$, no gradient. With latent |
| recurrence + the new \code{exact\_match\_bonus} reward, the model |
| occasionally samples a 2-element set, gets a much higher reward, and |
| that prompt gets a real gradient signal. |
| \end{enumerate} |
|
|
| \section{Reproducibility} |
|
|
| \noindent |
| Code repository: \url{https://github.com/Avra98/curriculum_cot} \\ |
| Latent checkpoints: \url{https://huggingface.co/Avra98/sudoku-latent-recurrent-hidden-20empty-stages} \\ |
| Baseline checkpoints: \url{https://huggingface.co/Avra98/sudoku-9x9-20empty-baseline-1p5b-sweep} |
|
|
| Key scripts: |
|
|
| \begin{itemize}[leftmargin=*,nosep] |
| \item Master orchestrator (latent, 9-phase warm-baseline pipeline): |
| \code{hard\_9x9\_stage1\_consistency\_queue/launch\_20empty\_warm\_baseline\_all\_latent\_modes\_stages123.sh} |
| \item Vanilla baseline pipeline: |
| \code{\_runs/baseline\_1p5b\_pipeline\_v4.sh} (with v6 launchers |
| \code{\_runs/launch\_baseline\_push\_v6.sh}) |
| \item SFT trainer (vanilla): |
| \code{multi\_output\_cell\_policy/sft\_multi\_output\_train.py} |
| \item GRPO trainer (vanilla): |
| \code{multi\_output\_cell\_policy/grpo\_multi\_output\_train.py} |
| \item SFT trainer (latent): |
| \code{latent\_multi\_output\_cell\_policy/sft\_latent\_multi\_output\_train.py} |
| \item GRPO trainer (latent): |
| \code{latent\_multi\_output\_cell\_policy/grpo\_residual\_projector\_latent\_train.py} |
| \item Reward function: \code{multi\_output\_cell\_policy/rewards.py} |
| \item Prompt builder: \code{multi\_output\_cell\_policy/prompt\_builder.py} |
| \item Stage-i consistency: |
| \code{multi\_output\_cell\_policy/shared\_multi\_output\_policy.py} |
| \item 100-puzzle evaluator: \code{analysis/eval\_stage2\_checkpoint.py} |
| \end{itemize} |
|
|
| To reproduce the latent champion (1.5B, 9-phase, $\sim 16$ GPU$\cdot$h |
| on $8\times$H100 80GB): |
|
|
| \begin{lstlisting} |
| export STAGE1_BASELINE_ADAPTER_DIR=/path/to/stage1_baseline_seed_adapter |
| bash hard_9x9_stage1_consistency_queue/launch_20empty_warm_baseline_all_latent_modes_stages123.sh |
| \end{lstlisting} |
|
|
| To reproduce the v6 baseline push (single-GPU per variant, $\sim 6$ |
| GPU$\cdot$h): |
|
|
| \begin{lstlisting} |
| bash _runs/launch_baseline_push_v6.sh |
| \end{lstlisting} |
|
|
| \appendix |
|
|
| \section{The reward fix as a one-line patch} |
|
|
| The single most consequential code change in this whole project, as a |
| self-contained patch on \code{multi\_output\_cell\_policy/rewards.py}: |
|
|
| \begin{lstlisting}[language=Python] |
| # new args (default 0 preserves legacy behaviour) |
| penalty_missing: float = 0.0 |
| exact_match_bonus: float = 0.0 |
| cardinality_mismatch_penalty: float = 0.0 |
|
|
| num_missing = max(0, len(target_set) - num_good) |
| is_exact = bool(predicted_values) and (set(predicted_values) == target_set) |
|
|
| # ... base reward (triangular_number(num_good)*reward_good_value |
| # - num_bad*penalty_bad_value) |
|
|
| if num_missing > 0: |
| reward -= num_missing * penalty_missing |
| if is_exact: |
| reward += exact_match_bonus |
| if len(predicted_values) < len(target_values) and len(target_values) > 1: |
| reward -= cardinality_mismatch_penalty |
| \end{lstlisting} |
|
|
| Defaults are zero so old runs are unaffected; the recipe sets |
| $(P_m, B_x, P_c) = (0.75, 2.0, 1.0)$ for the vanilla recipe and |
| $(1.0, 4.0, 3.0)$ for the ``sharp\_rwd'' variant. |
|
|
| \section{The warm-baseline trick as a sequence diagram} |
|
|
| \begin{lstlisting} |
| Stage 1 Stage 2 Stage 3 |
| --------- --------- --------- |
| [base Qwen] | | |
| | | | |
| v v v |
| S1 baseline SFT -> S2 baseline SFT -> S3 baseline SFT |
| (no latent, k=0) (no latent, k=0) (no latent, k=0) |
| | | | |
| v v v |
| S1 latent SFT -> S2 latent SFT -> S3 latent SFT |
| (k=1) (k=2) (k=3) |
| | | | |
| v v v |
| S1 latent GRPO -> S2 latent GRPO -> S3 latent GRPO |
| (k=1, b=0, lr 1e-6) (k=2, b=0) (k=3, b=0, lr 5e-6) |
| | | | |
| [final policy] |
| \end{lstlisting} |
|
|
| Every arrow is \code{init\_adapter\_dir = <previous output>}. Each row |
| is a ``slot in the curriculum''; the column adds reasoning capacity |
| ($k\mathrel{+}=1$) and moves to a harder target distribution |
| ($i\mathrel{+}=1$). The diagonal across the diagram is the actual |
| training trajectory. |
|
|
| \bigskip |
| \noindent\emph{End of report.} |
|
|
| \end{document} |
|
|
