Domain constraints improve risk prediction when outcome data is missing (2312.03878v3)

Abstract: Machine learning models are often trained to predict the outcome resulting from a human decision. For example, if a doctor decides to test a patient for disease, will the patient test positive? A challenge is that historical decision-making determines whether the outcome is observed: we only observe test outcomes for patients doctors historically tested. Untested patients, for whom outcomes are unobserved, may differ from tested patients along observed and unobserved dimensions. We propose a Bayesian model class which captures this setting. The purpose of the model is to accurately estimate risk for both tested and untested patients. Estimating this model is challenging due to the wide range of possibilities for untested patients. To address this, we propose two domain constraints which are plausible in health settings: a prevalence constraint, where the overall disease prevalence is known, and an expertise constraint, where the human decision-maker deviates from purely risk-based decision-making only along a constrained feature set. We show theoretically and on synthetic data that domain constraints improve parameter inference. We apply our model to a case study of cancer risk prediction, showing that the model's inferred risk predicts cancer diagnoses, its inferred testing policy captures known public health policies, and it can identify suboptimalities in test allocation. Though our case study is in healthcare, our analysis reveals a general class of domain constraints which can improve model estimation in many settings.

• The paper introduces a Bayesian model that accurately estimates risk by accounting for selective label biases using domain constraints.
• It incorporates prevalence and expertise constraints to refine parameter estimation for both tested and untested groups.
• Theory and synthetic experiments, including a breast cancer case study, validate the model's precision and practical relevance.

Improving Risk Prediction with Domain Constraints in Selective Labels Settings

The paper "Domain constraints improve risk prediction when outcome data is missing" addresses a critical challenge in machine learning when applied to real-world decision-making scenarios. Specifically, the paper focuses on settings characterized by selective label data, where outcomes are only observed if a human decision-maker decided to take an action, such as administering a medical test. This results in a biased dataset where the untested population might differ significantly from the tested one, both in observable and unobservable ways.

Core Contributions

The authors propose a novel Bayesian model to address issues arising from selective labels. The model is designed to estimate the true risk of a condition—such as disease presence—across both the tested and untested populations, acknowledging the potential biases introduced by selective testing policies. Key contributions of the paper include:

1. Bayesian Model for Selective Labels: The proposed model elucidates the data generation process by considering both observed features and latent unobservables. This allows for more accurate risk estimation by accounting for the unobserved variance between the tested and untested groups.
2. Domain Constraints: The paper introduces two domain-specific constraints to enhance parameter estimation:
   • Prevalence Constraint: Utilizes known disease prevalence rates in the population to constraint and correct risk predictions.
   • Expertise Constraint: Assumes that deviations from risk-based decision-making are bounded by certain known factors, reflecting the decision-maker's expertise.
3. Theoretical and Empirical Validation: The model is validated both theoretically and with synthetic experiments, demonstrating that incorporating domain constraints refines parameter estimates, boosting both precision and accuracy.
4. Case Study in Healthcare: The model's practical applicability is demonstrated through a case paper on breast cancer risk prediction using data from the UK Biobank. The model effectively predicts the likelihood of disease and tests allocation, identifying suboptimalities in historical testing decisions.

Theoretical Insights and Empirical Evidences

The theoretical analysis hinges on the Heckman correction model, traditionally used in econometrics for sample selection bias. The paper shows that its proposed model is a generalization, demonstrating that domain constraints can enhance the precision of parameter inference by reducing variance in the posteriors. The authors apply this framework through rigorous synthetic experiments, revealing empirical consistency with theoretical expectations.

The model application in breast cancer prediction highlights its ability to effectively utilize domain constraints. The approach mitigates the bias from selective testing practices by incorporating known prevalence rates, resulting in more plausible and accurate risk assessments. By comparing their approach with baseline models that don't consider domain constraints, the authors provide strong evidence supporting their strategy, showing improved prediction alignments with real-world policies and known epidemiological patterns.

Implications and Future Directions

This research has significant theoretical and practical implications. Theoretically, it advances understanding of how domain knowledge can constrain and guide machine learning models in biased settings, potentially impacting a wide range of applications beyond healthcare, such as criminal justice and finance.

Future work could explore further extensions to other domains, develop scalable implementations for high-dimensional data using variational approaches, and expand on the integration of more complex decision-making processes into the model framework.

In conclusion, this paper contributes a sophisticated approach to handling selective label bias in machine learning with concrete, domain-specific strategies, paving the way for more robust risk prediction models in practice.