FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models

Yinuo Ren; Haoxuan Chen; Grant M. Rotskoff; Jiequn Han; Lexing Ying

cs.LGSep 28, 2026

FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models

Authors: Yinuo Ren, Haoxuan Chen, Grant M. Rotskoff, Jiequn Han, Lexing Ying

Organizations: Kempner Institute Harvard University · ICME Stanford University · Department of Chemistry Stanford University · Center for Computational Mathematics Flatiron Institute · Department of Mathematics Stanford University

Abstract

Many inference-time tasks for pretrained discrete diffusion models and diffusion language models reduce to drawing samples from a tilted version of the pretrained distribution. Feynman-Kac sequential Monte Carlo (SMC) makes this correction exact in principle, but its prescribed weights routinely degenerate when the proposal dynamics are misaligned with the tilt, capping the practical gains from additional particles. We introduce FluxLite, a lightweight, training-free proposal-control framework for discrete diffusion. On the sparse directed graph of pretrained reverse rates, any sparse jump-rate perturbation can be exactly compensated by a $q_t$ -weighted graph-divergence term in the Feynman-Kac potential; the target path is therefore preserved while the residual reweighting variance becomes a local convex objective. We instantiate this principle as two practical samplers: a one-hop local reallocation rule (HEU) and a small nonnegative quadratic program over pretrained-rate bases (D-VCG). We further prove population stability under the standard score-entropy training loss, identifying a tilted-path coverage factor that governs robustness to score error, together with finite-particle convergence for a fixed controlled Feynman-Kac recursion. Empirically, FluxLite improves over standard Feynman-Kac SMC baselines by up to two orders of magnitude in terminal KL on an analytically tractable finite-state CTMC benchmark, and reduces row-correlation MSE on 2D Ising sampling by 5-7x in geometric mean and up to 55x at peak.

Figures & tables

	Continuous diffusion ( $\mathcal{X}=\mathbb{R}^{d}$ )	Discrete CTMC ( $\|\mathcal{X}\|=D$ )
Generator	$\widetilde{\mathcal{L}}_{t}^{\ast}\rho=-\nabla\cdot\big(\widetilde{\mathbf{v}}_{t}\rho\big)+\frac{\alpha_{t}^{\shortleftarrow}{}^{2}}{2}\Delta\rho$	$(\widetilde{\mathcal{L}}_{t}^{\ast}\rho)(x)=\sum_{y\neq x}\big(\widetilde{Q}_{t}(x,y)\rho(y)-\widetilde{Q}_{t}(y,x)\rho(x)\big)$
Feynman–Kac form	$\partial_{t}q_{t}=\widetilde{\mathcal{L}}_{t}^{\ast}q_{t}+g_{t}\,q_{t},\;\;\mathbb{E}_{q_{t}}[g_{t}]=0$
Control	drift perturbation $\mathbf{u}_{t}:\mathbb{R}^{d}\to\mathbb{R}^{d}$	rate perturbation $\mathbf{R}_{t}\in\mathbb{R}^{D\times D}$
Divergence	$\operatorname{div}_{q_{t}}\mathbf{u}_{t}(x):=\frac{\nabla\cdot(q_{t}(x)\mathbf{u}_{t}(x))}{q_{t}(x)}$	$\operatorname{div}_{q_{t}}\mathbf{R}_{t}(x):=\frac{\sum_{y\neq x}\big(R_{t}(y,x)q_{t}(x)-R_{t}(x,y)q_{t}(y)\big)}{q_{t}(x)}$
Equivalence	$(\widetilde{\mathbf{v}}_{t},g_{t})\sim(\widetilde{\mathbf{v}}_{t}+\mathbf{u}_{t},\;g_{t}+\operatorname{div}_{q_{t}}\mathbf{u}_{t})$	$(\widetilde{\mathbf{Q}}_{t},g_{t})\sim(\widetilde{\mathbf{Q}}_{t}+\mathbf{R}_{t},\;g_{t}+\operatorname{div}_{q_{t}}\mathbf{R}_{t})$
Optimality	$\min_{\mathbf{u}_{t}}\operatorname{Var}_{q_{t}}\!\left[g_{t}+\operatorname{div}_{q_{t}}\mathbf{u}_{t}\right]$	$\min_{\mathbf{R}_{t}}\operatorname{Var}_{q_{t}}\!\left[g_{t}+\operatorname{div}_{q_{t}}\mathbf{R}_{t}\right]$

Table 1: Continuous vs. discrete variance-control.

Figure 1: Finite-state CTMC benchmark: terminal

\KL(q_{T}\|\widehat{q}_{T}^{N})

for four

V=5

L=3

regimes with

N=4000

particles, 80 integration steps, and 10 seeds. D-VCG improves over D-FKC (geometric-mean ratio of terminal KL across seeds) up to

114.8\times

; settings and sweeps are in Appendix E .

Panel of Figure 2 — Figure 2: 2D Ising experiments against a SW reference, with a ghost spin for nonzero fields. Panels report magnetization/energy- $\mathsf{W}_{2}$ , row-correlation $\mathrm{MSE}(C(r))$ , and mean magnetization for PG , D-FKC , and 2-/4-basis D-VCG .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Finite-state CTMC ablations. Across both the parameter sweeps and the particle-count axis, D-VCG matches or improves on the dense oracle DEN without DEN ’s tractability assumptions.

Panel of Figure 4 — Figure 4: Further 2D Ising configurations. The same axes and methods as Fig. 2 confirm that the D-VCG advantage (i) holds for a second base-model training temperature ( a ) and (ii) survives in the pure-tilt limit ( $\gamma=1$ , $\lambda_{\rm eff}=0$ , panel b ), where the joint-mode anchor reduces to the free-controller D-VCG of Proposition 3.6 .

Figure 5: Particle-count sweep at

\beta_{r}=0.20

\beta_{\rm train}=\beta_{\rm target}=0.4

\gamma=1

\lambda_{\rm eff}=0

(Config 3). Across

N\in\{100,\dots,5000\}

, D-VCG (4-basis) reduces

\mathrm{MSE}(C(r))

over D-FKC by

15.21\times

in geometric mean, with peak

\mathbf{34.7\times}

N=1000

, showing that the advantage does not close as

N

grows.

	$\rho=40\%$		$\rho=60\%$
	Hellinger $\downarrow$	Meas. Error $\downarrow$	Hellinger $\downarrow$	Meas. Error $\downarrow$
SGDD	0.076	4.93	0.151	6.11
D-FKC	0.033	0.29	0.081	0.88
D-VCG	0.035	0.00	0.079	0.05

Appendix

Table 2: Quantitative evaluation on monophonic music infilling. Lower is better.

	MPS $\uparrow$	HPSv2 $\uparrow$
D-FKC	$15.282\pm 0.032$	$0.2980\pm 0.0003$
D-VCG	$\mathbf{15.422}\pm 0.032$	$\mathbf{0.3033}\pm 0.0003$

Appendix

Table 3: Quantitative evaluation on CFG text-to-image generation: mean

\pm

standard error over

100

prompts. Higher is better.

Prompt	An anime dragon rider above snowy mountains, scarf trailing behind in the wind.	A dramatic anime close-up of a masked hero under falling snow, city skyline behind.	Concept art of a giant mechanical whale swimming through clouds above a coastal town.
D-FKC
D-VCG
Prompt	Concept art of a hidden pirate cove with waterfalls, wooden bridges, and lantern-lit paths.	A soft acrylic painting of a lighthouse on cliffs during a calm blue evening.	A watercolor botanical illustration of orchids, ferns, and delicate green leaves.
D-FKC
D-VCG

Appendix

Table 4: Qualitative comparison of D-FKC and D-VCG on six representative CFG prompts.

Explore similar work

Aug 20, 2026stat.ML

Discrete Diffusion Inference-Time Control with Nested Sequential Monte Carlo

We study inference-time control for text generation in discrete diffusion language models, where the goal is to steer sampling toward sequence-level rewards without retraining. Prior work in this domain has focused on particle-based methods such as best-of-

n

sampling and bootstrap sequential Monte Carlo, which may suffer from overoptimism and weight degeneracy, respectively. We address these limitations using \emph{nested} sequential Monte Carlo methods. We formulate nested SMC (NSMC) and fully-adapted nested SMC (FA-NSMC) for Feynman--Kac steering, identifying and correcting errors in prior formulations that lead to biased final estimates. We evaluate these methods on toxicity and fluency steering tasks, showing that NSMC and FA-NSMC consistently outperform best-of-

n

and bootstrap SMC.

Lohithsai Yadala Chanchu, Hany Abdulsamad, Christian A. Naesseth

Inference-Time Guidance Diffusion Language Model Inference Diffusion Language Models

Ask PDF

Oct 1, 2026stat.ML

Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models

Inference-time scaling adapts pretrained diffusion models to new sampling tasks without additional training. Existing methods rely primarily on Monte Carlo sampling with more particles, yet are premised on the pretrained model being exact. In practice, data and training limitations make the model imperfect, and these methods inherit its error. More particles reduce Monte Carlo error but cannot remove the mismatch between the endpoint and the desired target or the error in tracking the prescribed probability path. We introduce the Energy-based Feynman-Kac Corrector (EBFKC), a framework for energy-based diffusion models that corrects these errors on the fly given a reference energy. We first derive Feynman-Kac dynamics that track a prescribed path exactly in the continuous-time population limit even when the model is imperfect, and approximate these dynamics using sequential Monte Carlo with variance-controlling guidance. To remove the endpoint mismatch, we use the pretrained energy as a surrogate along the diffusion path and progressively incorporate the discrepancy between the learned and target terminal energies. Experiments on Gaussian mixture models, particle systems, alanine dipeptide, and alanine tetrapeptide show that our method closely matches target distributions and molecular free-energy profiles under annealing and reward tilting, whereas standard inference-time scaling baselines retain substantial sampling errors.

Zuokai Wen, Louis Grenioux, Weinan E +1

School of Artificial Intelligence, Shanghai Jiao Tong University · Center for Computational Mathematics, Flatiron Institute · School of Mathematical Sciences, Peking University

Markov Chain Monte Carlo Diffusion Model Sampling Inference-Time Scaling Diffusion Models

Ask PDF

May 22, 2026cs.LG

Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion

Discrete diffusion models have emerged as powerful frameworks for generating structured categorical data. However, efficiently sampling from reward-tilted distributions remains a fundamental challenge. While Twisted Sequential Monte Carlo (SMC) offers asymptotic exactness for this task, estimating the optimal twist function in discrete state spaces necessitates costly Monte Carlo approximations, resulting a severe computational bottleneck at inference. To overcome this limitation, we introduce Contrastive Distribution Matching (CDM), a novel framework that amortizes the cost of SMC inference by learning a parameterized twist function via positive and negative samples. For efficient training, we reformulate the gradient estimator to leverage the closed-form forward kernels of discrete diffusion models. In practice, evaluating our learned twist function incurs less than 5% additional computational overhead compared to a single forward pass of the base model. Through extensive empirical evaluations, we demonstrate that CDM consistently outperforms existing baselines under matched wall-clock time. We validate the effectiveness and versatility of our approach across a diverse range of applications, including toxic text generation, regulatory DNA sequence design, protein designability, and diffusion large language model alignment.

Jaihoon Kim, Taehoon Yoon, Prin Phunyaphibarn +3

1KAIST · University of Michigan · 3NVIDIA

Markov Chain Monte Carlo Amortized Inference Diffusion Sampling Inference-Time Guidance Diffusion Models

Ask PDF