Partition Generative Modeling: Masked Modeling Without Masks
Authors: Justin Deschenaux
Lan Tran
Caglar Gulcehre
Venue: ICLR 2026 Oral (2026)
Keywords: masked generative modeling
discrete diffusion
masked diffusion language modeling
diffusion language modeling
Primary area: generative models
TL;DR: We show that it is possible to train masked generative models without using MASK tokens, resulting in efficiency gains at inference.

Abstract: Masked generative models (MGMs) can generate tokens in parallel and in any order, unlike autoregressive models (ARMs), which decode one token at a time, left-to-right. However, MGMs process the full-length sequence at every sampling step, including \mask tokens that carry no information. In contrast, ARMs process only the previously generated tokens.
We introduce ``Partition Generative Models'' (PGMs), which replace masking with partitioning. Tokens are split into two groups that cannot attend to each other, and the model learns to predict each group conditioned on the other, eliminating mask tokens entirely. Because the groups do not interact, PGMs can process only the clean tokens during sampling, like ARMs, while retaining parallel, any-order generation, like MGMs.
On OpenWebText, PGMs achieve $5-5.5\times$ higher throughput than MDLM while producing samples with lower Generative Perplexity. On ImageNet, PGMs reach comparable FID to MaskGIT with a $7.5\times$ throughput improvement. With twice as many steps, the FID improves to 4.56 while remaining $3.9\times$ faster than MGMs. Finally, PGMs remain compatible with existing MGM samplers and distillation methods.

Published as a conference paper at ICLR 2026
PARTITION GENERATIVE MODELING: MASKED MODELING WITHOUT MASKS
Justin Deschenaux
Lan Tran
Caglar Gulcehre
EPFL, Lausanne, Switzerland
ABSTRACT
Masked generative models (MGMs) generate tokens in parallel and in any order, unlike autoregressive models (ARMs), which decode one token at a time, left-to-right. However, MGMs process the full-length sequence every sampling step, including [MASK] tokens that carry no information. In contrast, ARMs pro- cess only the previously generated tokens. We introduce "Partition Generative Models" (PGMs), which replace masking with partitioning, Tokens are split into two groups that cannot attend each other, and the model learns to predict each group conditioned on the other, eliminating [MASK] tokens entirely. Because the groups do not interact, PGMs can process only the clean tokens during sampling, like ARMs, while retaining parallel, any-order generation, like MGMs. On Open- WebText, PGMs achieve 5-5.5x higher throughput than MDLM while producing samples with lower Generative Perplexity. On ImageNet, PGMs reach comparable FID to MaskGIT with 7.5x throughput improvement. With twice as many steps, the FID improves to 4.56 while remaining 3.9x faster than MGMs. Finally, PGMs remain compatible with existing MGM samplers and distillation methods.
$.69
7.84
Generative Perplexity (FP64)
X GPT22 Small (KV cache)
MDLM (dim-768)) 16/6, dim-1024)
180
Throughput (mmagesss)
FID (lower better)
PGM 18/8, dim
160
140
120
100
Halton (FID
Throughout
10a
10*
Throughput (tokisec)
32
32
64
Figure 1: (Left): On ImageNet, using the Halton sampler, PGM (ours), reaches similar FID as MaskGIT with 7.5x speedup. By sampling with twice as many steps, PGM reaches an FID of 4.56 while remaining 3.9x faster. (Right): On OpenWebText, PGM achieves a better Generative Perplexity with 5.3x higher sampling throughput compared to MDLM (an MGM for text), at context length of 1024.
INTRODUCTION
Masked generative models (MGMs) offer two key advantages over autoregressive models (ARMs): they can generate tokens in parallel and in any order, rather than one-by-one, left-to-right. These properties have led to strong results across images (Chang et al., 2022), video (Yu et al., 2023; Villegas et al., 2022), audio (Comunità et al., 2024), and language (Austin et al., 2023; Lou et al., 2024; Sahoo et al., 2024; Shi et al., 2025; Campbell et al., 2024; Gat et al., 2024). However, MGMs
"Correspondence just www.hhenaxxeppf1.coom

---

Published as a conference paper at ICLR 2026
slow at inference. Indeed, every sampling step, they process the full-length sequence, including [MASK] tokens that carry no information, whereas ARMs process only the previously generated tokens. This limits the practicality of MGMs in large-scale and real-time settings, and is a crucial disadvantage for test-time compute scaling (Snell et al., 2024; Wu et al., 2024) compared to ARMs.
Addressing the inference inefficiency of MGMs is trivial because training and sampling must be consistent. MGMs are trained with bidirectional architectures over the full sequence, so every hidden representation depends on all L positions, including masked ones. Furthermore, naively decoding tokens block-by-block means feeding the model shorter sequences at inference, which differs from training and leads to poor sample quality (Deschenaux & Gulcehre, 2024).
Prior work addresses the slow inference of MGMs from different angles. Decoding more tokens per step increases throughput, but degrades sample quality. Distillation (Deschenaux & Gulcehre, 2025; Zhu et al., 2025; Sahoo et al., 2025a) reduces the number of sampling steps, but each step remains equally expensive, and distillation can affect the sample diversity (Gandikota Bau, 2025). Block Diffusion (Arriola et al., 2025) enables partial K KV-caching by generating tokens block-by-block, but sacrifices the any-order generation capability. None of these approaches make individual sampling steps cheaper while preserving the full flexibility of MGMs.
We introduce Purtition Generative Models (PGMs), which replace masking with partitioning- Tokens split into two disjoint groups, and group- wise attention mechanism ensures that no information flows between them. The model learns to predict each group conditioned on the other, eliminating [MASK] tokens entirely. Because the two groups do not interact, PGMs process only the clean tokens during sampling. just like ARMs, while retaining the ability to generate tokens in parallel and in any order, like MGMs. We propose the Partition Transformer, dedicated architecture that prevents information flow between groups
Contributions (I) We introduce PGMs and the Partition Transformer, new architecture that enables MGM-style parallel, any-order generation without [MASK] tokens. PGMs are compatible with existing MGM samplers (Besnier al., 2025) and distillation methods (Deschenaux Gulcehre, 2025), making them drop-in replacement. (2) On Open WebText (Gokaslan & Cohen, 2019), PGMs generate samples with lower Generative Perplexity than MDLM (Sahoo et al., 2024) and reach similar downstream task performance, before and after distillation, while achieving 5 5-5.5x higher throughput. On ImageNet, PGMs reach comparable FID to MaskGIT (Chang et al., 2022) with a 7.5x throughput improvement.
(3) We show that PGM are trained with denser supervision than MGM. Since each group predicts the other, a single sequence yields two complementary training signals. This reduces gradient variance and yields 1.95 reduction in validation perplexity on LMIB (Chelba et al., 2014) compared to MDLM with the same number of layers.
2 BACKGROUND
2.1 SEQUENCE MODELING
We consider the task of generating sequences .II of length L over vocabulary
V = (O...., N - 1}. The training dataset 'D contains finitely many sequences drawn from an unknown data distribution Piaa over L Autoregressive models (ARMs) factorize the distribution as (x) I (x xx.). where x denotes the prefix before position i. Tokens are sampled sequentially, and because each conditional only depends on the prefix X ARMs process only the previously generated tokens instead of the whole sequence.
2.2 MASKED GENERATIVE MODELS
MGMs augment the vocabulary with a special [MASK] token, absent from the training data. Given
x E D, let z, denote a corrupted sequence where the token z at position is [MASK] with probability
Pt. or the clean value x otherwise. The masking probability Pt an increasing function of E (o, 1 with po 0 and pi = 1. MGMs train denoiser is modeled as factorized marginals: Ile=f= [MASK] (. z.). Because MGMs sample independently
VL - RLXN whose sampling distribution
from each p- 2t), they cannot model arbitrary joint distributions, unlike ARMs. However, they
2

---

Published as a conference paper at ICLR 2026
MGM
PGM (ours)
Thaining and Sampling
Trashiing
Samptnng
Target:
0
0
Learn fom masked posiiions only
Supervise posttions
POM
POM
Trendtomer Fronder
GroupSwop
GroupSwop
Srop positions
PIMM Errodder
PGM Erender
No masked tokons to process
Whooed Sops on meshod tahurt
Clean
Group
Group
Figure 2: Masked Generative Modeling (MGM) vs. Partition Generative Modeling (PGM), Left: MGMs learn at masked positions only, and must process [MASK] tokens at inference. Right: PGMs partition tokens in two groups, learn from all positions and process clean tokens only during sampling. PGMs achieve >5.3x higher throughput on Open (context length 1024, 128 sampling steps).
can decode tokens in parallel. The training objective for MGMs is:
LMGM == Ex-D,--010011 (m(())C(xx(Z); x)]
where w : [0. > Rg is a weighting function and CE( x) denotes the cross-entropy loss over masked positions. To generate samples, MGMs start from fully masked sequence and iteratively unmask subsets of positions over multiple evaluations of Xs, selecting positions at random or based on confidence scores. We now describe two instantiations used in this work.
MDLM Masked Diffusion Language Models (MDLM; Sahoo et al. (2024); Ou et al. (2025); Shi al. (2025)) are MGMs for language modeling. Analogously to continuous diffusion (Sehl-Dickstein et al., 2015; Song & Ermon, 2020; Ho et al., 2020; Kingma et al., 2023), MDLM defines forward process that corrupts clean data and a generative process that recovers samples from noise. The forward process is:
(-(xx) == Cat(:: 0xx + (1 - aa)
(2)
where x the one-hot representation, π m is the one-hot encoding of [MASK], and a is strictly decreasing noise schedule with an = 1. 01 = O. (2) is applied independently at every position. The posterior distribution is:
Cat(:: zz):
2: 7 m,
Paje( Ze,
Cat (1-a, (1-aa))+(aa-aa)))) (1-ae)
(3)
2: = m.
To generate samples, fix decreasing sequence of times 1 TT all [MASK] tokens sequence, and iteratively sample from
To 0, set Z,z to the
2T. 10. 2,,) Pr. (.22, Xo (x7, (z)),,,,,))) Tx))
(4)
MDLM optimizes a variational bound on log- likeliilioood that reduces to (1) with w(t) masked positions contribute the loss.
af T-as"
Only
MaskGIT MaskGIT (Chang et al., 2022) is an MGM that operates in the latent space of a pre- trained VQGAN (Esser al., 2021) tokenizer. MaskGIT proposes tokens x at masked position and uses the predicted likelihood of the sampled token X as confidence score: ef = (2;;;; t) A predefined schedule determines the number of positions to unmask, and the most confident positions are kept. Tokens generated in earlier steps are kept unchanged. This differs from MDLM, which denoises at random positions. Besnier et al. (2025) observed that confidence- based sampling tends to decode spatially clustered tokens, since the denoiser is most confident near previously generated positions. Because MGMs sample independently from a product of marginals Ilees Po(x' zr) rather than the joint token distribution, decoding nearby tokens increases the risk of generating
3

---

Published as a conference paper at ICLR 2026
inconsistent samples. By sampling according to low-discrepancy sequence (Halton, 1964), Besnier et al. (2025) enforce more uniform coverage of the space, which improves the FID and IS compared to confidence sampling.
Classifier-Free Guidance Let px c) denote a class-conditions distribution learned by an MGM, where {O,...; 13 is a class label and (s ø) denotes the baas-unncodiiional distribution. Let w 2 0 control the guidance strength. Classifier-Free Guidance (CFG; Ho & Salimans (2022);
Chang et al. (2022)) steers generation toward class c by replacing log (x c) during sampling with
log (x
c) = (1 + w) log c) w log Pe X Ø),
(5)
Self-Distillation Through Time Self-Distillation Through Time (SDTT; Deschenaux & Gulcehre (2025)) accelerates sampling from MGMs by distilling teacher trained for denoising with many steps into few-steps student. Let P9 (m) denote the distribution samples generated with m steps
using denoiser Xa, and let pe denote the distribution when using < m steps with student denoiser Xy. SDTT trains X, with the following objective:
min E a D.z. (2. [6(x, (Zc. t) X Zz, t, m/k)]]
gteacher
(6)
where 8 is a divergence measure (e.g., KLD) and x je (Ze, t, m//k) are the distillation targets. These targets are constructed using n/4k sampling steps with the teacher, starting from Z, and collecting the predicted proobabbbillties for each token at the step where a token was denoised. After training. one step of the student should match m// teacher steps. SDTT can be applied iteratively by reusing the student as teacher in each round (Salimans & Ho, 2022), progressively halving the number of required steps. Empirically, distilling 2 steps per round is most effective.
3 PARTITION GENERATIVE MODELING
At each sampling step, MGMs process the entire sequence, including many [MASK] tokens that will not be decoded yet. In contrast, ARMs process clean tokens only, but generate token at time. Partition Generative Models (PGMs) combine the strengths of both approaches, by generating multiple tokens parallel, like MGMs, while processing only the clean tokens. like ARMs. PGMs are a direct extension of the MGM paradigm. As a result, sampling algorithms, guidance mechanisms, and distillation methods developed for MGMs apply directly to PGMs. Only the neural network architecture must be adapted (Sec. 4).
3.1 TRAINING
From Masking to Partitioning Instead of replacing tokens with [MASK] PGMs partition the sequence into two complementary groups. Given D and t ~ 24(0, 1], each token is assigned to group with probability pt = 1 at and to group otherwise. Let g E 10, denote the group membership vector. We propose a Transformer variant in Sec. 4 that ensures that information cannot flow between groups. Predictions at positions in group 0 depend only on tokens in group 1, and vice-versa (Figure 2). This is consistent with MGMs, where masked tokens predicted from clean ones, except that PGMs Iearn from both groups
Connection the MDLM Variational Bound In MDLM. the forward process (2) masks each position independently, so at time t an expected fraction of tokens remain clean. PGMs assign an expected fraction ag of tokens to group 0, which plays the same role as the clean tokens in MDLM.
of By treating group 0 as clean and group as masked, the MDLM loss weight w(t) I-ar is applied to tokens in group 1. By symmetry, tokens in group 0 are weighted by w(1 - t). Therefore, in single forward pass, PGMs evaluate the MDLM training objective at two complementary masking rates. Hence, the training objective is
LPGM := E EX
PGM PGM (]CE(xx(x;85) x)]]
where
w(()) if g: ww( - t)
PGM
(g.
t);
(8)
if gg 1.
4

---

Published as a conference paper at ICLR 2026
crost crost-ofter
(m loyers)
Enc. outous
Qurceeint (ilfoiidaa
Swmpped group
Cross-a Cross-attertion
B
GroupSwap loyer
B
Erc ouso
Encoder (k toyers)
Group wioo ref-cctention
Figure 3: Partition Transformer. RoPE applied before every attention layer (omitted for clarity). (C) Encoder: group-wise self-atiention (no cross-group flow). (B) GroupSwap: cross-atiention that routes each position's 's representation to the opposite group. (A) Decoder: group-wise cross-attention the encoder output (no self-attention).
Variance Reduction Unlike MGMs, which compute the loss over masked positions only (Figure 2), PGMs compute the loss at every position, yielding two gradient contributions per training sample. By training on two complementary copies, PGMs reduce the variance. Empirically, training diffusion models with lower variance improves the validation likelihood (Kingma et al., 2023; Sahoo et al., 2024). We study the variance reduction in Sec. 5.3.
3.2 SAMPLING
During inference, PGMs process clean tokens only, like ARMs, yet decode tokens in parallel at arbitrary positions, like MGMs (Figure 2). Let C, S (1..... L}} denote the clean token indices step 7 E {I...., T),, with n, = C, and m, = L - n2- At each step, we select k, masked positions, sample from pa Xc, ). and add the decoded tokens Cr+1- For text, we find that using fixed schedule with k, = & tokens per step (Algo. 2) improves sample quality and throughput compared the MDLM posterior that decodes each position with probability T-a;
0,-a and requires padding for
batched generation (Algo. 3; Suppl. E.2). For images, experiment with both the confidence and Halton samplers (Suppl. B, Besnier et al. (2025)). The Halton sampler performs better empirically, so we report confidence-based sampler results in Suppl. D.3.
THE PARTITION TRANSFORMER
PGMs require careful architectural design. In particular, since goal is to process single group only during inference, tokens across groups should not attend to each other. As shown in Figure and Figure 3, build the Partition Transformer such that the predictions for tokens in group 0 based on tokens in group 1 only. The Partition Transformer implements a mechanism to swap the physical location of information across groups. During training, this allows using the input sequence x as target. During sampling, it moves information about the input tokens from the clean positions to the positions to predict. Our architecture consists of an encoder. GroupSwap layer, and decoder, which we describe below.
The encoder is made of partition sesitiioo--iieeeee self-attention blocks, which are similar to standard bidirectional transformer blocks except that tokens in separate groups do not attend each other.
Encoder
5

---

Published as a conference paper at ICLR 2026
Table Validation perplexity, sampling latency, and throughput (TP) on LMIB and OpenWeTText. PGM k m uses & encoder and m decoder layers. The best PGM per dataset is highlighted. Latency and TP are measured batch size 32. Trained with 2x larger batch size (Sec. 5.3). See Table for architecture ablations.
Model
#Params
Val. PPL +
Latency (sec) +
TP (tok/sec) 1
LMIB (ctx len. 128)
MDLM
170M
27.67
3.78
1'081.57
M (Compl. masking)
170M
25.72
3.78
1'081.57
PGM 6/6
171M
26.80
2.12
1'930.93
WebTet (ctx len. 1024)
MDLM
170M
23.07
31.41
1'043.22
MDLM (Compl. masking)
170M
22.98
31.41
1'043.22
PGM 8/8
203M
22.61
5.86
5'585.57
PGM 6 / 6 (dim. 1024)
268M
21.43
5.93
5'518.09
Decoder The decoder uses cross-attention layers, whose keys and values are computed based on the output of the encoder. In contrast, the queries computed using either the output of the GroupSwap layer (for the first block of the decoder) or the output of the previous decoder block (see Sec. 4. I). Importantly, there self-attention layer in the decoder, which allows efficient generation, can compute predictions solely at the positions that we will decode.
4.1 THE GROUPSWAF LAYER
In the encoder, information remains localized. If token belongs to group 0, its hidden representation only depends tokens in group O. For prediction, however, require the opposite: representations at positions in group must depend exclusively on group 1, and vice versa. To enforce this, we introduce the GroupSwap layer (Figure 3B), which exchanges information between groups. The GroupSwap layer is implemented using cross-attention, and to prevent information leakage, the queries used in cross-attention cannot depend tokens in the other group. We describe ways of initializing queries.
Data-Indepen Indepeneetnn Queries Let u IR a learnable vector. initialize the queries, we replicate across the sequence length, add fixed positional encodings, and apply layer normalization followed by a linear projection. The query matrix
lxll (where V,, is the i-th row) satisfies
Vg. = W [LN (u + pos... +b]
(9)
where W R b R are learnable parameters and LN denotes layer normalization (Ba et al., 2016). We use sinusoidal positional encoding (Vaswani et al., 2023):
pos
{
cos
10000
if < "122
(10)
sin
otherwise
10000977M
Let X E R be the encoder output. We first perform group-wise aggregation over the sequence length (e.g., logsumexp or mean) to obtain vectors Yo, Yi E R the aggregate representations of groups and 1, The queries V are then
Data-Dependent Queries
{ Yı,
if g; =
V = VgJ +
+1
(11)
Yo otherwise.
5 EXPERIMENTS
We compare PGM with MDLM (Sahoo et al., 2024) on standard language modeling datasets, training on LM1B (Chelba al., 2014) and OpenWeb'T (OWT; Gokaslan & Cohen (2019)) in Sec. 5.1.

---

Published as a conference paper at ICLR 2026
We evaluate them using the validation perplexity and downstream task aceuracy before and after distillation with SDTT (Deschenaux & Gulcehre, 2025). We compare PGM with MaskGIT (Chang et al., 2022) As described in Sec. 3, by predicting each group from the other, PGMs implement a mechanism akin to training on two complementary masked sequences per batch, while also introducing new architecture (Sec. 4). The effect of complementary masking is studied in isolation in Sec. 5.3. Our experiments show that, for both language and image modeling, and after distillation, PGMs are competitive with MDLM and MaskGIT, while providing a 5-5.5x throughput improvement for text and a 7.5x improvement for images. Find more experimental details in Suppl. C.
"QGAN--uaatiza
(Esser et al., 2021) ImageNet256 (Deng al., 2009) (Sec. 5.2).
5.1 LANGUAGE MODELING
Experimental settings We closely follow the settings of Sahoo et al. (2024). MDLM uses a modified Diffusion Transformer (Peebles & Xie, 2023; Lou et al., 2024) with RoPE (Su et al., 2023), with 12 layers and an embedding dimension of 768, without time conditioning. We train with a global batch size of 512 for 1M steps, dropout of 0.1, and the Adam optimizer with learning rate 3 x 10-1 and no weight decay. We maintain an Exponential Moving Average (EMA) of the weights with decay 0.9999. For PGM, we use the Partition Transformer architecture (Sec. 4) with 12 or 16 layers, embedding dimensions of 768 or 1024, and varying numbers of encoder and decoder layers. On LMIB, all models use a context length of 128, with shorter documents padded and tokenized using the bert-bass-uncasee
(Devlin et al., 2019) tokenizer. On OWT, we use context length of 1024 with sentence packing (Raffel et al., 2023) with the GPT-2 tokenizer and insert an [EOS] token between documents. Since the dataset lacks official validation split, the last 100k documents reserved for validation. To evaluate the sample quality, we use the Generative Perplexity (Gen. PPI.), computed using GPT-2 Large (Radford et al., 2019), following Sahoo et al. (2024). We cast the logits in float6 4 prior to sampling, following Zheng et al. (2025).
Likelihood Evaluation After 1M steps, PGMs with as many layers as MDLM achieve validation perplexity of 1.95 lower than MDLM on LMIB (Table 1). Table 5 (left) shows that balanced models with equal numbers of encoder and decoder layers outperform imbalanced variants. Interestingly, data-independent queries perform comparably to data-dependent queries, so we use the simpler, data- independent version in all subsequent experiments. On OpenWebText, PGMs with the same number of layers and embedding dimension as MDLM slightly underperform (Table 5, right). Increasing the number of encoder and decoder layers by two, increasing the embedding dimension to 1024, allows PGMs to surpass MDLM in validation perplexity, while achieving at least 5 x higher sampling throughput. This improved efficiency makes PGMs particularly attractive for scaling test-time computation (Madaan et al., 2023; Yao et al., 2023; Snell et al., 2024; Wu et al., 2024; Chen et al., 2024; Brown et al., 2024; Goyal et al., 2024).
Downstream Evaluation Following Deschenaux & Gulcehre (2024); Nie et al. (2025), we evaluate MDLM and PGMs trained on OpenWebTen using the 2024). As shown in Table 2, PGMs slightly outperform MDLM on six out of eight tasks, although the overall accuracy across models is similar. This suggests that PGM achieves faster inference without sacrificing downstream performance. Since 1m-eval -harness is originally designed for ARMs, must adapt it for MGMs. Fortunately, both MDLM and PGM can compute variational bound on the likelihood, which is used in place of the true likelihood to select the most probable answer in multiple-choice tasks. Additional details and tasks provided in Suppl. D.5.
suite (Gao et al.,
Distillation of PGMs After likelihood training, PGMs achieve 5.5x higher throughput than MDLM. To further accelerate sampling, we apply Delf--istillation Through Time (SDTT; Deschenaux & Gulcehre (2025)). To remain as faithful as possible to the implementation of Deschenaux & Gul- cehre (2025), we apply the distillation loss to single group while treating the other as [MASK] tokens. This shows that PGMs are compatible with distillation methods designed for MGMs. We leave the development of new distillation strategies for PGMs to future work. Hence, the setup naturally favors MDLM. Figure 4 (right) and Table 6 compare the Gen. PPL., unigram entropy, and sampling speed of PGM and MDLM. After five rounds of distillation, and with standard ancestral sampling. PGMs achieve higher Generative Perplexity and entropy than MDLM. With nucleus sampling (p = 0.9) (Holtzman al., 2020), PGMs produce samples with comparable perplexity and entropy. Due to
7

---

Published as a conference paper at ICLR 2026
the overhead of nucleus sampling, the speed advantage of PGMs decreases from least 5x to approximately 4.6x faster than MDLM for the same number of steps (Fig. 4). Generative perplexity alone does not fully capture language model performance, hence also evaluate distilled models on downstream tasks. As shown in Table 2, distillation slightly shifts accuracy across tasks, but overall performance remains similar. PGMs still achieve slightly higher accuracy than MDLM most tasks after distillation.
5.2 IMAGE MODELING
Experimental Settings We train MaskGIT
(Chang et al., 2022) and PGM on ImageNer256.
Images are cropped to a centered square along the longer side and then rescaled to 256 X 256. Besnier
MDLM
PGM
PGM-nudeus (p=0.3)
use the MaskGIT implementation et al. (2025), including their pre-trained VQ- GAN tokenizer. We train for 500k steps with batch size of 256 using AdamW (weight de- cay 0.03, learning rate le-4, cosine schedule with 2500 warmup steps). We a dropout of 0.1 in the Transformer. All models are class- conditional, with class-label dropout of 0.1 to enable classifier-free guidance (CFG) at sam- pling time. As Besnier al. (2025), we train with one register (Darcet et al., 2024) for the MaskGIT baseline, and two (one per group) for PGM, so that we can use one register during sampling. We sample with the confidence and Halton samplers.
15.50
15.53
(5.4)
65.4
50-
(3.4)
15.41
(5.4)
35
Throughput (lokens/sec.)
Figure 4: After distillation, PGM (6/ 6, dim. 1024) with nucleus sampling remains significantly faster than MDLM, at matching entropy and Gen. PPL.
Results In Figure (left), we compare the Fréchet Inception Distance (FID; Heusel et al. (2018)) of samples from MaskGIT (Chang et al., 2022) and PGM, using the Halton sampler and classifier-
free guidance with the guidance weight w E {0. .,..... .6} that yields the lowest FID. PGM 12/12 achieves 7.5x higher throughput with only slight FID degradation (5.54 vs. 5.35). Increasing
the sampling steps to 64 further improves the FID to 4.56, while remaining 3.9x faster than MaskGIT. See Suppl. D.3 for full results across guidance strengths.
5.3 ISOLATING THE EFFECT COMPLEMENTARY MASKING
Experimental Setup To disentangle the contributions of PGM, we isolate the effect of complemen- tary masking (Sec. 3) by training standard bidirectional Transformer with double the batch size. Fach input sequence is turned into two complementary masked copies: if the token at position & is masked one copy, it remains unmasked in the other. This setup provides an upper bound on the potential gains, it directly measures the benefit of complementary masks during training.
Results Table shows that complementary masking improves the validation perplexity on LMIB and OWT, though with smaller gains OWT. On both datasets, a gap remains between PGM and MDLM with complementary masking. This suggests that the current neural network architecture can be improved further. Because of the smaller improvement count to surpass MDLM. Nonetheless, recall that despite having more parameters, PGMs remain at least 5x faster than MDLM during sampling. In Suppl. D.I, we present preliminary experiments exploring why complementary masking improves performance on LM1B but not on Open Web
OWT, must increase the parameter
6 RELATED WORK
Discrete Diffusion Although autoregressive models currently dominate text generation, recent advances in discrete diffusion (Austin al., 2023; Lou et al., 2024; Shi et al., 2025; Sahoo et al., 2024; von Rütte et al., 2025; Schiff et al., 2025; Haxholli et al., 2025; Sahoo et al., 2025a) and discrete flow matching (Campbell al., 2024; Gat et al., 2024) have demonstrated that MGMs can
8

---

Published as a conference paper at ICLR 2026
Table 2: Accuracy on downstream tasks (Gao et al., 2024). HS: HellaSwag. OQA: OpenBook
QA. Are: Arc-easy. We select the tasks following Nie et al. (2025). We see that distillation slightly changes the downstream tasks performance, but that PGMs continue to outperform MDLM on most tasks. The best performance is bolded, while the second best is underlined.
LAMBADA
Arc
BoolQ
HS
OQA
PIQA
RACE
SIQA
Before Distillation MDLM
38.52
37.88
49.42
31,36
28.60
58.27
28.04
38.84
PGM 8/ 8/8
46.98
40.40
53.49
33,20
26.60
58.92
26.89
39.97
PGM 6 5/6 (1024)
41.39
39.98
49.82
34.27
25.40
59.19
27.37
40.28
After Distillation (SDTT)
MDLM
41.34
33.80
48.59
30.75
28.80
57.73
27.94
38.79
PGM 8 8/8
47.22
37.42
51.50
31.62
25.80
59.03
30.62
39.61
PGM 6/6 (1024)
44.48
36.70
49.36
32.55
25.00
59.85
27.37
39.25
approach AR models in generation quality. We propose a simple framework that allows sampling without processing any [MASK] tokens, but remains compatible with methods developed MGMs (such distillation and alternative samplers).
Variable Length Masked Diffusion Block Diffusion (BD; Arriola et al. (2025)) enables partial KV-caching (Pope al., 2022) by generating tokens block-by-block using discrete diffusion. BD improves throughput but sacrifices the any-order generation capabilities of MGMs. We do not experiment with integrating causal attention enable KV-caching. However, Ma al. (2025): Wu et al. (2025) show that KV caching can be integrated post-hoc into MGMs despite being trained without causal attention. FlexMDM (Kim et al., 2025) and Edit Flows (Havasi al., 2025) enable variable-length generation via insertion, deletion, and replacement. While promising, these depart from the simplicity of MGM and PGM. Finally, Eso-LMS (Sahoo al., 2025b) train with hybrid AR-MGM objective. Sahoo et al. (2025b) first sample draft in MGM mode, then fill in the remaining tokens autoregressively. During training, Eso-LMs must choose the fraction of examples to process in AR versus MGM mode, which adds hyperparameter to tune. Eso-LMs use [MASK] tokens during training in MGM mode, whereas PGMs do not because of the Partition Transformer.
Non-Autoregressive Language Models Any-order and any-subset autoregressive models (Yang et al., 2020; Pannatier et al., 2024; Shih et al., 2022; Guo & Ermon, 2025) factorize the sequence distribution autoregressively over permutations of tokens. Hence, these models use causal attention and generate tokens one by one. In contrast, MGMs use bidirectional attention and generate multiple tokens in parallel, which is the setting PGM builds on.
CONCLUSION
We introduce Partition Generative Modeling (PGM), a novel approach masked generative modeling that eliminates [MASK] tokens entirely. PGM achieves significant improvements in inference speed on both text and images, with minimal effect on quality. The significant improvements suggest that PGM might be suited for domains that benefit from test-time scaling, such as coding and reasoning. We show that PGMs can be distilled for further acceleration. Future work should explore optimizations the PGM architecture, investigate distillation techniques specifically designed for PGMs, and extend the approach to multimodal settings. In summary, PGM offers an alternative masked generative models, with particular advantages for applications where inference speed is critical.

---

Published as a conference paper at ICLR 2026
8 ACKNOWLEDGEMENTS
This work has received funding from the Swiss State Secretariat for Education, Research and Innovation (SERT). We are grateful to Razvan Pascanu, Sungjin Ahn, Jaesik Yoon, Mingyu Jo, Subham Sahoo and Zhihan Yang for insightful discussions and suggestions. We acknowledge the SCITAS team at EPFL for providing access to their cluster, and the Swiss National Supercomputing Centre for the Alps platform. We are grateful to Karin Gétaz for her administrative assistance.
REFERENCES
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang. Zhixuan Qi. Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregres- sive and diffusion language models, 2025. URL https: /ttp:///riiii arxiv. org/ abs 09573.
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2023. URL https: //aa:///////i/////// wwwpiv.vvvvv///.v///.aaaaaaaa http:/aawwv......aaaaaa/a abs//
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL https: /ttpiviivv../////// orgg wwwp/////...///////////////..w./ 04450.
Victor Besnier, Mickael Chen, David Hurych, Eduardo Valle, and Matthie Cord. Halton scheduler for masked generative image transformer, 2025. URL https:://aaaatttee https////axwvvv..aaaaaeea /503.. 17076. 6.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024. URL https::t///o arxiv. org/ abs 22002440.222 211877.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows discrete state-spaces: Enabling multimodal flows with applications to protein co-design, 2024. URL httph::////// //tttp:://rxxi.../// org/ 2402. 2492204997.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman. Maskgit: Masked generative image transformer, 2022. URL https: /ttpiv.xivvvvvv///////.. org/ abs///b//20 abs//202..2222 2202. 04200.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling, 2014. URL https::/// arxiv. org. aabs///////// 11111111111 3005..
Lingjiao Chen, Jared Quincy Davis, Boris Hanin, Peter Bailis, Ion Stoica, Matei Zaharia, and James Zou. Are more Ilm calls all you need? towards scaling laws of compound inference systems, 2024.
URL https:://// /ale://aaxxa /arxiviiiv..////vv///av////aavvv org/ abs 2bss/200 2403. 02419.
Marco Comunità, Zhi Zhong, Akira Takahashi, Shiqi Yang, Mengjie Zhao, Koichi Saito, Yukara lkemiya, Takashi Shibuya, Shusuke Takahashi, and Yuki Mitsufuji. Specmaskgit: Masked gen- erative modeling of audio spectrograms for efficient audio synthesis and beyond, 2024. URL
https: / arxii.v.. wwwpxivvoorrgooooo org/ abss /2006//24666
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers, 2024. URL
orgg www.oogg/////oooooooo/o...ot...r 2209 6588.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
Justin Deschenaux and Caglar Gulcehre. Promises, outlooks and challenges of diffusion language modeling, 2024. URL /atpss///ariii www..oorgg/./.... /2////bb//246..... 11473.
Justin Deschenaux and Caglar Gulcehre. Beyond autoregression: Fast Ilms via self-distillation through time, 2025. URL https: //arxii/////aa///a/a wwtpw//wwwwww//////////////////e 2335.
10

---

Published as a conference paper at ICLR 2026
Jacob Devlin, Ming- Ming-Wei bidirectional transformers for language understanding, 2019. URL httpp:://axxi...//// abs/1810 04805,
Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, and Jonas Adler. Continuous diffusion for categorical data, 2022. URL
/ttpss:/azzae org/ abs 20... 15089.
Patrick Esser, Robin Rombach, and Björ Ommer. Taming transformers for high-resolution image synthesis, 2021. URL hwwpt:tpp///o //aaariiiiiv///.i///aaaaa/aaaava org/ abs/ 2012. 09841.
Rohit Gandikota and David Bau, Distilling diversity and control in diffusion models, 2025. URL
https: //axivv org/ als 2503. 2
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. URL https:tt////mm...t www.zeeod.ooooooooo
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching, 2024. URL
15595.
Aaron Gokaslan and Vanya Cohen, Openwebtext corpus. http:://8kyLic:007.gttttt io// OpenWebTexttCorput 2019.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens, 2024. URL
https::/aae org/
wtggabs//////ooooooo/ggg///// absa///s111 2210..0... 02226.
Gabe Guo and Stefano Ermon. Reviving any-subset autoregressive models with principled parallel sampling and speculative decoding, 2025. URL http:: / /ttpxvavvvvvv///////eee/eaaaal wrtg/bs//////aa/e 20456
John H Halton. Algorithm 247: Radical-inverse quasi-random point sequence. ACM, 7(12):001-70,, 1964.
Marton Havasi, Brian Karrer, Itai Gat, and Ricky T. Q. Chen. Edit flows: Flow matching with edit operations, 2025. URL https: /ttp///avvvv//////////////////al org/ abs /200 09018.
Etrit Haxholli, Yeti Z. Gurbuz, Ogul Can, and Eli Waxman. Efficient perplexity bound and ratio match- ing in discrete diffusion language models. In The Thirteenth International Conference on Learning Representations, 2025, URL https: / wpwp/eveeeeeeeeee////ee .ww./nee///e.e htt::/foummmmddiam 9WI Sm.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium, 2018. URL https: arxiv. org/ abs//100 08500.
Jonathan Ho and Tim Salimans. Classifier free diffusion guidance, 2022. URL httptt/////ax
org/abs abs//ab//aa 2207.22222 12598.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020. URL https://aat /ttp:vi//////ve//a org/ abs 2200..1299.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text /atp////////////aa abs 09751.
degeneration, 2020. URL
Jaeyeon Kim, Lee Cheuk-Kit Carles Domingo-Enrich, Yilun Du, Sham Kakade, Timothy Ngotiaoco, Sitan Chen, and Michael Albergo. Any-order flexible length masked diffusion, 2025. URL https:/ /ttp//aavvvv///////////////a.avl wtt//bss/////oooooooooo///aaa 2509. 01025.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL
https://aran wwweiiioorr///// htrg//bb//// 6980.
11

---

Published as a conference paper at ICLR 2026
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models, 2023.
URL https: /awxivvvvvv../////a//aa////aaavv org/ abs/2101 00630.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution, 2024. URL https: http://aaxxvvv..aaaaooo/ 2ab.b///3311.18834
Xinyin Ma, Runpeng Yu, Gongfan Fang, and Xinchao Wang- dkv-cache: The cache for diffusion 15781.
language models, 2025. URL /ttttss//arxivv.or////aa /
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang. Shashank Gupta, Bodhisattwa Prasad Majumder, Katherine Hermann, Sean Welleck, Amir Yazdanbakhsh, and Peter Clark. Self-refine: Iterative refinement with self-feedback, 2023. URL / http//axxvv.rrro
2303. 17651.
Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling up masked diffusion models on text, 2025. URL https: /ttp:v//////vev/////eeeeeeeeevea org/ org/abs////aaaaoooo//goooobbaaae 1410.18514
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, 2025. URL httph:tp/////c. / arxivv org/ abs /206 03736.
Arnaud Pannatier, Evann Courdier, and François Fleuret. Sigma-gpts: A new approach to autoregres- sive models, 2024. URL https: /attts://azxiv..rrtt org/abs 2404.400000 09562.
William Peebles and Saining Xie. Scalable diffusion models with transformers, 2023. URL https:
www.xii.org/bb.... 22 09748.
Reiner Pope. Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference, 2022. URL https: /tattps://arivvy arrivv org/ abs//////////// 2212.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL wwwpttttt//// /ttettss://riiiiii /twpvv..ivvv/v////////..//..aav. /tgg/ab///a /1900/ 910. 10683.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models, 2024. URL https: wrw.a.avvvv.oo. www.wiiww../aaaa.. abs/ 2406 07524.
Subham Sekhar Sahoo, Justin Deschenaux. Aaron Gokaslan, Guanghan Wang, Justin Chiu, and Volodymyr Kuleshov. The diffusion duality, 2025a. URL htt.tttss//axi...ooot aass/aa///a 2206. 10892.
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, 2025b. URL. htl.httss:////e //t////axiv//v orgg abs / 6506. 01 928.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models, 2022. /t.bs/2200 00512.
URL bttts:://zziv..or
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Im- proved techniques for training gans, 2016. URL h 1606. 03498.
Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang. Sam Boshar, Hugo Dalla-torre, Bemardo P. de Almeida, Alexander Rush, Thomas Pierrot, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models, 2025. URL https::/ http:///aaxvvv.gg/a/ae/ 2412. 1412.1019..
Jiaxin Shi, Kehang Han, Zhe Wang, Amaud Doucet, and Michalis K. Titsias. Simplified and generalized masked diffusion for discrete data, 2025. URL https: http///axvvvcie org/abs/////.......... 2406. 04329.
12

---

Published as a conference paper at ICLR 2026
Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference onyer autoregressive models the right way, 2022. URL
abs / 13554.
Charlie Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters, 2024. URL https://: wrg. / abs/ 2408. 03314.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2256-2265, Lille, France, 07-09 Jul 2015. PMLR. URL https:// html.
proceedings mwm.m.r..... www.spress/vva httr///hhhaaccccseiiiesa
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, 2020. URL https:://///.. http://///aaa...../aa/////////// 056000000000
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023. URL https: 09864.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL /tttps://arxiv.oom abs/1706. 03762.
Ruben Villegas, Mohammad Babacizadeh, Picter-Jan Kindermans, Hernan Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan. Phenaki: Variable length video generation from open domain textual description, 2022. URL https: org/abs/ 2210 02399.
Dimitri von Rütte, Janis Fluri, Yuhui Ding, Antonio Orvieto, Bernhard Schölkopf, and Thomas Hofmann. Generalized interpolating discrete diffusion, 2025. abs/2503. 04482.
Chengyue Wu, Hao Zhang. Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training-free acceleration of diffusion 1lm by enabling kv cache and parallel decoding, 2025. URL https:/ /aaxiv.ivvvv..///////aa///aaaava orggiv.rr/// abs//////// 22005.55522 .2.2226222 22618.
Yangzhen Wu. Zhiqing Sun, Shanda Li. Sean Welleck, and Yiming Yang. An empirical analysis of compute-optimal inference for problem-solving with language models, 2024. URL https: //arxiv. htt.vvorg/abb.. ablgg/bb//440 /408 00724.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. Xlnet: Generalized autoregressive pretraining for language understanding, 2020. URL
https:/ //arxiv..vv//////// httg//abs/1900 08237.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023.
URL https::t/eeeeeec.t /tassv////xvv../// org/ abs/ 10601.
Lijun Yu, Yong Cheng, Kihyuk Sohn, José Lezama, Han Zhang, Huiwen Chang, Alexander G. Hauptmann, Ming-Hsuan Yang, Yuan Hao, Irfan Essa, and Lu Jiang. Magvit: Masked generative video transformer, 2023. URL https: arxiv. org/ abss 2212.222222 05199.
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling, 2025. URL https::/arxiv..rrgg
2409 02908.
Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière, and Vicky Kalogeiton. Di[Mo: Distilling masked diffusion models into one-step generator, 2025. URL pattps//aaree http://awpw.gg/////..////o arg/abs//50000 15457.
13

---

Published as a conference paper at ICLR 2026
Algorithm Building the Halton Unmasking Schedule
1: Input: Grid size H
2: Output: Ordered list of L - H² grid positions
3: schedule f- 0
4: seen t-
5: 2t] 1
6: while |schedule| < do
7: cell < (1) H,, [1) H))
8:
if' cell & seen then
9:
seen - seen U [cell)
10:
schedule..ppeend(eel)
end if
12:
i+i+1
end while
14:
return schedule
A LIMITATIONS
To match the validation perplexity the MDLM baseline at a context length of 1024, our models require slight increase in parameters. We attribute this to the GroupSwap layer, and future work will explore more efficient mechanisms for information exchange between groups in PGMs. While PGMs offer faster inference, their training is slightly more computationally expensive (Appendix E), as we use torch' 's default attention implementation ("sdpa") for simplicity. By reordering tokens according their group assignment, the self-attention matrices becomes block-diagonal. Future work will explore efficient kernel implementations that exploit this block-diagonal sparsity. Partition Generative Modeling is general framework, and its application to multimodal settings remains an open direction for future research.
B SAMPLING WITH HALTON SEQUENCES (MASKGIT)
Let S = fe independently at all the masked positions from a product of marginal predictions Iles Pe (x" 2z), not from the joint p(x 2.). The KI. divergence between the joint and the product marginals the mutual information (MI) (Besnier et al., 2025):
MASK] } denote the set of masked positions step MGMs samples
DEL p(ta 2z)
Il EES
po 2.)) MI(fa'Jess zT).
(12)
Empirically, the denoiser is most confident positions close previously generated tokens. There- fore, the decoded tokens tend to cluster together. This leads to a larger MI, with a higher risk of generating inconsistent tokens. Besnier et al. (2025) propose replace the confidence-basee ordering with Halton sequence (Halton, 1964), a low-discrepancy sequence whose consecutive points are far apart in space. Formally, let (i) denote the j-th digit of the representation of N4 in base so that E, a;(1) &., Let the radical-invers function o denote the inverse of in base b:
(b) Σ (1(())-(5t)) 10.
(13)
Recall that MaskGIT operates over a square grid of VQGAN tokens (Esser et al., 2021), of size H x H (L = H2 To build the unmasking order (Algo. 1), Besnier et al. (2025) iterate over positive integers and compute the cell coordinates (i)) HJ. ( HJ). If the cell was already visited for some i' < i, it is skipped. This continues until all cells are visited. The Halton scheduler consistently improves the FID and, unlike confidence- based sampling, continues to improve with more inference steps (Besnier et al., 2025). In our image experiments, we use both confidence and Halton schedulers.
14

---

Published as a conference paper at ICLR 2026
Hliiton Samppiing
Condidennee-aaeed Sampling
SAMPLING TIME
Figure 5: Comparison of unmasking schedules 2D grid. (a) Halton sampling covers space uniformly. (b) Random sampling can leave large areas empty. (c) With confidence- based sampling. new tokens are decoded close existing ones. leading high mutual information and a risk of inconsistent generation.
C EXPERIMENTAL DETAILS
We trained all models from scratch. Our baselines achieve similar performance as reported by Sahoo al. (2024). On LMIB, we obtain a validation perplexity 27.67 after IM steps (compared to MDLM's reported 27.04), while OWT, we reach 23.07 (versus MDLM's 23.21).
Minor differences can be expected since estimating the perplexity of diffusion language models involves Monte-Carlo approximation of the NELBO (1) with finitely many samples. Although used libraries (e.g., PyTorch) with the same version as MDLM, differences compute environments and underlying software stacks may also contribute to these variations. Since the performance gap is small, we are confident that we used the code of MDLM correctly.
C.1 LMIB
tokenizer with context length of 128 tokens, padding shorter sequences. Our architecture consisted a Diffusion Transformer (DiT) with 12 transformer blocks, 12 attention heads, hidden dimension of 768, and dropout rate of 0.1. We optimized the model using Adam (Kingma & Ba, 2017) (learning rate 3e-4, betas of (0.9, 0.999), epsilon le-8) without weight decay. We based our implementation We trained with a global batch size of 512 across GPUs (2 nodes with 4 GPUs), gradient clipping at 1.0, and constant learning with an EMA rate of 0.9999. Besides the neural network hyperparameters, the other parameters were unchanged when training the PGM.
For the LMIB dataset, we employed the bert See--baae-uncee
the official MDLM codebase.
with 2,500 steps of linear warmup. We trained for I million steps
C.2 OWT
For the Open WebnWet (OWT) dataset, we used the GPT-2 tokenizer with a context length of 1024 tokens. Our architecture consisted of Diffusion Transformer (DiT) with 12 transformer blocks,
15

---

Published as a conference paper at ICLR 2026
Table 3: Latency and throughput for a single forward+back pass of the MDLMs and PGMs. computed on a single A100-SXM4-80GB GPU. On LMIB, PGM introduces negligible overhead over MDLM. On OWT, our PGM with 6 encoder and decoder layers and an embedding dimension of 1024 achieves around 75% of the training throughput of MDLM. Recall that at inference, the same PGM around 5x faster than MDLM.
Forward Pass
Forward Backward
Model
Latency (ms)
Seq/sec
Latency (ms)
Seq/Sec
LMIB (context length 128, hatch size 64, trained on 8 GPUs)
MDLM
0.03 + 0.00
1'978.87 + 44.21
0.08 + 0.00
714.80 + 15.47 794.42 + 18.81
PGM 6/6
0.03 + 0.00
1'966.60 + 102.14
0.08 + 0.00
OpenWebText (context length 1024, batch size 32, trained on GPUs) MDLM
0.13 + 0.00
233.28 + 2.58
0.39 + 0.00
80.86 + 0.15
PGM 8 8
0.17 + 0.00
188.07 + 0.75
0.47 0.00
68.04 + 0.08
PGM 6/6 (dim. 1024)
0.18 + 0.00
176.47 + 0.65
0.50 + 0.00
62.85 + 0.19
12 attention heads, hidden dimension of 768, and dropout rate of 0.1. We optimized the model using Adam (Kingma & Ba, 2017) with a learning rate of 3e-4, betas of (0.9, 0.999), and epsilon of 1e-8, without weight decay. We trained with global batch size of 512 across 16 GPUs (4 nodes with 4 GPUs). We applied gradient clipping at 1.0 and used a constant learning rate schedule with 2,500 steps of linear warmup. The model was trained for 1 million steps with EMA rate of 0.9999. Besides the neural network hyperparameters, the other parameters were unchanged when training the PGM.
C.3 IMAGENET
For the ImageNet experiments, we used a pre-trained VQGAN tokenizer (Esser et al., 2021; Besnier al., 2025), following the setup of Besnier et al. (2025). The images are tokenized into sequences of 1024 tokens. This allowed for a direct comparison between PGM and MaskGIT, both trained in the codebase of Besnier et (2025) and the FID is evaluated using the Halton sampler and the confidence sampler. We compute the FID between 50k generated images and the validation set, following Besnier et al. (2025)
All models use 24 transformer blocks. For PGM, we add GroupSwap layer to enable information exchange between partition groups. We use the same hyperparameters as HaltonMaskGIT for all models, except reduce the training duration 500k steps (from 2M) due to computational constraints. All models are trained to be class-conditional, which enables the use classifier-free guidance to significantly improve performance.
C.4 IMPACT OF NUMERICAL PRECISION ON SAMPLING
Zheng et al. (2025) identified that Masked Diffusion Models often achieve lower Generative Perplexity results because of underflow in the logits when sampling using low precision. The resulting decrease in token diversity can make evaluations based solely on Generative Perplexity misleading. Hence, we always cast the logits to FP64 before sampling.
C.5 SAMPLF-BASED EVALUATION
Generative Perplexity We use the Generative Perplexity to evaluate the quality of samples, follow- ing prior work (Lou et al., 2024; Sahoo al., 2024; Deschenaux Gulcehre, 2025). The Generative Perplexity measures how well reference model (in our case, GPT-2 Large) can predict the next token generated sequences. Specifically, we generate 1'024 samples from each model being evaluated. For each generated sample, we compute the Generative Perplexity using GPT-2 Large as follows:
Perplexity = =cxp (-L
Σ
log Parr2 Large (I, |BEEs))
(14)
16

---

Published as a conference paper at ICLR 2026
MDLM 10
MDLM (Complementary masking)
PGM (6/6, dim-1024)
10
10
Ttairing loss
Trainingg loss
Training loss
4-
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Training stop
166
Training step
106
Training step
Figure 6: Training loss MDLM, MDLM with Complementary Masking (Section 5.3) and PGM. Complementary masking seems to introduce spikes in the loss, even though it did not cause the models to diverge.
where L is the length of the sequence, 2; is the i-th token, and PGPT-2 (z, (x3))) is the probability assigned by GPT-2 Large to token given the preceding tokens Zei.
Unigram Entropy Unfortunately, a low Generative Perplexity can be achieved by generating repetitive text. To catch such cases, we compute the average unigram entropy of the generated samples:
x(Uxx(1)) log
x())
Unigram Entropy
(15)
L
L
where V is the vocabulary, e is token of the vocabulary, and c(e, x) is the empirical appearance count of the token v in the sequence x. Low unigram entropy helps us to catch degenerate generation, shown by Dieleman al. (2022).
Fréchet Inception Distance and Inception Score On image generation tasks, we evaluate the quality of samples using the Fréchet Inception Distance (FID) (Heusel et al., 2018) and Inception Score (IS) (Salimans et al., 2016). Both metrics are computed using 50'000 images, following the standard practice.
D ADDITIONAL RESULTS
D.1 IMPACT OF CONTEXT LENGTH ON THE EFFECTIVENESS COMPLEMENTARY MASKING
There three key differences between experiments LMIB and OWT. First, used different tokenizers: bert-base-uncasee
for LMIB and GPT2's 's tokenizer for OWT, following the setup of MDLM (Sahoo et al., 2024). Second, the context lengths differ significantly: 128 tokens for LMIB versus 1024 for OWT. Third, we train on different datasets that might have different characteristics.
We observed that complementary masking helps when training on OWT using a shorter context length of 128 tokens with the GPT-2 tokenizer. Indeed, after 200k training step, the MDLM with complementary masking achieved validation PPL of 37.92, outperforming the standard MDLM, which reached 39.90. This suggests that PGMs may not need extra parameters when the sequence length short. Exploring the use of PGMs domains where the sequence length is short, such as modeling chemical sequences, is promising direction for future work.
D.2 MDLM+SDTT VS PGM+SDTT
The precision of logits during sampling can have significant effect sample quality, noted in Appendix C.4. Hence, we cast all logits to FP64 prior to sampling, unlike the original MDLM and SDTT implementations.
17

---

Published as a conference paper at ICLR 2026
Using higher precision also affects distillation, which compresses two sampling steps into one. As shown Table 7, models distilled with float32 achieve lower Generative Perplexity than those trained with mixed precision (bfloat16). We therefore report float32 results in the main body.
D.3 ADDITIONAL RESULTS ON IMAGENET
Table and Table 9 show the FID, IS, latency, and throughput for the Confidence and Halton samplers. Overall, the Halton sampler works best for both MaskGIT and PGM. With 32 steps and the confidence-basee sampler, PGM 12/12 gets a better FID than MaskGIT and is 3.58x faster. With 32 steps and the Halton sampler, PGM has slightly higher FID than MaskGIT (5.54 5.35), but is 7.5x faster. If we increase the number of sampling steps to 64, PGM achieves an FID of 4.56, which better than MaskGIT, and 3.92x faster. Generally, the 12/12 variant outperforms the 14/10 variant, which suggests that balanced number of layers in the encoder and decoder is beneficial, just as for language modeling (Table 5).
D.4 TRAINING STABILITY
Complementary masking introduces occasional spikes in the training loss in both MDLMs and PGMs, as shown in Figure 6. This phenomenon should kept in mind when scaling PGMs to larger sizes. Despite these spikes, all runs converged on the first attempt. We observed different precision requirements between models. For loss computations, MDLMs performed best with BFI6 precision, while PGMs achieved better results with FP32 precision. Both models use mixed precision within the neural network; the precision difference only affects computations performed outside the model, such the loss calculation.
D.5 ADDITIONAL DOWNSTREAM TASKS
Table 4 reports additional downstream results as in Deschenaux & Gulcehre (2025), where PGM outperforms MDLM all but one benchmark, with only small gap on the latter. We evaluate models with the lm-eva -aaallhhaeee harness library (Gao et al., 2024), originally designed for autoregressive LMs and adapted here for MDLM. For multiple-choice tasks, 1m-eva -ha rness computes the log- likelihood of each candidate answer y, given a prefix x, i.e., p\y,,l(,)))))))))))).p)))()) / and selects the answer with the highest score.
While 1m Lm-eval harness uses the log-likelihood of the continuation, the NELBO objective (I) bounds the log-likelihood of the complete sequence (x, yz). However, we only need know which continuation achieves the highest log- ligliikl, not to compute the exact log- liggliiklihood. Using Bayes' theorem, we note that
log (x) = log px, - log p(x) x log,(xx,5;
(16)
since log((x)) constant with respect to y.. Therefore, we can simply evaluate the variational bound on log p(x, y:) to select the most likely continuation y/-
D.6 PERFORMANCE ON LONGER CONTEXT LENGTH
Due to the high computational cost, we were unable to train models with context lengths greater than 1024. Nevertheless, report the latency and throughput of both MDLM and PGM at context length 4096. As shown in Table 10, PGM remains substantially faster than MDLM in this setting.
E COMPUTATIONAL COSTS
This section presents the computational costs associated with the models reported in this paper. We exclude costs associated with exploratory experiments that yielded inferior results and were not included in this manuscript.
E. TRAINING COSTS
Training PGMs is currently slower than training MGMs since we use torch. sdpa with dense tensor masks. Future work should explore efficient kemels address this limitation. We measure the
18

---

Published as a conference paper at ICLR 2026
Model
LAMBADA
ARC-c
ARC-c
HSwag
MathQA
PIQA
WinoG
MDLM
38.52
34.26
24.66
31.54
20.70
57.89
51.93
PGM 8 PGM 6/ 6 (1024)
46.98
37.37
24.06
33.10
21.24
59.09
51.30
41.39
38.80
22.95
33.92
21.71
61.43
54.30
Table 4: Accuracy on downstream tasks. We evaluate MDLM and PGM on LAMBADA, ARC Easy and Challenge, HellaSwag, MathQA, PIQA, and WinoGrande. Both models show comparable performance across tasks. PGM outperforms MDLM on all but one benchmark, where the difference between MDLM and PGM 8 //8 8 small.
latency and throughput using single NVIDIA Table We compute the mean and standard deviation over 100 batches after 2 warmup batches.
GPU, with results reported in
The total training duration approximately equals the per-step latency multiplied by the number of steps. Experiments with complementary masking required twice the computational resources due to larger batch sizes and gradient accumulation. Training times for 1M steps varied by dataset: approximately 22 hours for LM1B, 4.5 days for OWT, and 3.8 days for ImageNet.
Despite the current training overhead, we are confident that future work improve the training efficiency of PGMs, thanks to their block-diagonal attention patterns, once the tokens are grouped together along the sequence-length axis.
E.2 INFERENCE COSTS
We evaluate the inference efficiency of PGMs compared MDLMs and GPT-2 with KV caching. As shown in Figure 1, PGMs achieve around 5 - 5.5x improvements in throughput over MDLM while reaching superior Generative Perplexity. For inference measurements, we single NVIDIA A100- SXM4-80GB GPU. The efficiency gain stems from the ability of PGMs to process only unmasked tokens during inference, as illustrated in Figure 2. Table 6 compares MDLM and PGMs on the Generative Perplexity, unigram entropy, latency, and throughput. We compute the mean and standard deviation the latency and throughput over 20 batches after two warmup batches.
E.3 LICENSING
Our code and model artifacts will be released under the MIT license. The OWT dataset (Gokaslan & Cohen, 2019) is available under the Apache License 2.0. We were unable to identify a specific license for the LMIB dataset (Chelba et al., 2014). The images in ImageNet remain the property of their respective copyright holders.
Algorithm 2 Simplified Sampling for PGMs
1: Input: Batch size BS, number of steps K, model length L, special BOS index
2: Output: Generated samples x
b Initialize Set BOS as first token b Number of tokens to denoise at each step Keep track of already-decoded and positions to decode 7: positions_tc to decode + 1+ rand row Perm(BS, L-1) Each rows is permutation of {1,..., L}
3: X < empty_ Lenpty_tensor(BS, 1)
4: x[:, BOS
5: k < L/K
6: decoded pposttiossssssss + zBS,
8: for ra)ge(K) do
Random positions to be predicted 10: new_values + pgm agm_prredcc(ss decoded _otiided__pstioonss pos_to_dees co__ecode)
9: pos_to_ co__eeoddeeeeeeoee +- positions sociiioossttoeeeee do_deee :k]
11: + concat([x, new vaw_ealues], dim=l) & Add new values to the sequence length dimension
12: decoded decotie__ssstioons +- concat[[decoded pocoiioossttsss. pos to decode), dim=1)
13: positions to decode positions decode[:. k:]
Remove the k decoded positions
14: end for
15: out reorder(x, decoded_position
Sort based positions
16: return out
19

---

Published as a conference paper at ICLR 2026
Algorithm 3 MDLM-eqqivvlent sampling for PGMs.
1: Input: Batch size BS, number steps K, model length L, special BOS index
2: Output: Generated samples x
3: X < empty_ter emitt__tenoor(S, 1)
Initialize Set BOS as first token b Number of tokens to denoise at each step Keep track of clean and noisy positions 7: concrete_lengthe + ones(BS, 1) > Keep track of the actual length of each sequence (some are padded).
4: x(:, 0] <- BOS
5: k <- I/K
6: clean caait_posttooss zeros(BS,
8: joisy_postiions +- 1+ rand_ row Www_permBB,BB. L-1)
9: for range(K)
10: n_denoise_per_see
mussypossiippus
saleeeooooo(oosooooooeeee k)
Algorithm
Algo. 4
11:
values pgm predict(x, clean positiooss noisy_pos pac__pposspuuttt input)
12: x, ddean_positions, coiiy_positions concrete lengths +- extract preeiiciere((((
13:
Algorithm Algo.
14:
clean_positions noisy_positions,
16:
noisy_pos__ s___ntt,,
concrete lengths.
18:
___enoise__pes Bacis__pe________________ _eq,
new_values,
20: end for
out <- reorder(x, clean_ )
b Sort based on clean clean_leeeiieeeee positions
22: return out
Algorithm 4 Sample the number tokens to denoise from a binomial distribution and pad the input.
1: Input: Noisy positions tensor, probability of denoising prob deo_denoise, model length L, concrete lengths tensor
2: Output: Noisy positions to denoise
3: denoise Laciis____e______ per_seq < binomial(BS, L. prob denoise)
Sample from binomial distribution
4: n denoise lar__ee_sse___ <- min(n denoise_per, seq, L concrete lengths)
Don't t denoise more
than available
5: denoise_sed _etee_seeeeeeee_eeeeeeeeeeeeeeee +- max(n_ max(n_denoise_per___no.oo wew_dddee_ee_sseeee 0)
> Maximum number of tokens to denoise
6: if denoise_sedp seeeesseeleeeeeeeeeeeeeeeeeeeeee then
7: return empty_
Nothing to denoise
8: end if
9: noisy_por aappposspuuuuupp + noisy positioositi :denoise_see seeee_se___e__________e_eeeeeeee
Some predictions won't 't be used
return denoise Secnoise__errseee_ per_seq. noisy_ _aii_posss____ pos_input
20

---

Published as a conference paper at ICLR 2026
Algorithm 5 Extract the correct number of predictions per sequence
1: Input: concrete lengths, denoise_perr _er__eee- denoised_ values, clean positions
cossy_postiion,, noisy_pir _asssypps_______ jessinnputtt
2: Output: Updated x, clean blean_positions,
concrete_lengths 3: new new_concrete__ieeee concrete_lengths <- concrete_lengtiee + r__denosse_per_se
Update sequence lengths > Calculate padding needed
4: tok totto_add add <- feox(newwconcett_uueeee lengths) shape(x,
5: if n tok_to_add add > then
6:
pad <- zeros(BS, j__oo_te tok_to_add)
Create padding tensor Pad the sequences Pad the positions
7: X 6- concatxx pad, dim=1)
_rcan_positions +- concat(clean positions, pad, dim=l) 9: end if
8:
10: for in range(BBS) do
11: if' n_ dennoise_eer_aeeeeeeeeeee per__e___pp______ per_see[ii== = 0 then
continue
> Skip if tokens denoise
13: end if
x(i, concrete (earrrte_lengthssiimeeew [ea_l]/eett_ceeeeeee_e_e__ concreteeeengttieeeeeee lengths[i]] f-
concrete
15:
poooised_tokenn Ltt__ttookeeeeelltte values[i, in mo_denoise_per. per_seelli]
16:
clean_positions// concrete concrete_lengthene mangghssijnew_ccnnrreeeeeo
+
17:
noisy secy__oosiiputii :n lel_dddenneeeereeeeee (eatoiise_pe__sq_i]]]
18:
noisy coiyy_ysstionn(i, ishape positions, 1) n L_l_ddinn__eeere per per_seq[i]]
19:
sossy_positionsi. denoise Bacl__pe______ ver_sc][[]]]
20: end for
21: return x, clean__eeaate postttoostti noisy_positions, new_concrete__enngnn lengths
Model (LMIB)
Val. PPL +
Model (OWT)
Val. PPL +
200k steps
MDLM
34.29
200k steps
MDLM (Compl. masking)
30.87
MDLM
25.35
PGM 8 8/4
32.83
MDLM (Compl. masking)
25.32
PGM 10 2
33,55
PGM 6 6
26.96
PGM 4/8
32.84
PGM 8 8/8
25.10
PGM 6/6
32.69
PGM 10/6
25.19
PGM 6/ / (Ism)
32.70
PGM 6 6 (dim. 1024)
23.75
PGM 6 6 (mean)
33.89
IM steps
1M steps MDLM
MDLM MDLM (Compl. masking)
23.07
27.67
22.98
MDLM (Compl. masking)
25.72
PGM 8 8/8
22.61
PGM 6/6
26.80
PGM 6 6 (dim. 1024)
21.43
Table 5: Perplexity evaluations. Validation perplexity of the Masked Diffusion Language Model (MDLM) and PGMs (ours) on LMIB and OpenWebText (OWT). The row MDLM (Compl. masking) denotes MDLM trained with the complementary masking strategy discussed in Section 5.3. The row PGM k / m denotes a PGM with & encoder and m decoder layers, and we highlighted best PGM results in gray. Ism and mean denote the logsumexp and mean queries initializations (Section 4). Takeaway: using the same number layers in the encoder and decoder, and iata-independent queries performed best. On LMIB, PGM reaches 1.95 lower perplexity than MDLM after IM steps. On OWT, we grow the embedding dimension or the number of layers to outperform MDLM on OWT.
21

---

Published as a conference paper at ICLR 2026
Table Sample quality and efficiency on Open Teb with different numbers of sampling steps. We generate sequences of 1024 tokens with batch size of 32 to measure the latency and throughput. PGM 6 / 6 with a hidden dimension of 1024 and uniform sampling achieves at least 5x latency and throughput improvement over MDLM, with better Generative Perplexity and matching entropy.
Model
Gen. PPL +
Entropy
Latency + (ms)
Throughput f (tok/s)
MDLM
32 steps
192.31
5.73
8.037
0.01
4'077.08 3.06
64 steps
142.58
5.69
15.82 + 0.01
2'700.67 + 0.69
128 steps
122.89
5.67
31.41 0.01
1'043.22 = 0,16
256 steps
113.96
5.66
6..54
523.90 = 0.06
512 steps
109.05
5.64
1.4.94
262.26 + 0.33
1024 steps
106.75
5.64
249.31 0.11
131.42 + 0.05
PGM 8 (uniform sampling)
32 steps
189.02
5.73
1.55 0.01
64 steps
143.79
5.69
3.00
10'914.91
41.69
128 steps
122.21
5.66
5.86
5'585.57 + 24.49
256 steps
112.48
5.65
11.64
2'814.99 9.33
512 steps
(08.76
5.64
22.98 = 0.02
1'425.89 = 1.61
1024 steps
107.03
5.63
45.84 + 0.03
714.71 0.50
PGM 8 / 8 (non uniform sampling)
32 steps
194.09
5.73
2.07 0.02
15'764.09 192.12
64 steps
143.60
5.69
3.90 0.07
84405.14 158.01
128 steps
124.38
5.67
7.41 4.008
4'419.77
53.27
256 steps
116.85
5.66
14.73
2'223.63 + 28.47
512 steps
111.11
5.64
28.15
0.32
1'163.79
13.25
1024 steps
108.24
5.63
54.62 + 0.66
599.97 +
7.27
PGM / (dim. 1024, uniform sampling)
32 steps
185.16
5.73
1.59 + 0.01
90'569.99
64 steps
138.87
5.70
3.03 0.01
10'805.31 14.11
14.11
128 steps
116.95
5.67
5.93 - 0.01
5'518.09 + 13.46
256 steps
108.51
5.65
11.77 0.01
2'782.78 3.46
512 steps
101.94
5.63
23.25 + 0.01
1'408.88 + 1.05
1024 steps
99.64
5.62
46.31
0.02
707.52 - 0.34
PGM 6 6 (dim. 1024, non-uniform sampling)
32 steps
191.30
5.74
2.12 0.07
15'415.56 467.20
64 steps
138.67
5.69
3.940 0.06
+18.721135.44
128 steps
118.17
5.67
7.60
4/311.80
54.92
256 steps
108.93
5.65
1.20
2'207.71
29.71
512 steps
105.41
5.64
28.56
1'147.17
13.47
1024 steps
102.93
5.62
55.50 0.336
590.37+
3.85
22

---

Published as a conference paper at ICLR 2026
Table 7: Generative perplexity MDLM and PGM after distillation with varying precision.
Model
Gen. PPL
Entropy 1
Latency (ms)
Throughput f
(tok/s)
MDLM SDTT (loss in BF16)
32 steps
66.26
5.49
8.037 + 0.01
4'077.08
64 steps
53.98
5.46
15.82 + 0.01
2'070.67 0.69
128 steps
48.02
5.44
31.41 0.01
1'043.22
256 steps
45.86
5.42
62.54 0.01
523.90 0.06
512 steps
44.21
5.40
124.94 0.16
262.26 0.33
1024 steps
43.19
5.38
249.31
+ 0.11
131.42 # 0.05
MDLM + SDTT (loss in FP32)
32 steps
61.65
5.46
8,037
0.01
4'077.08 1000088:0 3.06
64 steps
50.65
5.43
15.82 0.00
2'070.67 0.69
0.69
128 steps
45.06
5.40
#1.01
1'043.22
256 steps
41.70
5.37
6.004
523.90 0.06
512 steps
40.63
5.36
124.9400.10
262.26 + 0.33
1024 steps
39.50
5.32
249.31 0.11
131.42 0.05
PGM 6/6 (dim. 1024)
SDTT (loss in BF16)
32 steps
91.61
5.56
20'569.99 95.63
steps
72.73
5.52
3.03 0.00
10'805.31 14.11
128 steps
63.83
5.49
5.93 0.001
5'318.09 + 13.46
256 steps
58.74
5.47
11.77
2'782.78 3.46
512 steps
58.77
5.47
23.25 0.00
1'408.88 1.05
1024 steps
56.47
5.46
46.31 + 0.02
707.52 = 0.34
PGM 6 (dim. 1024) nucleus (p=0.90
(p=0.9)+
9)+ SDTT
(loxs BF16)
32 steps
68.33
5.50
1.74 4: 0.01
18'866.12 1: 18.35
64 steps
53.88
5.45
3.18 + 0.01
10'307.16 +: 6.58
128 steps
16.99
5.42
6.10 + 0.01
5'375.20
2.40
256 steps
13.22
5.40
11.95
2'742.74
1.32
512 steps
12.79
5.39
23.63 0.01
1'386.79 + 0.69
1024 steps
40.99
5.38
46.83 + 0.02
699.80 + 0.24
PGM 6/6 (dim. 1024)
SDTT (loss in FP32)
32 steps
84.97
5.52
1.74 0.00
20'569.99 + 95.63
steps
67.60
5.49
3.18
10'805.31
128 steps
60.06
5.47
6.10 0.001
5'518.09 + 13.46
256 steps
55.97
5.45
11.95 0.01
2'782.78 + 3.46
512 steps
54.13
5.44
23.13
1'408.88
1.05
1024 steps
52.77
5.44
46.83
707.52 0.34
PGM 6 (dim. 1024) nucleus (p=0.9)+
SDTT
(loss FP32)
32 steps
63.46
5.45
1.59 1 0.01
18'866.12
18.35
64 steps
49.94
5.41
3.03 1: 0.01
10'307.16 1 6.58
128 steps
43.84
5.39
5.93 0.00
5'375.20 1 2.40
256 steps
40.76
5.36
11.77 0.01
2'742.74
1.32
512 steps
39.46
5.36
23.25 1 0.01
1'386.79 1 0.69
1024 steps
38.81
5.35
46.31 +: 0.02
699.80 1 0.24
PGM 8 8/8 + SDTT (loss in BF16)
32 steps
102.64
5.54
1.55 0.01
21 21'120.99 83.59
64 steps
82.93
5.50
0.001
10'914.91 41.69
128 steps
73.19
5.48
5.86 0.00
5'585.57 24.49
256 steps
70.30
5.47
11.64 0.00
2'814.99
+ 9.33
512 steps
68.07
5.46
22.98 0.02
1'425.89 + 1.61
1024 steps
65.87
5.44
714.71 0.50
PGM 8/8 SDTT (loss FP32)
steps
87.64
5.51
1.55
21'120.99 83.59
64 steps
70.47
5.48
3.00
10'914.91 + 41.69
128 steps
62.66
5.46
5.8640.00
5'585.57 24.49
256 steps
59.38
5.45
1..03
2'814.99
9.33
512 steps
57.57
5.44
22.98 1.002
1'425.89
1.61
1024 steps
56.12
5.44
45.84 0.00
714.71 + 0.50
23

---

Published as a conference paper at ICLR 2026
Table Sample quality and efficiency on ImageNet for different numbers of sampling sleps using the Confidence-basee sampler. We generate images in batches of 32 to measure throughput, and use a batch size of measure latency. Throughput is lower with CFG because each step requires two forward passes (conditional and unconditional). The throughput and latency are averaged over 10 batches.
Model
FID
IS
Latency (ms)
Throughput (img/s)
MaskGIT (32 steps; 458M)
W = 14.30
82.41
0.70
2.05
w 7.80
151.62
1.21
1.05
W = 2
6.78
208.92
1.21
1.05
w 7.37
255.69
1.21
1.05
w 4
7.46
289.93
1.21
1.05
w 9.61
250.86
1.21
1.05
= 6
21.68
149.61
1.21
1.05
MaskGIT (64 steps; 458M)
W =0
15.62
79.23
1.39
1.03
W 9.40
140.48
2.41
0.52
W = 8.06
195.35
2.41
0.52
w 8.19
239.89
2.41
0.52
w 4 7.61
267.26
2.41
0.52
w 11.44
202.92
2.41
0.52
W 26.41
113.63
2.41
0.52
PGM 12 / 12 (32
steps; 464M)
W = 18.77
67.22
1.04
6,86
W 8.96
135.03
1.08
3.76
W = 6.67
201.66
1.08
3.76
W 7.09
255.43
1.08
3.76
W 4 8.30
290.18
1.08
3.76
W = 9.59
307.52
1.08
3.76
W 10.84
313.27
1.08
3.76
PGM 12 / 12 (64
steps; 464M)
W = 19.45
64.44
2.04
3.52
W 10.08
124.90
2.04
1.90
W = 7.35
188.77
2.04
1.90
W 7.39
238.31
2.04
1.90
W 4 8.13
276.55
2.04
1.90
W =
9.18
297.44
2.04
1.90
W 10.38
302.29
2.04
1.90
PGM 14 / TO (32
steps; 464M)
W =
21.97
60.07
1.04
7.09
W 11.24
121.39
1.04
3.90
W =
8.05
183.20
1.04
3,90
W
7.76
232.62
1.04
3,90
W = 4
8.47
263.73
1.04
3.90
W = 9.39
288.60
1.04
3.90
W
10.40
291.46
1.04
3,90
PGM 14 / 10 (64
steps; 464M)
W = 22.74
56.77
2.03
3,63
W 12.32
112.98
2.04
1.97
W = 2
8.80
171.11
2.04
1.97
W
8.11
219.33
2.04
1.97
W = 4
8.44
253.10
2.04
1.97
W 6
9.12
270.16
2.04
1.97
9.90
279.00
2.04
1.97
24

---

Published as a conference paper at ICLR 2026
Table 9: Sample quality and efficiency on ImageNet for different numbers of sampling steps using the Halton sampler. We generate images in batches of 32 to measure throughput, and use batch size of to measure latency. Throughput is lower with CFG because each step requires two forward passes (conditional and unconditional). The throughput and latency are averaged over 10 batches.
Model
FID
IS f
Latency
Throughput (img/s)
(ms)
MaskGIT (32 steps; 458M)
W =
25.72
57.70
0.70
2.02
5.35
267.49
1.21
1.05
W = 12.82
365.36
1.21
1.05
= 17.24
408.30
1.21
1.05
w = 15.65
365,97
1.21
1.05
25.33
182.14
1.21
1.05
W = 6
48.97
74.74
1.21
1.05
MaskGIT (64 steps; 458M)
W = 18.61
69.68
1.39
1.03
1 6.76
283.96
2.41
0.52
W = 14.97
372.74
2.41
0.52
= 3 18.30
410.60
2.41
0.52
W = 16.18
312.75
2.41
0.52
32.69
126.69
2.41
0.52
W = 60.12
51.40
2.41
0.52
PGM 12 / 12 (32
steps; 464M)
W = 22.59
66.32
1.03
13.44
10.21
134.22
1.03
7.93
W =
6.04
203.28
1.03
7.93
33 5.54
263.53
1.03
7.93
W 4 6.38
311.53
1.03
7.93
7.58
345.18
1.03
7.93
W = 8.83
372.07
1.03
7.93
PGM 12 / 12 (64
steps; 464M)
W = 16.19
79.26
1.96
7.17
1 6.67
151.44
2.01
4.12
W = 4.56
218.98
2.01
4.12
=
4.98
276.16
2.01
4.12
W 4 6.47
322.53
2.01
4.12
7.95
352.56
2.01
4.12
W = 9.47
379.39
2.01
4.12
PGM 14 / TO (32
steps; 464M)
W = 25.42
62.11
1.01
12.70
11.57
128.22
1.00
7.37
W = 6.60
196.56
1.00
7.37
= 5.55
253.76
1.00
7.37
w = 4 6.00
301.66
1.00
7.37
W
7.04
334.29
1.00
7.37
W = 8.16
365.56
1.00
7.37
PGM 14 / 10 (64
steps; 464M)
W = 18.03
75.40
1.93
6.69
7.77
144.26
1.99
3.80
W = 4.76
212.39
1.99
3,80
3 4.85
268.55
1.99
3.80
W
5.88
309.88
1.99
3,80
W
7.23
342.40
1.99
3.80
W =
8.63
366.57
1.99
3.80
25

---

Published as a conference paper at ICLR 2026
Table 10: Throughput (TP) of MDLM and PGM with a context length of 4096, for varying number of inference steps. PGM is significantly faster than MDLM.
Model
TP (4096)
TP (1024)
TP (256)
TP (64)
MDLM
30.45 0.06
121.25 0.02
483.53 + 0.25
1'912.16 + 1.44
PGM 8/8
128.99 + 0.23
697. 697.36 = 32.83
2'216.91 + 3.06
8'203.82 + 6.60
PGM 6/6 (dim=1024)
129.01 + 0.67
706.65 + 36.23
2'146.60 + 15.12
8'175.69 + 7.85
26


Reviews:
[{"review_id": "LabNEsk09h", "reviewer": "Reviewer_YDst", "summary": "The authors introduce Partition Generative Models (PGMs), based on the observation that masked generative models (MGMs) waste compute on masked tokens, which contain no information. \nInstead of masking tokens, PGMs partition the input tokens into two disjoint groups and train the model to predict one group from the other.\nThis approach allows the model to process only unmasked tokens which eliminating the need for explicit masking and leads to significantly faster sampling.", "questions": "- How does the choice of partition ratio (t) affect convergence and quality? Is it dynamically sampled or fixed?\n- why cant you use KVcache that would reduce the time complexity from sampling in MGM? Would PGM will be still faster?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The GroupSwap layer and partition-aware transformer structure are well-motivated\n- Includes analyses of perplexity, latency, throughput, and ablations on masking vs. partitioning.\n- Strong empirical results across both text and image generation tasks, PGMs deliver substantial inference speedups (up to 7×) with little to no degradation in output quality.", "weaknesses": "- The architectural details (e.g., data-dependent vs. data-independent queries) are dense and could be clarified or simplified, the paper is a bit difficult to follow.\n- The largest experiments are modest in size (268M parameters). It remains unclear if PGMs scale favorably compared to state-of-the-art large AR or diffusion model\n- No comparison against recent SOTA model non-autoregressive language models beyond MDLM.", "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": "Ljcr4gb12z", "reviewer": "Reviewer_pTqP", "summary": "This paper introduces a new generative modeling framework for language modeling, termed Partition Generative Models (PGM), aimed at improving inference efficiency. Unlike Masked Generative Models (MGM), PGM avoids applying the forward process to masked tokens, thereby reducing computational cost. The authors present tailored architectural modifications, along with corresponding training and inference strategies, to enable efficient generation within this framework. Experimental results indicate that PGM achieves faster inference than existing MGM approaches while maintaining comparable generation quality.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The core idea of avoiding computation on masked tokens during inference, along with the corresponding training strategy, is interesting and effectively targets a key inefficiency in existing masked generative models.\n\n2. The empirical results demonstrate that PGM can significantly accelerate inference while maintaining generation quality comparable to other state-of-the-art generative models, supporting the practical value of the proposed approach.\n\n3. The paper is clearly written, well-structured, and easy to follow, making the technical contributions accessible to the reader.", "weaknesses": "I did not identify any major weaknesses in this paper. I do, however, have one question for clarification:\n\nThe proposed training pipeline includes two prediction components that operate on the same batch of data, which suggests that training efficiency could potentially be better than MDLM. Could the authors provide quantitative results or analysis regarding training efficiency, such as training speed, computational cost, or resource usage compared to MDLM?", "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": "MyKtVVbFCT", "reviewer": "Reviewer_DktD", "summary": "This paper proposes avoiding the repetitive computation of the `MASK` token in masked generative modeling (MGM), switching from the decoder-only MGM to an encoder-decoder architecture. The inference model is defined as:\n1. self-attention **only within known indices**.\n2. cross-attention swapping to unknown indices (with opposite-group masking to prevent leakage).\n3. cross-attention **only within unknown indices** (projecting to the embeddings of stage one).\n\nThe complexity is still O(L^2) (L = sequence length) but with a significantly smaller coefficient (encoder $(L-k)^2$ + decoder $k(L−k)$ per step when remaining sequence length = k). \nThe practical speedup is about 5x and is scalable, with comparable generation quality against MGM.\nDistillation-accelerated models maintains the acceleration against MGM.", "questions": "- Except for the top-k/nucleus confident tokens, the computations are wasted. I wonder if it is possible to reuse these noisy states instead of re-initializing decoder queries at each denoising step?\n\n- The current decoder architecture is cross-attention-only—which makes it easy to control parameter count c.f. MDLM, but lacks the standard self-attention component. Have you thought about this variant?\n\n- The information exchange from known to unknown indices entirely relies on the swap xattention layer. I wonder if it is possible to do the exchange in each decoder layer instead? (Of course this will make the complimentary masking trick not possible.)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The empirical benefit is strong: 5x faster than MGM (4.6x faster with nucleus sampling).\n\n- Complementary masking is a smart and original trick to let one training step effectively count as two steps.\n\n- Section 5.3: fair comparison against MDLM (MGM) by isolating the complementary masking trick.\n\n- The down-stream tasks spreads across image and language, and the evaluation is solid. Distillation is also explored, which improves the practical significance of the paper.", "weaknesses": "- The fairness of Table 2's comparison is not immediately visible—I believe the fairness should outweigh matching performance. Since the paper switches from decoder-only to encoder-decoder architecture, controlling hyperparameters (width, head, depth and MLP width multipliers) seems crucial to get a fair comparison.\nIn LM1B, it is a good idea controlling parameter counts and comparing with PGM(6/6)\\~170M, but in OWT, that model is missing in the main text (only the dim. 1024 model is shown). I don't understand why it only appears in the appendix.\n\n- I don't understand the labels (5.3) (5.4) (5.5) in Figure 4 (right).\n\n- Minor: \"sparse attention\" is used to describe the masking mechanism, but I believe it is an overuse of the term, as the mask is not actually sparse—perhaps group-wise attention is more suitable.", "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": "LP1ol2z26W", "reviewer": "Reviewer_m818", "summary": "This paper introduces the PGM, a new framework with core architecture innovation that combines the strengths of AR and MGM. PGM partitions tokens into two disjoint groups and constrains attention such that each group predicts the other. This removes the need for explicit MASK tokens while preserving parallel decoding and arbitrary generation order and address the training inefficiency of MGM.\nExperiments on LM1B, OpenWebText, and ImageNet show that PGMs achieve up to 5–7× faster inference throughput than MDLM and MaskGIT, with similar or better metrics including perplexity and FID. PGMs also support distillation for additional speedups.\nThis work is overall sound to me, but I am not an expert in architecture design.", "questions": "1. Does \"PGM 8 / 8\"  mean 8 layers of encoder and 8 layers of decoder?\n2. The origin of the training instability. Do authors still observe this when PGM is trained without complementary masking.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. Good novelty: replacing masking with partitioning is a simple yet powerful idea that effectively unifies the efficiency of AR models with the flexibility of MGMs.\n2. Solid architectural design: The GroupSwap mechanism and partition-wise attention are well-motivated and carefully engineered to achieve the partition mechanism.\n3. The experimental results are strong, with improved performance for both text and image generation.", "weaknesses": "1. While the PGM is motivated be the inefficiency of MDM training, the authors are encouraged to provide evidence to show faster learning/convergence of PGM than MDLM. This probably relates to the training stability.", "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 (Oral)