Pawsterior: Variational Flow Matching for Structured Simulation-Based Inference

TL;DR

P Pawsterior formalizes endpoint-induced affine geometric confinement, a principle that incorporates domain geometry directly into the inference process via a two-sided variational model, and enables SBI tasks involving discrete latent structure that are fundamentally incompatible with conventional flow-matching approaches.

Abstract

We introduce Pawsterior, a variational flow-matching framework for improved and extended simulation-based inference (SBI). Many SBI problems involve posteriors constrained by structured domains, such as bounded physical parameters or hybrid discrete-continuous variables, yet standard flow-matching methods typically operate in unconstrained spaces. This mismatch leads to inefficient learning and difficulty respecting physical constraints. Our contributions are twofold. First, generalizing the geometric inductive bias of CatFlow, we formalize endpoint-induced affine geometric confinement, a principle that incorporates domain geometry directly into the inference process via a two-sided variational model. This formulation improves numerical stability during sampling and leads to consistently better posterior fidelity, as demonstrated by improved classifier two-sample test performance across standard SBI benchmarks. Second, and more importantly, our variational parameterization enables SBI tasks involving discrete latent structure (e.g., switching systems) that are fundamentally incompatible with conventional flow-matching approaches. By addressing both geometric constraints and discrete latent structure, Pawsterior extends flow-matching to a broader class of structured SBI problems that were previously inaccessible.

Figures (3)

• Figure 1: Toy example illustrating the effect of geometric constraints. Standard FM transports probability mass through infeasible regions of the parameter space, whereas the proposed endpoint-based variational formulation restricts the flow to the feasible set.
• Figure 2: Comparison between FMPE and Pawsterior across sbibm tasks, evaluated using the C2ST metric. Lower values indicate better performance, with $0.5$ corresponding to samples that are indistinguishable from the reference posterior.
• Figure 3: C2ST performance on the SGM task as a function of model depth (number of residual blocks) for hidden dimensions $h=64$ (left) and $h=128$ (right), across data regimes of $10^3$, $10^4$, and $10^5$ simulations. Lower-opacity curves correspond to fewer simulations, while higher-opacity curves indicate larger datasets.