summaries: 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.

summaries: ##########################################################################  Summary: This works proposes dialogue policy learning by incorporating a hierarchical policy. The proposed hierarchical reinforcement learning-based approach has two models, i.e., one master model that instantiates low models. The low models have several symptoms checkers and disease classifier. Experiments are conducted on both real and synthetic datasets to demonstrate the efficacy of the proposed work.    ##########################################################################

summaries: 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.

summaries: This paper applies Hierarchical Reinforcement Learning (HRL) to automatic disease diagnosis in task-oriented dialogues setting. The authors argue that applying RL to automatic disease diagnosis is challenging because the action space (i.e., symptoms) is very large. Therefore, they propose to learn a hierarchical dialogue policy where the high level policy is for categorizing patents into different groups and the low level policy is for checking symptoms and classifying diseases within a group.  Experimental results on multiple datasets show that the proposed HRL strategy achieve higher disease diagnosis accuracy compared to existing RL systems.

abstractText: In this paper, we focus on automatic disease diagnosis with reinforcement learning (RL) methods in task-oriented dialogues setting. Different from conventional RL tasks, the action space for disease diagnosis (i.e., symptoms) is inevitably large, especially when the number of diseases increases. However, existing approaches to this problem typically works well in simple tasks but has significant challenges in complex scenarios. Inspired by the offline consultation process, we propose to integrate a hierarchical policy of two levels into the dialogue policy learning. The high level policy consists of a master model that is responsible for triggering a low level model, the low level policy consists of several symptom checkers and a disease classifier. Experimental results on both self-constructed real-world and synthetic datasets demonstrate that our hierarchical framework achieves higher accuracy and symptom recall in disease diagnosis compared with existing systems. Besides, both our datasets and codes will be released.

authors:
Y.E. Chen
K.F. Tang
Y.S. Peng
E.Y. Chang

title: Effective medical test suggestions using deep reinforcement

authors:
Finale Doshi-Velez
Yaorong Ge
Isaac Kohane

title: Comorbidity clusters in autism spectrum disorders: an electronic health record time-series analysis

venue: Pediatrics, 133(1):e54–e63,

authors:
Jun Feng
Minlie Huang
Yijie Zhang
Yang Yang
Xiaoyan Zhu

title: Relation mention extraction from noisy data with hierarchical reinforcement learning

venue: arXiv preprint arXiv:1811.01237,

authors:
Carlos Florensa
Yan Duan
Pieter Abbeel

title: Stochastic neural networks for hierarchical reinforcement learning

venue: arXiv preprint:

authors:
Carlos Florensa
David Held
Xinyang Geng
Pieter Abbeel

title: Automatic goal generation for reinforcement learning agents

title: Hierarchical reinforcement in continuous state and multi-agent environments

authors:
Jiaxian Guo
Sidi Lu
Han Cai
Weinan Zhang
Yong Yu
Jun Wang

title: Long text generation via adversarial training with leaked information

venue: In Thirty-Second AAAI Conference on Artificial Intelligence,

authors:
Weijie He
Xiaohao Mao
Chao Ma
José Miguel Hernández-Lobato
Ting Chen

title: Fit: a fast and accurate framework for solving medical inquiring and diagnosing

authors:
Siddhartha R Jonnalagadda
Abhishek K Adupa
Ravi P Garg
Jessica Corona-Cox
Sanjiv J Shah

title: Text mining of the electronic health record: An information extraction approach for automated identification and subphenotyping of hfpef patients for clinical trials

venue: Journal of cardiovascular translational research,

authors:
Hao-Cheng Kao
Kai-Fu Tang
Edward Y Chang

title: Context-aware symptom checking for disease diagnosis using hierarchical reinforcement learning

authors:
Tejas D. Kulkarni
Karthik R. Narasimhan
Ardavan Saeedi
Joshua B. Tenenbaum

title: Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

title: Human-level control through deep reinforcement learning

authors:
Ofir Nachum
Shixiang Gu
Honglak Lee
Sergey Levine

title: Data-efficient hierarchical reinforcement learning

authors:
Ronald Parr
Stuart J Russell

title: Reinforcement Learning with hierarchies of machines

venue: Advances in neural information processing systems,

authors:
Yu-shao Peng
Kai-Fu Tang
Hsuan-Tien Lin
Edward Y Chang

title: Refuel: Exploring sparse features in deep reinforcement learning for fast disease

authors:
Chaitanya Shivade
Preethi Raghavan
Eric Fosler-Lussier
Peter J Embi
Noemie Elhadad
Stephen B Johnson
Albert M Lai

title: A review of approaches to identifying patient phenotype cohorts using electronic health records

venue: Journal of the American Medical Informatics Association,

authors:
Richard S Sutton
Doina Precup
Satinder Singh

title: Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning

venue: Artificial Intelligence,

authors:
Ryuichi Takanobu
Tianyang Zhang
Jiexi Liu
Minlie Huang

title: A hierarchical framework for relation extraction with reinforcement learning

venue: arXiv preprint arXiv:1811.03925,

authors:
Da Tang
Xiujun Li
Jianfeng Gao
Chong Wang
Lihong Li
Tony Jebara

title: Subgoal discovery for hierarchical dialogue policy learning

venue: arXiv preprint arXiv:1804.07855,

authors:
Kai-Fu Tang
Hao-Cheng Kao
Chun-Nan Chou
Edward Y Chang

title: Inquire and diagnose: Neural symptom checking ensemble using deep reinforcement learning

venue: In Proceedings of NIPS Workshop on Deep Reinforcement Learning,

authors:
Xin Wang
Wenhu Chen
Jiawei Wu
Yuan-Fang Wang
William Yang Wang

title: Video captioning via hierarchical reinforcement learning

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Zhongyu Wei
Qianlong Liu
Baolin Peng
Huaixiao Tou
Ting Chen
Xuanjing Huang
Kam-Fai Wong
Xiangying Dai

title: Task-oriented dialogue system for automatic diagnosis

venue: In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers),

authors:
Yuan Xia
Jingbo Zhou
Zhenhui Shi
Chao Lu
Haifeng Huang

title: Generative adversarial regularized mutual information policy gradient framework for automatic diagnosis

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Lin Xu
Qixian Zhou
Ke Gong
Xiaodan Liang
Jianheng Tang
Liang Lin

title: End-toend knowledge-routed relational dialogue system for automatic diagnosis

venue: arXiv preprint: arXiv:1901.10623,

title: End-to-end knowledge-routed relational dialogue system for automatic diagnosis

authors:
Jiaping Zhang
Tiancheng Zhao
Zhou Yu

title: Multimodal hierarchical reinforcement learning policy for task-oriented visual dialog

venue: In Proceedings of the 19th Annual SIGdial Meeting on Discourse and Dialogue,

authors:
Jing Zhang
Bowen Hao
Bo Chen
Cuiping Li
Hong Chen
Jimeng Sund

title: Hierarchical reinforcement learning for course recommendation in moocs

heading: 1 INTRODUCTION

text:
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.

heading: 2 FRAMEWORK

text: 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.

heading: 2.1 HIERARCHICAL POLICY

text:
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.

heading: 2.1.1 MASTER

text:
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.

heading: 2.1.2 WORKER

text:
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.

heading: 2.1.3 DISEASE CLASSIFIER

text: 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.

heading: 2.1.4 INTERNAL CRITIC

text: 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.

heading: 2.1.5 USER SIMULATOR

text: Following (Wei et al., 2018) and (Xu et al., 2019a), we set up a user simulator to interact with the agent. At the beginning of each dialogue session, the user simulator samples a user goal from the training set randomly. The simulator will initialize the dialogue session base on the explicit symptoms, and interacts with the agent based on the implicit symptoms. One dialogue session will be terminated as successful if the agent make the correct diagnosis, or failed if the informed disease is incorrect or the dialogue turn reaches the maximal turn T .

heading: 2.2 REWARD SHAPING

text:
In practice, the number of symptoms a patient suffers from is much less than the size of symptom set S , which results in a sparse feature space. In other words, it is hard for the agent to locate the symptoms that the user truly suffers from. In order to encourage master to choose a worker that can discover more positive symptoms, we follow (Peng et al., 2018) and use the reward shaping method to add auxiliary reward to the original extrinsic reward while keeping the optimal reinforcement learning policy unchanged.
The auxiliary reward function from state st to st+1 is defined as f(st, st+1) = γλφ(st+1)− φ(st), φ(s) is the potential function, φ(s) counts the number of correct symptoms informed for a given state s, λ > 0 is a hyper-parameter which controls the magnitude of reward shaping and S⊥ is the terminal state set. Thus, the reward function for master will be changed into Rφt = rt + f(st, st+1).

heading: 2.3 TRAINING

text: Both the master policy πm and each worker’s policy πwi are parameterized via Deep QNetwork (Kulkarni et al., 2018; Mnih et al., 2015). In DQN, the action is often selected following an -greedy policy. In our hierarchical framework, both the master and the workers behave following their own -greedy policy for training and greedy policy for testing. During the training process, we store (st, amt , r m t , st+N ) in Bm and (sit, awt , rit, sit+1) in Bwi . At each training step, we perform experience replay to update the current networks for both master and workers in Bm and Bwi respectively, while the target networks are fixed during experience replay.

heading: 3 EXPERIMENTAL DATASETS

text: We evaluate our framework on one simulated public datasets: SymCat, and three medical dialogue datasets: MZ-4 (Wei et al., 2018), MZ-10 and Dxy (Xu et al., 2019b). MZ-10 comes from our expansion on the basis of MZ-4, so that it can be close to a more complex dialogue environment in real medical dialogue.

heading: 3.1 REAL-WORLD DATASET

text:
MZ-4 An existing dataset collected from real world for the evaluation of task-oriented DS (Wei et al., 2018). After proposed, it became one of the popular datasets for evaluating automatic medical systems. This dataset includes 4 diseases, 230 symptoms and 1733 user goals. Most of the existing RL-base methods have been tested on this dataset.
MZ-10 A newly constructed real-world dataset contains 4116 user goals that belong to 10 diseases, including typical diseases of digestive system, respiratory system and endocrine system. The raw data is collected from the pediatric department on a Chinese Online Healthcare Community1. Each user record consists of the self-report provided by the user and conversation text between the patient and a doctor. We hire experts with medical background to identify symptom expressions and label them with three tags (“True”, “False” or “UNK”) to indicate whether the user suffers this symptom. After that, experts manually link each symptom expression to a concept on SNOMED CT 2. Symptoms extracted from self-report are treated as explicit symptoms and the ones extracted from conversation are implicit symptoms. At the same time, to ensure that the dataset is not too sparse, we only keep the symptoms with the top-100 frequency. The real-world dataset we constructed contains 4116 user goals, of which 80% for training ,10% for validating and 10% for testing.
Dxy A Dialogue Medical dataset (Xu et al., 2019b) contains data from a prevalent Chinese online healthcare website3. This dataset contains 527 user goals, including 5 diseases and 41 specific symptoms. Both datasets are already in the form of disease-symptoms pairs.
For real-world dataset, we compare the similarities of the counts of symptoms in each disease on the training set, and divide the diseases with higher similarity (above 0.5) into a group.

heading: 3.2 SYNTHETIC DATASET

text:
However, the number of disease and user goals in real-world dataset is still limited. In order to show the effectiveness of the HRL method, we build a synthetic dataset (SD) following (Kao et al., 2018) for further analysis of the method. It is constructed based on symptom-disease database called SymCat4. There are 801 diseases in the database and we classify them into 21 departments (groups) according to International Classification of Diseases (ICD-10-CM)5. We choose 9 representative departments from the database, each department contains top 10 diseases according to the occurrence rate in the Centers for Disease Control and Pre-vention (CDC) database.
In SymCat database, each disease is linked with a set of symptoms, where each symptom has a probability indicating how likely the symptom is identified for the disease. Based on the probability distribution, we generate record one by one for each target disease. Given a disease and its related symptoms, the generation of a user goal follows two steps. First, for each related symptom, we sample the symptom based on the probability. Second, a symptom is chosen randomly to be the explicit one (same as symptoms extracted from self-report in RD) and rest of true symptoms are treated as implicit ones. The synthetic dataset we constructed contains 30,000 user goals, including 90 diseases and 267 symptoms, of which 80% for training and 20% for testing.
Table 1 shows the details of all dataset used in this paper.
1http://muzhi.baidu.com 2https://www.snomed.org/snomed-ct 3https://dxy.com/ 4www.symcat.com 5https://www.cdc.gov/nchs/icd/

heading: 4 PERFORMANCE ON REAL-WORLD DATASETS

text: We train our model following the setting in Appendix A.1, and compare the effectiveness of different models in terms of accuracy, average turns and match rate.

heading: 4.1 MODELS FOR COMPARISON

text:
We compare our model with some state-of-the-art reinforcement learning models for disease diagnosis.
• Flat-DQN: This is the agent of (Wei et al., 2018), which has one layer policy and an action space including both symptoms and diseases.
• HRL-pretrained: This is a hierarchical model from (Kao et al., 2018). The setting is similar to ours, however, the low level policy is pre-trained first and then the high level policy is trained. Besides, there is no disease classifier and the diagnosis is made by workers.
• REFUEL: This is a reinforcement learning method with reward shaping and feature rebuilding (Peng et al., 2018). It uses a branch to reconstruct the symptom vector to guide the policy gradient.
• GAMP: This is a GAN-based policy gradient network (Xia et al., 2020). It uses the GAN network to avoid generating randomized trials of symptom, and add mutual information to encourage the model to select the most discriminative symptoms.
• KR-DS: This is an improved method based on Flat-DQN (Xu et al., 2019b). It integrates a relational refinement branch and a knowledge-routed graph to strengthen the relationship between disease and symptoms.
• FIT: FIT is a novel non-RL framework for inquiring and diagnosing tasks (He et al., 2020). It propose a self-attention diagnosis model and use a multi-modal variational autoencoder (MVAE) as a probabilistic model dealing with partially observed data to improve the diagnostic accuracy.

heading: 4.2 OVERALL PERFORMANCE

text:
Table 2 show the overall performance of different models respectively. In MZ-4 dataset, most of the methods do not report the average turns and match rate of the dialogue, so we only compare the accuracy of the disease. Observed from MZ-4 dataset, our method outperforms all the methods by at least 9%. In MZ-10 dataset, the advantages of our model become more obvious. In comparison with flat-DQN, our method outperforms in both accuracy and match rates. When the number of symptoms increases, the symptom-disease and symptom-disease relationship becomes more complicated, which make Flat-DQN difficult to deal with. Our method makes the distribution of symptoms in each type of disease more logical by dividing diseases in different groups. So it can hit more symptoms during dialogue and provide more information for disease diagnosis.
Also, the results in Dxy dataset shows similar conclusions. Compared with other methods, our method can inquire more appropriate symptoms through interactions with patients, so as to make better decisions. In Dxy dataset, our methods outperforms other method in both accuracy and match rate by a larger magnitude, which means our method is more reasonable, efficient and logical.

heading: 4.3 ABLATION STUDIES

text:
To verify the effects of the main components of our method, we further conducted a series of ablation studies. Because our HRL method consists of two parts, master-worker structure and disease classifier, we try not to divide symptoms into different group (Denoted as HRL (w/o grouped)) and remove the separate disease discriminator (Denoted as HRL (w/o discriminator)). Note that when we remove the master-worker structure and leave the work of judging the disease to the master, the model will degenerate into a simple Flat-DQN model.
The result is shown in Table 2. From the results, it can be seen that dividing diseases into different groups can effectively improve the recall of symptoms. Comparing HRL (w/o grouped) and Flatdqn, a single disease classifier can greatly improve the efficiency of disease classification, while avoiding the dilemma caused by the uneven distribution of diseases and symptoms. Note that the implicit symptoms and the number of user goals in the Dxy dataset are limited, the advantages of the overall HRL method over its various parts are not very obvious.

heading: 5 PERFORMANCE ON SYNTHETIC LARGE-SCALE DISEASES

text: In current practice, most reinforcement learning methods are still powerless when faced with a large and sparse action space. However, the existing real-world datasets are limited in the number of diseases and the user goals. Therefore, in order to show the advantages and effects of our method on large-scale datasets, we will conduct validity experiments on Synthetic datasets.

heading: 5.1 OVERALL PERFORMANCE

text:
As shown in Table 3, like Flat-DQN, most reinforcement learning methods face a difficult choice when facing this problem. Due to the large difference in the distribution of diseases and symptoms in the SymCat dataset, when the reward for the symptoms is too high, the agent will continue to inquiry symptoms without judging the disease. On the contrary, the agent will directly predict the disease and give up to search the sparse symptom space (like the Flat-DQN model we reported in table). Both of these situations are of limited help to real-world medical practice, but the balance between the two is difficult to hold and the system is extremely unstable.
Compared to baseline models, our method takes more turns to have interactions with the user so that it can collect more information of their implicit symptoms. With more information about implicit symptoms, our model significantly outperforms the other baselines in the diagnosis success rate, which shows our method is more reasonable and stable.

heading: 5.2 STABILITY OF DISEASE CLASSIFIER

text:
Because different diseases may have the same symptoms, it is an impossible task to distinguish all diseases correctly. At the same time, the same type of disease usually have similar treatment regimen. Therefore, if the disease classifier can make the judgment that is as close as possible to the correct disease, it can also help doctors roughly determine the type of disease. In order to have deeper analysis of the user goals which have been informed the wrong disease by the agent, we collect all the wrong informed user goals and present the error matrix in Fig 3.
It shows the disease prediction result for all the 9 groups. We can see the color of the diagonal square is darker than the others, which means the wrongly predicted disease and the correct disease are in the same group. This is reasonable because diseases in the same groups usually share similar symptoms and are therefore difficult to be distinguished. On the other hand, it also proved that even if the model cannot make correct predictions, it can still find right type of diseases to assist in real consultations.

heading: 5.3 EFFICIENCY OF DIFFERENT WORKERS

text: While proving the effect of the disease classifier, we also hope that the workers in each group can also play a positive role in the judgment of the model. We evaluate the performance of workers in terms of success rate, average intrinsic rewards and match rate. Match rate means the proportion of actions which have requested about the implicit symptoms that the user has. The results can be seen in Appendix A.2.We found most of the workers can successfully exit by querying the symptoms which patient suffered. It proves that workers in different groups can learn the symptom characteristics of the group, and use the knowledge to guide the consultation process.

heading: 6 RELATED WORK

text:
This paper is related to two major research fields, namely hierarchical reinforcement learning and RL-based methods for disease diagnosis.
Hierarchical reinforcement learning HRL has a hierarchical policy and has been proposed to solve the problems with large action space. One classic framework of HRL is options framework (Sutton et al., 1999), which involves abstractions over action space. At each step, the agent chooses either a one-step “primitive” action or a “multi-step” action (option). Kulkarni et al. (2018) proposed a hierarchical-DQN, which integrates deep Q-learning into HRL, to learn policies for different levels simultaneously. HRL has been successfully applied to different tasks and reached promising results. (Zhang et al., 2019; Wang et al., 2018; Zhang et al., 2018; Takanobu et al., 2018; Feng et al., 2018; Guo et al., 2018). Most existing works decompose the corresponding task into two steps manually, where the first step is finished by high level policy while the second step is finished by the low level policy. The task-specific design limits the generalization of HRL to complex tasks. Florensa et al. (2017) proposed a general framework that first learns useful skills (high level policies) in an environment and then leverages the acquired skills for learning faster in downstream tasks. What’s more, some methods that generate or discover goals automatically when training the policies of two levels have been proposed. (Nachum et al., 2018; Florensa et al., 2018; Tang et al., 2018).
RL-based methods for disease diagnosis There are some previous works that applies the flat RL-based method in the medical dialogue system Wei et al. (2018); Xu et al. (2019a); Tang et al. (2016); Kao et al. (2018); Peng et al. (2018) and generates positive results. In the work of Wei et al. (2018) and Xu et al. (2019a), both symptoms and diseases are treated as actions and a flat policy is used for choosing actions. After that, the HTC Research & Healthcare team did a lot of interesting work on this field. Kao et al. (2018) present a context-aware hierarchical reinforcement method, using policy gradients to make decisions based on the patient’s age, gender, season, and explicit symptoms. Peng et al. (2018) proposed reward shaping and feature rebuilding techniques to help agent effectively learn a better strategy and Chen et al. (2019) introduced a new multiple action policy representation to help agent suggesting medical tests to facilitate disease diagnosis. While pursuing a higher accuracy of disease diagnosis, researchers also hope to make the inquiry process more efficient and logical. In order to encourage the agent to locate the symptoms more specifically, Xu et al. (2019a) introduce the knowledge graph into their dialogue system, Xia et al. apply GAN as policy network to capture the relations between different symptoms, these two works all reach a better accuracy than previous works.

heading: 7 CONCLUSIONS AND FUTURE WORK

text:
In this work, we formulate disease diagnosis as a hierarchical policy learning problem, where symptom acquisition and disease diagnosis are assigned to different kinds of workers in the lower level of the hierarchy. A master model is designed in the higher level that is responsible for triggering models in low level. We improve the previous model by introducing disease classifier and internal critic and then train them together to reach the optimal. Besides the exist dataset, we extend the Muzhi dataset and build a synthetic dataset to evaluate our hierarchical model. To the best of our knowledge, this is the first time both kinds of datasets are used for model evaluation and we make both datasets public. The experimental results on all datasets demonstrate that our hierarchical model outperforms other RL-based model in both disease accuracy and symptom recall.
In the future, we would like to continue our research in three directions. First, we would like to make contribution on real-world dataset construction by introducing more diseases and annotate more samples. Second, although we can get a higher symptom recall, the symptom inquiry sequence is still disorderly. We can try to use some special reward setting to make it more logical. Third, we only use MLP as out deep learning network in this paper, and it would be interesting to integrate our hierarchical structure with other deep learning models and training strategies in the online selfdiagnosis.

heading: A APPENDIX

text:
A.1 IMPLEMENTATION DETAILS
The for master and all the workers are all set to 0.1. For the master, the maximal dialogue turn T is set to 22, it will receive a extrinsic reward of +1 if the master inform the right disease. Otherwise, it will receive a extrinsic reward of -1 if the dialogue turn reaches the maximal turn or a wrong disease is informed. At non-terminal turns, the extrinsic reward is 0. Moreover, the sum of the extrinsic rewards (after reward shaping) over one subtask taken by a worker will be the reward for master. The maximal dialogue turn T sub is set to 5 for each worker.
For master and all the workers, the neural network of DQN is a three-layer network with two dropout layers and the size of hidden layer is 512, learning rate for the DQN network is set to 0.0005. All parameters are set empirically and settings for the datasets are the same. In addition, all the workers are trained every 10 epochs separately during the training process of master. For the disease classifier, the neural network is a two-layer network with a dropout layer and the size of hidden layer is 512, learning rate for the network is set to 0.0005. Moreover, it’s trained every epochs during the training process of master.
During the training process, it will take about 5000 epochs for the model to reach convergence, which takes about 18 hours given one GPU. For the best-performing model, γ is set to 0.95, discounted factor γw is set to 0.99, λ in reward shaping is set to +1.
A.2 THE PERFORMANCE OF DIFFERENT WORKERS IN SYMCAT DATASET.


summaries: This paper introduces the self-supervised learning (SSL) framework for FL. Different SSL methods are investigated to study their feasibility under the FL setting. With the popular SimSiam framework, personalized federated SSL is proposed. The performance comparison indicates that representation regularization-based personalization method is able to outperform other variants.

summaries: The authors merge self-supervised learning in personalized federated learning to solve the limited label and data heterogeneity problems in the local clients. They test several current algorithms under their framework. Then they propose an algorithm named Per- SSFL considering the balance between the consensus and personalization. Finally, they provide experiment results to support their claims and comprehensive analysis of what they find. The main contribution is to design a self-supervised FL framework with supportive experiment results. They also provide suggestions to choose appropriate algorithms and hyper-parameters under different settings.

summaries: The submission is not anonymized, the author's name chaoyanghe appears frequently in the attached SSFL_Sumpplementary/SSFL-Source-Code. Therefore, it should be directly desk rejected.  --- Logging error --- Traceback (most recent call last):   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 1025, in emit     msg = self.format(record)   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 869, in format     return fmt.format(record)   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 608, in format     record.message = record.getMessage()   File "/Users/chaoyanghe/opt/anaconda3/envs/fedml/lib/python3.7/logging/__init__.py", line 369, in getMessage     msg = msg % self.args

summaries: This work is concerned with a very practical scenario of Federated Learning where the participating agents may not have access to labelled data. SimSiam architecture to learn useful feature representations with extensions that incorporate personalisation for local client models. The learnt representation are evaluated against a KNN classifier for analysing their usefulness.

abstractText: Federated Learning (FL) is transforming the ML training ecosystem from a centralized over-the-cloud setting to distributed training over edge devices in order to strengthen data privacy, reduce data migration costs, and break regulatory restrictions. An essential, but rarely studied, challenge in FL is label deficiency at the edge. This problem is even more pronounced in FL, compared to centralized training, due to the fact that FL users are often reluctant to label their private data and edge devices do not provide an ideal interface to assist with annotation. Addressing label deficiency is also further complicated in FL, due to the heterogeneous nature of the data at edge devices and the need for developing personalized models for each user. We propose a self-supervised and personalized federated learning framework, named (SSFL), and a series of algorithms under this framework which work towards addressing these challenges. First, under the SSFL framework, we analyze the compatibility of various centralized self-supervised learning methods in FL setting and demonstrate that SimSiam networks performs the best with the standard FedAvg algorithm. Moreover, to address the data heterogeneity at the edge devices in this framework, we have innovated a series of algorithms that broaden existing supervised personalization algorithms into the setting of self-supervised learning including perFedAvg, Ditto, and local fine-tuning, among others. We further propose a novel personalized federated self-supervised learning algorithm, Per-SSFL, which balances personalization and consensus by carefully regulating the distance between the local and global representations of data. To provide a comprehensive comparative analysis of all proposed algorithms, we also develop a distributed training system and related evaluation protocol for SSFL. Using this training system, we conduct experiments on a synthetic non-I.I.D. dataset based on CIFAR-10, and an intrinsically non-I.I.D. dataset GLD-23K. Our findings show that the gap of evaluation accuracy between supervised learning and unsupervised learning in FL is both small and reasonable. The performance comparison indicates that representation regularization-based personalization method is able to outperform other variants. Ablation studies on SSFL are also conducted to understand the role of batch size, non-I.I.D.ness, and the evaluation protocol.

authors:
H. Brendan McMahan
Eider Moore
Daniel Ramage
Seth Hampson
Blaise Agüera y Arcas

title: Communication-Efficient Learning of Deep Networks from Decentralized Data

venue: arXiv e-prints, art

authors:
Jane Bromley
James W Bentz
Léon Bottou
Isabelle Guyon
Yann LeCun
Cliff Moore
Eduard Säckinger
Roopak Shah

title: Signature verification using a “siamese” time delay neural network

venue: International Journal of Pattern Recognition and Artificial Intelligence,

authors:
Han Cai
Chuang Gan
Ligeng Zhu
Song Han

title: Tinytl: Reduce memory, not parameters for efficient on-device learning

venue: Advances in Neural Information Processing Systems,

authors:
Mathilde Caron
Ishan Misra
Julien Mairal
Priya Goyal
Piotr Bojanowski
Armand Joulin

title: Unsupervised learning of visual features by contrasting cluster assignments

venue: arXiv preprint arXiv:2006.09882,

title: Unsupervised learning of visual features by contrasting cluster

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: In International conference on machine learning,

authors:
Xinlei Chen
Kaiming He

title: Exploring simple siamese representation learning

title: Sinkhorn distances: Lightspeed computation of optimal transportation distances

authors:
Alireza Fallah
Aryan Mokhtari
Asuman Ozdaglar

title: Personalized federated learning: A metalearning approach

venue: arXiv preprint arXiv:2002.07948,

authors:
Jean-Bastien Grill
Florian Strub
Florent Altché
Corentin Tallec
Pierre H Richemond
Elena Buchatskaya
Carl Doersch
Bernardo Avila Pires
Zhaohan Daniel Guo
Mohammad Gheshlaghi Azar

title: Bootstrap your own latent: A new approach to self-supervised learning

venue: arXiv preprint arXiv:2006.07733,

authors:
Chaoyang He
Songze Li
Jinhyun So
Mi Zhang
Hongyi Wang
Xiaoyang Wang
Praneeth Vepakomma
Abhishek Singh
Hang Qiu
Li Shen
Peilin Zhao
Yan Kang
Yang Liu
Ramesh Raskar
Qiang Yang
Murali Annavaram
Salman Avestimehr

title: Fedml: A research library and benchmark for federated machine learning

venue: arXiv preprint arXiv:2007.13518,

authors:
Andrew G Howard
Menglong Zhu
Bo Chen
Dmitry Kalenichenko
Weijun Wang
Tobias Weyand
Marco Andreetto
Hartwig Adam

title: Mobilenets: Efficient convolutional neural networks for mobile vision applications

venue: arXiv preprint arXiv:1704.04861,

authors:
Tzu-Ming Harry Hsu
Hang Qi
Matthew Brown

title: Measuring the effects of non-identical data distribution for federated visual classification

authors:
Sohei Itahara
Takayuki Nishio
Yusuke Koda
Masahiro Morikura
Koji Yamamoto

title: Distillationbased semi-supervised federated learning for communication-efficient collaborative training with non-iid private data

venue: arXiv preprint arXiv:2008.06180,

authors:
Wonyong Jeong
Jaehong Yoon
Eunho Yang
Sung Ju Hwang

title: Federated semi-supervised learning with inter-client consistency

venue: arXiv preprint arXiv:2006.12097,

authors:
far
O. Koyejo
Tancrède Lepoint
Yang Liu
Prateek Mittal
M. Mohri
R. Nock
A. Özgür
R. Pagh
Mariana Raykova
Hang Qi
D. Ramage
R. Raskar
D. Song
Weikang Song
Sebastian U. Stich
Ziteng Sun
A.T. Suresh
Florian Tramèr
Praneeth Vepakomma
Jianyu Wang
Li Xiong
Zheng Xu
Qiang Yang
F. Yu
Han Yu
Sen Zhao

title: Advances and open problems in federated learning

venue: Found. Trends Mach. Learn.,

authors:
Peter Kairouz
H Brendan McMahan
Brendan Avent
Aurélien Bellet
Mehdi Bennis
Arjun Nitin Bhagoji
Keith Bonawitz
Zachary Charles
Graham Cormode
Rachel Cummings

venue: arXiv preprint arXiv:1912.04977,

title: Pseudo-label : The simple and efficient semi-supervised learning method for deep neural networks

authors:
Tian Li
Anit Kumar Sahu
Manzil Zaheer
Maziar Sanjabi
Ameet Talwalkar
Virginia Smith

title: Federated optimization in heterogeneous networks

venue: arXiv preprint arXiv:1812.06127,

authors:
Tian Li
Shengyuan Hu
Ahmad Beirami
Virginia Smith

title: Ditto: Fair and robust federated learning through personalization

authors:
Xinle Liang
Yang Liu
Jiahuan Luo
Yuanqin He
Tianjian Chen
Qiang Yang

title: Self-supervised cross-silo federated neural architecture search

venue: arXiv preprint arXiv:2101.11896,

authors:
Yi Liu
Xingliang Yuan
Ruihui Zhao
Yifeng Zheng
Yefeng Zheng

title: Rc-ssfl: Towards robust and communication-efficient semi-supervised federated learning system

venue: arXiv preprint arXiv:2012.04432,

authors:
Zewei Long
Liwei Che
Yaqing Wang
Muchao Ye
Junyu Luo
Jinze Wu
Houping Xiao
Fenglong Ma

title: Fedsemi: An adaptive federated semi-supervised learning framework

authors:
Brendan McMahan
Eider Moore
Daniel Ramage
Seth Hampson
Blaise Aguera y Arcas

title: Communication-efficient learning of deep networks from decentralized data

venue: In Artificial Intelligence and Statistics,

authors:
Takeru Miyato
Shin-Ichi Maeda
Masanori Koyama
Shin Ishii

title: Virtual adversarial training: A regularization method for supervised and semi-supervised learning

authors:
Takeru Miyato
Shin-ichi Maeda
Masanori Koyama
Shin Ishii

title: Virtual adversarial training: a regularization method for supervised and semi-supervised learning

venue: IEEE transactions on pattern analysis and machine intelligence,

authors:
Sashank Reddi
Zachary Charles
Manzil Zaheer
Zachary Garrett
Keith Rush
Jakub Konečnỳ
Sanjiv Kumar
H Brendan McMahan

title: Adaptive federated optimization

venue: arXiv preprint arXiv:2003.00295,

authors:
Aaqib Saeed
Flora D Salim
Tanir Ozcelebi
Johan Lukkien

title: Federated self-supervised learning of multisensor representations for embedded intelligence

venue: IEEE Internet of Things Journal,

authors:
Tim Salimans
Diederik P. Kingma

title: Weight normalization: A simple reparameterization to accelerate training of deep neural networks

authors:
Mingxing Tan
Quoc Le

title: Efficientnet: Rethinking model scaling for convolutional neural networks

venue: In International Conference on Machine Learning,

authors:
Jianyu Wang
Qinghua Liu
Hao Liang
Gauri Joshi
H Vincent Poor

title: Tackling the objective inconsistency problem in heterogeneous federated optimization

venue: arXiv preprint arXiv:2007.07481,

authors:
Hang Qi
Sashank J. Reddi
Peter Richtárik
K. Singhal
Virginia Smith
M. Soltanolkotabi
Weikang Song
A.T. Suresh
Sebastian U. Stich
Ameet S. Talwalkar
Hongyi Wang
Blake E. Woodworth
Shanshan Wu
Felix X. Yu
Honglin Yuan
M. Zaheer
Mi Zhang
Tong Zhang
Chunxiang Zheng
Chen Zhu
Wennan Zhu

title: A field guide to federated optimization

authors:
Zhirong Wu
Yuanjun Xiong
Stella Yu
Dahua Lin

title: Unsupervised feature learning via nonparametric instance-level discrimination, 2018

authors:
Fengda Zhang
Kun Kuang
Zhaoyang You
Tao Shen
Jun Xiao
Yin Zhang
Chao Wu
Yueting Zhuang
Xiaolin Li

title: Federated unsupervised representation learning

venue: arXiv preprint arXiv:2010.08982,

authors:
Zhengming Zhang
Yaoqing Yang
Zhewei Yao
Yujun Yan
Joseph E Gonzalez
Michael W Mahoney

title: Improving semi-supervised federated learning by reducing the gradient diversity of models

venue: arXiv preprint arXiv:2008.11364,

authors:
Yuchen Zhao
Hanyang Liu
Honglin Li
Payam Barnaghi
Hamed Haddadi

title: Semi-supervised federated learning for activity recognition

venue: arXiv preprint arXiv:2011.00851,

heading: 1 INTRODUCTION

text:
Federated Learning (FL) is a contemporary distributed machine learning paradigm that aims at strengthening data privacy, reducing data migration costs, and breaking regulatory restrictions (Kairouz et al., 2021; Wang et al., 2021). It has been widely applied to computer vision, natural language processing, and data mining. However, there are two main challenges impeding its wider adoption in machine learning. One is data heterogeneity, which is a natural property of FL in which diverse clients may generate datasets with different distributions due to behavior preferences (e.g., the most common cause of heterogeneity is skewed label distribution which might result from instances where some smartphone users take more landscape pictures, while others take more photos of daily life). The second challenge is label deficiency at the edge, which is relatively less studied. This issue is more severe at the edge than in a centralized setting because users are reluctant to annotate their private and sensitive data, and/or smartphones and IoT devices do not have a user-friendly interface to assist with annotation.
To mitigate the data heterogeneity issue among clients, researchers have proposed algorithms for training a global model FedAvg (McMahan et al., 2017), FedProx (Li et al., 2018), FedNova (Wang et al., 2020), FedOPT (Reddi et al., 2020), as well as personalized FL frameworks (e.g., pFedMe, Ditto, Per-FedAvg). These algorithms all depend on the strong assumption that the data at the edge has sufficient labels. To address the label deficiency issue in FL, recent works (Liu et al., 2020; Long et al., 2020; Itahara et al., 2020; Jeong et al., 2020; Liang et al., 2021; Zhao et al., 2020; Zhang et al., 2020a;b) assume that the server or client has a fraction of labeled data and use semi-supervised methods such as consistency loss (Miyato et al., 2018) or pseudo labeling (Lee, 2013) to train a global model. A more realistic but challenging setting is fully unsupervised training. Although a recent work in FL (Saeed et al., 2020) attempts to address this challenge through Siamese networks proposed around thirty years ago (Bromley et al., 1993), its design does not tackle data heterogeneity for learning personalized models, and it only trains on small-scale sensor data in IoT devices. Moreover, these existing works in FL have not examined recent progress in the Self-Supervised Learning (SSL) community where methods such as SimCLR (Chen et al., 2020), SwAV(Caron et al., 2021), BYOL (Grill et al., 2020), and SimSiam (Chen & He, 2020) have shown tremendous improvement in reducing the amount of labeled data required to achieve state-of-the-art performance. As such, it remains still unclear how these SSL methods can be incorporated into FL and how well they would perform, especially when intertwined with the data heterogeneity challenge that does not exist in centralized training.
In this paper, we propose Self-Supervised Federated Learning (SSFL), a unified self-supervised and personalized federated learning framework, and a series of algorithms under this framework to address these challenges. As shown in Figure 1, this framework brings state-of-the-art self-supervised learning algorithms to the realm of FL in order to enable training without using any supervision, while also integrating model personalization to deal with data heterogeneity (Section 3.1). More specifically, under the SSFL framework, we analyze the compatibility of various centralized selfsupervised learning methods in the FL setting and demonstrate that SimSiam networks performs the best with the standard FedAvg algorithm (Section 3.2). Moreover, to address the data heterogeneity at edge devices, we have innovated a series of algorithms that broaden the reach of existing supervised personalization algorithms into the setting of self-supervised learning, including perFedAvg (Fallah et al., 2020), Ditto (Li et al., 2021), and local fine-tuning, among others. We further propose a novel personalized federated self-supervised learning algorithm, per-SSFL (Section 3.3), which balances personalization and consensus by carefully regulating the distance between the local and global representations of data (shown as the yellow block in Figure 1).
To provide a comprehensive and comparative analysis of the proposed algorithms, we also develop a distributed training system and evaluation protocol for SSFL. Using this training system, we conduct experiments on a synthetic non-I.I.D. dataset based on CIFAR-10 and a natural non-I.I.D. dataset GLD-23K. Our experimental results demonstrate that all algorithms in our framework work reliably. In FL, the gap of evaluation accuracy between supervised learning and unsupervised learning is small. Personalized SSFL performs better than FedAvg-based SSFL. We also conduct ablation studies to fully understand the SSFL framework, namely the role of batch size, different degrees of non-I.I.D.ness, and performance in more datasets. Finally, our unified API design can serve as a suitable platform and baseline, enabling further developments of more advanced SSFL algorithms.

heading: 2 PRELIMINARIES

text: SSFL builds upon two fundamental areas in machine learning: federated optimization and selfsupervised learning. Thus, we first introduce some basics and formulations in these areas.

heading: 2.1 FEDERATED OPTIMIZATION

text: Federated optimization refers to the distributed optimization paradigm that a network of K devices collaboratively solve a machine learning task. In general, it can be formulated as a distributed optimization problem with the form (McMahan et al., 2017): minθ ∑K k=1 |Dk| |D| L(θ,Dk). Here, each device k has a local dataset Dk drawn from a local distribution Xk. The combined dataset D = ∪Kk=1Dk is the union of all local datasets Dk. θ represents the model weight of a client model. L is the client’s local loss function that measures the local empirical risk over the heterogeneous dataset Dk. Under this formulation, to mitigate the non-I.I.D. issue, researchers have proposed algorithms such as FedProx (Li et al., 2018), FedNova (Wang et al., 2020), and FedOPT (Reddi et al., 2020) for training a global model, as well as personalized FL frameworks such as Ditto (Li et al., 2021), and Per-FedAvg (Fallah et al.). All of these algorithms have a strong assumption that data at the edge have sufficient labels, meaning that their analysis and experimental settings are based on a supervised loss function, such as the cross-entropy loss for image classification.

heading: 2.2 SELF-SUPERVISED LEARNING

text:
Self-supervised learning (SSL) aims to learn meaningful representations of samples without human supervision. Formally, it aims to learn an encoder function fθ : X 7→ Rd where θ is the parameter of the function, X is the unlabeled sample space (e.g. image, text), and the output is a d dimensional vector containing enough information for downstream tasks such as image classification and segmentation. The key to SSL’s recent success is the inductive bias that ensures a good representation encoder remains consistent under different perturbations of the input (i.e. consistency regularization).
One prominent example among recent advances in modern SSL frameworks is the Siamese architecture (Bromley et al., 1993) and its improved variants SimCLR (Chen et al., 2020), SwAV (Caron et al., 2021), BYOL (Grill et al., 2020), and SimSiam (Chen & He, 2020). Here we review the most elegant architecture, SimSiam, and defer the description and comparison of the other three to Appendix A. SimSiam proposes a two-head architecture in which two different views (augmentations) of the same image are encoded by the same network fθ. Subsequently, a predictor Multi Layer Perceptron (MLP) hθ and a stop-gradient operation denoted by ·̂ are applied to both heads. In the SSL context, “stop gradient” means that the optimizer stops at a specific neural network layer during the back propagation and views the parameters in preceding layers as constants. Here, θ is the concatenation of the parameters of the encoder network and the predictor MLP. The algorithm aims to minimize the negative cosine similarity D(·, ·) between two heads. More concretely, the loss is defined as
LSS(θ,D) = 1 |D| ∑ x∈D D(fθ(T (x)), ¤ hθ(fθ(T (x)))), (1)
where T represents stochastic data augmentation and D is the data set.

heading: 3 SSFL: SELF-SUPERVISED FEDERATED LEARNING

text: In this section, we propose SSFL, a unified framework for self-supervised federated learning. Specifically, we introduce the method by which SSFL works for collaborative training of a global model and personalized models, respectively.

heading: 3.1 GENERAL FORMULATION

text:
We formulate self-supervised federated learning as the following distributed optimization problem:
min Θ
{θk}k∈[K]
G (L (θ1,Θ;X1) , . . . ,L (θK ,Θ;XK)) (2)
where θk is the parameter for the local model (fθk , hθk); Θ is the parameter for the global model (fΘ, hΘ); L(θk,Θ;Xk) is a loss measuring the quality of representations encoded by fθk and fΘ on the local distribution Xk; and G(·) denotes the aggregation function (e.g. sum of client losses weighted by |Dk||D| ). To capture the two key challenges in federated learning (data heterogeneity and label deficiency), we hold two core assumptions in the proposed framework: (1) Xk of all clients are heterogeneous (non-I.I.D.), and (2) there is no label.
To tackle the above problem, we propose a unified training framework for federated self-supervised learning, as described in Algorithm 1. This framework can handle both non-personalized and personalized federated training. In particular, if one enforces the constraint θk = Θ for all clients k ∈ [K], the problem reduces to learning a global model. When this constraint is not enforced, θk can be different for each client, allowing for model personalization. ClientSSLOPT is the local optimizer at the client side which solves the local sub-problem in a self-supervised manner. ServerOPT takes the update from the client side and generates a new global model for all clients.
Algorithm 1: SSFL: A Unified Framework for Self-supervised Federated Learning input :K,T,Θ(0), {θ(0)k }k∈[K], CLIENTSSLOPT , SERVEROPT
1 for t = 0, . . . , T − 1 do 2 Server randomly selects a subset of devices S(t) 3 Server sends the current global model Θ(t) to S(t) 4 for device k ∈ S(t) in parallel do 5 Solve local sub-problem of equation 2:
θk,Θ (t) k ← CLIENTSSLOPT (θ (t) k ,Θ (t),∇L (θk,Θ;Xk))
6 Send ∆(t)k := Θ (t) k −Θ (t) back to server 7 Θ(t+1) ← SERVEROPT Ä Θ(t), {∆(t)k }k∈S(t) ä return :{θk}k∈[K],Θ(T )
Next, we will introduce specific forms of ClientSSLOPT and ServerOPT for global training and personalized training.

heading: 3.2 GLOBAL-SSFL: COLLABORATION TOWARDS A GLOBAL MODEL WITHOUT SUPERVISION

text:
To train a global model using SSFL, we design a specific form of ClientSSLOPT using SimSiam. We choose SimSiam over other contemporary self-supervised learning frameworks (e.g., SimSiam, SwAV, BYOL) based on the following analysis as well as experimental results (see Section 5.1).
The simplicity in neural architecture and training method. SimSiam’s architecture and training method are relatively simple. For instance, compared with SimCLR, SimSiam has a simpler loss function; compared with SwAV, SimSiam does not require an additional neural component (prototype vectors) and Sinkhorn-Knopp algorithm; compared with BYOL, SimSiam does not need to maintain an additional moving averaging network for an online network. Moreover, the required batch size of SimSiam is the smallest, making it relatively friendly for resource-constrained federated learning. A more comprehensive comparison can be found in Appendix A.
Interpretability of SimSiam leads to simpler local optimization. More importantly, SimSiam is more interpretable from an optimization standpoint which simplifieds the local optimization. In particular, it can be viewed as an implementation of an Expectation-Maximization (EM) like algorithm, meaning that optimizing LSS in Equation 1 is implicitly optimizing the following objective
min θ,η E T x∼X
î ‖fθ(T (x))− ηx‖22 ó . (3)
Here, fθ is the encoder neural network parameterized by θ. η is an extra set of parameters, whose size is proportional to the number of images, and ηx refers to using the image index of x to access a sub-vector of η. This formulation is w.r.t. both θ and η and can be optimized via an alternating algorithm. At time step t, the ηtx update takes the form η t x ← ET [fθt(T (x))], indicating that ηtx is assigned the average representation of x over the distribution of augmentation. However, it is impossible to compute this step by going over the entire dataset during training. Thus, SimSiam uses one-step optimization to approximate the EM-like two-step iteration by introducing the predictor hθ
to approximate η and learn the expectation (i.e. hθ (z) ≈ ET [fθ(T (x))]) for any image x. After this approximation, the expectation ET [·] is ignored because the sampling of T is implicitly distributed across multiple epochs. Finally, we can obtain the self-supervised loss function in Equation 1, in which negative cosine similarity D is used in practice (the equivalent L2 distance is used in Equation 3 for the sake of analysis). Applying equation 1 as ClientSSLOPT simplifies the local optimization for each client in a self-supervised manner.

heading: 3.3 PER-SSFL: LEARNING PERSONALIZED MODELS WITHOUT SUPERVISION

text:
In this section, we explain how SSFL addresses the data heterogeneity challenge by learning personalized models. Inspired by the interpretation in Section 3.2, we define the following sub-problem for each client k ∈ [K]:
min θk,ηk E T x∼Xk
ï ‖fθk(T (x))− ηk,x‖ 2 2 + λ
2 ‖ηk,x −H∗x‖ 2 2 ò s.t. Θ∗,H∗ ∈ arg min
Θ,H n∑ i=1 |Dk| |D| E T x∼Xi
î ‖fΘ(T (x))−Hx‖22 ó (4) Compared to global training, we additionally include Θ, the global model parameter, andH, the global version of η, and the expected representations which correspond to the personalized parameters θk and ηk. In particular, through the term ‖ηk,x −H∗x‖ 2 2, we aim for the expected local representation of any image x to reside within a neighborhood around the expected global representation of x. Therefore, by controlling the radius of the neighborhood, hyperparameter λ helps to balance consensus and personalization.
Algorithm 2: Per-SSFL input :K,T, λ,Θ(0), {θ(0)i }k∈[K], s: number of local iteration, β: learning rate
1 for t = 0, . . . , T − 1 do 2 Server randomly selects a subset of devices S(t) 3 Server sends the current global model Θ(t) to S(t) 4 for device k ∈ S(t) in parallel do 5
CLIENTSSLOPT
Sample mini-batch Bk from local dataset Dk, and do s local iterations /* Optimize the global parameter Θ locally */
6 Z1, Z2 ← fΘ(t)(T (Bk)), fΘ(t)(T (Bk)) 7 P1, P2 ← hΘ(t)(Z1), hΘ(t)(Z2) 8 Θ
(t) k ← Θ (t) − β∇Θ(t) D(P1,Ẑ2)+D(P2,Ẑ1) 2 , where ·̂ stands for stop-gradient
/* Optimize the local parameter θk */ 9 z1, z2 ← fθk (T (Bk)), fθk (T (Bk))
10 p1, p2 ← hθk (z1), hθk (z2) 11 θk ← θk − β∇θk
Ä D(p1,”z2)+D(p2,”z1)
2 + λD(p1,P1)+D(p1,P2)+D(p2,P1)+D(p2,P2) 4 ä 12 Send ∆(t)k := Θ (t) k −Θ (t) back to server
13 SERVEROPTΘ (t+1) ← Θ(t) + ∑ k∈S(t) |Dk| |D| ∆ (t) k
return :{θi}i∈[n],Θ(T )
We see that Equation 4 in the above objective is an optimization problem w.r.t. both θ and η. However, as the above target is intractable in practice, following an analysis similar to Section 3.2, we use the target below as a surrogate:
min θk
LSS (θk, Dk) + λ |Dk| ∑ x∈Dk D (hθk (fθk (T (x))) , hΘ∗ (fΘ∗ (T (x)))) (5)
s.t. Θ∗ ∈ arg min Θ LSS (Θ, D) (6)
In practice, Θ can be optimized independently of θk through the FedAvg (McMahan et al., 2017) algorithm. To make the computation more efficient, we also apply the symmetrization trick proposed in (Chen & He, 2020). We refer to this algorithm as Per-SSFL and provide a detailed description in Algorithm 2 (also illustrated in Fig. 1).
Regarding the theoretical analysis. To our knowledge, all self-supervised learning frameworks do not have any theoretical analysis yet, particularly the SimSiam dual neural network architecture.
Our formulation and optimization framework are interpretable, they are built based on an EM-like algorithm for SimSiam and minimizing the distance between the private model and the global model’s data representation.
Innovating baselines to verify SSFL. Note that we have not found any related works that explore a Siamese-like SSL architecture in an FL setting. As such, to investigate the performance of our proposed algorithm, we further propose several other algorithms that can leverage the SSFL framework. 1. LA-SSFL. We apply FedAvg (McMahan et al., 2017) on the SimSiam loss LSS for each client to obtain a global model. We perform one step of SGD on the clients’ local data for local adaption; 2. MAML-SSFL. This algorithm is inspired by perFedAvg (Fallah et al.) and views the personalization on each devices as the inner loop of MAML (Finn et al.). It aims to learn an encoder that can be easily adapted to the clients’ local distribution. During inference, we perform one step of SGD on the global model for personalization; 3. BiLevel-SSFL. Inspired by Ditto (Li et al., 2021), we learn personalized encoders on each client by restricting the parameters of all personalized encoders to be close to a global encoder independently learned by weighted aggregation. More details of these algorithms, formulation, and pseudo code are introduced in Appendix B. In Section 5.3, we will show the comparison results for these proposed SSFL algorithmic variants.

heading: 4 TRAINING SYSTEM AND EVALUATION PIPELINE FOR SSFL

heading: A Distributed Training System to accelerate the algorithmic exploration in SSFL framework.

text:
We also contributed to reproducible research via our distributed training system. This is necessary for two reasons: (1) Running a stand-alone simulation (training client by client sequentially) like most existing FL works requires a prohibitively long training time when training a large number of clients. In SSFL, the model size (e.g., ResNet-18 v.s. shallow CNNs used in the original FedAvg paper) and the round number for convergence (e.g., 800 epochs in the centralized SimSiam framework) is relatively larger than in FL literature. By running all clients in parallel on multiple CPUs/GPUs, we can largely accelerate the process. (2) Given that SSFL is a unified and generic learning framework, researchers may develop more advanced ways to improve our work. As such, we believe it is necessary to design unified APIs and system frameworks in line with the algorithmic aspect of SSFL. See Appendix C for more details on our distributed training system.
Evaluation Pipeline. In the training phase, we use a KNN classifier (Wu et al., 2018) as an online indicator to monitor the quality of the representations generated by the SimSiam encoder. For Global-SSFL, we report the KNN test accuracy using the global model and the global test data, while in Per-SSFL, we evaluate all clients’ local encoders separately with their local test data and report their averaged accuracy. After self-supervised training, to evaluate the performance of the trained encoder, we freeze the encoder and attach a linear classifier to the output of the encoder. For Global-SSFL, we can easily verify the performance of SimSiam encoder by training the attached linear classifier with FedAvg. However, for Per-SSFL, each client learns a personalized SimSiam encoder. As the representations encoded by personalized encoders might reside in different spaces, using a single linear classifier trained by FedAvg to evaluate these representations is unreasonable (see experiments in Section 5.4.3). As such, we suggest an additional evaluation step to provide a more representative evaluation of Per-SSFL’s performance: for each personalization encoder, we use the entire training data to train the linear classifier but evaluate on each client’s local test data.

heading: 5 EXPERIMENTS

text:
In this section, we introduce experimental results for SSFL with and without personalization and present a performance analysis on a wide range of aspects including the role of batch size, different degrees of non-IIDness, and understanding the evaluation protocol.
Implementation. We develop the SSFL training system to simplify and unify the algorithmic exploration. Details of the training system design can be found in Appendix C. We deploy the system in a distributed computing environment which has 8 NVIDIA A100-SXM4 GPUs with sufficient memory (40 GB/GPU) to explore different batch sizes (Section 5.4.1). Our training framework can run multiple parallel training workers in a single GPU, so it supports federated training with a large number of clients. The client number selected per round used in all experiments is 10, which is a reasonable setting suggested by recent literature (Reddi et al., 2020).
Learning Task. Following SimCLR (Chen et al., 2020), SimSiam (Chen & He, 2020), BYOL (Grill et al., 2020), and SwAV (Caron et al., 2020) in the centralized setting, we evaluate SSL for the image classification task and use representative datasets for federated learning.
Dataset. We run experiments on synthetic non-I.I.D. dataset CIFAR-10 and intrinsically non-I.I.D. dataset Google Landmark-23K (GLD-23K), which are suggested by multiple canonical works in the FL community (Reddi et al., 2020; He et al., 2020; Kairouz et al., 2019). For the non-I.I.D. setting, we distribute the dataset using a Dirichlet distribution (Hsu et al., 2019), which samples pc ∼ Dir(α) (we assume a uniform prior distribution) and allocates a pc,k proportion of the training samples of class c to local client k. We provide a visualization of the data distribution in Appendix E.1.
Model Architecture. For the model architecture, ResNet-18 is used as the backbone of the SimSiam framework, and the predictor is the same as that in the original paper.
Next, we focus on results from the curated CIFAR-10 dataset and defer GLD-23K to Appendix D.

heading: 5.1 COMPARISONS ON SIMSIAM, SIMCLR, SWAV, AND BYOL

text: Our first experiment determines which SSL method is ideal for FL settings. We run experiments using FedAvg for these four methods and obtain two findings: (1) SimSiam outperforms SimCLR in terms of accuracy; (2) BYOL and SwAV do not work in FL; we tested a wide range of hyper-parameters, but they still are unable to converge to a reasonable accuracy. These experimental results confirm our analysis in Section 3.2.

heading: 5.2 EVALUATION ON GLOBAL-SSFL

text:
The goal of this experiment is to understand the accuracy gap between supervised and self-supervised federated learning in both I.I.D. and non-I.I.D. settings where we aim to train a global model from private data from clients.
Setup and Hyper-parameters. We evaluate Global-SSFL using non-I.I.D. data from CIFAR-10: we set α = 0.1 for the Dirichlet distribution. For supervised learning, the test accuracy is evaluated on a global test dataset. For self-supervised training, we follow the evaluation protocol introduced in Section 4. We use SGD with Momentum as the client-side optimizer and a learning rate scheduler across communication rounds. We searched the learning rate on the grid of {0.1, 0.3, 0.01, 0.03} and report the best result. The experiment is run three times using the same learning rate with fixed random seeds to ensure reproducibility. The training lasts for 800 rounds, which is sufficient to achieve convergence for all methods. More hyperparameters are in Appendix E.2.
We display the training curves in Figure 10 which demonstrates that SSFL can converge reliably in both I.I.D. and non-I.I.D. settings. For the I.I.D. data, we find that SSFL can achieve the same accuracy as the centralized accuracy report in the SimSiam paper (Chen & He, 2020). For the non-I.I.D. data, SSFL achieves a reasonable accuracy compared to the centralized accuracy. The accuracy comparisons in different dimensions (supervised v.s. self-supervised; I.I.D. v.s. non-I.I.D.) are summarized in Table 1.

heading: 5.3 EVALUATION ON PER-SSFL

text:
Based on the results of SSFL with FedAvg, we further add the personalization components for SSFL introduced in Section 3.3 (Per-SSFL).
Setup and Hyper-parameters. For a fair comparison, we evaluate Per-SSFL on non-I.I.D. data from CIFAR-10 and set α = 0.1 for the Dirichlet distribution. For Per-SSFL training, we follow the evaluation protocol introduced in Section 4. Similar to SSFL, we use SGD with Momentum as the client-side optimizer and the learning rate scheduler across communication rounds. We search for the learning rate on a grid of {0.1, 0.3, 0.01, 0.03} and report the best result. For Per-SSFL and BiLevel-SSFL, we also tune the λ of the regularization term with a search space {1, 10, 0.1, 0.01, 0.001}. The experiments are run three times with the same learning rate and with fixed random seeds to ensure reproducibility. The training also lasts for 800 communication rounds, which is the same as Global-SSFL. Other hyperparameters can be found in Appendix E.2.
We illustrate our results in Figure 10 and Table 2. To confirm the convergence, we draw loss curves for all methods in Figure 10(b) (note that they have different scaled values due to the difference of their loss functions). Figure 10(b) indicates that Per-SSFL performs best among all methods. MAML-SSFL is also a suggested method since it obtains comparable accuracy. LA-SSFL is a practical method, but it does not perform well in the self-supervised setting. In Figure 10(b), the averaged personalized accuracy of LA-SSFL diverges in the latter phase. Based on BiLevel-SSFL’s result, we can conclude such a method is not a strong candidate for personalization, though it shares similar objective functions as Per-SSFL. This indicates that regularization through representations encoded by SimSiam outperforms regularization through weights.

heading: 5.4 PERFORMANCE ANALYSIS

text:
5.4.1 ROLE OF BATCH SIZE
FL typically requires a small batch size to enable practical training on resource-constrained edge devices. Therefore, understanding the role of batch size in SSFL is essential to practical deployment and applicability. To investigate this, we use different batch sizes and tune the learning rate to find the optimal accuracy for each one. The results in Figure 4 show that SSFL requires a large batch size (256); otherwise, it reduces the accuracy or diverges during training. Since system efficiency is not the focus of this paper, we use gradient accumulation, which is a simple yet effective method. We fix the batch size at 32 and use accumulation step 8 for all experiments. For an even larger batch size (e.g., 512), the memory cost is significant, though there is no notable gain in accuracy. Therefore, we discontinue the search for batch sizes larger than larger than 256. A more advanced method includes batch-size-one training
and knowledge distillation. We defer the discussion to Appendix F.

heading: 5.4.2 ON DIFFERENT DEGREES OF NON-I.I.D.NESS

text:
We investigate the impact of the degree of data heterogeneity on the SSFL performance. We compare the performance between α = 0.1 and α = 0.5. These two settings provide a non-negligible gap
in the label distribution in each client (see our visualization in Appendix E.1). Figure 5(a) and 5(b) shows the learning curve comparisons. It is clearly observed that the higher degree of data heterogeneity makes it converge more slowly, adversely affecting the accuracy.
0 20 40 60 80 100
round 0.3
0.4
0.5
0.6
0.7
0.8
Per-SSFL
LA-SSFL
Accuracy
Figure 6: Understanding the Evaluation Protocol

heading: 5.4.3 UNDERSTANDING THE LINEAR EVALUATION OF PERSONALIZED ENCODERS

text: As we discussed in 4, in SSFL, we can easily verify the quality of the SimSiam encoder using federated linear evaluation; however, in Per-SSFL, each client learns a personalized SimSiam encoder. Such heterogeneity in diverse encoders makes a fair evaluation difficult. To demonstrate this, we run experiments with naive federated linear evaluation on personalized encoders and surprisingly find that such an evaluation protocol downgrades the performance. As shown in Figure 6, the federated linear evaluation for Per-SSFL performs worse than even LA-SSFL. This may be attributed to the fact that the naive aggregation drags close to the parameter space of all heterogeneous encoders, making the encoder degenerate in terms of personalization.

heading: 6 RELATED WORKS

text: Federated Learning (FL) with Personalization. pFedMe (Dinh et al.), perFedAvg (Fallah et al., 2020), and Ditto (Li et al., 2021) are some representative works in this direction. However, these methods all have a strong assumption that users can provide reliable annotations for their private and sensitive data, which we argue to be very unrealistic and impractical. Label deficiency in FL. There are a few related works to tackle label deficiency in FL (Liu et al., 2020; Long et al., 2020; Itahara et al., 2020; Jeong et al., 2020; Liang et al., 2021; Zhao et al., 2020; Zhang et al., 2020b). Compared to these works, our proposed SSFL does not use any labels during training. FedMatch (Jeong et al., 2020) and FedCA (Zhang et al., 2020a) requires additional communication costs to synchronize helper models or public labeled dataset. (Saeed et al., 2020) addresses the fully unsupervised challenge on small-scale sensor data in IoT devices. However, compared to our work, it uses the Siamese networks proposed around thirty years ago (Bromley et al., 1993), lacking consideration on the advance in the past two years (i.e., SimCLR (Chen et al., 2020), SwAV(Caron et al., 2021), BYOL (Grill et al., 2020), and SimSiam (Chen & He, 2020)). Moreover, these works does not have any design for learning personalized models.

heading: 7 CONCLUSION

text: We propose Self-supervised Federated Learning (SSFL) framework and a series of algorithms under this framework towards addressing two challenges: data heterogeneity and label deficiency. SSFL can work for both global model training and personalized model training. We conduct experiments on a synthetic non-I.I.D. dataset based on CIFAR-10 and the intrinsically non-I.I.D. GLD-23K dataset. Our experimental results demonstrate that SSFL can work reliably and achieves reasonable evaluation accuracy that is suitable for use in various applications.