Expand model card with method and evaluation details

README.md

@@ -6,7 +6,8 @@ pipeline_tag: text-generation
 tags:
   - diffusion-language-model
   - reinforcement-learning
-  - math
+  - mathematical-reasoning
+  - remasking
   - drpo
 model-index:
   - name: DRPO
@@ -26,25 +27,56 @@ model-index:
 
 # DRPO
 
-DRPO is a full-weight math reasoning checkpoint derived from
+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 equips the SDAR block-diffusion language model with a learned remasking head
-and optimizes the remasking policy with same-state relative rollouts.
+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. Training-state files such as optimizer, scheduler, and RNG states are
-intentionally excluded.
+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
 
-| Benchmark | Split | Metric | Score |
-| --- | --- | --- | ---: |
-| MATH-500 | test (500 examples) | Accuracy | **64.4** |
+### 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** |
 
-The score corresponds to 322 correct answers out of 500 and was produced with
-the project's OpenCompass-based evaluator on 2026-06-15.
+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.
 
-The decoding configuration used for this result was:
+### MATH-500 decoding configuration
 
 | Setting | Value |
 | --- | ---: |
@@ -60,9 +92,21 @@ The decoding configuration used for this result was:
 | Prefix guard | 192 generated tokens |
 | Tail guard | 1 block |
 
-The reported result depends on the custom remasking decoder and these settings;
-it should not be interpreted as the result of the default Transformers text
-generation pipeline.
+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
 
@@ -86,24 +130,38 @@ model = AutoModelForCausalLM.from_pretrained(
 ```
 
 The training and evaluation environment used Transformers 4.52.4. Reproducing
-the reported MATH-500 score requires the project-specific remasking evaluation
-adapter in addition to the checkpoint.
+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.
 
-## Training summary
+## Intended use
 
-- Backbone: `JetLM/SDAR-1.7B-Chat`
-- Training data: 7,490 competition-math training examples
-- Fine-tuning: full parameters
-- Objective: learned remasking policy with same-state relative rollouts
-- Sampled revision branches per example: 8
-- Checkpoint: step 450
+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 mathematical reasoning. It has not
-been comprehensively evaluated for general knowledge, factual reliability,
-safety, bias, or multilingual use. Generated reasoning may be incorrect even
-when it is fluent. Applications should validate outputs independently.
+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.