Addressing Misspecification in Simulation-based Inference through Data-driven Calibration

Addressing Misspecification in Simulation-based Inference through Data-driven Calibration

Speakers: Antoine Wehenkel, Juan L. Gamella (first authors), Ozan Sener, Jens Behrmann, Guillermo Sapiro, Jörn Jacobsen, Marco Cuturi (all from Apple)

Conference: ICML 2025

YouTube: https://slideslive.com/39043997

Overview

This talk by Juan L. Gamella and Antoine Wehenkel, representing a collaborative effort from Apple, introduces Warped Posterior Estimator (WORP), a novel framework designed to enhance Simulation-Based Inference (SBI) in scenarios where simulators are inherently misspecified. The core challenge in many scientific and engineering domains is the need to infer complex, often unobservable parameters from readily available observations, a task complicated by a severe scarcity of labeled data. While simulators offer a powerful avenue for incorporating domain knowledge and generating synthetic data, their inherent inaccuracies—the "sim-to-real gap"—typically render standard SBI pipelines unreliable, leading to inaccurate and uncalibrated uncertainty estimates.

WORP directly confronts this problem by leveraging optimal transport to explicitly model and bridge the misspecification gap between simulated and real-world data. The approach is particularly impactful because it allows for robust inference and precise uncertainty quantification even when the underlying simulator is imperfect. The presenters demonstrate that WORP can operate effectively in a zero-shot learning context when the sim-to-real gap is moderate, and can be further extended with a minimal amount of labeled data for severe misspecification. This innovation significantly broadens the applicability of SBI, making it a viable tool in critical, data-constrained applications such as medical diagnostics and complex system control.

The significance of WORP lies in its ability to unlock the potential of domain-specific simulators, which are often expensive to develop but invaluable for encoding expert knowledge. By making SBI robust to simulator imperfections, WORP offers a path to accurate and reliable inference that is highly data-efficient, drastically reducing the need for large, costly labeled datasets. This work paves the way for new applications in fields where traditional supervised learning is impractical due to data scarcity, providing a principled method for extracting meaningful insights and quantifying uncertainty from imperfect models and limited real-world observations.

Background

▶ Watch: Introduction: Inferring parameters with little labeled data (0:00)

Many scientific and engineering problems revolve around an inference task: given some readily available observations, $X_o$, the goal is to infer the value of an underlying, often unobservable, parameter, $\theta$. A recurring challenge is that these parameters are frequently difficult, expensive, or invasive to measure, resulting in a severe scarcity of labeled data (i.e., pairs of $(X_o, \theta)$). For instance, inferring cardiovascular parameters of a patient's heart from non-invasive pulse measurements (e.g., from a smartwatch) presents such a dilemma, as direct measurement of heart parameters is typically invasive and costly.

In the absence of abundant labeled data, a common strategy in machine learning and systems design is to incorporate domain knowledge in the form of a simulator. A simulator is a computer program that, given the parameters $\theta$, generates a simulation of the expected observations, $X_s$. This provides a mechanism to generate an arbitrary amount of synthetic data $(X_s, \theta)$ pairs. However, the simulator operates in the "forward" direction ($\theta \to X_s$), whereas the inference problem requires "inverting" this process ($X_o \to \theta$). Furthermore, simulators are often complex, black-box programs.

To invert such simulators and perform inference, the field of Simulation-Based Inference (SBI) has emerged. SBI encompasses a family of techniques designed to estimate the posterior distribution $p(\theta|X_o)$ when the likelihood function $p(X_o|\theta)$ is intractable but sampling from it (via simulation) is possible. A common approach within SBI involves a neural network pipeline. Specifically, simulated measurements $X_s$ are first passed through an encoder-type network (also referred to as a neural statistic estimator) to obtain a low-dimensional embedding. This embedding is then fed into a neural density estimator, which is trained to approximate the posterior distribution $p(\theta|X_s)$. The entire pipeline is trained end-to-end by generating numerous pairs of parameters $\theta$ (sampled from a prior distribution) and their corresponding simulated measurements $X_s$, using maximum likelihood estimation (MLE). This approach has proven effective and robust under ideal conditions.

The critical problem arises when this SBI pipeline is applied to real-world observations, $X_o$. Simulators are, by their nature, approximations of reality; "simulation and reality are never exactly the same." This discrepancy is known as the sim-to-real gap or model misspecification. When real observations $X_o$ are fed into a pipeline trained solely on simulated data, both the encoder and, crucially, the neural density estimator operate in an out-of-distribution (OOD) regime. The neural density estimator is particularly sensitive to OOD inputs. This misspecification causes the inference pipeline to break down, leading to highly inaccurate and often collapsed posterior estimates that bear little resemblance to the true underlying parameters. The challenge, therefore, is to develop an SBI method that is robust to this inevitable sim-to-real gap, allowing for reliable inference and accurate uncertainty quantification from imperfect simulators.

Key Findings

▶ Watch: SBI's limitation: Failure with misspecified simulators (3:10)

The central contribution of this work is the introduction of the Warped Posterior Estimator (WORP), a robust framework for simulation-based inference that effectively addresses the problem of simulator misspecification. The key findings and contributions can be summarized as follows:

1. Robustness through Optimal Transport: WORP's foundational idea is to explicitly model the sim-to-real gap using optimal transport. This allows the inference pipeline to avoid operating the sensitive neural density estimator in an out-of-distribution regime. By matching unlabeled real-world data to simulated data, WORP effectively "warps" the real-world observations into a space where the pre-trained density estimator remains valid.

1. Key Modeling Assumption: A critical underlying assumption for WORP is that the misspecification (the sim-to-real gap) is not a direct function of the parameter of interest ($\theta$), but rather depends on $\theta$ only through the simulation ($X_s$). This allows for a clean separation of the simulator's forward model from the data-driven correction for misspecification.

1. Zero-Shot Inference Capability: In scenarios where the sim-to-real gap is moderate, WORP demonstrates the ability to perform zero-shot inference. This means it can produce informative and well-calibrated posterior distributions for real-world observations without requiring any labeled real-world data. It achieves this by leveraging readily available unlabeled real-world data and simulated data generated from the prior.

1. Addressing Extreme Misspecification with Minimal Labeled Data: For situations involving a much larger and more complex sim-to-real gap—where the initial neural statistic estimator (encoder) fails to generalize effectively—WORP introduces a calibration step. This step involves using a very small amount of labeled validation data to fine-tune the neural statistic estimator. The fine-tuning objective is to ensure consistency between the embeddings of real-world examples and their corresponding simulated counterparts, effectively realigning the embedding space.

1. Significant Data Efficiency: The calibrated version of WORP exhibits remarkable data efficiency. In a real-world experimental setup (the "causal chamber"), WORP achieved well-calibrated and informative posterior distributions using only 50 labeled samples for calibration. Crucially, a directly trained supervised model required 20 times more samples (1000 labeled samples) to achieve comparable performance, highlighting WORP's ability to leverage domain knowledge from the simulator to overcome data scarcity.

1. Uncertainty Quantification: Unlike traditional point estimation methods, WORP provides full posterior distributions, enabling comprehensive uncertainty quantification. The experimental results consistently show that WORP's posteriors are informative about the true parameter values and are well-calibrated, encompassing the true values within their credible intervals.

In essence, WORP provides a robust and data-efficient solution to a pervasive problem in scientific machine learning, making complex simulators actionable for reliable inference and uncertainty estimation even in the presence of significant model misspecification and limited labeled data.

Technical Deep Dive

▶ Watch: WORP's core idea: Modeling misspecification with optimal transport (4:15)

The technical foundation of WORP builds upon and significantly extends standard Simulation-Based Inference (SBI) to account for simulator misspecification.

Standard SBI Pipeline Revisited

A typical SBI pipeline, as described, begins with a simulator that maps parameters $\theta$ to simulated observations $X_s$. The goal is to learn the inverse mapping, $p(\theta|X_s)$. This is achieved using a neural network architecture:

1. Neural Statistic Estimator (Encoder): An encoder network, $f_\phi$, maps high-dimensional simulated observations $X_s$ to a lower-dimensional, informative embedding $z_s = f_\phi(X_s)$.
2. Neural Density Estimator: A second network, $g_\psi$, takes this embedding $z_s$ and estimates the conditional posterior distribution $p(\theta|z_s)$. This is often implemented using normalizing flows or mixture density networks.

The entire pipeline, parameterized by $\phi$ and $\psi$, is trained end-to-end. Training data consists of pairs $(\theta_i, X_{s,i})$, where $\theta_i$ is sampled from a prior distribution $p(\theta)$, and $X_{s,i}$ is generated by running the simulator with $\theta_i$. The loss function typically involves maximizing the likelihood of the parameters given the simulated observations, $\sum_i \log p(\theta_i | f_\phi(X_{s,i}))$. This process ensures that $g_\psi$ learns to provide accurate posteriors when fed in-distribution simulated embeddings.

The fundamental breakdown occurs when real observations $X_o$ are introduced. Due to the sim-to-real gap, $X_o$ is often out-of-distribution (OOD) for $f_\phi$. Even if $f_\phi$ somehow produces a meaningful embedding $z_o = f_\phi(X_o)$, this $z_o$ is likely OOD for $g_\psi$, which was trained exclusively on $z_s$ embeddings. This OOD input to the sensitive neural density estimator $g_\psi$ leads to unreliable and often degenerate posterior estimates.

WORP: Modeling Misspecification with Optimal Transport

WORP's core innovation is to explicitly model the misspecification to ensure the neural density estimator always operates on in-distribution data. This is achieved through optimal transport (OT).

1. Modeling Assumption: WORP makes a crucial assumption about the nature of misspecification: the discrepancy between real and simulated data is assumed to depend on the parameters $\theta$ only through the simulation $X_s$. Formally, this can be expressed as $p(X_o|\theta) = \int p(X_o|X_s) p(X_s|\theta) dX_s$. This allows for a two-stage modeling: first, the simulator's forward model $p(X_s|\theta)$, and second, a data-driven mapping $p(X_o|X_s)$ that captures the sim-to-real gap.

1. Optimal Transport for Matching Distributions: The objective is to find a mapping between the distribution of real-world observations $p(X_o)$ and the distribution of simulated observations $p(X_s)$. Since labeled real-world data (pairs of $(X_o, \theta)$) are scarce, WORP leverages abundant unlabeled real-world observations and a large dataset of simulated observations (generated from the simulator by sampling $\theta$ from the prior). The goal is to find, for each real-world example, a corresponding simulated example that is "closest" in some meaningful sense.

1. Cost Function in Embedding Space: A key design decision in optimal transport is the choice of a cost function that quantifies the dissimilarity between data points. Defining this cost in the raw data space (e.g., pixel space for images) can be challenging due to the high dimensionality and complex structure of the data. WORP proposes defining the cost function in the low-dimensional embedding space produced by the neural statistic estimator $f_\phi$. The rationale is that this embedding space is designed to capture the essential information for inference while discarding irrelevant artifacts, and is thus more likely to generalize to real-world data effectively. The cost is typically the Euclidean distance between embeddings: $c(z_o, z_s) = \|f_\phi(X_o) - f_\phi(X_s)\|_2^2$.

1. Optimal Transport Formulation: Given a set of unlabeled real-world embeddings $\{z_{o,i}\}_{i=1}^N$ and a set of simulated embeddings $\{z_{s,j}\}_{j=1}^M$, optimal transport seeks a coupling (or transport plan) $\pi$ that minimizes the total cost of moving "mass" from the real distribution to the simulated distribution. In practice, WORP employs a semi-balanced and entropy-regularized optimal transport formulation, solved efficiently using the Sinkhorn algorithm. This produces a soft-max coupling, meaning each real-world embedding $z_o$ is mapped not to a single $z_s$, but to a distribution over multiple $z_s$ with associated weights. This softness is crucial for robustly handling ambiguities and noisy data.

1. Posterior Estimation with WORP: Once the optimal coupling is found, inference for a specific real observation $X_o$ proceeds as follows:

• Compute its embedding $z_o = f_\phi(X_o)$.
• Use the optimal coupling to determine the weights $w_j$ that connect $z_o$ to each of the simulated embeddings $z_{s,j}$.
• Since the neural density estimator $g_\psi$ is trained and valid for simulated embeddings, we can compute the posterior $p(\theta|z_{s,j})$ for each relevant $z_{s,j}$.
• The final posterior for the real observation $X_o$ is then constructed as a mixture of these simulated posteriors, weighted by the coupling: $p(\theta|X_o) = \sum_j w_j p(\theta|z_{s,j})$. This ensures that the neural density estimator $g_\psi$ is always queried with in-distribution embeddings, resolving the OOD problem.

WORP with Calibration for Extreme Misspecification

While the zero-shot WORP works well for moderate sim-to-real gaps, severe misspecification can cause the initial neural statistic estimator $f_\phi$ to fail entirely in generalizing to real-world data. In such cases, $f_\phi(X_o)$ might be completely meaningless, leading the optimal transport to find a mostly uniform coupling, and consequently, an uninformative posterior that reverts to the prior.

To address this, WORP introduces a calibration step that leverages a small amount of labeled validation data ($X_o, \theta_{true}$ pairs).

1. Consistency Objective: The goal of calibration is to fine-tune the neural statistic estimator $f_\phi$ such that it becomes "consistent" across real and simulated domains. Specifically, for a given labeled real observation $X_o$ with true parameter $\theta_{true}$, we can generate its corresponding simulated observation $X_s$ using the simulator and $\theta_{true}$. The calibration loss then minimizes the Euclidean distance between the embedding of the real observation and the embedding of its corresponding simulated observation: $L_{calib} = \|f_\phi(X_o) - f_\phi(X_s(\theta_{true}))\|_2^2$.

1. Fine-tuning: This loss is used to fine-tune the parameters $\phi$ of the neural statistic estimator. This process "realigns" the real-world embeddings with the space where the neural density estimator was trained, ensuring that $f_\phi(X_o)$ produces meaningful, in-distribution embeddings for the subsequent OT and posterior estimation steps.

After this fine-tuning, the full WORP pipeline (optimal transport and mixture posterior) can be applied, now operating effectively even with large sim-to-real gaps. This elegant solution allows WORP to robustly handle a wide spectrum of misspecification challenges.

Experimental Setup & Results

▶ Watch: Designing optimal transport cost function in embedding space (5:00)

The effectiveness of WORP was demonstrated across two distinct experimental settings, showcasing its capabilities under varying degrees of simulator misspecification.

Cardiovascular Parameter Inference (Synthetic Misspecification)

Setup:

• Problem: Inferring cardiovascular parameters of a patient's heart from pulse measurements. This is a challenging problem due to the high cost and invasiveness of direct parameter measurement, leading to very little labeled data.
• Simulator: A physical simulator modeling blood flow through veins and interaction with a sensor was used to generate synthetic pulse data from given cardiovascular parameters.
• Misspecification: To create a "sim-to-real gap" in a controlled environment, the parameters of the simulator itself were slightly varied when generating the "real" observations compared to the parameters used for training the initial SBI pipeline. This created a moderate, yet significant, discrepancy between the simulated and "real" data.
• Data: No labeled real-world data was used for inference. The experiment relied on unlabeled real-world-like observations and a large pool of simulated data.

Results:

• Standard SBI (without WORP) completely failed. The posterior estimates collapsed to single, inaccurate points, demonstrating extreme sensitivity to even moderate misspecification.
• WORP, operating in a zero-shot learning mode (i.e., without any labeled real-world data), successfully produced informative and well-calibrated posterior distributions. Examples shown in the talk illustrated posteriors that accurately encompassed the ground truth values, providing meaningful uncertainty estimates. This demonstrated WORP's ability to leverage optimal transport to bridge the sim-to-real gap effectively when the misspecification is not overly severe.

Causal Chamber Experiment (Real-world Misspecification)

Setup:

• Hardware: A physical "causal chamber" – an optical experiment machine that takes physical inputs (parameters) and produces image data as output.
• Problem: Inferring the input parameters to the chamber from the captured images.
• Simulator: A simple simulator was developed to emulate the physical process within the chamber. This simulator produced rough approximations of the images, and critically, ignored some of the underlying physics, leading to a much larger and more complex sim-to-real gap compared to the cardiovascular example.
• Ground Truth: Because the inputs to the causal chamber were controlled, true parameter values were available, allowing for direct evaluation of the inference pipeline's accuracy.
• Misspecification Severity: The visual and physical discrepancy between simulated and real-world images from the causal chamber was substantial, representing an "extreme out-of-distribution regime" for the neural statistic estimator.

Results:

• Initial WORP Failure: When applied without calibration, the initial WORP approach (zero-shot) failed. The extreme misspecification caused the neural statistic estimator to not generalize at all to the real-world images. Consequently, the optimal transport algorithm produced a mostly uniform coupling, leading to posterior distributions that were essentially just the prior, offering no informative inference.
• WORP with Calibration Success: To overcome this, a small set of 50 labeled samples from the causal chamber was used for calibrating the neural statistic estimator. This fine-tuning step realigned the real-world embeddings with the simulated embedding space.
• Following calibration, the full WORP pipeline performed remarkably well. It yielded informative and well-calibrated posterior distributions for the four parameters being inferred. The example posteriors clearly encompassed the true parameter values.
• Data Efficiency Benchmark: A crucial result was the comparison to a purely supervised baseline model. This baseline was trained directly on labeled data from the causal chamber without leveraging the simulator. The supervised model, when trained on the same 50 labeled samples, simply returned the prior distribution, indicating that 50 samples were insufficient for it to learn anything meaningful from scratch. The supervised baseline only started to catch up to WORP's performance when provided with 20 times more labeled samples (1000 samples). This stark comparison highlights WORP's exceptional data efficiency by effectively leveraging the domain knowledge embedded in the (imperfect) simulator.

These experiments rigorously demonstrate that WORP, especially with its calibration mechanism, provides a robust and highly data-efficient solution for accurate inference and uncertainty quantification in scenarios characterized by significant simulator misspecification and scarce labeled data.

Practical Implications

▶ Watch: Demonstrating WORP's well-calibrated and informative posteriors (6:50)

The Warped Posterior Estimator (WORP) offers significant practical implications for various stakeholders in the AI/ML and scientific computing communities, particularly in domains where data scarcity and model imperfections are prevalent.

For Practitioners and Model Builders

• Enabling Inference in Data-Scarce Domains: WORP provides a robust framework for performing inference and quantifying uncertainty in fields where obtaining large labeled datasets is prohibitively expensive, time-consuming, or impossible. This includes critical areas like medical diagnostics (e.g., inferring physiological parameters from non-invasive measurements), climate modeling, materials science, robotics, and complex physical systems.
• Leveraging Imperfect Simulators: Model builders often invest heavily in developing sophisticated simulators, which encode valuable domain knowledge. WORP ensures that these simulators remain useful even when they are not perfectly accurate. This reduces the pressure to create "perfect" simulators, allowing for a more pragmatic approach where approximate simulators can still yield high-quality inference.
• Uncertainty Quantification: Beyond point estimates, WORP provides full posterior distributions. This is crucial for applications requiring reliable uncertainty estimates, such as risk assessment, decision-making under uncertainty, and scientific discovery, where understanding the confidence in an estimate is as important as the estimate itself.
• Reduced Development Cycles: By reducing the dependency on extensive labeled data collection, WORP can accelerate the development and deployment cycles of AI/ML models in specialized domains.

For Infrastructure and Deployment Teams

• Efficient Resource Utilization: The data efficiency demonstrated by WORP means that smaller, more manageable labeled datasets can be sufficient for calibration, reducing the computational and human resources required for data annotation and validation.
• Robustness in Production: Deploying models that are robust to real-world data shifts and discrepancies (sim-to-real gap) is a major challenge. WORP offers a principled way to build such robustness directly into the inference pipeline, potentially leading to more reliable and stable deployed systems.
• Scalability with Unlabeled Data: While some labeled data might be needed for calibration, WORP primarily relies on readily available unlabeled real-world data and easily generated simulated data. This makes it scalable in environments where unlabeled data streams are continuous.

Tradeoffs and Limitations

• Modeling Assumption: WORP's effectiveness hinges on the assumption that the misspecification depends on the parameters only through the simulation. While broadly applicable, there might be scenarios where the real-world observation is directly influenced by parameters in ways not captured by the simulator, which could challenge this assumption.
• Availability of Unlabeled Data: WORP still requires a sufficient amount of unlabeled real-world data to effectively perform optimal transport and model the sim-to-real gap.
• Computational Cost of Optimal Transport: While the Sinkhorn algorithm is efficient, computing optimal transport for very large datasets can still be computationally intensive, especially if the embedding space is high-dimensional or the number of real/simulated samples is extremely large.
• Quality of Neural Statistic Estimator: For severe misspecification, the initial neural statistic estimator (encoder) must be amenable to fine-tuning with a small labeled dataset. If its initial generalization is extremely poor, even calibration might struggle. The design of this encoder remains critical.
• "Meaningful" Embedding Space: The choice of the embedding space and the definition of the cost function within it are crucial. While the talk suggests Euclidean distance in the neural statistic estimator's output space, this might not always be optimal for all data types or inference tasks.

Despite these considerations, WORP represents a significant step forward in making simulation-based inference a more practical and reliable tool for real-world applications, especially where data labeling is a bottleneck.

Key Takeaways

• Simulators are valuable but often misspecified: While simulators provide crucial domain knowledge, their inherent inaccuracies (the "sim-to-real gap") cause standard Simulation-Based Inference (SBI) pipelines to fail catastrophically when applied to real-world data.
• WORP addresses misspecification via Optimal Transport: The Warped Posterior Estimator (WORP) explicitly models the sim-to-real gap using optimal transport, mapping real-world observations to the simulated data distribution. This ensures the sensitive neural density estimator always operates in-distribution.
• Zero-shot inference for moderate gaps: For moderate levels of misspecification, WORP can perform robust inference and uncertainty quantification in a zero-shot setting, requiring only unlabeled real-world data and simulated data.
• Calibration for extreme misspecification: For severe sim-to-real gaps, WORP can be effectively calibrated with a minimal amount of labeled data (e.g., 50 samples) to fine-tune the neural statistic estimator, enabling it to generalize and produce informative posteriors.
• Exceptional data efficiency: WORP demonstrates significant data efficiency, outperforming purely supervised baselines that require 20 times more labeled data (e.g., 1000 samples) to achieve comparable results.
• Enables uncertainty quantification in data-scarce domains: By leveraging imperfect simulators and minimal labeled data, WORP makes robust inference and reliable uncertainty quantification possible in critical applications where labeled data is scarce and expensive.

About the Speaker(s)

The talk was presented by Juan L. Gamella and Antoine Wehenkel, who share first authorship on the paper. Both are affiliated with Apple, indicating this work is a product of research conducted within the company. They were joined by several other colleagues from Apple: Ozan Sener, Jens Behrmann, Guillermo Sapiro, Jörn Jacobsen, and Marco Cuturi, highlighting a collaborative effort across multiple researchers. Their work focuses on advancing machine learning techniques, particularly in areas like simulation-based inference and robust model estimation under challenging data conditions.