AIMHI: Protecting Sensitive Data through Federated Co-Training
Abstract
Federated learning offers collaborative training among distributed sites without sharing sensitive local information by sharing the sites’ model parameters.
It is possible, though, to make non-trivial inferences about sensitive local information from these model parameters.
We propose a novel co-training technique called
AIMHI that uses a public unlabeled dataset to exchange information between sites by sharing predictions on that dataset.
This setting is particularly suitable to healthcare, where hospitals and clinics hold small labeled datasets with highly sensitive patient data and large national health databases contain large amounts of public patient data.
We show that the proposed method reaches a model quality comparable to federated learning while maintaining privacy to high degree.
Introduction
Can we collaboratively train models from distributed sensitive datasets while maintaining data privacy at a level required in healthcare?
Federated learning [12] allows distributed sites, e g., hospitals or clinics, to collaboratively train a joint model without directly disclosing their sensitive data by instead periodically sharing model parameters.
An attacker or curious observer can, however, make inferences about local data from model parameters [11] and model updates [21].
Differential privacy provides a rigorous and measurable privacy guarantee [5] that can be achieved by perturbing parameters appropriately [20].
This perturbation, however, can reduce model quality, resulting in a trade-off between privacy and quality.
As we show in our experiments, even with substantial perturbation one can infer membership of a training sample [17] with high probability, i.e., whether a data point is present in a local dataset from the model parameters shared in federated learning.
We propose to instead use a distributed co-training approach [9], where sites train local models and exchange predictions on a shared unlabeled dataset, instead of sharing model parameters.
By forming a consensus from shared predictions, one obtains pseudo-labels for the shared unlabeled dataset that can be used for local training.
Iterating this process improves the consensus, and thereby the quality of pseudo-labels, effectively nudging local models to come to an agreement.
We show in our experiments that this approach, which we call AIMHI (AI Models for Healthcare Improvement), achieves the same model quality as federated learning (FedAvg [12]), but protects privacy to a high level—membership inference is significantly less likely compared to vanilla federated learning and federated learning with differential privacy.
These results indicate that, given a public unlabeled dataset, this approach constitutes a more favorable trade-off between privacy and model quality.
The AIMHI approach requires a large shared unlabeled dataset, and is in spirit similar to distributed distillation [2].
While such unlabeled public datasets are not always available, in healthcare, large public health databases are quite common: the US NCHS databases, the UK’s NHS databases, the
UK Biobank [18], the MIMIC-III database [6], or the planned European EHDS contain vast amounts of patient data that can be used in many application scenarios.
Our contributions are
(i) a novel distributed co-training approach to collaboratively train models from privacysensitive distributed data sources, and (ii) a preliminary empirical evaluation of model quality and privacy on the CIFAR10 benchmark dataset, indicating high model quality and a substantial improvement in privacy.
Preliminaries
We assume learning algorithms
A : X × Y → H that trains a model h ∈ H using a dataset
D ⊂ X
× Y from an input space X and output space Y , i.e., h =
A ( D ) .
Given a set of m ∈ N clients with local datasets D 1 , . . .
, D m ⊂ X × Y drawn iid from a data distribution D and a loss function ℓ : Y × Y → R , the goal is to find a single model h ∗ ∈ H that minimizes the risk
[Equation]
In centralized learning, the datasets are pooled as D =
(cid:83)
i ∈ [ m ] D i and A is applied to D , usually to minimize the empirical risk
[Equation]
In federated learning (FL) we assume that the learning algorithm is iterative [cf.
Chp.
2.1.4 7], i.e.,
A : X × Y × H → H that updates a model h t +
=
A ( D, h t ) .
In this case, centralized learning means applying A to D until convergence.
Note that applying A on D can be the application to any random subset, e.g., as in mini-batch training, and convergence is measured in terms of low training loss, small gradient, or small deviation from previous iterate.
In standard federated learning [12], A is applied in parallel for b ∈ N rounds on each client locally to produce local models h 1 , . .
, h m .
These models are then centralized and aggregated using an aggregation operator a : H m → H , i.e., h = a ( h 1 , . .
h m ) .
The aggregated model h is then redistributed to local clients which perform another b rounds of training using h as a starting point.
This is iterated until convergence of h with the goal to minimize the empirical risk over all local datasets[12], i.e.,
[Equation]
When aggregating by averaging, this method is also known as federated averaging.
Next, we describe our proposed distributed co-training approach.
Algorithm
Input: communication period b , learning algorithm A , m clients with local datasets
D 1 , . . .
, D m , unlabeled shared dataset U , total number of rounds T
Output: final models h 1 T , . . .
, h mT 1 initialize local models h 10 , . . .
, h m2 P ← ∅ 3 Locally at client i at time t do h it ←
A ( D i ∪ P, h it − 1 ) if t % b =
b − 1 then
L i ← h
it ( U ) send L i to server receive L from server
P ← ( U, L ) end 11 At server at time t do receive local pseudo-labels L 1 , . . .
, L m
L ← consensus ( L 1 , . . .
, L m ) send L to all clients
AIMHI: Distributed Co-Training
We propose a semi-supervised, distributed co-training approach that collaboratively trains models via sharing predictions. It uses an unlabeled dataset U , producing pseudo-labels L for it by forming a consensus of the predictions of all local models. Unlabeled data and pseudo-labels form an additional public, shared dataset P that is combined with local data for training.
The details are described in
Alg. 1: at each client i , the local model is updated using the local dataset D i combined with the current pseudo-labeled public dataset P .
The updated model is used to produce improved pseudo-labels
L i for the unlabeled data U , which are sent to a server every b rounds.
At the server, as soon as all local prediction L 1 , . . .
, L m are received, a consensus L is formed and broadcasted back to the clients.
Forming a consensus is similar to obtaining a prediction from an ensemble [4].
For our classification experiments, we use vanilla majority voting [3].
Note that more elaborate consensus mechanisms offer a rich design space for improvements.
On receiving the new consensus labels L from the server, the client updates the public pseudo-labeled dataset P and performs another iteration of local training.
We perform a preliminary empirical evaluation of the performance of AIMHI in comparison to vanilla federated learning (i.e., FedAvg) and federated learning with differential privacy on the CIFAR image classification dataset [10].
For that, we measure their prediction accuracy on a test set, as well as their privacy vulnerability.
4.1 Privacy Vulnerability
We measure privacy vulnerability by performing membership inference attacks against AIMHI and FL according to two different attack scenarios per approach.
In both attacks, the attacker creates an attack model using a model it constructs from its training and test datasets.
Similar to previous work [17], we assume that the training data of the attacker has a similar distribution to the training data of the client.
Once the attacker has its attack model, it uses this model for membership inference.
In blackbox attacks (in which the attacker does not have access to intermediate model parameters), it only uses the classification scores it receives from the target model (i.e., client’s model) for membership inference.
On the other hand, in whitebox attacks (in which the attacker can observe the intermediate model parameters), it can use additional information in its attack model.
Since the proposed AIMHI does not reveal intermediate model parameters to any party, it is only subject to blackbox attacks.
Vanilla federated learning on the other hand is subject to whitebox attacks.
Each inference attack produces a [Figure] [Figure] membership score of a queried data point, indicating the likelihood of the data point being a member of the training set.
We measure the success of membership inference as ROC AUC of these scores.
The vulnerability (VUL) of a method is the ROC AUC of membership attacks over K runs over the entire training set (also called attack epochs) according to the attack model and scenario.
A vulnerability of 1 . 0 means that membership can be inferred with certainty,
whereas 0 .
5 means that deciding on membership is a random guess.
We assume the following attack model: clients are honest and the server may be semi-honest (follow the protocol execution correctly, but it may try to infer sensitive information about the clients).
The main goal of a semi-honest server is to infer sensitive information about the local training data of the clients.
This is a stronger attacker assumption compared to a semi-honest client since the server receives the most amount of information from the clients during the protocol, and a potential semi-honest client can only obtain indirect information about the other clients.
We also assume that parties do not collude.
The attack scenarios are 1 :
• AIMHI realistic : the attacker can send a (forged) unlabeled dataset to the clients and observe their predictions, equivalent to one attack epoch ( K = 1 ); • FL realistic : the attacker receives model parameters and can run an arbitrary number of attacks—we use K = 500 attack epochs; • AIMHI relaxed : the attacker can send a (forged) unlabeled dataset in each communication round, albeit to different models in each round—we simulate this by assuming pessimistically that models are sufficiently similar over all rounds and set K = T/b = 150 ; • FL relaxed : the attacker cannot copy model parameters on their machine, but can only perform an attack epoch in each communication round, thus being able to perform K = T/b = 150 attack epochs after T = 3000 rounds.
To measure vulnerability, we use the ML Privacy Meter tool [13].
This tool allows us to quantify the privacy risks associated with machine learning models by performing a range of membership inference attacks [14] on machine learning models and measuring attack success.
It simulates different levels of access and model knowledge for attackers, e.g., limiting attackers to only predictions, or loss values, or assuming they have access to the model’s parameters.
For our experiments, we assume that for AIMHI the attacker only has access to the predictions (i.e., output of the last layer), while for
FL, the attacker can access all layers.
4.2 Experimental Setup
We use the common CIFAR10 dataset [10] which consists of 50000 training and 10000 test images with 10 classes.
We use 10000 samples drawn iid from the training images as training set,images as unlabeled dataset and the 10000 test images as test set.
We use m = 5 clients, each with a local training set of n = 2000 samples.
Clients use a convolutional neural network (details are provided in the Appendix) and communicate every b = 20 rounds for FL and AIMHI.
Both FL and
AIMHI are run for T = 3000 rounds.
The experiments are implemented in TensorFlow [1], the code is publicly available 2 .
For our experiments, we use a simple CNN architecture given in Table 1.
As optimizer we use Adam with learning rate α = 0 .
001 , exponential decay rates β 1 = 0 .
9 and β 2 = 0 .
999 , and ϵ =
− 7 .
[Table]
4.3 Differential Privacy for Federated Learning
A common defense against membership inference attacks is applying appropriate clipping and noise before sending models.
This guarantees ϵ, δ -differential privacy for local data [20] at the cost of a slight-to-moderate loss in model quality.
This technique is also proven to defend against backdoor and poisoning attacks [19].
We compare AIMHI against federated learning with differential privacy through the Gaussian mechanism (DP-FL).
For that, each client i first clips their parameters w i to
[Equation]
for a constant C > 0 , and then add Gaussian noise ˜ w i = w ic
+ ϵ with ϵ ∼ N (0 , σ ) for σ ≥ 0 .
The level of privacy depends on the choice of C and σ [20].
4.4 Results
The results presented in Table 2 show that AIMHI and FL achieve comparable test accuracy after
T = 3000 rounds 3 .
Note that centralized training on the 10000 training samples achieves a test accuracy of 0 .
661 , so FL and AIMHI both achieve virtually optimal model quality.
At the same time, FL is vulnerable to membership attacks, both in the realistic (VUL) and the relaxed (VUL(rel.)) attack scenarios with a vulnerability of over 0 .
Note that since the attacker can run an arbitrary number of attacks, the privacy vulnerability is in principle 1 . 0 , we instead measure the practical vulnerability under a large, but limited number of attacks.
AIMHI on the other hand preserves privacy ( V UL = 0 . 51 ) in the realistic attack scenario, and still has very little vulnerability ( 0 . 56 ) under the relaxed scenario.
Adding noise for differential privacy reduces vulnerability considerably to around 0 . 85 , but at the cost of accuracy which drops to around 0 .
We investigate the convergence in terms of test accuracy in Figure 1: Both AIMHI and FL converge quickly, with AIMHI converging slightly faster.
Looking at the development of privacy vulnerability over time, we observe that vulnerability for FL is already high after the first communication round (i.e., after t = 20 rounds), as shown in
Figure 2, and increases to 0 .
94 after t = 1000 rounds.
This holds also for the relaxed scenario, although with slightly less vulnerability.
For AIMHI realistic instead the vulnerability remains low [Table] 5 to 0 . 51 ).
For AIMHI relaxed we observe an increase in vulnerability, as expected, since we assume in this scenario that the attacker performs a number of attacks proportional to the number of communication rounds.
Discussion and Conclusion
We presented a novel distributed co-training approach for privacy-preserving federated learning that protects sensitive local datasets, where vanilla federated learning is susceptible to privacy attacks, at the same time achieving similar model quality as FL.
Our initial experiments support the hypothesis that distributed co-training can be competitive with FL, given a large unlabeled dataset, while preserving data privacy to a much higher degree.
Co-training requires a shared unlabeled dataset, which is not available in all application scenarios.
In healthcare, however, it is not uncommon to have large quantities of unlabeled data points available.
A limitation of AIMHI is that local datasets must be sufficiently large to create useful local models [cf. 8]—otherwise, the poor quality pseudo-labeled dataset will not improve local training.
Choosing more elaborate consensus methods [15, 16] is an interesting direction for future work that can improve performance even with small local datasets.
An important advantage of co-training is that local models do not have to have the same architecture, as FL requires.
In fact, local models can be arbitrary.
This includes interpretable models, like decision trees or rule ensembles, for which no FL method exists so far.
Exploring distributed co-training for such model classes could open a whole new avenue for federated learning.