DRPO

DRPO (Denoising and Remasking Policy Optimization) is a full-weight reasoning checkpoint derived from JetLM/SDAR-1.7B-Chat. It augments a block-wise masked diffusion language model with a lightweight remasking head that learns whether each visible token should be kept or revised.

Unlike confidence-only remasking heuristics, DRPO learns revision decisions from rollouts. It rewards revisions that improve the final generation relative to an all-keep continuation from the same intermediate state, while penalizing unnecessary remasking.

This repository contains the inference checkpoint saved after 450 optimizer updates. Optimizer, scheduler, and RNG states are intentionally excluded.

Method

Training has two stages:

1. Local remasking warmup. The base denoiser is frozen while the remasking head learns a soft token-level target. Correct tokens are assigned a keep target; for mismatched tokens, the target reflects how much probability mass the denoiser assigns to alternatives ranked above the selected token.
2. Rollout optimization. Seven sampled keep/remask branches and one all-keep reference branch continue from the same parent state. Denoising uses a group-relative branch advantage. Remasking uses a keep-relative advantage based on improvement over the all-keep branch, with an explicit remasking cost. The two objectives are optimized jointly.

The rollout return combines task outcome with an optional step-level confidence signal derived from changes in the top-1/top-2 probability gap. The default revision window is three blocks.

Evaluation

Main results

All values are percentages; higher is better. Math tasks use exact match after answer extraction, and code tasks use unit-test pass rate.

Model	GSM8K	MATH-500	MMLU	HumanEval	MBPP	Average
SDAR-1.7B-Chat	80.10	63.20	62.90	61.60	61.10	65.36
SFT	80.53	59.80	63.56	60.37	61.87	65.23
DRPO	82.34	64.40	63.57	60.98	62.26	66.71

For MATH-500, DRPO answered 322 of 500 examples correctly. This result was produced with the project-specific OpenCompass evaluation adapter on 2026-06-15.

MATH-500 decoding configuration

Setting	Value
Block length	4
Maximum new tokens	1536
Sampling temperature	0.0
Confidence threshold	0.95
Remask policy	learned gap head
Remask threshold	0.50
Remask interval	2 blocks
Remask window	3 blocks
Remask start	192 generated tokens
Prefix guard	192 generated tokens
Tail guard	1 block

The reported score depends on the custom remasking decoder and these settings; it is not the result of the default Transformers generation pipeline.

Decoding behavior

Model	Tokens / denoising step	Net progress / step	Denoising steps	Remasking steps
SDAR-1.7B-Chat	2.75	2.75	96.52	—
SFT	2.63	1.62	101.17	48.30
DRPO	2.69	1.98	97.44	30.22

On GSM8K, learned remasking rescued 68 initially incorrect examples while damaging 38 initially correct examples, for a net gain of 30 examples. It also used fewer remasking steps than the evaluated random, low-confidence, margin, budgeted-confidence, and supervised remasking baselines.

Loading the checkpoint

The model contains custom SDAR code, so trust_remote_code=True is required.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Leotsia/DRPO"

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype="auto",
    device_map="auto",
)

The training and evaluation environment used Transformers 4.52.4. Reproducing the reported results requires the project-specific remasking decoder and evaluation adapter in addition to the checkpoint.

Training details

Item	Value
Backbone	JetLM/SDAR-1.7B-Chat
Parameter count	1.7B
Training data	7,490 competition-math training examples
Fine-tuning	Full parameters
Rollout branches	7 sampled revision branches + 1 all-keep reference
Default candidate window	3 blocks
Selected checkpoint	Optimizer step 450

In the branch-count ablation, seven sampled revision branches achieved the highest peak group accuracy among the tested settings of one, four, and seven branches.

Intended use

This checkpoint is intended for research on masked diffusion language models, mathematical reasoning, learned token revision, and rollout-based optimization. It is not designed as a drop-in replacement for an autoregressive chat model.

Limitations

This is a research checkpoint focused on reasoning benchmarks. It has not been comprehensively evaluated for factual reliability, safety, bias, multilingual use, or high-stakes applications. Generated reasoning can be incorrect even when fluent, and benchmark performance may vary with prompts, extraction rules, randomness, hardware, and decoder implementation.

Because loading executes repository-provided Python code, review the files and pin a repository revision in security-sensitive deployments.

License and attribution

The checkpoint is released under Apache-2.0, following the license metadata of the SDAR-1.7B-Chat base model. Please also cite and follow the attribution guidance of the SDAR project.