Robust prediction under missingness shifts

Response to reviewer ExJF

We thank the reviewer for their time and constructive feedback! We would like to take the opportunity to address the points raised in the review. Since they were requested by multiple reviewers, a discussion of Zhou et al. (2023) as well as further details regarding NeuMISE are provided in a joint response above. \ \ \ **Related work** For a detailed discussion of [1] Zhou et al. (2023), please see our joint response. We further added [2] Jeanselme et al. (2022) to the related work section. \ \ \ **Assumptions on the noise** The assumptions in Equation (2) are primarily needed for the formulation of the Bayes predictor in Equation (3), which under those assumptions simplifies to the expected value of a deterministic function of $X_{obs}$ and $X_{mis}$. However, all results still hold for weaker assumptions about the noise. We regret that this was not clear in our original submission and have revised the section accordingly. **Action:** We added additional text following Equation (2) as well as an intermediate step to Equation (3) to make the role of our assumptions clearer. \ \ \ **Real-world data** We agree with the reviewer that real-world data that contains a missingness shift would further strengthen our results. Unfortunately, it is impossible to ascertain whether any observed shifts are ignorable or not. We therefore opted to focus on (semi-)simulated datasets with known missingness shifts, leaving the benchmarking of models on a wide range of real-world data for future work. \ \ \ **Comparison to MLP with zero imputation** We included an MLP model with zero imputation in our early experiments. However, results from those early experiments showed such a model to be a weaker baseline than an MLP with ICE imputation. In the interest of manuscript space, a direct comparison of ICE and MICE, as well as already strained computational resources, we opted to omit zero imputation in favour of ICE. We believe that the most meaningful comparison of NeuMICE is to NeuMiss and ICE. \ \ \ **Additional points** *Explicit assumption of no covariate/label shift:* We added an explicit statement of $p(Y, X) = q(Y, X)$ to the beginning of Section 4. *Extended discussion of results:* We further agree with the reviewer regarding Appendix F and have added some more interpretation to the results section, particularly related to the decreased performance of NeuMISE. \ \ \ We believe that the above changes have strengthened the overall quality and contribution of the paper. We would like to kindly request that you reconsider the score assigned to our paper in light of the revisions. We believe the changes made have effectively addressed the issues you raised, and we hope this will be reflected in an updated evaluation. If you have any further questions or if there are specific areas you would like us to clarify, please do not hesitate to let us know.

Response to reviewer ntEa

We thank the reviewer for their time and constructive feedback. We would like to take the opportunity to address the points raised in the review. Since they were requested by multiple reviewers, a discussion of Zhou et al. (2023) as well as further details regarding NeuMISE are provided in a joint response above. \ \ \ **Experimental detail in the main text** Due to the rather stringent space limitations, we limited the description of experiments to the minimum judged necessary to understand the results, providing the details required for full reproducibility in the appendix. We acknowledge that the final version may have allocated too little space to experimental description. We moved some of the experimental details back into the main text. In particular, we define exactly what we mean by “high” and “low” correlation in our experiments. We hope that we were able to strike the right balance between concision and detail. **Action:** We revised Section 7.1 to include additional details on the experimental setup. \ \ \ **Clearer explanations of the role of Y** The Bayes predictor $\tilde f_m = E[Y \mid X_{obs}, M=m]$ is still well-defined for $Y$-dependent missingness. Furthermore, in the absence of a missingness shift, the Bayes predictor in source and target environment remain the same, i.e., $\tilde f_m = \tilde g_m$. Any method that successfully estimates $\tilde f_m$ without using $Y$ (i.e., omits $Y$) will therefore have unchanged performance in the target environment. However, the circumstances change, once we adjust for $Y$ by including it in an impute-then-regress model: we learn an estimator $\int \hat f (X_{mis}, X_{obs})p(X_{mis} \mid X_{obs}, Y)dX_{mis}$ that gives an unbiased estimate $\hat f$ of the complete data model $f^\star$ but requires knowledge about $Y$ — the outcome we want to predict. This is obviously problematic. To overcome this, others have proposed to learn $\hat f$ as above through the use of $p(X_{mis} \mid X_{obs}, Y)$ but then replace the imputation with $p(X_{mis} \mid X_{obs})$ in the target environment. We argue here that this does not estimate the Bayes predictor $\tilde g_m$. Please also see our reply to reviewer tBPA for some further explanations. **Action:** We revised the text in Section 6.1 (formerly 5.1) to clarify our argument. \ \ \ **Informative missingness vs. “unbiased” estimation** One of the primary motivations for this research was indeed a common claim that unbiased estimation is preferable for clinical prediction modelling, usually implying that it leads to more robust models. However, we find that Theorem 1 does not provide an a priori reason why unbiased methods should be preferred for robust prediction. Instead, robust prediction under ignorable missingness shifts depends on the precise estimation of the Bayes predictors in the source environment, which may depend on both the method (e.g., impute-then-regress vs. end-to-end learning) and the data (e.g., levels of missingness). **Action:** We followed the advice of the reviewer and added a Section 4.1 that introduces the above argument. We further revised our section on “Learning under missingness shifts” to discuss certain conditions that may influence the estimation of the Bayes predictor. \ \ \ **Additive noise and its independence of missing covariates** The reviewer is correct that neither the additive nature of the noise nor the assumption on its conditional mean are strictly necessary. We opted for this formulation, as it leads to a simple and clear expression of the optimal predictor given complete data in the form of $f^\star(X)$. However, our results remain valid for binary outcomes as well as non-additive noise that depends on the entire $X$. **Action:** We added additional text following Equation (2) in order to clarify the role of our assumptions. We also added Appendix A.2 to discuss the generality of our results. \ \ \ We believe that the above changes have strengthened the overall quality and contribution of the paper. We would like to kindly request that you reconsider the score assigned to our paper in light of the revisions. We believe the changes made have effectively addressed the issues you raised, and we hope this will be reflected in an updated evaluation. If you have any further questions or if there are specific areas you would like us to clarify, please do not hesitate to let us know.

Response to reviewer P8Mn

We thank the reviewer for their time and constructive feedback. We would like to take the opportunity to address the points raised in the review. Since they were requested by multiple reviewers, a discussion of Zhou et al. (2023) as well as further details regarding NeuMISE are provided in a joint response above. \ \ \ **Assumptions about noise** In our formulation, the noise may indeed depend on the missingness. That is, missingness may actively (and even causally) influence Y. Consider for example a situation in which a biomarker A is used by doctors to decide whether another test B should be performed, which requires referral to a specialist laboratory. If the patient is referred, the laboratory also measures Y. If the patient isn’t referred, the doctor themselves measures Y. If the laboratory has a more precise machine than the doctor then $Var(\epsilon \mid A, M_B=0) < Var(\epsilon \mid A, M_B=1)$ but $E(\epsilon \mid A, M) = 0$ may still be true. Our theory would allow for such a scenario. We would further like to point out that the assumptions in Equation (2) are primarily needed for the formulation of the Bayes predictor in Equation (3), which under those assumptions simplifies to the expected value of a deterministic function of $X_{obs}$ and $X_{mis}$. Please also refer to our reply to reviewer ntEa regarding the possibility of non-additive noise or noise that depends on $X_{mis}$. **Action:** We added additional text following Equation (2) as well as an intermediate step to Equation (3) to clarify the role of our assumptions. \ \ \ **Ignorability of Y-dependent missingness in the source environment** Since Y is observed at training time, Y-dependent missingness in the source environment can be MAR if the imputation model conditions on $Y$. The source missingness is thus ignorable. If a shift occurs through which the target missingness no longer depends on Y and only depends $X_{obs}$, then the target missingness is also ignorable. By Definition 2, the missingness shift is therefore ignorable. The main difficulty then lies in how to derive an $\tilde{f}^\star_m$ that does not require $Y$ in order to be used in the target environment (since we will not have $Y$ when we want to predict). Mathematically speaking, we need to find a way to derive $q(X_{mis} \mid X_{obs})$ from $p(X_{mis} \mid X_{obs}, Y)$. We added a sketch of how this might be achieved to the appendix, although we would like to stress that we haven’t been able to verify it empirically. **Action:** We revised Section 6.2 (formerly 5.2) of our paper to make this argument clearer. We further outlined a simple two-step impute-then-regress approach in the Appendix that may be used to derive a Bayes estimator for the target environment from the source data. \ \ \ **Additional points** *Real-world examples:* Our original introduction contained two examples of real-world missingness shifts: 1) a reduction in a test’s costs may increase its use, making it more readily available for a larger share of patients and 2) the deployment of the model in clinical practice may change how doctors collect data, for example by more commonly collecting the variables that are used by the model. We’ve added a third example: a change in clinical guidelines that govern how and when doctors ought to perform clinical examinations. Each of these cases changes the missingness in the data and represents an example of real-world missingness shift. All examples focus on healthcare, as this is our main domain of expertise. We believe that the importance of medical applications — and the role that AI can play within it — serve as sufficient motivation but we are certain similar examples could be found for other domains. *Focus on covariate shift:* We added a clarification to the abstract and introduction, highlighting that our work focuses on missingness in covariates. \ \ \ We believe that the above changes have strengthened the overall quality and contribution of the paper. We would like to kindly request that you reconsider the score assigned to our paper in light of the revisions. We believe the changes made have effectively addressed the issues you raised, and we hope this will be reflected in an updated evaluation. If you have any further questions or if there are specific areas you would like us to clarify, please do not hesitate to let us know.

Response to reviewer tBPA

We thank the reviewer for their time and constructive feedback. We would like to take the opportunity to address the points raised in the review. Since it was requested by multiple reviewers, further details regarding NeuMISE are provided in a joint response above. \ \ \ **Too much detail in descriptions** We aimed for maximum clarity in our description of the problem setting and background. We appreciate the reviewer’s feedback and tried to remove some of the detail in favour of other, more insightful, sections. **Action:** We removed some of the detail in Section 3. However, due to additional content requested by other reviewers, the section is only marginally shorter in the revised manuscript. \ \ \ **Adjusting/omitting Y** In Section 5.1, we used “adjusting Y” to refer to the inclusion of the outcome in the imputation model $p(X_{mis} \mid X_{obs}, Y)$. This is standard practice in statistical inference under Y-dependent missingness, as it results in a MAR setting and allows for unbiased estimation. In contrast, by “omitting Y” we mean an imputation model $p(X_{mis} \mid X_{obs})$ that does not include $Y$ in the conditioning set. We consider an estimator to be stable if it works equally well in the source and target environment. In the absence of a missingness change, $\tilde f_m = \tilde g_m$ is trivially true and any estimator that applies the function it learned from the source distribution to the target distribution will be stable. However, a method that tries to adjust for $Y$ during training will learn $\int \hat{f}(X_{mis}, X_{obs})p(X_{mis} \mid X_{obs}, Y)dX_{mis}$ but — since $Y$ cannot be observed in the target environment — would need to apply $\int \hat{f}(X_{mis}, X_{obs})q(X_{mis} \mid X_{obs})dX_{mis}$ at deployment. These two functions differ in general. The unobservable nature of $Y$ therefore introduces an “artificial” missingness shift, rendering a setting that is MAR at training time MNAR at deployment time. **Action:** We revised the text in Section 6.1 (formerly 5.1) to clarify our argument. \ \ \ We believe that the above changes have strengthened the overall quality and contribution of the paper. We would like to kindly request that you reconsider the score assigned to our paper in light of the revisions. We believe the changes made have effectively addressed the issues you raised, and we hope this will be reflected in an updated evaluation. If you have any further questions or if there are specific areas you would like us to clarify, please do not hesitate to let us know.

Continued joint response

**Motivation for NeuMISE** The reviewers have requested further detail on NeuMISE and why it might be more robust to missingness shifts. NeuMISE was inspired by the observation that pattern submodels like NeuMiss have greater difficulties extrapolating to previously unseen missingness patterns, especially if they are far away from the training data. In particular, when considering fully observed data, NeuMiss would still embed the fully observed vector $X$ as $\Sigma^{-1} X$, where $\Sigma^{-1}$ is the inverted observed-data covariance. In the case of ignorable missingness, however, it would be sufficient to forward the unchanged $X$. NeuMiss thus adds unnecessary noise to the input. By modifying the mask, we change this behaviour: rather than embedding with $\Sigma_{obs}^{-1}X_{obs}$ like NeuMiss, we impute with $\Sigma_{mis,obs}\Sigma_{obs}^{-1}X_{obs}$. Importantly, the observed $X_{obs}$ remain unchanged and are forwarded to the prediction layers as is. We further would like to point out that we do not intend to claim that NeuMISE remains robust across all possible missingness shifts, including those that are non-ignorable. Reviewer P8Mn is perfectly right in stating that “generalizing well to non-ignorable missingness shift should only be possible if the model is in some sense robust to a variety of non-ignorable missingness shift and I would presume that such a model may trade off some performance on the source data for better generalization across target environments”. In fact, our experiments on Y-dependent missingness clearly show this, and one would be much better off using MICE+Y. This is directly related to the trade-off mentioned by reviewer P8Mn, as MICE+Y performs much worse in the source environment. However, we observed empirically that NeuMISE consistently performed more robustly or on par with NeuMiss in both shifted and unshifted environments. We have revised some of the text to make this clearer and avoid false claims of robustness across any missingness shift. **Action:** We revised section 5 (formerly section 6), and in particular section 5.1 (formerly section 6.2) to provide further detail on NeuMISE and why modifying the mask may be well-motivated. **Changelog:** - We revised our abstract and introduction to adjust our claims regarding novelty. We also clarified our focus on missingness in covariates and added another real-world example of missingness shifts. - We added Zhou et al. (2023) and Jeanselme et al. (2022) to the related work in Section 2. - We explain the impact of our assumptions on the noise in Section 3. - We updated our description of Theorem 1 in Section 4, prominently referring to Zhou et al. - We added a discussion on the importance of unbiased estimation for robust prediction under missingness shifts to Section 4. - We moved former Section 4.2 “Note on changes in Bayes risk” to the Appendix in order to save space. - We switched former Sections 5 and 6 and revised them in response to the reviewers’ comments. - We slightly extended the discussion of the empirical results. - We changed the description of the missingness mechanism from Monotone MAR to Non-monotone MAR, as this was falsely labelled in the original submission. - We added further explanations and details to the appendix. We believe that the above changes have strengthened the overall quality and contribution of the paper. We would like to kindly request that you reconsider the score assigned to our paper in light of the revisions. We believe the changes made have effectively addressed the issues you raised, and we hope this will be reflected in an updated evaluation. If you have any further questions or if there are specific areas you would like us to clarify, please do not hesitate to let us know and we will still try to provide an answer before the deadline.