Reliable and Trustworthy Machine Learning for Health Using Dataset Shift Detection

1 Introduction . Advances in artificial intelligence and machine learning have made medical diagnostic and screening tools more accurate and accessible . AI-powered diagnostic tools [ 5 , 14 ] are intended to assist medical personnel by making unbiased decision based on thousands of examples . In recent years , these models [ 15 , 44 , 48 , 32 ] are even becoming available to consumers through the growth of mobile health with the intention of expediting diagnoses through increasingly frequent testing . Moreover , mobile health [ 4 , 50 ] aims to improve access to medical expertise for those who are uninsured or live far away from hospitals . Despite the potential benefits of health AI systems , there are concerns about their performance in real-world settings . Data-driven models learn from examples , making them heavily reliant on the data upon which they have been trained . However , datasets often fail to get complete coverage over a domain , particularly for emerging datasets ; when new pulmonary diseases ( e.g. , MERS and COVID-19 ) emerge , a pulmonary classifier trained on the existing lung sounds would not be able to interpret sound of the new diseases . Previous work [ 42 , 78 ] has found that machine learning models behave unpredictably on the unseen data . This problem [ 4 , 69 ] is especially critical for medical diagnostic and screening tools since there are significant repercussions for mistakes . Researchers have proposed methods to estimate the uncertainty of a machine learning models ’ predictions based on the input [ 29 , 41 , 40 , 64 , 43 ] . Out-of-distribution detection methods can distinguish whether the input lies within the distribution of the training dataset , with out-of-distribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . data leading to less reliable prediction results . However , such important information has not been widely explored in the context of health applications . When health applications are put into the hands of consumers with limited understanding of the underlying algorithms , they may upload poor quality data that lies outside the distribution of the data that was collected by experts . For example , consumers who are using a health application that involves image processing may take photographs in poor lighting conditions or framing of the target object . Even when the data is high-quality , it may be captured with a smartphone that has different hardware specifications than the devices that were used to collect the training dataset . Unless the models are explicitly designed or trained to detect invalid data , the models will incorrectly produce a diagnostically meaningless result . In this work , we explore the utility of out-of-distribution detection for improving model performance and user-perceived trustworthiness of health-related models . We first benchmark our approach using publicly available deep learning models relating to various medical challenges and sensing domains — images for skin lesion classification , motion data for Parkinson ’ s disease severity , and audio for lung sound classification . After demonstrating that these models are susceptible to dataset shift , we demonstrate that the state-of-the-art out-of-distribution detectors can effectively exclude such data with over 95 % detection accuracy in most cases . We then explore the implications of this detector on user-perceived trustworthiness of the health models . After translating the out-of-distribution score into a human-interpretable metric , CONFIDENCE SCORE , we found that showing this information to end-users improved the user-perceived trustworthiness of the models . Furthermore , participants stated that they were more willing to make medical decisions based on models when they were shown the certainty metric . Our contributions in this work are as follows : • We identify and quantify the limitations of current health deep learning models when encountered with unseen data , • We evaluate the utility of out-of-distribution detection on various data types ( e.g. , image , audio , motion ) for medical screening and diagnosis , and • We evaluate the impact that dataset shift information has on user-perceived trustworthiness of health diagnostic results . 2 Related Work . Machine Learning-based Health Screening and Diagnosis In recent years , machine learning has been widely used for medical diagnosis and screening tool to help doctors diagnosis patients easier , faster , and more accurately . Machine learning models that learn from large-scale medical datasets are able to detect various symptoms and conditions , including mental health [ 26 , 68 ] , retinal disease [ 14 ] , lung cancer [ 5 ] . With the increasing ubiquity of smartphone and advances in its computing power , machine learning-based health screening can be done on mobile devices . Various machine learningbased mobile health applications have been proposed to detect health conditions ( e.g. , traumatic brain injury [ 49 ] , pancreatic cancer [ 48 ] , jaundice [ 15 ] ) and vital signals ( e.g. , heart rate [ 44 ] , respiratory rate [ 44 ] , heart rate variability [ 32 ] , blood pressure [ 65 ] , SpO2 [ 31 ] ) . Such mobile health applications can benefit nurses , health workers , and the general population for easier medical screening . While the health machine learning models show high accuracy on their own test datasets , their performance is questionable in real-world settings where the input data can vary drastically , resulting in unreliable prediction results [ 70 , 63 ] . Researchers have investigated the dataset shift problem for medical imaging ( e.g. , x-ray [ 11 , 10 ] , fundus eye images [ 11 ] , CT scans [ 74 ] , dermatology [ 63 , 56 ] ) , focusing on developing and evaluating out-of-distribution detection methods for specific domains . However , as more consumer-facing health applications are available in the market , this issue can lead the users to make medical decision based on incorrect results . In this work , we aim to explore ways to leverage dataset shift information to make the health machine learning models more reliable and trustworthy to the users . Dataset Shift Detection Recently researchers have proposed various methods to estimate the models ’ uncertainty due to dataset shift . The proposed methods leverage the output of the models to effectively detect out-of-distribution input that are different from the known distribution , in-distribution . Softmax confidence [ 29 ] has been the baseline for the out-of-distribution detection . Several work has been proposed for out-of-distribution detection using deep ensemble [ 38 ] , Mahalanobis distance [ 40 ] , Gram matrices [ 64 ] , energy score [ 43 ] , temperature scaling [ 41 , 64 ] , input perturbation [ 41 , 40 ] , mean and variance of channels activations [ 58 ] . Alternate training strategies [ 30 , 39 , 47 , 52 ] have been proposed to enable model to detect out-of-distribution . Generative models [ 54 , 60 , 53 , 77 ] are proposed to detect out-of-distribution examples . However , many approaches require re-training and re-designing of the models and prior knowledge of out-of-distribution datasets ; it is not realistic to apply these methods to the existing models . In this work , we explore Mahalanobis distance- [ 40 ] , Gram matrices- [ 64 ] , and energy-based [ 43 ] out-of-distribution detection methods for reliable and trustworthy machine learning for health since these methods show reasonable out-of-distribution detection performance , do not require retraining or prior knowledge of out-of-distribution datasets , and work on pre-trained discrimitive classifiers . Trustworthy AI Machine learning systems are deployed in real-world settings to billions of users , making significant impacts on high-stake decision making such as healthcare , policy , economy , and transportation . Failures in machine learning systems can cause fatal consequences and building trustworthy AI is one of the most important problems in machine learning community . In recent years , there are active and ongoing efforts aimed at making machine learning systems causal [ 6 , 55 ] , explainable [ 2 , 36 , 45 , 34 , 35 ] , fair [ 3 , 59 , 17 ] , robust [ 18 , 21 , 28 , 19 , 72 ] , and privacy-preserving [ 1 , 57 , 51 , 71 ] . This work contributes to trustworthy AI by improving reliability and user-perceived trustworthiness of machine learning for health using estimated uncertainty . Bhatt et al . [ 8 ] proposed to leverage uncertainty for users making decision and placing trust in machine learning models . This work explores similar approach where we adopt out-of-distribution detection as a method to measure uncertainty . We took a step further to investigate and quantify its effect in improving reliability and trustworthiness in the context of machine learning for health . 3 Background : Dataset Shift Detection Methods . We aim to leverage state-of-the-art out-of-distribution detection methods [ 40 , 64 , 43 ] in the health domain for users to safely use health deep learning models . We selected three out-of-distribution methods that show high accuracy on different out-of-distribution datasets , do not require re-training and prior knowledge of out-of-distribution datasets , and work on pre-trained discrimitive classifiers . These characteristics are important to help developers or other stakeholders ( e.g. , regulators , auditors , platforms ) easily adopt out-of-distribution detectors to any pre-trained models . In this section , we provide background on each out-of-distribution detection methods . 3.1 Mahalanobis Distance-Based Out-of-Distribution Detection . Mahalanobis distance is used to measure the proximity of a point to a certain Gaussian distribution . In the Mahalanobis distance-based out-of-distribution detector [ 40 ] , this property is used to represent each class ’ s samples at each layer of a network as a class conditional Gaussian distribution with mean µ̂cl and co-variance ⌃̂cl , where c indicates the class and l indicates the layer in the model . Given a sample input x to a neural network , for each layer , it computes the minimum layer-wise class conditional Mahalanobis distances for x . That is , for each layer , it finds the Mahalanobis distance associated with the closest class to x . In other words , this is equivalent to M ( x ) = maxc ( f ( x ) µc ) T⌃ 1 ( f ( x ) µc ) . The authors have demonstrated that adding small noise to the input can help better distinguish between in-distribution and out-of-distribution data . As the authors suggested , for the real-world setting where the out-of-distribution datasets are generally not available , we obtain the input noise magnitude by generating adversarial samples generated by FGSM [ 27 ] . 3.2 Gram Matrices-Based Out-of-Distribution Detection . Gram matrices are used to compute pairwise feature correlations and encode stylistic attributes . For out-of-distribution detection [ 64 ] , higher order Gram matrices are used to compute class-conditional bounds of feature correlations at all hidden layers of the network as higher order shows better out-ofdistribution detection performance . Higher order Gram matrices is expressed as GPl = ( F P l F PT l ) 1 P , where Fl is feature map at l-th layer and P is order . All elements of Gram matrices of an input at each layer are compared against the prepossessed minimum and maximum Gram matrices element values from in-distribution dataset to obtain deviation . If the input data is predicted as a certain class , the minimum and maximum values of the corresponding class will be used for comparison . The comparison is done for each layer to obtain layerwise deviations . Then , the deviations are used to get a total deviation , which is defined by the normalized sum of layerwise deviations . Whether the input data is from out-of-distribution is determined with a threshold which is defined as 95 % percentile of the total deviations of in-distribution energy score distribution .

Summary: This work investigates the effect of out-of-distribution data on machine learning models on healthcare applications. The authors first display a decrease in model performance on near and far out-of-distribution samples for 3 health applications. They then assess the performance of 3 techniques for out-of-distribution detection and build a confidence score from this binary classifier. The confidence score is assessed in a user-study as a complementary piece of information to be presented to the user.