Domain constraints improve risk prediction when outcome data is missing

We thank the reviewer for their positive review, and are glad that they found that our paper is well-written and easy to follow, that our model is elegant and simple, that the theoretical results justify why the constraints help parameter estimation, and that the experiments are well thought-out and compelling. We will now address the reviewer’s questions and suggestions for improvement.

The reviewer first asks how the model generalizes to more complex inputs. To show our constraints are useful with more complex inputs, we ran two additional synthetic experiments. First, we demonstrated applicability to higher-dimensional features (Figure S4). Even after quadrupling the number of features (increasing runtime by a factor of three), both constraints still improve precision and accuracy. Secondly, we evaluate a more complex model with pairwise nonlinear interactions between features (Figure S5). Again both constraints generally improve precision and accuracy. We note our implementation relies on MCMC which is known to be less scalable than approaches like variational inference [1]. However, our approach does not intrinsically rely on MCMC (we are pursuing a follow-up paper investigating alternate approaches to fitting our models).

Then the reviewer states, “A trivial extension would be a partially linear model where some features are captured by a linear component and the rest are captured by a neural net.” We thank the reviewer for this comment and agree that this is a natural direction for future work: for example, the neural net could make a prediction from a mammogram while the linear component incorporates clinical or demographic features. Another option is to use a purely linear model, but include features which are precomputed functions of more complex inputs. Indeed, our current model does this through the “genetic risk score” feature, capturing each patient's polygenic risk score which is a function of many genetic variants.

Finally, the reviewer suggests that we reword the sentence: “Throughout, we generally refer to $Y_i$ as a binary indicator, but our framework extends to non-binary $Y_i$, and we derive our theoretical results in this setting” to make it more clear what “this setting” is. We thank the reviewer for this comment and will edit our manuscript to make this more clear. To clarify, we derive our theoretical results with the Heckman correction model which assumes a continuous $Y_i$. Some examples of non-binary $Y_i$ are tumor size or cancer stage.

[1] Martin J Wainwright and Michael I Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(1–2):1–305, 2008.

We thank the reviewer for their positive review that our draft was well-written. We now address their comments.

The reviewer first comments that our model may be overly simplistic. To show our constraints are useful with more complex models, we ran two additional synthetic experiments. First, we demonstrated applicability to higher-dimensional features (Figure S4). Even after quadrupling the number of features (increasing runtime by a factor of three), both constraints still improve precision and accuracy. Secondly, we evaluate a more complex model with pairwise nonlinear interactions between features (Figure S5). Again both constraints generally improve precision and accuracy. We note our implementation relies on MCMC which is known to be less scalable than approaches like variational inference [1]. However, our approach does not intrinsically rely on MCMC (we are pursuing a follow-up paper investigating alternate approaches to fitting our models).

The reviewer asks for specific predictive metrics for the results in section 5.2. To address this we report the AUC. (We report AUC instead of F1 score to allow comparison to past work.) The AUC amongst the tested population is 0.63 and amongst the untested population that attended a followup visit is 0.63. These AUCs are similar to past predictions which use similar feature sets [2] (and could be improved by using richer feature sets, though that is not the focus of this work). For instance, the Tyrer-Cuzick [3] and Gail [4] models achieved AUCs of 0.62 and 0.59.

The reviewer asks how our model compares to models trained solely on the tested population. We fit a logistic regression model only on the tested population. (To confirm that non-linear methods did not yield an improvement, we also fit random forest and gradient boosted classifiers; this yielded very similar results to the logistic regression model.) These baselines suffer from the same issue: they learn that cancer risk first increases and then decreases with age which, as discussed in section 5.4, is implausible in light of prior research in oncology (Figure S6). For the tested population, our model achieves similar AUCs to these other models.

The reviewer also asks how our model compares to treating the untested group as negative; this is equivalent to predicting $p(T=1, Y=1|X)$, an approach used in prior work [5,6]. Though this baseline no longer learns an implausible age trend, it underperforms our model in AUC (for both the tested and untested population AUC is 0.60 vs. 0.63 for our model) and quintile ratio (quintile ratio on the tested population is 2.4 vs. 3.3 for our model; quintile ratio on the untested population is 2.5 for both models). This baseline is a special case of our model with an implausibly low prevalence constraint $p(Y=1|T=0) = 0$. In light of this, it makes sense that this baseline learns a plausible age trend, but underperforms our model overall.

Though the reviewer did not explicitly request this, we also compare to two other selective labels baselines. First, we predict hard pseudo labels [7]: i.e., we train a classifier on the tested population and use its outputs as pseudo labels for the untested population. Due to the low prevalence of cancer, the pseudo labels are all $Y=0$, so this model is equivalent to treating the untested group as negative and similarly underperforms our model. Second, we use inverse propensity weighting [8]: i.e., we train a classifier on the tested population but reweight each sample by $\frac{1}{P(T=1|X)}$. This baseline learns the implausible age trend because it merely reweights the sample, without encoding that the untested patients are less likely to have cancer via a prevalence constraint. All of the baseline results have been added to Appendix E.1.

The reviewer asks, “How does the model perform on the older population?” To address this we run our model on the entire population. We find that performance when using the full cohort is better than when using the younger cohort. Specifically, AUC=0.67 and quintile ratio=4.6 among the tested population; AUC=0.66 and quintile ratio=7.0 among the untested population that attended a follow-up visit. We also evaluate this model using only adults over 50 (as opposed to the entire cohort) and performance remains better than our initial analysis. Specifically, AUC=0.67 and quintile ratio=5.1 among the tested population; AUC=0.80 and the quintile ratio is infinite (4.5% vs. 0%) among the untested population. (Performance for the untested population is noisily estimated because there are relatively few untested adults over 50 due to the testing guideline.) We have added these results to footnote 6 in Appendix D.

Finally, the reviewer asks “why is $\beta_\Delta$ negative for genetic risk score?” A negative $\beta_\Delta$ indicates that, controlling for cancer risk, patients with high genetic risk are under-tested. This is plausible because doctors frequently lack patient genetic information.

We thank the reviewer for their positive review and are glad that they found that our model makes sense and is novel, that our constraints are novel and practical, and that we verify the benefit of our constraints in both synthetic and real data. We will now address the reviewer’s questions and suggestions for improvement.

The reviewer first comments that the Heckman model introduced in Proposition 3.1 is too simple. We would like to clarify that our work investigates models beyond the Heckman model. In our paper we describe a broader class of models (described in equation 1), of which the Heckman model is only one special case (Proposition 3.1), and most of our experiments are run using models beyond the Heckman model. The reviewer also states, “Having the assumption that the unobservable always comes from an independent normal distribution can be too strong.” Indeed, this highlights a benefit of the general model class we describe, which works with alternate distributions of unobservables: in both our synthetic and real data experiments, we consider both uniform and normal distributions of unobservables (see Appendix C and Appendix E.2). Overall, we agree with the reviewer that the Heckman model is simple, and one of the strengths of our work is that we investigate models beyond the Heckman model.

Then the reviewer states, “When applying the model to the UK Biobank, filtering out individuals whose age is below 45 is not convincing.” While we focus on the younger cohort in our main analyses to create a challenging distribution shift, to address the reviewer’s concern we run our model on the entire population. We find that performance when using the full cohort is better than when using the younger cohort. Specifically, AUC=0.67 and quintile ratio=4.6 among the tested population; AUC=0.66 and quintile ratio=7.0 among the untested population that attended a follow-up visit. We have added these results to footnote 6 in Appendix D.

Finally, the reviewer asks why, for the model without the prevalence constraint, the $\beta_Y$ parameter decreases when the age variable increases in Figure 4. The model without the prevalence constraint learns this implausible age trend because it is learning the age trend which occurs among the tested population. Due to the age-based testing policy in the UK, patients under the age of 50 are tested only if they are of very high risk for breast cancer, so tested patients below the age of 50 have a higher risk than tested patients above the age of 50. The model without the prevalence constraint learns this trend, which is why its $\beta_Y$ decreases as age increases.

We are glad the reviewer found that the paper was well-written and easy to follow, that we applied our model to real data from a difficult medical setting, and that we addressed a significant distribution shift problem which is applicable to many domains. We will now address the reviewer’s questions and suggestions for improvement.

The reviewer first comments about the novelty of our work, stating “the linear risk setting has been considered before as cited in [(Hicks, 2021)].” We clarify that the novelty of our work is twofold. First, (Hicks, 2021) only considers the Heckman correction model; in our paper we describe a broader class of models (described in equation 1), of which the Heckman correction model is only one special case (Proposition 3.1). Secondly, we propose two novel constraints — the prevalence and expertise constraints — which are not considered in the Heckman model. These constraints are straightforward to implement and well-motivated in a medical setting. We validate both empirically and theoretically that these constraints improve parameter estimation.

The reviewer also provides individual comments for each of our constraints. For the expertise constraint, the reviewer states, “The expertise constraint sets one of the variables to 0 - this could be easily addressed by dropping that feature in the dataset.” We clarify that the expertise constraint only drops a subset of features when predicting the testing decision. Thus, $\beta_\Delta$ for these features are set to 0. However, these features are not dropped when predicting disease risk. Thus we still estimate $\beta_Y$ for all features (even for the features whose $\beta_\Delta$ is set to 0). Therefore the features for which we assume expertise cannot be completely dropped from the dataset. For the prevalence constraint, the reviewer states, “the prevalence constraint sets the expectation of the outcome - this could be addressed by normalizing the feature space and adding a bias term.” While this is correct for a linear model with an identity link, this is not true in general (e.g. for a Bernoulli-sigmoid link). Overall, we thank the reviewer for their comments on both constraints, and we will revise the manuscript to clarify these points. We agree with the reviewer that the constraints are straightforward to implement: this is a benefit which makes them compatible with a wide class of models.

The reviewer also asks what value our theoretical results, such as Proposition 3.2, add to the paper. The significance of the theoretical results is twofold. First, we show that the Heckman model is a special case of our model, providing an important connection to the econometrics literature and providing intuition about model identifiability. Secondly, we provide conditions under which fixing a parameter reduces the variance of the other parameters, improving the precision of parameter inference. While Proposition 3.2 is general, we also provide more specific results for the Heckman model (Proposition A.2). Ultimately, our theoretical results (i) provide a connection to the Heckman model, showing it is a special case of our general model and (ii) provide an explanation for why our constraints improve parameter estimation.

Finally, the reviewer asks what happens if we do not have follow-up data for certain patients. We use follow-up data to validate that our model’s inferred risk predictions indeed predict future breast cancer diagnoses even among the untested population, an approach also leveraged by prior work [1]. This is an improvement on merely assessing the model on the tested population, since it allows us to get some sense of whether we are able to accurately predict risk for the untested population, though the reviewer is correct that we are not guaranteed to have follow-up data for all untested patients (and we exclude patients with missing follow-up data from this analysis). To get around this limitation, we conduct three additional validations in section 5.2. Nevertheless, we thank the reviewer for their comment and we will acknowledge the limitations of this validation in the main text.

[1] Sendhil Mullainathan and Ziad Obermeyer. Diagnosing physician error: A machine learning approach to low-value health care. The Quarterly Journal of Economics, 137(2):679–727, 2022