Discrete Feynman-Kac Correctors

1 INTRODUCTION

• Figure 1 contrasts standard masked-diffusion inference with DFKC’s inference-time sampling from modified distributions.Its examples include annealed, product or geometric-average, and reward-tilted distributions.Source: Figure 1 · p. 2
• DFKC addresses the gap that Feynman-Kac Correctors derived for Fokker–Planck equations do not directly apply to discrete diffusion’s CTMC processes.The cited continuous-process approach uses weighted stochastic differential equations with SMC resampling (Skreta et al., 2025).Source: §1 · p. 2
• DFKC uses inference-time transformations to sample temperature-annealed, product, geometrically averaged, and reward-tilted distributions from trained discrete diffusion models.The framework also covers differently conditioned processes and requires no additional model training or fine-tuning.Source: §1 · p. 2
• DFKC is applied to efficient lower-temperature Ising sampling, language-model programming and amortized-learning tasks, and protein sequence generation with external rewards.Source: §1 · p. 2

2.1 SIMULATING FORWARD KOLMOGOROV EQUATION (FKE)

• A forward Kolmogorov equation specifies marginal evolution through a rate matrix and initial distribution, with off-diagonal rates nonnegative and diagonal rates nonpositive.The diagonal can be set to the negative sum of off-diagonal rates.Source: §2.1 · p. 3
• The FKE can be simulated by sampling successive categorical state transitions from the rate matrix over discretized time steps.Source: §2.1 · p. 3
• Adding weights to state trajectories supports weighted simulation and self-normalized importance sampling.The state transitions and log weights are updated jointly.Source: §2.1 · p. 4

2 BACKGROUND

• A discrete diffusion process maps the data distribution to a tractable marginal and learns a reverse process; masked diffusion reserves a discrete state for the mask.Source: §2 · p. 3, §2.2 · p. 4

2.2 DISCRETE MASKED DIFFUSION

• Masked diffusion generates data samples by simulating the reverse-time process from the tractable marginal; its reverse rates can use probability ratios or de-masking probabilities.In the multidimensional process, transitions occur only when all coordinates except one match.Source: §2.2 · p. 4

3 DISCRETE FEYNMAN-KAC CORRECTORS

• The framework transforms CTMC marginals and derives corresponding rate matrices to control masked diffusion at inference time, without additional training or fine-tuning.Its simulations require only marginal-density ratios, equivalently denoising conditional probabilities.Source: §3 · p. 5

3.1 TEMPERATURE ANNEALING1

• Temperature annealing evolves marginals proportional to p_t(i)^β using the original rate matrix and marginal-probability ratios.For masked diffusion, the rate matrix is scaled by β and the probability ratio raised to β; the weighting term needs no additional function evaluations.Source: §3.1 · p. 5

3.2 PRODUCT AND GEOMETRIC AVERAGING2

• Product and geometric-average transformations target samples likely under multiple CTMC marginals, with product-process rates and weights evaluable in one forward pass per network.The product interpretation requires agreement: a zero probability under any model makes the product zero.Source: §3.2 · p. 6

3.3 REWARD-TILTED MARGINALS3

• Reward tilting targets marginals proportional to p_t(i) exp(β_t r(i)), but rate and weight evaluation can require rewards over all states reachable from the mask.Computing these quantities may be expensive depending on the application; transition reward differences are proposed as a possible alternative for future work.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

• Across experiments, DFKC modifies the rate matrix and uses sample reweighting and resampling to control generation for the target densities.The experiments illustrate annealing, geometric averaging, and reward-tilting; inference uses SNIS resampling between update steps.Source: §4 · p. 7

4.2 PRODUCTS AND ANNEALING FOR LANGUAGE MODELLING

• DFKC achieves higher accuracy on coding tasks than Argmax sampling and Naive Annealing.The evaluation uses HumanEval and MBPP; Naive Annealing is equivalent to DFKC with one SMC sample, so the improvement demonstrates the importance of resampling.Source: §4.2 · p. 9, Figure 4 · p. 9
• As dataset prompts grow in length and complexity, joint prompting degrades while the DFKC product remains more stable; using five samples slightly improves performance over one.The task infers parameters for a noisy linear predictor from a synthetic dataset partitioned into five subsets.Source: §4.2 · p. 9, Figure 3 · p. 9

4.3 GUIDING PROTEIN SEQUENCE GENERATION WITH EXTERNAL REWARDS

• For protein generation, multiple-sample DFKC yields higher mean reward than unguided DPLM sampling and guidance without resampling.Likelihood-guided sampling produces more designable sequences at the cost of diversity; thermostability guidance maintains similar performance to the base model on other metrics.Source: §4.3 · p. 10, Table 2 · p. 10

Reward Fine-tuning

• Reward fine-tuning adjusts pretrained parameters with reinforcement learning, whereas DFKC leaves them fixed and avoids a costly fine-tuning stage; DFKC can also use reward-fine-tuned models.The cited discrete-diffusion reward-fine-tuning methods include Venkatraman et al. (2024), Rector-Brooks et al. (2024), and Wang et al. (2025).Source: §5 · p. 10

Inference Time Alignment

• Inference-time alignment methods may yield biased samples through approximation, while DFKC adapts an unbiased SMC strategy to smoothly annealed rewards and general products.The passage cites Vignac et al. (2022), Gruver et al. (2023), Nisonoff et al. (2024), and Tang et al. (2025) for approximation-based methods, and Singhal et al. (2025) for SMC reward sampling.Source: §5 · p. 10

Boltzmann distribution annealing

• Unlike samplers trained across temperatures or with energy access during training, DFKC reuses a trained model to generate samples at different unseen temperatures; more accurate convergence analysis remains open.The paper proves convergence of its time-discretization scheme and points to existing theoretical tools as possible aids to sharper analysis.Source: §5 · p. 11

• DFKC repurposes discrete diffusion models at inference time without retraining them.The framework combines learned probability ratios with SMC to sample from annealed, product, or reward-weighted distributions.Source: §6 · p. 11
• The empirical study supports DFKC for lower-temperature Ising sampling, programming tasks and parameter inference with language models, and protein-sequence control.Source: §6 · p. 11
• Future work includes extending DFKC to joint continuous and discrete models and combining it with reward fine-tuning.Source: §6 · p. 11

• The authors made the code publicly available.The passage provides a GitHub link.Source: REPRODUCIBILITY STATEMENT · p. 11
• The paper describes mathematical and algorithmic details needed to reproduce the results.Algorithm 1 is given as an example.Source: REPRODUCIBILITY STATEMENT · p. 11
• The authors state that the code facilitates reproducibility of the empirical results and algorithm.Source: REPRODUCIBILITY STATEMENT · p. 11

A.1 CORRESPONDENCE TO CONTINUOUS FEYNMAN-KAC CORRECTORS

• Table A1 compares state updates for discrete and continuous Feynman-Kac Correctors; its discrete reward-difference state replaces the token position being de-masked.The replacement token is drawn from the vocabulary.Source: Table A1 · p. 16
• For the reward-tilted case, the state used to compute the reward difference is the current state with the de-masked position replaced by a vocabulary token.Source: Table A1 · p. 16
• Table A2 compares the methods’ weight updates at line 5 of Algorithm 1.Source: Table A2 · p. 16

B.1 WEIGHTED FORWARD KOLMOGOROV EQUATION

• The weighted forward Kolmogorov equation is introduced with a weighting term and then rewritten to express its final term differently.Source: §B.1 · p. 17

B.2 DISCRETE FEYNMAN-KAC FORMULA

• The discrete Feynman-Kac formula relates weighted marginal evolution to trajectories; its proof uses a matrix ODE and a time-ordered exponential.The matrix exponential is split using a time-dependent Lie-Trotter formula, although the rate and weight matrices do not generally commute.Source: §B.2 · pp. 17–18

B.3 DISCRETE MASKED DIFFUSION

• The masked-diffusion transition schedule must satisfy composition relations across intermediate times to be consistent with the master equation.Source: §B.3 · p. 19

B.4 REVERSE-TIME MASKED DIFFUSION

• Reverse-time masked diffusion is derived by flipping the marginal time parameter and differentiating with respect to reverse time.Source: §B.4 · p. 20

B.5 DE-MASKING PARAMETERIZATION

• The de-masking parameterization is derived from conditional probabilities involving the masked state and the original token.Source: §B.5 · p. 20

B.6 MULTIDIMENSIONAL CASE

• For multidimensional masked diffusion, coordinates are masked independently, and the resulting rate matrix describes changes in one coordinate at a time.Source: §B.6 · p. 21

C.1 ANNEALING OF FKE

• Annealing transforms marginals to qt(i) proportional to pt(i)^β and yields corresponding dynamics and weights for reverse-time masked diffusion.Source: §C.1 · pp. 22–24

C.2 PRODUCT OF FKES

• For two processes, the product of their marginals has derived dynamics and weights, with a corresponding result for reverse-time masked diffusions.Source: §C.2 · pp. 24–25

C.3 GEOMETRIC AVERAGE OF FKES

• The geometric average of multiple process marginals is given its own evolution equation and corresponding masked-diffusion result.Source: §C.3 · pp. 26–27

C.4 REWARD-TILTED FKE

• Reward tilting defines marginals proportional to pt(i) exp(βt r(i)) and derives their dynamics, including a reverse-time masked-diffusion case.Source: §C.4 · pp. 28–29

D.1 ANNEALING THE ISING MODEL

• The Ising-model experiments specify training and sampling configurations and evaluate global energy and magnetization distributions alongside local correlation structure.The stated metrics include 2-Wasserstein distances for energy and magnetization distributions and MSE for the correlation function.Source: §D.1 · p. 30

D.2.1 THEORETICAL JUSTIFICATION

• With a uniform prior, the amortized-learning posterior over parameters is expressed as a product across data subsets.Source: §D.2.1 · p. 31

D.2.2 EXPERIMENTAL SETUP

• The synthetic linear-regression evaluation uses five data subsets, LLaDA inference, and five SMC samples selected using MSE.The generated data use a noisy linear predictor; predictions are made in one block with a 128-token cap.Source: §D.2.2 · p. 31

D.2.3 ADDITIONAL RESULTS FOR AMORTIZED LEARNING

• More SMC samples improve amortized-learning performance up to a threshold of eight, while product prompting has lower MSE than joint prompting and resampling improves it further.The outputs from joint prompting often fail the requested format or cannot be parsed, whereas product-prompt outputs did not show this issue.Source: §D.2.3 · p. 31, Figure A1 · p. 32

D.3 CODE GENERATION

• On HumanEval and MBPP, the code-generation comparison uses naive annealing as the one-sample, no-resampling version of DFKC, so its improvement over that baseline reflects the contribution of resampling.The results are tabulated in Table A4, whose supplied caption identifies accuracy and standard errors but gives no values.Source: §D.3 · p. 33, Table A4 · p. 33

D.4 MULTI-CONSTRAINT STORY GENERATION

• For multi-constraint story generation, the authors compare joint prompting with products over individual constraints using Qwen2.5 perplexity across different numbers of constraints.Source: §D.4 · p. 33, Figure A3 · p. 33

D.5 PROTEIN SEQUENCE GENERATION

• The protein-generation experiments use DPLM1 650M and evaluate ESM2-likelihood and thermostability rewards across sequences of several lengths.The ESM2-likelihood reward is based on an estimated pseudo log-likelihood; reward evaluation for partially masked sequences first denoises them to an unmasked sequence.Source: §D.5 · p. 34, §D.5.1 · pp. 34–35

D.5.1 REWARD MODELS

• The reward models include an ESM2 pseudo log-likelihood estimate and a fine-tuned DPLM-650M model whose log predicted thermostability is used as reward.Source: §D.5.1 · p. 34

D.5.2 BASELINES

• DG-Exact is equivalent to DFKC without resampling, while FK Steering uses the base model as proposal and a difference potential for resampling.Source: §D.5.2 · p. 35

D.5.3 DESCRIPTION OF PROTEIN METRICS

• Protein metrics cover sequence diversity, predicted structural confidence and quality, and novelty relative to known structures.The appendix defines diversity through alignment and clustering, and novelty using TM-score comparisons against a structure database.Source: §D.5.3 · pp. 35–36

D.5.4 HYPERPARAMETERS

• Reward scales and SMC sample counts are selected through validation experiments, with the best-performing sample count reported for each method.Source: §D.5.4 · p. 36

D.5.5 ADDITIONAL PROTEIN RESULTS

• Table 2 reports the highest-reward settings for each method.The supplied passages do not provide the corresponding values.Source: §D.5.5 · p. 36
• Using five particles instead of one improves protein reward, and DFKC generally outperforms both the base model and DG-Exact.Figure A4 reports ESM2 likelihood and thermostability across evaluated sequence lengths and particle counts.Source: Figure A4 · p. 36, §D.5.5 · p. 36
• Table A6 varies the number of SMC samples for FK Steering and DFKC on the ESM2 likelihood and Thermostability tasks.The table reports means and standard errors over five seeds.Source: §D.5.5 · p. 36, Table A6 · p. 38