Paper: Out-of-Variable Generalisation for Discriminative Models (ICLR-2024)

Reviewer comment (Theory):
It would be helpful to include more intuitive discussion throughout the paper providing more analysis on the sections. For example, more discussion on the assumptions of the settings/theorems would be helpful, and it's not clear exactly under what assumptions the proposed predictor is appropriate.

Author response:
> It would be helpful to include more intuitive discussion throughout the paper providing more analysis on the sections. For example, more discussion on the assumptions of the settings/theorems would be helpful, and it's not clear exactly under what assumptions the proposed predictor is appropriate.

We thank the reviewer for the suggestion and have included 
* a bullet list of assumptions on theorems and identifiability of the proposed predictors (Appendix E.2)
* a discussion on violations to each of the assumptions from causal graph, to function class and to noise distributions (Appendix E.2.1). 
* On experimental analysis side, in addition to experiments on real world dataset (Appendix D.4), we conducted additional experiments on the robustness of our approaches when noise level changes (Appendix D.5), and when the distribution of noise is heavy tailed (Appendix D.6) 

Overall, we thank you for taking time, which we carefully take into account your feedback and included both experimental results and theoretical discussions in the updated version. We hope that our response has addressed all your questions and concerns. We kindly ask you to let us know if you have any remaining criticism, and - if we have answered your questions - to consider reevaluating your score.

Reviewer comment (Experiments):
For future work, is there a more complicated/realistic dataset to validate the algorithm?

Author response:
> For future work, is there a more complicated/realistic dataset to validate the algorithm?

We thank the reviewer for the suggestion and included results on real-world dataset in Table 2 (Appendix D.4).

Reviewer comment (Theory):
Theorem 3 connects all moments of the residual distribution to the partial derivatives with respect to the unique variable of the target environment. If additional moments were to be calculated as part of the proposed algorithm, would it improve results (for the general function case)?

Author response:
> Theorem 3 connects all moments of the residual distribution to the partial derivatives with respect to the unique variable of the target environment. If additional moments were to be calculated as part of the proposed algorithm, would it improve results (for the general function case)?

Thank you for the question, yes incorporating higher-order moments would improve the result as it allows us to estimate general function beyond first-order Taylor approximations. The higher moments provide a space of system of equations that can help to identify the first order partial derivatives and higher order partial derivatives, which assist in estimating a more complex function class.

Reviewer comment (Experiments):
In general, since the paper's main claim is that in the real world, it is likely to encounter both aspects of OOD and OOV - How simple is it to combine state-of-the-art OOD methods with the proposed approach? I cannot imagine at the moment a straightforward way to do that.

Author response:
> In general, since the paper's main claim is that in the real world, it is likely to encounter both aspects of OOD and OOV - How simple is it to combine state-of-the-art OOD methods with the proposed approach? I cannot imagine at the moment a straightforward way to do that.

The reviewer asks re insights on combining both OOD and OOV approaches: we think sparse mechanism shift (SMS) might be a useful notion to characterize distributions shifts as SMS and our current approach both uses causal framework to assist analysis. It assumes the distribution shifts are caused by sparse causal mechanism shifts. We think depending on different variables to which the distribution shift happens would call for a different approach. For example, say, a functional relationship on $Y$ is $\alpha X_1 + X_2 + X_2X_3$. Changing $\alpha$ between the source and target environment does not affect our proposed approach on OOV generalization. However, if $\alpha$ is placed on $X_2X_3$ term, a more sophisticated method need to be developed. We do not have broader results yet, and we feel that the paper is already useful to the community as a starting point to tackle real-world problems that exhibit both OOD and OOV.

Reviewer comment (Experiments):
The main weakness is the applicability of the method. The authors only showed results for proof-of-concept, not for real-world usage.

Author response:
We thank the reviewer for the suggestion and here include results on a real-world dataset in Table 2 (Appendix D.4). We observe our method outperforms the other baselines. 

Overall, we thank the reviewer for taking the time and pushing us to be more explicit on our method’s robustness and applicability. Following your suggestions, we have included 3 additional experiments: on real-world dataset (Table 2), on robustness of changing noise level (Table 3), on robustness of heavy tailed distribution (Table 4). To further offer clarity in understanding the assumptions of our method, we provide a bullet list of assumptions on our theorems in Appendix E.2 and discusses robustness of our method to violation of each assumptions (Appendix E.2.1)

Reviewer comment (Novelty):
It is not yet clear what realistic problem can be well modeled by OOV generalization.

Author response:
We think OOV generalization is a general ability present in both human and animal’s navigation towards Nature. For example, AI in medicine is a relevant area where we often face strong limitations in guaranteeing dataset consistency. To begin with, patients have unique circumstances, and some diseases/symptoms are rare. More generally, the world is full of medical datasets with varying variable sets - diagnostic measurements greatly vary across patients (serum, various forms of imaging, genomics, proteomics, immunoassays, etc.). A good human doctor can generalize across patients even if the measured variables are not identical. 

To be more concrete, we present a realistic problem in Section 3.1: consider two medical labs collecting different sets of variables. Lab A collects $X_{1}= $ lifestyle factors and $X_2 = $ blood test; Lab B, in addition to $X_2$, collects $X_{3} = $ genomics. Lab A is hospital-based and can measure diseases $Y$, whereas B is a research lab and hope to utilize existing Lab A’s data. The OOV problem asks: given a model trained to predict $Y$ on Lab A’s data, how should Lab B use this model for its own dataset that differs in the set of input variables? 


Overall, there has been much talk about the potential of AI in medicine since machines can learn from millions of patients, whereas a doctor may only see a few thousand during their lifetime. In practice, data scientists end up using imputation or simply ignore rarely-measured features. To fully realize the potential of AI in medicine, we believe one needs to solve the OOV problem one way or another.

Reviewer comment (Theory):
It seems OOV fits very well the frame of missing-not-at-random and covariate-dependent missingness. Could the authors comment on that?

Author response:
We thank the reviewer for the question. Missing-not-at-random and covariate-dependent missingness refer to a scenario where whether a variable is observed or not contains information about other covariates, or about certain other properties of the data point. One can thus hope to exploit this assumption to recover more information about the dataset. We agree that such problems are related to OOV, though think OOV addresses a larger set of problems. For example, the setting of our paper (Fig. 1a) studies when covariates are independent. Such assumption is even more challenging as it provides less information on the remaining dataset. We surprisingly find that even in such challenging scenario, one can enable OOV generalization for certain cases.

Reviewer comment (Writing):
Theorem 2 is slightly confusing for me at first glance because I thought PA_Y by definition includes all parents of Y (so x1, x2, x3 in the example) and not just those in the target environment (x2, x3). It may be helpful to clarify.

Author response:
Thank you for your feedback on improving the paper. We have incorporated an explanation on the difference of Theorem 2 setting with the original setup in the paragraph between section 3.3 and section 3.3.1 in the updated version.

Reviewer comment (Evaluation):
How does the method fair with the oracle as the magnitude of the noise increases?

Author response:
Thank you for the question, we performed an additional systematic analysis similar in Table 1 with increase in noise standard deviation from 0.01 to 1 in an interval of 0.2. We averaged results over 5 repeated runs and observe our method are robust to increasing noise level and consistently outperforms remaining benchmarks despite noise increases. Table 3 (Appendix D.5) in the updated pdf shows the results.

Reviewer comment (Theory):
What if the noise is not gaussian but more heavy tailed?

Author response:
Thank you for your question, we included an additional systemic noise when noise is heavy tailed. Specifically, when noise follows a log-normal distribution with mean 0 and sigma 0.5. We repeated the experiment over 5 runs and averaged over a hyperparameter sweep. Table 4 (Appendix D.6) shows the results. We observe as expected by our theorems (Theorem 3), the skew of the noise pollutes the skew estimation of the residual distribution, which deteriorates performance.

Reviewer comment (Experiments):
Does the performance degrade or improve with increasing number of variables?

Author response:
Thank you for your question. When the number of variables increases in the source environment, we expect the performance remains the same as the base network $f_s$’s accuracy is indifferent to the number of variables given sufficient sample size; when the number of variables increases in the space of missing variables, we would expect our method would still be able to transfer under OOV setting. This is due to one can leverage multiple moments from the residual distribution to create a system of equations and solve multiple unknowns. We provide an additional theoretical analysis on a toy example facing two unobserved variables from the source environment in Appendix C.6.

Reviewer comment (Presentation):
Figure 3: I don’t quite understand the figure. It would be helpful to define OOV loss, be explicit about the number of samples on the y-axis being (x2, x3, y) or (x1, x2, y) or something else. I also don’t understand why relative loss is zero means the method is on par with the oracle predictor. Why not just show how the fine-tuning error compares with oracle training, which seems easier to interpret? Anyway, I am overall a bit confused about the figure, so my questions may not make sense.

Author response:
We apologize for the confusion, the number of samples on the x-axis refer to (x2, x3, y), and the relative loss is calculated as $log(loss_{pred}/loss_{oracle})$, the log ratio of predictor’s loss divided by the oracle’s loss. If the predictor achieves the same loss as the oracle loss then the ratio would be 1 and $log(1) = 0$. The purpose of the experiment is to show what is the number of samples needed in the target environment, at which we would prefer the MomentLearn over the joint predictor. Figure 3 shows that up to ~100 samples in the target environment, MomentLearn outperforms the joint predictor.

Reviewer comment (Writing):
Referring to Figure 1, in the first paragraph in page 3, the claim "it would seem all but impossible...(orange box)" could be better explained.

Author response:
We thank the reviewer for the suggestions to help us clarify Figure 1. We have incorporated the proposed changes in the updated version, by explaining the claim further and changing the border of the orange box.  We also thank the reviewer for their detailed observations on typos and have corrected them in the updated version.

Reviewer comment (Presentation):
In Figure 1, it is unclear whether "With $Y$ not observed in the target domain" is an assumption made or is somehow indicated in the diagram or earlier in the paper. Eventually I realized that it's an assumption made, but the illustration Figure 1a alone isn't enough to show this assumption. This ambiguity may clear for some or compound for some later in Section 3.

Author response:
We thank the reviewer for the suggestions to help us clarify Figure 1. We have incorporated the proposed changes in the updated version, by explaining the claim further and changing the border of the orange box.  We also thank the reviewer for their detailed observations on typos and have corrected them in the updated version.

Reviewer comment (Writing):
The abstract states "merely considering differences in data distributions is inadequate for fully capturing differences between learning environments." Doesn't out-of-variable technically fall under out-of-distribution, so shouldn't this be adequate? Perhaps more specificity is needed here.

Author response:
Here we refer to settings exhibit in-distribution if the environments share the same data-generating process. For example, in Fig 1a, though the environments share the same data generating process (in-distribution) but observed different sets of variables (out-of-variable). To give an example on OOD but not OOV, consider vaccine data from different countries: they have the same variables [Age, is_vaccined], but coming from different distributions due to country differences. Thank you for your question and we will make this point clearer in the next version.

Reviewer comment (Writing):
On page 2, should "modal" be "model?"

Author response:
We thank the reviewer for the suggestions to help us clarify Figure 1. We have incorporated the proposed changes in the updated version, by explaining the claim further and changing the border of the orange box.  We also thank the reviewer for their detailed observations on typos and have corrected them in the updated version.

Reviewer comment (Writing):
On page 6, do you mean "parentheses" instead of "brackets" between Eq (9) and Eq (10)?

Author response:
We thank the reviewer for the suggestions to help us clarify Figure 1. We have incorporated the proposed changes in the updated version, by explaining the claim further and changing the border of the orange box.  We also thank the reviewer for their detailed observations on typos and have corrected them in the updated version.

Reviewer comment (Theory):
Why is the joint predictor considered an oracle predictor if MomentLearn outperforms it?

Author response:
We apologize for the confusion. In Table 1, we refer to the joint predictor as oracle predictor in cases where there are enough data samples in the target environment such that regressing on the joint variable will lead to a near optimal predictor. In Figure 3, we showed that when the sample size is small (~100 samples) in the target environment, joint predictor can lead to estimation error. MomentLearn predictor thus is able to outperform it.

Reviewer comment (Evaluation):
Could you explain why MomentLearn is reliably more sample efficient than the oracle predictor for "few"-shot prediction?

Author response:
MomentLearn is more sample efficient than the joint predictor in “few”-shot prediction because with small sample size, the joint predictor would lead to estimation error, whereas MomentLearn leverages the information observed in the source environment (due to large sample size in source) are able to mitigate the problem of small sample size in the target environment.