FL Games: A Federated Learning Framework for Distribution Shifts
Abstract
Federated learning aims to train predictive models for data that is distributed across clients, under the orchestration of a server.
However, participating clients typically each hold data from a different distribution, which can yield to catastrophic generalization on data from a different client, which represents a new domain.
In this work, we argue that in order to generalize better across non-i.i.d.
clients, it is imperative to only learn correlations that are stable and invariant across domains.
We propose FL G AMES , a game-theoretic framework for federated learning that learns causal features that are invariant across clients.
While training to achieve the Nash equilibrium, the traditional best response strategy suffers from high-frequency oscillations.
We demonstrate that FL G AMES effectively resolves this challenge and exhibits smooth performance curves.
Further, FL G AMES scales well in the number of clients, requires significantly fewer communication rounds, and is agnostic to device heterogeneity.
Through empirical evaluation, we demonstrate that FL G AMES achieves high out-of-distribution performance on various benchmarks.
Introduction
With the rapid advance in technology and growing prevalence of smart devices, Federated Learning (FL) has emerged as an attractive distributed learning paradigm for machine learning models over networks of computers Kairouz et al.
[2019], Li et al.
[2020], Bonawitz et al.
In FL, multiple sites with local data, often known as clients , collaborate to jointly train a shared model under the orchestration of a central hub called the server while keeping their data private.
While FL serves as an attractive alternative to centralized training because the client data does not need to move to the server, statistical heterogeneity is a key challenge in its optimization.
While one of the most popular algorithms in this setup, Federated Averaging (F ED AVG) McMahan et al.
[2017] delivers huge communication gains in i.i.d.
(independent and identically distributed) setting, its performance on non-i.i.d.
clients is an active area of research.
As shown by Karimireddy et al.
[2020], client heterogeneity has direct implications on the convergence of F ED AVG since it introduces a drift in the updates of each client with respect to the server model.
While recent works Li et al.
Karimireddy et al.
[2020], Yu et al.
[2019], Wang et al.
[2020], Li et al.
[2020], Lin et al.
[2020], Li and Wang [2019], Zhu et al.
[2021] have tried to address client heterogeneity through constrained gradient optimization and knowledge distillation, most did not tackle the underlying distribution shift.
These methods can at best generalize to interpolated domains and fail to extrapolate well, i.e., generalize to new extrapolated domains 1 .
Over the past year, there has been a surge in interest in bringing the machinery of causality into machine learning Arjovsky et al.
[2019], Ahuja et al.
[2020], Schölkopf [2019], Ahuja et al.
Parascandolo et al.
[2020], Robey et al.
[2021], Krueger et al.
[2021], Rahimian and Mehrotra [2019].
However, despite their success, they suffer from key limitations which are unsuitable for deployment in a real-world setup.
Since FL typically consists of a large number of clients, it is natural for data at each client to represent different annotation tools, measuring circumstances, experimental environments, and external interventions.
Inspired by this idea and by the recent progress in causal machine learning, we draw connections between OOD generalization and robustness across heterogeneous clients in FL.
To date, only two scientific works Francis et al.
[2021], Tenison et al.
[2021] have incorporated the learning of invariant predictors in order to achieve strong generalization in FL.
The former adapts masked gradients Parascandolo et al.
[2020] and the latter builds on IRM to exploit invariance and improve leakage protection in FL.
While IRM lacks theoretical convergence guarantees, failure modes of Parascandolo et al.
[2020] like formation of dead zones and high sensitivity to small perturbations Shahtalebi et al.
[2021] are also issues when it is applied in FL, rendering it unreliable.
In this study, we consider IRM G AMES Ahuja et al.
[2020] from the OOD generalization literature since its formulation shows resemblance to the standard FL setup.
However, as discussed above, IRM
G AMES too encounters a few fundamental challenges not just specific to FL but also in a generic
ML framework.
We take a step towards fixing these limitations and addressing the challenge of client heterogeneity under distribution shifts in FL from a causal viewpoint.
Specifically, we propose
Federated Learning Games (FL G AMES ) for learning causal representations which are stable across clients.
We summarize the our main contributions below.
• We propose a new framework called FL G AMES for learning causal representations that are invariant across clients in a federated learning setup.
• The underlying sequential game theoretic framework in IRM G AMES causes the time complexity of FL algorithm to scale linearly with the number of clients.
Inspired from the game theory literature, we equip our algorithm to allow parallel updates across clients, further resulting in superior scalability.
• IRM G AMES exhibits large oscillations in the performance metrics as the training progresses, making it difficult to define a valid stopping criterion.
Using ensembles over client’s historical actions, we demonstrate that FL G AMES appreciably smoothens these oscillations.
• The convergence rate of IRM G AMES is slow and hence directly impacts systems with speed or communication cost as a primary bottleneck.
By increasing the local computation at each client, we show that FL G AMES exhibits high communication efficiency •
Empirically, we show that the performance of the invariant predictors found by our approach on unseen OOD clients improves significantly over state-of-art prior works.
Background 2.1 Causality in Machine Learning
Consider a multi-source domain generalization task, where the goal is to learn a robust set of parameters that generalize well to unseen (test) environments E all ⊃ E tr , given a set of m training domains (or environments) E tr .
A popular algorithm, Invariant Risk Minimization (IRM) Arjovsky et al.
[2019] does so by defining the Invariance principle
Definition 1 A data representation ϕ : X → Z elicits an invariant predictor w ◦ ϕ across environments E tr if there is a classifier w : Z → Y simultaneously optimal for all environments i.e.
[Equation]
Invariant Risk Minimization Games (IRM G AMES ) propose a game theoretic formulation for finding ( ϕ, w ) that satisfy the invariance principle.
It endows each environment with its own predictor w k ∈ H w and aims to train an ensemble model w av ( z ) =
[Equation]
|
[Equation]
tr | (cid:80) |
[Equation]
tr | k =1 w k ( z ) for each z ∈ Z s.t.
w av satisfies the following optimization problem
[Equation]
where H ϕ , H w are the hypothesis sets for feature extractors and predictors, respectively.
The constraint in equation 1 is equivalent to the Nash equilibrium of a game with each environment k as a player with action w k , playing to maximize its utility
− R k ( w av , ϕ ) .
The resultant game is solved using the best response dynamics (BRD) with clockwise updates (for more details, refer to the supplement) and is referred to as V-IRM G AMES .
Fixing ϕ to an identity map in V-IRM G AMES is also shown to be very effective and is called F-IRM G AMES .
Federated Learning Games (FL G AMES )
OOD generalization is often typified using the notion of data-generating environments.
Arjovsky et al.
[2019] formalizes an environment as a data generating distribution representing a particular location, time, context, circumstances and so forth.
This concept of data-generating environments can be related to FL by considering each client as producing data generated from a different environment.
However despite this equivalence, existing OOD generalization techniques can’t be directly applied to
FL.
Apart from the FL-specific challenges, these approaches also suffer from several key limitations in non-FL domains Rosenfeld et al.
[2020], Nagarajan et al.
[2020], further rendering them unfit for practical deployment.
Thus developing causal inference models for FL which are inspired from invariant prediction in OOD generalization, are bound to inherit the failures of the latter.
In this work, we consider one such popular OOD generalization technique, IRM G AMES as it formulation bears close resemblance to a standard FL setup.
However, as discussed, IRM G AMES too suffers from various challenges, which impede its deployment in a generic ML framework, specifically in FL.
In the following section, we elaborate on each of these limitations and discuss the corresponding modifications required to overcome them.
Further, inspired from the game theoretic formulation of IRM G AMES , we propose FL G AMES which forfeits it’s failures and can recover the causal mechanisms of the targets, while also providing robustness to changes in the distribution.
3.1 Challenges in Federated Learning
Data Privacy.
Consider a FL system with m client devices, S = { 1 , 2 , ..., m } .
Let N k denote the number of data samples at client device k , and
D k = { ( x ki , y ki ) }
N k i =1 as it’s labelled dataset.
The constraint of each environment in IRM G AMES can be used to formulate the local objective of each client.
In particular, each client k ∈ S now serves as a player, competing to learn w k ∈ H w by optimizing its local objective, i.e. w k ∈ arg min
w
[Equation]
k
[Equation]
w R k
[Equation]
∈ H
w
However, the upper-level objective of IRM G AMES requires optimization over the dataset pooled together from all environments.
Centrally hosting the data at the server or sharing it across clients contradicts the objective of FL.
Hence, we propose using F ED SGD McMahan et al.
[2017] to optimize ϕ .
Specifically, each client k computes and broadcasts g k = ∇ R k (cid:16) 1 |S| ( w ′ k
+ (cid:80) q ̸ = k w k ) ◦ ϕ (cid:17) to the server, which then aggregates these gradients and applies the update rule ϕ t +1 =
ϕ t − η (cid:80) k ∈S N k N g k where g k is computed over C % batches locally and N = (cid:80) k ∈ S N k .
We call this approach V-FL
G AMES .
The variant with ϕ =
I is called F-FL G AMES .
Sequential dependency As discussed, IRM G AMES poses IRM objective as finding the Nash equilibrium of an ensemble game across environments and adopts the classic best response dynamics (BRD) algorithm to compute it.
This approach is based on playing clockwise sequences wherein players take turns in a fixed cyclic order, with only one player being allowed to change its action at any given time t (details in the supplement).
This linear scaling of time complexity with the number of players poses a major challenge solving the game in FL.
We modify the classic BRD algorithm by allowing simultaneous updates at any given round t .
However, a client now best responds to the optimal actions played by it’s opponent clients in round t − 1 instead of t .
We refer to this approximation of BRD in F-FL G AMES and V-FL G AMES as parallelized F-FL G AMES and, parallelized V-FL G AMES respectively.
Oscillations.
As demonstrated in Ahuja et al.
[2020], when a neural network is trained using the IRM
G AMES objective (equation 1), the training accuracy initially stabilizes at a high value and eventually starts to oscillate.
The explanation for these oscillations attributes to the significant difference among the data of the training environments.
With model’s performance metrics oscillating to and from at each step, defining a reasonable stopping criterion becomes challenging.
As shown in various game theoretic literature Herings and Predtetchinski [2017], Barron et al.
[2010], Fudenberg et al.
[1998], Ge et al.
[2018], BRD can often oscillate.
Computing the Nash equilibrium for general games is non-trivial and is only possible for a certain class of games (e.g., concave games)
Zhou et al.
Thus, rather than alleviating oscillations completely, we propose solutions to reduce them significantly to better target valid stopping points.
We propose a two-way ensemble approach wherein apart from maintaining an ensemble across clients ( w av ), each client k also responds to the ensemble of historical models (memory) of its opponents.
Formally, we maintain queues (a.k.a. buffer) at each client which store its historically played actions.
In each iteration, a client best responds to a uniform distribution over the past strategies of its opponents.
The global objective at the server remains unchanged.
Mathematically, the new local objective of each client k ∈ E tr can be stated as
[Equation]
where B q denotes the buffer at client q and w qj denotes the j th historical model of client q .
As the buffer reaches its capacity, it is renewed based on first in first out (FIFO) manner.
Note that this approach doesn’t result in any communication overhead since a running sum over historical strategies can be calculated in O (1) time by maintaining a prefix sum.
This variant is called F-FL G AMES (S MOOTH ) or V-FL G AMES (S MOOTH ) based on the constraint on ϕ .
Convergence speed.
As discussed, FL G AMES has two variants: F-FL G AMES and, V-FL G AMES with the former being an approximation of the latter ( ϕ = I ).
While both the approaches exhibit superior performance on a variety of benchmarks, the latter has shown its success in a variety of large scale tasks like language modeling.
Peyrard et al.
Despite being theoretically grounded,
V-IRM G AMES suffers from slower convergence due to an additional round for optimization of ϕ .
Typically, clients (e.g. mobile devices) posses fast processors and computational resources and have datasets that are much smaller compared to the total dataset size.
Hence, utilizing additional local computation is essentially free compared to communicating with the server.
To improve the efficiency of our algorithm, we propose replacing the stochastic gradient descent (SGD) over ϕ by a full-batch gradient descent (GD).
Intuitively, now at each gradient step, the resultant ϕ takes large steps in the direction of its global optimum, resulting in fast training.
Note that the classifiers at each client are still updated over one mini-batch.
This variant of FL G AMES is referred to V-FL G AMES (F AST ).
Experiments and Results
In Ahuja et al.
[2020], IRM G AMES was tested over a variety of benchmarks including C OLORED
MNIST, C OLORED F ASHION MNIST and C OLORED D SPRITES dataset.
We utilize the same datasets for our experiments.
Additionally, we create another benchmark, S PURIOUS CIFAR10.
Details on each of the datasets can be found in the supplement.
We report the mean performance of various baselines over 5 runs.
We compare these algorithmic variants across fixed and variable ϕ separately as shown in the Table 1.
Clearly, across all benchmarks, the FL baselines F ED SGD, F ED AVG McMahan et al.
F ED BN Li et al.
[2021] and F ED P ROX Li et al.
[2020] are unable to generalize to the test set, with
F ED BN exhibiting superior performance compared to the others.
Intuitively, these approaches latch onto the spurious features to make prediction, hence leading to poor generalization over novel clients.
From Table
[Table]
we observe that all modifications in FL G AMES individually achieve high testing accuracy, hence eliminating the spurious correlations unlike state-of-the-art FL techniques.
Further, on C OLORED D SPRITES , we discover that the mean performance of V-FL G AMES is superior to all algorithms with fixed representation ( ϕ = I ).
This accentuates the importance of V-FL G AMES over
F-FL G AMES , especially over complex and larger datasets where learning ϕ becomes imperative
[Table] [Table]
In all the above experiments, both of our end approaches: parallelized V-FL G AMES (S MOOTH +F AST ) and parallelized F-FL G AMES (S MOOTH ) are able to perform better than or at par with the other variants.
These algorithms were primarily designed to overcome the challenges faced by causal FL systems while retaining their original ability to learn causal features.
Hence, the benefits provided by these approaches in terms of 1) robust predictions; 2) scalability; 3) fewer oscillations and 4) fast convergence are not at the cost of performance.
While the former is demonstrated by
Table 1, the latter three are detailed in Section 4.2.
4.1 Interpretation of Learned Features
In order to explain the predictions of our parametric model, we use LIME Ribeiro et al.
[2016] to learn an interpretable model around each prediction.
Figure 1(a) shows LIME masks for both F ED AVG (ii) and our method (iii).
Clearly, the former focuses solely on the background to make the model prediction.
However, the latter uses causal features in the image (like shape, stroke, edges and curves) along with some pixels from the background (noise) to make predictions.
This demonstrates the reasoning behind robust generalization of FL GAMES as opposed to state-of-the-art FL techniques.
4.2 Ablation Analysis
We analyze the effect of each of our algorithmic modifications on the C OLORED MNIST dataset.
The results on other datasets are similar and are provided in the supplement.
4.2.1 Effect of Simultaneous BRD
We examine the effect of replacing the classic best response dynamics as in Ahuja et al.
[2020] with the simultaneous best response dynamics.
For the same, we use a more practical environment: (a) more clients are involved, and (b) each client has fewer data.
Similar to Choe et al.
[2020], we extended the C OLORED MNIST dataset by varying the number of clients between 2 and 10.
For each setup, we vary the degree of spurious correlation between 70% and 90%) for training clients and merely 10% in the testing set.
A more detailed discussion of the dataset is provided in the supplement.
For F-FL G AMES , it can be observed from Figure 1(b.i), as the number of clients in the FL system increase, there is a sharp increase in the number of communication rounds required to reach equilibrium.
However, the same doesn’t hold true for parallelized F-FL G AMES .
Further, parallelized F-FL G AMES is able to reach a comparable or higher test accuracy as compared to F-FL
G AMES with significantly lower communication rounds (refer to Figure 1(b.ii)).
4.2.2 Effect of a memory ensemble
As shown in Figure 2(b), compared to F-FL G AMES , F-FL G AMES (S MOOTH ) reduces the oscillations significantly.
While in the former, performance metrics oscillate at each step, the oscillations in the later are observed after an interval of roughly 50 rounds.
Further, F-FL G AMES (S MOOTH ) seems to envelop the performance curves of F-FL G AMES .
As a result, apart from reducing the frequency of oscillations, F-FL G AMES (S MOOTH ) also achieves higher testing accuracy compared to F-FL
G AMES .
The observations are consistent across the parallelized variants.
Further, as observed from
Figure 2(a), both sequential and parallel F-FL G AMES (S MOOTH ) significantly reduce the number of rounds required to reach equilibrium, further underscoring the efficacy of our proposed methodology.
4.2.3 Effect of using Gradient Descent (GD) for ϕ
Communication costs are the principal constraints in FL setup.
Edge devices like mobile phones and sensor are bandwidth constrained and require more power for transmission and reception as compared to remote computation.
As observed from Figure 2(c), V-FL G AMES (S MOOTH +F AST ) is able to achieve significantly higher testing accuracy in fewer rounds as compared to V-FL G AMES .
Consistent results are also reported in Figure 2(a), where both sequential and parallel variants of
V-FL G AMES
(S MOOTH +F AST ) result in a significant improvement ( ∼ 2×) in the number of rounds.
Conclusion
In this work, we develop a novel framework based on Best Response Dynamics (BRD) training paradigm to learn invariant predictors across clients in Federated learning (FL).
Inspired from Ahuja et al.
[2020], the proposed method called Federated Learning Games (FL G AMES ) learns causal representations which have good out-of-distribution generalization on new train clients or test clients unseen during training.
We investigate the high frequency oscillations observed using BRD and equip our algorithm with a memory of historical actions.
This results in smoother performance metrics with significantly lower oscillations.
FL G AMES exhibits high communication efficiency as it allows parallel computation, scales well in the number of clients and results in faster convergence.
Given the impact of FL in medical imaging, we plan to test our framework over medical benchmarks.
Future directions include theoretically analyzing the smoothed best response dynamics, as it might have potential implications for other game-theoretic based machine learning frameworks.
A Appendix B Game Theory Concepts
We define some basic game theory notations that will be used later.
Let Γ =
( N, { S k } k ∈ N , { u k } k ∈ N ) be the tuple representing a normal form game, where N denotes the finite set of players.
For each player k , S k = { s k 0 , s k 1 , ...
s km } denotes the pure strategy space with m strategies and u k ( s k , s − k ) denotes the payoff function of player k corresponding to strategy s k .
Here, an environment for player k is s − k , a set containing strategies taken by all players but k and S − k denotes the space of strategies of the opponent players to k .
S =
(cid:81)
i ∈ N S
i denotes the joint strategy set of all players.
A game Γ is said to be finite if S is finite and is continuous if S is uncountably infinite.
While a pure strategy defines a specific action to be followed at any time instance, a mixed strategy of player k , σ k = { p k ( s k 0 ) , p k ( s k 1 ) , ...
p k ( s km ) } is a probability distribution over a set of pure strategies, where (cid:80) mj
=1 p k ( s kj ) = 1 .
The expected utility of a mixed strategy u k ( σ k , σ − k ) for player k is the expected value of the corresponding pure strategy payoff i.e.
[Equation]
where ˜ S k corresponds to the mixed strategy space of player k .
Best response (BR).
A mixed strategy σ ∗ k for player k is said to be a best response to it’s opponent strategies σ − k if
[Equation]
Nash equilibrium.
A mixed strategy profile σ ∗ = { σ ∗ 1 , σ ∗ 2 , ...σ ∗ N } is a Nash equilibrium if for all players k , σ ∗ k is the best response to the strategies played by it’s opponent players i.e σ ∗− k .
Best response dynamics (BRD).
BRD is an iterative algorithm in which at each time step, a player myopically plays strategies that are best responses to the most recent known strategies played by it’s
opponents previously.
Based on the playing sequence across layers, BRD can be classified into three broad categories: BRD with clockwise sequences, BRD with simultaneous updating and BRD with random sequences.
For this study, we focus only on the first two playing schedules.
Let the function seq : N → P [ N ] denote a playing sequence which determines the set of players whose turn it is to play at each time period t ∈ N .
Here, P [ N ] denotes the power set of { 1 , 2 , ...N } players and N be the set of natural numbers { 1 , 2 , ... } .
By BRD, at each time step t , ∀ i ∈ seq ( t ) , action taken by player i i.e. a ti is the best response to it’s current environment i.e a t − i .
• BRD with clockwise sequences:
In this playing sequence, players take turns according to a fixed cyclic order and only one player is allowed to change it’s action at any given time t .
Specifically, the playing sequence is defined by seq ( t )
= 1 + ( t − 1) mod n .
Since only a single player is allowed to play at any given time t , a t -seq ( t ) =
a t − 1 -seq ( t ) .
• BRD with simultaneous updating:
In this playing sequence, seq ( t ) chooses a non empty subset of players to participate in round t .
However, for each player i ∈ seq ( t ) , the optimal action chosen a ti depends on the knowledge of the latest strategy of it’s opponents
C Datasets
The MNIST dataset consists of handwritten digits, with a total of 60,000 images in the training set and 10,000 images in the test set Deng [2012].
These images are black and white in colour and form a subset of the larger collection of digits called NIST.
Each digit in the dataset is normalized in size to centre fit in the fixed size image of size 28 × 28.
It is then anti-aliased to introduce appropriate gray-scale levels.
C. C OLORED MNIST
We modify the MNIST dataset in the exact same manner as in Arjovsky et al.
Specifically,
Arjovsky et al.
[2019] creates the dataset in a way that it contains both the invariant and spurious features according to different causal graphs.
Spurious features are introduced using colors.
Digits less than 5 (excluding 5) are attributed with label 0 and the rest with label 1.
The dataset is divided across three clients, out of which two serve as training and one as testing.
The 60,000 images from
MNIST train set are divided equally amongst the two training clients i.e. each consists of 30, samples.
The testing set contains the 10,000 images from the MNIST test set.
Preliminary noise is added to the label to reduce the invariant correlation.
Specifically, the initial label ( ˜ y ) of each image is flipped with a probability δ k to construct the final label y .
The final label y of each image is further flipped with a probability p k to construct it’s color code ( z ).
In particular, the image is colored red, if z = 1 and green if z = 0 .
The flipping probability which defines the color coding of an image, p k is 0.2 for client 1, 0.1 for client 2 and 0.9 for the test client.
The probability δ k is fixed to 0.
for all clients k .
The above choice is defined in a way that the mean degree of label-color (spurious) correlation ( 1 − p k , ∀ k ) is more than the average degree of invariant correlation ( 1 − δ k , ∀ k ).
A sample batch of images elucidating the above construction is shown in Figure 3.
C. C OLORED F ASHION MNIST
We use the exact same environment for creating C OLORED F ASHION MNIST as in Ahuja et al.
The data generating process of C OLORED F ASHION MNIST is motivated from that of
C OLORED MNIST in a way that it possesses spurious correlations between the label and the colour.
Fashion MNIST consists of images from a variety of sub-categories under the two broad umbrellas of clothing and footwear.
Clothing items include categories like: “t-shirt", “trouser", “pullover", “dress", “shirt" and “coat" while the footwear category includes “sandal", “sneaker", “bag" and “ankle boots".
Similar to C OLORED MNIST, the train dataset is equally split across two clients (30,000 images each) and the entire test set is attributed to the test client.
Preliminary labels for binary classification are constructed such ˜ y = 0 for “t-shirt”, “trouser”, “pullover”, “dress”, “coat”, “shirt” and ˜ y
= 1 : “sandle”, “sneaker” and “ankle boots”.
Next, we add noise to the preliminary label by flipping ˜ y with a probability δ k = 0 .
25 , ∀ k to construct the final label y .
We next flip the final label with a probability p k to designate a color ( z ), with p 1 = 0 .
2 for the first client, p 2 = 0 .
1 for the second client and p 3 = 0 .
9 for the test client.
The image is colored red, if z = 1 and green if z = 0 .
A sample batch of images elucidating the above construction is shown in Figure 4.
C. S PURIOUS CIFAR
In this setup, we modify the CIFAR-10 dataset similar to the C OLORED MNIST dataset.
Instead of coloring the images, we use a different mechanism based on the spatial location of a synthetic feature to generate spurious features.
CIFAR10 dataset consists of 60,000 images from 10 classes including “airplane", “automobile", “bird", “cat", “deer", “dog", “frog", “house", “ship", “truck".
The original dataset is relabelled to create a binary classification task between motor and non-motor objects.
All images corresponding to the label “frog" are discarded to ensure a similar samples count for the two classes.
Similar to C OLORED MNIST, the train dataset is equally split across two clients and the entire test set is attributed to the test client.
Preliminary labels for binary classification are constructed such ˜ y = 0 for “airplane”, “automobile”, “ship”, “truck” and ˜ y
= 1 : “bird”, “cat”, “deer",“dog" and
“horse”.
Next, we add noise to the preliminary label by flipping ˜ y with a probability δ k = 0 .
25 , ∀ k to construct the final label y .
We next flip the final label with a probability p k to designate a positional index ( z ), with p 1 = 0 .
2 for the first client, p 2 = 0 .
1 for the second client and p 3 = 0 .
9 for the test client.
This index defines the spatial location of a 5 × 5 black patch in the image.
An index value of 0 ( z = 0 ) specifies the patch at the top left corner of the image, while an index value of 1 ( z = 1 ) corresponds to a black patch over the top-right corner of the image.
Based on the choice of flipping probabilities, images in the training set are spuriously correlated to the position of the patch in the image.
Such correlation does not exist while testing.
A sample batch of images elucidating the above construction is shown in Figure 5.
C. C OLORED D SPRITES
To construct this FL setup, we modify the Desprites dataset 2 (737,280 images) as inAhuja et al.
The data generating process of C OLORED D SPRITES is motivated from that of C OLORED MNIST in a way that it possesses spurious correlations between the label and the colour.
Desprites dataset consists of 2D shapes procedurally generated from 6 ground truth independent latent factors.
These factors are color, shape, scale, rotation, x and y positions of a sprite.
The shape further belongs to a set of three different categories: square, ellipse or heart.
We convert this multi-label and multi-class task into binary classification problem.
Specifically, we consider classification between a circle and a square.
Similar to C OLORED MNIST, the train dataset is equally split across two clients and the entire test set is attributed to the test client.
Preliminary labels for binary classification are constructed such ˜ y = 0 for circle and ˜ y
= 1 for square.
Next, we add noise to the preliminary label by flipping ˜ y with a probability δ k = 0 .
25 , ∀ k to construct the final label y .
We next flip the final label with a probability p k to designate a color ( z ), with p 1 = 0 .
2 for the first client, p 2 = 0 .
1 for the second client and p 3 = 0 .
9 for the test client.
The image is colored red, if z = 1 and green if z = 0 .
A sample batch of images elucidating the above construction is shown in Figure 6.
C.5 Extended Datasets
We extend the two datasets C OLORED MNIST, and C OLORED F ASHION MNIST as in Choe et al.
[2020] to robustly test our approach across multiple clients.
Analogous to C OLORED MNIST with two training environments ( N = 2 ), we extend the datasets to incorporate N = 2 , 3 , 5 and 10 training clients.
In particular, we attribute each client with a unique flipping probability p k .
For each value of N , the maximum value of p k , ∀ k is 0.3, while the minimum value is 0.1.
The values of p k for each client are spaced evenly between this range.
For example, for the case of N = 5 clients, the flipping probabilities of clients are p 1 = 0 .
3 , p 2 = 0 .
2 , and p 3 = 0 .
The flipping probability δ k which decides the final label y is fixed to 0.25 for all clients.
The maximum and the minimum values of p k are chosen in a way that the average spurious correlation which is 0.8 is more than the invariant correlation i.e. 0.75.
Further, since all the previous settings were binary classification tasks, we extend the standard datasets
C OLORED MNIST and C OLORED F ASHION MNIST over multi class classification Choe et al.
Specifically, we extend the number of classes from 2 to 5 and 10.
Specifically, a unique color is assigned to each output class such that the label is highly correlated ( ∼ 80% − 90% ) with the color in the training set.
In the test set, these correlations are significantly reduced ( 10%) by allowing high spurious correlations with the color of the following class.
For instance, in testing, the color corresponding to class label 8 would be the one which was heavily corrected with class label 9 while training.
This reduces the original spurious correlations and is hence useful for evaluating the extent to which the trained model has learned the invariant features.
A sample batch of images elucidating the above construction is shown in Figure 7.
We do not construct a multi-class classification setup for S PURIOUS CIFAR10 and C OLORED
D SPRITES .
For the former, it is difficult to find a unique spatial location in the image corresponding to each class (e.g. in case of 10-class classification).
For the latter, the Desprites dataset can be classified into three categories on the basis of shape present in the image: square, ellipse, and heart.
Hence, introducing more than three classes requires categorization of other latent factors like orientation, scale or position.
D Experimental Setup D.1 Architecture Details
For all the approaches using a fixed representation i.e. F-FL G AMES , F-FL G AMES (S MOOTH ), parallelized F-FL G AMES and parallelized F-FL G AMES (S MOOTH ), we use the architecture mentioned below.
The architecture used to train a predictor at each client is a multi-layered perceptron with three fully connected (FC) layers.
The details of layers are as follows:
• Flatten:
A flatten layer that converts the input of shape (Batch Size, Length, Width, Depth) into a tensor of shape (Batch Size, Length * Width * Depth) • FC1: A fully connected layer with an output dimension of 390, followed by ELU non-linear activation function • FC2: A fully connected layer with an output dimension of 390, followed by ELU non-linear activation function • FC3: A fully connected layer with an output dimension of 2 (Classification Layer)
This same architecture is used across approaches with fixed representation.
For the approaches with trainable representation i.e. V-FL G AMES , V-FL G AMES (S MOOTH ), parallelized V-FL G AMES and parallelized V-FL G AMES (S MOOTH ), we use the following architecture for the representation learner
• Flatten:
A flatten layer that converts the input of shape (Batch Size, Length, Width, Depth) into a tensor of shape (Batch Size, Length * Width * Depth) • FC1: A fully connected layer with an output dimension of 390, followed by ELU non-linear activation function
The output from FC1 is fed into the following architecture, which is used as the base network to train the predictor at each client:
• FC1: A fully connected layer with an output dimension of 390, followed by ELU non-linear activation function • FC2: A fully connected layer with an output dimension of 390, followed by ELU non-linear activation function • FC3: A fully connected layer with an output dimension of 2 (Classification Layer)
This architecture is used across approaches with variable representation.
D.2 Optimizer and other hyperparameters
We use a different set of hyper-parameters based on the dataset.
In particular,
• C OLORED MNIST and C OLORED F ASHION MNIST:
For fixed representation, we
use
Adams optimizer with a learning rate of 2.5e-4 across all experiments (sequential or parallel).
For the variable representation, we use Adams optimizer with a learning rate of 2.5e-5 for the representation learner and the same optimizer with a learning rate of 2.5e-4 for the predictor at each client.
• S PURIOUS CIFAR10: For fixed representation, we use Adams optimizer with a learning rate of 1.0e-4 across all experiments (sequential or parallel).
For the variable representation,
we use Adams optimizer with a learning rate of 9.0e-4 for the representation learner and the same optimizer with a learning rate of 1.0e-4 for the predictor at each client.
• C OLORED D SPRITES : For fixed representation, we use Adams optimizer with a learning rate of 9.5e-4 across all experiments (sequential or parallel).
For the variable representation, we use Adams optimizer with a learning rate of 2.5e-4 for the representation learner and the same optimizer with a learning rate of 9.0e-4 for the predictor at each client
For all the experiments, we fix the batch size to 256 and optimize the Cross Entropy Loss.
We use the same termination criterion as in Ahuja et al.
Specifically, we stop training when the observed oscillations become stable and the ensemble model is in a lower training accuracy state.
We choose a training threshold and terminate the training as soon as the training accuracy drops below this value.
In order to ensure stability of oscillations, we set a period of warm start.
In this period, the training is not stopped even if the accuracy drops below the threshold.
For variable representation, the duration of this warm start period is set to the number of training steps in an epoch i.e. (training data size/ batch size).
However, for the approaches with fixed representation, this period is fixed to N rounds where N is the number of training clients.
In particular, for two clients, the warm start period ends as soon as the second client finishes playing its optimal strategy for the first time.
D.3 Computing Environment
The experiments were done on an Intel Xeon E5 Processor with 2133MHz DDR4 and NVIDIA Tesla v100, 32GB GPU.
E Algorithm
Algorithm 1 represents the pseudo code for V-FL G AMES (S MOOTH +F AST ).
In the following analysis, the terms ‘Sequential’ and ‘Parallel’ denote BRD with clockwise playing sequences and simultaneous updates respectively (Lines 20 and 26 of Algorithm 1).
We use FL G AMES as an umbrella term that constitutes all the discussed algorithmic modifications.
F-FL G AMES and V-FL
G AMES refer to the privacy preserving variants of FL G AMES .
The approach used to smoothen out the oscillations (Line 4 of Algorithm 1) is denoted by F-FL G AMES (S MOOTH ) or V-FL G AMES (S MOOTH ) depending on the constraint on ϕ .
The fast variant with high convergence speed is typified as V-FL G AMES (S MOOTH +F AST ))
(Line 9 of Algorithm 1).
F Additional Results and Analysis F.1 Robustness to the number of Outcomes
[Table]
As described in Section C.5, we test the robustness of our approach to an increase in number of output classes.
We compare F-FL G AMES and parallelized F-FL G AMES across 2-digit, 5-digit and 10-digit classification for C OLORED MNIST and C OLORED F ASHION MNIST.
As shown in Tables 2 and 3, both the sequential and the parallel version of FL G AMES i.e. F-FL
G AMES and parallelized F-FL G AMES respectively are robust to an increase in the number of output classes.
For both the datasets, parallelized F-FL G AMES performs at par or better than F-FL G AMES .
[Table]
Insert w k to
[Table]
k
[Table]
in
[Table]
k if full 6: return w k 7: RepresentationUpdate(k): 8: /
* Gradient
[Table] [Table]
) over entire
[Table]
at client k */ 9: for every batch b ∈ B do 10:
Compute ∇ ℓ k ( w avcur ◦ ϕ cur ; b ) ; Add in ∇ ϕ k 11: return ∇ ϕ k 12: Server executes: 13:
Initialize w k , ∀ k ∈ S and ϕ 14: while round ≤ max-round do 15: /* Update representation ϕ
at even round parity */ 16: if round is even then 17: if Fixed-Phi then 18: ϕ cur =
I 19: if Variable-Phi then 20: for each client k ∈ S in parallel do 21: ∇ ϕ k = RepresentationUpdate(k) 22: /* Update representation ϕ
ϕ next = ϕ cur − η (cid:16) (cid:80) k ∈S
N k (cid:80) j ∈S N j ∇ ϕ k (cid:17) 24: ϕ cur = ϕ
next 25: else 26: for each client k ∈ S in parallel do 27: w k ← PredictorUpdate(k) 28: /
* Client k updates its information set P k by
[Table] [Table]
F.2 Effect of Simultaneous BRD
Similar to the experiments conducted for C OLORED MNIST
[Table]
where we compared F-FL G AMES and parallelized F-FL G AMES across an increase in the number of clients
[Table]
we replicate the same setup for C OLORED F ASHION MNIST
[Table]
S PURIOUS CIFAR10 and C OLORED D SPRITES .
We report the results on the three datasets in Figures 8, 9 and 10.
Consistent with the results on C OLORED
MNIST
[Table]
as the number of clients in the FL system increases
[Table]
there is a sharp increase in the number of communication rounds required to reach equilibrium.
However, the same doesn’t hold true for parallelized F-FL G AMES .
Further, the accuracy achieved by parallelized F-FL G AMES is comparable or better than that achieved by F-FL G AMES .
Moreover, as observed from Table 4,
[Table]
[Figure] [Figure]
across the three benchmarks, all variants of parallel FL G AMES show significantly faster convergence compared to their corresponding sequential counterparts.
This improvement is as large
6 × when comparing
[Figure]
[Table]
F.3 Effect of Memory Ensemble
Similar to the experiments conducted for C OLORED MNIST, where we compared F-FL G AMES and
F-FL G AMES
(S MOOTH ), we replicate the same setup for C OLORED F ASHION MNIST, S PURIOUS
CIFAR10 and C OLORED D SPRITES .
We report the results on all the three datasets in Figures 11,and 13.
Consistent with the results on C OLORED MNIST performance curves oscillate at each step for F-FL
G AMES while the oscillations in F-FL G AMES (S MOOTH ) are observed after an interval of roughly 40 rounds for all the three datasets.
Further, F-FL G AMES (S MOOTH ) also achieves high testing accuracy.
This implies that it does not rely on the spurious features to make predictions.
As observed from Table 4, apart from achieving high accuracy, both F-FL G AMES (S MOOTH ) and parallelized
F-FL G AMES (S MOOTH ) require significantly fewer communication rounds as compared to FFL G AMES and parallelized F-FL G AMES respectively.
Moreover, parallelized F-FL G AMES (S MOOTH ) exhibits fastest convergence rate across all the variants of fixed FL G AMES .
Similar performance evolution curves are also observed for parallelized F-FL G AMES (S MOOTH ) with an added benefit of faster convergence as compared to F-FL G AMES and F-FL G AMES (S MOOTH ).
F.4 Effect of using Gradient Descent (GD) for ϕ
Similar to the experiments conducted for C OLORED MNIST, where we compared V-FL G AMES and V-FL G AMES (S MOOTH +F AST ), we replicate the same setup for C OLORED F ASHION MNIST,
S PURIOUS CIFAR10 and C OLORED D SPRITES .
We report the results on all the three datasets in
Figures 14, 15 and 16.
Consistent with the results on C OLORED MNIST, V-FL G AMES (S MOOTH +F AST ) is able to achieve significantly higher testing accuracy in fewer communication rounds as compared to V-FL G AMES on all the three benchmarks.
Further, as evident from Table 4, parallelized
V-FL G AMES
(S MOOTH +F AST ) has the fastest convergence rate across both S PURIOUS CIFAR and C OLORED D SPRITES .
On a rather simpler dataset i.e. C OLORED F ASHION MNIST, parallelized
V-FL G AMES
(S MOOTH +F AST ) exhibits fastest performance across all the variable variants of FL
G AMES .
As discussed before, the performance of variable FL G AMES is superior to that of it’s fixed variant over complex and large-scale datasets like C OLORED D SPRITES .
Additionally, as conspicuous from
Table 4, parallelized V-FL G AMES (S MOOTH +F AST ) requires significantly fewer communication rounds to reach the Nash equilibrium.
This further underscores the importance of parallelized V-FL
G AMES both in terms of (i) it’s ability to recover the causal mechanisms of the targets, while also providing robustness to distribution shift across clients; (ii) communication efficiency.
F.5 Effect of exact best response
F ED AVG provides the flexibility to train communication efficient and high quality models by allowing more local computation at each client.
This is particularly detrimental in scenarios with poor network connectivity, wherein communicating at every short time span is infeasible.
Inspired by F ED AVG, we study the effect of increasing the amount of local computation at each client.
Specifically, in
F-FL G AMES
, each client updates its predictor based on a step of stochastic gradient descent over its mini-batch.
We modify this setup by allowing each client to run a few steps of stochastic gradient descent locally ( C % ).
When the number of local steps at each client reaches is maximum (training data size/ mini-batch size) or C = 100% ), the scenario becomes equivalent to a gradient descent (GD) over the training data.
For C OLORED MNIST, it is evident from Figure ??
(ii) that as the number of local steps increase i.e. each client exactly best responds to its opponents, the testing accuracy at equilibrium starts to decrease.
When the local computation reaches 100%, i.e. each client updates its local predictor based on a GD over its data, F-FL G AMES exhibits converges (as shown in Figure 17(i)).
Similar to the experiments conducted for C OLORED MNIST, where we studied the influence of increasing the amount of local computation at each client, we replicate the setup for C OLORED
F ASHION MNIST, S PURIOUS CIFAR10 and C OLORED D SPRITES .
Consistent with the results on
C OLORED MNIST and as shown in Figures 18, 19 and 20, even when the number of local steps at each client reaches its maximum i.e. (training data size/ mini-batch size)), the trained models are able to achieve high testing accuracy at equilibrium.
Further, this setup achieves faster convergence with higher training accuracy at equilibrium.
As observed from the empirical results and discussed by
Ahuja et al.
[2021], FL G AMES is guaranteed to exhibit convergence and good out-of-distribution generalization behavior Ahuja et al.
[2021] despite increasing local computations.
Although the testing accuracy at convergence is lower compared to the standard setup, this approach opens avenues for practical deployment of the approach in FL.