Mixed neural posterior estimation for simulators with discrete and continuous parameters

Technical Summary: Mixed Neural Posterior Estimation for Simulators with Discrete and Continuous Parameters

Problem Statement
Neural Posterior Estimation (NPE) has become a standard tool for Simulation-Based Inference (SBI), enabling rapid, amortized parameter inference for complex simulators with intractable likelihoods. However, standard NPE methods assume continuous parameter spaces, typically modeling the posterior as a continuous probability density using normalizing flows or diffusion models. This assumption fails for a broad class of scientific simulators where the parameter space is mixed, containing both continuous dimensions (e.g., rates, conductances) and discrete dimensions (e.g., categorical variables, model selection indices, change-point locations).

Existing approaches to mixed inference face significant limitations. Standard Markov Chain Monte Carlo (MCMC) methods require careful tuning, do not amortize inference costs across observations, and often struggle with discrete variables (e.g., requiring analytical marginalization or specialized discontinuous Hamiltonian dynamics). Recent diffusion-based SBI methods attempt to embed discrete parameters into continuous spaces or treat them as model indices, but these can be computationally heavy or require complex Riemannian embeddings. Consequently, simulators with mixed parameter spaces remain insufficiently supported by current SBI toolboxes.

Methodology: Mixed Neural Posterior Estimation (MNPE)
The authors propose Mixed Neural Posterior Estimation (MNPE), an extension of NPE that jointly handles discrete and continuous parameters through a factorized inference network.

1. Factorization Strategy:
   The joint posterior $p(\theta_d, \theta_c | x)$ over discrete parameters $\theta_d$ and continuous parameters $\theta_c$ is factorized as:
   $$p(\theta_d, \theta_c | x) = p(\theta_d | x) \cdot p(\theta_c | \theta_d, x)$$
   This separates the problem into two sub-tasks: classification over the discrete space and conditional density estimation in the continuous space, conditioned on the discrete state.

2. Network Architecture:

   • Discrete Component: The term $p(\theta_d | x)$ is modeled using a Masked Autoregressive Density Estimator (MADE). This autoregressive network outputs class probabilities for each discrete dimension, conditioned on preceding discrete dimensions and the observation $x$. This allows handling arbitrary numbers of discrete parameters with varying class counts.
   • Continuous Component: The term $p(\theta_c | \theta_d, x)$ is modeled using a standard conditional generative model (e.g., a neural spline flow or diffusion model) conditioned on both the observation $x$ and the discrete parameters $\theta_d$.
   • Training Objective: The network is trained jointly to minimize the negative log-probability of the mixed parameters:
     $$\mathcal{L} = -\log q(\theta_d, \theta_c | x) = -\log q(\theta_c | \theta_d, x) - \log q(\theta_d | x)$$
     The continuous part uses the standard flow likelihood, while the discrete part reduces to a cross-entropy loss.

3. Calibration Diagnostics:
   To assess the quality of mixed posteriors without reference solutions, the authors propose a combined diagnostic framework:

   • Continuous Dimensions: Uses Simulation-Based Calibration (SBC) via rank statistics to check for uniformity in the marginal rank distributions.
   • Discrete Dimensions: Uses Expected Calibration Error (ECE) and reliability diagrams, treating the discrete inference as a classification problem.
   • Baselines: The authors derive empirical finite-sample baselines for both metrics to distinguish actual miscalibration from statistical noise, acknowledging that perfect calibration is unattainable with finite samples.

Key Results
The authors validate MNPE on three simulators of increasing complexity:

1. Tractable Gaussian Toy Example:

   • A simulator with one continuous and one discrete latent variable coupled through a shared observation.
   • MNPE accurately recovers the analytical ground-truth posterior, including bimodal regimes where the discrete state is uncertain.
   • The Classifier Two-Sample Test (C2ST) score converges to chance level (~0.5) with approximately 1,000 simulations, indicating indistinguishability from the true posterior.

2. Tandem Queueing Simulator:

   • A simulator with three continuous rates and two discrete server counts ($|D|=25$).
   • MNPE is compared against MCMC references (PyMC and NumPyro) where discrete dimensions are analytically marginalized.
   • With 100,000 simulations, MNPE posteriors match the MCMC reference in both continuous and discrete dimensions.
   • Calibration diagnostics (SBC rank histograms and ECE reliability diagrams) confirm the posterior is well-calibrated, with continuous dimensions approaching the finite-sample baseline and discrete dimensions matching the expected error of a perfectly calibrated classifier.

3. Hodgkin–Huxley Biophysical Simulator:

   • A "black-box" simulator for neural activity with intractable likelihoods, involving continuous conductances and discrete channel identities (e.g., presence of specific ion channels).
   • MNPE successfully infers joint posteriors for continuous conductances and categorical channel types.
   • The method produces well-calibrated posteriors across training budgets, enabling the analysis of compensatory mechanisms between ion channels (e.g., the relationship between L-type calcium channels and potassium conductance).

Significance and Claims
The paper claims that MNPE fills a critical gap in SBI by providing a unified, amortized framework for simulators with mixed parameter spaces. Its primary contributions are:

• Generalizability: It supports arbitrary numbers of discrete and continuous parameters and is compatible with any continuous generative network (flows, diffusion).
• Efficiency: Unlike MCMC, MNPE amortizes the cost of simulation and training, allowing for rapid inference on new observations via a single forward pass.
• Calibration: It provides a principled method for assessing calibration in mixed settings, combining rank-based and classification-based metrics with empirical baselines.
• Implementation: The method is implemented in the publicly available sbi Python package, making it accessible to domain scientists without requiring extensive tuning.

The authors note limitations inherent to the NPE family, such as the need for large training datasets for complex posteriors and susceptibility to model misspecification. They also clarify that their calibration procedure assesses marginal calibration per dimension rather than full joint calibration, which remains challenging for mixed spaces.