July 01, 2026
While generative models have enabled training-free reward alignment, current methods typically excel in local exploration within narrow regions of the underlying distribution. These approaches struggle when preferences are unknown a priori and only revealed through sequential feedback—a scenario demanding broad exploration to uncover high-utility regions. To address this, we propose Sequentially-Controlled Interactive Multi-Particle Flow-Maps (IMPFM), a framework for sample-efficient online feedback-driven search. IMPFM progressively transports a group of interactive particles toward the target distribution, maintaining the broad coverage essential for heterogeneous preference alignment. IMPFM introduces a principled and efficient posterior sample sharing mechanism across particles powered by flow maps. By correcting individual particle drift with the collective posterior samples of the entire ensemble at each resampling step, the framework maximizes sample utility to enable global exploration while actively mitigating reward over-optimization, typical of standard control frameworks. Paired with a principled exploration-exploitation reweighting mechanism involving multi-particle interaction, this sequentially corrected multi-particle dynamics explicitly preserves structural diversity and overcomes the weight degeneracy inherent to standard SMC samplers. Crucially, we prove that the resulting sampling framework yields a multi-particle interaction-aware Feynman-Kac corrector that progressively steers the multi-particle system toward a KL-tilted target distribution, facilitating global exploration and preventing mode collapse. Extensive empirical evaluations and rigorous ablations across diverse search and alignment tasks confirm the efficacy of IMPFM over existing baselines.
Across diverse scientific and engineering disciplines, the pace of discovery is fundamentally constrained by the high cost of evaluation. This bottleneck gives rise to the ubiquitous challenge of online feedback-driven search, where finding an optimal solution requires iteratively proposing candidates and adapting based on sequential feedback. Whether designing life-saving therapeutics or optimizing interactive visual recommendation systems, these real-world applications often rely on costly physical experiments or human engagement. Consequently, minimizing the number of such interactions is a central requirement. Accelerating breakthroughs in these domains demands a robust sampling framework capable of strategic, global exploration across complex, high-dimensional spaces—such as molecular structures or natural images. Crucially, such a framework must seamlessly integrate sequential feedback to optimize the search trajectory under strict sampling budgets.
Existing methods remain fundamentally ill-equipped for this challenge. For instance, online RL fine-tuning of diffusion models [1] relies on a mode-seeking KL-divergence objective that inherently fails to capture diverse, high-utility regions. Furthermore, these approaches depend on online-trained reward models, whose early-stage biases frequently misdirect fine-tuning and degrade sample quality. Conversely, recent Sequential Monte Carlo (SMC) inference-time scaling methods [2], [3] avoid fine-tuning by leveraging fixed generative priors. Yet, they suffer from severe weight degeneracy, leading to proposal collapse and diminished diversity [4]. Tree-based samplers [5], [6] attempt to mitigate these evaluation bottlenecks by backpropagating only terminal rewards, thus bypassing intermediate evaluations. However, existing tree- and SMC-based samplers rely rigidly on prior generative dynamics, failing to actively correct search trajectories using sequential feedback. This passive exploration inevitably yields suboptimal performance.
Recently, Feynman-Kac correctors [7] (FKCs) have emerged as a powerful alternative—progressively correcting dynamics at every resampling step—making them a highly promising paradigm for rapid, feedback-driven adaptation. However, these approaches rely on reward values or gradients at every resampling step; in practice, these signals are unavailable and must be approximated, introducing systematic bias into the sampler and degrading sample quality. Crucially, standard FKC samplers transport particles in strict isolation. This lack of inter-particle interaction severely bottlenecks broad strategic exploration. To overcome these bottlenecks, we introduce Interactive Multi-Particle Flow-Maps (IMPFM), a principled sampling framework that rapidly adapts to online feedback, maintains broad distributional coverage, and maximizes sample efficiency—the core prerequisites for budget-constrained, online feedback-driven search. Figure [fig:teaser] provides a conceptual overview of IMPFM.
IMPFM is built on the premise that discovery thrives not in isolation, but through collective intelligence. It recognizes that collaboration is the ultimate engine for expanding information diversity and unlocking truly sample-efficient search.
At the heart of IMPFM is a particle interaction mechanism that drives a dual-force exploration dynamic, combining an attractive pull toward high-utility regions with a repulsive push that prevents mode collapse. Powered by Flow Maps, IMPFM enables each particle to exploit the ensemble’s collective posterior to optimally correct its drift at every step, and this information reuse mechanism facilitates highly-informed, feedback-efficient search. IMPFM’s power is substantially enhanced when paired with a principled reweighting mechanism at each resampling step. This reweighting encodes rich multi-particle interactions and enforces a dual-force dynamic—attracting particles to high-density targets while simultaneously exerting a repulsive push to prevent mode collapse and tame reward overoptimization, the classic pitfall of reweighting-based samplers. Critically, we show that IMPFM induces interaction-aware Feynman–Kac correctors that rapidly correct the search dynamics based on online feedback and drive the particle system toward the KL-tilted marginal distribution, with the true target as its reference measure. These findings show that, rather than treating samples independently and targeting a reward-tilted distribution directly, it is more powerful to reason about the particle ensemble as a whole and minimize a distributional distance. This ensemble-level view provides a broader picture, promotes global convergence, and yields substantially wider coverage.
Interactive FKC Sampler: We introduce IMPFM, a Multi-Particle Interaction-aware Feynman–Kac Corrector for online feedback-driven search.
Flow-Map-Driven Information Reuse: IMPFM introduces a collective posterior-sharing mechanism powered by Flow Maps, maximizing the utility of every drawn sample, transforming isolated particle updates into a highly-informed, globally coordinated search.
Rapid Adaptation and Coverage: Driven by a dual-force (attractive-repulsive) exploration dynamic, and by coupling collaborative drift correction with dynamic reweighting, IMPFM ensures rapid adaptation to online feedback and expansive search space coverage.
Rigorous Empirical Validation: We validate the effectiveness of each component of IMPFM through comprehensive quantitative and qualitative ablation studies.
Background: The cornerstone of effective online feedback-driven search is a policy capable of efficiently sampling from the look-ahead value-tilted distribution. Formally, under optimal control, the marginal density \(p_t^*\) of \(X_t^*\) at timestep \(t\) satisfies: \[\small{p_t^*(x) \propto p_t(x) e^{V_t(x)}, \quad \text{where} \quad V_t(x) = \log \mathbb{E} \left[ e^{r(X_1)} \mid X_t = x \right]}\] Here, \(p_t\) denotes the marginal density associated with the unconditional prior process. This formulation ensures that the terminal marginal distribution, \(p_1^*\), aligns exactly with the target reward-tilted distribution \(p_{\text{reward}}(x) \propto p_1(x) e^{r(x)}\). Leveraging Doob’s \(h\)-transform [8], [9], the optimal drift \(\smash{b_t^*}\) can be expressed as: \[\label{eq:oc} \small{\dot{x}_t^* = \underbrace{b_t(x_t^*) + \frac{\sigma_t^2}{2} \nabla V_t(x_t^*)}_{b_t^*(x_t^*)}, \quad \text{with} \quad \text{Law}(x_t^*) = p_t^*, \quad x_0^* \sim p_0}\tag{1}\] While \(b_t\) is directly accessible from the pre-trained Flow Matching model, the primary algorithmic bottleneck lies in efficiently estimating the value function gradient, \(\nabla V_t(x)\). A consistent Monte Carlo estimator of \(\nabla V_t(x)\) is: \[\label{eq:grad95v} \small{\nabla_{x_t} V_t(x_t) = \nabla_{x_t} \log \left( \frac{1}{N} \sum_{i=1}^N \exp \left[ r(\hat{x}^i_1) \right] \right), \quad \tcbox[on line, boxsep=0pt, left=2pt, right=2pt, top=2pt, bottom=2pt, colback=gray!10, colframe=gray!50, arc=2pt]{\hat{x}^i_1 \overset{\text{iid}} {\sim} p_{1|t}(\cdot|x_t)}}\tag{2}\] Hence, given the advantages of ODE sampling, the primary challenge boils down to efficiently drawing samples from \(p_{1|t}(\cdot|x_t)\) via an ODE-based flow model. Conditioned on \(x_t\), one can theoretically define an ODE that transports the prior \(p_0\) directly to the posterior \(p_{1|t}(\cdot|x_t)\). The associated drift field, \(b_s(\cdot;x_t)\), is derived via standard flow matching [10], with the crucial distinction that the objective targets the conditional posterior \(p_{1|t}(\cdot|x_t)\) rather than the marginal data distribution \(p_1\): \[\small{{b}_s(x; x_t) = \mathbb{E}\left[\dot{\alpha}_sI_1 + \dot{\beta}_sI_0 \mid I_s = x\right], \quad I_s = \alpha_sI_1 + \beta_sI_0, \quad I_0 \sim \mathcal{N}(0, I_d) , \quad I_1 \sim p_{1|t}(\cdot|x_t)}\] Consequently, the probability flow associated with \(b_s(\cdot; x_t)\) satisfies: \[\small{\frac{d}{ds}x_s = b_s(x_s; x_t), \quad x_0 \sim \mathcal{N}(0, I_d) \implies \text{Law}(x_1) = p_{1|t}(\cdot|x_t)} \label{eq:ode95cond}\tag{3}\] Where \(\alpha_t\) and \(\beta_t\) denote time-dependent coefficients with boundary conditions \(\alpha_0 = \beta_1 = 0\) and \(\alpha_1 = \beta_0 = 1\). Under the assumption of a Gaussian prior \(p_0\), the conditional drift \(b_s\) admits an analytical derivation [11] via a reparameterization of the unconditional drift \(b_t\): \[b_s(x_s; x_t) = w_1x_s + w_2 b_{t^*}(S(x_s, x_t)) + w_3 x_t\] Where, \(w_1, w_2,\) and \(w_3\) represent scalar weights. Furthermore, \(S\) acts as a linear sufficient statistic, while \(t^*\) represents a reparameterized time variable: \[\small{S_{s,t}(x_s, x_t) = \frac{\alpha_s \sigma^2_t x_s + \alpha_t \sigma^2_s x_t}{\sigma^2_t \alpha^2_s + \alpha^2_t \sigma^2_s}, \qquad t^*(s, t) = g^{-1}\left( \frac{\sigma^2_t \sigma^2_s}{\sigma^2_t \alpha^2_s + \alpha^2_t \sigma^2_s} \right), \quad g(t) = \frac{\sigma^2_t}{\alpha^2_t}} \label{eq:combined95reparam}\tag{4}\] While this reparameterization renders the conditional drift \(b_s\) tractable, integrating the ODE trajectories to draw posterior samples from \(p_{1|t}(\cdot|x_t)\) remains a significant computational bottleneck. To bypass this computational bottleneck, a standard strategy is teacher-distillation, which leverages the exact analytical field \(b_s(x_s; x_t)\) as the teacher for the distillation process. The student velocity model, \(\hat{v}\), is optimized via a joint objective function. The first component is an instantaneous loss that minimizes the mean squared error between the student’s predicted velocity, \(\hat{v}_{s,s}\), and the exact analytical teacher field: \[\small{\mathcal{L}_{inst}^{distill}(\hat{v}) := \int_0^1 \int_0^1 \mathbb{E} ||\hat{v}_{s,s}(x; x_t) - b_s(x; x_t)||^2 ds dt}\] This is paired with a consistency loss [12] (\(\mathcal{L}^{distill}_{cons}\)) to ensure the student reliably integrates this velocity across arbitrary time intervals, defined mathematically as: \[\small{\int_{0}^{1} \int_{0}^{u} \mathbb{E} \left\| \hat{v}_{s,u}(I_{s}; I'_t) - \text{sg}\left( b_u (\hat{X}_{s,u}(I_{s}; I'_t)); I'_t \right) - (u - s)\partial_u \hat{v}_{s,u}(I_{s}; I'_t) \right\|^2 ds \, du}\] Here, \(X_{s,u}(\cdot \, ; x_t) : \mathbb{R}^d \to \mathbb{R}^d\) denotes the solution operator for the context-dependent ODE defined in 3 . This distillation process effectively collapses the iterative, multi-step ODE integration into a single amortized map via \(\hat{v}\), enabling computationally efficient, one-step sampling directly from the conditional posterior \(p_{1|t}(\cdot|x_t)\), thereby streamlining the evaluation of \(\nabla V_t(x)\) as defined in 2 .
Proposed Approach: Merely repeatedly simulating the optimal drift \(b_t^*\) is inadequate for efficient feedback-driven search: solely value-guided trajectories are inherently exploitative and lack an explicit exploratory mechanism, causing them to miss disjoint, high-utility modes within the target landscape. Furthermore, this approach is highly vulnerable to reward over-optimization unless stabilized by an impractically large ensemble of Monte Carlo posterior samples, forcing an untenable trade-off under strict sampling budgets. In particular, [13] shows that the distance to the true target distribution is bounded by the inverse of the number of Monte Carlo posterior samples utilized to compute \(\nabla V_t\) at each correction step.
To resolve this, we introduce an interactive multi-particle system that bypasses the sampling bottleneck by sharing posterior samples across an ensemble. By pooling this global information to calibrate the drift of individual particles, our framework ensures broad exploration and effectively mitigates reward over-optimization. Specifically, we introduce a system of interacting particles, where the dynamics of each particle (say \(i\)-th particle) is actively steered according to the following rule: \[\label{eq:dynamics} \small{\dot{x}_t^i = b_t(x_t^{i}) + \underbrace{\frac{\sigma_t^2}{2} \nabla_{x^i_t} V_t(x_t^{i})}_{\text{\boldsymbol{Exploitation}}} + \tcbox[on line, boxsep=0pt, left=2pt, right=2pt, top=5pt, bottom=5pt, colback=gray!10, colframe=gray!50, arc=2pt]{\displaystyle \underbrace{\frac{\sigma_t^2}{2}\frac{1}{n-1} \sum_{j=1, j \neq i}^n \left[ \underbrace{k(x^j_t, x^{i}_t) \nabla_{x^j_t} V(x^j_t)}_{\text{Pulling towards high utility region}} + \underbrace{\nabla_{x^j_t} k(x^j_t, x^{i}_t) }_{\text{Repulsive force induce diversity}}\right]}_{\text{\boldsymbol{Exploration}: Interaction between particles encourage broad coverage of high utility regions}}}}\tag{5}\] Here, \(n\) represents the number of particles, \(k(x^i_t, x^j_t)\) denotes the Radial Basis Function (RBF) kernel, defined as \(k(x^i_t, x^j_t) = \exp \left( -\frac{1}{h} \| x^i_t - x^j_t \|_2^2 \right)\), where \(h\) is the bandwidth parameter, and \(\nabla V_t(x)\) is efficiently computed via Flow-Map following 2 . In Proposition 4.2 below, we demonstrate that the interplay between the exploitation and exploration terms in Eq. 5 induces a joint steering of the particle dynamics that achieves a steepest descent in the KL divergence between the empirical particle distribution and the true target density.
Let \(q(x)\) denote the empirical distribution induced by the interacting particles, and define the Boltzmann distribution as \(p_{\text{boltz}}(x) \propto e^{V(x)}\). The steepest descent direction that minimizes the Kullback-Leibler divergence, i.e., \(\nabla_x D_{\text{KL}}(q \parallel p_{\text{boltz}})\) within the unit ball of a Reproducing Kernel Hilbert Space (RKHS) is given by: \[\label{eq:stein95} \small{\mathbb{E}_{x^i \sim q} \left[ \nabla_{x^i} V(x^{i}) + \frac{1}{n-1} \sum_{j=1, j \neq i}^n \left[ k(x^j, x^{i}) \nabla_{x^j} V(x^j) + \nabla_{x^j} k(x^j, x^{i}) \right] \right]}\tag{6}\]
Note that the exploration mechanism is driven by a dual-force mechanism. The first component executes a kernel-weighted gradient ascent, steering the ensemble toward high-utility manifolds. Complementing this, the second component exerts a repulsive force that prevents particle collapse into local modes, ensuring the ensemble maintains a broad, diverse coverage of the target distribution’s landscape. Equation 5 reveals a powerful collaborative mechanism: the drift of the \(i\)-th particle is calibrated by the value gradients of its peers. Since computing the value gradient fundamentally requires drawing multiple Monte Carlo samples from a particle’s corresponding posterior. Hence, by aggregating Monte Carlo posterior samples across the entire ensemble, the framework effectively pools global information to refine individual trajectories. This collective sharing transforms the estimation of \(\nabla V_t\) from a local, high-variance task into a robust, ensemble-wide operation, ensuring highly informed corrections and accelerated convergence to the target distribution.
While highly efficient, ODE sampling inherently restricts exploration, necessitating stochastic transitions for effective search. A standard approach introduces this stochasticity by converting the ODE into a corresponding SDE, steering each particle via the following augmented dynamics: \[\label{eq:sde} \small{ dx^i_t = b_t(x^i_t)\>dt + \sigma_t^2\nabla_{x^i_t} V_t(x_t^{i}) + \sigma_t^2\frac{1}{n-1} \sum_{j=1, j\neq i}^n \left[ k(x^j_t, x^{i}_t) \nabla_{x^j_t} V(x^j_t) \> dt + \nabla_{x^j_t} k(x^j_t, x^{i}_t) \> dt \right] + \sigma_tdW_t}\tag{7}\] While standard SDE conversion enables stochasticity, it sacrifices the efficiency of ODE and lacks the dynamic transition steps necessary for online feedback-efficient search. Consequently, we seek a mechanism to inject arbitrary-duration stochastic transitions while strictly preserves ODE-level efficiency. To achieve this, we first sample standard noise \(\epsilon \sim \mathcal{N}(0, I)\) and push it forward through the distilled flow map \(\hat{v}(\epsilon; x_t)\). By definition, this directly yields exact samples from the true posterior, \(x_1 \sim p_{\hat{v}_{\boldsymbol{IMPD}}}(x_1 | x_t)\). With the trajectory’s boundary points (\(\epsilon\) and \(x_1\)) firmly anchored, we induce a DDPM-style transition via the Sufficient Statistic [11]. Specifically, we perform the following steps: \[\begin{align} & \small{\epsilon \xrightarrow{\hat{v}(\epsilon; x_t)} x_1 \sim p_{\hat{v}_{\boldsymbol{IMPD}}(\epsilon;x_t)}(x_1 | x_t) \xrightarrow{\text{Interp.}} x_{r^\ast} = \bar\alpha_{r^\ast} x_1 + \bar\sigma_{r^\ast} \varepsilon_1 \xrightarrow{\text{Trans.}} x_{t'} = \alpha_{t'} \frac{ \bar{\alpha}_{r^*} \sigma_t^2 x_{r^*} + \alpha_t \bar{\sigma}_{r^*}^2 x_t}{\bar{\alpha}_{r^*}^2 \sigma_t^2 + \alpha_t^2 \bar{\sigma}_{r^*}^2},} \\[-6pt] & \small{r^*(t, t') = g^{-1}\left( \frac{g(t)g(t')}{g(t) - g(t')} \right), \quad g(t) = \frac{\sigma^2_t}{\alpha^2_t}, \quad \varepsilon \sim \mathcal{N}(0, I), \quad \varepsilon_1 \sim \mathcal{N}(0, I).} \end{align} \label{eq:bss}\tag{8}\]
Crucially, we first generate the posterior sample (i.e., \(x_1\)) by evolving each particle according to the corrected dynamics \(\hat{v}_{\boldsymbol{IMPD}}\) defined in Eq. 5 (IMPD stands for Interactive Multi Particle Dynamics). The sampled \(x_1\) then serves as fixed boundary points to compute the stochastic interpolant \(x_{r^*}\), which acts as the essential foundation for constructing the final DDPM-style transition sample \(x_{t'}\). Armed with DDPM-style transitions, the interactive drift-correction mechanism naturally facilitates a FKC-style reweighting scheme, further amplifying the sampler’s capability. While standard value-based reweighting greedily preserves high-reward trajectories at the expense of diversity—inevitably leading to weight degeneracy and mode collapse—we introduce a principled, interaction-aware alternative. By leveraging ensemble-wide multi-particle dynamics, our proposed reweighting mechanism actively preserves the structural diversity of high-utility samples, encouraging broad exploration and global convergence. Specifically, we introduce the following scoring function to weight each particle \[\label{eq:weight} \small{\underbrace{\boldsymbol{\langle} \nabla_{x_t^i} V(x_t^{i}), \boldsymbol{b_t(x_t^i)} \boldsymbol{\rangle} + \frac{1}{n-1} \sum_{j=1, j \neq i}^n \left[ \boldsymbol{\langle} k(x_t^j, x_t^{i}) \nabla_{x_t^j} V(x_t^j), \boldsymbol{b_t(x_t^i)} \boldsymbol{\rangle} + \boldsymbol{\langle} \nabla_{x_t^j} k(x_t^j, x_t^{i}), \boldsymbol{b_t(x_t^i)} \boldsymbol{\rangle} \right]}_{\boldsymbol{weight of particle i (i.e. w^{(i)}_t)}}}\tag{9}\] Intuitively, 9 prevents off-manifold degradation by projecting the value and diversity gradients onto the prior model dynamics \(b_t(x_t^i)\). A particle is weighted highly only if its pursuit of high reward—both locally and via the ensemble—aligns with the base model dynamics, preventing it from straying into low-likelihood, off-manifold regions. Similarly, the repulsive interaction term guarantees that the drive for sample diversity does not compromise sample quality, as exploratory trajectories are strictly penalized if they violate the prior model dynamics. To maintain sampling efficiency, we trigger a standard resampling mechanism whenever the effective sample size (ESS) falls below a defined threshold: \(\text{ESS}_t^n = \frac{\left( \sum_{i=1}^n w_t^{(i)} \right)^2}{\sum_{i=1}^n \left( w_t^{(i)} \right)^2} \in [1, n]\). Crucially, this resampling procedure leverages importance sampling to ensure the resulting samples remain asymptotically unbiased. Next, we demonstrate that coupling the interactive drift correction (Eq. 5 ) with our reweighting mechanism (Eq. 9 ) yields an interaction-aware Feynman-Kac corrector, as formalized in Proposition 4.3.
Consider a Flow model with drift \(b_t(x)\) trained to sample from \(q_t(x)\). Sampling from the KL-tilted marginals \(p_t^{\text{kl}}(x) \propto q_t(x) \exp\Big( -D_{\text{KL}}\big( q_t(x) \parallel p_{\text{boltz}}(x) \big) \Big)\) is performed by an weighted SDE with an equivalent corresponding weighted ODE: \[\small{dx^i_t = b_t(x^i_t)\>dt + \frac{\sigma_t^2}{2}\frac{1}{n} \sum_{j=1}^n \left[ k(x^j_t, x^{i}_t) \nabla_{x^j_t} V(x^j_t) \> dt + \nabla_{x^j_t} k(x^j_t, x^{i}_t) \> dt \right]}\] \[\small{dw^i_t = \frac{1}{n} \sum_{j=1}^n \left[ \boldsymbol{\langle} k(x_t^j, x_t^{i}) \nabla_{x_t^j} V(x_t^j), b_t(x_t^i) \boldsymbol{\rangle} \> dt + \boldsymbol{\langle} \nabla_{x_t^j} k(x_t^j, x_t^{i}), b_t(x_t^i) \boldsymbol{\rangle}\> dt \right]}\]
Thus, as shown in Proposition 4.3, our framework jointly steers the particle ensemble toward the true reward-tilted distribution while maintaining high likelihood under the prior to preserve intrinsic sample quality. By leveraging multi-particle interactions to explicitly minimize the KL divergence to the target, this holistic approach fundamentally circumvents the severe vulnerabilities of reward over-optimization and mode collapse inherent in isolated trajectory optimization. We term our framework Sequentially Controlled Interactive Multi-Particle Flow-Maps (IMPFM) (See Fig. 2). To broaden the applicability of the IMPFM to settings lacking an off-the-shelf stochastic flow map, we seamlessly adapt our approach using a standard deterministic flow matching model—governed by instantaneous drift \(u_t(x_t)\)—paired with a consistency model enabling a single-step sample (i.e., \(x_1\)) generation from any \(x_t\). To compute the corrected dynamics (Eq. 5 ) and reweighting scores (Eq. 9 ), we substitute the baseline drift \(b_t\) with \(u_t\) and leverage the consistency model to compute an unbiased, gradient-free MC estimator [14] for the required value gradient \(\nabla V_t\), formulated as: \[\small{ \frac{\sigma_t^2}{2} \nabla_{x_t} V_t(x_t) = \left( \dot{\beta}_t - \frac{\dot{\alpha}_t}{\alpha_t} \beta_t \right) \frac{\sum_{i=1}^N x_1^{(i)} \exp(r(x_1^{(i)}))}{\sum_{i=1}^N \exp(r(x_1^{(i)}))} + \frac{\dot{\alpha}_t}{\alpha_t} x_t - u_t(x_t), \quad x_1^{(i)} \stackrel{\text{i.i.d.}}{\sim} p_{1|t}(\cdot|x_t).} \label{eq:value95gradient}\tag{10}\] This gradient-free formulation is highly compute-efficient, entirely bypassing the need to backpropagate through the consistency model, and unlocks the ability to optimize non-differentiable, black-box reward functions. However, the inherent determinism of consistency models precludes sampling multiple \(x^{(i)}_1\) from the posterior \(p_{1|t}(\cdot|x_t)\), a critical requirement for estimating the value gradient \(\nabla V_t\). To bypass this limitation, we introduce an iterative posterior sampling scheme: we alternate between predicting \(x_1\) from \(x_t\) via the consistency model, and injecting noise via a stochastic interpolant to produce a perturbed \(x_t\). Repeating this sequence effectively yields valid samples (i.e., \(x_1^{(1)}, x_1^{(2)}, \ldots x_1^{(K),}\)) from the desired posterior, as depicted below (where \(\epsilon_1, \epsilon_2, \ldots, \epsilon_k \stackrel{\text{i.i.d.}}{\sim} \mathcal{N}(0, I)\)): \[x_t^{(0)} \xrightarrow{\text{Consistency}} \textcolor{red}{x_1^{(1)}} \xrightarrow{x_t^{(1)} = \alpha_t x_1^{(1)} + \sigma_t \epsilon_1} \textcolor{red}{x_t^{(1)}} \xrightarrow{\text{Consistency}} \textcolor{blue}{x_1^{(2)}} \xrightarrow{x_t^{(2)} = \alpha_t x_1^{(2)} + \sigma_t \epsilon_2} \textcolor{blue}{x_t^{(2)}} \cdots \xrightarrow{\text{Repeat}} \textcolor{purple}{x_1^{(K)}} \label{eq:explicit95iterative}\tag{11}\]
We evaluate IMPFM across two complementary settings. In the online feedback-driven search setting, a target-class-conditioned prior operates without a priori knowledge of the fine-grained target prompt, relying entirely on iterative feedback to progressively steer generated samples toward the target specification revealed only through sequential interaction. In the online feedback-driven alignment setting, a prior model, conditioned on the target prompt, leverages online feedback to iteratively refine its outputs toward satisfying complex, fine-grained constraints — including compositional and quantitative reasoning. We benchmark IMPFM against a diverse set of competitive, state-of-the-art baselines. These include Feynman-Kac Steering (FKS) [2], the SMC-based approach DAS [3], Best-of-N, and the recently proposed Meta-Flow-Map (MFM) method [13]. To simulate online feedback, we employ widely adopted reward models, i.e., VQA score measured with CLIP-FlanT5 [15], and reward score using InstructBLIP [16], ImageReward [17], and Pick-Score [18]. Sample diversity is evaluated using LPIPS [3] score.
To evaluate the search capabilities of IMPFM, we utilize ImageNet-1,000. For the prior, we employ a distilled flow map model optimized jointly using the \(\mathcal{L}^{\text{distill}}_{inst}\) and \(\mathcal{L}^{\text{distill}}_{const}\) loss objectives. We evaluate on 5 randomly selected target classes, revealing their fine-grained specifications strictly via iterative feedback (see Appendix for full class and prompt details).
Fig. [fig:search-result] reports the VQA scores, reward scores, and diversity evaluated across different feedback budgets (\(\mathcal{B}\)). Experimental results demonstrate that IMPFM achieves higher reward scores with considerably fewer feedback interactions than competing approaches. This highlights its efficacy for feedback-driven search, a setting where maximizing reward while ensuring sample efficiency is paramount. IMPFM achieves superior sample efficiency through a highly effective information-sharing mechanism. Unlike prior approaches, the posterior samples generated by any single particle are shared across the entire ensemble for drift correction and reweighting, maximizing information utilization and significantly boosting the framework’s overall sample efficiency.
Notably, IMPFM demonstrates superior resistance to reward over-optimization. While both IMPFM and the MFM baseline utilize ImageReward to steer the dynamics, IMPFM consistently outperforms the baselines across held-out evaluation reward models, e.g., Pick-Score [18]. We attribute this robustness to our ensemble-based reward guidance mechanism. Rather than evaluating each particle against local reward signals in isolation, IMPFM aggregates reward feedback across the full ensemble’s posterior samples. This global aggregation effectively smooths the reward landscape, establishing a broader consensus for the guidance that prevents individual particles from collapsing into the narrow, high-reward local maxima that typically emerge under purely local evaluation. Crucially, IMPFM unlocks a striking increase in semantic diversity among the generated samples, as reflected in the LPIPS score. This underscores the value of our explicit, diversity-inducing repulsive force—demonstrating that actively enforcing particle separation for reweighting and drift correction is key for diversity. We visualize the generated samples across various approaches in Fig. [fig:search-visu] (See Appendix for more viz).
Weight degeneracy remains a fundamental bottleneck in traditional SMC/FKC frameworks. To isolate the contribution of particle interaction, we conduct two ablation tests. First, we remove the interaction term from both the drift correction (Eq. 5 ) and reweighting (Eq. 9 ), reducing the sampler to a standard reward-tilted FKC. Second, we take a purely value-guided optimal control sampler — specifically MFM — and equip it with our proposed particle interaction term (Eq. 5 ), forming a variant we refer to as MPC, to assess the performance gains attributable to interaction alone. As shown in Fig. [fig:fkc-reward], the interaction-aware formulation in both drift correction and reweighting is critical: it promotes substantially greater sample diversity, sustains a persistently high ESS across all correction steps— failure modes endemic to the reward-tilted FKC. Similarly, as depicted in Fig. [fig:oc], incorporating particle interaction into the MFM framework confirms that it is key to overcoming mode collapse and reward over-optimization even within a purely value-based optimal control setting, as evidenced by richer sample diversity and consistently higher VQA scores measured via InstructBLIP. Additional analysis is in the Appendix.
SS serves a dual purpose: it introduces DDPM-like stochasticity to enable exploration while retaining ODE-based generative fidelity, and supports arbitrary-duration stochastic transitions essential for budget-efficient search.
A natural alternative is to induce stochasticity by converting the ODE to an SDE via Eq. 7 . A direct comparison across varying feedback budgets (Fig. 6) reveals that the SS-based formulation yields consistently higher sample quality — even under severe feedback-budget constraints — confirming its importance for efficient feedback-driven search. Additional analysis on the role of SS is in the Appendix.
We isolate the effect of the reweighting step by comparing the full IMPFM framework against an ablated variant, referred to as Multi Particle Control (MPC), in which only the reweighting mechanism is removed. As reported in [fig:reweight], the reweighting step proves critical — it drives faster convergence to higher reward scores (evaluated using the VQA score) across all feedback budgets.
To assess the effectiveness of IMPFM on online, feedback-driven alignment problems, we focus on two challenging tasks. First, for compositional alignment, we evaluate using 50 randomly selected prompts from GenAI-Bench [19] that involve at least three advanced compositional skills. Second, for quantity-aware alignment, we evaluate using 50 prompts randomly sampled from the numeracy category of T2I-CompBench++ [20]. For the quantity-aware alignment task, the reward is also evaluated as the negative Residual Sum of Squares (RSS) between the target object counts and the detected counts, where detections are obtained using DINO [21] and SAM [22]. We employ an off-the-shelf consistency model [23] as the prior and sample from the posterior (\(p_{1|t}(\cdot|x_t)\)) using the proposed iterative posterior sampling scheme. We provide a detailed derivation of the reparametrization of the noise schedule used in the state-of-the-art consistency model [23] to facilitate its integration within the IMPFM framework.
The Alignment results (as depicted in Fig. [fig:alignment-result]) reveal a strikingly similar trend. IMPFM converges to high reward scores with significantly fewer feedback interactions than competing approaches, while simultaneously maintaining greater sample diversity. The rapid convergence further underscores the effectiveness of the particle interaction mechanism, which enables maximum information utilization for drift correction. The sustained diversity, in turn, confirms that IMPFM does not suffer from the weight degeneracy issues that commonly afflict practical inference — a property made possible by the explicit repulsive force between particles, which actively discourages collapse and promotes population-level diversity. We provide a qualitative comparison of generated samples across competing approaches in Fig. 8. Furthermore, these findings underscore the effectiveness of the iterative posterior sampling scheme.
Inference-time scaling has garnered significant attention recently, with particle-based SMC emerging as a highly versatile approach that maintains a population of samples to approximate the target distribution. SMC [2], [3], [24]–[26] leverages potential functions, which usually approximate the soft value function, to assign weights to particles and resample them at every step. While numerous adaptations have been developed specifically for diffusion model alignment [3], [26]–[29], practical implementations face significant hurdles. Although SMC provides theoretical guarantees of exact sampling in the infinite-particle limit, empirical applications are hindered by weight variance and inaccurate value estimates, which degrade sample diversity. Furthermore, these approaches are bottlenecked by their reliance on a pre-trained reward model to derive the necessary particle weights. Within our online framework, early-stage reward models inherently suffer from poor calibration. This initial bias cascades through the resampling step, ultimately misdirecting the reverse denoising trajectory and compromising final sample quality. Recently, several FKC techniques [4], [7] have been introduced to alleviate particle weight degeneracy and debias particle-based guidance. However, these methods struggle in high-dimensional domains like images. Because they rely on reward gradients at every resampling step—which are typically unavailable and require approximation—they inadvertently inject systematic bias into the sampling process. Although numerous gradient-free alignment methods avoid computing reward gradients by employing BON sampling [30]–[32], this generate-and-filter approach introduces substantial computational overhead. Furthermore, because these methods only intervene post-hoc at the final output stage, they cannot actively guide the intermediate generative trajectory or fundamentally alter the underlying distribution.
Fine-tuning generative models with human feedback, such as user preferences, has become increasingly prevalent [33], [34]. However, these approaches typically assume a static reward model, treating rewards as fixed ground truth and not accommodating online queries. In contrast, we study an online, interactive setting in which target properties are initially unknown and are gradually revealed through sequential feedback, enabling continual target discovery rather than one-shot alignment to a fixed reward proxy. [35] introduces a general online learning framework for aligning generative models, but it is not tailored to diffusion models. More recently, [1] proposed a feedback-efficient online fine-tuning method for diffusion models; however, it learns a separate online reward model to guide sampling, which amplifies bias in the early stages of fine-tuning, and due to the KL-regularized fine-tuning objective, it remains vulnerable to mode collapse. While tree-based methods [6] provide a gradient-free alternative for search, they lack a mechanism for active trajectory correction. By relying strictly on prior dynamics, these samplers fail to leverage acquired feedback effectively, resulting in inefficient search under limited budget constraints.
In this work, we introduce IMPFM, a principled, interactive Feynman-Kac corrector (FKC) sampler designed for sample-efficient online feedback-driven search and alignment. Our rigorous empirical analysis demonstrates its scalability and efficacy, particularly in high-dimensional settings. In the future, we aim to extend the applicability of IMPFM across diverse scientific domains.
Page
Derivation of Optimal Drift Correction via Interactive Ensemble Dynamics
IMPFM Yields Fynman-Kac Corector to KL-Tilted Target
Additional Analysis on the Impact of the Interaction-Aware Feynman-Kac Corrector
Visualizing the Impact of Sufficient Statistic
Additional Qualitative Analysis: Taming Weight Degeneracy Through Particle Interaction-Guided Drift Correction
Details of Class and Prompt Detail in the Search Experiment
Additional Comparative Visualization on Search Task
Details about the Evaluation Dataset for Alignment Tasks
Implementation Details
Derivation of Noise Schedule to Adapt IMPFM on Alignment Tasks
Additional Visualizations on the Alignment Task
Let \(q(x)\) denote the empirical distribution induced by the interacting particles, and define the Boltzmann distribution as \(p_{\text{boltz}}(x) \propto e^{V(x)}\). The steepest descent direction that minimizes the Kullback-Leibler divergence, i.e., \(\nabla_x D_{\text{KL}}(q \parallel p_{\text{boltz}})\) within the unit ball of a Reproducing Kernel Hilbert Space (RKHS) is given by: \[\mathbb{E}_{x^i \sim q} \left[ \nabla_{x^i} V(x^{i}) + \frac{1}{n-1} \sum_{j=1, j \neq i}^n \left[ k(x^j, x^{i}) \nabla_{x^j} V(x^j) + \nabla_{x^j} k(x^j, x^{i}) \right] \right]\]
Proof. We aim to derive the steepest direction minimizing the KL-Divergence within the unit ball of an RKHS.
Let \(q(x)\) be the current distribution of particles. We apply a smooth transformation \(T(x) = x + \epsilon \phi(x)\), where \(\epsilon > 0\) is a small step size and \(\phi \in \mathcal{H}^d\) is the velocity field in a Reproducing Kernel Hilbert Space (RKHS) with kernel \(k(\cdot, \cdot)\). Let \(q_{[\epsilon\phi]}\) denote the pushforward distribution of \(q\) under \(T\).
We want to find the direction \(\phi\) that maximizes the decrease of \(D_{KL}(q_{[\epsilon\phi]} \parallel p_{boltz})\). By the results of [36], the first variation of the KL divergence at \(\epsilon = 0\) evaluates to the expectation of the Stein operator:\[\left. \frac{d}{d\epsilon} D_{KL}(q_{[\epsilon\phi]} \parallel p_{boltz}) \right|_{\epsilon=0} = - \mathbb{E}_{x \sim q} \big[ \text{trace}(\mathcal{A}_{p_{boltz}} \phi(x)) \big]\]where the Stein operator \(\mathcal{A}_p\) is defined as:\[\mathcal{A}_{p_{boltz}} \phi(x) = \phi(x)^\top \nabla_x \log p_{boltz}(x) + \nabla_x \cdot \phi(x)\]Given \(p_{boltz}(x) \propto e^{V(x)}\), we have \(\nabla_x \log p_{boltz}(x) = \nabla_x V(x)\). Substituting this in yields:\[\left. \frac{d}{d\epsilon} D_{KL} \right|_{\epsilon=0} = - \mathbb{E}_{x \sim q} \big[ \phi(x)^\top \nabla_x V(x) + \nabla_x \cdot \phi(x) \big]\]
By the reproducing property of the RKHS, any \(\phi \in \mathcal{H}^d\) satisfies \(\phi(x) = \langle \phi, k(x, \cdot) \rangle_{\mathcal{H}}\). We can pass the linear inner product outside the expectation:\[\mathbb{E}_{x \sim q} \big[ \phi(x)^\top \nabla_x V(x) + \nabla_x \cdot \phi(x) \big] = \left\langle \phi(\cdot), \; \mathbb{E}_{x \sim q} \big[ k(x, \cdot) \nabla_x V(x) + \nabla_x k(x, \cdot) \big] \right\rangle_{\mathcal{H}}\]
To find the steepest descent direction within the unit ball of the RKHS, we solve:\[\phi^* = \arg\max_{\|\phi\|_{\mathcal{H}} \le 1} \left( -\left. \frac{d}{d\epsilon} D_{KL} \right|_{\epsilon=0} \right)\]By the Cauchy-Schwarz inequality in Hilbert spaces (\(\langle \phi, u \rangle \le \|\phi\| \|u\|\)), the linear functional is maximized when \(\phi\) is strictly proportional to the second argument of the inner product. Therefore, the optimal vector field is:\[\phi^*(\cdot) = \mathbb{E}_{x \sim q} \big[ k(x, \cdot) \nabla_x V(x) + \nabla_x k(x, \cdot) \big]\]
Given \(n\) interacting particles \(\{x^1, \dots, x^n\}\), evaluating the optimal velocity field \(\phi^*\) at a specific particle \(x^i\) using the empirical distribution \(q(x) \approx \frac{1}{n} \sum_{j=1}^n \delta(x - x^j)\) yields:\[\phi^*(x^i) = \frac{1}{n} \sum_{j=1}^n \big[ k(x^j, x^i) \nabla_{x_j} V(x^j) + \nabla_{x_j} k(x^j, x^i) \big]\] To match the proposed formulation in Proposition 4.1, we isolate the self-interaction term (\(j=i\)). For standard translation-invariant kernels (like the RBF kernel), \(k(x^i, x^i) = 1\) and \(\nabla_{x_i} k(x^i, x^i) = 0\). By separating the \(j=i\) term and replacing the generic \(\frac{1}{n}\) weight of the interaction sum with the unbiased leave-one-out U-statistic estimator \(\frac{1}{n-1}\) for the other particles, the update direction at \(x^i\) becomes exactly:\[\nabla_{x_i} V(x^i) + \frac{1}{n-1} \sum_{j \neq i} \big[ k(x^j, x^i) \nabla_{x_j} V(x^j) + \nabla_{x_j} k(x^j, x^i) \big]\]Taking the expectation over all particles \(x^i \sim q\) yields the exact expression as stated in Proposition 4.1. ◻
Consider a Flow model with drift \(b_t(x)\) trained to sample from \(q_t(x)\). Sampling from the KL-tilted marginals \(p_t^{\text{kl}}(x) \propto q_t(x) \exp\Big( -D_{\text{KL}}\big( q_t(x) \parallel p_{\text{boltz}}(x) \big) \Big)\) is performed by an weighted SDE with an equivalent corresponding weighted ODE: \[\begin{align} dx^i_t &= b_t(x^i_t)\>dt + \frac{\sigma_t^2}{2}\frac{1}{n} \sum_{j=1}^n \left[ k(x^j_t, x^{i}_t) \nabla_{x^j_t} V(x^j_t) \> dt + \nabla_{x^j_t} k(x^j_t, x^{i}_t) \> dt \right] \end{align}\] \[\begin{align} dw^i_t = \frac{1}{n} \sum_{j=1}^n \left[ \boldsymbol{\langle} k(x_t^j, x_t^{i}) \nabla_{x_t^j} V(x_t^j), b_t(x_t^i) \boldsymbol{\rangle} \> dt + \boldsymbol{\langle} \nabla_{x_t^j} k(x_t^j, x_t^{i}), b_t(x_t^i) \boldsymbol{\rangle}\> dt \right] \end{align}\]
Proof. We first show the weighted SDE form that corresponds to sampling from \(p_t^{kl}(x)\), then we convert the resulting weighted SDE to the corresponding weighted ODE.
Assume \(f^{kl}(x) = -D_{KL}(q_t(x) || p_{boltz}(x))\). The generation process can be defined as the family of denoising SDEs in the opposite time direction as below: \[\label{eq:sde951} dx_t = \left( -f_t(x_t) + \sigma_t^2 \nabla \log q_t(x_t) \right) dt + \sigma_t dW_t\tag{12}\] The SDE in Eq. 12 can be viewed as the composition of a flow and diffusion terms, where the corresponding Fokker-Planck PDE describes the combined evolution: \[\frac{\partial q_t(x)}{\partial t} = -\underbrace{\langle \nabla, q_t(x) v_t(x) \rangle}_{\text{Continuity Eqn. }} + \underbrace{\frac{\sigma_t^2}{2} \Delta q_t(x)}_{\text{Diffusion Eqn.}}.\] We want to find the PDE for the reward-tilted density: \[\label{eq:marginal95kl} \hat{p_t}^{kl}(x) = \frac{q_t(x) \exp(\beta_t f^{kl}(x))}{\int dx \, q_t(x) \exp(\beta_t f^{kl}(x))}.\tag{13}\] By applying the logarithm and then taking the partial derivative with respect to \(t\), we get: \[\frac{\partial}{\partial t} \log \hat{p_t}^{kl}(x) = \frac{\partial}{\partial t} \log q_t(x) + \frac{\partial \beta_t}{\partial t} f^{kl}(x) - \int dx \, \hat{p_t}^{kl}(x) \left[ \frac{\partial}{\partial t} \log q_t(x) + \frac{\partial \beta_t}{\partial t} f^{kl}(x) \right].\] For the first term, we have: \[\begin{align} \frac{\partial}{\partial t} \log q_t(x) &= -\langle \nabla, v_t(x) \rangle - \langle \nabla \log q_t(x), v_t(x) \rangle + \frac{\sigma_t^2}{2} \Delta \log q_t(x) + \frac{\sigma_t^2}{2} \|\nabla \log q_t(x)\|^2 \nonumber \\ &= -\langle \nabla, v_t(x) \rangle - \langle \nabla \log \hat{p_t}^{kl}(x), v_t(x) \rangle + \frac{\sigma_t^2}{2} \Delta \log \hat{p_t}^{kl}(x) + \frac{\sigma_t^2}{2} \|\nabla \log \hat{p_t}^{kl}(x)\|^2 \nonumber \\ &\quad + \left\langle \beta_t \nabla f^{kl}(x), v_t(x) - \sigma_t^2 \nabla \log q_t(x) - \frac{\sigma_t^2}{2} \beta_t \nabla f^{kl}(x) \right\rangle - \beta_t \frac{\sigma_t^2}{2} \Delta f^{kl}(x). \end{align}\] The last equality is obtained using the fact that \[\nabla \log q_t(x) = \nabla \log \hat{p_t}^{kl}(x) - \beta_t \nabla f^{kl}(x) \quad\] and \[\Delta \log q_t(x) = \Delta \log \hat{p_t}^{kl}(x) - \beta_t \Delta f^{kl}(x) \quad\] Both of which follow from 13 .
Thus, we have: \[\frac{\partial \hat{p_t}^{kl}(x)}{\partial t} = -\langle \nabla, \hat{p_t}^{kl}(x) v_t(x) \rangle + \frac{\sigma_t^2}{2} \Delta \hat{p_t}^{kl}(x) + \hat{p_t}^{kl}(x) \big( g_t(x) - \mathbb{E}_{\hat{p_t}^{kl}(x)}[g_t(x)] \big),\] \[g_t(x) = \left\langle \beta_t \nabla f^{kl}(x), v_t(x) - \sigma_t^2 \nabla \log q_t(x) - \frac{\sigma_t^2}{2} \beta_t \nabla f^{kl}(x) \right\rangle - \beta_t \frac{\sigma_t^2}{2} \Delta f^{kl}(x) + \frac{\partial \beta_t}{\partial t} f^{kl}(x).\] Furthermore, we can add the gradient of the tilted function as an additional drift term \(a \nabla f^{kl}(x)\), i.e., \[\frac{\partial \hat{p_t}^{kl}(x)}{\partial t} = -\langle \nabla, \hat{p_t}^{kl}(x) (v_t(x) + a \nabla f^{kl}(x)) \rangle + \frac{\sigma_t^2}{2} \Delta \hat{p_t}^{kl}(x) + \hat{p_t}^{kl}(x) \big( g_t(x) - \mathbb{E}_{\hat{p_t}^{kl}(x)}[g_t(x)] \big),\] \[\begin{align} g_t(x) &= a \Delta f^{kl}(x) + a \langle \nabla \log p_t(x), \nabla f^{kl}(x) \rangle - \beta_t \frac{\sigma_t^2}{2} \Delta f^{kl}(x) + \frac{\partial \beta_t}{\partial t} f^{kl}(x) \nonumber \\ &\quad + \left\langle \beta_t \nabla f^{kl}(x), v_t(x) - \sigma_t^2 \nabla \log q_t(x) - \frac{\sigma_t^2}{2} \beta_t \nabla f^{kl}(x) \right\rangle. \end{align}\] Taking \(v_t(x) = -f_t(x) + \sigma_t^2 \nabla \log q_t(x)\) and \(a = \beta_t \sigma_t^2 / 2\), we have: \[\begin{align} \frac{\partial \hat{p_t}^{kl}(x)}{\partial t} = -\left\langle \nabla, \hat{p_t}^{kl}(x) \left( -f_t(x) + \sigma_t^2 \nabla \log q_t(x) + \beta_t \frac{\sigma_t^2}{2} \nabla f^{kl}(x) \right) \right\rangle \\ + \frac{\sigma_t^2}{2} \Delta \hat{p_t}^{kl}(x) + \hat{p_t}^{kl}(x) \big( g_t(x) - \mathbb{E}_{\hat{p_t}^{kl}(x)}[g_t(x)] \big), \end{align}\] \[\begin{align} g_t(x) &= \left\langle \beta_t \nabla f^{kl}(x), \frac{\sigma_t^2}{2} \nabla \log \hat{p_t}^{kl}(x) \right\rangle + \frac{\partial \beta_t}{\partial t} f^{kl}(x) + \left\langle \beta_t \nabla f^{kl}(x), -f_t(x) - \frac{\sigma_t^2}{2} \beta_t \nabla f^{kl}(x) \right\rangle \nonumber \\ &= \frac{\partial \beta_t}{\partial t} f^{kl}(x) + \left\langle \beta_t \nabla f^{kl}(x), \frac{\sigma_t^2}{2} \nabla \log q_t(x) - f_t(x) \right\rangle. \end{align}\] This can be simulated as: \[dx_t = \left( -f_t(x_t) + \sigma_t^2 \nabla \log q_t(x_t) + \beta_t \frac{\sigma_t^2}{2} \nabla f^{kl}(x_t) \right) dt + \sigma_t dW_t,\] \[dw_t = \left[ \frac{\partial \beta_t}{\partial t} f^{kl}(x_t) + \left\langle \beta_t \nabla f^{kl}(x_t), \frac{\sigma_t^2}{2} \nabla \log q_t(x_t) - f_t(x_t) \right\rangle \right] dt.\]
Assuming \(\beta_t = 1\), the above weighted SDE can be written as: \[dx_t = \left( -f_t(x_t) + \sigma_t^2 \nabla \log q_t(x_t) + \frac{\sigma_t^2}{2} \nabla f^{kl}(x_t) \right) dt + \sigma_t dW_t,\] \[dw_t = \left[ \left\langle \nabla f^{kl}(x_t), \frac{\sigma_t^2}{2} \nabla \log q_t(x_t) - f_t(x_t) \right\rangle \right] dt.\]
By assuming \(b_t(x_t)= -f_t(x_t) + \frac{\sigma^2_t}{2} \nabla \log q_t(x_t)\), we can write the following as, \[dx_t = \left( b_t(x_t) + \frac{\sigma_t^2}{2} \nabla \log q_t(x_t) + \frac{\sigma_t^2}{2} \nabla f^{kl}(x_t) \right) dt + \sigma_t dW_t,\] \[dw_t = \left[ \left\langle \nabla f^{kl}(x_t), b_t(x_t) \right\rangle \right] dt.\]
Now, by following the standard SDE to ODE conversion rule, we can write as follows: \[dx_t = \left( b_t(x_t) + \frac{\sigma_t^2}{2} \nabla f^{kl}(x_t) \right) dt,\] \[dw_t = \left[ \left\langle \nabla f^{kl}(x_t), b_t(x_t) \right\rangle \right] dt.\]
Now, utilizing the results from Proposition 4.2 regarding \(- \nabla D_{KL}(q_t(x) || p_{boltz}(x))\) and utilizing the formulation of \(f^{kl}(x_t)= -D_{KL}(q_t(x) || p_{boltz}(x))\), we get the final expression for the \(i\)-th particle as: \[\begin{align} dx^i_t &= b_t(x^i_t)\>dt + \frac{\sigma_t^2}{2}\frac{1}{n} \sum_{j=1}^n \left[ k(x^j_t, x^{i}_t) \nabla_{x^j_t} V(x^j_t) \> dt + \nabla_{x^j_t} k(x^j_t, x^{i}_t) \> dt \right] \end{align}\] \[\begin{align} dw^i_t = \frac{1}{n} \sum_{j=1}^n \left[ \boldsymbol{\langle} k(x_t^j, x_t^{i}) \nabla_{x_t^j} V(x_t^j), b_t(x_t^i) \boldsymbol{\rangle} \> dt + \boldsymbol{\langle} \nabla_{x_t^j} k(x_t^j, x_t^{i}), b_t(x_t^i) \boldsymbol{\rangle}\> dt \right] \end{align}\] This completes the Proof. ◻
While in the main paper, we evaluate IMPFM on online feedback-driven search tasks with relatively short horizons, this section extends our comparative analysis to longer-horizon tasks, with findings reported in Figure [fig:search-result-longer]. Notably, the performance trends in these extended settings mirror those observed in the shorter-horizon tasks. These additional results reinforce that IMPFM is highly sample-efficient—achieving higher rewards with fewer interactions—while simultaneously maintaining a robust search capability over longer horizons that consistently outperforms the baselines.
Building on the qualitative and quantitative analysis presented in the main paper, we provide here additional results to further substantiate the impact of the interaction-aware Feynman-Kac corrector. Following the same ablation protocol, we remove the interaction terms — comprising both the repulsive term and the kernel-weighted value gradient — from the drift correction and reweighting mechanisms, yielding a standard reward-tilted Feynman-Kac corrector, referred to as FKC-Reward-Tilted. We benchmark this against our full Interaction-Aware Feynman-Kac corrector and report the results in Fig. [fig:fkc-comp]. As illustrated, the interaction-aware formulation substantially enhances sample diversity, directly alleviating mode collapse. The key limitation of the FKC-Reward-Tilted sampler lies in its purely pointwise optimization objective: by neglecting global distributional constraints, it greedily pursues high-reward regions and consequently becomes trapped in local neighborhoods, rendering it susceptible to reward over-optimization. In contrast, the interaction-aware Feynman-Kac corrector incorporates a KL-divergence tilting factor that enforces distributional alignment, while the explicit inter-particle interaction mechanism actively promotes global coverage by aggregating information across the entire particle ensemble. This dual mechanism naturally steers particles away from local optima, fostering both diversity and robustness against reward over-optimization, as clearly reflected in Fig. [fig:fkc-comp]. Furthermore, we analyze the specific impact of the interaction term within the drift correction mechanism (Eq. 5 ). To isolate this effect, we visually compare the samples generated by our fully interactive formulation (MPC) against an ablated, non-interacting baseline (MFM). As shown in Figure [fig:oc-mpc], incorporating particle interaction significantly improves both sample diversity and target alignment. This highlights the interaction term as a critical component for overcoming the severe mode collapse and reward over-optimization typical of standard optimal control frameworks.
In a nutshell, these additional qualitative analyses reinforce the importance of the interaction-aware Feynman-Kac corrector for efficient online feedback-driven search.
In the main paper (Section 3.0.0.4), we analyze the impact of the Sufficient Statistic (SS). Here, we present additional visualizations highlighting that the SS is crucial for enabling feedback-efficient search, especially under severe budget constraints. As depicted in Fig. 13, when the feedback step is less than 50, the alternative approach fails to generate high-quality samples. In contrast, by leveraging dynamic transition steps enabled via the SS mechanism, IMPFM successfully generates high-quality, target-aligned samples even with a very small feedback step (e.g., 20).
In this section, we present a deeper qualitative investigation into how the proposed multi-particle interaction-driven drift correction mechanism combats weight degeneracy. To achieve this, we visualize the denoised samples generated via flow-map for each particle in the batch and compute the Effective Sample Size (ESS) at every single transition step. By contrasting our fully interacting dynamics (Eq. 5 ) against a naive baseline—where particles are treated in strict isolation and corrected purely based on individual value—we directly highlight the stabilizing power of our approach. We report the findings in Figure 14 and 15. Crucially, we demonstrate that multi-particle interaction acts as a fundamental driver of sample diversity. Unlike isolated particle dynamics—which rapidly collapse into semantically redundant outputs—our interacting system sustains rich, diverse generation across every correction step. This structural stability is corroborated by a drastic, consistent improvement in the Effective Sample Size (ESS), effectively eradicating the pervasive weight degeneracy problem. Importantly, this interactive framework is made computationally viable by flow-map; by enabling highly efficient evaluation of value gradients (\(\nabla V_t\)), flow-maps unlock the ability to compute complex, multi-particle interactions without prohibitive overhead.
In the following table 1, we detail the randomly selected ImageNet target classes and their fine-grained specifications used to evaluate the competing methods on the online feedback-driven search task.
| Target Class | Fine-Grained Prompt |
|---|---|
| Car | A Red Sports Car |
| Cat | A Black Cat. |
| Dog | A Dog in the Snow. |
| Daisy | A Red Daisy. |
| Pepper | A Green Bell Pepper. |
In Fig. 16, we present additional comparative visualizations for search tasks involving diverse classes from ImageNet. These figures further highlight the impact of IMPFM in enabling efficient, online feedback-driven search.
We provide the evaluation prompts for both the compositional and quantity-aware alignment tasks in separate Excel files, available via the GitHub link.
For all experiments in the online feedback-driven search setting, we use 32 particles, with each particle contributing one posterior sample for feedback guidance, resulting in 32 feedback samples per resampling step. Notably, when configured with a single particle, IMPFM reduces to the MFM baseline. However, as the particle count increases, IMPFM outperforms MFM under the same feedback budget due to its posterior sample sharing mechanism. Other implementation details are available in the provided code. Finally, all experiments are implemented in PyTorch and conducted on 3 NVIDIA A100 GPUs. Our training and inference code will be made public.
In this section, we describe how we adapt the trigonometric noise schedule from TrigFlow [23] and the consistency model (SANA-Sprint) [23] to our IMPFM framework, which operates under a different noise schedule. We first review the TrigFlow parameterization, then describe the consistency model reparameterization, and finally derive the adaptation of the noise schedule to our setting.
TrigFlow introduces a trigonometric reparameterization of the Flow Matching timestep, establishing a bijective mapping between the standard FM timestep \(t_{FM} \in [0, 1]\) and a cosine-based schedule parameterized by \(t_{cm}\). Specifically, the mapping is defined as: \[t_{FM} = \frac{\sin(t_{cm})}{\sin(t_{cm}) + \cos(t_{cm})}, \label{eq:tFM}\tag{14}\] where \(t_{cm}\) controls the balance between noise and signal through trigonometric interpolation. This mapping ensures that \(t_{FM} = 0\) corresponds to a clean sample and \(t_{FM} = 1\) corresponds to pure noise, while providing a smoother, geometry-aware interpolation path compared to a linear schedule.
Given the mapping in Eq. 14 , the noisy latent \(X_{t_{FM}}\) is expressed in terms of the original latent \(X_{t_{cm}}\) as: \[X_{t_{FM}} = \frac{X_{t_{cm}}}{\sigma_d} \sqrt{t_{FM}^2 + (1 - t_{FM})^2}, \label{eq:Xtfm}\tag{15}\] where \(\sigma_d\) is a scaling factor that normalizes the latent representation. The factor \(\sqrt{t_{FM}^2 + (1-t_{FM})^2}\) accounts for the norm of the trigonometric interpolation coefficients, ensuring consistent scaling across timesteps.
Under this reparameterization, the score network \(\hat{F}_\theta\) is reformulated to operate on the rescaled latent \(X_{t_{cm}} / \sigma_d\) at timestep \(t_{cm}\). The relationship between the TrigFlow score network and the velocity field \(V_\theta\) of the standard Flow Matching model is: \[\begin{align} \hat{F}_\theta\!\left(\frac{X_{t_{cm}}}{\sigma_d},\, t_{cm},\, y\right) &= \frac{1}{\sqrt{t_{FM}^2 + (1 - t_{FM})^2}} \\ &\quad \times \Big[ (1 - 2t_{FM})\, X_{t,FM} + \left(1 - 2t_{FM} + 2t_{FM}^2\right) V_\theta\!\left(X_{t_{FM}},\, t_{FM},\, y\right) \Big], \end{align} \label{eq:Ftheta}\tag{16}\] where \(y\) is the conditioning input (e.g., a text prompt). This reformulation shows that the TrigFlow score network can be expressed as a linear combination of the noisy latent and the predicted velocity field, weighted by time-dependent coefficients that are derived analytically from the trigonometric schedule.
SANA-Sprint builds upon TrigFlow by introducing a consistency-style denoising function \(f_\theta\) that directly maps a noisy latent \(X_t\) to a clean estimate \(\hat{X}_0\) in a single step. Under the TrigFlow time parameterization, this function takes the form: \[f_\theta(X_t, t) = \cos(t_{cm}) \cdot X_t - \sin(t_{cm}) \cdot \sigma_d \cdot F_\theta\!\left(\frac{X_t}{\sigma_d},\, t_{cm}\right), \label{eq:ftheta95sana}\tag{17}\] where \(X_{cm} = X_t / \text{Scale}\) is the rescaled latent. Intuitively, Eq. 17 can be understood as a trigonometric decomposition: the \(\cos(t_{cm})\) term retains the signal component of \(X_t\), while the \(\sin(t_{cm})\) term removes the noise component as predicted by \(F_\theta\). This yields the denoised estimate: \[\hat{X}_0 = \cos(t_{cm})\, X_{t_{cm}} - \sin(t_{cm})\, \hat{F}_\theta, \label{eq:X0hat}\tag{18}\] which provides a clean, geometry-consistent reconstruction of the original sample from any noisy intermediate state.
Our IMPFM framework uses a different noise schedule compared to TrigFlow. Specifically, we adopt the \(S\) schedule, which defines a signal-to-noise ratio (SNR) function that differs from the trigonometric parameterization used in TrigFlow. To adapt the SANA-Sprint framework to our setting, we must establish a consistent mapping between the two schedules.
Under the \(S\) schedule, the signal and noise coefficients at timestep \(t = t_{cm}\) are defined as: \[\alpha_t = \cos t = \sqrt{s}, \qquad \sigma_t = \sin t = \sqrt{1 - s}, \label{eq:alpha95sigma}\tag{19}\] where \(s \in [0, 1]\) is the schedule variable. This gives the forward noising process: \[X_t = \left(\alpha_t\, \hat{X}_0 + \sigma_t\, \frac{\varepsilon}{\sigma_d}\right) \cdot \sigma_d = \cos(t)\, \hat{X}_0 \cdot \sigma_d + \sin(t)\, \varepsilon, \label{eq:forward95process}\tag{20}\] which is a standard affine interpolation between the clean sample \(\hat{X}_0\) and Gaussian noise \(\varepsilon \sim \mathcal{N}(0, I)\), scaled by the latent normalization factor \(\sigma_d\).
Under the \(S\) schedule, the SNR function \(g(s)\) and its transformed counterpart \(\tilde{g}(\alpha_t)\) are defined as: \[g(s) = \frac{1}{s} - 1, \qquad \tilde{g}(\alpha_t) = \frac{1}{\alpha_t} - 1 = \frac{1}{\cos^2(t_{cm})} - 1, \label{eq:snr95functions}\tag{21}\] where \(\tilde{g}(\alpha_t)\) is the SNR expressed as a function of the signal coefficient \(\alpha_t = \cos(t_{cm})\). This function is monotonically decreasing in \(t_{cm}\), consistent with the intuition that noise increases as the timestep progresses.
A key step in adapting the noise schedule is determining the optimal intermediate timestep \(r^*\) at which to perform the consistency transition between two noise levels \(t_{cm}\) and \(t_{cm'}\) (with \(t_{cm'} < t_{cm}\), i.e., \(t_{cm'}\) is closer to the clean data). The optimal \(r^*\) is derived by matching the SNR levels under the \(S\) schedule, yielding: \[\tilde{g}(r^*) = \frac{\tilde{g}(\cos t_{cm'}) \cdot \tilde{g}(\cos t_{cm})}{\tilde{g}(\cos t_{cm}) - \tilde{g}(\cos t_{cm'})}, \label{eq:r95star95snr}\tag{22}\] which can be solved by applying the inverse of \(g\) to both sides: \[\hat{r}^* = g^{-1}\!\left( \frac{\tilde{g}(\cos t_{cm'}) \cdot \tilde{g}(\cos t_{cm})}{\tilde{g}(\cos t_{cm}) - \tilde{g}(\cos t_{cm'})} \right). \label{eq:r95star}\tag{23}\] Eq. 23 gives the optimal SNR level \(r^*\) in the \(s\)-domain. To recover the corresponding trigonometric timestep \(r^*_{cm}\), we use the relation \(\alpha_t = \cos(t_{cm}) = \sqrt{s}\), which gives: \[\cos(r^*_{cm}) = d\sqrt{r^*} \quad \Longrightarrow \quad r^*_{cm} = \arccos\!\left(\sqrt{r^*}\right). \label{eq:r95star95cm}\tag{24}\] This recovers the optimal intermediate timestep \(r^*_{cm}\) in the trigonometric domain, which can be directly used within the TrigFlow/SANA-Sprint framework.
Given the optimal transition timestep \(r^*_{cm}\), the DDPM transition sample \(X_{t_{cm'}}\) at the new timestep \(t_{cm'}\) is obtained by applying the Sufficient Statistic scheme as defined in Equation 9 in the main paper.
In Fig. 17, 18, [fig:align-visu4], we provide additional comparative visualizations for the quantitative alignment tasks. These results underscore the effectiveness of IMPFM in complex, online feedback-driven alignment tasks that demand robust reasoning capabilities.
Equal Contribution↩︎