arXiv subjects

Zachary Bezemek

Publications and source records attributed to Zachary Bezemek.

6 recordsLinked to original sources

arXiv2026

Path Planning for Masked Diffusion Model Sampling

Any order generation of discrete data using masked diffusion models (MDMs) offers a compelling alternative to traditional autoregressive models, especially in domains that lack a natural causal ordering of data. However, current popular MDMs depart from their successful continuous diffusion model counterparts with simplified masked inference wherein unmasked tokens cannot be iteratively refined -- even if there is a mistake. In this paper, we extract the full power of MDMs by introducing a novel inference sampling strategy termed Path Planning (P2) that decomposes each generation step into two sub-stages: planning and denoising. Under P2, the planner at every step selects appropriate tokens that are marked to be updated, which can then be sampled using the denoiser. We demonstrate that P2 generalizes all existing sampling strategies for MDMs and critically enhances generative quality through the new capability of refining and updating existing unmasked tokens. We theoretically prove that P2 establishes a (new) expanded evidence lower bound (ELBO) on the log marginal likelihood of data. We instantiate P2 with a family of planners including: 1.) Self-Planning, 2.) BERT-Planning, and 3.) Trained-Planning with a learned planner leading to SOTA generative performance for MDMs on a suite of domains. Specifically, solely using P2 inference, we observe relative improvements of 22% in protein sequence foldability, 8% in RNA sequence pLDDT, 4% in math reasoning, 68% in story generation (ROUGE score), and 33% in code generation for the challenging pass@1 metric.

cs.LG↗

arXiv2026

Planner Aware Path Learning in Diffusion Language Models Training

Diffusion language models have emerged as a powerful alternative to autoregressive models, enabling fast inference through more flexible and parallel generation paths. This flexibility of sampling is unlocked by new engineered sampling strategies, or planners, that select more favorable generation paths by iteratively planning - versus uniformly at random - where to denoise along the sequence. However, by modifying the reverse paths via planning, planners create an irrevocable mismatch between the uniformly random denoising paths assumed during training and planning-based inference. In this paper, we systematically investigate the mismatch of discrete diffusion training and inference under planning and theoretically prove that the standard discrete diffusion training evidence lower bound (ELBO) does not accurately describe a denoiser that uses a non-uniform planner. To address this gap, we derive a new planned evidence lower bound (P-ELBO) that incorporates planner-based reverse dynamics directly into the training objective. Using the P-ELBO, we introduce Planner Aware Path Learning (PAPL), a novel training scheme that aligns training and inference under a planned denoiser. PAPL is implemented as a simple yet effective modification to the standard masked discrete diffusion loss, making it widely applicable and easy to adopt. Empirically, we show PAPL delivers consistent gains across domains, including a 40% relative improvement in protein sequences, improved text generation with up to a 4x relative MAUVE gain, and 23% relative improvement in code generation HumanEval pass@10. Code is available at github.com/pengzhangzhi/PAPL .

cs.LG↗

arXiv2023

Importance Sampling for the Empirical Measure of Weakly Interacting Diffusions

We construct an importance sampling method for computing statistics related to rare events for weakly interacting diffusions. Standard Monte Carlo methods behave exponentially poorly with the number of particles in the system for such problems. Our scheme is based on subsolutions of a Hamilton-Jacobi-Bellman (HJB) Equation on Wasserstein Space which arises in the theory of mean-field (McKean-Vlasov) control. We identify conditions under which such a scheme is asymptotically optimal. In the process, we make connections between the large deviations principle for the empirical measure of weakly interacting diffusions, mean-field control, and the HJB Equation on Wasserstein Space. We also provide evidence, both analytical and numerical, that with sufficient regularity of the HJB Equation, our scheme can have vanishingly small relative error in the many particle limit.

math.PR↗

arXiv2023

Moderate deviations for fully coupled multiscale weakly interacting particle systems

We consider a collection of fully coupled weakly interacting diffusion processes moving in a two-scale environment. We study the moderate deviations principle of the empirical distribution of the particles' positions in the combined limit as the number of particles grow to infinity and the time-scale separation parameter goes to zero simultaneously. We make use of weak convergence methods, which provide a convenient representation for the moderate deviations rate function in terms of an effective mean field control problem. We rigorously obtain equivalent representations for the moderate deviations rate function in an appropriate "negative Sobolev" form, which is reminiscent of the large deviations rate function for the empirical measure of weakly interacting diffusions obtained in the 1987 seminal paper by Dawson-Gärtner. In the course of the proof we obtain related ergodic theorems and we rigorously study the regularity of Poisson type of equations associated to McKean-Vlasov problems, both of which are topics of independent interest.

math.PR↗

arXiv2022

Large deviations for interacting multiscale particle systems

We consider a collection of weakly interacting diffusion processes moving in a two-scale locally periodic environment. We study the large deviations principle of the empirical distribution of the particles' positions in the combined limit as the number of particles grow to infinity and the time-scale separation parameter goes to zero simultaneously. We make use of weak convergence methods providing a convenient representation for the large deviations rate function, which allow us to characterize the effective controlled mean field dynamics. In addition, we obtain equivalent representations for the large deviations rate function of the form of Dawson-Gärtner which hold even in the case where the diffusion matrix depends on the empirical measure and when the particles undergo averaging in addition to the propagation of chaos.

math.PR↗

arXiv2022

Rate of homogenization for fully-coupled McKean-Vlasov SDEs

We consider a fully-coupled slow-fast system of McKean-Vlasov SDEs with full dependence on the slow and fast component and on the law of the slow component and derive convergence rates to its homogenized limit. We do not make periodicity assumptions, but we impose conditions on the fast motion to guarantee ergodicity. In the course of the proof we obtain related ergodic theorems and we gain results on the regularity of Poisson type of equations and of the associated Cauchy-Problem on the Wasserstein space that are of independent interest.

math.PR↗