Context-weighted Discrete Flow Matching

2026-07-23

AI Summary

Standard training objectives in discrete flow matching conflate easily predictable tokens with high-entropy, hard-to-predict ones, limiting both generation efficiency and quality. This work proposes an improved framework based on continuous-time Markov chains that, for the first time, incorporates local context density as a weighting factor in discrete flow matching. The approach introduces a context-weighted sampler and a scaled cross-entropy loss, significantly enhancing training and sampling efficiency with negligible computational overhead. Evaluated on OpenWebText, the method reduces generation perplexity by up to 63%, achieving generation quality comparable to strong semi-autoregressive block diffusion models while preserving the capability for arbitrary-order generation.

Abstract

Discrete flow matching provides a flexible framework for generative modeling on discrete structures. However, the standard factorized training objective exposes the model to targets of varying difficulty, mixing well-conditioned, predictable tokens with ambiguous, high-entropy ones. We empirically demonstrate that the uncertainty over the value of each token is closely related to the density of available context in its neighborhood. Motivated by this observation, we propose a simple modification to the underlying continuous-time Markov chain (CTMC) that incorporates local context information. Our context-weighted sampler improves generation quality with negligible computational overhead, while our scaled cross-entropy loss function reweights the training signal from different tokens and reduces generative perplexity by up to 63% on OpenWebText. Moreover, our approach matches a strong semi-autoregressive block diffusion baseline in quality while retaining the ability to perform generation in any order. These results highlight the role of local context as an important factor in discrete generative modeling and show that simple context-aware modifications can significantly improve both sampling and training efficiency.

Problem

• discrete flow matching
• generative modeling
• token uncertainty
• context density
• training objective

Innovation

• discrete flow matching
• context-weighted sampling
• continuous-time Markov chain
• scaled cross-entropy loss
• generative modeling