Prediction models that learn to avoid missing values

We thank the reviewer for their valuable feedback and appreciate their recognition of the clear explanations and sensible evaluation criteria. We respond to specific comments below.

Re: real-world importance of reducing missingness reliance

Thanks for raising this point—we understand its importance in demonstrating the improved interpretability of our MA method. In the last paragraph of Section 6.2, we refer to MA trees fitted to the Life dataset, shown in Figure 6 in the appendix. This example demonstrates how MA-DT behaves under different values of the missingness regularization parameter $\alpha$. For instance, when $\alpha$ is tuned to balance accuracy and reliance on missing values, MA-DT reduces reliance on missing data by 33% without sacrificing accuracy compared to a standard decision tree. We will consider including these figures in the main text in a potential camera-ready version. Additionally, we will consider adding examples of MA trees fitted to the ADNI dataset. These examples illustrate how MA learning avoids splitting on features that are likely to be missing, instead favoring more complete features that maintain predictive performance. At test time, such trees are more interpretable since they do not rely on features with missing values.

Re: different strategies for choosing $\alpha$

In the submitted version of the work, we primarily considered two strategies for selecting $\alpha$. We chose $\alpha=\alpha^{*}$ as part of the overall model selection process, selecting the candidate model with the lowest missingness reliance among those achieving at least 95% of the maximum AUC. If predictive performance is a priority, this percentage could be increased (at the potential cost of increased missingness reliance). Conversely, we chose $\alpha=\infty$ by selecting the model with the highest AUC among those achieving near-zero missingness reliance (ensuring that the range of $\alpha$ included sufficiently high values). A more general approach could involve setting a threshold for missingness reliance, such as 20%, and then maximizing AUC among models meeting this threshold. Another alternative is to maximize $\mathrm{AUC} – \gamma\cdot\hat{\rho}$, where $\gamma$ controls the trade-off between AUC and missingness reliance and needs to be set for the specific problem at hand.

We thank the reviewer for raising this point and we have added the following discussion to the revised paper: Alternative strategies for selecting $\alpha^*$ should consider the specific application needs. For instance, rather than focusing solely on the acceptable trade-off in predictive performance for achieving the lowest possible $\alpha$, we might prioritize limiting the number of missing features per individual. Relying on an average measure across the dataset may not provide sufficient insight. Additionally, domain knowledge can help define an acceptable level of missingness, but quantifying this threshold is often challenging and subjective.

We thank the reviewer for their insightful feedback and appreciate their recognition of our work addressing a practically important yet underexplored challenge in the healthcare domain—enhancing model reliability under test-time missingness. We address the specific comments below.

Re: MA learning in neural networks

Standard feed-forward neural networks typically rely on all input features to a non-zero degree unless the learning objective is regularized to yield sparse weights. Even then, the reliance is generally not contextual but remains the same for every input, regardless of the observed features. In principle, contextual reliance could be achieved by combining the MA penalty with attention modules that attend only to observed features. However, our focus is on tabular data, where neural networks tend to underperform. We instead target model classes like decision trees and linear models, which not only perform well on tabular data but also offer interpretability—an essential requirement in our primary application domain, healthcare. Exploring MA learning with attention modules for tabular representations is an interesting direction for future work.

Re: dataset sizes

We thank the reviewer for their observation and agree that the number of features in our datasets is relatively small (ranging from 16 to 42). However, we include a diverse set of datasets in terms of sample size—from a few hundred (e.g., Pharyngitis: 676, ADNI: 1,337, Breast Cancer 1,756, LIFE 2,864) to over 10,000 samples (e.g., FICO: 10,549, NHANES: 10,000). These dataset sizes are typical in the tabular ML literature and reflect realistic constraints in many application domains.

Re: advanced imputation methods as baselines

Thank you for sharing the two works. We agree that the methods could be compared to more advanced imputation approaches. However, our goal for the baselines was to include representative methods from different strategies for handling missing data: impute-then-predict, model-specific handling of missingness, and treating missingness as an informative signal. Each of these serves a slightly different purpose compared to the MA approach. We chose zero imputation and MICE because they are both widely studied in the literature and commonly used in practice, making them strong reference points for comparison. More advanced imputation methods are likely to yield similar reliance on missing values when combined with similar regressions or classifiers. Thus, the choice of imputation method is secondary in our study. To enable a fair comparison with our method, we also calculate the missingness reliance for the baselines. Since the baseline methods were not designed with the goal of learning to avoid missing values, but rather to provide accurate imputations, they may be incentivized to rely on well-imputed values—potentially leading to a higher missingness reliance metric.

We thank the reviewer for their thoughtful and valuable feedback, including their appreciation of our methods, empirical effectiveness, and the helpful discussion. We have addressed all suggestions to improve clarity and respond to specific comments below.

Re: motivation of avoiding missing values (w4)

MA learning is designed for prediction tasks with missing data where interpretability matters—clinical risk scores are a key example. Standard approaches (imputation, missingness indicators, default rules) often yield models that rely on unavailable inputs, as shown by their high missingness reliance in our experiments, which can reduce interpretability. For example, it may be difficult to justify that a physician or model should guess the value of a missing test (if using imputation) or that a patient’s risk of readmission should go down because the test result is missing (if using indicators or default rules in trees)—especially if the test is irrelevant based on the values of other features, such as a diagnosis. Though such tests may be predictive in other contexts, many learning tasks are underspecified, and models with similar accuracy may differ in missingness reliance. MA learning favors models that use information that is both predictive and available. Finally, we argue that MA learning and reasoning about missingness patterns go hand in hand. Increasing the penalty $\alpha$ encourages trees to split first on commonly observed features (e.g., demographics), and then on contextually available ones (e.g., tests for certain conditions). As discussed in Section 5 on ODDC rules, these observation patterns may be correlated with predictiveness, e.g., when the outcome reflects a diagnosis based on observed information.

Re: relevant literature (w2)

We thank the reviewer for highlighting the work by Beaulac & Rosenthal and have added it to the related work section. As noted, BEST suits ODDC settings with clear knowledge of the data-generating process, but defining gating variables is less straightforward with limited knowledge. In this sense, we view the MA framework as more general. A key difference between our methods is that MA trees do not explicitly split on the missingness mask. On the one hand, BEST can explicitly build trees using mask splits—often at the top—approaching the pattern submodel (PSM) of Mercaldo & Blume. On the other hand, MA trees cannot split based on non-MAR patterns, which may limit their expressiveness in such cases. However, avoiding explicit mask splits improves interpretability by preventing the tree structure from being dominated by missingness logic. Comparing MA-DT to BEST on the dataset from Beaulac and Rosenthal would be interesting, but the dataset does not appear to be publicly available.

Re: informative missingness (w3)

We introduce informative missingness in Section 3 but agree that Appendix C felt somewhat isolated. The reviewer is correct: when missingness is informative, adding missingness indicators is generally preferable to using an MA approach. For example, Beaulac & Rosenthal show that when missingness depends on the label, indicators can be predictive. In such cases, MA trees cannot leverage this information and may underperform, depending on the imputation. We appreciate the feedback and address this in the revised manuscript (Section 5.2).

Re: convergence (w5)

The reviewer raises an important question: if an optimal model $h^*$ exists, can our algorithm recover it with enough data? This depends on the model class, optimization, and data distribution. For example, if $h^*$ is a sparse logistic model, our L1 objective can likely recover it—similar to LASSO. The optimization in Section 4 supports Corollary 1: under ODDC rules, we can learn models with low missingness reliance and strong performance. We will add this discussion to the final paper.

Re: clarity (w1), notation, and regularization

• We agree it's important to clarify the relationship between $h^*$ and $\mathcal{H}$ and now state explicitly that $h^*(x) = \mathbb{E}[Y \mid x]$ may lie outside $\mathcal{H}$. This is benign for universal approximators (e.g., decision trees) but more significant for restricted classes like linear models.
• We agree that introducing $\sigma_{i,j}$ in Section 4.1 may be premature; we now assume equal feature contribution initially and introduce $\sigma_{i,j} \in {0,1}$ later in Section 4.3, where it supports the generalization to tree ensembles.
• We also acknowledge that starting with L0 regularization to penalize missing features would have been natural; however, we focused on L1 (LASSO) due to its compatibility with widely used implementations (e.g., scikit-learn, statsmodels). We will clarify this choice, along with a discussion of L0-based alternatives.

Lastly, we thank the reviewer for their valuable feedback and excellent suggestions, and for noting inconsistencies and typos, which have been corrected in the revision.

We thank the reviewer for their insightful feedback and appreciate their recognition of our well-motivated and novel MA framework. We have addressed all suggestions to improve clarity and respond to specific comments below.

Re: relevance of MA learning

We would like to direct the reviewer to Sec. 5, where we outline when MA learning is expected to perform well (Sec. 5.1) and when it may face challenges (Sec. 5.2)—a discussion also highlighted as a strength by Reviewer go58. MA learning is expected to work well when the missingness mask follows a clear structure, such as when ODDC rules govern the data-generating process. When there is no structure in the missingness patterns related to observed features, such as in an MCAR or a fully MNAR setting, we expect it to be harder for MA learning to achieve low missingness reliance with maintained accuracy. As the reviewer noted, our experiments show that MA learning outperforms baselines across datasets while greatly reducing reliance on missing values. Finally, MA learning may benefit from being combined with missingness indicators when missingness itself is highly informative—the reliance penalty will just make sure to keep such features to a minimum.

Re: guidelines for practitioners

To apply the MA framework in practice, we recommend starting by selecting a model class that aligns with the interpretability requirements of the application. Linear models and decision trees offer greater transparency, while ensembles may yield better performance with reduced interpretability. Once the model type is chosen, the trade-off between predictive performance and reliance on missing features can be adjusted through the regularization parameter $\alpha$. A higher $\alpha$ encourages the model to avoid relying on features that are often missing at test time, while a lower $\alpha$ prioritizes performance. The appropriate choice of $\alpha$ depends on how much missingness reliance is acceptable in the given context. We describe different strategies in the response to Reviewer eoMs. To evaluate this, practitioners can compute the reliance score ($\hat{\rho}$), which quantifies how frequently a model depends on missing features. Importantly, MA learning requires no assumptions about the underlying missingness mechanism (although the achievable accuracy-reliance tradeoff will be affected by its structure), making it a practical and robust approach across a wide range of real-world settings.

Re: qualitative analysis

We agree that qualitative analysis of MA trees is valuable. In Section 6.2, we refer to Figure 6 (appendix), which shows how MA-DT behavior varies with different values of the missingness regularization parameter $\alpha$ on the Life dataset. For instance, when $\alpha$ is tuned to balance accuracy and reliance on missing values, MA-DT reduces reliance on missing data by 33% without sacrificing accuracy compared to a standard decision tree. For the camera-ready version, we will consider including these figures in the main text and adding MA tree examples from the ADNI dataset. These illustrate how MA learning favors complete features over those with missingness, enhancing interpretability without compromising accuracy at test time.

Re: theoretical foundations and stacked models

Thank you for the insightful comment. We agree that strengthening the theoretical foundations is important. As for stacked models, we're unsure how they would address the challenges in our setting. Stacking typically involves training a meta-model on the outputs of several base models (e.g., decision trees, logistic regression). The reviewer suggests that these may be trained using complete data, but this is not generally available without imputation, which is what we try to avoid in the first place. Alternatively, we could train base models on different data subsets based on their missingness pattern, but this poses the same problem—it seems to us that the stacking question is orthogonal to reducing missingness reliance, but could be an alternative to boosting. We’re very open to considering this direction and would greatly appreciate any additional insights the reviewer can share on their suggestion.

Re: clarification on MA learning

Many model classes have large Rashomon sets—sets of near-optimal models that have similar predictive performance but differ on other properties. By adding the missingness reliance penalty, we prioritize models within this set that rely less on missing values. This can exploit nonlinear patterns (as with MA trees, where reliance is contextual) or linear ones (as with variable selection in MA-LASSO). We will include this in our Section 5 discussion.

Re: performance metrics

We agree that multiple metrics are important to capture different aspects of performance. Due to space constraints, we focused on AUROC and missingness reliance in the main text and will consider including additional metrics (e.g., F1 score) in the appendix.