Source-linked AI summary
Promises and Pitfalls of Generative Masked Language Modeling: Theoretical Framework and Practical Guidelines
Yuchen Li, Alexandre Kirchmeyer, Aashay Mehta, Yilong Qin, Boris Dadachev, Kishore Papineni, Sanjiv Kumar, Andrej Risteski
TL;DR
Autoregressive generation requires sequential decoding, while the paper investigates how to characterize the learning and inference properties of GMLMs. It develops theory relating statistical complexity to Markov-chain mixing and adapts T5 for iterative-refinement parallel decoding. Larger masks improve statistical efficiency, and parallel decoding can produce reasonable machine translations in one pass.
Problem
Autoregressive generation requires N sequential decoding steps, and the paper seeks a framework for characterizing GMLM learning and inference.
Source: §1 · pp. 1–2Method
The paper relates masked-prediction statistical complexity to Markov-chain mixing and adapts T5 for iterative-refinement parallel decoding.
Source: §1 · p. 2, §5 · p. 13Results
Larger masks improve statistical efficiency, while parallel-decoding GMLMs often produce reasonable machine translations in one pass.
Source: §1 · p. 2Takeaways & Limitations
The authors offer rules of thumb for GMLM design and connect empirical findings with the theoretical framework.
Source: §5 · p. 13Larger-block samplers can mix faster, but parallel inference may find it harder to re-randomize those coordinates from their conditional distribution.
Source: §2.4 · p. 8
Abstract
from arXiv · showhide
Autoregressive language models are the currently dominant paradigm for text generation, but they have some fundamental limitations that cannot be remedied by scale-for example inherently sequential and unidirectional generation. While alternate classes of models have been explored, we have limited mathematical understanding of their fundamental power and limitations. In this paper we focus on Generative Masked Language Models (GMLMs), a non-autoregressive paradigm in which we train a model to fit conditional probabilities of the data distribution via masking, which are subsequently used as inputs to a Markov Chain to draw samples from the model, These models empirically strike a promising speed-quality trade-off as each step can be typically parallelized by decoding the entire sequence in parallel. We develop a mathematical framework for analyzing and improving such models which sheds light on questions of sample complexity and inference speed and quality. Empirically, we adapt the T5 model for iteratively-refined parallel decoding, achieving 2-3x speedup in machine translation with minimal sacrifice in quality compared with autoregressive models. We run careful ablation experiments to give recommendations on key design choices, and make fine-grained observations on the common error modes in connection with our theory. Our mathematical analyses and empirical observations characterize both potentials and limitations of this approach, and can be applied to future works on improving understanding and performance of GMLMs. Our codes are released at https://github.com/google-research/google-research/tree/master/padir
1 Introduction
Autoregressive generation is sequential and cannot use right-hand context to revise earlier predictions; GMLMs instead fit masked conditionals and use them in Markov-chain sampling. The paper develops a framework for GMLM learning and inference and investigates where parallel decoding works and where it struggles.
- Training with larger masks always improves statistical efficiency, and the framework relates masked-prediction sample complexity to the mixing time of a corresponding sampling chain.Source: §1 · p. 2
- Finite-sample bounds connect how well conditional distributions are learned to how well the joint distribution is learned, given capacity control such as covering-number bounds.Source: §1 · p. 2
- Transformers’ coordinate-factorized decoding steps can impede sampling distributions with strong correlations, while stuttering is a common error mode associated with those dependencies.Source: §1 · p. 2
- Experiments identify large masking ratios, custom vocabulary, distillation from autoregressive models, and positional attention as empirically critical components.Source: §1 · p. 2
- Parallel-decoding GMLMs work well on machine translation, where even one forward pass can often produce reasonable translations.The authors relate this result to translation outputs being lower-entropy and less multimodal than outputs in other language-generation tasks.Source: §1 · p. 2
- The combined findings motivate co-designing Markov chains that mix quickly under strong correlations and losses that retain good statistical behavior.Source: §1 · p. 2
2 Theoretical framework
The theoretical framework relates GMLM learning complexity and finite-sample error to inference-time Markov-chain mixing, including the effects of mask choice and adaptive masking. It also analyzes how parallel transition constraints and dependent block updates affect decoding, particularly in strongly dependent settings.
2.1 Setup and notation
- Masked prediction fits conditionals of selected variable subsets; smaller subsets can reduce computation, while maximum likelihood serves as the sample-efficiency benchmark.For discrete data, matching these conditionals gives the pseudolikelihood objective; the paper notes that conditional computation can be cheaper than computing the full likelihood.Source: §2.1 · p. 3, Setup and notation · pp. 3–4
2.2 Asymptotic sample efficiency via functional inequalities
- The asymptotic framework bounds sample complexity for broad masked-prediction objectives and assesses estimator quality through parameter recovery.It relies on asymptotic normality under mild technical conditions.Source: §2.2 · p. 4
2.2.1 Masking more is (statistically) better
- Increasing the predicted subset size k strictly improves k-pseudolikelihood statistical efficiency, while increasing optimization cost; k = N recovers maximum likelihood.Under the stated assumptions, the asymptotic variance improves with larger k, yielding a better asymptotic l2 bound for learning θ.Source: §2.2.1 · pp. 4–5
2.2.2 Statistical efficiency bounds via mixing time bounds
- A Poincaré constant for the corresponding weighted block dynamics bounds weighted MPLE asymptotic variance by C times the inverse Fisher information.The block dynamics corresponds to the Markov chain used for inference, linking statistical efficiency to its mixing behavior.Source: §2.2.2 · p. 6
2.2.3 Adaptive masking: masked positions depend on the sequence
- The asymptotic-variance bound extends to adaptive masking, where mask probabilities depend on the current sequence, using the corresponding adaptive block dynamics.Under the stated assumptions and a Poincaré inequality with constant C, adaptive MPLE variance is bounded by C times the inverse Fisher information.Source: §2.2.3 · pp. 6–7
2.3 Finite sample bounds and distributional distance
- Finite-sample analysis replaces asymptotic parameter closeness with distributional distance and replaces the Poincaré constant with approximate tensorization and model-class complexity.The paper states a generalization bound for learning the joint distribution under its support-margin and covering-bound/Lipschitz assumptions.Source: §2.3 · pp. 7–8
- The distributional bound compares total variation from p with a CAT constant and class complexity, replacing the parameter-distance, Poincaré, and Fisher-information terms in the earlier bound.The complexity term captures generalization error through C_ε(Θ).Source: §2.3 · p. 8
2.4 Inference-time limitations due to parallelism
- Larger update blocks can mix faster, but parallel re-randomization becomes harder as more coordinates must be sampled jointly; independent parallel updates can also be slower to reach modes than large-k Gibbs updates.The section identifies this as a limitation of parallel decoding for GMLM inference.Source: §2.4 · pp. 8–9
2.4.1 Which Markov Chains are implementable via parallel decoding?
- Transformers with sufficient depth and width can implement deterministic Markov-chain transitions, but parallel decoding implements only transitions whose next-state distributions factor across positions.This restriction does not preclude a non-product stationary distribution: under mild conditions, 1-Gibbs updates have p as a stationary distribution, whereas k-Gibbs transitions for k > 1 are generally not implementable.Source: §2.4.1 · p. 9
2.4.2 Markov Chains with dependent transitions can be (much) faster
- In the strongly dependent Ising-model setting, k-Gibbs can reach the high-probability mode R1 in one step while independent parallel sampling may fail to reach it for arbitrarily many steps.Jointly updating dependent variables enables faster mode-finding; reaching R1 represents finding a high-probability sample, not necessarily mixing to the full distribution.Source: §2.4.2 · pp. 9–11
3 Experiments
Experiments find that iterative parallel decoding can combine similar translation quality with lower latency, while exposing stuttering as a frequent error. Synthetic and attention analyses test links between masking, statistical efficiency, and target-position dependencies.
3.1 Synthetic experiments on Ising model
- With the same training data size, larger k in Ising-model k-pseudolikelihood yielded lower parameter-estimation error, validating Theorem 1.The experiments used synthetic data from a ground-truth Ising model.Source: §3.1 · p. 11
3.2 Parallel Decoding by Iterative Refinement (PaDIR)
- PaDIR uses an encoder-decoder whose decoder updates all target positions in parallel, iteratively refining a random initial target hypothesis.Each additional decoder forward pass provides another refinement step.Source: §3.2 · p. 12
3.3 Evaluation
- Word-level stutter was 2 or more times more frequent for non-autoregressive models across all datasets, motivating evaluation beyond BLEU's n-gram overlap.The authors introduce stutter as repeated consecutive output words absent from the reference, noting BLEU does not capture readability well.Source: §3.3 · p. 12
- >2x median and >5x 99th-percentile end-to-end speedups accompanied similar BLEU and BLEURT scores for PaDIR versus autoregressive models.Speed measurements used TPU v3, four decoding steps, and batch size 1.Source: §3.3 · p. 12
- Four decoding steps require 7 to 8 times fewer decoder passes for target lengths of 28 to 33 tokens, but speed depends on the task and some in-painting is slower than autoregressive decoding.Fixed costs and optimized autoregressive decoding help explain why observed end-to-end speedups are smaller than the decoder-pass reduction.Source: §3.3 · p. 12
3.4 Connecting to theory: quantifying dependency via attention scores
- Stuttering positions did not show larger average decoder self-attention, while their cross-attention distributions to source tokens were more similar than those of non-stuttering adjacent positions.The first finding makes Hypothesis 1 unlikely; the second makes Hypothesis 2 potentially promising.Source: §3.4 · p. 13
4 Related works
The related-work discussion situates the paper's theory and experiments alongside sampling, iterative parallel decoding, and discrete diffusion research.
- The theory draws on connections between pseudolikelihood and approximate tensorization of entropy, including Marton (2013; 2015), Caputo et al. (2015), Caputo & Parisi (2021), and Koehler et al. (2023).Source: §4 · p. 13
- Prior iterative-refinement parallel decoding work is described as at least twice faster than autoregressive approaches with a small quality drop on tasks such as machine translation, including Lee et al. (2018), Ghazvininejad et al. (2019; 2020), Kasai et al. (2020), and Savinov et al. (2022).Source: §4 · p. 13
- The paper distinguishes its masked-autoencoder-like training, which reconstructs the whole target each decoding step, from diffusion training that reverts a small amount of noise per step, as in Austin et al. (2021), Li et al. (2022), Gong et al. (2023), Zheng et al. (2023), and Lou et al. (2023).Source: §4 · p. 13
5 Conclusion
The paper develops theory and empirical guidance for the potentials and limitations of generative masked language models. It adapts T5 for iterative parallel decoding and connects practical design recommendations to theoretical and empirical findings.
- The paper introduces a framework for analyzing GMLM learning and inference through parameter sample complexity, finite-sample distribution generalization, and Gibbs-like sampling efficiency.Source: §5 · p. 13
- The authors adapt T5 to parallel decoding by iterative refinement, report practical rules of thumb for design choices, and relate empirical observations to theory.Source: §5 · p. 13
- The authors hope the framework and observations will inform future training objectives, inference algorithms, and architectures suited to parallel decoding.Source: §5 · p. 13
Supplementary Material
The supplementary material provides proof details for information-matrix and masking results, then extends the framework to adaptive masking and block dynamics.
A Proof of Lemma 2: Generalized information matrix equality
- Under Assumptions 1 and 2, Lemma 2 states that the weighted pseudolikelihood loss satisfies a generalized information matrix equality.Its proof exchanges expectations and derivatives before relating the Hessian terms to the covariance expression.Source: A Proof of Lemma 2: Generalized… · pp. 24–25
B Proof of Theorem 1: Masking more is (statistically) better
- The proof of Theorem 1 uses a positive-semidefinite term and monotonicity of the matrix inverse to establish the statistical benefit of larger masks.Source: B Proof of Theorem 1: Masking more is… · p. 27
C.1 Conditioning on K
- With adaptive masking, the conditional distribution of masked variables given the unmasked variables and mask can differ from the conditional without the mask; the appendix gives a counterexample and a marginalization identity.Source: §C.1 · p. 28
C.2 Information matrix equality for adaptive masking
- Under Assumptions 1 and 2, Lemma 5 extends the generalized information matrix equality to the weighted pseudolikelihood loss for adaptive masking.Source: §C.2 · p. 29
C.3 Proof of Proposition 1: Dirichlet form for adaptive block dynamics
- Proposition 1 gives the Dirichlet form corresponding to adaptive weighted block dynamics.The proof expands the form using the dynamics' transition matrix and combines symmetric terms.Source: §C.3 · pp. 30–31
C.4 Proof of Theorem 3: Asymptotic sample complexity for adaptively-weighted MPLE
Theorem 3 bounds the asymptotic variance of adaptively weighted MPLE under a Poincaré inequality. Its proof relates the estimator’s variance to Fisher information.
C.4 Proof of Theorem 3: Asymptotic sample complexity for adaptively-weighted MPLE
- Theorem 3 generalizes Theorem 2; the proof gives only the more general result to avoid duplication.Source: §C.4 · p. 31
- The asymptotic variance of adaptively weighted MPLE is bounded by C I^-1 when the distribution satisfies the stated Poincaré inequality and assumptions.Here, C is the inequality constant for the adaptively weighted block dynamics, and I is the Fisher information matrix.Source: §C.4 · p. 31
- The proof relates a term in the variance bound to Fisher information using the Poincaré inequality and matrix-inverse monotonicity.It separately evaluates a conditional variance term and then applies monotonicity of the matrix inverse.Source: §C.4 · pp. 31–32
D Proof of Theorem 4: Generalization bound for learning the joint distribution
Theorem 4 provides a high-probability generalization bound for learning the joint distribution. Its proof connects conditional-distribution losses to joint-distribution distance and combines bounds over masks and training sequences.
D Proof of Theorem 4: Generalization bound for learning the joint distribution
- Under Assumptions 3 and 4, Theorem 4 gives a high-probability generalization bound on total variation distance for the learned joint distribution.Source: D Proof of Theorem 4: Generalization… · p. 33
- The proof controls empirical mask frequencies and sequence sampling, then combines the resulting bounds with a union bound and triangle inequality.The mask-frequency argument uses the convergence of observed mask probabilities as the number of sampled masks increases.Source: D Proof of Theorem 4: Generalization… · pp. 33–35
- The proof connects conditional-distribution loss to joint-distribution distance through approximate tensorization of entropy and Pinsker’s inequality.The conditional-distribution guarantee uses finite samples of sequences and masked positions.Source: D Proof of Theorem 4: Generalization… · p. 33
- The conditional KL terms are generally not independent across masks, a qualification relevant to the proof’s concentration arguments.Source: D Proof of Theorem 4: Generalization… · p. 35
- The bound’s mask-sampling and sequence-sampling terms decrease as their respective sample counts increase, while its first term grows with sequence length.The passage identifies more mask configurations per sequence and more training sequences as ways to reduce the corresponding terms.Source: D Proof of Theorem 4: Generalization… · p. 37
E Proof of Proposition 8: Modes of the strongly ferromagnetic Ising model
For the strongly ferromagnetic Ising model under Assumption 5, Proposition 8 characterizes the relative probabilities of two regions and states a pairing relation between them.
E Proof of Proposition 8: Modes of the strongly ferromagnetic Ising model
- Under Assumption 5, every state in R1 has greater probability than every state in R−1, whose probabilities exceed those outside both regions by the stated factor.Source: E Proof of Proposition 8: Modes of the… · p. 40
- There is a bijection from R1 to R−1 such that paired states have probabilities related by the factor e^2hG.Source: E Proof of Proposition 8: Modes of the… · p. 40
- The proposition’s claims concern the Ising model in Equation (6) under Assumption 5.Source: E Proof of Proposition 8: Modes of the… · p. 40
F Proof of Proposition 4: k-Gibbs sampler can reach the mode fast
Proposition 4 establishes that a k-Gibbs sampler can reach the mode region R1 within a stated number of steps with high probability when k is large enough to include CG.
F Proof of Proposition 4: k-Gibbs sampler can reach the mode fast
- For k ≥ |CG|, the sampler visits R1 within the stated T steps with probability at least 1 − δ, from any initial state.The supplied proposition statement does not specify the value of T.Source: F Proof of Proposition 4: k-Gibbs… · p. 42
- The proof considers updates whose selected coordinate set K contains all of CG, allowing that group to be resampled jointly.Source: F Proof of Proposition 4: k-Gibbs… · p. 42
- The proof uses Proposition 8 to show that one of its cases occurs with probability at least a constant.Source: F Proof of Proposition 4: k-Gibbs… · p. 42
G Proof of Proposition 5 independent parallel sampling stuck in bad samples
The proposition establishes that independent parallel sampling can become stuck in bad samples for the considered Ising model. Its proof controls coordinate-wise indicators and derives a probability bound for the event that the chain leaves a specified region.
G Proof of Proposition 5 independent parallel sampling stuck in bad samples
- Proposition 5 establishes that independent parallel sampling can become stuck in bad samples for the considered Ising model.The proof treats the coordinate indicators as independent Bernoulli variables and uses Hoeffding’s inequality to control their sum.Source: G Proof of Proposition 5 independent… · p. 43
- The proof aligns probabilities to bound when the chain lies outside R1 for a specified initial state and failure probability.The supplied proof passage states a probability of at least 1 − δ after a specified number of steps, but does not preserve the full bound.Source: G Proof of Proposition 5 independent… · pp. 43–44
H Proof of Corollary 1: Separation between N-Gibbs sampler and independent parallel sampling
The corollary separates N-Gibbs sampling from independent parallel sampling on the specified Ising model under stated assumptions. Its proof derives bounds for the two procedures and combines them with a union bound.
H Proof of Corollary 1: Separation between N-Gibbs sampler and independent parallel sampling
- Corollary 1 states a separation between N-Gibbs sampling and independent parallel sampling on the specified Ising model under Assumptions 5–7.The corollary also requires an initial-state condition and quantifies over δ and M.Source: H Proof of Corollary 1: Separation… · p. 45
- The proof bounds the N-Gibbs case using Proposition 4 and contrasts it with independent parallel sampling using Proposition 5.The supplied proof passages do not include the complete bounds.Source: H Proof of Corollary 1: Separation… · pp. 45–46
- A union bound combines the two probability statements to establish that both hold with probability at least 1 − δ.The proof refers to Equations (H.49) and (H.52).Source: H Proof of Corollary 1: Separation… · p. 46
I.1 Technical setup and proofs
This appendix formalizes the expressivity of parallel-decoding Transformers and supplies supporting results on prior methods, estimator behavior, convexity, and PaDIR’s training and inference setup. Its central boundary is that the model’s independently sampled positions permit product transitions, but not exact non-product transitions.
I.1 Technical setup and proofs
- Parallel-decoding Transformers can represent product transition distributions, but cannot exactly implement non-product transitions because position-wise sampling is independent.They can also simulate deterministic sequence maps for T steps; the appendix gives a Transformer construction with depth O(log T).Source: §I.1 · pp. 48–49
I.2 Connection to prior works in GMLM
- Prior iterative-refinement methods differ from both N-Gibbs and independent-parallel sampling, and their learned conditionals may lack a consistent joint distribution.The appendix notes that some methods can incorporate inter-position dependencies through intermediate Transformer representations.Source: §I.2 · pp. 49–50
J Regularity conditions for asymptotic behavior of M-estimators
- Under mild regularity conditions, the appendix states that M-estimators, including maximum likelihood and the specified estimators, converge in distribution to a normal distribution as samples become infinite.The accompanying lemma gives local uniqueness of an empirical minimizer for sufficiently large sample size.Source: J Regularity conditions for asymptotic… · p. 51
K Convexity of pseudolikelihood for Ising models
- For an Ising model, the weighted pseudolikelihood objective is convex in its parameters.The appendix includes this known property for completeness and connects it to efficient algorithms for bounded-degree Ising models.Source: K Convexity of pseudolikelihood for… · p. 52
L.1 Training and inference approach for PaDIR
- The PaDIR appendix provides a formal description of its training and inference strategy.The supplied passage introduces that description without further procedural detail.Source: §L.1 · p. 53
L.1.1 Inference
- At inference, PaDIR predicts candidate target lengths, initializes a target hypothesis, and iteratively refines it with one decoder forward pass per step.The procedure uses inference-time length-beam and refinement-step hyperparameters and can stop early if a step leaves the sequence unchanged.Source: §L.1.1 · p. 53
L.1.2 Training
- One-stage training does not expose later decoding steps to the model’s own earlier predictions, so the appendix uses a multi-stage training objective.The stated concern is that later-step inference inputs are out of distribution relative to one-stage training.Source: §L.1.2 · p. 53
L.2 Details of training recipe
- The recipe evaluates both directions of WMT14 DE↔EN, WMT16 RO↔EN, and WMT17 ZH↔EN, using a 67M-parameter encoder-decoder and two training stages.The recipe also uses autoregressive-model distillation data and remasks token-level stutter across iterations.Source: §L.2 · pp. 53–54
L.3 Discussion on modeling and metrics
The discussion and appendices pair translation-evaluation choices with empirical probes of attention, masked estimation, and sampling, then situate the framework among prior parallel-generation and sampling work.
L.3 Discussion on modeling and metrics
- The encoder-decoder design lets iterative target refinement reuse the encoder outputs and supports comparison with prior benchmarked approaches.Source: §L.3 · p. 54
- The paper uses SacreBLEU without preprocessing because preprocessing can affect scores and may be irreversible, while differing BLEU implementations hinder direct comparison with much prior work.Source: §L.3 · p. 54
L.4 Quantitative experimental results on machine translation task
- Translation appendices report Table 1 test SacreBLEU on three WMT datasets, Table 2 baseline BLEU with implementation and preprocessing differences, and Table 3 model BLEURT.Source: Table 1 · p. 55, Table 2 · p. 55, Table 3 · p. 55
L.5 Quantifying dependency via attention scores
- Stuttering adjacent pairs have comparable average last-layer self-attention, but more similar average last-layer cross-attention than non-stuttering pairs.Source: Table 4 · p. 56, Table 5 · p. 56
L.6 Masking more is statistically better for learning synthetic Ising models
- Larger k lowers parameter-estimation error at equal training-data size, including under the more peaked, two-mode Ising distribution.The experiments use k-pseudolikelihood on synthetic Ising data.Source: Figure 2 · p. 57, §L.6 · p. 57, Figure 1 · p. 57
L.7 Markov Chains with dependent transitions can be (much) faster in sampling Ising models
- With initialization in the smaller mode, k-Gibbs often reaches the larger mode within budget; larger k is faster, while the independent parallel sampler does not reach it within budget.More peaked distributions, corresponding to larger J, take longer to sample.Source: Figure 3 · p. 58, §L.7 · p. 58
M Additional related works
- The related-work discussion positions the framework alongside iterative parallel token refinement and distinguishes its improper-learning generalization bound from Koehler et al. (2023)’s proper-learning bound.Source: M Additional related works · p. 59