ACDiR: Actor-Critic Diffusion Repair

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.

The released critic checkpoint is hosted in this Hugging Face repository: NYCU-MLLab/acdir. The base actors remain upstream downloads: GSAI-ML/LLaDA-8B-Instruct and nvidia/Nemotron-Labs-Diffusion-8B. The lightweight OKR evidence package, including converted predictions and analysis notes, is maintained in NYCU-MLLab/NSTC-OKR.

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:

steps=256                 gen_length=512
block_length=32           block_steps=16
batch_size=1              temperature=0
remask_method=count_set   remask_temperature=0.4
lookback_blocks=1         remask_min_age_current=2
remask_max_age_lookback=6
reforward_after_remask=True
deterministic_joint_argmax=False
force_remask_window=0

Hardware and software

The release path was validated with:

• Linux and Python 3.11;
• PyTorch 2.5.1 with CUDA 12.1 wheels;
• 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:

cd acdir

python3.11 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel ninja packaging
python -m pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu121
python -m pip install -r requirements.txt
python -m pip install flash-attn==2.8.3 --no-build-isolation

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:

export LLADA_EXACT_BACKEND=lmdeploy
export LLADA_LMDEPLOY_FAST_MODE=full_window
export LLADA_LMDEPLOY_CUDAGRAPH=0
export LLADA_LMDEPLOY_VARLEN_FLASH=1

python eval_math500.py \
  --batch_size 1 \
  --nproc_per_node 1 \
  --result_dir outputs/math500_44

To use an existing actor copy:

python eval_math500.py \
  --base_model /absolute/path/to/LLaDA-8B-Instruct \
  --batch_size 1 \
  --nproc_per_node 1 \
  --result_dir outputs/math500_44

The expected final summary is:

guided   44.00%  (220/500)
base     39.60%  (198/500)  delta=+4.40%

Every run writes:

outputs/math500_44/eval_command.txt
outputs/math500_44/predictions/predictions_rank000.json

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:

export LLADA_EXACT_BACKEND=lmdeploy
export LLADA_LMDEPLOY_FAST_MODE=full_window
export LLADA_LMDEPLOY_CUDAGRAPH=0
export LLADA_LMDEPLOY_VARLEN_FLASH=1

python eval_gsm8k.py \
  --batch_size 1 \
  --nproc_per_node 1 \
  --result_dir outputs/gsm8k

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.

Outputs follow the MATH500 layout:

outputs/gsm8k/eval_command.txt
outputs/gsm8k/predictions/predictions_rank000.json

Consolidated report

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:

sbatch -A YOUR_ACCOUNT -p YOUR_GPU_PARTITION \
  --export=ALL,CONDA_MODULE=miniconda3,CUDA_MODULE=cuda/12.6,RUN_FLASH_ATTN_BUILD=1 \
  scripts/batch/setup_env.batch

sbatch -A YOUR_ACCOUNT -p YOUR_GPU_PARTITION \
  --export=ALL,CONDA_MODULE=miniconda3,CUDA_MODULE=cuda/12.6,LLADA_LMDEPLOY_VARLEN_FLASH=1 \
  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:

CUDA_VISIBLE_DEVICES="" python -m unittest discover -s tests -v

These tests do not replace the full 500-example GPU reproduction.

Serve locally

The same pinned actor download and critic can be exposed through a small OpenAI-compatible endpoint:

python serve_openai.py \
  --host 127.0.0.1 \
  --port 23333

The endpoint is /v1/chat/completions.

Artifact integrity

Artifact	SHA256
weights/critic-ckpt-000040.pt	b8f86493bfd629968e18ed362f47614affc869af3ab01827255e4025ba26e68c
datasets/MATH500/test/data-00000-of-00001.arrow	ff2663846092b986df3026f53904030cbaf8e061c9f978d319a7bd86b3a04ea4
configs/math500_44.json	b430248e1199d8d0d27dbf73f277fe3d436a4a2dd4866e41615bea410c1828c9
datasets/gsm8k/test/data-00000-of-00001.arrow	45965b000311d1550e5619b60b5bf31cf76edebfd8b8eddc62a876fbf8c9be95
configs/gsm8k.json	f79d1860488d53d6a9b79d4ce65c27921d36c536e2d7ccf14778378fb8283c0f

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:

@misc{acdir2026,
  title        = {ACDiR},
  author       = {{ACDiR Project}},
  year         = {2026},
  howpublished = {Hugging Face model release}
}

@article{nie2025large,
  title   = {Large Language Diffusion Models},
  author  = {Nie, Shen and Zhu, Fengqi and You, Zebin and Zhang, Xiaolu and
             Ou, Jingyang and Hu, Jun and Zhou, Jun and Lin, Yankai and
             Wen, Ji-Rong and Li, Chongxuan},
  journal = {arXiv preprint arXiv:2502.09992},
  year    = {2025}
}