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: The paper tackles the problem of automatic disease diagnosis through reinforcement learning under a setting of task-oriented dialogues. The authors proposed to integrate hierarchical policies with two levels (one high-level master model, and one low-level policy) into the dialogue policy learning. Experiments on both real-life and synthetic data suggest the proposed approach is effective.