Few-shot Cross-country Generalization of Tabular Machine Learning and Foundation Models for Childhood Anemia Prediction under Distribution Shift

Technical Summary: Few-shot Cross-country Generalization of Tabular Machine Learning and Foundation Models for Childhood Anemia Prediction under Distribution Shift

Problem Statement
Childhood anemia affects approximately 40% of children aged 6–59 months globally, driven by heterogeneous nutritional, infectious, and socioeconomic factors that vary significantly across geographic settings. This heterogeneity creates a "distribution shift" challenge for predictive machine learning models: models trained in one country often degrade when applied to another due to differing epidemiological profiles and data characteristics. Furthermore, many low-resource settings where anemia surveillance is most needed suffer from data scarcity, limiting the efficacy of traditional supervised learning methods that require large labeled datasets. Existing literature has largely focused on single-country studies using standard classifiers (e.g., Logistic Regression, Random Forest) without rigorous assessment of cross-population generalization, uncertainty quantification, or the potential of emerging foundation models in tabular domains.

Methodology
The authors conducted a multi-country predictive modeling study using harmonized Demographic and Health Surveys (DHS) children's recode data from 16 countries spanning Africa, Asia, Latin America, the Caucasus, and the Middle East. The final analytic cohort included 68,856 children with valid hemoglobin measurements. Anemia was defined as a binary outcome (hemoglobin < 11.0 g/dL) based on WHO age and altitude-adjusted thresholds.

The study evaluated four modeling approaches:

1. Logistic Regression (LogReg): A linear baseline with L2 regularization.
2. XGBoost and LightGBM: Gradient-boosted tree algorithms representing state-of-the-art supervised tabular learning.
3. TabPFN v2.6: A transformer-based tabular foundation model that utilizes in-context learning (ICL) without gradient-based training updates during inference.

The evaluation framework was comprehensive, employing:

• Internal Validation: Stratified 5-fold cross-validation within each country.
• External Validation (LOCO): Leave-One-Country-Out validation, where models were trained on pooled data from 15 countries and tested on the held-out 16th.
• Reverse LOCO: Training on a single source country and testing on the remaining 15 to assess directional transferability.
• Few-Shot Learning: Evaluating performance as a function of training sample size (from near-zero to full data) to simulate data-scarce regimes.
• Metrics: Primary metric was AUC-ROC. Calibration was assessed via Brier score and Expected Calibration Error (ECE). Uncertainty was quantified using 5,000-iteration bootstrap resampling to derive 95% confidence intervals.
• Interpretability: SHAP values (TreeSHAP for tree models, standard SHAP for LogReg, and shapiq for TabPFN) were used to assess feature importance across models and countries.

Key Results

• Few-Shot Advantage: In data-scarce regimes (specifically below ~200 within-country training samples), TabPFN consistently outperformed all supervised models (LogReg, XGBoost, LightGBM). In countries like Guatemala, Haiti, Malawi, and Gambia, TabPFN achieved AUC-ROC values 0.10–0.15 higher than logistic regression in the near-zero-shot regime. As sample sizes increased, the performance gap narrowed, and supervised models eventually reached comparable discrimination.
• Calibration Superiority: TabPFN demonstrated superior probabilistic calibration across all 16 countries, achieving the lowest mean Brier score (0.203) and Expected Calibration Error (0.042). In contrast, LightGBM exhibited the highest miscalibration, particularly at higher predicted probabilities.
• Generalization Limits: Under LOCO validation, all models experienced a modest decline in performance (AUC-ROC 0.58–0.69), indicating that distribution shift remains a significant barrier. The performance was driven more by the target country's characteristics than by the choice of model architecture.
• Asymmetric Transferability: Reverse LOCO experiments revealed that transferability is asymmetric. Epidemiologically diverse and larger populations (e.g., Ghana, Guatemala, Haiti) served as effective training sources for other countries. Conversely, models trained on smaller or atypical populations (e.g., Armenia) transferred poorly. Certain target populations remained difficult to predict regardless of the training source.
• Feature Importance: Across all models and countries, child age was the dominant predictor, followed by altitude and height-for-age z-score. Wealth index and maternal education were also prominent. The ranking of features was largely consistent across models, validating the identified predictor hierarchy.
• Subgroup Analysis: Performance was broadly similar across demographic strata (age, sex, education, wealth, residence) within countries, with no model exhibiting consistent bias toward specific subgroups. Variation was primarily driven by country context rather than model architecture.

Significance and Claims
The paper claims that predictive performance in childhood anemia modeling is currently constrained more by cross-population heterogeneity and data characteristics than by model architecture. Within these constraints, the study establishes that transformer-based tabular foundation models (specifically TabPFN) offer a meaningful and practically relevant advantage in data-scarce public health settings.

The authors argue that TabPFN's ability to leverage pre-trained representations allows it to function as a form of implicit regularization, stabilizing inference when local labeled data is insufficient. This is critical for low-resource surveillance contexts where data paucity is the norm. Furthermore, the superior calibration of TabPFN suggests it provides more reliable risk estimates, which is essential for screening and resource allocation decisions.

The study shifts the emphasis from purely selecting the best-performing algorithm to addressing structural challenges in global health data, such as population heterogeneity and data quality. It posits that foundation models represent a promising pathway for anemia prediction in data-scarce environments, provided that the limitations of cross-population transferability and the need for local data infrastructure are acknowledged. The findings are modest regarding the "ceiling" of performance, noting that even the best models struggle in settings with high missingness or unique etiologies (e.g., endemic malaria), reinforcing that algorithmic innovation alone cannot resolve all generalization challenges.