Pawsterior: Variational Flow Matching for Structured Simulation-Based Inference

TL;DR

Pawsterior introduces a variational flow-matching framework that effectively handles structured, constrained, and discrete latent variables in simulation-based inference, improving stability and fidelity.

Contribution

It generalizes geometric constraints into the inference process and enables flow-matching for discrete latent structures, broadening SBI applicability.

Findings

Improved classifier two-sample test performance on SBI benchmarks.

Enhanced posterior fidelity with geometric confinement.

Successfully models discrete latent structures in SBI.

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 provides a principled way to apply flow-matching in a broader range of structured SBI settings.