Prediction Models That Learn to Avoid Missing Values

TL;DR

The paper addresses the challenge of making accurate, interpretable predictions when test-time data contain missing values. It introduces missingness-avoiding (MA) learning, a framework that adds a missingness-reliance regularization term to standard learning objectives, and instantiates it across decision trees (MA-DT), sparse linear models (MA-LASSO), and tree ensembles (MA-RF, MA-GBT). The authors formalize missingness reliance with $\rho(h)$ and optimize $\mathbb{E}_p[L(Y,h(X))] + \alpha\rho(h)$, demonstrating that MA variants can maintain competitive AUROC while substantially reducing reliance on missing features across six real-world datasets; they also analyze when MA is feasible, and how distribution shifts or MNAR can affect performance. The results show that MA models can provide strong predictive performance with enhanced interpretability in the presence of test-time missing data, offering practical guidance for practitioners on tuning $\alpha$ and selecting model classes. Overall, this work advances the ability to deploy transparent, robust predictions under structured missingness and outlines directions for broader applicability and robustness.

Abstract

Handling missing values at test time is challenging for machine learning models, especially when aiming for both high accuracy and interpretability. Established approaches often add bias through imputation or excessive model complexity via missingness indicators. Moreover, either method can obscure interpretability, making it harder to understand how the model utilizes the observed variables in predictions. We propose missingness-avoiding (MA) machine learning, a general framework for training models to rarely require the values of missing (or imputed) features at test time. We create tailored MA learning algorithms for decision trees, tree ensembles, and sparse linear models by incorporating classifier-specific regularization terms in their learning objectives. The tree-based models leverage contextual missingness by reducing reliance on missing values based on the observed context. Experiments on real-world datasets demonstrate that MA-DT, MA-LASSO, MA-RF, and MA-GBT effectively reduce the reliance on features with missing values while maintaining predictive performance competitive with their unregularized counterparts. This shows that our framework gives practitioners a powerful tool to maintain interpretability in predictions with test-time missing values.

Figures (6)

• Figure 1: Two decision trees built to diagnose cognitive impairment (CI) in adult patients. The left side illustrates a regular decision tree fit solely for accuracy, while the right shows a missingness-avoiding (MA) tree, which incorporates a regularization parameter $\alpha$ to reduce reliance on missing values. The regular tree initially splits on a positive MRI scan outcome, determined by hippocampal volume ($V_h$). However, this feature is missing for many patients who did not undergo a scan, resulting in high reliance on missing values (indicated by a large $\rho$). In contrast, the MA tree achieves a comparable AUC while avoiding reliance on missing data entirely ($\rho = 0$). Orange nodes indicate missing values along the decision path, while blue denotes no missingness reliance. Note that zero imputation is used, resulting in complete observations in the right branch of the regular tree.
• Figure 2: Missing values can be avoided in several ways. Sparse models (left) can be trained to not use features that are frequently missing. Disjunctive rule models (middle) can be fit to include rules that exploit redundancy in the variable set. Trees (right) can be fit so that missing values rarely occur on the decision paths.
• Figure 3: Average test performance across cross-validation folds for MA decision trees fit to LIFE using varying values of $\alpha$ and . A shallow tree of depth 3--4 captures the complexity of the data effectively. Missingness regularization with $\alpha\le1.0$ does not significantly reduce predictive performance but helps decrease missingness reliance from approximately 30% to approximately 10%.
• Figure 4: (a) and (b): Test-set AUROC and missingness reliance ($\hat{\rho}$) for MA estimators in ADNI when transitioning from $\alpha=\infty$ to $\alpha=0$. Removing missingness regularization improves predictive performance, but the increase in missingness reliance is much more pronounced, especially for and . (c): Test-set AUROC versus $\hat{\rho}$ for in Breast Cancer, where 50% synthetic missingness is added to an increasing proportion of input features. Missingness not at random (MNAR) is more challenging than missingnes at random (MAR), but demonstrates robust performance for large fractions of missingness.
• Figure 5: LIFE: $\mathrm{AUROC}-\hat{\rho}$ vs. max. tree depth.

Theorems & Definitions (8)

• Example 1
• Definition 1
• Proposition 1
• Corollary 1
• Proposition 2
• proof
• proof
• proof