One-sentence Summary

Researchers from the University of Amsterdam and Meta FAIR propose a context-weighted discrete flow matching method that augments the continuous-time Markov chain with local context information, reducing generative perplexity by up to 63% on OpenWebText, matching a strong semi-autoregressive block diffusion baseline while preserving any-order generation, and demonstrating that simple context-aware modifications significantly boost both sampling and training efficiency.

Key Contributions

• The paper empirically demonstrates that token-level prediction entropy in discrete flow matching is closely tied to the density of unmasked neighbors, with denser context reducing uncertainty.
• A context-weighted modification to the CTMC transition dynamics is introduced that adjusts update probabilities based on local neighborhood density while preserving the per-token marginal unmasking probability, shifting intermediate steps toward better-conditioned tokens.
• Two practical context-aware mechanisms are derived: an inference-time sampler that improves generation quality without fine-tuning, and a scaled cross-entropy objective that reweights training tokens, reducing generative perplexity on OpenWebText by up to 63%, boosting MAUVE by 24%, and increasing valid and novel molecular samples up to 2.8× and 1.9×, while matching a semi-autoregressive baseline and preserving any-order generation.

Introduction

Generative models for discrete data, such as text and molecules, increasingly adopt discrete flow matching and diffusion, which learn continuous-time Markov processes (CTMCs) to generate tokens in parallel and in any order. Unlike autoregressive approaches that condition on a fixed prefix, these methods marginalize over all possible token orderings, so the model must predict tokens under a wide range of partial observations. This leads to highly variable prediction difficulty: some tokens are strongly constrained by surrounding context while others remain ambiguous. Prior work attempts to handle this imbalance by enforcing structured generation orders or adjusting the sampler at inference time, but does not integrate prediction uncertainty directly into the generative dynamics.

The authors observe that prediction difficulty strongly correlates with the number of already unmasked tokens in a token’s local neighborhood. They modify the CTMC transition probabilities to depend on this local context while preserving the per-token marginal probability of being unmasked along the probability path. This yields two practical mechanisms: a context-weighted sampler that can be applied to pre-trained models without fine-tuning, and a scaled cross-entropy training objective that reweights tokens based on their local context. The approach improves generation quality across domains, boosting text generation scores by up to 24% and dramatically increasing validity and novelty in molecule generation while matching the performance of block diffusion baselines without sacrificing any-order generation.

Method

The authors build upon Discrete Flow Matching (DFM), which models finite discrete sequences using a Continuous-Time Markov Chain (CTMC). In this framework, a conditional probability path is defined between a source noise sample and a data sample. The transition kernel relies on factorized velocities, meaning transitions change one coordinate at a time. This factorization enables efficient coordinate-wise simulation. The standard training objective is the factorized Conditional Matching objective, which decomposes into coordinate-wise prediction problems where all tokens are weighted uniformly.

However, the authors observe that uniform weighting is suboptimal because the reducible learning signal varies across tokens. They empirically find that local context serves as a simple proxy for prediction difficulty. Specifically, the number of unmasked neighbors around a token is strongly correlated with both the prediction entropy and the negative log-likelihood of the true token.

As shown in the figure below:

Both metrics decrease consistently as more local context becomes available, indicating that nearby observed tokens make the prediction less uncertain and more accurate. This variation matters because the expected loss of the model predictor decomposes into an irreducible uncertainty term and a reducible model error term. For high-entropy coordinates, a large part of the loss comes from intrinsic ambiguity, providing a less informative training signal. Low-entropy coordinates correspond to more determined targets and offer sharper supervision. Ideally, the weighting scheme should emphasize the reducible component of the loss rather than uniformly weighting raw per-token losses.

To address this, the authors propose a context-aware modification to the standard CTMC process called Context-weighted Discrete Flow Matching. They introduce a context weight function $\alpha$α that maps the current state to positive weights for each token position. This function is integrated into the CTMC by defining a new factorized velocity that scales the standard probability velocity by the context weight. This new velocity generates a context-weighted conditional probability path. To ensure the path connects the noise and data distributions and preserves the marginal distribution of the standard convex path, the authors establish specific conditions on the context weight function. These conditions require the weights to average to 1 over the currently masked coordinates and to be equivariant under circular shifts.

During inference, this formulation allows for a purely inference-time modification of the sampling procedure. The authors inject a multiplicative weight into each coordinate of the Euler solver, obtaining a context-weighted update rule. They introduce two specific context-weighted solvers: a Neighbor-weighted solver, which uses the number of unmasked tokens in a local window, and an Entropy-weighted solver, which uses the token-level predictive entropy available at test time. In both cases, the weights are normalized so that locally well-contextualized or low-entropy positions receive larger weights.

For training, since the context-weighted path is not available in an analytic form for general context weights, simulation-free sampling is impossible. Instead of relying on costly time-wise simulation, the authors leverage the known distribution of the number of unmasked tokens to construct an exact token-wise sampler with linear complexity. They first sample the number of unmasked tokens and then update the tokens sequentially, eliminating the need for time-axis simulation while returning exact samples from the context-weighted path.

Finally, motivated by the empirical link between prediction uncertainty and available local context, the authors introduce Scaled Cross-Entropy as a plug-in alternative to the standard objective. This loss function uses the context weight function to reweight the contribution of each coordinate, upweighting better-conditioned coordinates and downweighting ambiguous ones. This approach leaves the train-time path sampling unchanged, retaining the efficiency of standard training and any-order sampling at inference time, while effectively capturing the context-weighted importance of different token updates.

Experiment

The method is evaluated on text generation with OpenWebText and molecular generation with QM9, using a discrete flow matching architecture. Inference-time context-weighted sampling, particularly neighbor-weighted updates, improves generation quality and is especially effective in low-data regimes, while train-time scaled cross-entropy brings the model close to strong semi-autoregressive baselines. Training with the context-weighted path further benefits small-data settings, and an intermediate neighborhood radius provides the best quality-diversity trade-off.

Context-weighted Discrete Flow Matching with Scaled Cross-Entropy lifts MAUVE from 0.628 to 0.777 and cuts generative perplexity by 63%, bringing it close to strong semi-autoregressive block diffusion baselines. It clearly outperforms non-autoregressive baselines SEDD and MDLM, while retaining any-order generation flexibility. Scaled Cross-Entropy raises MAUVE to 0.777, approaching BD3LM with L'=4 (0.784) and surpassing L'=8 (0.734), and leaves SEDD (0.574) and MDLM (0.685) far behind. Generative perplexity drops from 82.17 to 30.20, a 63% reduction, placing it between BD3LM L'=8 (29.47) and L'=16 (32.95) and well ahead of SEDD (110.23) and MDLM (41.73).

Scaled cross-entropy (SCE) greatly raises generation quality over standard cross-entropy, NELBO, and Bregman divergence for discrete flow models. On the OpenWebText benchmark with a uniform source, SCE lifts MAUVE and cuts generative perplexity by about two-thirds while entropy remains close to that of the baseline. For molecular generation on QM9, SCE produces the most valid molecules among the compared losses and stays competitive in novelty. With a uniform source, SCE increases MAUVE from 0.628 to 0.777 and drops generative perplexity from 82.2 to 30.2 on OpenWebText. Entropy under SCE (7.66) stays near the cross-entropy level (7.93), indicating that the quality gain does not come at the cost of diversity collapse. On QM9, SCE yields 556 valid molecules, substantially more than the 475 from standard cross-entropy, while novelty remains similar (297.6 vs. 287.0). NELBO and Bregman losses offer only moderate improvements over CE, whereas SCE provides a much larger jump in both MAUVE and perplexity.

Switching from the mixture path to the neighbor-weighted path yields clear improvements in the masked-source setting on both OWT and QM9, while offering negligible benefit for the uniform source on OWT. On QM9, the neighbor path consistently raises validity and novelty, nearly doubling valid molecules for the masked source. For the uniform source on OWT, the neighbor path does not improve MAUVE or generative perplexity, only slightly increasing entropy. For the masked source on OWT, the neighbor path lifts MAUVE from 0.690 to 0.751 and reduces generative perplexity. On QM9, the neighbor path improves both validity and novelty across source types, with masked-source validity nearly doubling from 177 to 298. The neighbor path is most impactful in the masked-source regime, aligning with its enhanced effectiveness in low-data settings.

The evaluation spans text generation on OpenWebText and molecular generation on QM9, examining discrete flow matching models with different loss functions and source distributions. Scaled cross-entropy substantially outperforms standard cross-entropy, NELBO, and Bregman losses, delivering major gains in generation quality without sacrificing diversity. A neighbor-weighted masking path further improves generation quality under masked-source conditions, particularly in low-data molecular settings, but does not enhance uniform-source performance.