Task-oriented Dialogue System for Automatic Disease Diagnosis via Hierarchical Reinforcement Learning

1 INTRODUCTION . With the development of electronic medical records ( EMRs ) , researchers explore different machine learning approaches for automatic diagnosis ( Shivade et al. , 2013 ) . Although impressive results have been reported for the identification of various diseases ( Jonnalagadda et al. , 2017 ; Doshi-Velez et al. , 2014 ) , they rely on well established EMRs which are labor-intensive to build . Moreover , each diagnosis requires EMR of the patient , while those systems lack the ability to obtain information through interaction with patients . To relieve the pressure for constructing EMRs , researchers ( Wei et al. , 2018 ; Xu et al. , 2019a ) introduce task-oriented dialogue system to request symptoms automatically from patients for disease diagnosis . They formulate the task as Markov Decision Processes ( MDPs ) and employ reinforcement learning ( RL ) based methods in the policy learning for the agent . Existing frameworks utilize the setting of flat policy that treats diseases and all related symptoms equally . However , although RL-based approaches show positive results for symptom acquisition , when it comes to hundreds of diseases in real environment , flat policy is quite impractical . In general , a particular disease is often related to a certain group of symptoms . As shown in figure 1 , we present the correlation between diseases and symptoms , where x-axis represents symptoms and y-axis is the proportion of diseases related . From figure 1 , there are obvious differences in the distribution of symptoms for different diseases . In offline consultation , doctors also do the preexamination and triage according to the different symptoms that patient suffered , then doctors in different departments will make more detailed diagnosis . This method significantly reduces the workload of individual doctors and enables them to be more specialized in a certain field . In the reinforcement learning algorithm , we found that Hierarchical Reinforcement Learning ( HRL ) ( Parr & Russell , 1998 ; Sutton et al. , 1999 ) conforms to this problem-solving logic . Recently , HRL , in which multiple layers of policies are trained to perform decision making , has been successfully applied to different scenarios , including course recommendation ( Zhang et al. , 2019 ) , visual dialogue ( Zhang et al. , 2018 ) , relation extraction ( Feng et al. , 2018 ) , etc . The natural hierarchy of target tasks are modeled either manually or automatically , which motivates us to classify diseases 0 % 50 % 100 % Connective tissue disorder Flat feet Ganglion cyst De Quervain disease Gout Chondromalacia of the patella Fibromyalgia Adhesive capsulitis of the shoulder Chronic back pain Ankylosing spondylitis ( a ) ( b ) 0 % 50 % 100 % children ’ s bronchitis upper respiratory infection children functional dyspepsia infantile diarrhea symptoms symptoms pr op or tio ns pr op or tio nsF gure 1 : The disease distribution over symptoms in the real-world dataset ( see section 3.1 ) . x-axis stands for symptoms and y-axis is the proportion . Each bar describes the disease distribution given a symptom . into different groups following the setting of departments in the hospital , and design a hierarchical structure for symptom acquisition and disease diagnosis . In this paper , we classify diseases into several groups and build a dialogue system with a hierarchy of two levels for automatic disease diagnosis using HRL methods . The high level policy consists of a model named master and the low level policy consists of several workers and a disease classifier . The master is responsible for triggering a model in the low level . Each worker is responsible for inquiring symptoms related to a certain group of disease while disease classifier is responsible for making the final diagnosis based on information collected by workers . The proposed framework imitates a group of doctors from different departments to diagnose a patient together . Among them , each worker acts like a doctor from a specific department , while the master acts like a committee that appoints doctors to interact with this patient . When information collected from workers are sufficient , the master would activate a separate disease classifier to make the diagnosis . Models in the two levels are trained jointly for better disease diagnosis . We test our model in three large real-world datasets and a synthetic dataset for the evaluation of our model . Experimental results demonstrate that the performance of our hierarchical framework outperforms other state-of-the-art approaches . We summary our contribution as follows : 1 ) We propose a new common-sensible method for automatic diagnosis . It simulates the real scene of clinical practice , and assigns patients to different workers through high-level policy , thereby reducing the action space and improving training efficiency . Also , the method can be compatible with different network models and training strategies ; 2 ) We evaluate our method on three public datasets and a synthetic datasets to demonstrate the superiority of our framework with higher accuracy of disease and recall of symptoms ; 3 ) We construct a detailed analysis of the effectiveness of each part of this method to confirm that it is reasonable and effective . 2 FRAMEWORK . In this section , we introduce our hierarchical reinforcement learning framework for disease diagnosis . We start with the hierarchical policy setting and then introduce every part of method . We further introduce reward shaping techniques which help us improve the model performance . 2.1 HIERARCHICAL POLICY . To reduce the large action space , we extend the RL formulation to a hierarchical structure with twolayer policies for automatic diagnosis . Following the options framework ( Sutton et al. , 1999 ) , our framework is illustrated as in Figure 2 ( a ) . There are five components in our framework : master , workers , disease classifier , internal critic and user simulator . Specifically , we divide all the diseases in D into h subsets D1 , D2 , . . . , Dh , where D1 ∪D2 ∪ · · · ∪ Dh = D and Di ∩ Dj = ∅ for any i 6= j . Each Di is associated with a set of related symptoms Si ⊆ S . While worker wi is responsible for collecting information about symptoms of Si . At turn t , The state st will be encoded as one-hot vectors that reflects the status of each symptom , i.e. , st = [ b > 1 , b > 2 , · · · , b > |S | ] > , and the master decides whether to collect symptom information ( picking one worker to interact with user for several turns ) or inform the diagnosis result ( picking disease classifier to output the predicted disease ) . When agent inquires the symptom i , the patient will respond with true/false/unknown and agent will update the corresponding symptom statuses bi ∈ R3 . When agent informs a disease , the dialogue will be terminated as success/fail in terms of the correctness of diagnosis . Note that the not-requested symptoms are all encoded as b = [ 0 , 0 , 0 ] . After each action is executed , the agent will receive a reward r based on the result of the action . An illustration of the diagnosis process with interactions between models in two levels are presented in Figure 2 ( b ) . As for internal critic , it is responsible for both returning intrinsic reward to the worker and telling whether the subtask of the invoked worker is finished . In addition , the user simulator is applied to interact with our model and return extrinsic reward . 2.1.1 MASTER . The action space of master isAm = { wi|i = 1 , 2 , · · · , h } ∪ { d } . The action wi indicates activating worker wi while d is a primitive action which means activating the disease classifier . At each turn t , the master takes the dialogue state st ∈ S as input and takes an action amt ∈ Am according to it ’ s policy . An extrinsic reward ret will be returned to master from the environment . Once the master activates a worker , this worker will interact with user for N turns until the subtask is terminated . Only after that master can take a new action and observe a new dialogue state . So , the learning problem of master can be formulated as a Semi-Markov Decision Process ( SMDP ) , where the extrinsic rewards returned can be accumulated as the immediate rewards for the master ( Ghavamzadeh , 2005 ) . That is to say , after taking an action amt , the reward r m t for master can be defined as : rmt = { ∑N t′=1 γ t′ret+t′ , if a m t = w i ret , if a m t = d ( 1 ) Where s′ is the observed dialogue state of master after it takes an action am , am′ is the next action when the state is s′ . The objective of master is to maximize the extrinsic reward , thus we can write the master ’ s loss function as follows : L ( θm ) = Es , am , rm , s′∼Bm [ ( y −Qm ( s , am ; θm ) ) 2 ] ( 2 ) where y = rm + γN maxam′ Qm ( s′ , am′ ; θ−m ) , θm is the network parameter at current iteration , θ−m is the network parameter of previous iteration and Bm is the fixed-length buffer of samples for master . 2.1.2 WORKER . The worker wi corresponds to the set of diseases Di and the set of relevant symptoms Si , and the action space of worker wi is Si . At turn t , if worker wi is invoked , the current state of master st will be passed to worker wi , then worker wi will extract the corresponding state representationsit from st and take sit as input to generate an action a i t ∈ Awi . After taking action ait ∈ Awi , the dialogue is updated into st+1 and worker wi will receive an intrinsic reward rit from the module of internal critic . So the objective of worker is to maximize the expected cumulative discounted intrinsic rewards . The loss function of worker wi can be written as : L ( θiw ) = Esi , ai , ri , si′∼Bwi [ ( yi −Q i w ( s i , ai ; θiw ) ) 2 ] ( 3 ) where yi = ri + γwmaxai′ Qiw ( s i′ , ai ′ ; θi−w ) , γw is the discounted factor of all the workers , θ i w is the network parameter at current iteration , θi−w is the network parameter of previous iteration and Bwi is the fixed-length buffer of samples for worker wi . 2.1.3 DISEASE CLASSIFIER . Once the disease classifier is activated by master , it will take the master state st as input and output a vector p ∈ R|D| , which represents the probability distribution over all diseases . The disease with highest probability will be returned to the user as the diagnosis result . Two layers of Multi-Layered Perceptron ( MLP ) is utilized here for the disease diagnosis . 2.1.4 INTERNAL CRITIC . The internal critic is responsible for generating intrinsic reward rit to worker w i and judging the termination condition for the worker . In our task , a worker is terminated as failed when the number of subtask turns reaches T sub , and terminated as successful when a correct symptom is requested by the agent . After that , a positive reward will be returned if the worker requests a correct symptom or disease . If there are repeated actions generated or the number of subtask turns reaches T sub , rit would be negative . Otherwise , rit would be 0 .

Summary: This paper introduces hierarchical reinforcement learning (HRL) into automatic disease diagnosis, which reduces the action space and improves training efficiency. Besides, the authors also expand an existing public dataset and build a synthetic dataset for evaluation. The Experimental results show that their proposed hierarchical framework achieves higher accuracy and symptom recall in disease diagnosis than existing several baselines.