README.md

---
license: apache-2.0
library_name: transformers
base_model: JetLM/SDAR-1.7B-Chat
pipeline_tag: text-generation
tags:
  - diffusion-language-model
  - reinforcement-learning
  - mathematical-reasoning
  - remasking
  - drpo
model-index:
  - name: DRPO
    results:
      - task:
          type: text-generation
          name: Text Generation
        dataset:
          type: HuggingFaceH4/MATH-500
          name: MATH-500
          split: test
        metrics:
          - type: accuracy
            name: Accuracy
            value: 64.4
---

# DRPO

DRPO (Denoising and Remasking Policy Optimization) is a full-weight reasoning
checkpoint derived from
[`JetLM/SDAR-1.7B-Chat`](https://huggingface.co/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.

```python
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](https://github.com/JetAstra/SDAR).