Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models
Authors: Zemin Huang
Yuhang Wang
Zhiyang Chen
Guo-Jun Qi
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion language model
discrete diffusion
masked diffusion model
language model
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose RemeDi, a new diffusion-based text generation model that introduces remasking allowing model to detect and resample low-confidence tokens during generation.

Abstract: Mask-based Diffusion Language Models (DLMs) struggle to revise incorrect tokens: once a token is generated, it typically remains fixed. The key challenge is to identify potential errors in the inputs. In this paper, we propose Remasking-enabled Diffusion Language Model (RemeDi), a mask-based DLM that introduces remasking as another fundamental mechanism, enabling more flexible text refinement in diffusion-based text generation. To achieve this, RemeDi jointly predicts token distributions and per-token confidence scores at each step. The confidence scores determine which tokens to be unmasked after the current step, allowing the model to identify tokens with low quality and remask them. These remasked tokens can be resampled with richer context in subsequent steps. We design a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model to detect and remask incorrect tokens in addition to predict mask tokens, and reinforcement learning which optimizes full generation trajectories toward higher rewards. Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.

Published as a conference paper at ICLR 2026
DON'T SETTLE Too EARLY: SELF-REFLECTIVE RE- MASKING FOR DIFFUSION LANGUAGE MODELS
Zemin Huang
Yuhang Wang²
Zhiyang Chen Guo-Jun Qi
Zhejiang University
MAPLE Lab, Westlake University
3 Matterwave Intelligence
Institute Advanced Technology, Westlake Institute for Advanced Study (huangzemi n, wangyuhang, chenzhi @west .et.uuukeeeeee.......... cn, guo.uun@@gma 1. com
lab/ DemeDi
ABSTRACT
Mask-based Diffusion Language Models (DLMs) struggle to revise incorrect to- kens: once token is generated, it typically remains fixed. The key challenge is to identify potential errors in the inputs. this paper, we propose Remasking- enabled Diffusion Language Model (RemeDi). mask- DLM that introduces remasking as another fundamental mechanism, enabling more flexible text refine- in diffusion-baser text generation. To achieve this, RemeDi jointly predicts token distributions and per-token confidence scores at each step. The confidence scores determine which tokens to be unmasked after the current step, allowing the model to identify tokens with low quality and remask them. These remasked tokens can be resampled with richer context in subsequent steps. We design a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model detect and remask incorrect tokens addition to predict mask tokens, and reinforcement learning which optimizes full generation trajectories to- ward higher rewards. Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.
INTRODUCTION
Diffusion Language Models (DLMs) have recently emerged as a promising alternative to autore- gressive language models (Nic et al. 2025 Ye et al. 2025 Lou et al.) 2024: Arriola et al.) 2025). A DLM defines a forward process that gradually corrupts text into a noise prior, and learns reverse process to recover clean text (Campeell al. 2022 Lou et al. 2024). Unlike autoregressive models, DLMs do not commit to a fixed left-to-right order, offering greater flexibility in generation and an inherent ability to predict multiple tokens in parallel.
A dominant variant is the mask-based DLM (Nie et al. 2025 Ye ct al. 2025), where the noise is represented by a special mask token. Under this formulation, the model learns to recover masked tokens during training, while assuming that once tokens are unmasked, they are supposed to be cor- rect without having to clean them later. This assumption is problematic: the model may generate wrong tokens, which should be revealed and corrected in later steps when more contexts are avail- able. However, most existing DLMs (Ne et al. 2025 Ye et al. (025) keep already unmasked tokens fixed, preventing them from being revised by self self-refiecting
errors.
To address this, several works have explored methods to revise generated tokens. von Rütte et al. (202)) defines revision of wrong tokens on small-scale models. [Wang et al.] (2025a) applies predictor-correctn samplers by stochastically remasking subset of tokens only at inference time, where the remasking is performed randomly without training the model how to find and remask incorrect tokens. For large-scale DLMs, Seed Diffusion (Song al. 2025) allows all tokens to be resampled at every step.
noise schedule by interpolating between masking and uniform noise, enabling
"Corresponding authors
This project initiated and supported by MAPLE Lab Westlake University.

---

Published as a conference paper at ICLR 2026
-L.DA RemeDi
Code
HumanEval
KBPP
denol steps
Calculate nany grapes Left for the. XX...
Calculate how grapes were left for the pies...
Calculate how grapes were Left for the pies...
ARG-C
Calculate how nany grapes were used for the pies...
CSNBK
Clncerall Tacks
(a) Green and Red tokens are already unmasked in the input quence of current step. Red tokens are remasked. Blue tokens are unmasked in the outputs. Gray tokens keep masked in the out- puts, and we display the token with the highest probability these positions.
AlpacaEval
(b) Radar plot comparing the perfor- mance of RemeDi with other DLMs across various evaluation benchmarks.
Figure I: (a) Illustration of remasking for quality improvement: RemeDi initially predicts the token "left", but later identifies semantic mismatch the phrase "left for the pies". The model then remasks this token and corrects it to the more appropriate "used". (b) RemeDi outperforms existing DLMs in various tasks, including math, code and general benchmark.
However, lacks mechanism to ensure that the number of mask tokens decreases monotonically it is a key feature for diffusion models to ensure decreasing noise levels over steps (Guo et al.) 2025.. so that the mask tokens will eventually vanish the final step to complete the generation.
In this paper, we propose self-reflective remasking approach to train DL.Ms. As illustrated in Fig. Ta aims to train DL.Ms with the ability of finding wrong tokens and turning them back mask ones so that they can be resampled with richer context in later steps. Based on this, we introduce Beemasking-enablee
Diffission Language Model (RemeDi), mask-based DLM that incorporates self-reflective remasking revise already generated but incorrect tokens. RemeDi jointly predicts token distributions and per-token confidence scores. At each diffusion step, high-confidence tokens unmasked while low-confidence ones are (re-) rre)mms regardless of whether they have been previously unmasked.
The key challenge is train the model how to remask incorrect tokens in seff-reflective manner. To this end, we design a remask-aware training pipeline two stages: 1) Remask SFT, where the model learns identify and remask incorrect tokens, while predicting masked tokens. We construct an input sequence for Remask SFT by randomly masking its tokens or replacing them with random alternatives to simulate the noise. The noise schedule deciding how many tokens are masked or randomly replaced is designed to follow the criterion that the noise level should monotonically decrease over steps. The model is then trained remask and revise incorrect tokens over noisy input sequences. 2) Remask RL, where the model is further fine-tuned with outcome- based reinforcement learning. seeks to optimize the entire generation trajectories toward final outputs with higher rewards by considering how remask and predict tokens in each step.
As shown in Fig. lb RemeDi achieves the state-of-the-at performance among open-source DLMs, achieving competitive results on various benchmark datasets, including math problems (89.1% on GSM8K (Cobee et al. (0.T), 52.9% on MATH (Hendrycks et al. (021)). code generation 73.2% on HumanEval (Chen et al. (021), 1), 59.4% MBPP (Austin al. 202 (), and general tasks (24.5%
on AlpacaEval (Dubois et al. 2024), 85.4% on IFEval (Zhou et al. 2023), and 87.7% on ARC-C (Clark al. 2018)).
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
2.1 MASK-BASED DIFFUSION LANGUAGE MODELS
Diffusion language models (DLMs) have emerged as promising alternatives auto-regressive (AR) models for text generation. Among them, mask-based DLMs (Nie et al.) 2025 Ye et al.; 2025
Zheng et al., 2023 Ou
2024 dominate, generating text by progressively denoising mask
tokens. Recent studies Arriola et al. 2025 Fathi et al. 2025 Huang & Tang 2025 Sahoo et al.
2025 Wang et al. 2025b Gat al. 2025) have increasingly explored the fusion of AR and diffusion models, often through an iterative block-wits wise decoding strategy: inference proceeds by iteratively appending block of mask tokens to the input sequence and denoising it, repeating until the EOS token is generated. This paradigm inherits the strengths - flexible generation order and parallel decoding from DI DLMs, and cache efficiency from AR yielding faster inference without sacrificing quality. In work, we adapt LLaDA-8BB MaDDA-BB-Instie Instruct (Nie et al.) (2025) to variable-length block-by- block generation, serving as the backbone for our remasking mechanism.
2.2 REVISING ERRORS IN DIFFUSION LANGUAGE MODELS
A key limitation of mask-based DLMs is their inability to revise tokens once unmasked, even if they are incorrect. Existing efforts to address this fall into two categories. The first category (Campbell et al., 2022 Wang et al. 2025a) applies predictor-corrector samplers without training, for example by stochastically remasking a subset of tokens during inference. These methods lack a mechanism identify which tokens are actually wrong. As a result, they have to rely on many extra sampling steps to take effect, which are inefficient and hard to optimize. The second category modifies the diffusion process to enable revision during the reverse diffusion process, e.g., combining mask diffu- sion process with either the uniform diffusion process (von Rütte et al. (025) or edit-based diffusion process (Havasi et al. 2025 Song et al. 2025).
In short, none of these approaches provides a principled way to detect and selectively correct errors during generation. In contrast, RemeDi fulfills self-reflection by identifying and remasking error- prone tokens through two-stage learning pipline, and jointly training the model to resample the remasked tokens in later steps.
METHODS
3.1 PRELIMINARIES: MASK-BASED DIFFUSION LANGUAGE MODELS
Diffusion Language Models (DLMs) aim to model text generation by approximating the probabil- ity distribution Paata over finite vocabulary process in which the unknown data distribution Pasta at t = 0 gradually evolves into a simple prior distribution Porior at t T (Lou et al. (2024). At intermediate times f, we denote the distribution as Pr. Formally, this diffusion process can be described by linear ODE involving a diffusion matrix Q
21, 2,.... VJ. They define a discrete diffusion
dp: Q.Pt, Pa dt
Pdata
PT = Pprior-
(1)
While can be defined continuously as in Eq. (1), practice work with discrete timesteps to.N.
In this paper, we focus on mask-based DLMs, where Ppwior is a distribution that puts all its mass on the mask state, denoted [M]. Given clean sequence To ~ Pdata corrupted sequence 2i is obtained by randomly replacing part of the tokens with the mask token [M]. The model is trained to recover zo by predicting each mask token zi with the output probability p (wglar)) The objective is:
| 1 [M]) log (
Ldiffusion (0) Ee
where x4 is a sequence of length L, sampled from the forward process, 1()) is an indicator function ensuring that the loss is computed only on mask tokens, following Nie et al. (2025).
During inference, the reverse diffusion process begins with a sequence of only mask tokens and proceeds for N steps monotonically decreasing timesteps lo.v. At step 1,-1, 1, the model takes the
3

---

Published as a conference paper at ICLR 2026
Sample 22 Seartt-Taeeeehaaat
M
[
2
sample xi, from plin (igg- V Ul xin [M]
ha
Po.
Lincar
RemeDi
Lincar & Softmax
/TPS ".r.
fup
TPS
Transformer Block
UPS
zero init
Transformer Block
Transformer Block
Transformer Block
zero init
Transformer Block
Transformer Block
Input Emmeeding
X
Figure 2:
The structure of RemeDi, including Unmasking Policy Stream (UPS) to predict confi- dences hg for selecting the set of unmasking tokens Un and Token Prediction Stream (TPS) to predict the token value when unmasking a masked position.
partially masked sequence Ti,,- as input and predicts all mask tokens simultaneously. subset of tokens is unmasked to obtain Tt,, according to the noise schedule and the unmasking policy (c.g., unmasking tokens with the highest confidence), while the remaining predictions are remasked and deferred to later steps.
A limitation of this paradigm is that once a token unmasked, it remains fixed in subsequent steps. In early stages, limited unmasked tokens often lead to unreliable predictions, resulting in errors that persist through the remainder of the generation process. As generation progresses, additional context may reveal these errors, but current paradigm offers way correct them. This motivates the ability to remask tokens, allowing the model to remask earlier predictions back to the mask token and predict them again using richer context in later steps.
3.2 REMEDI
We propose RemeDi, DLM that can identify and remask low-confidence tokens during generation enable iterative self-reflection. We extend the standard transformer into a dual-stream transformer architecture as shwon in Fig. 2 which comprises:
Token Prediction Stream (TPS): stack of transformer blocks that predict probabilities pb[ (-1,,) for masked tokens as in a typical DLM (Nie et al. 2025).
Unmasking Policy Stream (UPS): Another stack of transformer blocks that output token- wise confidence score h3- represents the model's confidence over the output tokens, indicating if they should be unmasked with high confidence. Otherwise, if the confidence is too low for a token, it should be kept masked or remasked so that it could be sampled or resampled later.
The two streams run in parallel. During inference, UPS is inserted periodically and receives hidden states from TPS as input, producing auxiliary representation fups for confidence scoring. The two streams perform bidirectional feature sharing: UPS layers are conditioned on firs. and their outputs also feed back into TPS to enrich its representations. At the final layer, p and h are produced simultaneously using two independent linear heads applied to fres and fups, respectively. More details about the model structure for these two streams can be found in Appendix B.1
4

---

Published as a conference paper at ICLR 2026
The token generation proceeds through iterative denoising steps. Given 2t,, the imput, UPS first predicts confidence score n at each position and select subset of positions Un to unmask the current step. Then, for the positions selected to be unmasked, if they have already been unmasked in - they remain unchanged; otherwise they are sampled from ()(-(tt,,.1) predicted by TPS. Unlike existing mask-based DLMs where tokens are fixed once being unmasked, RemeDi re-decides a token to be unmasked or (re- mee-)aakeed at each step by its trained confidence score. Thus, it is possible that an already generated token is assigned with low confidence and remasked. allowing to be resampled in later steps. A noise schedule controls that the total number of unmasked tokens increases linearly from to T. (Nie et al., 2025), that the number of mask tokens approaches at the final step.
In the following sections, we elaborate on how to train RemeDi with Remask SFT and Remask RL algorithms.
3.2.1 REMASK SFT
Traditional bassed DLMs conduct SFT with randomly masked input sequences (Nie et al. 2025 Lou et al. 2024), while RemeDi needs to detect and remask possible incorrect tokens that arise during the reverse diffusion process, so they can be resampled later steps. To achieve this, in SFT treat such incorrect tokens second noise type in addition to the first noise type of mask tokens in mask-based DLMs, and train the model to recover mask tokens as well as identify unmasked tokens that should be remasked.
To simulate inference inputs at a diffusion time f. we construct training samples 2y from clean text To by applying two types of noise: given randomly sampled diffusion time (0, 1), we set the corresponding mask ratio Pr,mask alongside the incorrect token ratio Pr.incoment- With both ratios, we randomly mask tokens with ft.mask. Then, among the remaining unmasked positions, sample subset with the ratio Pt. Pt,incenett and replace each selected token with a random alternative to simulate the incorrect tokens that may occur in the reverse diffusion process.
As aforementioned, during the reverse diffusion process, the noise level, defined as the number of mask tokens, should decrease monotonically (Guo et al., 2025). Since all incorrect tokens in an input sequence of length L must be remasked as designed below for the SFT, we require:
(3) to ensure a monotonically decreasing number of mask tokens as outputs. Otherwise, remasking all incorrect tokens would increase the total number masks in the next step, violating the principle that the number of mask tokens should decrease at each diffusion step.
Tp iinoorete pt.skk) L] [[t.mmsk L]
Considering the above inequality, we choose (r. mask and pr,jinnrrest 4r-1(1-1)) (r is a constant) following (Ni et al. 2025 yon Rütte et al. 2025). We set 0.1 in our experiments, under which it is not hard to see that the incquality always holds on t E [0,
Remask SFT Algorithm. During training, in addition to the typical diffusion loss in Eq, 2 we supervise the unmasking score he with a binary cross-entropy (BCE) objective across all token positions. We construct the training label y based different token types:
A clean token (i Seleun {i ot indicating they should remain unmasked.
xt1) receives a positive unmask label y² 1,
An incorrect token (i Sincoreet = {i of a zg. of * [M]}) receives a negative unmask label y = 0, indicating that they should be remasked.
A mask token (i E Smusk = {i x = [M]]] is assigned soft unmask label y = p (af (x,), equal to the predicted probability of the ground-truth token z A higher probability indi- cates a higher likelihood that the predicted token is correct and thus should be unmasked.
With unmask labels assigned above, we seek to minimize
CUPS = Σ BCE(((h)), v²)
(4)
where of-) is the sigmoid function. Thus, the overall Remask SFT objective is: [() Laittusion (0) + Xues (ups
(5)
5

---

Published as a conference paper at ICLR 2026
Algorithm Input sequence construction and loss calculation in Remask SFT
Require: Clcan sequence zo [ar..... ................................ .x6] length L. Model M with learnable parameters 8
1: Sample E (0, 1) according to the noise schedule, obtaining Pt.mask and iiniorcrreat-
2: Construct noisy input I=:
3: For each position i, replace zó with [M] w.p. Pt.makk
4: Among remaining positions, replace zb with random alternative token w.p. Ptiincomeet
5:
Define index sets:
Smadk {i = = [M]] Siivenert = (i ue ui A we [M]] Selean = (i a/ abd
6: Get model outputs: [pe, ho] = M(x;:0)
7: Calculate the diffusion loss, mask tokens only: Lamaun (0)
Smnkk EieSmni log (6(
E Satein
8: Get labels for UPS: y
0
E Sincoreet
stopgrad (pà (=))==))
E Smisk
9: BCE loss:
(o)) represents the sigmoid function
Lups (0) - - Σ
logo (1-y')laa log (1 -o(bl)
Total loss: (()) Camtuian (0) Aups (0)
where Aues balances the two losses.
Finally, we summarize the Remask SFT Algorithm where we elaborate on how to construct the input sequence and calculate the loss function for the Remask SFT.
3.2.2 REMASK RT.
After training with Remask SFT, we further fine-tune the model with outcome-based reinforcement learning (RL) to optimize the full generation trajectory (Huang al.) (025). Specifically, we rein- force the generation process with N denoising steps, beginning from an all-mask prior Tig Ppeior to and proceeding through timesteps to.N-
At each step t,., RemeDi generates Zt., from Pt,-1 by invoking two coupled policies: an unmasking policy that chooses subset of positions Un = [u, (1)...., un (K,)] ]] to unmask, and a token predic- tion policy that samples tokens at the chosen positions. Unlike standard DLMs, which never remask revealed tokens, RemeDi allows previously unmasked tokens to be remasked, enabling revision of earlier predictions.
Unmasking policy. The UPS produces a per-token confidence score hig indicating how strongly the model believes token position is correct (if unmasked) or predictable (if masked), At infer- ence, we rank tokens by their confidence scores and prioritize high-confidence ones to unmask. The number of unmasked tokens Kn, each diffusion step is determined by linearly increasing from to L.. During RI training, we construct unmasking policy to sample Un [u,, (1), ) ..... "n, (K..) using the Plackett-Luce model (Plackett 1975): based on hgg we use multinomial distribution and sequentially sample Kn, positions from of sampling Un is:
L.} without replacement. Formally, the probability
K
exp(h (1:)
unmnskk T.n
(Un Dt., .)
(6)
Eje(un (1).... Ejg(u. exp(1i).))
Token prediction policy. For each position E Un, if mt -|,_, [M]. the model samples token from (1(()tt..,)) otherwise, the token remains unchanged as in the input. The probability of generating St. given Zt,-, and Un, is:
token (4n)
II
pb(at. (to-r)-
iEUi., - A..-;-IMM =IM]

---

Published as a conference paper at ICLR 2026
Joint policy. Thus, the probability of transitioning from Pt,, to 2t., is the product of the unmask- ing probability and the token prediction probability:
To (Ztt Ita-.)
7g.n (Un (ta-.) Tokee T0, (Zt, (
With the probability defined in Eq. we apply outcome- based reinforcement learning to encourage generation trajectories Fto.w that lead to correct final responses Tty- Specifically, we adopt GRPO (Shao et al. 2024), scalable RL paradigm for language models. The reward is defined according to task type: verifiable correctness for math and code, and reward-model evaluation for open-ended questions. Further details datasets and reward design provided in Appendix [.2.4
As shown in Fig. T1 of Appendix A after Remask SFT and RL training, the learned hg serves reliable indicator assess the quality of input tokens. Tokens already unmasked in the input typically receive high confidence scores. However, when certain tokens are assigned low confidence, they are more likely be inadequate and are remasked for re-prediction in subsequent steps. suggests that the UPS-prest preepicdd confidence scores provide reliable estimate of per-token quality for the unmasking policy.
EXPERIMENTS
RemeDi enables remasking DLM capable of variable-length block-wise generation (Arriol et al., 2025) to support variable-length generation, a key feature for enabling the real-world DLM generate unfixed number of blocks (see Appendix [B.2.2 for details). Since there source large-scale variable-length block-wise DLMs, we adapt our model from LLaDA, a widely used benchmark DLM. Starting from LLaDA's model weights initialization, RemeDi undergoes two stages of supervised fine-tuning and RL. We detail the training configurations in Appendix B.2 and the evaluation metrics in Appendix B.3.2
open-
Table 1: Model performance math and code generation benchmarks. We highlight the best- performing model among compared DLMs bold. "_" indicates unknown cases not mentioned in original papers.
Math MATH
Code
Method
GSMSK
GPQA
HumanEval
MBPP
Diffusion Language Models
Dream Ye ct al.
2025)
82.1
49.6
30.6
59.8
59.6
LLaDA Nic
2025
78.3
38.9
28.1
45.7
39.0
LLaDA ReMDM Wang
2025al
81.4
38.5
44.5
37.8
d1-LLaDA Zhao
2025
82.1
37.8
44.7
wd1-LLaDA Tang el al.
2025
82.3
LLaDA 1.5 (Zhu el al. 2025
83.3
42.6
36.9
$2.4
42.8
LLaDOU (uaang
88.1
44.6
59.1
51.6
RemeDi (+ Remask SFT)
86.3
51.4
32.6
71.3
57.8
RemeDi (++ Remask RL)
89.1
52.9
29.5
73.2
59.4
Auto-regressive Models
LLaMA2 7B Touvron et al.
2023
14.6
2.5
28.4
12.8
20.8
MetaMath 7B Yu et al. 2023
66.5
19.8
Cad 7B (Rzzerr al.
2023
34.8
44.4
Deepseek 7B (Bi et al.
2024)
63.0
15.8
48.2
35.2
DeepseckMath 7B (Shao ct al.
2024
88.2
51.7
DeepseekCoder Coder 7B (Guo el
2024
66.1
65.4
LLaMA3 8B Dubey
2024
78.3
29.6
31.9
59.8
57.6
Gemma2 9B
Team 2024
76.7
44.3
32.8
68.9
74.9
4.1
RESULTS
To evaluate various capabilities of RemeDi in different aspects, we conducted detailed comparisons against existing large language models of comparable scale in Tab. and Tab. 2 including both

---

Published as a conference paper at ICLR 2026
Table 2: Model performance on general tasks. We highlight the best-performing model among compared DLMs in bold. "_" indicates unknown cases not mentioned in original papers.
Method
Hellaswag
ARC-C
IFEval
AlpacaEval
Diffusion Language Models
Dream Ye
2025
70.3
79.2
67.5
5.9
LLaDA Nie ct al. 2025)
69.7
83.9
70.0
11.2
LLaDA 1.5 [hu el al. 2025
70.5
83.5
73.5
13.9
RemeDi (+ Remask SFT)
71.1
85.2
81.9
12.5
RemeDi (++ Remask RL)
72.2
87.7
85.4
24.8
Auto-regressive Models
LLaMA2 7B
Touvron et al.
2023
51.5
57.3
Deepseek 7B [Bi al. 2024
68.5
49.4
LLaMA3 8B
(ubey al. 2024
75.5
82.4
DLMs and auto-regressive models. We select nine popular benchmarks across general tasks, math- ematics, coding, and human preference domains.
After Remask SFT, RemeDi demonstrates high performance on almost all these benchmarks. not only achieves the state of the art performance among existing DLMs, but also outperforms auto- regressive models similar model size. On math benchmarks, RemeDi after Remask SFT achieves 86.3% on GSM8K and 51.4 11.4% on MATH, surpassing MetaMath with math-specifee instruction tun- ing. even on par with DeepseckMath using math-specific reinforcement Icarning. On code generation benchmarks, RemeDi achieves 71.3% HumanEval, outperforming LadA and Deepseek Coder. For general natural language tasks, RemeDi also demonstrates strong performance in common knowledge answering (85.2% on ARC-C) and instruction following (81.9% on IFEval) tasks. It also aligns well with human preference (12.5% on AlpacaEval), outperforming other DLMs such Dream and LLaDA.
After Remask RL, RemeDi achieves further improvements across wide range of math, coding and general tasks. For example, the accuracies on GSM8K and MATH reach 89. 59.1% and 52.9% respec- tively, outperforming all compared DLMs and AR models. Among all the benchmarks, RemeDi achieves its most substantial improvement on the AlpacaEval (Dubois et al.) (024) benchmark, with +12.3% gain over the Remask SFT model. This demonstrates the effectiveness of our approach in post-training the model's s ability for a broad range of tasks.
4.2 VISUALIZATION AND ANALYSIS
We visualize how remasking improves text generation in RemeDi in Fig. 3 erated the token "making." After generating the object "tests and estimators," it found that "making" not the proper verb this verb-object structure. Thus, the model remasks it and opts for the more appropriate "developing." This example shows RemeDi's ability to iteratively refine its output content. We provide more examples in Appendix A demonstrating that RemeDi able to perform a variety of operations such as replacing, inserting and deleting with the remask mechanism.
The model initially gen-
To provide quantitative analysis, we calculate the frequencies of remasking in block of length 32 on MATH-500 (Lightman et al.) 2023). HumanEval (Chen et al.) (021). and AlpacaEval (Dubois) al., 2024). In Fig- 4 we can see that remasking occurs most frequently in code generation, fol- lowed by mathematical reasoning, and general tasks. This pattern may be attributed to differences in structural constraints: code requires strict syntactic correctness, and mathematical solutions demand formally structured derivations, whereas responses open-ended problems allow more flexibility.
We also analyzed remasking frequencies across different difficulty levels on MATH-500, as shown in Tab. 3 RemeDi tends remask more frequently as the difficulty increases, rising from about tokens per block at level 1-2 to nearly 14 tokens at level 4-5. This pattern suggests that iterative refinement becomes increasingly necessary for harder problems.
8

---

Published as a conference paper at ICLR 2026
...A statistical nodel is nathematical representation that explains how data is generated, often in simplified and idealized way. It forms the foundation for understanding the data, naking hypotheses about, and making and population...
step
statistical model mathematical representation explains how is generated, in simplified and idealized way. It forms the foundation for understanding the data, -making tests and estinators, and naking the inference inference...
step 424
...A statistical model nathematical representation that explains how data is generated, often in simplified and idealized way. It forms the foundation for understanding the data, developing tests and estinators, and making the basis statistical...
step 425
Figure 3: An example of the step-by-step generation process. Green and Red are already unmasked in the inputs. Red tokens are remasked. Blue tokens are unmasked in the outputs. Gray tokens remain masked, and we display the token with the highest probability these positions. More examples can be found in Appendix
033
030
019
300
0200
Coalts per
Aamask
(a) MATH-500 (11.81 + 10.23)
(b) HumanEval (28.52 + 12.04) (c) Alpaca-Eval (2.78 5.33)
Figure 4: Distribution of remasking frequencies per block across different evaluation datasets. The numbers in parentheses indicate the mean and standard deviation for each dataset.
4.3 ABLATION STUDIES
Remask SFT We compare the improvement brought by the Remask SFT (introduced in Sec. 3.2.1 with that of vanilla SFT, under the same training configuration detailed in Appendix [B2.55] We start from baseline model that has already completed the warm-up phase tuning for variable-length block-wise generation, and perform training on the full code-category dataset and the open2math- 1M-gpt-4.1-man 1-mini dataset mentioned in Appendix B.2.T As shown in Tab. 4 Remask SFT outper- forms vanilla SFT on all benchmarks, especially on MATH-500 (+2.6%) and HumanEval (+1 (+1.8%) demonstrating that Remask SFT is an effective training method to improve DLM''s performance.
Remask RL We compare Remask RL with LLaDOU RL (Huang al., 2025). another algo- rithm that also reinforces the whole generation trajectories in the reverse diffusion process. Since LLADOU RI. is developed on LLaDA, we also implement Remask RL on LLaDA for the sake of fair comparison. All experiments
conducted on GSM8K with a generation length of 256, 64 denoising steps, a block length of 64, and temperature 0.7, while other hyperparameters follow the LLaDOU setup (see Appendix B.2.5).
Remask RL demonstrates advantages in both convergence speed and performance. As shown in Tab. Remask RL achieves higher final accuracy of 83.33%, with particularly noticeable im- provement in early training stages (e.g., 80.00% vs. 77.58% at step 50). This indicates that the more flexible remask process contributes to both faster convergence and stronger model performance.

---

Published as a conference paper at ICLR 2026
Table Statistics of the remasking frequencies per block (block size is fixed to 32) when generating responses to questions with different difficulty levels in MATH-500.
Difficulty Level
Remasking Frequencies / Block Accuracy
9.13 + 9.54
86.04%
2
8.91 7.29
80.21%
3
10.13 + 8.64
64.48%
13.91 11.44
50.00%
5
13.95 = 11.12
19.25%
Table 4: Experiment results after supervised tuning with different algorithms. The baseline model already tuned to variable-length block-wiss generation DLM (see Appendix B.2.2)
Method
GSM8K
MATH-500
HumanEval
MBPP
Baseline
80.3
34.7
41.5
42.6
Vanilla SFT
83.1
40.1
48.2
43.4
Remask SFT
83.6
42.7
50.0
44.0
Table 5: GSM8K pass @
accuracy comparison between Remask and LLaDOU RL
Training Steps
Remask RL
LLaDOU RL
50
80.00%
77.58%
100
81.40%
78.86%
150
81.59%
80.00%
200
83.33%
82.35%
5 CONCLUSION
In this paper, introduce the Remaskingeenabled Diffusion Language Model (RemeDi), self-reflective remasking mechanism to address the limitation of existing mask-based DLMs that they cannot revise generated tokens. In RemeDi, remasking is achieved by predicting confidence score to identify noisy tokens, allowing them to be remasked and then resampled with richer context in later steps.
Through a two-stage training pipeline of Remask SFT and Remask RI., RemeDi achieves the state- of-the-art performance among open-source DLMs. Our analysis further shows that the learned con- fidence scores provide a reliable signal of per-token quality during generation. RemeDi opens promising direction for self-reflective text generation, further releasing the full potentials of DLMs to solve complex tasks with higher quality.
REPRODUCIBILITY STATEMENT
We provide link containing the inference code and model weights, details of the datasets and configurations used in both Remask SET and RL Appendix B.2 and the evaluation settings in Appendix B.3
ACKNOWLEDGEMENTS
This work supported by the National Natural Science Foundation of China under Grant No. 92467104 104, and Zhejiang Leading Innovative and Entrepreneur Team Introduction Program (2024R01007).
10

---

Published as a conference paper at ICLR 2026
REFERENCES
ajibawa 2023. Maths-college. Hugging Face Datasets.
URL htttps::/hugginggggegte/ /
datasets/ ja jibawa-2223 2000222at22222MMCO lege
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autore- gressive and diffusion language models. arXiv preprini arXiv:2503. 099773 2025.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry. Quoe Le, et al. Program synthesis with large language models. arXiv preprint arXiv:21080077:0 2021.
Akhiad Bercovich, Itay Levy, Izik Golan, Mohammad Dabbah, Ran El- El-Yaniv, Omri Puny, Ido Galil, Zach Moshe, Tomer Ronen, Najeeb Nabwani, Ido Shahaf, Oren Tropp, Ehud Karpas, Ran Zil- berstein, Jiaqi Zeng, Soumye Singhal, Alexander Bukharin, Yian Zhang, Tugrul Konuk, Ger- ald Shen, Ameya Sunil Mahabaleshwarkar, Bilal Kartal, Yoshi Suhara, Olivier Delalleau, Zijia Chen, Zhilin Wang. David Mosallanezhad, Adi Renduchintala, Haifeng Qian, Dima Rekesh, Fei Jia, Somshubra Majumdar, Vahid Noroozi, Wasi Uddin Ahmad, Sean Narenthiran, Alek- sander Ficek, Mehrzad Samadi, Jocelyn Huang, Siddhartha Jain, Igor Gitman, Ivan Moshkov, Wei Du, Shubham Toshniwal, George Armstrong, Branislav Kisacanin, Matvei Novikov, Daria Gitman, Evelina Bakhturina, Jane Polak Scowcroit, John Kamalu, Dan Su, Kezhi Kong, Markus Kliegl, Rabeeh Karimi, Ying Lin, Sanjeev Satheesh, Jupinder Parmar, Pritam Gundecha, Bran- don Norick, Joseph Jennings, Shrimai Prabhumoye, Syeda Nahida Akter, Mostofa Patwary, Abhinav Khattar, Deepak Narayanan, Roger Waleffe, Jimmy Zhang, Bor-Yiing Su, Guyue Huang, Terry Kong, Parth Chadha, Sahil Jain, Christine Harvey, Elad Segal, Jining Huang. Sergey Kashirsky, Robert McQueen, Izzy Putterman, George Lam, Arun Venkatesan, Sherry Wu, Vinh Nguyen, Manoj Kilaru, Andrew Wang, Anna Warno, Abhilash Somasamudramath, Sandip Bhaskar, Maka Dong, Nave Assaf, Shahar Mor, Omer Ullman Argov, Scot Junkin, Olek- sandr Romanenko, Pedro Larroy, Monika Katariya, Marco Rovinelli, Viji Balas, Nicholas Edel- man, Anahita Bhiwandiwalla, Muthu Subramaniam, Smita Ithape, Karthik Ramamoorthy, Yut- ing Wu, Suguna Varshini Velury, Omri Almog, Joyjit Daw, Denys Fridman, Erick Galinkin, Michael Evans, Katherine Luna, Leon Derezynski, Nikki Pope, Bileen Long, Seth Schneider, Guillermo Siman, Tomasz Grzegorzek, Pablo Ribalta, Monika Katariya, Joey Conway, Trisha Saar, Ann Guan, Krzysztof Pawelec, Shyamala Prayaga, Oleksii Kuchaiev, Boris Ginsburg, Oluwatobi Olabiyi, Kari Briski, Jonathan Cohen, Bryan Catanzaro, Jonah Alben, Yonatan Geif- man, Eric Chung, and Chris Alexiuk. Llama-nemotron: Efficient reasoning models, 2025. URL /tttp:://arii.. www.orggaas....... 2ttp://0/002055004999999
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprini urXiv:2401.022) 002954, 2024.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnad Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 011222266-22799 2022.
Mark Chen, Jerry Tworek, Heewoo Jun Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger. Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss. William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Mc- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code. 2021.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv, hos/1803.05557. 2018.
11

---

Published as a conference paper at ICLR 2026
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems, 2021. URL https://arxiv. org/abs/2110.14106 9, 2021.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv e-prinis, pp. arXiv-2407, 2024.
Yann Dubois, Balázs Galambosi, Percy Liang, and Tatsunori B Hashimoto. Length-controlled al- pacaeval: A simple way to debias automatic evaluators. arXiv preprint arXiv:2404.044: 044422 2024.
Nima Fathi, Torsten Scholak, and Pierre-Andrs Noël. Unifying autoregressive and diffusion- sequence generation. arXiv preprint arXiv:2504.00 06446616666 2025.
Itai Gat. Heli Ben- Ben-Hamu, Marton Havasi, Daniel Haziza, Jeremy Reizenstein, Gabriel Synnaeve, David Lopez-Paz, Brian Karrer, and Yaron Lipman. Set block decoding is language model inference accelerator. arXiv preprint urXiv:2509.04//5 2025.
Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Yu Wu, YK Li, et al. Deepseek-coder When the large language model meets programming- the rise of code intelligence. arXiv preprint arXiv:2401. 14196, 2024.
Zhehao Guo, Jiedong Lang. Shuyu Huang, Yunfei Gao, and Xintong Ding. A comprehensive review noise control of diffusion model. arXiv preprint wrxv:2502.046669 2025.
Marton Havasi, Brian Karrer, Itai Gat, and Ricky TQ Chen. Edit flows: Flow matching with edit operations. arXiv preprint mxXiv:25060090/8 2025.
Zhiwei He, Tian Liang. Jiahao Xu, Qiuzhi Liu, Xingyu Chen, Yue Wang, Linfeng Song, Dian Yu, Zhenwen Liang, Wenxuan Wang, Zhuosheng Zhang. Rui Wang. Zhaopeng Tu, Haitao Mi, and Dong Yu. Deepmath- 103k: A large-scale, challenging, decontaminated, and verifiable mathemat- ical dataset for advancing reasoning. 2025. URL
hbtp:/bb//22500 250.. 11456
Dan Hendrycks. Collin Burns. Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurlPS, 2021.
Chihan Huang and Hao Tang. Ctrldiff: Boosting large diffusion language models with dynamic block prediction and controllable generation. arXiv preprini urXiv:2505./44455 2025.
Siming Huang, Tianhao Cheng, Jason Klein Liu, Jiaran Hao, Liuyihan Song. Yang Xu, J. Yang,
J. H. Liu, Chenchen Zhang, Linzheng Chai, Ruifeng Yuan, Zhaoxiang Zhang, Jie Fu, Qian Liu, Ge Zhang, Zili Wang, Yuan Qi, Yinghui Xu, and Wei Chu. Opencoder: The open cookbook for top-tier code large language models. 2024. htll:htts://arr wtt.aaawwww.oogpp///////./ 04905
Zemin Huang, Zhiyang Chen, Zijun Wang, Tiancheng Li, and Guo-Jun Qi. Reinforcing the diffusion chain of lateral thought with diffusion language models. arXiv preprint arXiv:2505.10444 2025.
Jacyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. 06666 2025.
arXiv preprint
Hynek Kydlíček. Math-verify: A robust mathematical expression evaluation system.
https:
//githus www.//gg/.oooooooooooo..o......o ng face/ http//Mttt-Veiiii Verify 2025.
Jia LI. Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. Numinamath. [https: http::://hhhgggnggeeeee
c0 Sc//Al-MO/Kuminaaathh http:://tttaa.....///g..u...... aimo wtt.ppppooreess-ppieebo..o.o-. httpa///....o.mmmmecceeeetttnnnn nww.aanmiiaaaiaiiiamaim dataset. http://ttt/muiiiaaaaatttttteeede
numina/ 2024.
Jijie Li, Li Du, Hanyu Zhao, Bo wen Zhang, Liangdong Wang, Boyan Gao, Guang Liu, and Yonghua Lin. Infinity instruct: Scaling instruction selection and synthesis to enhance language models,
2025. URLpsttps: httphttss///axxx
12

---

Published as a conference paper at ICLR 2026
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, llya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference Learning Representations, 2023.
Chris Yuhao Liu, Liang Zeng, Jiacai Liu, Rui Yan, Jujie He, Chaojie Wang, Shuicheng Yan, Yang Liu, and Yahui Zhou. Skywork-reward: Bag of tricks for reward modeling in llms. arXiv preprint arXiv:2410. 18451, 2024.
Chris Yuhao Liu, Liang Zeng, Yuzhen Xiao, Jujie He, Jiacai Liu, Chaojie Wang, Rui Yan, Wei Shen, Fuxiang Zhang, Jiacheng Xu, Yang Liu, and Yahui Zhou. Skywork-rewardd-2: Scaling preference data curation via human-ai synergy. arXiv preprint arXiv:2507.01352, 2025.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Forty-first International Conference on Machine Learning, 2024.
Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Tang. Manan Roongta, Colin Cai, Jeffrey Luo, Tianjun Zhang, Erran Li, Raluca Ada Popa, and Ion Stoica. Deepscaler: Surpassing ol ol-preview with a 1.5b model by scaling rl, 2025. Notion Blog.
Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xiubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang. Wizardcoder: Empowering code large language models with evol-instruct, 2023.
Konstanty Marczak. science.gaa qa. Hugging Face Datasets, 2023. URL URLAtttss//huugigfaaa co//aatasets/Koostattttyyo http///w/ttttttt/////eeneeeeeeee science_ga
Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orca-math: Unlocking
the potential of slms in grade school math, 2024.
mifoundations dev, open2math- 1-mp-mmath-mmm-gp-----mmmm Hugging Face Datasets, 2025. URL https: http://eeeettttttttooooottttttt.
1M gpt - 4.
mini
Dhruv Nathawani, Shuoyang Ding, Vitaly Lavrukhin, Igor Gitman, Somshubra Majumdar, Evelina Bakhturina, Boris Ginsburg, and Jane Polak Scowcroft. Nemotron--Post-Trainngg Dataset-v2, August 2025. URL htttp::/ wig.hhgggggacg....oococ.o co//attasets/vvidall
Nemotron Post Training Dataset-v2
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai
Lin, Ji-Rong Wen, and Chongxuan Li.
Large language diffusion models. arXiv preprint
arXhv:2502.09999, 2025.
Jingyang Ou Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint wxXiv:2406.037336 2024.
Robin L Plackett. The analysis of permutations. Journal of the Royal Statistical Society Series C:
Applied Statistics, 24(2):193-200 1975.
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. Gpqa: A graduate-level google-proof q&a benchmark, 2023.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, et Code llama: Open foundation models for code. arXiv preprint arXiv:2308. 12950, 2023.
Subham Sekhar Sahoo, Zhihan Yang. Yash Akhauri, Johnna Liu, Deepansha Singh, Zhoujun Cheng, Zhengzhong Liu, Eric Xing, John Thickstun, and Arash Vahdat. Esoteric language models. arXiv preprint arXiv:2506.0199 2025.
Zhihong Shao, Peiyi Wang. Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi. Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseckmath: Pushing the limits mathematical reasoning in open language models. arXiv preprint arXiv:2402.033300 033000 2024.
13

---

Published as a conference paper at ICLR 2026
Yuxuan Song, Zheng Zhang. Cheng Luo, Pengyang Gao, Fan Xia, Hao Luo, Zheng Li, Yuehang Yang, Hongli Yu, Xingwei Qu, et al. Seed diffusion: A large-scale diffusion language model with high-speed inference. arXiv preprint vrXiv:2508.021993 2025.
Xiaohang Tang, Rares Dolga, Sangwoong Yoon, and Ilija Bogunovic. wdl: Weighted policy opti- mization for reasoning in diffusion language models. arXiv preprint arXiv:2507.088386 2025.
Gemma Team. Gemma. 2024. doi: 10.47740KAGGLL///3330 BtLLhttp::/// http:://wwwwwwgggll..o
com/m /3001
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Dimitri von Rütte, Janis Fluri, Yuhui Ding, Antonio Orvieto, Bernhard Schölkopf. and Thomas Hofmann. Generalized interpolating discrete diffusion. arXiv preprint arXiv:2503.0444 2025.
Guanghan Wang. Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference- scaling. arXiv preprint ar arXiv:2503 00307, 2025a.
Xu Wang, Chenkai Xu, Yijie Jin, Jiachun Jin, Hao Zhang, and Zhijie Deng. Diffusion Ilms can do faster-than-ar inference via discrete diffusion forcing. arXiv preprint arXiv:2508.091 2025b.
Yunhui Xia, Wei Shen, Yan Wang, Jason Klein Liu, Huifeng Sun, Siyue Wu, Jian Hu, and Xiaolong Xu. Leetcodedataset: A temporal dataset for robust evaluation and efficient training of code llms. arXiv preprint arXiv:2504 2025.
Tengyu Xu, Eryk Helenowski, Karthik Abinav Sankararaman, Di Jin, Kaiyan Peng, Eric Han, Shao- liang Nie, Chen Zhu, Hejia Zhang, Wenxuan Zhou, et al. The perfect blend: Redefining rihf with mixture of judges. arXiv preprint arXiv:2409.23370, 2024.
Zhangchen Xu, Yang Liu, Yueqin Yin, Mingyuan Zhou, and Radha Poovendran. Kodcode: A diverse, challenging, and verifiable synthetic dataset for coding. arXiv preprint 2rXiv:2503.0295.. 2025.
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong, Dream 7b, 2025. URL putttss:/hkunne github. http.////////////eeeeeeeeeeeeeee
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309. 12284, 2023.
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309. 05653, 2023.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1836-3847, 2023.
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint rrXiv:2504.12216 2025.
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. arXiv preprint arXiv:2302.05737. 2023.
Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. Instractionnfollloiing evaluation for large language models, 2023. URL https:
// /aaxxivvorge http://gg/abb//////// 2231000111111111111111
Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu. Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, et al. Llada 1.5: Variance-reduced preference optimization for large language diffusion models. arXiv preprint arXiv:2505. 2025.
14

---

Published as a conference paper at ICLR 2026
A GENERATION PROCESS OF REMEDI
To better understand how RemeDi leverages the remasking mechanism, we visualize intermediate steps when solving math, code, and open-ended problems. Since the full responses are usually long, we focus on the token segment where critical remasking occurs. In the following figure, green boxes indicate already generated tokens, blue boxes represent tokens unmasked in this step, red boxes denote tokens remasked in this step, and gray boxes represent tokens that remain masked, showing the token with the highest probability. Key tokens are highlighted with bounding boxes.
We find that remasking enables diverse forms of revision beyond simple correction, including:
Correcting calculation errors: Remask can correct calculation errors. shown in Fig. the model initially predicted "\div" as the most probable operator and generated "0" However, since the actual operator generated was "\mod", the model remasked the previ- ous "0" and regenerated "5" as the correct result.
Refining text quality: Remasking allows more precise wording. In Fig 6 the initially generated phrase "methyl group" is not adequate when answering this problem. RemeDi replaces them with more precise term "secondary carbon" by remasking.
When two consecutive tokens correspond to a single vocabu- lary token, RemeDi may remask them and merge into one, thereby freeing a slot for sub- sequent generation. In Fig. 7 the separate tokens ". and "\" were remasked and merged into the single token ". ".\", releasing one token slot.
Merging adjacent tokens:
Splitting tokens: Conversely, the model can split token into smaller parts to fill idle positions, ensuring that no tokens remain unused and the denoising process can complete. In Fig. 8 to fill in the slot before "Mish", RemeDi remasks "Mish" and regenerates it as two tokens, "M" and "ish".
Inserting tokens: Remasking also supports insertion. In Fig. to add the word "again" before "bounces". the model first remasked the two tokens "b" and "ounces", and then regenerated the sequence with the insertion.
Deleting tokens: Finally, remasking can delete tokens and replace them with nothing or control symbols. Fig. 10, the phrase "per hour - was removed and replaced with a line break.
These cases illustrate that remasking gives RemeDi considerable freedom to revise its outputs in multiple ways, greatly extending the flexibility of diffusion-based text generation.
Moreover, we illustrate the predicted confidence scores in Fig. I1 In general, unmasked tokens receive higher confidence scores he than masked tokens, unless the model judges them as unreliable and decides to remask them. For example, see the tokens "say" in Fig. [Tb and "in" in Fig. [[ Interestingly, these tokens that are eventually remasked already exhibit relatively low confidence at the step when they were first predicted, as reflected by the lighter background green shading. This suggests that the model was uncertain about them from the start. Once more context revealed in subsequent steps, RemeDi is able to revise such low-confidence tokens into more appropriate alternatives.
B EXPERIMENT DETAILS
B.1 DUAL-STREAM MODEL. STRUCTURE
We construct TPS with the same transformer structure as LLaDA (Nie et al. 2025). comprising 32 transformer blocks. The model weights in this stream are initialized with -L-DA----Inssitn For the UPS, we stack four transformer blocks with the same hidden dimension to construct smaller network with random initialization. These two streams are weakly coupled via bi-directional con- nections at TPS blocks 1, 11, 21, and 31. At each connection point, the output from the previous TPS block is added to the current UPS feature to form the input for the next UPS block, while the output of the current UPS block is added to the output of the corresponding TPS block before it is passed onward. Thus, both streams enrich their representation with features from each other. To pre- serve the TPS's original capability inherited from the pretrained weights, add zero-initialized
15

---

Published as a conference paper at ICLR 2026
W
V4
W
div
step 237
W
step 238
W
modi
step
MX
W
M
step 240
Figure 5: An example of correcting calculation errors with remasking. Question: A group of N students, where N < 50, is field trip. If their teacher puts them in groups of 8, the last group has students. If their teacher instead puts them in groups of 6, the last group has students. What is the sum all possible values of N?
group
position
brom
step 83
attack
either
methry
pesition
the carbon group
position
The al group Is at
position
brom
step 84
can
methyl
the
methyl carbon position
group
position
step 91
Ins
can
attack
aither
the
sacondary
carbon
position
Figure 6: An example of refining text quality with remasking. Question: What is the major outcome of the reaction between 4.4-dimeehyylcyoopen---eeed and bromine?
16

---

Published as a conference paper at ICLR 2026
Given
the
-potentiak
fr
theta
step 11
)
step 12
W
Free
<
H
2
W
Given
the
potential
YE
In
thata
step 20
frac X
kr
free
Figure 7: An example of merging adjacent tokens. Question: A quantum mechanical particle of mass m moves in two dimensions in the following potential, as function of (r,0) (r, 1 1/2kr² + 3/2kr² cos²(6) Find the energy spectrum.
calculate
step 25
Step
Calculate the cost for shorts
Mish
Ics
bought
pairs
Let
calculate
the total cost for each item
An
step 26
Calculate cost shorts
bought
pairs
Let
calculate
step 27
Step
Calculate
cost
shorts
M
Ith
ka
bought
pairs
Figure 8: An example of splitting tokens. Question: Mishka bought 3 pairs of shorts, 3 pairs of pants, and pairs of shoes. One pair of shorts costs $16.50. One pair of pants costs $22.50 and one pair of shoes costs $42, How many dollars did Mishka spend on all the clothing items?
17

---

Published as a conference paper at ICLR 2026
Step
Calculate the height after the second
step 198
bounce
The hall cunnees V00
Step
Calculate the height after second
step 199
bounce
In
Sn
The ball oercees again
to
Step
Calculate
height
step 201
bounce
again
Step
Calculate the height after the
second
step 283
bounce
In
The ball
Step
height
after
step 209
bounce
In
The
ball
again
ounces
to
Figure 9: An example of inserting tokens. Question: Nathan has bouncy ball that bounces to 2/3rds of its starting height with each bounce. If he drops it from the third-floor balcony in the mall, where each story is 24 feet high, how high does the ball go its second bounce?
18

---

Published as a conference paper at ICLR 2026
19
step 44
19
step 61
step 64
Figure 10: An example of deleting tokens. Question: When Billy was first hired, he was paid at rate $10 per hour. After months, he was given a raise of $0.50 per hour. On his first anniversary at work, he was given raise of $1.00 per hour. Sally just started working at different business. and her starting salary is $0.50 more per hour than Billy's 's starting salary was. If both Billy and Sally work 20 hours, how much more money will Billy earn than Sally, in dollars?
projection (Zang al. 2023) the connections from UPS to TPS. This ensures that the model's token prediction behavior unchanged at the beginning of training, and gradually learns to predict the confidence for unmasking tokens in diffusion process. The model comprises total of 8.9B parameters.
B.2 TRAINING CONFIGURATION AND DATASETS
B.2.1 DATASETS
We use both high quality public datasets and in-house data for training, including four major cate- mathematics, code, general conversation, and science. The mathematics category includes NuminaMath-CCT eL al. 20224), MetaMathQA YYu al. (023), -rce-math-woor-probbhe
gories:
Mitra et al. 2024), Gaths-College (ajibawa 2023). Deep (He et al. 2025). MathInstruct Yue al. 2023), and open2math- 1aah---M-pp-...-- 1-mini (mlioundations dev 2025 The code category comprises evol codelaadallp (Luo et al. 2023), opc-sft-stagel (Huang et al. (024), and KodCode- VI-SFT-4o (Xu al. 2025). General conversation data is drawn from open- open-perfecthlend (Xu al,) 2024) and Infinit-Inssin Instruct (Li et al.) 2025). The science category incorporates the science _qa (Mar-) czak 2023) dataset. In total, the public datasets provide roughly 18 8.8M samples. The in-house data, containing about 140K samples of prompt-response pairs, was primarily generated by GPT-4. and covers mathematics, code, science, and instruction followiinnnooowinngg tasks. These generated samples went through human check to ensure overall quality.
B.2.2 VARIABLE-LENGTT BLOCK-WISE GENERATION
We first adapt LLaDA into a DLM capable of variable-length generation. The underlying archi- tecture remains unchanged, but we fine-tune it for variable-length block-wise generation. During inference, generation proceeds block by block, where each block consists of L 32 tokens. For each block, the model runs a full reverse diffusion process until the block is fully denoised, after which the completed block is appended to the context. The next block is then generated in the same manner, continuing until an <eos> token appears. Similar to augo-segressive (AR) models, we en- force causality using bisck-wise causal mask in the self-artention layers: each token attend to
19

---

Published as a conference paper at ICLR 2026
W C.J. W the number of chairs Jenna Aewe ASte Step Find Ophelia's number of aaaaaaaia
step 88
Ste Find Ophelia's
step 81
Fewer2 then...
Let W C_J W be the number of chairs Jenna has, Step 1: Find Ophelia's number of chairs. *tinin step 85
They knowla has fewer fewer than than...
W J
Anve Step Ophelia's
step 87
They each have fever sofas than chairs:...
(a) Question from GSM8K: Ophelia and Jenna living in the same apartment building. They each have fewer sofas than chairs. Jenna has 3 times many chairs as Ophelia. Ophelia has 20 sofas, calculate the total number of sofas and chairs that they have.
step 389
Examlle Calculation
KnvaLet's "ariionnn
side Length-
Side
Length
ietakatetakn....
step 482
Exaxple calculation"" (nnnteen say An se Length: <*
cain"
Apothenn 1.5 cn/atn...
step Exaxple calevlationnyynvile'''' assune:\nian
Length:*
"Apothes:** anvaVa....
step 411
Exasple Galcatation assune:\nka'
siiide Length: **
an Ap then:
ogs anknha..........................
(b) Question from Alpaca-Eval: Hi, need to find the area of a 20-sided die for my math homework. Can you
help me do that?
Some individuals nay feel confortable and valued in platonic relationships, while others step feel pressure attraction in ronantic romantic.. wwwp///wwwww......../.aaaaaaaaaa and and and friendship friendshipta...
... individuals nay confortable valued platonic relationships, while may feel pressure attraction-án ronantic situations/Anya "Iporrttt considerations::*"all Connunication =* should to
step 108
and...
Sone individuals nay feel confortable and valued in platonic relationships, while others may feel pressure or attraction towards romantic relationships, Inporttnt considerations:=>nnn step *Comunication::** Open's for conmunicate comunicate comunicate communicate about about about feelings...
Some individuals nay feel confortable and valued in platonic relationships, while others may feel pressure attraction towards romantic "tartiomsiips.taaoo"""ttttrt Inporttt considerations::"Miii step **Comunication:**
Open's important for comnunicate about about feelings feelings feelings
feelings...
(c) Question from Alpaca-Evar Eval: Can boy and girl ever just be best friends?
Figure 11: Visualization of per-token confidence scores predicted by UPS. The darkness of the background color indicates the value of hig - darker means larger. The font color indicates different type of tokens: Green and red tokens already unmasked in the input sequence of the current step. Red tokens are remasked. Blue tokens are newly unmasked in the outputs. Gray tokens keep masked in the outputs, and display the token with the highest probability these positions.
20

---

Published as a conference paper at ICLR 2026
all tokens within its current block and all tokens in previously generated blocks, but to tokens in future blocks.
For supervised finetuning, the response part is divided into blocks of length L 32, with the last incomplete block padded by <eos> tokens. The training objective is to recover a noised version of each block conditioned all previous clean blocks. Following Arriola et al. (2025). concatenate both the clean blocks and their corresponding noised versions into a single input sequence, allowing all blocks to be trained jointly in forward-backward pass.
In this stage, we train variable-length block-wise generation with learning rate 2 x 10-6, batch size 160, and gradient threshold of 1.0. The baseline model in Table reports the results of the variable-Iength block-wise model on several datasets, alongside the comparisons when Remask SFT and Remask RL added this baseline model.
B.2.3 REMASK SFT
After finetuning LLaDA for variable-length block-wise generation, we attach the Unmasking Policy Stream (UPS) to construct the model architecture shown Fig. [2] We then further train RemeDi with Remask SFT with Ps 0.3. For optimization, apply learning rate of 2.0 x 10 to the newly introduced parameters in UPS, while keeping the learning rate for the original parameters
2.0 10-6
B.2.4 REMASK RL
To enable effective RL training. we curated dataset spanning mathematics, coding, STEM, instruction-folloiing and preference-aligmment tasks:
Math: GSMbe al. 2021). et al. 2025)
MATH HHedryys et al.
2021 DeepScaleR Luo
Code: KodCode-V1-SFT--e Xu et al. 2025
LeetCodeDaaaseet Xia et
2025
General: Skywork-Rewarde PrefraadrPPeffeeeee-------...-. et al.; 2024).
Llama-Nemotron
SegtTTaininn-Datasetaa (instruct-following)
et al. 2025). Nemotron-Por
Craining-Datasete (Nathawani et al. 2025)
To ensure quality, we applied the following filters:
Length: Since our RL training limits generation length to 1024 tokens, we discard any sample-questinn plus answer question exceedix this bound.
Verifiable: For math data, When both a short answer and a detailed response are available, we keep the sample only if the two answers match; for code data, we require that the provided solution passes all test cases.
Deduplication: Given the diverse sources, we perform global deduplication using Min- Sash
Reward Design. Our reward function incorporates two distinct types of reward signals:
Verifiable Reward: Verifiable rewards are widely used in mathematics, code, and STEM domains, where the answer is first extracted and then verifies it: math and STEM via Math- Verify (Kydličck[200) (225)), and code via executing test cases and computing the pass rate. We cerruction-follwwiini samples with IFEval (Zhou et 2023) format to further improve the model's 's ability to follow instructions.
also incorporated verifiable
Model-based Reward: We incorporated the Skywork- -ewe---wwwddd--eeeeee V2-LLama--311-81 (Liu et al, 2025), which was trained human preference data, to evaluate response quality. Each response is assigned a scalar reward score reflecting human preference, thereby enhancing the model's ability produce outputs that better align with human preference during RL training.
We optimized the model using the Adam W optimizer with a learning rate of 5.0 x 10 10-6 3 (0.9, 0.999) and maximum gradient norm of 1.0, for total of 100 training steps.
21

---

Published as a conference paper at ICLR 2026
Algorithm 2 Remask RL
Require:
Model parameters 0, a dataset D, and reward func.
1: while not converged and maximum epochs not reached do
2: Sample questions '>>
for g to G Initialize s%, with and mask tokens.
P Generate group G trajectories
5: for n = 1 to N do
b. N denotes the number of denoising steps
6:
Calculate the ranking score ho.n for each token
7:
Sample Kn positions to unmask in this step: Uo
Plackett-Luce K..)
8:
Sample =t, pa... (-10r
vi E Un 2
[M]
9:
end for r reward xxx)
10:
P Compute the rewards
11:
end for
12:
for g G do 1:C) stder
Compute the advantages in GRPO
A
end
for
15:
for 11 = 1 to N do
P. Compute my and losses for each denoising step
16:
To. (2),, 10
p/m-1 = samask (21 22 tken (x7, |og 1 U()
see Eq
To,n (08, 1a?
17:
Cen - DG
A9
Taan (o 1=?
Calculate the gradient VoLo,
19:
end for
20:
Update 0 with accumulated gradients E VeLo,, along the descent direction
21:
end while
B.2.5 ABLATION SETUP
We provide detailed configurations for ablation studies in Sec. 4.3
Remask SFT Both the vanilla SFT model and the Remask SFT model are trained with identical
hyper-parameters: block size 32, maximum generation length 1024, global batch size 80, and total of 3000 training steps. We use the AdamW optimizer with a learning rate of 2.0 10 10-5 in newly added parameters of UPS and of 2.0 10-4 in original TPS parameters, 3 (0.9, 0.999), and maximum gradient norm of 1.
Remask RL Both Remask RL and LLaDOU RL trained LLaDA using identical RL hyper- parameters. Specifically, roll-outs are generated with temperature 0.7, generation length 256, block length 64, and N 64 denoising steps. Each batch consists of 16 prompts, with each prompt generating G = 16 roll-outs, for total of 200 training steps. We use the Adam W optimizer with learning rate of 5.0 10 6 B (0.9, 0.999). and maximum gradient norm of 1. The complete training algorithm is elaborated in Alg.
B.3 EVALUATION DETAILS
B.3.1 INFERENCE SETTINGS
RemeDi For evaluation, RemeDi uses maximum generation length of 2048 MATH and 1024 on all other datasets. At each step, only one token is unmasked, with a block size of 32 for genera- tion. Both TPS and UPS adopt greedy sampling.
LLaDA The evaluation of LLaDA largely follows (Nie et al. 2025). On GSM8K and MATH, set the generation length to 256 with a block length of 8, unmasking token per step in semi-autoregressivee manner with greedy sampling. On HumanEval and MBPP, we use a generation length of 512 with block length of 32, while keeping all other settings unchanged.
LLaDA ReMDM
We implemented ReMDM (Wang et al. 2025a) on top of LLaDA et al. Specifically, we adopted the "ReMDM-cap Switch" configuration with Mepa 0.04 and
2025
22

---

Published as a conference paper at ICLR 2026
wwitch 0.55. For evaluation, set the generation length 256/256/512/512 and the block length to 8/8/32/32 for GSM8K, MATH, HumanEval, and MBPP, respectively.
B.3.2 BENCHMARKS
Here we provide the detailed input prompts and how the metries are computed for different bench- marks:
GSM8K GSM8K evaluates multi-step mathematical reasoning in elementary problems (Cobbe al. 2021). We illustrate below a zero-shot prompt used to evaluate the model. After generation, we extract the answer "boxed () from the response, and check it is equivalent to the ground truth with the scripts developed by Hendryykk et al. (2021). We report the accuracy on this benchmark.
Janet's ducks lay eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market? (Please put the final answer \boxed() tag, i.c. $1boxed here)
MATH MATH contains 5,000 challenging competition mathematics problems (Hendrycks et al., 2021]. We evaluate the model zero-shot setting, with prompts like the one below. After gener- ation, we extract the answer in "boxed( " from the response, and verify if it is equivalent to the ground truth with the scripts developed by Hendrycks al. 2021).
Convert the point $) in rectangular coordinates to polar coordinates. Enter your answer in the form (r, \theta), where Sr 0$ and $0 Mee \theta Api.$ (Please put the final answer in \boxed tag, i.e. Aboxed h)
GPQA GPQA is challenging multiple-choice benchmark for testing LLM's complex scientific reasoning and specialized knowledge domains (Rein et al. 2023). used all 448 questions from the main version of GPQA and evaluated the model in zero-shot setting with the prompt shown below. We select the token with the highest probability at <mom_mask> position as the final answer, and report the pass @ I on this benchmark.
<|stattoftett >>< >
>user< (nddereeen_
What is the correct answer to this question: large gene has dozens of exons, of which the central ones code for folded triple helical repeats that connect the cytoskeleton with sarcolemma and extracellular space.
Each exon usually codes for one folded triple alpha helix. The most conmon mutations of the gene are central exon deletions that create out-of-frane peptides and progressive degenerative organ waste. A solution is to deliver Morpholing ine that recognizes the 5' end of the out-of-frame exon in pre- mRNA. The molecule prevents binding of the spl ceosomee and creates exon ski skipping and in-frame joining. Several missing exons are well tolerated by an organism. Which structure belew is not involved in the proposed therapy? Choices:
(A) lariat
(B) R-loops
(C) antisense
(D) polyß tail
fnswer:Your answer should the format 'The best is @ete_aaswer__eteen where the aashe__ssve__ettee is one of (A), (B), (C) or
jnd__hederr_iaa
The best answer is <mdd__maasaaa mask>
MBPP MBPP consists of 500 python programming problems for entry level programmers (Austin et al. 2021). We evaluate the model with the prompt below in zero-shot setting. After generation, we extract the python code from the response, and check if it passes all test cases associated with this problem, and report pass @ this benchmark.
23

---

Published as a conference paper at ICLR 2026
You expert Python programmer Your task is to complete the implementation function named "rmove_Cce".
** TARGET FUNCTION **
Write a pythen function to remove first and last occurrence of a given character from the string.
** UNIT TESTS
Your code should pass unit tests like:
assert remove_Occ ("hello", "1") == "heo" assert remove.0cc ("abeda", "a") "bod"
assert remove_Occ ("PHP", "p") "H"
Here is the function to complete:
ppyhhon
def remove_0cc (input.param_l, input.param_2):
"""rite a python function to remove first and 1ast occurrence a given character from the string. "M*
HumanEval HumanEval consists of 164 hand-written programming problems (Chen et al. 2021 We evaluate the model on it with the prompt below zero-shot setting. After generation, we extract the Python code from the response, and check whether the output function passes all the provided test cases; we then report the pass@) on this benchmark.
You are expert Python programmer, Python code should be placed between the line of "python and the line of for easy extraction later, and here is your task:
... "python from typing import List
def has.close.eiementa (numbers: List(float]] threshold: float) -> bool: """Check if in given list of numbers, any two numbers are closer to each
other than the given threshoid.
Examples:
>>> has cls..loosee.... elements ([1.0, 2.0, 3.0], 0.5)
False
>>> has.close.elements ([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
True
uan
Hellaswag Hellaswag is benchmark dataset specifically designed to evaluate machine common- sense reasoning capabilities (Zellers et al. 2019). It primarily assesses a model's ability infer most plausible subsequent event based on given contextual information in natural language under- standing tasks. We evaluate the model in zero-shot setting on Hellaswag. We follow the approach of (Nie et al.) (025), incorporating Classifier-Free Guidance (CFG) and set the CFG weight to 0.5. Under CFG intervention, the model simultaneously computes conditional predictions (based on the given context) and unconditional predictions (absent specific context), guiding the generation pro- cess by scaling the difference between them. The final accuracy is calculated based on the model's normalized probability assigned to the correct option.
ARC-C ARC-C is a highly challenging benchmark dataset specifically designed to evaluate ma- chine abstract reasoning and scientific problem-solvin capabilities (Clark et al., 2018). We evaluate the model in a zero-shot setting on ARC-C, with the prompt below. After generation, we extract the answer after "The best answer is' and report the pass @ rate this benchmark.
24

---

Published as a conference paper at ICLR 2026
Given the following question and four candidate (A, B, and D), choose the best answer.
question: An astronomer observes that a planet rotates faster after metcorite impact. Which is the most likely effect of this increase in rotation?
A. Planetary density will decrease.
B. Planetary years will become lcnger.
C. Planetary days will become shorter.
D. Planetary gravity will pecome stronger.
Your response should end with "The best answer is [tne_ answer letter] where the [the_answer_letter is
A, B,
1FEval IFEval evaluates the model's instruction-followii capability with verifiable instructions (hoou al. 2023). We use the official evaluation code provided by the IFEval Benchmark, and compute the model's accuracy based on the loose metric.
AlpacaEval AlpacaEval evaluates the model's instraction-followin capability with the LLM- as-a-Judge methodology (Dubois al.] 2024). officially recommended by the AlpacaE- val benchmark, we use GPT4 1106 preview the baselinc/reference model and the weig
4_turbo as the evaluator/annotttor, and assess the win rate of the responses generated by RemeDi under length-controlled conditions to eliminate the confounding effect of response length.
C MORE EXPERIMENTS
C.I COMPARISON WITH SEED DIFFUSION
Since the official Seed Diffusion model and implementation are not publicly available, we do our best to reproduce it. For fair comparison, train seed diffusion under the same base model and identical training configuration as in Appendix [.2.5 As shown in Table 6 Remask SFT consis- tently outperforms Seed Diffusion, demonstrating the advantage of learning explicit remasking policy during training.
Table 6: Unified head-to-head comparison with other training algorithms under identical settings.
Method
GSMSK
MATH-500
HvmanEval
MBPP
Baseline
80.3
34.7
41.5
42.6
Vanilla SFT
83.1
40.1
48.2
43.4
Seed Diffusion
63.9
28.0
5.4
Remask SFT
83.6
42.7
50.0
44.0
C.2 PREDICTOR-CORRECTT VS. LEARNED REMASK POLICY
We apply representative predictor-coree corrector sampler, ReMDM (Wang et al., 2025a), to RemeDi- Instruct. The evaluation setup is the same in Appendix B.3.1 The results in Table show that our learned remasking policy (via Remask SFT) is more effective than the random remasking strategy employed in predictor-corrector samplers.
Table 7: Comparison between our learned remask policy in RemeDi and the ReMDM predictor- corrector, both evaluated with Restuti-Instruc.
Method
GSM8K
MATH-500
HumanEval
MBPP
RemeDi predictor-corrector
58.3
38.7
39.6
54.2
RemeDi (Ours)
86.3
52.2
71.3
57.8
25

---

Published as a conference paper at ICLR 2026
C.3 EFFECT OF DIFFERENT SAMPLERS
To isolate the effect of our multi-task objective and incorrect-token augmentation, we compared vanilla samper, adaptive samper(Kim al. 2025) and remask sampler under the same Remask SFT model in Table g
Table Effect of different samplers under the same Remask SFT model.
Sampler
GSM8K
MATH--50
HumanEval
MBPP
IFEval
Vanilla
86.3
38.3
43.9
55.4
69.2
A Kim et al. 2025)
86.6
40.3
43.9
55.4
69.2
Remask (Ours)
86.3
52.2
71.3
57.8
81.9
C.4 MATCHED-COMPULT ABLATION: EXTRA SFT vs. REMASK RL
To directly address whether performance gains from RL merely additional training time, we conduct a matched-compun ablation: we continue training the RemeDi-Instruct model for an 2,000 steps (which consumes approximately 32 H800-days, the same compute as used in the RL stage) and evaluate its performance. The results are summarized in Table
Table 9: Matched-compute ablation between extra SFT training and Remask RL.
GSM8K
MATH-500
HumanEval
MBPP
Demeei--nstrrnnn r~32 H800-days SFT training
86.3
52.2
71.3
57.8
83.6
52.6
62.8
57.8
~23 H800-days RL training
89.1
53.2
73.2
59.4
C.5 UPS STRUCTURE ABLATIONS
Table 10 reports ablations on UPS components, showing that removing either the bi-residual nections or the zero-init bridge leads to clear performance degradation. We train all models under the same base model and identical training configuration in Appendix B.2255
Table 10: UPS structure ablations. Removing either the bi-residual connections or the zero-init bridge degrades performance
GSM8K
MATH-500
HumanEval
MBPP
Baseline
83.6
42.7
50.0
44.0
w/o bi-residuals
76.6
43.1
45.7
42,2
w/o zero-init
75.2
42.7
45.1
44.8
C.6 EFFECTIVENESS OF RI.
To assess the effect of RL, we compare the per-token remask frequencies between the SFT-only model and the RL-trained model. Interestingly, the RL-trained model performs more remasking on average. As shown in Table TI higher remask frequency correlates with improved performance, suggesting that additional remasking provides more opportunities for RemeDi to detect and correct early-step errors.
D EXTRA VISUALIZATION
Fig. compares the por-token remask frequency the SFT-only model and the RL-trained model GSM8K. We observe that RI. training consistently increases the remask frequency, suggesting
26

---

Published as a conference paper at ICLR 2026
Table 11:
Average remask frequency (ARF) and performance across tasks. ARF measures how
many times each token is remasked average during decoding.
GSMSK
HumanEval
AlpacaEval
Model
Acc
ARF
Acc
ARF
Acc
ARF
RemeDi (+ Remask SFT)
86.3
0.16
71.3
0.070
12.5
0.012
RemeDi (++ Remask RL)
89.1
0.56
73.2
0.89
24.8
0.086
stribution of Remask Frequency on GSMBK (SFT only)
Distribution of Remask Frequency on GSMBK (SFT RL)
1.0
2.0
0.0
Rado
Rado
0.2
0.2
Rernask
Remmak
Figure 12: Comparison of per-token remask frequency between the SFT-only (left) and the RL- trained (right) model on GSM8K. The RL-trained model performs remasking more frequently on average, indicating that RL encourages more remasking behavior.
that the RL objective explicitly encourages the model to revise uncertain tokens more often, which correlates with the observed improvement final answer quality.
Fig. 13
reports the remask ratios different timesteps when inferencing GSM8K. The ratio first increases and then steadily decreases as the denoising process converges. Since the number of unmasked tokens Kn, is explicitly controlled each step, the process naturally terminates without spikes of remasking in late stage.
Fig. 14 shows the prrrghput-eeffrrmncce trade-off of RemeDi compared to AR and DLM base- lines. By adjusting the number of tokens denoised per step, RemeDi achieves acceleration in decod- ing speed with only small drops in GSMSK accuracy. forms better Pareto front than LLaDA, Dream, and other auto-regressive models.
shows the reward curve for Remask- RL and LLaDOU-RL on GSMSK. Due the larger action space introduced by the remask operation, Remask-RL starts with lower initial reward. However, it quickly surpasses LLaDOU-RL within the early stages of training and maintains consistently higher reward thereafter. This trend is consistent with the accuracy comparison in Table
Fig. 15
5 where Remask-Ri RI. demonstrates both faster convergence and a higher final performance.
THE USE OF LARGE LANGUAGE MODELS
We used large language models (LLMMs) in two limited ways: (1) for minor English polishing of the paper text, and (2) required by the AlpacaEval benchmark, where LLMs are invoked for automatic evaluation.
27

---

Published as a conference paper at ICLR 2026
Remask Ratio
0.030
0.025
0.020
Ramask Rato
0.015
0.010
0.005
0.000
10 12
13
15
16
17
18
19
20
21
22
23
25
26
27
28
29
30
31
(within
block)
Figure 13: Average remask ratio across diffusion timesteps when inferencing GSM8K. The re- mask ratio rises during early steps and then gradually declines the process converges. Because the number of unmasked tokens Kn, is explicitly defined at every diffusion step, the procedure ensures stable remask termination and avoids late-stage re-emergence of remasking.
Performance Throughput
100
RemeDi11 tok/step)
RemeDit2 tok/step)
GSMBK Accuracy (%)
DeepSeekMath
RemeDi(4 takistep)
Dream[1 tok/step)
LLaDA(1 tokistepp)
Gemma
LaDAI2 tok/step)
RemeDi
LLaDA
Dream
DeepSeekMeth
MetaMath
MetaMath
LLaDA(4 tok/step)
Gemma
Poreto Front
60
25
50
75
100
125
150
175
Throughput (tokens/sec)
Figure 14: Throughput--eeffrmmnce trade-off of RemeDi compared with other AR and DLM mod- els. By increasing the number of denoised tokens per step, RemeDi provides a smooth qual- ity-latency trade-off. All results measured with batch size I and sequence length 1024 single H800 GPU.
28

---

Published as a conference paper at ICLR 2026
GSMBK (Remasi-RL LLADCU-RL)
(L)) Herrank-itL, (1.M)
(uan)
11 (FNA)
Step
Figure 15: GSM8K training reward curves comparing Remask-RL. and LLaDOU-RL. Solid lines
show EMA-smoothed rewards, and faint lines denote raw step-wise values. Remask-RL exhibits faster carly-stage improvement and ultimately reaches higher reward, consistent with the perfor- mance gains summarized in Table
29


Reviews:
[{"review_id": "L2Q2DFJkqS", "reviewer": "Reviewer_4oJS", "summary": "The paper diagnoses a structural limitation of mask-based diffusion language models (DLMs)—namely that once a token is unmasked it becomes effectively immutable, which prevents the model from performing self-reflection and post-hoc correction of its outputs. To mitigate this, the authors propose a self-reflective remasking mechanism (RemeDi) that augments the diffusion process with learnable per-token confidence scores and an explicit remasking operation that can reintroduce previously unmasked tokens for further refinement. They train the mechanism with a two-stage procedure—supervised Remask SFT followed by policy-style Remask RL—so the model learns when and how to re-mask and revise. Crucially, RemeDi preserves the diffusion model’s noise-monotonicity (the monotone decrease of corruption during denoising) while enabling iterative edit-and-refine behavior, and the paper reports substantial empirical improvements on mathematical reasoning, code generation, and general language tasks.", "questions": "1. I suggest moving some experimental results to the appendix and placing the main architecture figure in the main text.\n2. Why did you run experiments with LLaDA rather than LLaDA 1.5?\n3. Is inference latency substantially slower compared to the baseline?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. RemeDi enables self-reflective iterative refinement by jointly predicting per-token confidence scores and remasking low-confidence tokens so that previously revealed tokens can be selectively re-sampled without breaking the diffusion noise schedule.\n\n2. The paper introduces a principled two-stage training pipeline, with supervised Remask SFT to teach detection and remasking and outcome-based Remask RL to optimize whole-generation trajectories, which yields stronger performance.\n\n3. The remasking mechanism supports a wide range of edit behaviors such as replacement, insertion, deletion, merging and splitting, and produces empirical gains on mathematics, code generation, and general instruction benchmarks compared to prior open-source diffusion language models.", "weaknesses": "1. No reproduction code is provided; please supply an anonymous repository link or include the code in the supplementary materials.\n2. The model diagram in the appendix is unclear — please clarify how $p$ and $h$ are predicted simultaneously and whether any new modules were introduced.\n3. Please add a baseline that performs RL directly on LLaDA using the same datasets you used.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N9KtOzedfD", "reviewer": "Reviewer_arzr", "summary": "RemeDi augments diffusion LMs with a learnable remasking policy: a per-token confidence head (UPS) decides which tokens to unmask or re-mask, trained via remask-SFT then trajectory-level RL, yielding consistent gains on math/code/general tasks under block-wise variable-length decoding.", "questions": "1. Report UPS overhead (params/FLOPs/memory) and equal-quality latency/tokens-per-second vs. AR/strong DLM baselines, with quality–latency Pareto curves on fixed hardware.\n2. Provide ablations for UPS attachments and components (bi-residuals, zero-init bridge, layer choice) plus sensitivity sweeps of (\\rho_{t,\\text{incorrect}}) and ratio (r).\n3. Decompose reward contributions, test cross-domain transfer without retuning, and include robustness/calibration (e.g., ECE) under distribution shift.\n4. Add unified head-to-head comparisons with edit-flow, seed-diffusion, and predictor-corrector using identical checkpoints, datasets, and decoding budgets.\n5. Quantify remask oscillation and compare termination/annealing (fixed/decayed thresholds, hysteresis) to show stability–quality trade-offs.\n6. Summarize engineering cost (LoC, stages, time, hardware) and release a minimal plug-and-play UPS recipe with an inference-only lightweight variant.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Turns remasking from an inference hack into a learned policy; per-token confidence (UPS) guides when to unmask or re-mask.\n2. Two-stage training (remask-SFT + RL) aligns token-level corrections with sequence-level rewards; compatible with monotonic denoising.\n3. Works atop block-wise variable-length decoding (LLaDA-style); integrates without changing the base noise family.\n4. Confidence head provides interpretable signals and a knob to trade exploration vs. commitment.\n5. Orthogonal to many decoding/acceleration tricks; likely to stack with other DLM advances.", "weaknesses": "1. No clear accounting of UPS params/FLOPs/memory or train/infer throughput; lacks equal-quality step/latency comparisons to AR/other DLMs.\n2. Missing UPS structure/attachment ablations (bi-residuals, zero-init bridge, layer choice) and sensitivity to ρ_(t,\"incorrect\" )  and ratio r.\n3. RL uses preference and verifiable rewards; robustness and transfer to new domains remain uncertain.\n4. Sparse head-to-head with edit-flow/seed-diffusion/predictor-corrector under identical settings.\n5. Potential remask oscillation (low-confidence flip-flops); termination/annealing not systematically reported.\n6. Higher engineering complexity (dual streams + two-stage training) raises adoption barriers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZElyofPZEl", "reviewer": "Reviewer_Z76m", "summary": "The paper investigates remasking as a way to improve masked diffusion language models. Remasking works by replacing some already-unmasked tokens with mask tokens, which subsequently get unmasked again. This way, uncertain or incorrect tokens can be revised and potentially fixed, thereby improving accuracy with minimal additional cost.\n\nThe proposed method, called RemeDi, introduces two new policies to identify the tokens to be remasked, one based on supervised fine-tuning (SFT) and one based on reinforcement learning (RL). This is done via a dual-stream transformer architecture: The token prediction stream (TPS) predicts unmasked tokens while the unmasking policy stream (UPS) predicts per-token confidence scores (for both masked and unmasked tokens). The resulting model achieves significant gains on a variety of benchmarks compared to baselines.", "questions": "- Q1. What is the effect of the multi-task objective (Eq. 5) together with incorrect token augmentation (Eq. 3)? More specifically, what is the performance of Remask SFT when using standard vanilla/adaptive sampling (Kim et al., 2025) compared to full Remask SFT? (also see W1)\n- Q2. What is the difference between Remask RL and LLaDOU RL (Huang et al., 2025)? Specifically, is the sampling policy in Eq. 8 not identical to the one from Eq. 9 in LLaDOU? If so, where do the improvements of Remask RL over LLaDOU come from? (also see W1)\n- Q3. Will the model weights and/or training code be open-sourced? (also see W2)\n- Q4. How does the proposed dual-stream architecture compare to more naive baselines, e.g. a single unmasking policy head or a single additional Transformer block (as employed by LLaDOU)? (also see W3)\n- Q5. What is the performance overhead associated with RemeDi compared to baselines? (also see W3)\n\nNits (not considered for final score):\n- L31: Austin et al. (2022) should be cited for discrete diffusion.\n- L34: Sahoo et al. (2024) and Shi et al. (2024) should be cited for masked diffusion.\n- L103: typo: missing space after “Recent studies”\n- L160: typo: “paradigm offers” -> “paradigms offer”\n- L203: typo: extra space after “(Guo et al., 2025)”\n\n---\n\n### References\n- Austin et al. (2022): https://arxiv.org/abs/2107.03006\n- Huang et al. (2025): https://arxiv.org/abs/2505.10446\n- Kim et al. (2025): https://arxiv.org/abs/2502.06768\n- Sahoo et al. (2024): https://arxiv.org/abs/2406.07524\n- Shi et al. (2024): https://arxiv.org/abs/2406.04329", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "### S1. Strong benchmark performance\nThe proposed model achieves strong benchmark performance, beating many diffusion and autoregressive baselines of similar size.\n\n### S2. Flexible post-training method\nThe proposed method is flexible enough to be generally applicable to many existing masked diffusion models, and can be combined with other optimizations such as block diffusion.\n\n### S3. Analysis of remasking behavior\nThe analysis of the remasking behavior of the trained policies provides valuable insights. It sheds light into task difficulty as well as which part of the sampling process may be most prone to errors.", "weaknesses": "(ordered by decreasing severity)\n\n### W1. Source of performance improvements is unclear\nIt is somewhat unclear which improvements stem from the continued (multi-task) pretraining and which come from the remasking during sampling. This question is partially addressed by comparing Remask SFT with vanilla SFT and Remask RL with LLaDOU RL, but some confounders remain. The multitask objective employed by Remask SFT may be beneficial even in a vanilla (or adaptive) inference setting where the model confidence is obtained heuristically. Similarly, the Remask RL approach is very similar to the one proposed by LLaDOU (Huang et al., 2025), so it’s unclear where exactly the improvements come from.\nThe former can easily be addressed by reporting the performance of Remask SFT using vanilla and adaptive masked diffusion sampling (Kim et al., 2025). The latter can be addressed by providing a detailed analysis of the differences between Remask RL and LLaDOU RL and, if necessary, ablating atomic changes to measure their individual impact (e.g. training vs. sampling improvements).\n\n### W2. Reproducibility\nAs far as I can tell, there will be no model weights or codebase accompanying the paper, which is a major concern for reproducibility, especially given the clarity concerns regarding the source of the performance gains (W1).\n\n### W3. Motivation and effect of dual-stream architecture\nThe motivation for using a dual-stream architecture over conceptually simpler approaches (e.g. a dedicated unmasking policy head) is unclear, and no ablations are performed on this. Similarly, the computational overhead of the proposed architecture (more parameters, slower forward pass) may provide an unfair advantage compared to baselines. The former can be addressed by providing appropriate ablations, whereas the latter calls for reporting inference speed of the proposed method compared to baselines in addition to benchmark accuracies.\n\n### Conclusion\nAs presented, reasons to reject the paper outweigh reasons to accept: While the benchmark performance of the proposed model is impressive and beats state-of-the-art methods (S1), the lack of clarity regarding the source of those gains (W1) together with concerns about reproducibility (W2) and inference speed (W3) amount to significant concerns regarding soundness. I will be happy to increase my final score if these weaknesses can be addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jvJ5g39hks", "reviewer": "Reviewer_AeNa", "summary": "The paper proposes **RemeDi (Remasking-enabled Diffusion Language Model)**, which adds a self-reflective remasking mechanism to diffusion language models (DLMs). Traditional DLMs fix tokens once unmasked, making them unable to revise early mistakes. RemeDi introduces a dual-stream transformer: a **Token Prediction Stream (TPS)** predicts token distributions, while an **Unmasking Policy Stream (UPS)** predicts per-token confidence. Low-confidence tokens are re-masked and regenerated in later steps, allowing iterative self-correction. Training proceeds in two stages: **Remask SFT** (supervised fine-tuning) teaches the model to identify and remask incorrect tokens using mixed masking and random replacement; **Remask RL** fine-tunes full trajectories using task-specific rewards. Experiments on code, math, and reasoning benchmarks show that RemeDi surpasses prior DLMs (Dream, LLaDA, LLaDOU) and even matches or exceeds autoregressive baselines of similar size.\n\nHowever, despite the promising empirical results, the paper exhibits several methodological and reporting weaknesses. Many experimental details are under-specified: the reward design for Remask RL is not clearly defined, and critical hyperparameters or normalization schemes are absent, even in the appendix. The paper does not report any compute or training-time comparisons, leaving unclear whether the observed gains stem from algorithmic innovation or simply greater computational cost. Similarly, convergence curves and efficiency analyses are missing, despite claims of faster training. The method section, while conceptually sound, omits low-level implementation details—such as how the UPS interacts with TPS during training or how remask thresholds are selected—making reproduction difficult. While the writing is clear at a high level, the technical exposition lacks sufficient depth and cross-references to equations and figures, resulting in incomplete methodological transparency.", "questions": "1. **Comparison scope:** Why were models like *ReMDM* and *Seed Diffusion* not included in quantitative comparisons? Could a smaller-scale reimplementation or proxy be feasible?\n2. **Reward design:** Please provide explicit formulas or examples of the reward functions used in the RL phase and clarify normalization or scaling strategies.\n3. **Effectiveness of RL:** How does the RL-trained model differ behaviorally from the SFT-only version (e.g., fewer unnecessary remasks, faster correction)?\n4. **Visualization and convergence:** Could you move the dynamic remasking visualization (currently in appendix) to the main paper and add training loss curves to verify efficiency claims?\n5. **Compute accounting:** Please provide compute cost (GPU hours, wall time, token count) for both stages. How does SFT-only compare to SFT→RL at equal compute?\n6. **Necessity of two-stage design:** Did you test a single-stage joint objective or stronger SFT baseline under matched compute?\n7. **Stability:** What regularization methods (e.g., KL penalty, reward whitening, gradient clipping) were applied to stabilize RL?\n8. **Matched compute ablation:** Could you run an SFT-only baseline using the same compute as SFT→RL to confirm that RL brings genuine benefit?\n9. **Generalization to other DLMs:** Since the model is based on LLaDA-8B, could this remasking mechanism generalize to other diffusion LMs such as Dream? If so, what architectural adjustments are needed?\n10. **Random alternative tokens:** In SFT, are random alternatives drawn uniformly from the vocabulary or sampled from the model’s top-k predictions?\n11. **Inference efficiency:** The paper states RemeDi achieves faster convergence; could you include **inference latency comparisons** (e.g., tokens per second vs. LLaDA), given that TPS + UPS may add computational overhead?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. **Addresses a concrete limitation of DLMs** — RemeDi directly tackles the inability to revise early decoding errors through a principled, learnable remasking mechanism.\n2. **Clear and interpretable design** — The dual-stream architecture (TPS + UPS) separates token prediction from confidence estimation, enabling transparent remasking decisions.\n3. **Comprehensive experiments** — Evaluations cover math (GSM8K, MATH), code (HumanEval, MBPP), and reasoning tasks (ARC-C, AlpacaEval), consistently improving over Dream and LLaDA.\n4. **Empirical evidence for self-correction** — Qualitative examples and remask frequency analysis show that RemeDi learns to remask more frequently on harder or more structured tasks.\n5. **Training stability and convergence** — Compared to LLaDOU RL, Remask RL converges faster and achieves higher accuracy under identical settings.", "weaknesses": "1. **Limited quantitative comparison to other edit-based diffusion models** (e.g., ReMDM, Seed Diffusion). The paper cites prior edit-based diffusion models (e.g., ReMDM, Seed Diffusion) but omits quantitative comparisons, leaving unclear how RemeDi’s improvements scale relative to these baselines.\n2. **Under-specified reward formulation in Remask RL** — lacks details on normalization, weighting, or how rewards interact with confidence-based remasking.\n3. **Compute and efficiency concerns** — The two-stage SFT → RL training likely increases overall cost, but compute parity with baselines is unreported, making it unclear whether improvements stem from algorithmic design or added compute.\n4. **Figure placement and clarity** — Key visualizations (like the dynamic remasking behavior) appear in the appendix instead of the main paper.\n5. **Missing convergence curves** — The claim that Remask RL converges as efficiently as other DLMs lacks quantitative support.\n6. **Necessity of two-stage training unclear** — It is not shown whether a stronger SFT baseline or a joint training objective could achieve similar effects.\n7. **Stability and variance during RL** — The paper lacks reporting on reward variance, remask rate dynamics, or regularization methods.\n8. **Ablation under matched compute not provided** — It is unclear whether RL brings improvement beyond additional training time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)