MDM-correction — TCT reveal-rule ablation checkpoints
Checkpoints for the NeurIPS 2026 rebuttal of submission 27582, "Random Remasking Scales: A Unified View of Masked Diffusion Inference."
Each model is a 6.85M-parameter diffusion-LLaMA (8 layers, hidden 256,
vocab 42), trained with the TCT progressive-edit recipe. These are the exact
weights behind the Algorithm-2 reveal-rule ablation (reviewer mRYZ Q7).
Optimizer/scheduler state has been stripped — these are inference-only weights
(model_state_dict + config + global_step).
The published top-k + argmax baseline is not in this repo; it lives at
zzy1123/sudoku-hard-models
(progressive_edit-s123/step460000).
Files
| File | Task | Reveal rule | Ablation | Reported accuracy (M=256) |
|---|---|---|---|---|
sudoku_tct_random-pos_argmax_s123_step460000.pt |
Sudoku-Hard | random position + argmax | Alg-2 line-11 | 85.5 |
sudoku_tct_topk_sampled_s123_step460000.pt |
Sudoku-Hard | top-k conf. + sampled | Alg-2 line-9 | 96.8 |
sat_tct_random-pos_argmax_s0_step100000.pt |
3-SAT | random position + argmax | Alg-2 line-11 | see report |
sat_tct_topk_sampled_s0_step100000.pt |
3-SAT | top-k conf. + sampled | Alg-2 line-9 | see report |
(Published baseline for reference: Sudoku top-k+argmax = 98.3 @ M=256.)
Loading
import torch
ck = torch.load("sudoku_tct_topk_sampled_s123_step460000.pt", map_location="cpu")
model.load_state_dict(ck["model_state_dict"])
# ck["config"] holds the full training/model/data config; ck["global_step"] the step.
Provenance
Trained in mdm_correction/ (progressive_edit, K=5/8, confidence_threshold=0.9).
Evaluated with the gibbs_edit_v2 corrector (confidence=markovian,
edit_freq=1, early_exit=5), sweeping M ∈ {16,32,64,128,256}.