📊 statistics

Improving Survey Inference in Two-phase Designs Using Bayesian Machine Learning

This paper proposes a Bayesian tree-based multiple imputation approach that incorporates complex survey design features to improve population mean estimation in two-phase sampling, demonstrating through simulations and a real-world COVID-19 vaccination study in Uganda that it outperforms traditional weighted estimators in terms of bias, accuracy, and confidence interval coverage.

Original authors: Xinru Wang, Anyu Zhu, Lauren Kennedy, Abigail Greenleaf, Qixuan Chen

Published 2026-04-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xinru Wang, Anyu Zhu, Lauren Kennedy, Abigail Greenleaf, Qixuan Chen

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 Picture: The "Two-Step" Survey Problem

Imagine you want to know the average height of every adult in a massive city.

Step 1 (Phase I): You knock on every door and ask for their age, gender, and zip code. This is cheap and easy. You get data on 10,000 people.
Step 2 (Phase II): You want to know their exact height, but measuring height is expensive and takes time. So, you only measure the height of 1,000 people you randomly picked from your first list.

The Problem: The 1,000 people you measured might not be perfectly representative of the whole city. Maybe you accidentally picked more tall people, or more people from a specific neighborhood. If you just average the 1,000 people, your guess for the whole city will be wrong.

The Old Solution (Weighting):
Traditionally, statisticians try to fix this by giving "weights" to the people they measured. If the 1,000 people have too many tall men, they might say, "Okay, we'll count each tall man as only 0.5 of a person, and each short woman as 1.5 people" to balance the math.

The Flaw: Sometimes, to fix the math, the weights get crazy. You might end up saying one person counts as 50 people. This makes the final answer very shaky and unstable, like trying to balance a house of cards on a windy day.

The New Solution: The "Magic Imputation" (BART)

The authors of this paper propose a smarter way. Instead of trying to stretch the 1,000 people to represent 10,000, they say: "Let's guess the heights of the 9,000 people we didn't measure."

They use a powerful computer brain called BART (Bayesian Additive Regression Trees). Think of BART as a super-smart detective that looks at the 10,000 people from Step 1. It notices patterns: "Oh, people in this zip code with this job and this age tend to be 5'8"."

The Detective Work: The BART model looks at the 1,000 people whose heights were measured. It learns the rules connecting age, location, and gender to height.
The Guess: It then uses those rules to "impute" (guess) the heights for the 9,000 people who weren't measured.
The Final Count: Now, instead of just 1,000 people, the computer has a complete list of 10,000 people (1,000 real, 9,000 guessed). It averages them all out.

Why is this better?

Stability: It doesn't rely on crazy "weights" that swing wildly.
Accuracy: It uses all the information from the first step (the rich data) to make better guesses, rather than just trying to fix the second step.
Flexibility: It can handle thousands of clues (variables) at once without getting confused.

The Uganda Example: The "Cell Phone" Survey

The authors tested this idea in real life in Uganda.

The Big Survey: They did a massive health survey (UPHIA) where they visited thousands of homes.
The Sub-Survey: Later, they called a smaller group of people on their cell phones to ask about COVID-19 vaccinations.
The Bias: People with cell phones are usually different from people without them (often wealthier, more urban, different ages). If you just look at the phone survey, you might think everyone in Uganda is vaccinated, because the phone owners are.

The Result:
When they used the old "weighting" method, their guess was too high (overestimating vaccination). When they used the new BART Imputation method, they used the rich data from the big survey to "fill in the blanks" for the people without phones.

Outcome: Their new estimate was much closer to the truth and more stable. It correctly identified that vaccination rates were likely lower than the phone survey suggested.

The "Secret Sauce" (Why it works)

The Tree Analogy: Imagine the data is a giant forest. Traditional methods try to draw a straight line through the trees. BART builds a complex, branching tree structure that follows the twists and turns of the data perfectly, capturing hidden patterns.
Rubin's Rules: Since the computer is making guesses, there is some uncertainty. The authors use a special statistical rule (Rubin's rules) to say, "We are 95% sure the answer is between X and Y." This ensures they don't pretend to be more certain than they actually are.
The "Black Box" isn't so black: Even though BART is a complex machine learning model, the authors showed you can still look inside to see which clues mattered most (e.g., "Living in the city" was a huge predictor of vaccination).

Summary in One Sentence

Instead of trying to stretch a small, biased sample to represent a whole population (which creates shaky math), this paper suggests using a smart computer model to "fill in the missing pieces" of the data, resulting in a more accurate and stable picture of the whole population.

1. Problem Statement

Two-phase sampling designs are widely used in public health and social sciences to reduce costs when collecting expensive or difficult-to-obtain outcome data. In this design:

Phase I: A large sample is drawn from a finite population using a complex survey design (e.g., stratified multistage sampling), collecting inexpensive auxiliary variables ( $x, z$ ).
Phase II: A subsample is selected from Phase I to collect the expensive outcome variable ( $y$ ).

The Core Challenge:
Standard analysis of Phase II data typically relies on weighting (creating subsample-specific weights by multiplying Phase I weights by subsample adjustment factors). This approach faces three major limitations:

High Variance: Compounding weights leads to extreme variability, causing unstable estimates.
Inefficiency: Weighting focuses on correcting distributional differences between samples rather than leveraging the predictive power of Phase I variables to estimate outcomes.
Model Misspecification: If the propensity model used to calculate weights is misspecified, the resulting estimates are biased.

Furthermore, existing machine learning approaches for survey inference (e.g., Liu et al., 2023) were designed for single-phase samples and do not adequately account for the complex design features (strata, clusters) inherent in Phase I of a two-phase design.

2. Methodology

The authors propose a Bayesian Tree-based Multiple Imputation (MI) framework that treats the missing outcomes in the Phase I sample (those not selected for Phase II) as a missing data problem rather than a weighting problem.

Key Components:

Imputation Model (BART/rBART):
- The authors utilize Bayesian Additive Regression Trees (BART) to model the relationship between the outcome $y$ and the auxiliary variables.
- To handle complex survey designs, they incorporate design features directly into the model:
  - Standard BART: Includes strata, cluster indicators, and survey weights as covariates.
  - Random Intercept BART (rBART): Explicitly models cluster effects using a random intercept ( $\delta_j$ ) to account for intra-cluster correlation.
- The model also includes the subsample weighting adjustment ( $a_i$ ) as a predictor to guard against model misspecification.
Multiple Imputation Process:
1. Generate $D$ imputed datasets for the Phase I sample using posterior draws from the BART/rBART model.
2. For each imputed dataset, calculate the population mean using the Phase I sample weights ( $w_{ci}$ ).
3. Combine the results using Rubin's Rules to obtain the final point estimate and variance.
Variance Estimation:
- The variance for each imputed dataset is estimated using Taylor series linearization to account for the Phase I complex design (strata/clusters).
- The total variance combines within-imputation and between-imputation variances, with degrees of freedom adjusted via a Satterthwaite approximation.

3. Key Contributions

Extension to Two-Phase Designs: The paper extends tree-based imputation methods (previously limited to single-phase samples) to complex two-phase designs, explicitly integrating stratification and clustering into the imputation model.
Superiority over Weighting: The authors demonstrate that treating the problem as imputation yields more efficient estimates than traditional weighting, particularly when high-dimensional auxiliary variables are available.
Robustness to Model Misspecification: By using BART, the method avoids the need to pre-specify linear or interaction terms, making it robust to non-linearities and complex interactions between covariates and outcomes.
Validation of Rubin's Rules: The study confirms that applying Rubin's variance estimation rules to BART-imputed data in two-phase designs produces valid statistical inference with coverage rates close to nominal levels.

4. Results

Simulation Study

The authors conducted extensive simulations across five scenarios (varying dimensionality and assumption violations) comparing:

Weighting Methods: WT-LGM (Logistic Regression), WT-CHAID, WT-BART, WT-rBART.
Imputation Methods: MI-BART, MI-rBART.
Benchmark: Full Phase I sample (ideal but unattainable).

Findings:

Bias and RMSE: The MI-based estimators (MI-BART and MI-rBART) consistently outperformed all weighting methods, exhibiting lower bias and Root Mean Squared Error (RMSE).
Confidence Intervals: MI methods produced narrower 95% confidence intervals compared to weighting methods.
Coverage Rates: MI methods achieved coverage rates close to the nominal 95% level. Weighting methods, particularly in scenarios with model misspecification or high dimensionality, suffered from poor coverage.
Assumption Violations: Even when Assumption A3 (comparable ranges of auxiliary variables) was violated (sparse data in tails), MI methods remained superior to weighting, though performance degraded for both.
Sensitivity: Including the estimated subsample propensity as a covariate in the imputation model slightly reduced bias compared to excluding it.

Real-World Application: Uganda Subcohort Survey

The method was applied to estimate COVID-19 vaccination coverage using a cellphone subcohort drawn from the 2020 Uganda Population-based HIV Impact Assessment (UPHIA2020).

Data: 124 covariates (household and individual level) were used.
Comparison: The MI-BART/rBART estimates were compared against weighted estimates (WT-BART/rBART).
Results:
- Weighted estimates suggested vaccination coverage of ~60%, which was higher than the CDC's reported ~57% for a later period, suggesting potential upward bias in the weighting approach.
- MI-rBART estimated 53.5% (95% CI: 50.8%, 56.3%), which was closer to external benchmarks and had a narrower confidence interval.
- A validation study using "health facility visits" confirmed that MI methods produced estimates closer to the full-parent-cohort benchmark than weighting methods.

5. Significance

This paper provides a robust, cost-effective framework for analyzing two-phase survey data. By shifting from weighting to Bayesian Multiple Imputation, researchers can:

Leverage rich Phase I data to improve the precision of Phase II estimates.
Avoid the instability caused by extreme weights in complex designs.
Handle high-dimensional auxiliary data without manual variable selection.
Obtain valid statistical inference even when the sampling mechanism is complex (stratified/multistage).

The approach is particularly valuable for public health research where follow-up surveys (like cellphone surveys) are conducted on subcohorts of large national surveys, allowing for efficient estimation of population parameters without the prohibitive cost of re-surveying the entire population.