📊 statistics

Overcoming Selection Bias in Statistical Studies With Amortized Bayesian Inference

This paper introduces a bias-aware simulation-based inference framework that embeds selection mechanisms directly into neural posterior estimation to enable tractable-likelihood-free, amortized Bayesian inference that corrects for selection bias and provides calibrated uncertainty quantification.

Original authors: Jonas Arruda, Sophie Chervet, Paula Staudt, Andreas Wieser, Michael Hoelscher, Isabelle Sermet-Gaudelus, Nadine Binder, Lulla Opatowski, Jan Hasenauer

Published 2026-04-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jonas Arruda, Sophie Chervet, Paula Staudt, Andreas Wieser, Michael Hoelscher, Isabelle Sermet-Gaudelus, Nadine Binder, Lulla Opatowski, Jan Hasenauer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Distorted Mirror"

Imagine you want to know the average height of everyone in a city. But, the only way you can measure people is by standing on a street corner and asking anyone who walks by to step on your scale.

Here's the catch: Only tall people are walking by. Why? Because the short people are all at home watching TV, or maybe the street is too high for them to reach the door.

If you just take the average of the people on the street, you'll think the whole city is giant. This is Selection Bias. The data you have is a "distorted mirror" of reality because the way you collected it (who walked by) was linked to what you were measuring (height).

In the real world, this happens all the time:

Medical studies: Only sick people show up at a clinic, so the disease looks more common than it is.
Surveys: Only people with strong opinions answer the phone, so the results don't represent the quiet majority.

The Old Way: Trying to Fix the Mirror with Math

Traditionally, statisticians try to fix this by doing complex math. They say, "Okay, we know 20% of short people didn't show up, so let's add a 'weight' to the tall people to balance it out."

This works well if the rules of the game are simple and you can write them down in a neat formula. But in modern science, things are messy. Imagine if the reason people didn't show up depended on a secret mix of their height, their mood, the weather, and a hidden genetic trait. The math becomes so incredibly complicated that it's impossible to solve. It's like trying to untangle a knot that has been tied by a thousand different hands.

The New Solution: The "Virtual Reality" Simulator

The authors of this paper propose a brilliant new way to solve this. Instead of trying to untangle the knot with math, they build a Virtual Reality (VR) simulator.

Here is how their method, called Amortized Bayesian Inference, works:

1. Build the "World" and the "Glitch"

First, they create a computer simulation of the entire population (the "World"). Then, they program the specific "Glitch" that causes the bias.

Example: They simulate a city where short people are less likely to walk down the street.
They run this simulation millions of times, creating millions of "fake cities" where they know exactly who is tall, who is short, and who got left out.

2. Train the "AI Detective"

They take these millions of fake, biased datasets and feed them into a powerful AI (a Neural Network). They teach the AI: "Here is a biased street sample. Here is the true average height of the whole city. Learn the pattern of how the bias distorts the truth."

The AI learns to look at a distorted sample and instantly guess what the real population looks like, effectively "un-biasing" the data.

3. The "Amortized" Superpower

This is the cool part. Usually, training an AI takes forever. But once this AI is trained, it becomes amortized.

Analogy: Think of it like learning to drive. You spend months learning (training). But once you have your license, you can drive any car, on any road, instantly, without re-learning how to drive every time you get behind the wheel.
The AI can now look at real data from a study and instantly give a corrected answer, even if the study design is weird or complex.

How Do We Know It's Not Hallucinating?

You might ask, "How do we trust the AI? What if it's just guessing?"

The authors built a "Lie Detector" test into the system.

They use a Classifier (another AI) to play a game of "Spot the Fake."
The Classifier is shown two groups of data: one from the AI's corrected guesses and one from the original simulation.
If the AI did its job perfectly, the Classifier shouldn't be able to tell them apart (it should guess 50/50).
If the Classifier can easily spot the difference, it means the AI is still biased, and the system flags it.

Real-World Examples They Tested

The team tested this on three tricky real-life scenarios:

The "Missing" Virus Test (KoCo19 Study):
In a study about COVID antibodies, many people stopped showing up for tests as the study went on (missing data). The old math methods got confused and gave wrong answers. The AI simulator, however, learned that "people dropping out" was part of the pattern and corrected the prevalence rates accurately.
The "Death" Trap (Framingham Heart Study):
They studied dementia. But if a person died before they could be diagnosed with dementia, their data was lost. This creates a huge bias because the sickest people disappear from the record. The AI learned to simulate the "death" process and realized, "Ah, the people who died were probably the ones who would have gotten dementia," and corrected the estimates.
The "Child-First" Rule (PedCovid Study):
A study only included households if a child tested positive first. This meant they missed households where an adult got sick first. The math to fix this is a nightmare. The AI simulator simply learned the rule "Only look for child-first households" and instantly figured out how to estimate the virus spread for the whole population.

The Bottom Line

This paper gives scientists a new tool to stop worrying about "perfect" data. In the real world, data is always messy, biased, and incomplete.

Instead of trying to force the messy data into a perfect mathematical box, this method says: "Let's build a virtual world that mimics the messiness, train an AI to understand the mess, and then let the AI clean up the real data for us."

It turns the problem of "fixing bad data" into a game of "simulating reality," making it possible to get accurate answers even when the study design is flawed.

1. Problem Statement

Selection bias occurs when the probability of an observation entering a dataset depends on variables related to the quantities of interest (e.g., outcomes, covariates, or latent states). This leads to systematic distortions in parameter estimation and uncertainty quantification.

Limitations of Classical Methods: Traditional corrections (e.g., inverse-probability weighting, Heckman correction) and standard Bayesian approaches rely on tractable likelihoods. They require explicit parametric forms for both the outcome model and the selection mechanism. This makes them infeasible for complex stochastic models involving latent dynamics, high-dimensional hidden states, or nonlinear interactions where likelihoods are intractable.
Limitations of Current Simulation-Based Inference (SBI): While SBI (specifically Neural Posterior Estimation, or NPE) allows inference without tractable likelihoods, standard implementations assume data is missing at random (MAR) or that the observed data represents an unbiased realization of the generative process. Applying standard SBI to selectively observed data results in biased posterior estimates and miscalibrated uncertainty.
The Gap: There is a lack of an amortized (reusable) framework that explicitly accounts for outcome-dependent or dynamically coupled selection mechanisms while providing tools to detect bias and assess posterior calibration in likelihood-free settings.

2. Methodology

The authors propose a Bias-Aware Amortized Bayesian Inference Framework that embeds the selection mechanism directly into the generative simulation process.

Core Concept: Embedding Selection in the Forward Process

Instead of treating selection as a post-hoc correction, the framework treats it as a structured observation operator acting on latent population-level processes.

Generative Model: The simulation pipeline generates parameters ( $\theta$ ) and covariates ( $C$ ), produces outcomes ( $Y$ ) via a mechanistic model, and then applies a selection operator ( $S$ ) to determine which individuals are observed.
Training Data: The Neural Posterior Estimator (NPE) is trained on pairs of (Observed Data $Y_{obs}$ , Selection Indicator $S=1$ , Covariates $C$ ) and their corresponding ground-truth parameters ( $\theta$ ).
Target Distribution: The NPE learns to approximate the conditional distribution $p(\theta \mid Y, S=1; C)$ , effectively learning to "invert" the biased observation process.
Amortization: Once trained, the NPE can perform inference on new datasets with the same selection mechanism instantly, without retraining.

Diagnostic Framework

To ensure the validity of the bias correction, the framework integrates two simulation-based diagnostics:

Simulation-Based Calibration (SBC): Checks if the posterior ranks of ground-truth parameters are uniformly distributed across repeated simulations. This verifies that the posterior uncertainty is correctly quantified under the selection mechanism.
Classifier Two-Sample Test (C2ST): A binary classifier is trained to distinguish between:
1. Pairs of $(\theta, Y_{obs})$ drawn from the learned posterior.
2. Pairs drawn directly from the joint generative distribution (including selection).
- If the classifier performs at chance level (accuracy $\approx 0.5$ ), the posterior is consistent with the generative model. If it performs significantly better, bias or model misspecification is detected.

3. Key Contributions

Likelihood-Free Bias Correction: The method enables Bayesian inference under selection bias without requiring tractable likelihoods for either the outcome or the selection process. It relies solely on forward simulations.
Amortized Efficiency: Unlike standard SBI which might require retraining for different datasets, this approach trains once on simulated biased data and applies globally to real-world datasets sharing the selection mechanism.
Explicit Bias Detection: By integrating C2ST and SBC, the framework provides a rigorous, testable strategy to detect whether the assumed selection mechanism matches the real data, a feature often missing in likelihood-free methods.
Handling Complex Dependencies: The framework successfully handles scenarios where selection depends on unobserved outcomes (Missing Not At Random - MNAR), latent states (e.g., death censoring dementia onset), and coupled covariates.

4. Results and Applications

The framework was evaluated across three distinct epidemiological applications:

A. Prevalence Estimation (KoCo19 Study)

Scenario: Estimating SARS-CoV-2 seroprevalence with non-representative sampling and outcome missingness.
Comparison: Compared against unadjusted estimators and Inverse Probability Weighting (IPW).
Outcome: The bias-aware NPE outperformed IPW, particularly in rounds with high missingness. IPW failed when missingness was informative (dependent on unobserved outcomes), whereas the NPE recovered accurate prevalence estimates and well-calibrated uncertainty. C2ST confirmed the posterior consistency.

B. Time-to-Event Analysis (Framingham Heart Study)

Scenario: Estimating dementia incidence where death acts as a competing risk, censoring the observation of dementia onset (MNAR).
Comparison: Compared against a standard NPE (trained on full data) and a penalized likelihood approach (Illness-Death Model).
Outcome: The standard NPE failed significantly, producing biased hazard estimates. The bias-aware NPE recovered transition hazards with accuracy comparable to the full-data standard NPE. C2ST diagnostics clearly distinguished the biased posterior of the standard NPE from the joint distribution, while the bias-aware NPE passed the test.

C. Transmission Modeling (PedCovid Study)

Scenario: Modeling household transmission of SARS-CoV-2 where study inclusion depended on a child testing positive (outcome-dependent selection).
Complexity: Involved a complex stochastic continuous-time transmission model with intractable likelihoods.
Comparison: Compared against Markov Chain Monte Carlo (MCMC) using a simplified likelihood assumption (random inclusion).
Outcome: MCMC produced systematically biased estimates for transmission parameters (e.g., age-specific infectivity) because it ignored the selection mechanism. The bias-aware NPE recovered ground-truth parameters in simulations and provided coherent estimates on real data. MCMC took ~12 hours (96 CPUs), while the amortized NPE inference took seconds (single GPU).

5. Significance

Paradigm Shift: The paper reframes selection bias correction from a post-hoc statistical adjustment problem to an explicit generative modeling problem.
Scalability: It makes bias-aware inference feasible for complex, high-dimensional models (e.g., agent-based models, multi-state processes) where deriving analytical corrections is impossible.
Robustness: The integration of C2ST allows researchers to validate their assumptions about the selection mechanism even when ground truth is unknown, preventing the deployment of biased models.
Practical Utility: By amortizing the inference, the method enables rapid sensitivity analysis and experimental design, allowing researchers to explore how different data collection strategies affect inference without re-deriving likelihoods.

In summary, this work establishes a principled, practical, and testable framework for Bayesian inference in the presence of complex selection biases, bridging the gap between simulation-based inference and rigorous bias correction.