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: In this work, the authors propose a new federated learning algorithm by adopting a neural scheduling technique. In particular, the neural scheduler is trained without needing to access the local data. Preliminary experiments demonstrate the effectiveness of the proposed algorithm.