PaperlayerView exact revision on arXiv

Source-linked AI summary

FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models

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

arXiv:2609.35947v1cs.LGstat.COstat.ML

TL;DR

Inference-time Feynman–Kac SMC can lose practical value when proposal dynamics misalign with the tilt and weights degenerate. FluxLite controls discrete proposals while preserving the target path, and reports gains over D-FKC on finite-state CTMC and 2D Ising benchmarks.

  • Problem

    Feynman–Kac SMC can suffer weight degeneracy when proposal dynamics misalign with the tilt, limiting the practical gains from adding particles.

    Source: §1 · p. 2

  • Method

    FluxLite uses sparse jump-rate control with a q_t-weighted graph-divergence correction to preserve the tilted path, instantiated by HEU and D-VCG.

    Source: §1 · p. 2, §1.1 · p. 2

  • Results

    D-VCG improves terminal KL over D-FKC by up to 114.8× on finite-state CTMCs, while FluxLite reduces 2D Ising row-correlation MSE over D-FKC by 5–7× geometrically and up to 55× at peak.

    Source: §1.1 · p. 2, §5.1 · p. 9

  • Takeaways & Limitations

    FluxLite provides training-free proposal control on the pretrained transition graph, with reported gains on finite-state CTMC and 2D Ising sampling.

    Source: §1 · p. 2, §1.1 · p. 2

    The finite-particle convergence result concerns a fixed-grid recursion and excludes adaptive ESS resampling, time discretization, score error, and controls estimated on the same particles.

    Source: §4 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Inference-time tasks often require sampling from a tilted pretrained distribution, but Feynman–Kac SMC can suffer weight degeneracy when its proposal and weights are misaligned. FluxLite controls proposals on the pretrained discrete transition graph without retraining, preserving the target path while reducing residual reweighting variance.

1 Introduction

  • Many inference-time objectives, including reward-aligned generation and posterior sampling, reduce to sampling from a tilted pretrained distribution.The paper situates this pattern across autoregressive LLMs and continuous and discrete diffusion, amid a shift toward allocating more compute at inference time.Source: §1 · p. 1
  • Weight degeneracy can concentrate SMC weights on few particles, so adding particles may not add information or sustain gains over selection baselines.Existing diffusion and LLM methods that seek better proposals often require extra training or modality-specific heuristics.Source: §1 · p. 2
  • FluxLite preserves the tilted path by compensating sparse jump-rate perturbations with a q_t-weighted graph-divergence correction to the Feynman–Kac potential.This turns proposal design into sparse flux reallocation with a variance-control objective.Source: §1 · p. 2

1.1 Contributions

  • FluxLite uses HEU and D-VCG, proves population stability and finite-particle convergence, and improves CTMC terminal KL by up to 114.8× and Ising row-correlation MSE by 5–7× geometrically and up to 55× at peak.The CTMC comparison is against D-FKC; the Ising comparison is also against D-FKC.Source: §1.1 · p. 2

2 Preliminaries

The preliminaries frame inference-time objectives as tilted paths and review the Feynman–Kac representations used to sample them. They connect continuous drift control to discrete CTMC rate control on a fixed transition graph.

2 Preliminaries

  • The preliminaries position Feynman–Kac SMC as a sampling-based form of inference-time width scaling, propagating particles and reweighting them by an incremental potential.The continuous discussion recalls DriftLite’s non-unique representations, where probability transport can shift between a proposal drift and residual potential through a weighted divergence.Source: §2 · p. 2, §2.1 · p. 3

2.2 Discrete Feynman–Kac flow

  • For discrete diffusion, the pretrained reverse process defines a CTMC on a finite state space, with transition rates and a reweighting potential representing the tilted path.The rate convention takes x as the source and y as the destination; the learned model estimates the reverse rates from score estimates.Source: §2.2 · p. 3
  • The tilt changes weights on the reverse process’s directed edges but introduces no new graph edges, and baseline Feynman–Kac SMC propagates under those rates before reweighting particles.A typical target path is proportional to the pretrained marginal raised to γ and multiplied by the exponential of a reward, likelihood, or guidance potential.Source: §2.2 · p. 3

3 Methodology

FluxLite searches within equivalent Feynman–Kac representations for sparse jump dynamics with lower residual reweighting variance. It develops a local heuristic and a basis-restricted quadratic program to make this proposal control practical.

3.1 Graph calculus

  • Any admissible sparse jump-rate perturbation can be compensated by a q_t-weighted graph-divergence term in the potential, preserving the same target path.This equivalence class lets proposal control search for a representation with smaller reweighting variance.Source: §3.1 · p. 4

3.2 Variance control

  • FluxLite minimizes the variance of the residual centered potential; a zero minimum means the target path can be realized by a pure CTMC without reweighting.Without sparsity constraints, the discrete zero-variance objective has an explicit solution, unlike the continuous case’s Poisson-equation solution.Source: §3.2 · p. 5
  • The dense zero-variance construction is generally impractical on discrete-diffusion graphs, while sparse variance reduction is a constrained weighted least-squares problem.The full sparse optimization has no closed-form solution and requires iterative solvers or coarse approximations that add per-step inference cost.Source: §3.2 · pp. 5–6

3.3 Practical sparse controls

  • The practical controls produce sparse effective rates and an associated residual centered potential, which supplies the particle weights while retaining the target path.The section presents HEU and D-VCG as practical constructions within the equivalence class.Source: §3.3 · p. 6

3.3.1 HEU: one-hop local reallocation

  • HEU restricts pairwise low-to-high flux reallocation to a local neighborhood and uses a local effective neighborhood size in place of global normalization.With reciprocal neighborhoods and matching local normalizations, it acts as a damped average of nearby source terms, removing local fluctuation while leaving the local mean.Source: §3.3.1 · p. 6

3.3.2 D-VCG: discrete variance-controlling guidance

  • D-VCG restricts controlled rates to a small nonnegative span of pretrained-rate bases, yielding a nonnegative quadratic program in the basis coefficients.For nonsingular A_t, the unconstrained minimizer is unique; the practical routine approximately solves the constrained problem using weighted particles.Source: §3.3.2 · pp. 6–7

4 Theoretical Analysis

The analysis separates population bias from finite-particle error in a controlled Feynman–Kac flow. It relates population robustness to tilted-path coverage and particle error to residual-potential oscillation, while stating limits on the guarantees.

Population stability under local-ratio error

  • The population analysis separates local-ratio score error from particle error and bounds guided-flow bias using score-entropy loss and a tilted-path edge-coverage factor.The coverage factor is not a score-accuracy condition: bias can be larger when the tilted path places mass on edges with large score error.Source: §4 · p. 8
  • When the tilted path shifts mass toward poorly covered edge endpoints, the same score-entropy loss can yield larger terminal bias.The coverage scalar measures whether the tilted path remains covered by the original reverse path on trained edges.Source: §4 · p. 8

Particle error for a fixed controlled representation

  • The particle theorem covers every-step bootstrap SMC with deterministic fixed-grid controls, not adaptive ESS resampling, score or time-discretization error, or same-particle control feedback.The paper describes its results as modular rather than a full end-to-end bound for adaptive controls estimated from the particle cloud.Source: §4 · p. 8
  • For a fixed controlled grid model, the particle approximation error is governed by residual-potential oscillation and is analyzed through a finite-particle Feynman–Kac recursion.The theorem concerns bounded test functions and its proof separates fixed-model Monte Carlo error from population score-ratio perturbation.Source: §4 · p. 8

5 Experiments

FluxLite is evaluated on an exact finite-state CTMC benchmark and a learned 2D Ising benchmark. D-VCG substantially improves terminal KL and row-correlation MSE over D-FKC in the reported settings.

5 Experiments

  • The experiments compare FluxLite samplers with D-FKC and guided propagation without reweighting across exact CTMC and learned 2D Ising settings.Additional large-scale discrete inverse-problem and classifier-free-guidance experiments are reported in the appendix.Source: §5 · pp. 8–9

5.1 Finite-state CTMC

  • Up to 114.8×, D-VCG improves over D-FKC in geometric-mean terminal KL across seeds on the four finite-state CTMC regimes.The benchmark uses N = 4000 particles, 80 integration steps, and 10 seeds; D-VCG matches or improves on the population zero-variance oracle DEN across all four cells.Source: Figure 1 · p. 9, §5.1 · p. 9

5.2 2D Ising Model

  • D-VCG reduces row-correlation MSE over D-FKC by 5.33× on annealing and 6.77× on joint anneal plus tilt in geometric mean.Per-cell peaks are 24.5× and 55.4×, respectively, with improvements across the plotted sweeps.Source: §5.2 · p. 9

6 Discussion and Conclusion

FluxLite controls discrete Feynman–Kac SMC by reallocating jump-rate flux on the pretrained graph while compensating with a graph-divergence potential, preserving the tilted marginals. The discussion places this proposal control alongside inference-time scaling and related methods, and identifies future extensions.

6 Discussion and Conclusion

  • FluxLite preserves tilted marginals while selecting sparse graph-flux reallocations to reduce reweighting variance; its theory addresses score-ratio stability and fixed-grid particle error.The paper identifies hybrid controls, alignment or pretraining couplings, and distillation as future directions, alongside extensions to adaptive ESS and particle-dependent controls.Source: §6 · p. 9

A.1 Markov generative models

  • Discrete diffusion connects iterative diffusion dynamics with finite- or countable-state reverse processes, including diffusion-language and masking models.The paper focuses on the CTMC structure of those reverse dynamics and proposal control without changing the pretrained graph.Source: §A.1 · p. 23

A.2 Inference-time scaling

  • Sampling-based width scaling targets tilted model distributions with Feynman–Kac SMC, while FluxLite changes proposal rates within an exact representation class to address proposal-weight misalignment.The nearest discrete corrector works target tilted laws through particle corrections [59] [81]; FluxLite instead controls sparse CTMC graph rates.Source: §A.2 · p. 23

A.3 Fine-tuning and alignment

  • FluxLite keeps the pretrained discrete diffusion model fixed, optimizing only its inference-time Feynman–Kac proposal and reweighting representation, unlike parameter- or guidance-network alignment.The passage frames it as complementary to training-time alignment methods, including RL-style objectives [187] and DPO variants [113] [154].Source: §A.3 · p. 24

A.4 Broader future directions

  • The paper points to hybrid discrete–continuous controls, alignment and preference-optimization couplings, and distillation into transport maps or parallel-reasoning architectures as adjacent directions.These future directions are presented as possibilities, not evaluated results.Source: §A.4 · p. 24

B Continuous Proposal Control (DriftLite)

  • DriftLite shifts probability transport between a continuous control drift and residual potential through a weighted divergence; FluxLite adapts that proposal-control idea to sparse jump-rate graphs.For a given tilted path, multiple controls can represent the same flow [119].Source: B Continuous Proposal Control… · p. 24

C Notation, Special Cases, and Proofs for Sections 2–3

C.3 Proof of Proposition 2.1

The appendix derives the sparse rate-control identities and optimization used by FluxLite, then states the coverage and regularity conditions for its theory results.

C.4 Proof of Proposition 3.1

  • For any admissible sparse rate perturbation, adding its q_t-weighted graph divergence to the centered potential preserves the Feynman–Kac dynamics.This is the compensation identity used to change rates without changing the represented path.Source: §C.4 · p. 27

C.6 Proof of Proposition 3.3

  • The dense zero-variance construction cancels the centered potential using flux, but sparse control restricts reallocation to allowed edges.The dense construction's flux is derived from the source vector; the sparse surrogate uses only the permitted graph.Source: §C.6 · p. 28

C.7 Proof of Proposition 3.4

  • Minimizing the sparse variance objective is equivalent to a constrained weighted least-squares problem.The residual flux form expresses the centered potential through the divergence of the candidate rate family.Source: §C.7 · p. 28

C.8 Proof of Proposition 3.6

  • The D-VCG objective is a convex quadratic over nonnegative coefficients, with normal equations characterizing the unconstrained minimizers.If the quadratic matrix is nonsingular, the unconstrained minimizer is unique.Source: §C.8 · p. 29

C.9 HEU details

  • HEU damps deviations of the source term from a local neighborhood average, and larger neighborhoods can better approximate the dense construction at higher computational cost.The damping parameter can further stabilize reallocation.Source: §C.9 · pp. 29–30

D Assumptions and Proofs for Section 4

  • The theory appendix treats population score-ratio stability and finite-particle convergence for a fixed deterministic-grid Feynman–Kac recursion.The stated results use specific assumptions and distinguish the score-ratio analysis from the fixed-grid particle result.Source: D Assumptions and Proofs for Section 4 · p. 30

D.1 Assumptions

  • Score-error robustness depends on tilted-path coverage: large endpoint coverage weights can amplify local-ratio error even when average training loss is low.The coverage factor weights edges by endpoint coverage rather than replacing the profile with a worst-case supremum.Source: §D.1 · p. 32

D.2 Formal score-ratio theorem and proof

The appendix formalizes score-ratio stability under known forward rates and gives a separate particle-only convergence result for a fixed-grid recursion. It also relates residual-potential control to one-step weight variance.

D.2 Formal score-ratio theorem and proof

  • The score-ratio theorem assumes known forward rates and bounds the difference between exact and implemented guided flows under a local-ratio perturbation.Its score-entropy form follows under the additional coverage assumption, and the time-integrated form is the main-text bound.Source: §D.2 · pp. 32–33
  • The theorem proofs use separate conditions: score-ratio stability concerns known-forward local-ratio perturbations, while particle convergence concerns a fixed recursion.The particle result isolates Monte Carlo error rather than score estimation or time-discretization error.Source: §D.2 · p. 32, §D.3 · p. 35
  • The score-entropy loss enters the operator-error bound edgewise, without a state-wise supremum over the score ratio.The proof converts this internal bound to the score-entropy form using the coverage factor.Source: §D.2 · pp. 34–35

D.3 Formal particle theorem and proof

  • The particle theorem analyzes an every-step bootstrap recursion on a fixed deterministic time grid, with conditionally independent particle updates.The proof uses this conditional i.i.d. structure for its particle-only bound.Source: §D.3 · pp. 35–36
  • The fixed-grid particle proof decomposes terminal error into initial-sampling error and later resampling-and-mutation errors, then sums their transported contributions.Its argument applies a positive-transform empirical-error bound at each grid step.Source: §D.3 · pp. 36–37

D.4 Incremental weight variance

  • The one-step weight-variance lemma connects the variance-control objective to a bound on residual centered-potential weight variance.The centered incremental weight is defined as the exponential of the grid-step residual centered potential.Source: §D.4 · p. 37

D.5 Discussion of the theorems

The discussion clarifies the scope of the theory and describes the finite-state CTMC benchmark setup and ablations. In those ablations, D-VCG matches or improves on the dense oracle across parameter and particle-count sweeps.

D.5 Discussion of the theorems

  • The score-ratio theorem covers known-forward local-ratio perturbations, not misspecified forward kernels or independently learned rates outside that form.Its regularity and coverage quantities are conditions, not score-accuracy measures.Source: §D.5 · p. 37
  • The particle theorem is limited to consistency for a fixed Feynman–Kac model and requires every-step multinomial resampling.It does not bound score-estimation, training, time-discretization, simulation, adaptive-resampling, or same-cloud control-estimation errors.Source: §D.5 · p. 37

E.2 Finite-state CTMC

  • Across parameter sweeps and the particle-count axis, D-VCG matches or improves on dense oracle DEN without DEN’s tractability assumptions.The ablations also report regime-dependent performance; D-VCG improves over PR in both main uniform-state configurations.Source: Figure 3 · p. 38, §E.2 · p. 39

E.1 Shared implementation

  • The nonnegative D-VCG coefficients are chosen by active-set enumeration with a diagonal ridge, making the implementation an approximate solve of the stated objective.Candidate generation and scoring use distinct regularization terms.Source: §E.1 · p. 38

E.2 Finite-state CTMC

  • The four CTMC configurations share V = 5, L = 3, 4000 particles, an 80-step grid, and 10 seeds, while varying diffusion family and tilt strength.The reported comparisons use a fixed discretization; masked-absorbing annealing is limited to γ = 1.3 because larger values destabilize standard Feynman–Kac SMC.Source: §E.2 · p. 38

E.3.1 Model, training, and reference

  • The Ising reference estimates magnetization, energy, and correlation from Swendsen–Wang samples, using an augmented ghost-spin graph for nonzero fields.The target is an Ising model in a uniform external field.Source: §E.3.1 · p. 39

E.3.2 Sampler

The sampler details specify the time grids, D-VCG basis variants, and reward ramp used for Ising experiments. They also identify a limitation of the empirical-cloud HEU implementation.

E.3.2 Sampler

  • Annealing uses 5000 time steps, while reward-tilt and joint experiments use 2000; frozen per-site Ising generators are exponentiated exactly at each step.This removes Euler error for the frozen site update.Source: §E.3.2 · p. 39
  • The reward strength ramps linearly from zero at the noisy end to β_r at the data end; pure annealing sets β_t to zero.The increasing reverse-time ramp gives data-end potential r_T = β_r r.Source: §E.3.2 · p. 39
  • D-VCG uses a small library of nonnegative basis rate families, with the chosen bases varying by tilt regime.The “2-basis” and “4-basis” labels are controller variants, not uniform basis counts across regimes.Source: §E.3.2 · p. 40
  • The empirical-cloud HEU evaluation is omitted for Ising because most candidate single-flip destinations have zero empirical mass; an analytical expansion is a proposed next step for DLM-scale use.The expansion resolves the issue without extra score-network evaluations.Source: §E.3.2 · p. 40

E.3.3 Metrics and aggregation

The Ising evaluation uses distributional distances, row-correlation error, and mean observables against Swendsen–Wang references. Results are averaged across seeds, with improvement factors aggregated within each setting and then across settings.

E.3.3 Metrics and aggregation

  • W2(|m|) for annealing or W2(M) for reward-tilt measures the one-dimensional Wasserstein-2 distance from the corresponding reference magnetization distribution.The reference consists of Swendsen–Wang samples.Source: §E.3.3 · p. 41
  • W2(Etotal) measures the one-dimensional Wasserstein-2 distance for total energy H(σ), which is sensitive to the joint distribution.Source: §E.3.3 · p. 41
  • MSE(C(r)) compares mean uncentered row correlations across r, excluding a one-site boundary margin and without subtracting magnetization.The squared difference between generated and reference estimates is averaged over r.Source: §E.3.3 · p. 41
  • Mean |m| or M is plotted against the Swendsen–Wang reference as a visual sanity check.Source: §E.3.3 · p. 41
  • Ising curves show three-seed means with across-seed standard-deviation bands; improvement factors use within-setting seed-mean error ratios before geometric means or peaks across settings.Source: §E.3.3 · p. 41

E.3.4 Configurations and additional results

The appendix extends the Ising tests across configurations and particle counts, then evaluates D-VCG on larger-scale settings, including monophonic music infilling. The reported comparisons show advantages over D-FKC or SGDD in the tested regimes.

E.3.4 Configurations and additional results

  • At N=5000, D-FKC still loses to D-VCG at N=500 by 5.4× on MSE(C(r)); the D-VCG advantage also holds at a second training temperature and in pure tilt.The pure-tilt configuration sets γ=1 and λeff=0, removing the joint-mode anchor and regularizer.Source: Appendix experiments · p. 41, §E.3.4 · p. 41, Figure 4 · p. 42

Appendix experiments.

  • Across N=100–5000, D-VCG (4-basis) reduces MSE(C(r)) over D-FKC by 15.21× geometrically, peaking at 34.7× at N=1000.The particle-count sweep uses Config 3; the gap does not close over the tested range.Source: Figure 5 · p. 42
  • At N=500, D-VCG (4-basis) has per-step wall-clock within roughly 5% of D-FKC on a single A100.The matrix-exponential integrator and score-network forward dominate cost; the active-set QP solve is negligible.Source: Appendix experiments · p. 42

F Additional Large-Scale Experiments

F.1 Discrete inverse problem on monophonic music (reward-tilted sampling)

  • On monophonic music infilling, D-VCG and D-FKC substantially outperform SGDD in Hellinger distance at both masking ratios, while D-VCG dominates both baselines on measurement error.The comparison fixes total function evaluations at 1024; Table 2 reports the quantitative results, with lower values better.Source: §F.1 · p. 43, Table 2 · p. 43

F.2 Text-to-image generation (classifier-free guidance)

The text-to-image experiment tests D-VCG with classifier-free guidance against D-FKC using 100 prompts and two alignment scores. D-VCG scores higher on both metrics, with qualitative examples also reporting closer prompt adherence.

F.2 Text-to-image generation (classifier-free guidance)

  • The evaluation uses Meissonic with 100 prompts across four art categories, comparing D-VCG only with D-FKC because prior work found D-FKC already dominates earlier CFG baselines.Both methods use 64 denoising steps and 8 particles.Source: §F.2 · pp. 43–44
  • D-VCG uses CFG-specific bases, including the unconditional reverse rate and target-aligned anchor; despite γ<0, its construction yields a nonnegative effective rate.The CFG-anchored Feynman–Kac proposal is recovered from a local-ratio decomposition.Source: §F.2 · pp. 43–44
  • D-VCG scores higher than D-FKC on both metrics: mean MPS is 15.422 vs. 15.282 and mean HPSV2 is 0.3033 vs. 0.2980.Table 3 reports means and standard errors over 100 prompts; higher is better.Source: Table 3 · p. 44, §F.2 · p. 44
  • Across six representative prompts, D-VCG generally adheres more closely to prompts; for the dragon-rider prompt, it renders the rider and scarf omitted by D-FKC.Source: §F.2 · p. 44, Table 4 · p. 45

Related papers

Loading 2609.35947v1…