Masked diffusion models (MDMs) have emerged as a compelling alternative to autoregressive generation for discrete sequences, particularly in domains like protein design where fixed left-to-right ordering is an arbitrary constraint. Yet a persistent tension exists in how these models are decoded: sequential unmasking preserves quality but is slow, while naive parallel unmasking is fast but breaks structural consistency. This paper (arXiv:2605.20187) addresses that tension directly by asking whether a lightweight neural estimator can predict pairwise conditional mutual information (MI) from an MDM's hidden states, and whether that MI signal can guide smarter parallel decoding.
The core insight is clean: an MDM already encodes dependency structure implicitly in its hidden representations, but exposes only marginal conditional distributions at inference time. If you could read out pairwise MI from those hidden states cheaply, you could identify which tokens are conditionally independent and unmask them together without sacrificing coherence. The paper builds exactly this machinery.
Key Contributions
- A neural MI estimator trained as a lightweight prediction head on top of a frozen MDM backbone, outputting a full symmetric N x N pairwise MI matrix in a single forward pass.
- A supervision strategy that computes ground-truth MI directly from the pretrained MDM's own conditional distributions, meaning the estimator learns the model's internal belief about dependency structure rather than some external ground truth.
- MI-guided parallel decoding, which groups conditionally independent tokens for simultaneous unmasking, reducing the number of forward passes by roughly 3-5x compared to sequential decoding.
- Evaluation on two structurally distinct domains: Sudoku puzzles (hard combinatorial constraints) and protein sequences using ESM-Cambrian (300M parameters), demonstrating generality beyond a single toy setting.
Methodology
The estimator is architecturally minimal, a small head of roughly 100K parameters for Sudoku and 810K for proteins, attached to one of the final layers of the frozen MDM. Given hidden states h from a partially masked input, it predicts the full MI matrix in one pass. The symmetry constraint is enforced by construction.
Supervision comes from computing exact pairwise MI using the MDM's own marginal conditionals, which is tractable for small vocabularies but expensive to do at scale. The training loop therefore computes ground-truth MI on-the-fly per example, which the authors acknowledge is a significant practical overhead during training. This is an important caveat: the method trades expensive training setup for cheap inference-time MI estimation.
The decoding strategy itself uses the predicted MI matrix to identify subsets of tokens with low mutual dependence, unmasking those subsets in parallel. A temperature-like parameter gamma controls the aggressiveness of parallelisation, trading off speed against quality in a predictable way.
The comparison baselines are sensible. Sequential decoding is the quality ceiling. Naive parallel decoding (fixed batch size k) is the speed ceiling with no dependency awareness. The entropy-based sampler (EB-Sampler) represents the current state of practice, using per-token confidence rather than pairwise MI to decide decoding order. The MI-guided method is positioned as a principled upgrade to EB-Sampler.
Results
On Sudoku, the results are encouraging. At comparable numbers of forward passes (~15), MI-guided sampling achieves 63.6% solution accuracy versus 61.0% for EB-Sampler and 52.4% for naive parallel decoding. At higher parallelism (~9-10 passes), MI-guided sampling holds at 56.2% while EB-Sampler drops to 51.2%. Sequential decoding at 53.9 passes achieves 61.6%, so the MI-guided method at 15 passes actually slightly exceeds sequential quality, which is a meaningful result.
The Sudoku MI maps themselves are qualitatively interpretable: the estimator recovers row, column, and box constraint structure without any explicit programming of those rules. This is a useful interpretability result independent of the decoding application, suggesting MDMs genuinely internalise structural constraints in their hidden representations.
For proteins with ESM-C, the paper reports 3-5x reductions in forward passes while preserving generative quality, though the protein results are less exhaustively benchmarked in the tables presented. The application to a 300M parameter production model (ESM-Cambrian) with only 810K additional parameters is practically significant; the overhead is negligible relative to the backbone.
Limitations and Implications
The most honest limitation is the training cost. Computing ground-truth MI from the MDM's conditional distributions is expensive, and the authors note the predictor requires substantial preparation. This makes the method more accessible to researchers with existing pretrained MDMs and compute budgets than to practitioners training from scratch. The quality of the estimator is also imperfect by the authors' own admission, and the architecture of the prediction head is not systematically optimised.
There is also a question of generalisation across masking rates. The MI structure of a sequence changes depending on how much is masked, and it is not entirely clear how robust the estimator is across the full masking trajectory during generation. The paper trains and evaluates at specific conditions, but a more thorough ablation across masking levels would strengthen confidence in the approach.
The broader implication is worth taking seriously. The dominant approach to parallel decoding in MDMs has been entropy-based, essentially a per-token confidence heuristic with no explicit model of inter-variable dependence. This paper demonstrates that explicitly modelling those dependencies, even approximately, yields consistent gains. Given that protein and RNA design are increasingly reliant on MDMs for their parallelism advantages over autoregressive models, a principled way to guide that parallelism without sacrificing structural validity is practically valuable.
Future directions the paper does not fully explore include whether the MI estimator could be trained jointly with the MDM rather than post-hoc, whether attention patterns in transformer-based MDMs could serve as a cheaper proxy for MI without a separate head, and whether the framework extends naturally to continuous diffusion models through analogous mutual information quantities. Each of these represents a tractable research direction building directly on this work.
Overall, this is a focused and technically sound contribution. It does not overstate its claims, the experimental design is appropriate for the scope, and the dual framing as both an interpretability tool and an inference acceleration method gives it broader relevance than either application alone would warrant.