Paper: Robust prediction under missingness shifts (ICLR-2024)

Reviewer comment (Theory):
Equation (2) relies on the assumption: $\mathbb{E}[\epsilon \mid X_{obs}, M] = 0$. Authors should detail this assumption, its meaning and its real-world relevance. Particularly, does it not imply some independence of Y upon the missing data?

Author response:
**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.

Reviewer comment (Theory):
In Section 4, there is an implicit assumption of no covariate or concept shift. Making this explicit would enhance clarity for readers.

Author response:
**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.

Reviewer comment (Evaluation):
It would be beneficial to delve further into the discussion of Appendix F, particularly in addressing the unclear aspects of why the model performs less effectively when uncertainty increases.

Author response:
**Additional points**

*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.

Reviewer comment (Experiments):
While the empirical results employ real-world covariates, the analysis lacks a study of real-world missingness shifts. Considering a dataset with changing missingness processes would strengthen the experimental analysis.

Author response:
**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.

Reviewer comment (Experiments):
Comparing the proposed approach with an end-to-end neural network that uses zero imputed data and mask as input seems a natural comparison to demonstrate the superiority of the proposed NeuMICE.

Author response:
**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.

Reviewer comment (Experiments):
Important aspects of the experiments are not presented clearly enough or with enough detail in the main text. For example, it is not explained what “low correlation” and “high correlation” corresponds to in the experiments.

Author response:
**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.

Reviewer comment (Writing):
Section 5 on the role of Y is interesting, but is not presented particularly clearly. In particular, please elaborate on how adjusting for Y in a source domain can induce missingness shift, but omitting Y results in a stable estimator.

Author response:
**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.

Reviewer comment (Writing):
The discussion focuses strongly on comparing methods that leverage informative missingness vs. “unbiased” estimators that do not. This seems to be a critical point for the paper overall (e.g. related to the second of the three contributions listed in the introduction section), but it did not come through clearly in the writing. Furthermore, it is not clear how (or if) this was evaluated in the experiments. This could perhaps be improved with additional exposition earlier in the paper that sets up this argument more clearly with specific hypotheses to be evaluated.

Author response:
**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.

Reviewer comment (Theory):
Can the results be generalized to binary outcomes? Naively, it seems like the additive noise model limits the direct applicability of this framework to binary outcomes.

Author response:
**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.

Reviewer comment (Theory):
Is the assumption that the error term in equation (2) depends only on the observed X and M limiting the generalizability of this work? Which aspects of the results would no longer hold if the error term were to depend on the full X?

Author response:
**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.

Reviewer comment (Novelty):
The main weakness of this work is that it does not cite or discuss its connection to [1], which is another work that studies robustness to missingness shift. This paper describes that when missing data indicators are available, domain adaptation under missingness shift reduces to a covariate shift problem. This finding seems to be related to one of the central contributions of this paper, which is that the optimal predictor remains unchanged if missingness only depends on observables in both the training and test environment.

Author response:
Dear reviewer, 

As the discussion period is drawing to a close in a few hours, we would like to thank you once again for your thoughtful consideration. We greatly appreciate the time and effort you have already dedicated to evaluating our work. 

If you could spare a moment to revisit our rebuttal, we would be immensely grateful. We carefully addressed your suggestions and we believe that these adjustments have significantly strengthened the overall quality of our paper. In particular, we **clarified the connection of our work to Zhou et al. (2023)** and **improved our discussion of NeuMISE**. 

In light of the above changes, we would greatly appreciate your feedback so your reviews and score may reflect the updated material.

Kind regards, 

The authors

Reviewer comment (Evaluation):
It’s not clear to me what advantage NeuMISE (the authors’ proposed method) has compared to the existing baseline in the presence of non-ignorable missingness shift. While the authors have some empirical results that NeuMISE performs outperforms other methods in the presence of non-ignorable missingness shift, I’m skeptical that such a result holds in general. To my understanding, 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. Is that the type of result that we observe for NeuMISE? What is the reason that NeuMISE is more robust? Furthermore, what benefit does NeuMISE offer compared to existing baselines.

Author response:
Dear reviewer, 

As the discussion period is drawing to a close in a few hours, we would like to thank you once again for your thoughtful consideration. We greatly appreciate the time and effort you have already dedicated to evaluating our work. 

If you could spare a moment to revisit our rebuttal, we would be immensely grateful. We carefully addressed your suggestions and we believe that these adjustments have significantly strengthened the overall quality of our paper. In particular, we **clarified the connection of our work to Zhou et al. (2023)** and **improved our discussion of NeuMISE**. 

In light of the above changes, we would greatly appreciate your feedback so your reviews and score may reflect the updated material.

Kind regards, 

The authors

Reviewer comment (Writing):
It would be helpful to add a few concrete examples where missingness shifts occur in the real-world to motivate the research.

Author response:
*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.

Reviewer comment (Theory):
It would be helpful if the authors add a line after Equation 2 that explains what the assumption $\mathbb{E}[ \epsilon \mid X_{obs}, M] = 0$ means concretely – i.e., to what extent can the noise $\epsilon$ depend on the missingness $M$? The current presentation does not require $\epsilon$ to be independent of $M$ – is that the desired interpretation? To the best of my understanding, in the current presentation, the variance of the noise $\epsilon$ could depend on the missingness mechanism.

Author response:
**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.

Reviewer comment (Theory):
Could the authors explain why the following is true: ``If $Y$ only influences missingness in the source environment, shifts may still be ignorable.”

Author response:
**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.

Reviewer comment (Novelty):
The paper uses quite a lot space to discuss the missingness shift. Although such descriptions are complete and clear, it seems to be relatively elementary and do not provide enough new intellectual insights. Under ignorable condition, Theorem 1 "equivalence" is also straightforward and hence is not surprising, at least to me.

Author response:
**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.

Reviewer comment (Writing):
Last few sentences in Section 5.1 confuse me. What is "adjusting Y", "omitting Y" and definition of "stable estimator"?

Author response:
**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.

Reviewer comment (Evaluation):
Section 6 is rather short. It should be expanded to explain why NeuMISE is more effective from a deeper viewpoint.

Author response:
Dear reviewer, 

As the discussion period is drawing to a close in a few hours, we would like to thank you once again for your thoughtful consideration. We greatly appreciate the time and effort you have already dedicated to evaluating our work. 

If you could spare a moment to revisit our rebuttal, we would be immensely grateful. We carefully addressed your suggestions and we believe that these adjustments have significantly strengthened the overall quality of our paper. In particular, we **expanded our discussion of NeuMISE**, discussing its motivation and principles in more depth. 

In light of the above changes, we would greatly appreciate your feedback so your reviews and score may reflect the updated material.

Kind regards, 

The authors