Technical Summary: Prediction Models That Learn to Avoid Missing Values

Problem Statement

The deployment of machine learning models in real-world scenarios is frequently hindered by missing values, particularly at test time. Traditional approaches, such as "impute-then-regress" (imputing missing data before training) or using missingness indicators, often introduce bias, increase model complexity, or obscure interpretability. While some specialized models (e.g., XGBoost with default paths, Bayesian networks) handle missingness natively, they are often restricted to specific model classes.

The core challenge addressed is the trade-off between predictive accuracy and the reliance on missing features. In safety-critical fields like healthcare, a model that requires imputed values or features that are frequently unavailable at test time can be unreliable or difficult to interpret. The authors propose a framework where models are trained to strategically avoid using features that are missing for a specific input, rather than forcing the model to rely on imputed values or missingness masks.

Methodology: Missingness-Avoiding (MA) Learning

The paper introduces Missingness-Avoiding (MA) machine learning, a general framework that modifies the learning objectives of standard algorithms to penalize reliance on missing values. The objective function minimizes a trade-off between expected predictive loss and expected missingness reliance:

$$\min_{h \in \mathcal{H}} \mathbb{E}_p[L(Y, h(X))] + \alpha \rho(h)$$

Where:

• $L$ is the prediction loss.
• $\rho(h)$ is the expected missingness reliance, defined as the probability that computing a prediction $h(x)$ requires access to a missing feature $x_j$ (where $x_j = na$).
• $\alpha \ge 0$ is a hyperparameter controlling the strength of the penalty.

The authors derive specific algorithms for three model classes:

1. MA Decision Trees (MA-DT)

Decision trees are naturally suited for this task because their reliance on features is contextual; a feature is only required if it appears on the decision path for a specific input.

• Mechanism: The authors modify greedy splitting algorithms (like C4.5 or ID3). During the split selection process, the criterion (e.g., Gini impurity) is augmented with a regularization term.
• Regularization: For a leaf $\ell$ and a candidate split $(j, \tau)$, the cost function includes a penalty proportional to the fraction of samples in that leaf where feature $j$ is missing.
• Contextual Avoidance: The algorithm learns to route samples with missing values down branches that do not require the missing feature, effectively exploiting the "contextual missingness" of the data.

2. MA Sparse Linear Models (MA-LASSO)

Linear models cannot contextually avoid missing values; if a coefficient $\theta_j \neq 0$, the model always requires $x_j$.

• Mechanism: The framework adapts the LASSO (L1 regularization) approach.
• Regularization: Instead of a uniform penalty, the penalty for each feature $j$ is scaled by its empirical missingness rate $m_j$. The objective minimizes:
  $$\sum L(y_i, \theta^T x^I_i) + \sum_{j=1}^d (\lambda + \alpha m_j) |\theta_j|$$
• Effect: Features with high missingness rates receive a larger penalty, encouraging the model to set their coefficients to zero, thereby reducing reliance on missing data.

3. MA Tree Ensembles (MA-RF and MA-GBT)

• MA Random Forest (MA-RF): Trains independent MA trees in parallel. Each tree minimizes the regularized splitting criterion.
• MA Gradient Boosted Trees (MA-GBT): Uses a boosting approach where subsequent trees fit the residuals of previous ones. Crucially, the regularization weights are updated dynamically. If a feature $j$ is already used by a previous model for a specific observation $x_i$ (and was missing), subsequent models are penalized less for relying on it again. This prevents the ensemble from accumulating excessive reliance on missing features across different trees.

Key Contributions

1. Framework Introduction: The proposal of a general "Missingness-Avoiding" framework that reconciles high accuracy with minimal reliance on missing features at test time.
2. Algorithmic Implementation: The development of specific, fit-predict compatible algorithms for Decision Trees (MA-DT), Sparse Linear Models (MA-LASSO), and Tree Ensembles (MA-RF, MA-GBT).
3. Theoretical Analysis:
   • Identification of Observed Deterministic Data Collection (ODDC) rules, where specific combinations of observed variables guarantee the presence of other variables. The authors prove that if such rules exist, models can achieve zero missingness reliance without sacrificing predictive accuracy.
   • Analysis of limitations, showing that if missingness is completely random (MCAR) or if the test-time distribution shifts significantly, the trade-off between accuracy and reliance may be unavoidable.
4. Empirical Validation: Demonstration across six real-world datasets (including NHANES, ADNI, and FICO) that MA models significantly reduce missingness reliance (often from 100% to near 0%) while maintaining predictive performance (AUROC) competitive with unregularized baselines.

Experimental Results

The authors evaluated MA-DT, MA-LASSO, MA-RF, and MA-GBT against standard baselines (imputation-based models, missingness indicators like M-GAM, and native handlers like XGBoost and NeuMiss) on six datasets.

• Performance vs. Reliance: MA models consistently achieved lower missingness reliance ($\hat{\rho}$) than their unregularized counterparts. For example, MA-DT reduced reliance from ~11.7% to 0.1% on the ADNI dataset while maintaining comparable AUROC.
• Comparison to Baselines:
  • MA models outperformed or matched standard imputation methods (Zero Imputation, MICE) in terms of reliance reduction.
  • MA-DT and MA-RF achieved significantly lower reliance than XGBoost and Random Forests, which often rely heavily on missing features.
  • MA-LASSO reduced reliance significantly compared to standard LASSO, with minimal loss in AUROC.
  • While MINTY (a related rule-based model) achieved low reliance, MA-LASSO was significantly faster to train (under 1 second vs. up to 292 seconds).
• Trade-off: The experiments confirmed that a favorable trade-off is achievable. By tuning $\alpha$, practitioners can select models that operate on the Pareto frontier of accuracy and missingness avoidance.

Significance and Claims

The paper claims that MA learning provides a powerful tool for practitioners who require interpretability and robustness in the presence of missing data.

• Interpretability: By avoiding reliance on missing features, the decision paths of MA models remain transparent and do not depend on the opaque logic of imputation or the complex interactions introduced by missingness indicators.
• Practicality: The approach does not require specialized architectures (like NeuMiss) or complex imputation pipelines. It can be integrated into standard workflows using existing tree and linear model libraries.
• Feasibility: The authors argue that in many real-world scenarios (particularly those governed by ODDC rules, such as clinical guidelines), it is possible to build models that are both highly accurate and entirely free of missingness reliance, challenging the notion that missing data must always be imputed or indicated to achieve good performance.

The work concludes that while avoiding missing values entirely may not be feasible in all scenarios (e.g., when missingness is purely random or informative in a way that cannot be exploited), the MA framework offers a principled method to minimize this reliance without compromising predictive utility.