ACDiR (Actor-Critic Diffusion Repair) is the discrete repair method
introduced in Learning Towards Unmasking and Remasking by Actor-Critic
Diffusion. This release evaluates the frozen
LLaDA actor on MATH-500 and GSM8K, and reports a cross-model transfer check on
Nemotron-Labs-Diffusion-8B. It combines:
the frozen GSAI-ML/LLaDA-8B-Instruct actor;
weights/critic-ckpt-000040.pt;
count-set critic-guided remasking; and
the vendored exact LLaDA runner derived from LMDeploy.
This is not a standalone Transformers checkpoint and cannot be loaded with
AutoModel.from_pretrained() as a complete actor-plus-critic model.
For the training implementation and its release boundary, see
TRAINING.md. In particular, the training-loop launcher used
for the released checkpoint is not represented as an end-to-end public command
in this inference release; it would be misleading to claim otherwise.
Reported result
Dataset
Examples
Guided
Actor-only
Delta
HuggingFaceH4/MATH-500 test
500
44.00% (220/500)
39.60% (198/500)
+4.40 pp
The included test set is from
HuggingFaceH4/MATH-500 revision
6e4ed1a2a79af7d8630a6b768ec859cb5af4d3be. It is not the reordered/
reprocessed ankner/math-500 split used by some JustGRPO evaluations, so the
scores must not be compared as though they used identical examples and
grading.
The 44.00% result intentionally retains the original ACDiR legacy answer
extractor and grader. Replacing the grader changes the evaluation protocol and
is outside exact reproduction of this result.
Selected text summarization results
ACDiR is also evaluated as a targeted remask--unmask repair policy for meeting
and long-context summarization. The rows below report only positive, matched
end-to-end comparisons against the base LLaDA actor; they are not claims that
the critic-only marginal effect is significant in every setting.
Dataset
Metric
Base LLaDA
ACDiR
Improvement
MeetingBank
ROUGE-L
22.20
38.29
+16.09
MeetingBank
BERTScore
85.19
88.41
+3.22
L-Eval MeetingSumm
ROUGE-1
4.70
4.89
+0.19
L-Eval MeetingSumm
ROUGE-2
1.13
1.47
+0.35
L-Eval MeetingSumm
ROUGE-L
3.67
4.12
+0.45
For the 862-example MeetingBank paired judge evaluation, ACDiR is preferred
over the actor-only output on faithfulness (70.19%), relevance (63.34%),
coverage (60.09%), conciseness (55.45%), and overall quality (52.44%). The
judge is applied after generation and is not used as an online training reward.
GSM8K
datasets/gsm8k (1319 test problems) and eval_gsm8k.py are included so that
GSM8K runs through the same actor, critic, decoding schedule and reporting path
as MATH500. configs/gsm8k.json mirrors configs/math500_44.json exactly
except for the dataset, so the two benchmarks stay directly comparable.
Dataset
Examples
Guided
Actor-only
Delta
GSM8K test
1319
81.73% (1078/1319)
80.21% (1058/1319)
+1.52 pp
Paired over the same 1319 problems the critic rescues 80 and harms 60, so the
GSM8K gain on its own is not significant (McNemar p = 0.108). MATH500 is
(+4.40 pp, p = 0.014), and pooled over both benchmarks the effect is
significant (n = 1819, 128 rescues against 86 harms, p = 0.005). Report the
pooled result rather than claiming each benchmark separately.
Cross-model transfer (external)
The released critic was trained on LLaDA only. Applied unchanged to
nvidia/Nemotron-Labs-Diffusion-8B — a different diffusion LM — it still
improves both benchmarks:
Benchmark
Baseline
+ ACDiR critic
Delta
Relative
MATH-500 (500)
51.00%
54.20%
+3.20 pp
+6.27%
GSM8K (1319)
80.14%
83.24%
+3.10 pp
+3.87%
Those runs used the Nemotron evaluation harness (dlm mode, llada eval
style, batch size 1, block length 32, block steps 16, confidence threshold
0.85, 512 max new tokens, BF16) with the critic restricted to current-block
remasking. That harness is not vendored here, so these numbers are reported for
reference and are not reproducible from this repository alone.
The gain costs decoding compute. Average NFE rises from 117.9 to 188.9 on
MATH-500 and from 92.1 to 132.0 on GSM8K, with tokens-per-forward dropping from
3.30 to 2.07 and from 2.92 to 2.04 respectively.
Important decoding settings are stored in configs/math500_44.json:
a CUDA 12.6 toolkit for building FlashAttention 2.8.3;
BF16 inference; and
one NVIDIA H200 GPU for the strict reference run.
The full actor is approximately 16 GB in BF16 before runtime buffers. GPUs with
less memory have not been validated for exact reproduction. The reported path
uses LMDeploy full-window inference and varlen FlashAttention. Disabling varlen
flash or changing the backend is useful for portability testing but is not the
strict 44.00% configuration.
Install
Clone this Hugging Face model repository, enter its root directory, and create
a clean Python 3.11 environment:
If a different PyTorch CUDA wheel is required by the host driver, install that
wheel first and treat the run as a portability reproduction unless the final
outputs match the reference counts.
The evaluator automatically downloads the pinned base-model revision
08b83a6feb34df1a6011b80c3c00c7563e963b07 into
.cache/huggingface/hub. Pass a local model directory with --base_model to
avoid the download.
Reproduce 44.00%
Run on exactly one visible GPU; batch size 1 is part of the reported protocol:
The prediction JSON contains the prompt, full guided sequence, full baseline
sequence, gold solution, extracted answers, correctness flags, level, and
per-sample ACDiR remask statistics. With multiple ranks, one file is produced
per rank; the reference 44.00% run uses one rank.
Evaluate GSM8K
Same environment variables, same one-GPU/batch-size-1 protocol:
GSM8K has 1319 test problems, so a full run takes roughly 2.6x a MATH500 run.
Use --max_eval_samples N for a cheaper subset run; a subset result is not a
GSM8K score and must be labelled as such.
scripts/make_report.py is the reporting template. It recomputes every number
from the prediction files under outputs/ and records the model revisions,
decoding settings and hardware next to the results, so a reader can see which
comparisons are matched and which are not:
python scripts/make_report.py --out RESULTS.md
RESULTS.md covers the reasoning benchmarks with their actor-only and strong
baselines, the ablations, MT-Bench, and the repair statistics. Per-run
summaries and sample-level agreement between two runs come from
scripts/compare_runs.py.
Ablations
Both evaluators expose the decoding knobs as flags or EVAL_* environment
variables, so ablations need no code changes. Each of these isolates one design
component against the released configuration:
Ablation
Flag
Cap total remasks per sample
--max_total_remask_per_sample 1
Disable immediate re-forward
--reforward_after_remask False
Remask only at block end
--remask_timing block_end
Joint instead of independent argmax
--deterministic_joint_argmax True
Narrow the lookback horizon
--lookback_blocks 0
Change the minimum token age
--remask_min_age_current 1 | 4
--remask_temperature has no effect on its own: select_count_set_action
applies it only when sampling, and the released configuration selects actions
by argmax. Pair it with --sample_remask True or it is a no-op.
scripts/batch/ablation_math500.batch runs the whole sweep in one allocation.
The actor-only arm needs no ablation flag: compare_with_baseline is on by
default, so every run already reports guided and actor-only side by side.
The released configuration is deterministic (no_sample: true,
temperature: 0.0, sample_remask: false), so repeated runs of the same
command return identical results and a seed sweep changes nothing. To obtain a
spread instead of a point estimate, enable stochastic gate decisions with
--sample_remask True and vary eval.seed in the config. Report that as a
separate robustness protocol; the headline numbers stay on the deterministic
one.
Slurm
Portable Slurm helpers are included. Account, partition, environment-module,
and CUDA-module names are cluster-specific and should be supplied at submit
time. For example:
bash
1sbatch -A YOUR_ACCOUNT -p YOUR_GPU_PARTITION \2 --export=ALL,CONDA_MODULE=miniconda3,CUDA_MODULE=cuda/12.6,RUN_FLASH_ATTN_BUILD=1\3 scripts/batch/setup_env.batch
45sbatch -A YOUR_ACCOUNT -p YOUR_GPU_PARTITION \6 --export=ALL,CONDA_MODULE=miniconda3,CUDA_MODULE=cuda/12.6,LLADA_LMDEPLOY_VARLEN_FLASH=1\7 scripts/batch/eval_math500.batch
If the cluster does not use environment modules, activate or expose conda,
nvcc, and the CUDA libraries before submission and omit the module variables.
Regression checks
The CPU suite checks the release policy and rollout invariants without loading
the 8B actor:
scripts/upload_to_hf.py excludes local caches, generated outputs, temporary
files, and Slurm logs while retaining the portable batch scripts.
License and upstream components
The ACDiR release code and critic are provided under the Apache License 2.0;
see LICENSE. Vendored LMDeploy code retains its Apache-2.0 license in
lmdeploy/LICENSE. The LLaDA base model is not redistributed here and remains
under its upstream MIT license. The included evaluation records originate from
HuggingFaceH4/MATH-500 and GSM8K and remain subject to the upstream
dataset/source terms.
Citation
Please cite the ACDiR release and LLaDA when using this checkpoint:
bibtex
1@misc{acdir2026,
2 title = {ACDiR},
3 author = {{ACDiR Project}},
4 year = {2026},
5 howpublished = {Hugging Face model release}
6}
78@article{nie2025large,
9 title = {Large Language Diffusion Models},
10 author = {Nie, Shen and Zhu, Fengqi and You, Zebin and Zhang, Xiaolu and
11 Ou, Jingyang and Hu, Jun and Zhou, Jun and Lin, Yankai and
12 Wen, Ji-Rong and Li, Chongxuan},
13 journal = {arXiv preprint arXiv:2502.09992},
14 year = {2025}
15}
