Towards Scheduling Federated Deep Learning using Meta-Gradients for Inter-Hospital Learning

1 INTRODUCTION . Federated learning is a field that has emerged recently due to the abundance of data available today and the risks that this poses to individuals . Privacy ( particularly of personal information ) is of great importance and should be protected by researchers working in machine learning . In parallel , the emergence of electronic health records ( EHRs ) has allowed the digitisation of much personal information pertaining to the health conditions of individuals . EHRs are often used in many machine learning research projects Shailaja et al . ( 2018 ) . This often involves the transfer and storage of very sensitive information which increases the risk of data leakage . As a result , federated learning can minimise this risk by utilising the data at it ’ s source rather than transferring it to the researchers servers to be processed . Machine learning researchers working with EHR data will be all too familiar with the difficulty of gaining access to this information in the first place . It can be a very lengthy and exhausting process ( for good reason ) to gain access and utilise the EHRs of one healthcare institution let alone accessing the datasets of many . Federated learning offers an alternative in that it allows the data of the patients to be utilised while reducing the risk of their privacy being compromised . Federated learning is not without its own limitations however . The different datasets that are stored in the different nodes may have different underlying distributions due to their data collection processes which can make machine learning across multiple domains difficult . There is also the possibility of data at each node being corrupted either maliciously or accidentally , leading to data that is undesirable to use for training . These issues can all lead to difficulties in training and convergence of the overall model being trained . To overcome these limitations , in this work we propose the following setup . Firstly , We use federated learning to i ) protect the privacy of patients by minimising the movement of their data and ii ) improve the quality of our machine learning model by utilising a diverse dataset sourced from different hospitals . As we are ‘ blind ’ to the data at the nodes , we propose the use of a student-teacher network setup . The teacher ( a reinforcement learning agent ) will have access to the local servers and be able to select the appropriate data for the ‘ student ’ ( our model ) to be trained on at that time . The ‘ scheduler ’ will be responsible for directing the teacher to a given data centre to select data for training on . To summarise : ( The Student : ) - the machine learning model we are training . ( The Teacher : ) - a reinforcement learning agent that selects data from a data centre based on the state of the student . This essentially defines a curriculum at each training step for the student . ( The Scheduler : ) - directs the teacher to the appropriate data centre for training . This is also based on the state of the student at each training iteration . Section 2 discusses the related work that has been carried out and Section 3 provides a detailed description of how our algorithm works . Section 4 details the datasets we benchmark our method against . We then present the results in Section 5 and discuss their significance and interesting behaviours of our model in Section 6 . 2 RELATED WORK . Federated learning has been used by researchers to exploit larger pools of data for training Bonawitz et al . ( 2019 ) , preserve data privacy Xu et al . ( 2019 ) and distributing computational resource requirements Yuan & Ma ( 2020 ) . Federated learning has also been used for healthcare applications to simultaneously utilise multiple datasets to train a model on patient data . In this work we create a model that learns in a federal fashion through the interaction of a scheduler that is trained using meta-gradients and a student-teacher algorithm that is trained using reinforcement learning . Meta-learning has been used effectively in Such et al . ( 2019 ) where the loss of a student model on a validation set was used as a signal to update the weights of a generative model . This work demonstrated the rapid and effective training of methods that exploit meta-gradients . Meta-learning was also used in Zahavy et al . ( 2020 ) , where the meta-gradients are used to tune the parameters of an actor-critic algorithm . As a result of the efficacy of this method in these domains we choose to use meta-gradients in order to schedule which data centre the gradients to update are student model will come from . The meta-learned scheduler chooses a node representing a data centre where a student-teacher algorithm is used to sample data . Student-teacher algorithms have been used in multiple works , with the general premise that one algorithm ( teacher ) is trained to train another ( student ) Fan et al . ( 2018 ) ; Liu et al . ( 2017a ) . These methods have also been used with a curriculum Bengio et al . ( 2009 ) , where the curriculum is either pre-defined and exploited by the teacher El-Bouri et al . ( 2020 ) or implicitly learned by the teacher during training Graves et al . ( 2017 ) . Federated learning is a method of training a model ( in our case a deep neural network ) by using data from multiple centres , without having central access to each of them McMahan et al . ( 2017 ) . Local models at each of the data centres are iteratively updated and aggregated to form a global model . At each round of iteration , a central coordinator samples a subset , m , of local models , Sm , and sends them the current global model Gt . Each member of Sm then updates this global model using their local data to create an updated model Lt+1 . These models are then aggregated and are sent back to update the global model as : Gt+1 = Gt + η n m∑ i=1 ( Lt+1i −G t ) ( 1 ) where n is the number of nodes ( i.e. , data centres ) and η acts as a learning rate for replacing the global model with the aggregate of the local models . While this has been shown to work in many cases , Wang et al . ( 2020a ) make the argument that there is inherent difficulty in updating neural networks in this manner . They argue that the permutation invariance of the summation operand renders averaging in the parameter space a naive approach . For meaningful averaging to be done , the permutation must first be undone . 2.1 COMPROMISING FEDERATED LEARNING . One of the vulnerabilities of federated learning is that nodes being compromised can significantly affect the training of the global model Bhagoji et al . ( 2018 ) . Attacks of these sort can either ‘ poison ’ the data found at one of the nodes ( known as an adversarial attack ) or bias the model that is trained at one of the nodes significantly , leading it to highly skew the aggregation step Tolpegin et al . ( 2020 ) . There is also the possibility of the attack being a single-shot attack or a repeated attack Fang et al . ( 2020 ) . In the single shot case , only one of the nodes is compromised whereas in the repeated case , multiple nodes can be compromised at any given time . Many works have been produced in discussing how federated learning can be compromised by introducing a backdoor into the training process Gu et al . ( 2017 ) ; Bagdasaryan et al . ( 2020 ) . A backdoor is an attack that causes a classifier to produce unexpected behaviour if a specific trigger is added to the input space . An example is a sticker being added to an image and associating this with the incorrect label Gu et al . ( 2017 ) . Defences against these attacks have been developed with some authors using pruning of redundant neurons for the core classification task Liu et al . ( 2018 ) , using outlier detection to detect potential triggers Wang et al . ( 2019 ) , and re-training and preprocessing inputs Liu et al . ( 2017b ) . In this work we aim to overcome these limitations and build defence into the training procedure through the use of a student-teacher network that actively selects which data to train on . 3 METHODOLOGY . Our method is comprised of three agents in the training setup , the student , the teacher and the scheduler . 3.1 THE OVERALL SETUP . The overall setup of our federated learning training routine is as follows . We have a scheduler that controls which node we will be learning from ( this can be one-hot or we can select multiple nodes ) . The teacher at the node can then select a batch of data according to the state of the student . The student at the node is a copy of the global student . We use the student to forward pass the batch of data selected by the teacher and return the loss . In the one-hot scheduler scenario , we send back the loss to the global student model to update the weights via backpropagation . In the multi-node learning scenario , we aggregate the losses from all nodes selected and feed these back to the global model for updating . 3.2 DATA PREPROCESSING . The first step we must take in order to exploit our teacher setup is to rank our data according to some metric . Using Wang et al . ( 2020b ) as a guide , we choose to use the Mahalanobis distance expressed as : d ( xn ) = ( ( xn − µ ) T S−1 ( xn − µ ) ) 1 2 ( 2 ) for medical datasets , and the cosine similarity as our similarity metric for image datasets . As the tabular data found in electronic health record systems consist of multiple data types , we encode these using a denoising autoencoder . This trained encoder is distributed to all the nodes so that the data in each node is processed in the same way for consistency . 3.3 THE TEACHER . For the student-teacher interaction we follow the setup in El-Bouri et al . ( 2020 ) . The task of the teacher is to select a batch of data from the curriculum by selecting the index along the curriculum and the ‘ width ’ around that index to include in the selection . The following sequential steps are implemented : • The data at each node is organised into N curriculum batches according to some metric H . • The teacher selects one or more batches for feeding into the student . • A pre-trained autoencoder is used to create a latent representation of the batch . • The student is trained on this batch and it ’ s performance on a separate validation set is recorded . In this work we use the Mahalanobis distance for H for medical data , and cosine similarity for image data as summarised in Wang et al . ( 2020b ) . The teacher is a reinforcement learning agent and therefore is tasked with minimising the Bellman loss function given by : L ( θi ) = ( r + γmax a′ Q ( s′ , a′ ; θ−i ) −Q ( s , a ; θi ) ) 2 ( 3 ) where r is the reward of a state-action ( s , a ) tuple , γ the discount factor , Q is the q-value defining the value of taking an action given a state and θ and θ− are the parameters of the prediction and target ( the version of the teacher that is held constant for K steps to stabilise training as described in Mnih et al . ( 2015 ) ) networks respectively . As we also choose to use an actor-critic setup for the teacher , the action space and Q-function are separately parameterised . This allows a continuous action space and the actor that selects actions is updated using the following loss : ∇θµJ ≈ 1 N ∑ i ∇aQ ( s , a | θQ ) | s=si , a=µ ( si ) ∇θµµ ( s | θ µ ) ( 4 ) where Q is the Q-function and µ is the policy . The teacher can either be pre-trained , or jointly trained with the scheduler . The teacher can also either be trained on the dataset of one node and distributed to the rest or independently trained at each node . The latter is preferable due to the ability of the teacher to adapt to the dataset at hand . However , the former is useful when not all nodes in the federated system have access to computational power . The intuition is that the curriculum strategy learned by the teacher should be general for the task at hand and thereby provide strong performance .

Summary: This paper proposes a novel federated learning framework for model training across multiple hospital data centers. There are mainly three components in the framework including a student machine learning model, a teacher reinforcement learning agent, and a scheduler algorithm that directs the teacher to specific data centers. The setup and the details of the whole framework are well presented in the paper. Corresponding experiments are carried out to show the performance of the proposed method.