Transporting Predictions via Double Machine Learning: Predicting Partially Unobserved Students' Outcomes

1. Problem Statement

The paper addresses the challenge of covariate shift in predictive modeling within social sciences. Specifically, it tackles the scenario where a predictive model is trained on a source population (where outcomes are observed) but must be applied to a target population (where outcomes are missing or unobserved) that has a different distribution of covariates (features).

• Context: Educational policymakers often lack outcome data (e.g., standardized test scores) for specific subgroups or regions.
• The Issue: Standard machine learning models assume the training and test data come from the same distribution. When covariate distributions differ between the source and target (e.g., different socioeconomic characteristics between Belgian regions), direct application of the model can lead to biased predictions, reduced accuracy, and reliance on unreliable extrapolation.
• Goal: To develop a framework that transports predictions from a source to a target population while correcting for distributional shifts, and to identify "outlier" predictions (students with extreme predicted outcomes) to inform policy.

2. Methodology: Double Machine Learning (DML) for Transportability

The authors propose a Double Machine Learning (DML) framework adapted for covariate shift, utilizing importance weighting to align the source distribution with the target distribution.

A. Core Assumptions

1. Conditional Sample Exchangeability (Explainable Covariate Shift): The conditional distribution of the outcome given covariates is invariant across populations ($f(Y|X, S=1) = f(Y|X, S=0)$). This implies the outcome is Missing At Random (MAR) conditional on observed covariates.
2. Positivity: Every covariate pattern in the target population has a non-zero probability of appearing in the source population ($Pr(S=1|X) > 0$).

B. The Two-Stage DML Framework

The method involves estimating two unknown quantities: the overlap scores (nuisance function) and the prediction function.

1. Stage 1: Overlap Score Estimation

   • A classification model is trained to estimate the probability that an observation belongs to the source population given its covariates: $\hat{P}(S=1|X)$.
   • This is referred to as the overlap score.
   • The model uses flexible machine learning methods (e.g., Logistic Regression, Random Forests, or BART) to avoid parametric misspecification.

2. Stage 2: Importance Weighted Prediction

   • Balancing Weights: The overlap scores are converted into balancing weights (inverse odds) to reweight the source data:
     $$\hat{\omega}_i = \frac{\hat{P}(S=0|X_i)}{\hat{P}(S=1|X_i)} = \frac{1 - \hat{P}(S=1|X_i)}{\hat{P}(S=1|X_i)}$$
   • Weighted Loss Optimization: The prediction model is trained on the source data by minimizing a weighted loss function. Observations in the source that resemble the target population are upweighted, while those that are dissimilar are downweighted.
   • Implementation with BART: The authors implement this using Bayesian Additive Regression Trees (BART).
     • In the Bayesian framework, weighting an observation $i$ by $\omega_i$ is equivalent to scaling its variance in the likelihood function: $\sigma^2_i = \sigma^2 / \omega_i$.
     • This allows the model to naturally incorporate the weights while maintaining the benefits of Bayesian inference (uncertainty quantification).

C. Outlier Detection

Using the posterior predictive distribution from BART, the authors identify "outliers" (students with extreme predicted outcomes).

• Definition: Observations where the predicted value falls outside the 95% credible interval of the mean posterior prediction (specifically, $< 2$ standard deviations from the mean).
• Purpose: To identify vulnerable subpopulations (e.g., students with very low predicted financial literacy) for targeted policy interventions.

3. Key Contributions

1. Novel Framework: Introduces a Double Machine Learning approach specifically for transporting predictions in the presence of covariate shift, bridging the gap between causal inference (generalizability) and predictive modeling.
2. Bayesian Implementation: Demonstrates how to integrate importance weighting into BART, leveraging its ability to provide principled uncertainty quantification (credible intervals) and handle non-linear relationships without extensive tuning.
3. Empirical Application: First application of supervised machine learning to predict Financial Literacy Scores (FLS) in the PISA dataset, specifically predicting missing scores for the Walloon region of Belgium using data from Flanders.
4. Diagnostic Insight: Provides a rigorous analysis of when covariate shift correction is necessary versus when it introduces noise, offering practical guidance for social science applications.

4. Results

A. Simulation Study

• Setup: Simulated data with varying degrees of overlap between source and target populations (High, Moderate, Low overlap).
• Findings:
  • High/Moderate Overlap: Importance weighting degraded performance compared to unweighted models. The variance introduced by the weights outweighed the bias reduction.
  • Low Overlap (Severe Shift): Importance weighting improved performance, significantly reducing RMSE compared to unweighted models.
  • Conclusion: Weighting is beneficial only when covariate shift is pronounced; otherwise, it introduces unnecessary noise.

B. Empirical Application (PISA Data: Flanders vs. Wallonia)

• Data: 5,675 students from Flanders (Source, FLS observed) and 3,976 from Wallonia (Target, FLS missing).
• Predictive Performance:
  • Unweighted BART: Achieved an adjusted $R^2 \approx 0.73$ with low RMSE/MAE.
  • Weighted BART: Showed no improvement over the unweighted model; in some proxy tests (Math/Reading scores), performance slightly declined.
  • Reasoning: The overlap between Flanders and Wallonia was sufficient (covariate shift was mild). The weighting scheme introduced variance without significant bias reduction.
• Outlier Analysis:
  • Identified students with the lowest predicted FLS.
  • Key Drivers of Low FLS:
    • Low scores in Math and Reading (PISA tests).
    • Speaking a language at home different from the school language.
    • Low parental education levels.
    • First-generation immigrant status (particularly in Flanders).
  • Policy Implication: These factors define the "vulnerable" student profile, allowing for targeted interventions.

5. Significance and Conclusion

• Methodological Robustness: The paper validates that while Double Machine Learning provides a theoretically principled way to handle covariate shift, its practical utility depends on the severity of the shift. In settings with moderate overlap (common in regional education data), standard flexible models like BART may be robust enough without complex reweighting.
• Policy Relevance: The framework offers a tool for policymakers to:
  1. Predict missing outcomes in under-sampled populations.
  2. Quantify uncertainty in these predictions.
  3. Identify specific subgroups at risk of poor outcomes based on interpretable decision trees derived from the model.
• Generalizability: The approach is broadly applicable to social and health sciences where data is often missing for specific subpopulations, and distributional differences between groups are a primary concern.

Final Takeaway: The study successfully demonstrates a pipeline for transporting predictions and detecting outliers. However, it crucially highlights that covariate shift correction is not a panacea; in cases of sufficient overlap, simpler unweighted models may outperform complex weighted counterparts due to the variance inflation caused by importance weights.