arXiv CSopen access

FUSE: FK-Steered Multi-Modal Flow Matching for Efficient Simulation-Based Posterior Estimation

FUSE: FK-Steered Multi-Modal Flow Matching for Efficient Simulation-Based Posterior Estimation Weichen Qin 1 Yufan Xie 1 Peihao Wang 2 Chia-Jui Chou 1 Minghui Du 3 4 Peng Xu 3 4 Ziren Luo 3 4 Yi Yang 1 Jingyi Yu 1 Bo Liang 3 4 Jiakai Zhang 1 5

arXiv:2607.05252v1 [cs.LG] 6 Jul 2026

Abstract

1. Introduction The rapid evolution of generative models (Kingma & Welling, 2014; Lipman et al., 2023; Ho et al., 2020) has recently unlocked transformative capabilities in multi-modal domains, enabling seamless translation between text, images, videos, and 3D assets (Zhang et al., 2024). Beyond their applications in entertainment, these powerful crossmodal transformation capabilities hold immense potential for scientific discovery, particularly in solving inverse problems. In such tasks, the objective is to infer the posterior distribution of latent parameters based on observed data challenge omnipresent in astrophysics, ranging from orbital parameter estimation (Liang et al., 2026b) to the characterization of gravitational wave sources from detectors either on the earth (Dax et al., 2021) or in space (Srinivasan et al., 2025; Du et al., 2024; Liang et al., 2024; 2026a; 2025a;b).

Simulation-Based Inference (SBI) is critical for scientific discovery, with generative models offering a promising path toward efficient inference. However, existing methods struggle with effective multimodal modeling. They often rely on brute-force fusion strategies that ignore the structural disparities between parameters and observations, thus limiting estimation fidelity. In this work, we introduce FUSE (Feynman-Kac steered mUlti-modal flow matching for efficient Simulation-based posterior Estimation). Unlike prior work, FUSE employs a dual-track architecture that preserves the distinct features of multimodal inputs while facilitating dynamic interaction. Additionally, we propose an FK-steered sampling strategy that leverages intermediate observation likelihoods to guide the generative trajectories, effectively improving the sample quality during inference. Our approach outperforms stateof-the-art baselines on standard SBI benchmarks, producing posteriors that closely match reference MCMC posterior samples. Furthermore, in a realworld exoplanet orbital estimation task, FUSE successfully resolves complex parameter degeneracies that challenge existing methods, highlighting its potential to accelerate complex scientific discoveries in astrophysics and beyond.

Despite their potential, estimating the posterior distribution in this context is non-trivial due to the intrinsic complexity of scientific data. First, a primary obstacle is the heterogeneous dimensionality between the input observations and the target parameters. Learning a robust mapping from high-dimensional sensor data to specific physical properties—often within a vast and complex search space, leading to a requirement for the customized design of model architectures. Second, the fidelity requirements are far more stringent than in general vision tasks; the model must accurately estimate the full posterior distribution, capturing multi-modal structural information and subtle correlations to avoid biased parameter estimation. Moreover, practical deployment demands high-throughput inference, requiring models that can generate accurate samples rapidly to keep pace with large-scale survey data or time-sensitive astronomical events.

1

ShanghaiTech University, Shanghai, China 2 The University of Texas at Austin, Austin, TX, USA 3 Center for Gravitational Wave Experiment, National Microgravity Laboratory, Institute of Mechanics, Chinese Academy of Sciences, Beijing 100190, China 4 Taiji Laboratory for Gravitational Wave Universe (Beijing/Hangzhou), University of Chinese Academy of Sciences (UCAS), Beijing 100049, China 5 Cellverse, Co., Ltd. Correspondence to: Bo Liang <[email protected]>, Jiakai Zhang <[email protected]>.

Existing solutions face a fundamental trade-off between accuracy and efficiency. While traditional algorithms like Markov Chain Monte Carlo (MCMC) (Foreman-Mackey et al., 2013) provide asymptotically exact posteriors, they incur prohibitive computational costs, often requiring weeks for a single event. Neural approaches (Wildberger et al., 2023; Gonçalves et al., 2020) offer rapid inference but lack customized architectural designs, leading to suboptimal per-

Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).

1

FUSE: FK-Steered Multi-Modal Flow Matching

formance. Specifically, FMPE (Wildberger et al., 2023) compresses observations into static embeddings and employs a brute-force fusion strategy. This approach ignores the structural disparity between time, parameters, and sensory data, significantly constraining effective multimodal generation. Consequently, these methods often fail to faithfully recover the complex posterior distribution in real-world applications.

discovery. The code repository, including reproduction instructions and artifact documentation, is publicly available at https://github.com/qinwch/FUSE.

2. Related Work Traditional posterior estimation. Traditional likelihoodfree inference approximates the Bayesian posterior without evaluating the likelihood by repeatedly simulating data and retaining (or weighting) parameters whose simulations match the observation under a prescribed discrepancy, typically defined on summary statistics with a tolerance threshold. Historically, this line of work progressed from Rejection ABC—a direct accept/reject procedure driven by discrepancy-based filtering (Tavaré et al., 1997)—to ABCMCMC, which embeds the same acceptance criterion within a Metropolis–Hastings kernel to concentrate computation in higher-posterior regions (Marjoram et al., 2003). In parallel, Synthetic Likelihood methods introduced a parametric likelihood (often Gaussian) for summary statistics, enabling standard likelihood-based inference while retaining simulation as the primitive (Wood, 2010). Despite these advances, classical approaches often exhibit poor mixing in high-dimensional parameter spaces and are inherently non-amortized, requiring costly simulator calls for each new observation and each step of the chain. In contrast, our method leverages FK-steered flow matching to enable amortized inference.

To address these limitations, we propose a novel framework, FUSE, which adopts a Feynman-Kac steered mUltimodal flow matching for efficient Simulation-based posterior Estimation. Inspired by the multi-modal diffusion transformer (MMDiT) (Esser et al., 2024), our model embeds parameters and observations in a dual-track framework to better preserve their distinct structural features, while enabling bidirectional interaction through a subsequent fusion module. This allows the generative flow to dynamically “revisit” the raw observation data at every integration step, ensuring maximal information extraction. Furthermore, to narrow the gap between fast amortized SBI and likelihood-based posterior correction, we introduce an FK-steered sampling strategy. Amortized generative models can place probability mass in physically implausible regions when the learned transport is imperfect, whereas asymptotically exact likelihood-based samplers are often too slow for time-sensitive scientific inference. FK-steering injects simulator likelihood information into the generative trajectory itself: by propagating multiple concurrent trajectories and resampling them at intermediate times, FUSE allocates more computation to high-density posterior regions. In the experiments below, this improves posterior fidelity on difficult inverse problems while adding only modest inference overhead relative to the amortized sampler.

Neural posterior estimation. Neural SBI has sought to modernize likelihood-free inference by pairing simulators with learned surrogates that improve practicality and scalability. One direction augments classical MCMC by amortizing likelihood(-ratio) information with neural estimators to drive Metropolis–Hastings transitions (Hermans et al., 2020) (and related direct ratio formulations that yield more informative transition signals) (Cobb et al., 2024). A second direction replaces chains altogether with amortized inference networks, using variational, adversarial, or score/diffusion-based training to generate posterior samples directly (Glöckler et al., 2022; Ramesh et al., 2022; Geffner et al., 2023; Gloeckler et al., 2024). Despite progress, MCMC-coupled methods remain limited by mixing and per-step costs, while many MCMC-free approaches underexploit the multimodal structure of SBI when fusing parameters and observations, which can hinder MCMC-level fidelity and MAP-oriented estimation. Our method addresses both issues by adapting MMDiT-style joint-attention fusion to parameter–observation inference, using flow matching to learn an amortized multimodal transport, and leveraging inference-time scaling to more reliably sample high-quality parameter estimates.

In our experiments, we comprehensively validate our method on a widely used SBI benchmark (Lueckmann et al., 2021). Compared to state-of-the-art neural methods, our model produces posterior distributions that most closely match reference MCMC results across diverse scenarios, demonstrating superior approximation capabilities over MLP-based baselines. To evaluate the efficacy of our FK-steered sampling strategy, we compare performance with and without this component on the challenging SLCP task. Results show that FK-steered inference yields more compact and high-fidelity posteriors, provided by higher mean and peak likelihood scores. Subsequently, we apply our model to a challenging real-world scientific problem: exoplanet orbital parameter estimation. Leveraging our advanced architecture and sampling strategy, our method successfully resolves complex parameter degeneracies that existing SBI approaches may fail to produce a meaningful posterior distribution. We believe this work presents a scalable, high-precision framework for future scientific

2

FUSE: FK-Steered Multi-Modal Flow Matching

Multimodal generative models Generative modeling has rapidly improved the quality-efficiency trade-off, with dominant paradigms evolving from latent-variable models (including multimodal VAEs) to score-based models and, more recently, flow-matching objectives that learn transport dynamics directly (Kingma & Welling, 2014; Suzuki et al., 2016; Wu & Goodman, 2018; Shi et al., 2019; Sutter et al., 2021; Ho et al., 2020; Song et al., 2021b; Lipman et al., 2023). Meanwhile, sampling has been steadily accelerated from ancestral DDPM to DDIM and modern fast samplers such as rectified flow and ODE/SDE solvers that deliver high-quality samples in far fewer steps (Ho et al., 2020; Song et al., 2021a; Liu, 2022; Lu et al., 2022; Karras et al., 2022; Salimans & Ho, 2022; Song et al., 2023). Architectures have also shifted from U-Nets to multimodal Diffusion transformer, which promote joint attention over modality tokens for stronger cross-modal interaction (Ronneberger et al., 2015; Peebles & Xie, 2023; Bao et al., 2023; Chen et al., 2024; Esser et al., 2024). On the conditioning side, guidance and control mechanisms have evolved from classifier(-free) guidance to explicit multimodal controls (e.g., edges/layouts/instructions) and personalization, supported by scalable tokenization and alignment (Dhariwal & Nichol, 2021; Ho & Salimans, 2022; Zhang et al., 2023; Li et al., 2023; Brooks et al., 2023; Ruiz et al., 2023; Oord et al., 2017; Esser et al., 2021; Radford et al., 2021). In contrast to text or images, parameter estimation involves more complex simulator-driven dependencies and multimodal structures. We handle these by adapting a dual-track fusion architecture and using inference-time scaling to better sample high-density regions.

While powerful, the efficiency of this process depends heavily on the computational cost of the likelihood function p(x|θ), where x = {x1 , . . . , xn } is a dataset consisting of n independent observations with xi ∈ Rdx . Assuming the measurement noise follows a multivariate Gaussian distribution with covariance matrix Σi ∈ Rdx ×dx , the loglikelihood is explicitly given by the sum of squared Mahalanobis distances: n

1X ⊤ (xi − fi (θ)) Σ−1 i (xi − fi (θ))+C, 2 i=1 (3) where fi (θ) ∈ Rdx denotes the theoretical prediction generated by a forward model for the i-th observation. In highdimensional parameter spaces, evaluating the forward model f (θ) and computing these quadratic forms for every proposal often leads to significant computational bottlenecks. log p(x|θ) = −

3.2. Flow Matching Flow matching (FM) is a generative modeling framework that transforms a simple prior distribution p1 (e.g., N (0, I)) into a complex data distribution p0 by learning a deterministic transport map. This transformation is modeled as an ordinary differential equation (ODE) defined by a timedependent vector field vt : Rd → Rd : d θ t = vt (θ t ), dt

(4)

where θ t represents the trajectory of a sample at time t ∈ [0, 1]. The goal is to parameterize vt with a neural network vϕ (θ, t) such that the flow at t = 0 generates samples from the data distribution p0 .

3. Preliminaries

To train vϕ , we employ conditional flow matching (CFM), which regresses the vector field onto a target conditional path defined between a data sample θ 0 ∼ p0 and a noise sample θ 1 ∼ p1 . We specifically use the linear interpolation path (optimal transport path), defined as θ t = (1 − t)θ 0 + tθ 1 . This path implies a constant target velocity: θ̇ t = θ 1 − θ 0 . The model is trained to approximate this target velocity by minimizing the expected mean squared error: h i 2 LFM = Et,θ0 ,θ1 vϕ (θ t , t) − (θ 1 − θ 0 ) , (5)

3.1. Problem Setting We consider the inverse problem of estimating a set of parameters θ ∈ RN given observed data x ∈ RM . The goal is to compute the posterior distribution p(θ|x), which, according to Bayes’ theorem, is proportional to the product of the likelihood and the prior: p(θ|x) ∝ p(x|θ)p(θ).

θ 1 ∼ p1 ,

(1)

Standard approaches for estimating this posterior rely on Markov Chain Monte Carlo (MCMC) methods, such as the Metropolis-Hastings algorithm. These algorithms generate a sequence of samples that asymptotically converges to the target distribution. At each iteration, given the current state θ, a new candidate θ ′ is drawn from a proposal distribution q(θ ′ |θ). The candidate is accepted with probability α(θ ′ |θ), defined as:   p(x|θ ′ )p(θ ′ ) q(θ|θ ′ ) α(θ ′ |θ) = min 1, . (2) p(x|θ)p(θ) q(θ ′ |θ)

where t ∼ U (0, 1). During inference, we generate samples by drawing θ 1 ∼ p1 and numerically integrating Eq. (4) backwards from t = 1 to t = 0 using the learned field vϕ .

4. Methodology As illustrated in Figure 1, we propose a unified framework for efficient and precise parameter estimation, synergizing a multimodal architecture for deep context-parameter fusion with inference-time scaling for enhanced sampling. 3

FUSE: FK-Steered Multi-Modal Flow Matching (a) Model Architecture Reward

(b) FK-Steered Inference

Timestep Embedding Modulation

One Feynman-Kac Steering Step : Estimated Parameter (at t = 0) Trajectories : Resampling

…

…

𝜽𝟏𝟎

𝐩(𝛉), 𝐱

⨀

Fusion Module

Pojection

⨀ Embedder

෡𝟐𝒕 −∆𝑡 𝛉 𝟏

෡𝟏𝒕 𝛉 𝟏 Posterior Predictor

Projection

Embedder

Prior/Observation

෡𝟐𝒕 𝛉 𝟏 ෡𝟑𝒕 𝛉 𝟏 ෡𝟒𝒕 𝛉 𝟏

log𝑝 𝐱 𝜃መ𝑡11 ) + log𝑝(𝜃መ𝑡11 )

෡𝟏𝒕 −∆𝑡 𝛉 𝟏

𝜽𝟒𝟎

log𝑝 𝐱 𝜃መ𝑡21 ) + log𝑝(𝜃መ𝑡21 )

෡𝟒𝒕 −∆𝑡 𝛉 𝟏

𝜽𝟑𝟎

log𝑝 𝐱 𝜃መ𝑡31 ) + log𝑝(𝜃መ𝑡31 )

෡𝟒t −∆𝑡 𝛉 𝟏

log𝑝 𝐱 𝜃መ𝑡41 ) + log𝑝(𝜃መ𝑡41 ) Estimated Posterior Distribution

⨀

Multimodal DiT Blocks

Simulator

×𝐿

𝑡=1

𝜽𝟐𝟎

𝑡 = 𝑡1

Timestep

𝑡 = 𝑡1 − ∆𝑡

t =0

Figure 1. Pipeline of FUSE. (a) FUSE’s Architecture.The architecture employs an independent embedding interface to map heterogeneous inputs into a shared space, followed by an MM-DiT-based fusion module for multimodal integration. Finally, the posterior predictor estimates the denoising velocity based on refined, parameter-wise token representations. (b) FK-Steered Inference. We steer the FUSE sampling process via a guidance mechanism that operates throughout the entire generative trajectory. By leveraging the log-likelihood from the simulator and the prior density of the parameters, the framework dynamically evaluates and rectifies the generative flow.

fxℓ are first processed via modality-specific linear projections to preserve their distinct features. These projections are then concatenated within the joint self-attention mechanism:

Specifically, we adopt a Multimodal Diffusion Transformer (MMDiT) to facilitate bidirectional information exchange between noisy parameter tokens and observational contexts (Section 4.1). We train the model using a conditional rectified flow objective to learn optimal transport trajectories (Section 4.2). Finally, to address challenging scientific inverse problems, we introduce Feynman-Kac steered inference, which leverages test-time compute to dynamically navigate generative trajectories toward high-likelihood regions (Section 4.3). We kindly refer to Appendix B for more implementation details.

[fθℓ+1 , fxℓ+1 ] = Attention ([Qθ , Qx ], [Kθ , Kx ], [Vθ , Vx ]) . (6) where Qm , Km , Vm denote the projected queries, keys, and values for each modality m ∈ {θ, x}. This mechanism enables parameter tokens to actively query relevant sensory features, capturing the cross-modal correlations necessary for high-fidelity posterior estimation. The resulting sequence is finally split back into fθℓ+1 and fxℓ+1 for subsequent layer-wise refinement. This iterative message passing refines the parameter representations layer-by-layer, effectively disentangling the multi-modal structures within the posterior distribution while preserving the specific token grouping size K.

4.1. Model Architecture Input embedding. To address the structural disparity between the physical parameters and sensory data, we design a specialized multi-modal embedding interface that projects these heterogeneous inputs into a shared latent space. Let θ ∈ RN denote the unknown parameters and x ∈ RM the real observation. We represent both as highdimensional embedding sequences with hidden width D. Specifically, we independently embed each scalar parameter θi into K tokens via a simple MLP layer, yielding a sequence of parameter embeddings fθ0 ∈ RNθ ×D where Nθ = N × K. Also, the observation is embedded into Nx context tokens fx0 ∈ RNx ×D , where Nx is a hyperparameter. To encode time information, we use a sinusoidal embedding e(t) ∈ RD together with a context projection to produce modulation parameters for adaptive normalization: a scaling vector γ(t, x) ∈ RD and a shifting vector β(t, x) ∈ RD .

Posterior parameter predictor. Upon obtaining the refined latent representations fθL from the final transformer layer, we decode the predicted flow velocity. We denote the feature sub-sequence for the n-th parameter as L fθ,n ∈ RK×D . To map these high-dimensional features back to the scalar velocity domain, we employ a shared, lightweight linear readout head followed by a global aggregation mechanism: vϕ,n (θ t , x, t) =

1 X Linear(f ). K L

(7)

f ∈fθ,n

Here, the linear layer projects each token independently, while the mean pooling synthesizes the distributed information into a robust scalar estimate. Finally, the full velocity vector is assembled as vϕ = [vϕ,1 , . . . , vϕ,N ]⊤ . This decoupled readout strategy ensures that the complex joint

Multimodal Diffusion Transformer. To bridge the heterogeneous dimensionality of inputs, we adopt a dual-track architecture inspired by MM-DiT (Esser et al., 2024). At each layer ℓ, the parameter tokens fθℓ and observation tokens 4

FUSE: FK-Steered Multi-Modal Flow Matching

is evaluated away from the endpoints of [0, 1], with the denominator 1 − t clipped by a small constant for numerical stability.

dependencies between parameters are fully resolved within the deep transformer layers via self-attention. Thus, the final projection focuses solely on local feature extraction and dimensionality reduction.

Potentials from simulator-based scoring. Given a simulator and a prior, we define the particle reward function rϕ (θ t , t; x) as the unnormalized log-posterior of the clean parameter estimated from the current state. Specifically,

4.2. Training Objective We train a conditional rectified-flow velocity field vϕ (θ t , x, t) with the standard RF regression loss. For each paired sample (x, θ) with θ ∼ pdata (θ | x), draw ϵ ∼ N (0, Idθ ) and t ∼ U (0, 1), and set θ t = t ϵ+(1−t) θ. We minimize h i 2 L(ϕ) = E vϕ (θ t , x, t) − (ϵ − θ) 2 , (8)

rϕ (θ t , t; x) = log p(x | θ̂ t ) + log p(θ̂ t ), with θ̂ t = θ t − t vϕ (θ t , x, t),

where the sign follows the interpolation θt = (1 − t)θ0 + tϵ and velocity target ϵ − θ0 . We evaluate the simulator at this denoised proxy rather than at the noisy intermediate state, since θ t need not satisfy the physical constraints of the simulator. The resulting potentials are best understood as a tractable likelihood-based tilt of the learned path measure; the finite-particle implementation is an approximate posterior-correction mechanism, not an asymptotically exact replacement for MCMC. A more detailed theoretical perspective is provided in Appendix I.

where the expectation is over (x, θ), ϵ, and t. 4.3. Feynman–Kac Steered Inference ODE sampler. For efficient parameter inference, we solve the associated probability flow ordinary differential equation (PF-ODE) to transport samples from a simple base distribution to the posterior. We employ the first-order Euler method (Liu, 2022) to discretize the time horizon t ∈ [1, 0] into T uniform steps. Let {tn }Tn=0 denote the time sequence where t0 = 1 and tT = 0, with ∆t = −1/T . The iterative update is θ tn+1 = θ tn + vϕ (θ tn , x, tn ) ∆t,

(11)

We employ this reward to construct potentials Gt that guide generation. Let R = {tr1 , . . . , trJ } be the intervalresampling schedule; equivalently, non-resampling steps carry the neutral potential Gt = 1. At the j-th resampling time trj , we use   λ Grj (θ trj ) := exp rϕ (θ trj , trj ; x) , (12) J

(9)

where θ t0 is drawn from the base noise distribution and θ tT is the generated posterior sample. FK steering. In complex scientific inverse problems, approximation errors in the learned velocity field can leave amortized samples with excess mass in physically implausible regions. We therefore use Feynman–Kac (FK) steering (Del Moral, 2004; Chopin & Papaspiliopoulos, 2020; Singhal et al., 2025) as a likelihood-guided correction during inference. The role of this step is not to make a finite particle sampler equivalent to MCMC. Instead, it uses simulator likelihoods to bias the learned generative trajectory toward plausible posterior regions while preserving the speed of amortized flow matching. Concretely, we maintain B parallel particles and resample them at selected intermediate times according to FK potentials.

where λ is a global scaling factor and 1/J corresponds to a uniform annealing schedule. This follows the intervalresampling form of FK steering: particles are propagated by the learned sampler and are resampled using normalized potential scores. Since our proposal is the learned reverse transition itself, the transition ratio in the general FK importance weight reduces to one; an alternative proposal would require the corresponding pθ /τ correction. Resampling strategy. To focus computation on highlikelihood regions without turning the sampler into a deterministic post-hoc selector, we resample particles at the predefined times rj . The normalized weights are computed from the FK scores themselves:

To reduce particle depletion after resampling, we use a stochastic sampler (Liu et al., 2025). With h = 1/T and   tn+1 = tn − h, the reverse-time update is written as Grj (θ ktr ) λ j k k  = softmaxk rϕ (θ tr , trj ; x) .  wrj = PB ℓ j J σ2  ℓ=1 Grj (θ trj ) θt−h = θt − vϕ (θt , x, t) − t θt + (1 − t) vϕ (θt , x, t) h 2t (13) √ We then perform multinomial resampling with replacement + σt h ϵ, (10) to obtain the next particle population. High-score trajecp where ϵ ∼ N (0, I), σt = α t/(1 − t), and α controls tories may be copied several times, while trajectories with the stochastic perturbation. In implementation, the update negligible posterior reward are likely to be discarded. The 5

FUSE: FK-Steered Multi-Modal Flow Matching 1.0 1.0

GL-U

GM

Two Moons

0.9

0.8 0.8

GL

0.7

SLCP

C2ST

0.6 0.6 0.5

1.0

B-GLM

0.4

SIR

0.9

0.8

LV

0.2 0.7

B-GLM-Raw

SLCP-D

Simformer NPE FMPE Ours

0.6

0.5 0.0 0.0 3

10

104

105

0.2

103

104

105

0.4

104

103

105

0.6

0.8

104

103

105

103

104

1.0

105

Number of Simulations

Figure 2. Sample efficiency comparison on 10 SBIBM tasks. We evaluate posterior fidelity using observation-wise ℓ-C2ST across simulation budgets from 103 to 105 , where 0.5 indicates an ideal match to the reference posterior. Curves report means over the official SBIBM observations, and shaded regions indicate the corresponding standard deviations. FUSE is evaluated without FK-steering in this benchmark to isolate the contribution of the MM-DiT architecture. While the high-capacity backbone is less competitive in the extremely low-budget regime on some tasks, it scales strongly with simulator coverage and reaches the best data-rich performance on the most challenging settings. M ETHOD NPE FMPE S IMFORMER FUSE

P OSTERIOR F IDELITY

S TATISTICAL D ISCREPANCY

ℓ-C2ST (↓)

MMD (↓)

KL (↓)

S INKHORN (↓)

PME (↓)

E STIMATION ACCURACY PVR (→ 1)

M ED D IST (↓)

0.64 0.66 0.63 0.59

0.046 0.033 0.040 0.016

0.89 0.94 0.66 0.28

0.228 0.115 0.085 0.077

0.19 0.16 0.20 0.06

9.14 3.95 1.55 1.27

1.60 1.62 1.66 1.55

Table 1. Quantitative performance on the SBIBM benchmark. All numbers are averaged over the 10 SBIBM tasks at the 105 simulation budget. FUSE is evaluated without FK-steering to isolate the contribution of the MM-DiT architecture. We report observation-wise ℓ-C2ST averaged over benchmark observations; close to 1 is better for PVR, lower is better for all other metrics. Best results are highlighted in bold.

stochastic propagation steps following resampling reintroduce local diversity, which empirically reduces the underdispersion one would obtain from deterministic best-of-N filtering.

variance calibration; lower values are better for the remaining discrepancy and error metrics. Further details on these metrics can be found in Appendix C. 5.1. Comparison on SBI benchmark

5. Experiments

Dataset. We evaluate on the Simulation-Based Inference Benchmark (SBIBM), implemented in the open-source sbibm framework (Lueckmann et al., 2021). It provides a standardized suite of 10 simulation-based inference tasks spanning simple statistical toy models and more challenging scientific simulators.

Metrics. We evaluate our method and baselines using seven standard metrics: local observation-wise Classifier Two-Sample Test (ℓ-C2ST) (Linhart et al., 2023); Maximum Mean Discrepancy (MMD) (Gretton et al., 2012); and Posterior Mean Error (PME). We further assess the fidelity of posterior spreads using the Posterior Variance Ratio (PVR) and compute the Median Distance (MEDDIST) between predicted samples and the true parameters. To capture higher-order distributional discrepancies, we also report Kullback-Leibler (KL) divergence (Wang et al., 2009) and Sinkhorn distance (Cuturi, 2013). For ℓ-C2ST, a value of 0.5 represents an optimal match for each fixed observation before averaging; for PVR, values closer to 1 indicate better

Across tasks, parameter dimensionalities range from dim(θ) ∈ [2, 10] and observation dimensionalities from dim(x) ∈ [2, 100]. For each task, sbibm ships 10 fixed observations and corresponding 104 reference posterior samples for quantitative evaluation, enabling a direct evaluation of posterior samples through established statistical distances and calibration-focused metrics.

6

FUSE: FK-Steered Multi-Modal Flow Matching 1.0

0

0.6 0.4 0.2

tasks. Our method consistently achieves strong performance across all seven evaluation metrics. In particular, it substantially improves distribution-level metrics such as MMD and KL divergence compared to existing baselines.

MCMC Density

0.8

0.0

Generalization capability.

0 0 0

Figure 2 shows the scalability and sample-efficiency behavior of our method. While all models improve as the simulation budget increases from 103 to 105 , FUSE is not uniformly best in the extremely low-budget regime: the MMDiT backbone has higher capacity and therefore requires sufficient simulator coverage to learn stable parameterobservation interactions. On low-difficulty tasks like Two Moons, GL, and GL-U, most methods eventually reach similar results near the ideal 0.5 threshold, suggesting that simple distributions can often be handled by lighter baselines given enough data. The advantage of FUSE becomes more pronounced on high-difficulty tasks such as LV, SLCP, and SLCP-D. In these challenging settings, which involve more complex posterior structures, existing baselines often plateau above the optimal ℓ-C2ST score, whereas FUSE continues to improve with additional simulations. This supports the intended use case of FUSE: when offline simulator budgets are available, the high-capacity MM-DiT backbone can better exploit complex parameter-observation dependencies and deliver stronger data-rich posterior fidelity.

θ1

θ2

θ3

θ4

2

2

0

−

2

2

0

−

2

2

0

−

2

2 −

0

−

2

θ5

−

θ4

2

−

2

θ3

2

− 2

Overall, these results suggest that the proposed MM-DiT backbone provides a more effective mechanism for modeling complex parameter-observation relationships than standard MLP-based or existing transformer-based architectures. Additional benchmark evaluation details, task-wise statistics, simulation-budget analyses, and ablation studies are provided in Appendices D and G.

2

2

θ2

2

θ1

Ref (MCMC) Ref 68% CI Ref 95% CI NPE FMPE Simformer FUSE w/o FK-Steering FUSE

θ5

Figure 3. FK-steering improves sample quality. We compare our method against several baselines with the posterior samples and marginal distributions on the SLCP task. The addition of FK-steering further concentrates the samples into the high-density regions of the posterior.

Baselines. We compare our approach against three state-of-the-art amortized inference baselines: (1) NPE (NSF) (Papamakarios & Murray, 2016; Durkan et al., 2019), a densityestimation standard utilizing Neural Spline Flows to model complex posteriors; (2) FMPE (Wildberger et al., 2023), a flow-matching baseline that typically employs ResNet-based vector field parameterization; and (3) Simformer (Gloeckler et al., 2024), a transformer-based framework leveraging self-attention to learn robust summary statistics. Since the official implementations of these baselines in the sbibm benchmark follow standardized hyperparameter protocols, we adopt their recommended settings to ensure a fair comparison. Note that we do not incorporate the FK-steering sampling strategy during evaluation by default to isolate the performance gains resulting from architectural differences alone.

FK-Steering improves sample quality. To evaluate the effectiveness of the FK-Steering sampling strategy, we test our model on the challenging SLCP task. Further details on SLCP’s score design can be found in Appendix A. As shown in Figure 3, even without the steering strategy, our model already provides a more accurate approximation of the reference posterior than the baselines.

Evaluation protocol.

When we apply FK-Steering, the samples are further guided toward high-density regions of the posterior. This improves mode localization on SLCP, but it should not be read as a guarantee of tail coverage: a likelihood-tilted particle system can be penalized by C2ST if it concentrates on highprobability regions while covering less of the reference tails. We therefore interpret the FK result as evidence that trajectory-level likelihood guidance improves high-density posterior fidelity and point-estimation behavior in this setting. Additional quantitative and qualitative analyses of FK-steering are provided in Appendix H.

Unless otherwise stated, Table 1 and Figure 2 evaluate FUSE without FK-steering, so the comparison isolates the MMDiT architecture from test-time likelihood guidance. Table 1 reports metrics at the 105 simulation budget, averaged across the 10 SBIBM tasks and the official observations. The C2ST values are computed locally for each observation and then averaged, so we denote them as ℓ-C2ST. Additional comparisons with SMC-ABC, NPSE, SNPE, SNLE, and task-wise budget analyses are provided in Appendix G. Quantitative Results. We present a comprehensive quantitative analysis of the SBI benchmark in Table 1, where all reported metrics are averaged across the ten benchmark 7

FUSE: FK-Steered Multi-Modal Flow Matching

a

b

Figure 4. Comprehensive evaluation of orbital parameter estimation for β Pictoris b. (a) Quantitative comparison: A heatmap of the Normalized Sinkhorn Divergence across all 8 orbital parameters. FUSE achieves the lowest error compared to baselines and the ablation version (w/o FK). Note the significant failure of NPE/FMPE in constraining orientation parameters (i, Ω), which our method resolves effectively. (b) Posterior Reconstruction: Marginalized posterior distributions for all parameters. The FUSE posteriors closely track the PTMCMC reference in high-probability regions, recovering sharp peaks and complex degeneracies that baselines miss.

employ an expanded MM-DiT architecture with D = 256, L = 12, and H = 8. Detailed hyperparameter configurations, including the training schedule and specific FKsteering window settings, are provided in Appendix E. We benchmark FUSE against two state-of-the-art simulationbased inference methods: NPE (Tejero-Cantero et al., 2020; Ruth, 2024) and FMPE (Wildberger et al., 2023), as aforementioned.

5.2. Case Study: Real-time Orbital Characterization of β Pictoris b We evaluate the performance of FUSE by applying it to a real-world orbital parameter estimation of β Pictoris b, a benchmark system in exoplanet science (Lagrange et al., 2010). Located approximately 63 light-years away, this system hosts a young, massive gas giant orbiting within a circumstellar debris disk (Nielsen et al., 2020). Accurately constraining the orbit of β Pictoris b is scientifically critical for understanding dynamic planet-disk interactions (Lacquement et al., 2025) and scheduling follow-up observations with instruments like JWST (Lightsey et al., 2012).

To ensure a fair comparison, all estimators are trained on an identical budget of simulated orbits. We then evaluate their ability to generalize from the simulation before the real observational likelihood. In this real-world case, FUSE uses FK-steering sampling strategy to enhance result quality. As a reference, we employ a long-run Parallel Tempered MCMC (PTMCMC) chain (Vousden et al., 2016), which is designed to asymptotically sample the target posterior under standard convergence assumptions but requires substantial computational resources.

Real-world astrometry and challenge. The inference target is an 8-dimensional state space comprising the Keplerian elements and system properties, specifically total mass and parallax. This inverse problem is notoriously difficult because the available data often span only a small fraction of the total orbital period , leading to complex parameter covariances and highly non-Gaussian, multimodal posteriors (Blunt et al., 2020).

Results. We first quantitatively assess the estimation quality using the Normalized Sinkhorn Divergence between the predicted marginals and the PTMCMC reference. As shown in the heatmap (Figure 4a), FUSE consistently yields the lowest divergence scores (indicated by the lightest col-

Baselines. We demonstrate the practical performance of FUSE by applying it to the orbital parameter estimation of β Pictoris b. For this high-dimensional real-world task, we 8

FUSE: FK-Steered Multi-Modal Flow Matching

ors) across all physical parameters. In contrast, baselines suffer from catastrophic errors in orientation parameters (e.g., i, Ω), marked by high divergence values (red blocks), whereas FUSE maintains high fidelity across the entire parameter space. Qualitatively, the corner plot (Figure 4b) visually corroborates these metrics. The FUSE contours closely align with the high-probability modes of the PTMCMC chain. Notably, our method successfully resolves the sharp peaks of the inclination (i) and the longitude of the ascending node (Ω), whereas baseline methods exhibit significant dispersion. The FUSE posteriors provide sharper constraints than the baselines, reflecting the integrated FK module’s ability to filter out low-probability sampling noise and focus the estimate on the high-likelihood manifold.

act replacement for MCMC. Second, the MM-DiT backbone is data-hungry: when simulator calls are extremely scarce, lighter amortized models or sequential observationspecific methods may remain preferable. Third, while FK steering mitigates approximation errors and improves high-density fidelity, accurately capturing extremely lowprobability tails remains non-trivial. Finally, the scalability of FUSE from single-planet systems (e.g., β Pictoris b) to higher-dimensional problems characterized by complex interactions—such as multi-planet systems or ground-based gravitational wave detection (e.g., LIGO, Virgo, KAGRA)— remains to be empirically validated. Because the MMDiT backbone and FK-steering module are modular, a promising future direction is to adapt FUSE to sequential or observation-specific SBI, where the learned multimodal transport initializes posterior proposals and FK likelihood steering prunes implausible regions during iterative refinement.

An ablation against FUSE without FK-steering confirms that this improvement is not merely due to the amortized backbone. The unsteered model remains more mass-covering and less accurate around sharp degeneracies, whereas FKsteering improves both 95% credible-region overlap and mode accuracy: FUSE improves the posterior overlap from 0.52 to 0.62 and reduces the Mode L2 Distance from 6.60 to 3.85, while FMPE and NPE obtain substantially worse Mode L2 distances of 132.41 and 75.74, respectively. Full quantitative details are given in Appendix H.1.

Conclusion. We have introduced FUSE, a framework integrating multi-modal flow matching with Feynman-Kac steering. By combining a dual-track architecture with likelihood-guided particle steering, FUSE improves the fidelity-efficiency trade-off of amortized SBI. In the β Pictoris b experiment, it recovers the dominant PTMCMC posterior structure at substantially lower inference cost under the reported metrics. Improving robustness to low-SNR observations and non-stationary artifacts remains an important step toward broader astrophysical applications, including gravitational wave astronomy.

We further compare FUSE with a Naive Best-of-N selection strategy to show that applying likelihood and prior scores along the trajectory gives better high-density posterior fidelity on complex degeneracies than selecting only at the final step (see Appendix F for details). Efficiency: Towards Real-time Astronomy.

Acknowledgements

In terms of computational efficiency, while the reference PTMCMC chain typically requires hours (i.e., computation took ≈ 8.5 hours using a 128-core CPU) to converge due to millions of sequential likelihood calculations, FUSE is capable of completing high-precision orbital inference within three minutes. The FK correction uses a small particle ensemble and lightweight simulator-based scoring, so its overhead remains modest relative to the cost of retraining or running a long observation-specific MCMC chain. A scaling breakdown over particle count and scoring frequency is reported in Appendix E.1. This dramatic acceleration enables “on-the-fly” orbital updating, allowing astronomers to refine ephemerides instantaneously as new astrometric data points are received.

This work was supported in part by the National Natural Science Foundation of China under Grant W2431046, the National Key R&D Program of China under Grant 2025YFA1309603, the Central Guided Local Science and Technology Foundation of China under Grant YDZX20253100001001, the Program for Grand Challenges in Basic Research under Grant No. 025GJHZ2025073GC, the MoE Key Lab of Intelligent Perception and HumanMachine Collaboration (ShanghaiTech University), and the Shanghai Frontiers Science Center of Human-centered Artificial Intelligence.

Impact Statement This paper introduces a Flow Matching framework designed to solve complex inverse problems, with a specific focus on applications in astrophysics. By enabling rapid and accurate posterior estimation from high-dimensional observational data, our work has the potential to significantly accelerate scientific discovery and maximize the utility of large-scale astronomical surveys.

6. Discussion Limitations. Despite its high fidelity, FUSE, as an amortized inference method, theoretically lacks the asymptotic exactness guarantees of MCMC, and the implemented finiteparticle FK-steering sampler should be interpreted as a tractable likelihood-guided correction rather than an ex9

FUSE: FK-Steered Multi-Modal Flow Matching

However, applying generative models to scientific inference carries inherent risks. If the model is not properly calibrated or if the simulation-to-real gap is significant, there is a risk of generating biased posteriors or ”hallucinated” physical parameters, which could lead to incorrect scientific conclusions. Therefore, practitioners must rigorously validate these models against domain-specific baselines before deployment.

3468. doi: 1 0 . 1 6 0 9 / a a a i . v 3 8 i 1 8 . 3 0 0 1 8. URL https://ojs.aaai.org/index.php/AAA I/article/view/30018. Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 26, 2013. URL https: //papers.nips.cc/paper_files/paper/2 013/hash/af21d0c97db2e27e13572cbf59e b343d-Abstract.html.

References

Dax, M., Green, S. R., Gair, J., Macke, J. H., Buonanno, A., and Schölkopf, B. Real-time gravitational wave science with neural posterior estimation. Physical Review Letters, 127(24):241103, 2021. doi: 10.1103/PhysRevLett.127. 241103. URL https://doi.org/10.1103/Ph ysRevLett.127.241103.

Bao, F., Nie, S., Xue, K., Cao, Y., Li, C., Su, H., and Zhu, J. All are worth words: A ViT backbone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 22669–22679, 2023. doi: 10.1109/CVPR52729.2023.0 2171. URL https://openaccess.thecvf.co m/content/CVPR2023/html/Bao_All_Are_ Worth_Words_A_ViT_Backbone_for_Diffu sion_Models_CVPR_2023_paper.html.

Del Moral, P. Feynman-Kac Formulae: Genealogical and Interacting Particle Systems with Applications. Springer, 2004. doi: 10.1007/978-1-4684-9393-1. URL https: //doi.org/10.1007/978-1-4684-9393-1.

Blunt, S., Wang, J. J., Angelo, I., Ngo, H., Cody, D., De Rosa, R. J., Graham, J. R., Hirsch, L., Nagpal, V., Nielsen, E. L., Pearce, L., Rice, M., and Tejada, R. orbitize!: A comprehensive orbit-fitting software package for the high-contrast imaging community. The Astronomical Journal, 159(3):89, feb 2020. doi: 10.3847/1538-3 881/ab6663. URL https://dx.doi.org/10.38 47/1538-3881/ab6663.

Dhariwal, P. and Nichol, A. Diffusion Models Beat GANs on Image Synthesis. In Advances in Neural Information Processing Systems, volume 34, pp. 8780–8794. Curran Associates, Inc., 2021. URL https://proceeding s.neurips.cc/paper/2021/hash/49ad23d 1ec9fa4bd8d77d02681df5cfa-Abstract.ht ml.

Brooks, T., Holynski, A., and Efros, A. A. InstructPix2Pix: Learning to follow image editing instructions. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 18392–18402, 2023. doi: 10.1109/CVPR52729.2023.01764. URL https://openaccess.thecvf.com/conten t/CVPR2023/html/Brooks_InstructPix2P ix_Learning_To_Follow_Image_Editing_ Instructions_CVPR_2023_paper.html.

Du, M., Liang, B., Wang, H., Xu, P., Luo, Z., and Wu, Y. Advancing space-based gravitational wave astronomy: Rapid parameter estimation via normalizing flows. Science China Physics, Mechanics & Astronomy, 67(3): 230412, 2024. doi: 10.1007/s11433-023-2270-7. URL https://doi.org/10.1007/s11433-023-2 270-7. Durkan, C., Bekasov, A., Murray, I., and Papamakarios, G. Neural spline flows. In Advances in Neural Information Processing Systems, volume 32, 2019. URL https: //proceedings.neurips.cc/paper/2019/ hash/7ac71d433f282034e088473244df8c0 2-Abstract.html.

Chen, J., Yu, J., Ge, C., Yao, L., Xie, E., Wu, Y., Wang, Z., Kwok, J., Luo, P., Lu, H., and Li, Z. PixArt-α: Fast training of diffusion transformer for photorealistic text-toimage synthesis. In International Conference on Learning Representations, 2024. URL https://openreview .net/forum?id=eAKmQPe3m1.

Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12873–12883, 2021. doi: 10 .1109/CVPR46437.2021.01268. URL https: //openaccess.thecvf.com/content/CVPR 2021/html/Esser_Taming_Transformers_ for_High-Resolution_Image_Synthesis_ CVPR_2021_paper.html.

Chopin, N. and Papaspiliopoulos, O. An Introduction to Sequential Monte Carlo. Springer, 2020. doi: 10.1007/ 978-3-030-47845-2. URL https://doi.org/10 .1007/978-3-030-47845-2. Cobb, A. D., Matejek, B., Elenius, D., Roy, A., and Jha, S. Direct Amortized Likelihood Ratio Estimation. Proceedings of the AAAI Conference on Artificial Intelligence, 38(18):20362–20369, March 2024. ISSN 237410

FUSE: FK-Steered Multi-Modal Flow Matching

Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, Z., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 12606– 12633. PMLR, 2024. URL https://proceedings. mlr.press/v235/esser24a.html. Foreman-Mackey, D., Hogg, D. W., Lang, D., and Goodman, J. emcee: The MCMC hammer. Publications of the Astronomical Society of the Pacific, 125(925):306–312, 2013. doi: 10.1086/670067. URL https://doi.or g/10.1086/670067.

Machine Learning, volume 119 of Proceedings of Machine Learning Research. PMLR, 2020. URL https: //proceedings.mlr.press/v119/hermans 20a.html. Ho, J. and Salimans, T. Classifier-Free Diffusion Guidance, July 2022. URL http://arxiv.org/abs/2207 .12598. arXiv:2207.12598 [cs]. Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840–6851, 2020. URL https://proceedings.neurips.cc/p aper/2020/hash/4c5bcfec8584af0d967f1 ab10179ca4b-Abstract.html.

Geffner, T., Papamakarios, G., and Mnih, A. Compositional score modeling for simulation-based inference. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research. PMLR, 2023. URL https: //proceedings.mlr.press/v202/geffner 23a.html.

Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems, volume 35, 2022. doi: 10.52202/068431-1926. URL https://proceedings.neurips.cc/paper _files/paper/2022/hash/a98846e9d9cc0 1cfb87eb694d946ce6b-Abstract-Confere nce.html.

Glöckler, M., Deistler, M., and Macke, J. H. Variational methods for simulation-based inference. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=kZ0U YdhqkNY.

Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014. URL https://openreview.net /forum?id=33X9fd2-9FyZd.

Gloeckler, M., Deistler, M., Weilbach, C. D., Wood, F., and Macke, J. H. All-in-one simulation-based inference. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 15735–15766. PMLR, 2024. URL https://proceedings.mlr.press/v235/g loeckler24a.html.

Lacquement, A., Beust, H., Faramaz-Gorka, V., and Duchêne, G. Dynamics of the beta pictoris planetary system and possibility of an additional planet. Astronomy & Astrophysics, 694:A236, February 2025. ISSN 1432-0746. doi: 10.1051/0004-6361/202452632. URL http://dx.doi.org/10.1051/0004-6361/ 202452632.

Gonçalves, P. J., Lueckmann, J.-M., Deistler, M., Nonnenmacher, M., Öcal, K., Bassetto, G., Chintaluri, C., Podlaski, W. F., Haddad, S. A., Vogels, T. P., Greenberg, D. S., and Macke, J. H. Training deep neural density estimators to identify mechanistic models of neural dynamics. eLife, 9:e56261, September 2020. ISSN 2050-084X. doi: 1 0 . 7 5 5 4 / e L i f e . 5 6 2 6 1. URL https://doi.org/10.7554/eLife.56261. Publisher: eLife Sciences Publications, Ltd.

Lagrange, A.-M., Bonnefoy, M., Chauvin, G., Apai, D., Ehrenreich, D., Boccaletti, A., Gratadour, D., Rouan, D., Mouillet, D., Lacour, S., and Kasper, M. A giant planet imaged in the disk of the young star β pictoris. Science, 329(5987):57–59, 2010. doi: 10.1126/science.1187187. Li, Y., Liu, H., Wu, Q., Mu, F., Yang, J., Gao, J., Li, C., and Lee, Y. J. GLIGEN: Open-set grounded text-to-image generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 22511–22521, 2023. URL https://openaccess .thecvf.com/content/CVPR2023/html/Li _GLIGEN_Open-Set_Grounded_Text-to-Ima ge_Generation_CVPR_2023_paper.html.

Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723–773, 2012. URL https://jmlr.org/papers/v13/gret ton12a.html.

Liang, B., Du, M., Wang, H., Xu, Y., Liu, C., Wei, X., Xu, P., e Qiang, L., and Luo, Z. Rapid parameter estimation for merging massive black hole binaries using continuous normalizing flows. Machine Learning:

Hermans, J., Begy, V., and Louppe, G. Likelihood-free MCMC with amortized approximate ratio estimators. In Proceedings of the 37th International Conference on 11

FUSE: FK-Steered Multi-Modal Flow Matching

Science and Technology, 5(4):045040, nov 2024. doi: 1 0 . 1 0 8 8 / 2 6 3 2 - 2 1 5 3 / a d 8 d a9. URL h t t p s : //dx.doi.org/10.1088/2632-2153/ad8da9.

Liu, J., Liu, G., Liang, J., Li, Y., Liu, J., Wang, X., Wan, P., Zhang, D., and Ouyang, W. Flow-GRPO: Training flow matching models via online RL. In Advances in Neural Information Processing Systems, volume 38, 2025. URL https://proceedings.neurips.cc/paper _files/paper/2025/hash/3a10c46572628 d58cb44fb705f25cbbf-Abstract-Confere nce.html.

Liang, B., Guo, H., Zhao, T., Wang, H., Evangelinelis, H., Xu, Y., Liu, C., Liang, M., Wei, X., Yuan, Y., Du, M., Xu, P., Qian, W., and Luo, Z. Unlocking new paths for efficient analysis of gravitational waves from extreme-mass-ratio inspirals with machine learning. Chinese Physics Letters, 42(8):081101, aug 2025a. doi: 10.1088/0256-307X/42/8/081101. URL https://do i.org/10.1088/0256-307X/42/8/081101.

Liu, Q. Rectified Flow: A Marginal Preserving Approach to Optimal Transport, September 2022. URL http:// arxiv.org/abs/2209.14577. arXiv:2209.14577 [stat].

Liang, B., Liu, C., Song, H., Lyu, Z., Du, M., Xu, P., Luo, Z., He, S., Gu, H., Zhao, T., Liang, M., Xu, Y., e Qiang, L., Sun, M., and Qian, W.-L. Toward efficient and accurate emri parameter estimation: A machine learning-enhanced mcmc framework. Research, 8:1055, 2025b. doi: 10.3 4133/research.1055. URL https://spj.science. org/doi/abs/10.34133/research.1055. Liang, B., Liu, C., Zhao, T., Du, M., Liang, M., Shi, R., Guo, H., Xu, Y., Qiang, L.-e., Xu, P., Qian, W.-L., and Luo, Z. Accelerating stochastic gravitational wave backgrounds parameter estimation in pulsar timing arrays with flow matching. Physics Letters B, 878:140551, 2026a. doi: 10.1016/j.physletb.2026.140551. URL https://doi. org/10.1016/j.physletb.2026.140551. Liang, B., Song, H., Liu, C., Zhao, T., Xu, Y., Xiao, Z., Liang, M., Du, M., Qian, W.-L., Qiang, L.-e., Xu, P., and Luo, Z. Estimating orbital parameters of direct imaging exoplanet using neural networks. Physical Review Research, 2026b. doi: 10.1103/ykfh- cdzk. URL https://doi.org/10.1103/ykfh-cdzk.

Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. DPMSolver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. Advances in Neural Information Processing Systems, 35:5775–5787, December 2022. URL https://proceedings.neurip s.cc/paper_files/paper/2022/hash/260 a14acce2a89dad36adc8eefe7c59e-Abstrac t-Conference.html. Lueckmann, J.-M., Boelts, J., Greenberg, D. S., Gonçalves, P. J., and Macke, J. H. Benchmarking simulation-based inference. In Banerjee, A. and Fukumizu, K. (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 343–351. PMLR, 13–15 Apr 2021. URL https://proceedings.ml r.press/v130/lueckmann21a.html. Marjoram, P., Molitor, J., Plagnol, V., and Tavaré, S. Markov chain Monte Carlo without likelihoods. Proceedings of the National Academy of Sciences, 100(26):15324– 15328, December 2003. ISSN 0027-8424, 1091-6490. doi: 10.1073/pnas.0306899100. URL https://pnas .org/doi/full/10.1073/pnas.0306899100.

Lightsey, P. A., Atkinson, C., Clampin, M., and Feinberg, L. D. James webb space telescope: large deployable cryogenic telescope in space. Optical Engineering, 51(1): 011003, 2012. doi: 10.1117/1.OE.51.1.011003. URL https://doi.org/10.1117/1.OE.51.1.01 1003.

Nielsen, E. L., De Rosa, R. J., Wang, J. J., Sahlmann, J., Wallace, J. K., Ward-Duong, K., Wiktorowicz, S., and Wolff, S. The Gemini Planet Imager Exoplanet Survey: Dynamical Mass of the Exoplanet β Pictoris b from Combined Direct Imaging and Astrometry. The Astronomical Journal, 159(2):71, February 2020. doi: 10.3847/1538-3881/ab5b92.

Linhart, J., Gramfort, A., and Rodrigues, P. L-C2ST: Local diagnostics for posterior approximations in simulationbased inference. In Advances in Neural Information Processing Systems, volume 36, 2023. URL https: //papers.nips.cc/paper/2023/hash/b03 13c2f4501a81d0e0d4a1e8fbf4995-Abstrac t-Conference.html.

Oord, A. v. d., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Advances in Neural Information Processing Systems, volume 30, 2017. URL https://proceedings.neurips.cc/paper /2017/hash/7a98af17e63a0ac09ce2e96d0 3992fbc-Abstract.html.

Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In International Conference on Learning Representations, 2023. URL https://openreview.net/forum ?id=PqvMRDCJT9t.

Papamakarios, G. and Murray, I. Fast ϵ-free inference of simulation models with bayesian conditional density 12

FUSE: FK-Steered Multi-Modal Flow Matching

estimation. In Advances in Neural Information Processing Systems, volume 29, 2016. URL h t t p s : //proceedings.neurips.cc/paper/201 6/hash/6aca97005c68f1206823815f66102 863-Abstract.html.

ence on Learning Representations, 2022. URL https: //openreview.net/forum?id=TIdIXIpzhoI. Shi, Y., Siddharth, N., Paige, B., and Torr, P. H. S. Variational mixture-of-experts autoencoders for multi-modal deep generative models. In Advances in Neural Information Processing Systems, volume 32, 2019. URL https://proceedings.neurips.cc/pap er/2019/hash/0ae775a8cb3b499ad1fca94 4e6f5c836-Abstract.html.

Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 4195–4205, 2023. URL https://openaccess.thecvf.co m/content/ICCV2023/html/Peebles_Scal able_Diffusion_Models_with_Transform ers_ICCV_2023_paper.html.

Singhal, R., Horvitz, Z., Teehan, R., Ren, M., Yu, Z., Mckeown, K., and Ranganath, R. A general framework for inference-time scaling and steering of diffusion models. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pp. 55810–55827. PMLR, 2025. URL https://proceedings.mlr.pres s/v267/singhal25b.html.

Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8748–8763. PMLR, 2021. URL https://proceedings.mlr.press/v139/r adford21a.html.

Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a. URL https://openrevi ew.net/forum?id=St1giarCHLP.

Ramesh, P., Lueckmann, J.-M., Boelts, J., Tejero-Cantero, Á., Greenberg, D. S., Goncalves, P. J., and Macke, J. H. GATSBI: Generative Adversarial Training for Simulation-Based Inference. In International Conference on Learning Representations, 2022. URL https: //openreview.net/forum?id=kR1hC6j48Tp. Ronneberger, O., Fischer, P., and Brox, T. U-Net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, volume 9351 of Lecture Notes in Computer Science, pp. 234–241. Springer, 2015. doi: 10.1007/978-3-319-24574-4 28. URL https://doi. org/10.1007/978-3-319-24574-4_28. Ruiz, N., Li, Y., Jampani, V., Pritch, Y., Rubinstein, M., and Aberman, K. DreamBooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 22500–22510, 2023. URL https://openaccess.thecvf.co m/content/CVPR2023/html/Ruiz_DreamBo oth_Fine_Tuning_Text-to-Image_Diffusi on_Models_for_Subject-Driven_Generati on_CVPR_2023_paper.html.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b. URL https://openreview.net/forum?id=PxTI G12RRHS. Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consistency models. In International Conference on Machine Learning, 2023. URL https://openreview.net /forum?id=FmqFfMTNnv. Srinivasan, R., Barausse, E., Korsakova, N., and Trotta, R. Simulation-based population inference of LISA’s galactic binaries: Bypassing the global fit. Physical Review D, 112(10):103043, 2025. doi: 10.1103/shym-w46f. URL https://doi.org/10.1103/shym-w46f. Sutter, T. M., Daunhawer, I., and Vogt, J. E. Generalized multimodal ELBO. In International Conference on Learning Representations, 2021. URL h t t p s : //openreview.net/forum?id=5Y21V0RDBV. Suzuki, M., Nakayama, K., and Matsuo, Y. Joint Multimodal Learning with Deep Generative Models, November 2016. URL http://arxiv.org/abs/1611.0 1891. arXiv:1611.01891 [stat].

Ruth, M. Exoplanet Orbital Characterization Using Simulation-Based Inference. PhD thesis, University of Liège, Liège, Belgium, June 2024. URL https://ma theo.uliege.be/handle/2268.2/20393.

Tavaré, S., Balding, D. J., Griffiths, R. C., and Donnelly, P. Inferring Coalescence Times From DNA Sequence Data. Genetics, 145(2):505–518, February 1997. ISSN 1943-2631. doi: 10.1093/genetics/145.2.505. URL

Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. In International Confer13

FUSE: FK-Steered Multi-Modal Flow Matching

assets. ACM Transactions on Graphics, 43(4):1–20, 2024. doi: 10.1145/3658146. URL https://doi.org/10 .1145/3658146.

https://academic.oup.com/genetics/ar ticle/145/2/505/6018089. Tejero-Cantero, A., Boelts, J., Deistler, M., Lueckmann, J.-M., Durkan, C., Gonçalves, P. J., Greenberg, D. S., and Macke, J. H. sbi: A toolkit for simulation-based inference. Journal of Open Source Software, 5(52):2505, 2020. doi: 10.21105/joss.02505. URL https: //doi.org/10.21105/joss.02505. Vousden, W. D., Farr, W. M., and Mandel, I. Dynamic temperature selection for parallel tempering in Markov chain Monte Carlo simulations. Monthly Notices of the Royal Astronomical Society, 455(2):1919–1937, January 2016. doi: 10.1093/mnras/stv2422. Wang, Q., Kulkarni, S. R., and Verdú, S. Divergence estimation for multidimensional densities via k-nearestneighbor distances. IEEE Transactions on Information Theory, 55(5):2392–2405, 2009. doi: 10.1109/TIT.2009 .2016060. Wildberger, J., Dax, M., Buchholz, S., Green, S. R., Macke, J. H., and Schölkopf, B. Flow matching for scalable simulation-based inference. In Advances in Neural Information Processing Systems, volume 36, pp. 16837–16864. Curran Associates, Inc., 2023. URL https://proc eedings.neurips.cc/paper_files/paper /2023/hash/3663ae53ec078860bb0b9c660 6e092a0-Abstract-Conference.html. Wood, S. N. Statistical inference for noisy nonlinear ecological dynamic systems. Nature, 466(7310):1102–1104, August 2010. ISSN 1476-4687. doi: 10.1038/nature09319. URL https://www.nature.com/articles/ nature09319. Publisher: Nature Publishing Group. Wu, M. and Goodman, N. Multimodal generative models for scalable weakly-supervised learning. In Advances in Neural Information Processing Systems, volume 31, 2018. URL https://proceedings.neurips.cc/p aper/2018/hash/1102a326d5f7c9e04fc3c 89d0ede88c9-Abstract.html. Zhang, L., Rao, A., and Agrawala, M. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 3836–3847, 2023. doi: 10.1109/ICCV 51070.2023.00355. URL https://openaccess.t hecvf.com/content/ICCV2023/html/Zhan g_Adding_Conditional_Control_to_Text -to-Image_Diffusion_Models_ICCV_2023_ paper.html. Zhang, L., Wang, Z., Zhang, Q., Qiu, Q., Pang, A., Jiang, H., Yang, W., Xu, L., and Yu, J. CLAY: A controllable large-scale generative model for creating high-quality 3D 14

FUSE: FK-Steered Multi-Modal Flow Matching

A. Exact log-posterior scores for SLCP This appendix specifies the score functions used to evaluate candidate parameters via an unnormalized log-posterior. We consider the SBIBM task SLCP. The score is designed to be an exact evaluation of log p(θ | xobs ) = log p(θ) + log p(xobs | θ) − log p(xobs ),

(14)

and we drop the constant log p(xobs ) w.r.t. θ. A.1. Generic posterior objective and the score used in sampling Given an observation xobs , Bayesian inference targets p(θ | xobs ) ∝ p(θ) p(xobs | θ).

(15)

A common objective for posterior maximization (MAP) is θMAP ∈ arg max log p(θ) + log p(xobs | θ).

(16)

θ

Likewise, MCMC methods (e.g., Metropolis–Hastings) only require an unnormalized log target: S(θ; xobs ) := log p(θ) + log p(xobs | θ),

(17)

since log p(xobs ) cancels in acceptance ratios: log α(θ → θ′ ) = S(θ′ ; xobs ) − S(θ; xobs ) + log q(θ | θ′ ) − log q(θ′ | θ).

(18)

We therefore implement task-specific log p(xobs | θ) and combine it with the task prior log p(θ) to form S. A.2. SLCP score Forward simulator. SLCP uses a 5D parameter θ = (θ1 , θ2 , θ3 , θ4 , θ5 ) ∈ R5 and produces N = 4 i.i.d. 2D observations. Define the transformed quantities µ(θ) = (θ1 , θ2 )⊤ ∈ R2 ,

(19)

s1 (θ) = θ32 ,

s2 (θ) = θ42 ,

(20)

ρ(θ) = tanh(θ5 ) ∈ (−1, 1),

(21)

and the covariance matrix  s1 (θ)2 ρ(θ) s1 (θ) s2 (θ) + εI2 , ρ(θ) s1 (θ) s2 (θ) s2 (θ)2

 Σ(θ) = The simulator draws

i.i.d. xn | θ ∼ N µ(θ), Σ(θ) ,

ε > 0.

n = 1, . . . , N,

(22)

(23)

and returns x = {xn }N n=1 , optionally flattened. Exact log-likelihood.

Given xobs = {xobs,n }N n=1 , the SLCP likelihood factorizes: log p(xobs | θ) =

N X

log N xobs,n ; µ(θ), Σ(θ) .

(24)

n=1

This is evaluated analytically via a multivariate normal density. Prior and score.

SLCP uses an independent uniform prior (as implemented by the benchmark): p(θ) =

5 Y

Unif(θd ; −3, 3),

⇒

log p(θ) =

d=1

5 X

log Unif(θd ; −3, 3).

(25)

d=1

The score is then SSLCP (θ; xobs ) = log p(θ) + log p(xobs | θ), with SSLCP (θ; xobs ) = −∞ if θ is outside the prior support or if numerical checks fail. 15

(26)

FUSE: FK-Steered Multi-Modal Flow Matching

Implementation notes. We (i) reshape flattened observations to [B, N, 2] for a batch of B candidates, (ii) build Σ(θ) as in (22) with a small diagonal jitter ε to guarantee positive semi definiteness, and (iii) sum per-point log densities as in (24). The score is computed in the physical parameter space; if candidates are produced in a standardized space, we first invert the standardization map before evaluating SSLCP .

B. Implementation details We evaluate on the SBIBM benchmark using a unified lightweight Multi-Modal Diffusion Transformer (MM-DiT) with D = 128, L = 6, H = 4 by default. Motivated by their distinct physical semantics, our architecture treats parameters θ, context x, and time t as separate modalities rather than a concatenated vector. Context x is projected to 4 tokens. For θ, we employ individual parameter tokenization, mapping each scalar dimension to 2 tokens (reduced to 1 for the simpler Two Moons and Gaussian tasks). Training utilizes 100, 000 simulations via Adam (batch 512, cosine annealing schedule 2 × 10−4 → 1 × 10−6 ) for 100 epochs. Basic Inference uses an Euler ODE solver with N = 200 steps. For FK-steered inference on SLCP task, we maintain a beam width B of 8 particles. Resampling is performed every 5 steps within the window of step 20 to 200, with a noise scale α = 0.3.

C. Detailed Definitions of Metrics In this section, we provide the formal mathematical definitions for the seven metrics used to evaluate the performance of our method and the baselines. In the following definitions, we assume a set of N samples generated by the model, {θ̂i }N i=1 , and a set of N reference posterior samples, {θi }N i=1 . Classifier Two-Sample Test (C2ST). C2ST measures the similarity between two distributions by training a binary classifier to distinguish between generated and reference posterior samples. The dataset is divided into training and testing sets with equal representation from both distributions. The metric is defined as the classification accuracy on the test set: Ntest 1 X C2ST = I(ŷj = yj ), Ntest j=1

(27)

where I is the indicator function, yj is the true label, and ŷj is the predicted label. An ideal match results in an accuracy of 0.5, indicating that the classifier cannot distinguish between the two sets. Maximum Mean Discrepancy (MMD). MMD quantifies the distance between two distributions P and Q by comparing their embeddings in a Reproducing Kernel Hilbert Space (RKHS). Given a kernel function k(·, ·), the squared MMD is defined as: MMD2 (P, Q) = Eθ,θ′ ∼P [k(θ, θ′ )] + Eθ̂,θ̂′ ∼Q [k(θ̂, θ̂′ )] − 2Eθ∼P,θ̂∼Q [k(θ, θ̂)]. (28) In our implementation, we use a multi-scale Gaussian kernel following the standard protocol in the sbibm benchmark. Posterior Mean Error (PME). PME evaluates the accuracy of the posterior mean by calculating the L2 distance between the average of the predicted samples and the average of the reference posterior samples: PME = ∥E[θ̂] − E[θ]∥2 .

(29)

This metric reflects the bias of the model in identifying the center of the posterior distribution. Posterior Variance Ratio (PVR). PVR assesses how well the model captures the spread or uncertainty of the posterior. It is defined as the ratio of the predicted variance to the reference posterior variance, averaged over all parameter dimensions d: D

PVR =

1 X Var(θ̂d ) . D Var(θd ) d=1

A PVR value close to 1 indicates that the predicted distribution has a similar scale to the reference posterior.

16

(30)

FUSE: FK-Steered Multi-Modal Flow Matching

Median Distance (MEDDIST). MEDDIST provides a robust measure of the distance between the generated samples and the true parameters θ∗ used to generate the observation: MEDDIST = median(∥θ̂i − θ∗ ∥2 ).

(31)

Here, θ∗ denotes the simulator ground-truth parameter associated with the observation when it is available. This metric captures the typical error of the predicted samples in the parameter space. Kullback-Leibler (KL) Divergence. KL divergence measures the information discrepancy between the true distribution P and the predicted distribution Q: Z p(θ) DKL (P ∥Q) = p(θ) log dθ. (32) q(θ) Since the density q(θ) is typically not available in a closed form for generative models, we use a k-nearest neighbor estimator to calculate this value directly from the samples. Sinkhorn Distance. Sinkhorn distance provides a geometric measure of the discrepancy between the predicted and reference distributions by approximating the optimal transport cost. We first compute the cost matrix C using the squared Euclidean distance between the predicted samples θ̂ and the reference samples θ, such that Cij = ∥θ̂i − θj ∥22 . To ensure numerical stability across different tasks, we normalize C by the maximum squared distance found within the reference samples. The distance is then obtained by solving an entropy-regularized optimal transport problem: X Sε = Pij Cij , (33) i,j

where P is the optimal coupling matrix and ε is the regularization parameter. In our implementation, we solve for P using Sinkhorn iterations in the log-domain, which prevents numerical underflow and ensures more stable convergence during the optimization process.

D. Ablation Study We conduct a systematic ablation study on the SLCP task to evaluate the impact of different architectural choices. The SLCP task is chosen for this analysis because its complex posterior structure effectively highlights the performance differences between configurations. All experiments are performed on a single H20 GPU using a fixed budget of 105 simulations and trained for 100 epochs. D

Individual

Tokens/θ

Time

C2ST

KL

128 128 128 64 256

False True True True True

1 1 4 2 2

9.37 9.28 20.32 6.28 32.49

0.6734 0.6458 0.6581 0.6553 0.7583

0.3711 0.3547 0.4336 0.3765 0.8721

128

True

2

12.48

0.6336

0.2820

Table 2. Ablation study on the SLCP task. We analyze the influence of the hidden dimension (D), the individual embedding strategy (Individual), and the number of tokens per parameter (Tokens/θ). We report the average inference time for 104 samples(Time), C2ST, and KL divergence. The optimal configuration in the final row achieves the best balance between inference efficiency and posterior quality. Best results are highlighted in bold.

As shown in Table 2, we first compare embedding strategies. Mapping each parameter component individually (Individual=True) consistently yields better C2ST and KL scores than shared embeddings. For token allocation (Tokens/θ), using 2 tokens per parameter provides the best efficiency; increasing this to 4 tokens significantly slows down inference without improving accuracy. Regarding the hidden dimension (D), the 128-dimensional setting proves optimal. While D = 64 is faster, it lacks the capacity for accurate estimation, and D = 256 leads to overfitting and higher computational costs. Based on these findings, we select D = 128, individual embeddings, and 2 tokens per parameter as our default configuration.

17

FUSE: FK-Steered Multi-Modal Flow Matching

E. Implementation details of Real-time Orbital Characterization of β Pictoris b We evaluate on Real-time Orbital Characterization of β Pictoris b using a unified large Multi-Modal Diffusion Transformer with Hiddendimensions = 256, Layers = 12, N umberheads = 8 by default. Motivated by their distinct physical semantics, our architecture treats parameters θ, context x, and time t as separate modalities rather than a concatenated vector. Context x is projected to 12 tokens. For θ, we employ individual parameter tokenization, mapping each scalar dimension to 1 tokens. Training utilizes 800, 000 simulations via Adam (batch 4096, cosine annealing schedule 5 × 10−4 → 1 × 10−6 ) for 200 epochs. Basic inference uses an Euler ODE solver with N = 200 steps. For FK-steered inference, we maintain a beam width B of 8 particles. Resampling is performed every 5 steps within the window of step 20 to 200, with a noise scale α = 0.3. E.1. Inference-Time Compute and FK Overhead The orbital-characterization experiment uses a compact FK configuration: B = 8 particles, resampling every 5 steps in the window from step 20 to step 200, and noise scale α = 0.3. This adds simulator-based scoring during inference, but the scoring calls are lightweight compared with a long observation-specific PTMCMC chain. In our reported setting, PTMCMC requires approximately 8.5 hours on a 128-core CPU, whereas FUSE completes high-precision inference within three minutes under our setup. Particles

FK-scoring steps

Time

8 8 4 16

17 81 17 17

152.77s 172.02s 71.11s 278.20s

Table 3. Inference-time overhead of FK-steering on the orbital-characterization experiment. Increasing the particle count has the dominant effect on runtime, while increasing the number of FK-scoring steps from 17 to 81 introduces a smaller additional cost in this setup.

These measurements indicate that, for this simulator and implementation, FK-steering scales primarily with the number of particles, while the simulator-based scoring frequency contributes a smaller overhead. This makes the chosen FK configuration practical for the reported real-time orbital-characterization setting.

F. Comparison with Naive Best-of-N Selection To further demonstrate the efficacy of our FK steered sampling, we compare FUSE against a “Naive Best-of-N” baseline. In this baseline, N independent samples are generated using the standard flow-matching decoder, and only the sample with the highest likelihood (or lowest discrepancy) according to the simulator is selected at the final step. It is important to note that such a post-generation selection strategy is generic and could be readily applied to existing methods like NPE and FMPE. As illustrated in Figure 5, while the Naive Best-of-N approach can identify high-probability candidates, it often fails to capture the intricate topological structures and sharp degeneracies of the reference posterior. In contrast, FUSE uses FK-steering to apply likelihood guidance at multiple intermediate timesteps. This trajectory-level correction concentrates computation on high-density regions more effectively than a single final selection. The results indicate that dynamic steering is more closely aligned with the MCMC reference in the reported metrics, particularly for complex parameter dependencies where static selection remains suboptimal.

18

FUSE: FK-Steered Multi-Modal Flow Matching

Figure 5. Posterior distribution comparison between FUSE and Naive Best-of-N selection. The corner plot displays the marginalized 1D and 2D posteriors for the β Pictoris b orbital parameters. The Naive Best-of-N strategy (orange) performs selection only at the final generation step and does not recover several sharp degeneracies and correlations present in the reference. In contrast, FUSE (blue) uses sequential FK-steering to guide trajectories during generation, yielding a posterior that more closely follows the MCMC reference (black dashed/grey) in high-probability regions.

G. Additional Benchmark Evaluation Details G.1. Evaluation Protocol We provide additional implementation and evaluation details for the SBIBM benchmark experiments. All experiments follow the official SBIBM evaluation protocol introduced by Lueckmann et al. (2021). For Table 1 in the main paper, all reported metrics are computed at the 105 simulation budget and averaged across the 10 19

FUSE: FK-Steered Multi-Modal Flow Matching

official observations provided by SBIBM for each task. The reported standard deviations correspond to the variability across these official observations. For the SBIBM benchmark experiments, we disable FK-steering and evaluate FUSE purely as an amortized flow-matching posterior estimator. This setting isolates the effect of the proposed MM-DiT architecture from the additional inference-time guidance introduced by FK-steering. In addition, the originally reported C2ST metric was computed independently for each observation and averaged afterward, which corresponds to the local classifier two-sample test metric (ℓ-C2ST). To avoid ambiguity, we explicitly denote this metric as ℓ-C2ST in the revised manuscript. G.2. Stability Across Official Observations Table 4 reports task-wise ℓ-C2ST results at the 105 simulation budget. Results are reported as mean ± standard deviation across the 10 official SBIBM observations. We select SIR, SLCP, and LV as representative tasks covering different posterior complexities and dynamical behaviors. Method

SIR

SLCP

LV

Simformer FMPE NPE SMC-ABC NPSE SNPE SNLE FUSE w/o FK

0.6878 ± 0.0533 0.5971 ± 0.0429 0.5661 ± 0.0322 0.6244 ± 0.0322 0.5653 ± 0.0204 0.5506 ± 0.0238 0.5847 ± 0.0331 0.5534 ± 0.0122

0.7229 ± 0.0488 0.8836 ± 0.0443 0.8364 ± 0.0430 0.9589 ± 0.0222 0.6949 ± 0.0980 0.6656 ± 0.0608 0.5781 ± 0.0293 0.6290 ± 0.0468

0.9040 ± 0.0493 0.9162 ± 0.0495 0.9456 ± 0.0449 0.9950 ± 0.0034 0.8718 ± 0.0699 0.9281 ± 0.0796 0.6951 ± 0.1025 0.8540 ± 0.0888

Table 4. Task-wise ℓ-C2ST results at the 105 simulation budget. Results are reported as mean ± standard deviation across the 10 official SBIBM observations. Lower values indicate better posterior fidelity.

G.3. Performance Across Simulation Budgets We further analyze ℓ-C2ST on the challenging SLCP task under different simulation budgets. FUSE exhibits reduced performance in the extremely low-data regime, which is consistent with the larger model capacity of the MM-DiT architecture. However, its performance improves steadily as the simulation budget increases, leading to competitive posterior fidelity at larger scales. Method

103

104

105

Simformer FMPE NPE SMC-ABC NPSE SNPE SNLE FUSE w/o FK

0.9147 ± 0.0281 0.9695 ± 0.0148 0.9524 ± 0.0203 0.9771 ± 0.0097 0.9695 ± 0.0157 0.9652 ± 0.0194 0.9206 ± 0.0251 0.9997 ± 0.0002

0.7902 ± 0.0549 0.9636 ± 0.0199 0.8940 ± 0.0420 0.9624 ± 0.0163 0.8255 ± 0.0689 0.8445 ± 0.0451 0.7131 ± 0.0524 0.8537 ± 0.0747

0.7229 ± 0.0488 0.8836 ± 0.0443 0.8364 ± 0.0430 0.9589 ± 0.0222 0.6949 ± 0.0980 0.6656 ± 0.0608 0.5781 ± 0.0293 0.6290 ± 0.0468

Table 5. SLCP ℓ-C2ST under different simulation budgets. Results are reported as mean ± standard deviation across the 10 official observations. Lower values indicate better posterior fidelity.

H. Additional Analysis of FK-Steering H.1. Quantitative Analysis on Exoplanet Posterior Estimation To further evaluate the effect of FK-steering on complex posterior estimation, we compare posterior overlap and mode accuracy on the exoplanet orbital characterization task. As shown in Table 6, the unsteered model (FUSE w/o FK-Steering) achieves relatively broad posterior coverage but suffers 20

FUSE: FK-Steered Multi-Modal Flow Matching

from substantially worse mode estimation accuracy. By incorporating likelihood-based trajectory correction during sampling, FK-steering improves posterior concentration around the high-likelihood regions and produces posterior samples that more closely match the high-density structure of the reference MCMC solution. Method

95% Credible Region IoU ↑

Posterior Mode L2 Distance ↓

FMPE NPE FUSE w/o FK FUSE

0.33 0.29 0.52 0.62

132.41 75.74 6.60 3.85

Table 6. Quantitative comparison on the exoplanet orbital characterization task. We report the intersection-over-union (IoU) between the estimated and reference 95% credible regions together with the Euclidean distance between posterior modes. Higher IoU and lower mode distance indicate better agreement with the long-run PTMCMC reference posterior samples. Best results are highlighted in bold and second-best results are underlined.

H.2. Why FK-Steering Helps on Complex Posteriors The benefit of FK-steering becomes more pronounced in scientific inverse problems with highly multimodal or degenerate posterior structures. In such settings, standard amortized SBI models may exhibit mass-covering behavior, producing overly broad posterior approximations that fail to accurately capture narrow high-likelihood regions. FK-steering mitigates this issue by incorporating intermediate likelihood evaluations during the sampling process. By periodically reweighting and correcting the generative trajectories, the method suppresses physically implausible samples and encourages exploration toward regions with higher posterior density. For relatively simple SBIBM tasks, the base flow-matching model already provides sufficiently accurate posterior estimation, and therefore FK-steering yields only marginal improvements. In contrast, for complex real-world tasks such as exoplanet orbital characterization, FK-steering substantially improves posterior concentration and mode accuracy, as demonstrated in Table 6. Tail coverage remains an empirical property of the finite-particle sampler rather than a formal guarantee.

I. Theoretical Perspective on FK-Steered Inference This section provides a theoretical interpretation of the FK-steered sampler used in FUSE. The goal is not to claim that the finite-particle implementation is an asymptotically exact replacement for MCMC. Rather, FK-steering can be viewed as a tractable posterior-tilting correction applied to the learned amortized proposal. This perspective clarifies three design choices: the target path measure, the use of a denoised proxy in simulator-based scoring, and the stochastic rejuvenation used to mitigate particle degeneracy. I.1. Posterior-Tilted Path Measure For a fixed observation x, let Qx denote the path measure induced by the learned reverse-time sampler conditioned on x. A sample path is written as θ1:0 = {θt : t ∈ [1, 0]}, where θ1 is initialized from the base distribution and θ0 is the final generated parameter. Let q0x (θ) be the terminal marginal density of Qx , and assume q0x (θ) > 0 whenever γx (θ) > 0. If this density were available, an ideal posterior-corrected path measure P⋆x could be defined through the Radon–Nikodym derivative dP⋆x γx (θ0 ) (θ1:0 ) = , γx (θ) = p(x | θ)p(θ), (34) dQx Zx q0x (θ0 ) where Zx is the normalizing constant. For any measurable set A, the terminal marginal then satisfies Z Z γx (θ) ⋆ Px (θ0 ∈ A) = dθ = p(θ | x) dθ. A Zx A

(35)

Thus, the ideal change of measure has the exact posterior as its terminal marginal. This construction is a reference density-ratio correction, not the finite-particle algorithm used by FUSE. In practice, however, the terminal proposal density q0x of an implicit flow sampler is unavailable. FUSE therefore uses the unnormalized log posterior as a tractable reward, producing a posterior-tilted refinement of the learned proposal rather than 21

FUSE: FK-Steered Multi-Modal Flow Matching

an exact density-ratio correction. At resampling times {tr1 , . . . , trJ }, the incremental potential is implemented as   λ Grj (θtrj ) = exp rϕ (θtrj , trj ; x) , J

(36)

where rϕ is defined in Eq. 11. The resulting particle system is a bootstrap approximation of this Feynman–Kac flow: PB particles are propagated by the learned sampler and resampled according to the normalized scores Gkrj / ℓ=1 Gℓrj . Because the implementation uses the learned reverse transition as the proposal, the transition correction in the general FK importance weight is one; if a different proposal were used, the corresponding transition-density ratio would have to be included. Even with infinitely many particles, these chosen potentials target the associated surrogate FK-tilted path measure rather than the ideal density-ratio-corrected posterior path measure above; finite B adds the usual SMC approximation error. I.2. Why Scoring the Denoised Proxy is Principled A key design choice is to evaluate the simulator likelihood at the denoised proxy θ̂t , rather than at the noisy intermediate state θt . This is important because θt is a mixture of clean parameters and noise and may not lie on the physical parameter manifold. Feeding such noisy states into a scientific simulator can violate parameter constraints and produce meaningless likelihood values. Under the rectified-flow interpolation used in training, θt = (1 − t)θ0 + tϵ, with velocity target ϵ − θ0 . In the idealized case where vϕ equals the optimal conditional velocity, rearranging the interpolation gives E[θ0 | θt , x] = θt − t v ⋆ (θt , t, x),

(37)

which motivates the plug-in proxy θ̂t = θt − t vϕ (θt , t, x) used in Eq. 11. The ideal intermediate FK lookahead potential would be Ut (θt ) = log EQx [γx (θ0 ) | θt ] , (38) which is generally intractable because it integrates over all terminal states reachable from θt . A local Taylor expansion around the denoised conditional mean yields the plug-in approximation Ut (θt ) ≈ log p(x | θ̂t ) + log p(θ̂t ).

(39)

This derivation explains why the simulator-based reward in FUSE is evaluated at θ̂t . The approximation is most accurate near the clean data manifold and should be interpreted as a tractable likelihood-guided correction rather than an exact posterior sampler. I.3. SDE Rejuvenation and Under-Dispersion Resampling concentrates computation on high-reward trajectories, but repeated resampling can reduce particle diversity. FUSE mitigates this by using a stochastic sampler after resampling. The role of the SDE perturbation is to restore local exploration while preserving the learned probability path as much as possible. p The stochastic term in Eq. 10 is controlled by σt = α t/(1 − t). In practice the sampler avoids the endpoints and clips the denominator for numerical stability. Larger α increases exploration but may weaken likelihood concentration; smaller α makes the sampler closer to deterministic FK selection and may increase under-dispersion. Thus, particle count, resampling frequency, and noise scale jointly control the fidelity-diversity trade-off. Empirically, this trajectory-level correction is more effective than deterministic post-hoc selection on the β Pictoris b task, as shown in Appendices H.1 and F. In summary, FK-steering should be viewed as a likelihood-guided correction to an amortized posterior sampler. Its ideal path-measure formulation connects to posterior tilting, while the implemented algorithm uses a denoised plug-in reward and finite-particle SMC approximation. This provides a practical mechanism for improving high-density posterior fidelity without sacrificing the speed advantage of amortized inference.

22