Pawsterior: Support-Aware Flow Matching for SBI

• Pawsterior is a variational flow-matching framework that adapts simulation-based inference to structured domains such as bounded intervals, simplices, and mixed discrete–continuous spaces.
• It replaces direct velocity regression with an endpoint-based variational inference approach, enforcing affine geometric confinement to respect support constraints.
• The framework effectively manages mixed discrete–continuous latent variables, demonstrating improved posterior fidelity and reduced misallocation of probability mass compared to standard FMPE.

Pawsterior is a variational flow-matching framework for simulation-based inference (SBI) that is designed for posterior distributions supported on structured domains, including bounded intervals, simplices, and mixed discrete–continuous spaces (Carrasco-Pollo et al., 14 Feb 2026). It was introduced to address a mismatch in standard flow-matching approaches, particularly Flow Matching for Posterior Estimation (FMPE), which learn an unconstrained velocity field over $\mathbb{R}^D$ even when the posterior support is a strict subset of that space (Carrasco-Pollo et al., 14 Feb 2026). In structured SBI settings, this mismatch leads to capacity wasted on infeasible directions, probability mass traversing regions where the simulator is undefined, and failure to represent discrete or hybrid latent variables. Pawsterior resolves these issues by replacing direct velocity regression with endpoint-based variational inference, enforcing endpoint-induced affine geometric confinement, and parameterizing joint endpoint distributions over mixed supports within a continuous-time flow-matching framework (Carrasco-Pollo et al., 14 Feb 2026).

1. Problem setting and motivation

In amortized neural posterior estimation, one typically has a prior $p(\theta)$ over simulator parameters $\theta\in\Omega\subseteq\mathbb{R}^D$, a simulator $x\sim p(x\mid\theta)$ generating observations $x$, and the goal of approximating $p(\theta\mid x)$ amortized across many $x$ (Carrasco-Pollo et al., 14 Feb 2026). Pawsterior is motivated by the observation that many posterior distributions of interest in SBI live on structured domains rather than unconstrained Euclidean space.

The framework targets settings in which $\Omega$ is a strict subset of $\mathbb{R}^D$, such as $\Omega=[a,b]^D$, a probability simplex, or a Cartesian product of one-hot regimes (Carrasco-Pollo et al., 14 Feb 2026). Under the standard FMPE parameterization, $p(\theta)$0 is treated as unconstrained in $p(\theta)$1, so learned dynamics can wander off $p(\theta)$2, violate physical constraints, or push mass into invalid regions of parameter space. In purely discrete settings, learning a continuous vector field is described as fundamentally incompatible with the geometry of the support (Carrasco-Pollo et al., 14 Feb 2026).

This framing places Pawsterior within a broader effort to adapt continuous-time generative modeling to inference problems with explicit structural constraints. A plausible implication is that its main significance lies less in replacing flow matching wholesale than in extending it to posterior families that standard unconstrained formulations cannot represent faithfully.

2. Relation to standard FMPE

Flow Matching for Posterior Estimation inserts amortized posterior estimation into a continuous-time generative model by sampling a “noise” endpoint $p(\theta)$3 such as $p(\theta)$4, a “data” endpoint $p(\theta)$5, defining the straight-line interpolation $p(\theta)$6, and learning a velocity field $p(\theta)$7 via the objective (Carrasco-Pollo et al., 14 Feb 2026)

$p(\theta)$8

The critique developed by Pawsterior is that this objective is intrinsically unconstrained. When the support $p(\theta)$9 is structured, unconstrained velocity regression can allocate model capacity to invalid directions and can produce trajectories that are inconsistent with the support of the target posterior (Carrasco-Pollo et al., 14 Feb 2026). The paper explicitly identifies three consequences: capacity wasted on infeasible directions, probability mass traversing regions where the simulator is undefined, and failure to represent discrete or hybrid latent variables.

The framework is described as generalizing the geometric inductive bias of CatFlow (Carrasco-Pollo et al., 14 Feb 2026). This suggests that Pawsterior should be understood not merely as a numerical modification of FMPE, but as a reformulation in which support geometry becomes part of the inference parameterization itself.

3. Endpoint-induced affine geometric confinement

A central concept in Pawsterior is endpoint-induced affine geometric confinement (Carrasco-Pollo et al., 14 Feb 2026). Let $\theta\in\Omega\subseteq\mathbb{R}^D$0 be a convex feasible set, for example a product of bounded intervals or a simplex. Consider the affine interpolation

$\theta\in\Omega\subseteq\mathbb{R}^D$1

with smooth scalars $\theta\in\Omega\subseteq\mathbb{R}^D$2 satisfying $\theta\in\Omega\subseteq\mathbb{R}^D$3 and $\theta\in\Omega\subseteq\mathbb{R}^D$4 (Carrasco-Pollo et al., 14 Feb 2026). The instantaneous velocity is

$\theta\in\Omega\subseteq\mathbb{R}^D$5

Using the conditional expectation in the probability-flow ODE,

$\theta\in\Omega\subseteq\mathbb{R}^D$6

and the identity $\theta\in\Omega\subseteq\mathbb{R}^D$7, one obtains

$\theta\in\Omega\subseteq\mathbb{R}^D$8

Because the posterior of $\theta\in\Omega\subseteq\mathbb{R}^D$9 given $x\sim p(x\mid\theta)$0 is supported on $x\sim p(x\mid\theta)$1, $x\sim p(x\mid\theta)$2, and therefore

$x\sim p(x\mid\theta)$3

This confinement statement is the geometric core of Pawsterior (Carrasco-Pollo et al., 14 Feb 2026). It formalizes how feasible support constraints can be preserved at the level of the induced flow. The paper also notes that, in finite-sample learning, unless the model explicitly enforces $x\sim p(x\mid\theta)$4, the learned field may break this confinement. A common misconception would therefore be to treat confinement as an automatic consequence of interpolation alone; in the Pawsterior formulation, it depends on explicit parameterization of the endpoint posterior so that the conditional mean lies in $x\sim p(x\mid\theta)$5.

4. Two-sided variational formulation

Instead of directly regressing $x\sim p(x\mid\theta)$6, Pawsterior learns a two-sided variational approximation to the joint endpoint posterior conditioned on the intermediate state and observation (Carrasco-Pollo et al., 14 Feb 2026):

$x\sim p(x\mid\theta)$7

Under the mean-field factorization

$x\sim p(x\mid\theta)$8

the training objective becomes

$x\sim p(x\mid\theta)$9

The predicted conditional means are

$x$0

The velocity is then recovered as

$x$1

The paper emphasizes that this velocity never requires dividing by $x$2 or $x$3 (Carrasco-Pollo et al., 14 Feb 2026). Because $x$4 by construction, through appropriate parameterization of $x$5, the induced velocity remains confined to $x$6. The claimed consequences are more stable sampling and higher-fidelity posterior approximations that respect physical or probabilistic bounds (Carrasco-Pollo et al., 14 Feb 2026).

The two-sided formulation is therefore both a statistical and geometric modification. Statistically, it replaces direct velocity regression by variational estimation of endpoint conditionals. Geometrically, it ties the learned flow to the feasible support through the conditional mean of the terminal endpoint.

5. Parameterization for mixed discrete–continuous structure

Pawsterior is designed to enable SBI tasks involving discrete latent structure, including switching systems, that are described as fundamentally incompatible with conventional flow-matching approaches (Carrasco-Pollo et al., 14 Feb 2026). To handle categorical or switching-regime models such as a Switching Gaussian Mixture (SGM), each discrete parameter $x$7 is represented as a one-hot vector in $x$8.

Within the variational factors, continuous endpoints $x$9 or continuous coordinates of $p(\theta\mid x)$0 are modeled with Gaussian likelihoods, yielding MSE losses, while discrete endpoints $p(\theta\mid x)$1 are modeled via logits $p(\theta\mid x)$2 defining

$p(\theta\mid x)$3

For such coordinates, the training loss is the cross-entropy against the true one-hot target (Carrasco-Pollo et al., 14 Feb 2026).

At sampling time, the ODE for $p(\theta\mid x)$4 is integrated from $p(\theta\mid x)$5 to $p(\theta\mid x)$6 via an Euler solver. At each step, the network produces logits for the discrete endpoint posterior, and these logits are converted to a point estimate of the conditional mean on the simplex via a softmax, ensuring $p(\theta\mid x)$7 lives in the probability simplex (Carrasco-Pollo et al., 14 Feb 2026). After integration to $p(\theta\mid x)$8, the continuous coordinates are projected normally, while discrete coordinates are mapped to hard one-hot draws with $p(\theta\mid x)$9.

The following table summarizes the support-dependent parameterization described in the paper.

Coordinate type	Variational form	Training signal
Continuous endpoints or continuous coordinates	Gaussian likelihoods	MSE losses
Discrete endpoints $x$0 (one-hot)	Logits $x$1 with softmax	Cross-entropy against the true one-hot

This parameterization is central to Pawsterior’s claim of extending flow matching to mixed-support inference. A plausible implication is that the method’s expressivity derives less from a new transport law than from a support-aware representation of endpoint uncertainty.

6. Empirical evaluation, limitations, and scope

The reported empirical evaluation covers two settings (Carrasco-Pollo et al., 14 Feb 2026). On sbibm continuous-parameter benchmarks, comprising 10 tasks and training regimes of $x$2, $x$3, and $x$4 simulations, performance is evaluated via the classifier two-sample test (C2ST), where $x$5 is indistinguishable from the true posterior and higher is worse. On bounded-support tasks, Pawsterior reduces C2ST substantially, often from $x$6, and it also yields modest gains on unbounded tasks.

The paper also reports results on a synthetic Switching Gaussian Mixture task with $x$7 regimes, $x$8 transitions, and continuous state dimension $x$9 (Carrasco-Pollo et al., 14 Feb 2026). In that setting, FMPE fails, with C2ST$\Omega$0 even at $\Omega$1 simulations, while Pawsterior attains C2ST$\Omega$2 with comparable compute.

The sbibm average C2ST values reported in the paper are as follows (Carrasco-Pollo et al., 14 Feb 2026).

Task category	$\Omega$3	$\Omega$4	$\Omega$5
bounded-support	$\Omega$6 (FMPE) / $\Omega$7 (Paw)	$\Omega$8	$\Omega$9
unbounded	$\mathbb{R}^D$0	$\mathbb{R}^D$1	$\mathbb{R}^D$2

These results are presented as evidence for two distinct claims: improved posterior fidelity in structured continuous domains and practical viability in mixed discrete–continuous settings that standard FMPE cannot capture (Carrasco-Pollo et al., 14 Feb 2026). At the same time, the method has explicit limitations. Pawsterior assumes convex supports $\mathbb{R}^D$3; non-convex or manifold supports, including a torus or Stiefel manifold, require further geometric extensions. Mapping network outputs to bounded intervals uses smooth squashing such as $\mathbb{R}^D$4, and edge-case numerical issues such as saturation may arise. Sampling cost scales with the number of ODE steps, so adaptive solvers or learned solvers could improve efficiency (Carrasco-Pollo et al., 14 Feb 2026).

Future directions identified in the paper include systematic study of non-convex and Riemannian supports, hybrid continuous–discrete flows beyond the Cartesian-product case, and principled time-prior schedules that adapt to local flow complexity (Carrasco-Pollo et al., 14 Feb 2026). This suggests that Pawsterior is best understood as a support-aware extension of flow matching whose current formulation is strongest on convex structured domains and mixed-support Cartesian products.