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

Discrete Diffusion Inference-Time Control with Nested Sequential Monte Carlo
Aug 20, 2026Lohithsai Yadala Chanchu, Hany Abdulsamad, Christian A. NaessethInference-Time Guidance Diffusion Language Model Inference
University of Amsterdam
Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models
Oct 1, 2026Zuokai Wen, Louis Grenioux, Weinan E +1Markov Chain Monte Carlo Diffusion Model Sampling
School of Artificial Intelligence, Shanghai Jiao Tong University · Center for Computational Mathematics, Flatiron Institute · School of Mathematical Sciences, Peking University
Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion
May 22, 2026Jaihoon Kim, Taehoon Yoon, Prin Phunyaphibarn +3Markov Chain Monte Carlo Amortized Inference
1KAIST · University of Michigan · 3NVIDIA

Aug 20, 2026, stat.MLJ/K move · Enter open · S save

Discrete Diffusion Inference-Time Control with Nested Sequential Monte Carlo

Lohithsai Yadala Chanchu, Hany Abdulsamad, Christian A. Naesseth

University of Amsterdam

Inference-Time Guidance Diffusion Language Model Inference Diffusion Language Models

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Discrete Diffusion Inference-Time Control with Nested Sequential Monte Carlo

Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models

Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion