Original: F EDERATED M IXTURE OF E XPERTS
Revised: F EDERATED M IXTURE OF E XPERTS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Federated learning (FL) has emerged as the predominant approach for collaborative training of neural network models across multiple users, without the need to gather the data at a central location.
Revised: Federated learning (FL) has emerged as the predominant approach for collaborative training of neural network models across multiple users, without the need to gather the data at a central location.

Original: One of the important challenges in this setting is data heterogeneity; different users have different data characteristics.
Revised: One of the important challenges in this setting is data heterogeneity; different users have different data characteristics.

Original: For this reason, training and using a single global model might be suboptimal when considering the performance of each of the individual user’s data.
Revised: For this reason, training and using a single global model might be suboptimal when considering the performance of each of the individual user’s data.

Original: In this work, we tackle this problem via Federated Mixture of Experts, FedMix , a framework that allows us to train an ensemble of specialized models.
Revised: In this work, we tackle this problem via Federated Mixture of Experts, FedMix , a framework that allows us to train an ensemble of specialized models.

Original: FedMix adaptively selects and trains a user-speciﬁc selection of the ensemble members.
Revised: FedMix adaptively selects and trains a user-speciﬁc selection of the ensemble members.

Original: We show that users with similar data characteristics select the same members and therefore share statistical strength while mitigating the effect of non-i.i.d data.
Revised: We show that users with similar data characteristics select the same members and therefore share statistical strength while mitigating the effect of non-i.i.d data.

Original: Empirically, we show through an extensive experimental evaluation that FedMix improves performance compared to using a single global model while requiring similar or less communication costs.
Revised: Empirically, we show through an extensive experimental evaluation that FedMix improves performance compared to using a single global model while requiring similar or less communication costs.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: An ever-increasing amount of devices are being connected to the internet, sensing their environment, and generating vast amounts of data.
Revised: An ever-increasing amount of devices are being connected to the internet, sensing their environment, and generating vast amounts of data.

Original: The term federated learning (FL) has been established to describe the scenario where we aim to learn from the data generated by this “federation” of devices (McMahan et al., 2016).
Revised: The term federated learning (FL) has been established to describe the scenario where we aim to learn from the data generated by this “federation” of devices (McMahan et al., 2016).

Original: Not only does the number of sensing devices increase, but also their processing power is increasing continuously to the point that it becomes viable to perform inference and training of machine learning models on device.
Revised: Not only does the number of sensing devices increase, but also their processing power is increasing continuously to the point that it becomes viable to perform inference and training of machine learning models on device.

Original: In federated learning, the goal is to learn from these client devices’ data without collecting the data centrally, which naturally allows for more private exchange of information.
Revised: In federated learning, the goal is to learn from these client devices’ data without collecting the data centrally, which naturally allows for more private exchange of information.

Original: Several challenges arise in the federated scenario.
Revised: Several challenges arise in the federated scenario.

Original: Federated devices are generally resourceconstrained, both in their computational capacity as well as in communication bandwidth and latency.
Revised: Federated devices are generally resource-constrained, both in their computational capacity as well as in communication bandwidth and latency.

Original: In a practical example, a smartphone has limited heat dissipation capacity and must communicate via Wi-Fi.
Revised: In a practical example, a smartphone has limited heat dissipation capacity and must communicate via Wi-Fi.

Original: From a global perspective, devices’ processing power and network connection can be highly heterogeneous across geographical regions and socio-economical status of device owners, causing practical issues (Bonawitz et al., 2019) and raising questions of fairness in FL (Li et al., 2019; Mohri et al., 2019).
Revised: From a global perspective, devices’ processing power and network connection can be highly heterogeneous across geographical regions and socio-economical status of device owners, causing practical issues (Bonawitz et al., 2019) and raising questions of fairness in FL (Li et al., 2019; Mohri et al., 2019).

Original: One of the key challenges in FL that we aim to address in this work is the non-i.i.d nature of the shards of data that are distributed across devices.
Revised: One of the key challenges in FL that we aim to address in this work is the non-i.i.d nature of the shards of data that are distributed across devices.

Original: In non-federated machine learning, assuming independent and identically distributed data is generally justiﬁable and not detrimental to model performance.
Revised: In non-federated machine learning, assuming independent and identically distributed data is generally justiﬁable and not detrimental to model performance.

Original: In FL however, each client performs a series of parameter updates on its own data shard to amortize the costs of communication.
Revised: In FL however, each client performs a series of parameter updates on its own data shard to amortize the costs of communication.

Original: Over time, the direction of progress across shards with non-i.i.d data starts diverging (as shown in Figure 1), which can set back training progress, signiﬁcantly slow down convergence and decrease model performance (Hsu et al., 2019).
Revised: Over time, the direction of progress across shards with non-i.i.d data starts diverging (as shown in Figure 1), which can set back training progress, signiﬁcantly slow down convergence and decrease model performance (Hsu et al., 2019).

Original: To this end, we propose Federated Mixture of Experts ( FedMix ), an algorithm for FL that allows for training an ensemble of specialized models instead of a single global model.
Revised: To this end, we propose Federated Mixture of Experts ( FedMix ), an algorithm for FL that allows for training an ensemble of specialized models instead of a single global model.

Original: In FedMix , expert models are learning to specialize in regions of the input space such that, for a given expert, each client’s progress on that expert is aligned.
Revised: In FedMix , expert models are learning to specialize in regions of the input space such that, for a given expert, each client’s progress on that expert is aligned.

Original: FedMix allows each client to learn which experts are relevant for its shard and we show how it can be extended for inference on a previously unseen client.
Revised: FedMix allows each client to learn which experts are relevant for its shard and we show how it can be extended for inference on a previously unseen client.

Original: FedMix shows competitive performance against the established standard in FL, FedAvg (McMahan et al., 2016; Deng et al., 2020) across a range of visual classiﬁcation tasks.
Revised: FedMix shows competitive performance against the established standard in FL, FedAvg (McMahan et al., 2016; Deng et al., 2020) across a range of visual classiﬁcation tasks.

Original: Code will be released upon publication.
Revised: Code will be released upon publication.

Original:  F EDERATED M IXTURE OF E XPERTS
Revised:  F EDERATED M IXTURE OF E XPERTS

Original: Federated learning (McMahan et al., 2016) deals with the problem of learning a server model with parameters w , e.g. , a neural network, from a dataset D
Revised: Federated learning (McMahan et al., 2016) deals with the problem of learning a server model with parameters w , e.g. , a neural network, from a dataset D

Original: = { ( x 1 , y 1 ) , . . .
Revised: = { ( x 1 , y 1 ) , . . .

Original: , ( x N , y N ) } of N datapoints that is distributed across S shards, i.e. , D
Revised: , ( x N , y N ) } of N datapoints that is distributed across S shards, i.e. , D

Original: =
Revised: =

Original: D 1 ∪ · · · ∪ D S , without accessing the shard speciﬁc datasets directly.
Revised: D 1 ∪ · · · ∪ D S , without accessing the shard speciﬁc datasets directly.

Original: By deﬁning a loss function L s ( D s ; w ) per shard, the total risk can be written as
Revised: By deﬁning a loss function L s ( D s ; w ) per shard, the total risk can be written as

Original: [Equation]
Revised: [Equation]

Original: It is easy to see that this objective corresponds to empirical risk minimization over the joint dataset D with a loss L ( · ) for each datapoint.
Revised: It is easy to see that this objective corresponds to empirical risk minimization over the joint dataset D with a loss L ( · ) for each datapoint.

Original: In federated learning one is interested in reducing the communication costs; for this reason McMahan et al.
Revised: In federated learning one is interested in reducing the communication costs; for this reason McMahan et al.

Original: (2016) propose to do multiple gradient updates for w in the inner optimization objective for each shard s , thus obtaining “local” models with parameters w s .
Revised: (2016) propose to do multiple gradient updates for w in the inner optimization objective for each shard s , thus obtaining “local” models with parameters w s .

Original: These multiple gradient updates are denoted as “local epochs”, i.e. , amount of passes through the entire local dataset, with an abbreviation of E .
Revised: These multiple gradient updates are denoted as “local epochs”, i.e. , amount of passes through the entire local dataset, with an abbreviation of E .

Original: Each of the shards then communicates the local model w s to the server and the server updates the global model at “round” t by averaging the parameters of the local models w t
Revised: Each of the shards then communicates the local model w s to the server and the server updates the global model at “round” t by averaging the parameters of the local models w t

Original: = (cid:80) s N s N w ts .
Revised: = (cid:80) s N s N w ts .

Original: This constitutes federated averaging ( FedAvg ) (McMahan et al., 2016), the standard in federated learning.
Revised: This constitutes federated averaging ( FedAvg ) (McMahan et al., 2016), the standard in federated learning.

Original: One of the main challenges in federated learning is the fact that usually the data are non-i.i.d.
Revised: One of the main challenges in federated learning is the fact that usually the data are non-i.i.d.

Original: distributed across the shards S , that is p ( D| s i ) (cid:54) =
Revised: distributed across the shards S , that is p ( D| s i ) (cid:54) =

Original: p ( D| s j ) for i (cid:54) =
Revised: p ( D| s j ) for i (cid:54) =

Original: j .
Revised: j .

Original: On the one hand, this can make learning a single global model from all of the data with the classical FedAvg problematic.
Revised: On the one hand, this can make learning a single global model from all of the data with the classical FedAvg problematic.

Original: On the other hand, there is one extreme that does not suffer from this issue; learning S individual models, i.e. , only optimizing w s on D s .
Revised: On the other hand, there is one extreme that does not suffer from this issue; learning S individual models, i.e. , only optimizing w s on D s .

Original: Although these individual models by deﬁnition do not suffer from non-i.i.d data, clearly we should aim to do better and exchange meaningful information between clients to learn more robust and expressive models.
Revised: Although these individual models by deﬁnition do not suffer from non-i.i.d data, clearly we should aim to do better and exchange meaningful information between clients to learn more robust and expressive models.

Original: 2. T HE F ED M IX ALGORITHM
Revised: 2. T HE F ED M IX ALGORITHM

Original: y x z s
Revised: y x z s

Original: With FedMix , we propose to strike a balance between the two aforementioned extremes; learning a single global model and learning S individual models.
Revised: With FedMix , we propose to strike a balance between the two aforementioned extremes; learning a single global model and learning S individual models.

Original: For this reason, we revisit an old model formulation, the Mixture of Experts (MoE).
Revised: For this reason, we revisit an old model formulation, the Mixture of Experts (MoE).

Original: The classical formulation of a MoE model (Jacobs et al., 1991; Jordan & Jacobs, 1994) contains a set of K experts and a gating mechanism that is responsible for choosing an expert for a given data-point.
Revised: The classical formulation of a MoE model (Jacobs et al., 1991; Jordan & Jacobs, 1994) contains a set of K experts and a gating mechanism that is responsible for choosing an expert for a given data-point.

Original: A MoE model for a data point ( x , y ) can generally be described by
Revised: A MoE model for a data point ( x , y ) can generally be described by

Original: [Equation]
Revised: [Equation]

Original: where z is a categorical variable that denotes the expert, w k are the parameters of expert k and θ are the parameters of the selection mechanism.
Revised: where z is a categorical variable that denotes the expert, w k are the parameters of expert k and θ are the parameters of the selection mechanism.

Original: The MoE was proposed as a model for datasets where different subsets of the data exhibit different relationships between input x and output y .
Revised: The MoE was proposed as a model for datasets where different subsets of the data exhibit different relationships between input x and output y .

Original: Instead of training a single global model to ﬁt this relationship everywhere, each expert performs well on a different subset of the input space.
Revised: Instead of training a single global model to ﬁt this relationship everywhere, each expert performs well on a different subset of the input space.

Original: The gating function models the decision boundary between input regions, assigning data-points from subsets of the input region to their respective experts.
Revised: The gating function models the decision boundary between input regions, assigning data-points from subsets of the input region to their respective experts.

Original: In this work, we show that, in the federated scenario, sub-dividing the input region through a MoE can alleviate the consequences of non-i.i.d data by aligning gradient updates across experts (Figure 1).
Revised: In this work, we show that, in the federated scenario, sub-dividing the input region through a MoE can alleviate the consequences of non-i.i.d data by aligning gradient updates across experts (Figure 1).

Original: In Federated Mixture of Experts ( FedMix ) we enrich this model by conditioning the gating mechanism on the shard assignment s .
Revised: In Federated Mixture of Experts ( FedMix ) we enrich this model by conditioning the gating mechanism on the shard assignment s .

Original: Whatever characteristics make shard s different from other shards can manifest in learning a different, localized gating mechanism that does not need to be communicated to the server.
Revised: Whatever characteristics make shard s different from other shards can manifest in learning a different, localized gating mechanism that does not need to be communicated to the server.

Original: In choosing K = 1 , FedMix recovers the standard setting of federated averaging.
Revised: In choosing K = 1 , FedMix recovers the standard setting of federated averaging.

Original: K = S in combination with ﬁxing p ( z = s | x , s ) =
Revised: K = S in combination with ﬁxing p ( z = s | x , s ) =

Original: 1 recovers S independent models.
Revised: 1 recovers S independent models.

Original: From a global perspective, we are interested in maximizing the following single objective:
Revised: From a global perspective, we are interested in maximizing the following single objective:

Original: [Equation]
Revised: [Equation]

Original: Given the graphical model decomposition depicted in Figure 2, the objective in Eq.
Revised: Given the graphical model decomposition depicted in Figure 2, the objective in Eq.

Original: 3 corresponds to a federated MoE, where we have omitted the generative models p ( x | s ) .
Revised: 3 corresponds to a federated MoE, where we have omitted the generative models p ( x | s ) .

Original: We will brieﬂy touch upon the role of learning generative models in Appendix C but focus on the discriminative part of the model, i.e. , the MoE, in this paper.
Revised: We will brieﬂy touch upon the role of learning generative models in Appendix F but focus on the discriminative part of the model, i.e. , the MoE, in this paper.

Original: While it is possible to optimize Eq.
Revised: While it is possible to optimize Eq.

Original: 3 directly, we have found empirically that it is hard to achieve both: avoiding collapse to a single expert, thus obtaining FedAvg , and specialization of the experts.
Revised: 3 directly, we have found empirically that it is hard to achieve both: avoiding collapse to a single expert, thus obtaining FedAvg , and specialization of the experts.

Original: Instead, we propose to form a variational lower-bound on Eq.
Revised: Instead, we propose to form a variational lower-bound on Eq.

Original: 3 with a global variational approximation q φ ( z | . . . )
Revised: 3 with a global variational approximation q φ ( z | . . . )

Original: to the true posterior p ( z | x , y, s ) with parameters φ .
Revised: to the true posterior p ( z | x , y, s ) with parameters φ .

Original: At test time, p ( y | x ∗ , s ) = (cid:80)
Revised: At test time, p ( y | x ∗ , s ) = (cid:80)

Original: Kk =1 p ( y | x ∗ , z ) p ( z | x ∗ , s ) can be readily evaluated without requiring q .
Revised: Kk =1 p ( y | x ∗ , z ) p ( z | x ∗ , s ) can be readily evaluated without requiring q .

Original: This allows us to condition q φ ( z | . . . )
Revised: This allows us to condition q φ ( z | . . . )

Original: on any available side-information at training time that might result in better specialization in the non-i.i.d federated scenario.
Revised: on any available side-information at training time that might result in better specialization in the non-i.i.d federated scenario.

Original: In this paper we mainly consider classiﬁcation tasks whose non-i.i.d nature predominantly stems from the non-i.i.d distribution of labels y .
Revised: In this paper we mainly consider classiﬁcation tasks whose non-i.i.d nature predominantly stems from the non-i.i.d distribution of labels y .

Original: Other or additional known sources of misalignment could be included to further improve this approximation, such as a manufacturer-id for a medical device in a medical scenario, a geographic identiﬁer, or general domain-speciﬁc information.
Revised: Other or additional known sources of misalignment could be included to further improve this approximation, such as a manufacturer-id for a medical device in a medical scenario, a geographic identiﬁer, or general domain-speciﬁc information.

Original: We show one such additional example in 4.2.
Revised: We show one such additional example in 4.2.

Original: The lower bound to be maximized in FedMix therefore is as follows:
Revised: The lower bound to be maximized in FedMix therefore is as follows:

Original: [Equation]
Revised: [Equation]

Revised: - mental setup is identical to what is described in Section 4.

Original: Conditioning only on y allows us to efﬁciently parameterize the variational approximation, incurring only a small communication overhead.
Revised: Conditioning only on y allows us to efﬁciently parameterize the variational approximation, incurring only a small communication overhead.

Original: While it would be possible to condition q φ ( z | y ) on s , thus having localized approximations with parameters φ s that do not need to be communicated, we found a global approximation to help align the gating mechanisms across shards.
Revised: While it would be possible to condition q φ ( z | y ) on s , thus having localized approximations with parameters φ s that do not need to be communicated, we found a global approximation to help align the gating mechanisms across shards.

Original: A global q φ ( z | y ) encourages shards that contain data with the same label to assign them to the same expert.
Revised: A global q φ ( z | y ) encourages shards that contain data with the same label to assign them to the same expert.

Original: Specialization of the experts is a key ingredient for
Revised: Specialization of the experts is a key ingredient for

Original: FedMix to be successful; with specialization, the gradients for each expert become aligned across shards (see Figure 1), the hold-out accuracy improves (see Figure 3), and the communication costs decrease as each shard may only need to access a subset of the experts.
Revised: FedMix to be successful; with specialization, the gradients for each expert become aligned across shards (see Figure 1), the hold-out accuracy improves (see Figure 3), and the communication costs decrease as each shard may only need to access a subset of the experts.

Original: We ﬁnd that performing maximum a-posteriori (MAP) inference for z generally leads to better and more personalized models.
Revised: We ﬁnd that performing maximum a-posteriori (MAP) inference for z generally leads to better and more personalized models.

Original: By removing the entropy term from equation 4, q φ ( z | y ) and therefore p θ s ( z | x , s ) are encouraged to concentrate and select only one expert for a given data point.
Revised: By removing the entropy term from equation 4, q φ ( z | y ) and therefore p θ s ( z | x , s ) are encouraged to concentrate and select only one expert for a given data point.

Original: In the extreme case where a client’s shard contains only data that is assigned to the same expert, we can reduce communication by receiving and sending updates for that single expert only.
Revised: In the extreme case where a client’s shard contains only data that is assigned to the same expert, we can reduce communication by receiving and sending updates for that single expert only.

Original: We show in Section 4 that communicating and evaluating experts based on thresholding the aggregate q φ ( z | s ) =
Revised: We show in Section 4 that communicating and evaluating experts based on thresholding the aggregate q φ ( z | s ) =

Original: E y ∼ D s [ q φ ( z | y )] can reduce communication and computation overhead.
Revised: E y ∼ D s [ q φ ( z | y )] can reduce communication and computation overhead.

Original: Figure 3 compares FedMix with and without the entropy term to standard FedAvg as a function of communication steps.
Revised: Figure 3 compares FedMix with and without the entropy term to standard FedAvg as a function of communication steps.

Original: With the entropy term, FedMix develops no expert specialization and collapses to an ensemble of K = 4 models.
Revised: With the entropy term, FedMix develops no expert specialization and collapses to an ensemble of K = 4 models.

Original: One drawback of the heavy specialization with MAP inference is that sometimes FedMix prematurely completely prunes experts, i.e. , p θ s ( z = k | x , s )
Revised: One drawback of the heavy specialization with MAP inference is that sometimes FedMix prematurely completely prunes experts, i.e. , p θ s ( z = k | x , s )

Original: ≈ 0 ∀ x , s .
Revised: ≈ 0 ∀ x , s .

Original: This can be undesirable as we lose model capacity that can be used for better modeling the data.
Revised: This can be undesirable as we lose model capacity that can be used for better modeling the data.

Original: As q φ ( z | y ) is one of the main training signals of p θ s ( z | x , s ) , we introduce the marginal entropy term in the server, H ( E p ( y )
Revised: As q φ ( z | y ) is one of the main training signals of p θ s ( z | x , s ) , we introduce the marginal entropy term in the server, H ( E p ( y )

Original: [ q φ ( z | y )]) , as a regularizer.
Revised: [ q φ ( z | y )]) , as a regularizer.

Original: Notice that this leads to different training dynamics than locally optimizing the lower bound with the entropy included and we, empirically, found that it alleviates premature pruning, while still leading to specialized models.
Revised: Notice that this leads to different training dynamics than locally optimizing the lower bound with the entropy included and we, empirically, found that it alleviates premature pruning, while still leading to specialized models.

Original: Figure 9 in Appendix D visualizes the development of q φ ( z | y ) over time from initially uniform to high specialisation for the experiment depicted in Figure 3.
Revised: Figure 14a in Appendix H visualizes the development of q φ ( z | y ) over time from initially uniform to high specialisation for the experiment depicted in Figure 3.

Original: Server Side Updates In a general federated learning algorithm, a central server selects a subset
Revised: Server Side Updates In a general federated learning algorithm, a central server selects a subset

Original: S (cid:48)
Revised: S (cid:48)

Original: ⊂ { 1 , . . .
Revised: ⊂ { 1 , . . .

Original: , S } of clients at time t and transmits the current estimate of the global parameters w t to them.
Revised: , S } of clients at time t and transmits the current estimate of the global parameters w t to them.

Original: These clients perform a series of mini-batch gradient updates with data from their shard
Revised: These clients perform a series of mini-batch gradient updates with data from their shard

Original: D s on a local loss function, which can come at the price of each client moving in possibly different directions in parameter space.
Revised: D s on a local loss function, which can come at the price of each client moving in possibly different directions in parameter space.

Original: In generalized FedAvg (Reddi et al., 2020), the server interprets ∆ ts = w t − w t
Revised: In generalized FedAvg (Reddi et al., 2020), the server interprets ∆ ts = w t − w t

Original: +1 s as a single-step gradient update from client s , averages those gradients and applies an optimizer such as Adam (Kingma & Ba, 2014) to receive w t +1 .
Revised: +1 s as a single-step gradient update from client s , averages those gradients and applies an optimizer such as Adam (Kingma & Ba, 2014) to receive w t +1 .

Original: In light of non-i.i.d data across clients, this averaging strategy can result in slow progress since averaging updates in a highly non-convex parameter space can be sub-optimal.
Revised: In light of non-i.i.d data across clients, this averaging strategy can result in slow progress since averaging updates in a highly non-convex parameter space can be sub-optimal.

Original: In FedMix , this effect is mitigated since for a given expert, the data that is used to update its parameters are aligned better across shards.
Revised: In FedMix , this effect is mitigated since for a given expert, the data that is used to update its parameters are aligned better across shards.

Original: FedMix offers a second way to improve convergence speed by modifying the server-side updates.
Revised: FedMix offers a second way to improve convergence speed by modifying the server-side updates.

Original: In generalized FedAvg , the individual gradients returned by the subset S (cid:48) of clients are averaged according to
Revised: In generalized FedAvg , the individual gradients returned by the subset S (cid:48) of clients are averaged according to

Original: [Equation]
Revised: [Equation]

Original: In FedMix , we can speed up convergence by considering expert-speciﬁc updates ∆ tk,s = w tk − w t
Revised: In FedMix , we can speed up convergence by considering expert-speciﬁc updates ∆ tk,s = w tk − w t

Original: +1 k,s .
Revised: +1 k,s .

Original: If a client s pruned away expert k from its local gating mechanism, ∆ tk,s will be zero.
Revised: If a client s pruned away expert k from its local gating mechanism, ∆ tk,s will be zero.

Original: We propose to normalize the effective magnitude of the resulting update ∆ k by up-weighing the updates of all other clients that do consider expert k for their local mixture:
Revised: We propose to normalize the effective magnitude of the resulting update ∆ k by up-weighing the updates of all other clients that do consider expert k for their local mixture:

Original: [Equation]
Revised: [Equation]

Original: Computing p ( z | s ) =
Revised: Computing p ( z | s ) =

Original: E x ∼D s [ p θ s ( z | s, x )] prior to sending updates to the server involves evaluating potentially large neural network models.
Revised: E x ∼D s [ p θ s ( z | s, x )] prior to sending updates to the server involves evaluating potentially large neural network models.

Original: Therefore we choose to approximate p ( z | s )
Revised: Therefore we choose to approximate p ( z | s )

Original: ≈ q φ ( z | s ) =
Revised: ≈ q φ ( z | s ) =

Original: E y ∼D s [ q φ ( z | y )] , which involves just a single matrix multiplication.
Revised: E y ∼D s [ q φ ( z | y )] , which involves just a single matrix multiplication.

Revised: We discuss the implications of sending q φ ( z | s ) to privacy and how these fare relative to FedAvg in Appendix C.

Original: Pruning experts Maintaining the cheap-to-compute marginal posterior per shard offers an additional opportunity to increase computation speed during local shard iterations and reduce overall communication costs.
Revised: Pruning experts Maintaining the cheap-to-compute marginal posterior per shard offers an additional opportunity to increase computation speed during local shard iterations and reduce overall communication costs.

Original: We propose to “prune away” experts locally from the MoE if q φ ( z | s ) does not surpass a threshold η/K .
Revised: We propose to “prune away” experts locally from the MoE if q φ ( z | s ) does not surpass a threshold η/K .

Original: In order to still optimize a valid bound, we need to re-normalize q φ ( z | y ) before evaluation of the loss function (Pal et al., 2005).
Revised: In order to still optimize a valid bound, we need to re-normalize q φ ( z | y ) before evaluation of the loss function (Pal et al., 2005).

Original: We evaluate the same threshold prior to sending updates to the server in order to avoid communicating parameters that have not changed during the client’s iterations.
Revised: We evaluate the same threshold prior to sending updates to the server in order to avoid communicating parameters that have not changed during the client’s iterations.

Original: Once the server selects a client for another round, it provides only those experts to the client that were updated by the client in the previous round.
Revised: Once the server selects a client for another round, it provides only those experts to the client that were updated by the client in the previous round.

Original: We empirically ﬁnd that the entropy of q ( z | s ) decreases steadily and we prune away experts k with probability q ( z = k | s ) <
Revised: We empirically ﬁnd that the entropy of q ( z | s ) decreases steadily and we prune away experts k with probability q ( z = k | s ) <

Original: η/K without signiﬁcant drop in performance.
Revised: η/K without signiﬁcant drop in performance.

Original: We explore the consequences of pruning experts in the experiment section and in Appendix B. Algorithm 2 shows how FedMix can be enriched by pruning.
Revised: We explore the consequences of pruning experts in the experiment section and in Appendix E. Algorithm 2 shows how FedMix can be enriched by pruning.

Original: function S ERVER SIDE
Revised: function S ERVER SIDE

Original: Initialize φ and K vectors W = [ w 1 , . . .
Revised: Initialize φ and K vectors W = [ w 1 , . . .

Original: , w K ] for round t in 1 , . . .
Revised: , w K ] for round t in 1 , . . .

Original: T do
Revised: T do

Original: S (cid:48) ← random subset of the clients
Revised: S (cid:48) ← random subset of the clients

Original: Initialize ∆ t W = 0 , ∆ t φ =for s in S (cid:48) do
Revised: Initialize ∆ t W = 0 , ∆ t φ =for s in S (cid:48) do

Original: W ts , φ ts , p ( z | s ) ← C LIENT SIDE
Revised: W ts , φ ts , p ( z | s ) ← C LIENT SIDE

Original: ( s, φ , W ) end for p ( s | z ) ← p ( z | s ) p ( s ) / (cid:80) s ∈ S (cid:48) p ( z | s ) p ( s ) for s in S (cid:48)
Revised: ( s, φ , W ) end for p ( s | z ) ← p ( z | s ) p ( s ) / (cid:80) s ∈ S (cid:48) p ( z | s ) p ( s ) for s in S (cid:48)

Original: do ∆ t w k + = p
Revised: do ∆ t w k + = p

Original: ( s | z = k )( w t − 1 k − w ts,k ) ∀ k ∆ t φ + = N s N S (cid:48)
Revised: ( s | z = k )( w t − 1 k − w ts,k ) ∀ k ∆ t φ + = N s N S (cid:48)

Original: ( φ t − 1 − φ ts ) end for ∆ t φ − = ∇ φ H ( (cid:80)
Revised: ( φ t − 1 − φ ts ) end for ∆ t φ − = ∇ φ H ( (cid:80)

Original: c
Revised: c

Original: q φ ( z | y =
Revised: q φ ( z | y =

Original: c ) p ( y = c ))
Revised: c ) p ( y = c ))

Original: w t +11: K ← A DAM ( ∆ t w 1: K , β ) φ t +1 ← A DAM ( ∆ t φ , β ) end for end function function C LIENT SIDE ( s, φ , W )
Revised: w t +11: K ← A DAM ( ∆ t w 1: K , β ) φ t +1 ← A DAM ( ∆ t φ , β ) end for end function function C LIENT SIDE ( s, φ , W )

Original: Get local parameters θ s for epoch e in 1 , . .
Revised: Get local parameters θ s for epoch e in 1 , . .

Original: , E do for batch b ∈ B do
Revised: , E do for batch b ∈ B do

Original: L s ← E q φ ( z | y b ) [log p w z ( y b | x b , z ) p θ s ( z | x b , s )]
Revised: L s ← E q φ ( z | y b ) [log p w z ( y b | x b , z ) p θ s ( z | x b , s )]

Original: φ + = α ∇ φ L s
Revised: φ + = α ∇ φ L s

Original: W + = α ∇ W L s θ s + = α ∇ θ s L s end for end for q ( z | s ) ← E y
Revised: W + = α ∇ W L s θ s + = α ∇ θ s L s end for end for q ( z | s ) ← E y

Original: ∼D s [ q φ ( z | y )] return w 1: K , φ , q ( z | s ) end function
Revised: ∼D s [ q φ ( z | y )] return w 1: K , φ , q ( z | s ) end function

Original: Designing robust gates In the federated scenario, N s is often much smaller than N and especially small in relation to the complexity of the data we try to model.
Revised: Designing robust gates In the federated scenario, N s is often much smaller than N and especially small in relation to the complexity of the data we try to model.

Original: Any localized parameters therefore are prone to overﬁtting.
Revised: Any localized parameters therefore are prone to overﬁtting.

Original: On the other hand, the global parameters of an expert are trained using all data-points assigned to that expert across all shards, allowing to learn more robust features.
Revised: On the other hand, the global parameters of an expert are trained using all data-points assigned to that expert across all shards, allowing to learn more robust features.

Original: We can make use of the robustness of these experts’ features for the gating mechanism by conditioning on them instead of training an entirely separate model for p θ s ( z | x , s ) .
Revised: We can make use of the robustness of these experts’ features for the gating mechanism by conditioning on them instead of training an entirely separate model for p θ s ( z | x , s ) .

Original: Let us deﬁne h k ( x ) as intermediary features of expert k .
Revised: Let us deﬁne h k ( x ) as intermediary features of expert k .

Original: Since not all experts might be used for a given shard and in order to scale with K , we average over the marginal posterior of the training set at that shard before applying a linear transformation to compute the input to the softmax gates:
Revised: Since not all experts might be used for a given shard and in order to scale with K , we average over the marginal posterior of the training set at that shard before applying a linear transformation to compute the input to the softmax gates:

Original: [Equation]
Revised: [Equation]

Original: where θ s =
Revised: where θ s =

Original: ( A s , b s ) are local learnable parameters and SM represents the softmax function.
Revised: ( A s , b s ) are local learnable parameters and SM represents the softmax function.

Original: Inference at test time We consider three variants for test-time evaluation of FedMix .
Revised: Inference at test time We consider three variants for test-time evaluation of FedMix .

Original: In the ﬁrst case, a client s that participated in training is presented with a new data point ( x ∗ , s ) .
Revised: In the ﬁrst case, a client s that participated in training is presented with a new data point ( x ∗ , s ) .

Original: Predictions can then be straightforwardly done by selecting the y that maximizes (cid:80) Kz
Revised: Predictions can then be straightforwardly done by selecting the y that maximizes (cid:80) Kz

Original: =1 p ( y | x ∗ , z ) p ( z | x ∗ , s ) .
Revised: =1 p ( y | x ∗ , z ) p ( z | x ∗ , s ) .

Original: In the second, more challenging, scenario a new client s ∗ is introduced together with a new labelled local data set D s ∗ .
Revised: In the second, more challenging, scenario a new client s ∗ is introduced together with a new labelled local data set D s ∗ .

Original: Here we propose to instantiate and train the local gating mechanism by optimizing the parameters θ s of p θ s ( z | x , s ∗ ) via MAP inference at the local objective.
Revised: Here we propose to instantiate and train the local gating mechanism by optimizing the parameters θ s of p θ s ( z | x , s ∗ ) via MAP inference at the local objective.

Original: Afterwards, predictions can be made in a manner similar to the ﬁrst case.
Revised: Afterwards, predictions can be made in a manner similar to the ﬁrst case.

Original: Finally, we consider the case in which a new client s ∗ has no labelled dataset available.
Revised: Finally, we consider the case in which a new client s ∗ has no labelled dataset available.

Original: Without a local gating function, simply ensembling experts exhibits almost random behaviour since experts can be overly conﬁdent on out-of-distribution data (Snoek et al., 2019).
Revised: Without a local gating function, simply ensembling experts exhibits almost random behaviour since experts can be overly conﬁdent on out-of-distribution data (Snoek et al., 2019).

Original: We therefore propose to ensemble across local gating mechanisms to compute p ( z | x ∗ ) =
Revised: We therefore propose to ensemble across local gating mechanisms to compute p ( z | x ∗ ) =

Original: (cid:80) Ss =
Revised: (cid:80) Ss =

Original: 1 p θ s ( z | x ∗ , s ) p ( s | x ∗ ) ; a method which works well in practice.
Revised: 1 p θ s ( z | x ∗ , s ) p ( s | x ∗ ) ; a method which works well in practice.

Original: In Appendix C we discuss results for new shard inference as well as a more principled approach which makes use of the graphical model formulation in Figure 2.
Revised: In Appendix F we discuss results for new shard inference as well as a more principled approach which makes use of the graphical model formulation in Figure 2.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: FedMix has similarities to many recent works in the topic of federated learning.
Revised: FedMix has similarities to many recent works in the topic of federated learning.

Original: Two methods closely related to ours are described in (Sattler et al., 2019; Briggs et al., 2020).
Revised: Two methods closely related to ours are described in (Sattler et al., 2019; Briggs et al., 2020).

Original: The authors propose to perform hierarchical clustering on the updates returned from each shard in order to incrementally create separate models for groups of users, with a cluster assignment mechanism based on handcrafted heuristics.
Revised: The authors propose to perform hierarchical clustering on the updates returned from each shard in order to incrementally create separate models for groups of users, with a cluster assignment mechanism based on handcrafted heuristics.

Original: FedMix instead takes a different approach; it starts with a ﬁxed set of K models
Revised: FedMix instead takes a different approach; it starts with a ﬁxed set of K models

Original: Another closely related work is presented by Mansour et al.
Revised: Another closely related work is presented by Mansour et al.

Original: (2020), where the authors propose to similarly create an ensemble of K -models and assign to each shard the model that achieves the lowest training loss on the local dataset.
Revised: (2020), where the authors propose to similarly create an ensemble of K -models and assign to each shard the model that achieves the lowest training loss on the local dataset.

Original: This is closer to the assignment that happens in FedMix with one main difference; FedMix takes into account the uncertainty in the selection mechanism as well with p ( z | x , s ) instead of selecting the top performing component during training.
Revised: This is closer to the assignment that happens in FedMix with one main difference; FedMix takes into account the uncertainty in the selection mechanism as well with p ( z | x , s ) instead of selecting the top performing component during training.

Original: This is beneﬁcial early in training where the models have not fully specialized yet.
Revised: This is beneﬁcial early in training where the models have not fully specialized yet.

Original: Using local and global model parts has also been explored by Liang et al.
Revised: Using local and global model parts has also been explored by Liang et al.

Original: The authors propose to have a local feature extractor at each shard and a global classiﬁer on top of those features as opposed to having
Revised: The authors propose to have a local feature extractor at each shard and a global classiﬁer on top of those features as opposed to having

Original: K -separate models and a local selection mechanism as in FedMix .
Revised: K -separate models and a local selection mechanism as in FedMix .

Original: This setup yields improvements upon the vanilla federated averaging algorithm, however there are two potential drawbacks; ﬁrst, empirically, the authors had to start their procedure from a pre-trained model with FedAvg and secondly, they have to ensemble all of the different feature extractors for predictions in new shards.
Revised: This setup yields improvements upon the vanilla federated averaging algorithm, however there are two potential drawbacks; ﬁrst, empirically, the authors had to start their procedure from a pre-trained model with FedAvg and secondly, they have to ensemble all of the different feature extractors for predictions in new shards.

Original: In our experiments, we omit the pre-training step and show that the ensembling strategy fails.
Revised: In our experiments, we omit the pre-training step and show that the ensembling strategy fails.

Original: Federated learning in the non-i.i.d setting can also be improved upon in other ways.
Revised: Federated learning in the non-i.i.d setting can also be improved upon in other ways.

Original: Li et al.
Revised: Li et al.

Original: (2018) propose to employ a proximal regularizer at the shard level in order to prevent the local models from drifting too far from the global model, thus making federated learning more robust.
Revised: (2018) propose to employ a proximal regularizer at the shard level in order to prevent the local models from drifting too far from the global model, thus making federated learning more robust.

Original: Jiang et al.
Revised: Jiang et al.

Original: (2019) notice that FedAvg and Reptile (Nichol et al., 2018), a meta-learning algorithm, are essentially the same algorithm and thus propose ﬁne tuning with Reptile in order to improve the personalized performance of the global model.
Revised: (2019) notice that FedAvg and Reptile (Nichol et al., 2018), a meta-learning algorithm, are essentially the same algorithm and thus propose ﬁne tuning with Reptile in order to improve the personalized performance of the global model.

Original: In a similar vein, there are promising new works that explore the meta-learning view of federated learning (Chen et al., 2018; Khodak et al., 2019;
Revised: In a similar vein, there are promising new works that explore the meta-learning view of federated learning (Chen et al., 2018; Khodak et al., 2019;

Original: Fallah et al., 2020).
Revised: Fallah et al., 2020).

Original: Improving the personalized performance of the global model has also been done without meta-learning in works such as by Deng et al.
Revised: Improving the personalized performance of the global model has also been done without meta-learning in works such as by Deng et al.

Original: ; Mansour et al.
Revised: ; Mansour et al.

Original: In general, such improvements are complementary to FedMix and can be used to further enhance its performance.
Revised: In general, such improvements are complementary to FedMix and can be used to further enhance its performance.

Original: We refer the interested readers to the recent surveys by (Kairouz et al., 2019; Kulkarni et al., 2020).
Revised: We refer the interested readers to the recent surveys by (Kairouz et al., 2019; Kulkarni et al., 2020).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate FedMix on three datasets: Cifar10, Cifar100 (Krizhevsky et al., 2009) and Femnist (Caldas et al., 2018), a 62-way image classiﬁcation problem on hand-written digits and letters that is naturally non-i.i.d due to different writing styles of 3500 users.
Revised: We evaluate FedMix on three datasets: Cifar10, Cifar100 (Krizhevsky et al., 2009) and Femnist (Caldas et al., 2018), a 62-way image classiﬁcation problem on hand-written digits and letters that is naturally non-i.i.d due to different writing styles of 3500 users.

Original: In Appendix B we  provide additional ablation studies. For Cifar10, we replicate the federated data split of Hsu et al.
Revised: In Appendix A we detail the experimental setup and provide additional ablation studies in Appendix E.

Original: The dataset is split across 100 clients, whose data-points are drawn according to their label from a Dir ( α = 1 . 0) distribution without replacement.

Original: For the base model, we use a LeNet-5 architecture (LeCun et al., 1998).

Original: We use a SGD optimizer with a learning rate of 0 .

Original: 05 and a batch size B = 64 locally and the

Original: Adam (Kingma & Ba, 2014) optimizer with its default hyperparameters at the server by interpreting the difference of the local from the global model as a gradient (Reddi et al., 2020).

Original: We sample 10 clients without replacement on each round (but with replacement across rounds) and train for

Original: E = 1 local epochs.

Original: For Cifar100 we replicate the data split of Reddi et al.

Original: The dataset is split into 500 clients by using a hierarchical model over the coarse and ﬁne labels, with the same hyperparameters as the ones provided by Reddi et al.

Original: The other hyperparameters are the same as Cifar10 with the exception of the batch size, where we use B = 20 , as well as the architecture, where we use a ResNet-20 with group normalization (Wu & He, 2018) layers instead of batch normalization (Ioffe & Szegedy, 2015).

Original: We augment the data by random cropping from apixel padded image and horizontal ﬂipping.

Original: Finally, for the Femnist dataset 1 , we similarly followed the setup of Reddi et al.

Original: (2020) with the same LeNet-5 architecture and hyperparameters of Cifar with the exception of the batch size where we used B = 20 .

Original: For Cifar10 and Cifar100, h k ( x ) is deﬁned as the input to an expert’s output layer.

Original: For Femnist, we observed signiﬁcant overﬁtting with the same deﬁnition, so we set h k ( x ) to be the output logits for each expert.

Original: 4. P ERFORMANCE VS . COMMUNICATION
Revised: 4. P ERFORMANCE VS . COMMUNICATION

Original: We compare FedMix along several dimensions to baselines such as (generalized) FedAvg (Reddi et al., 2020), biased FedAvg , and the Local/Global approach of Liang et al.
Revised: We compare FedMix along several dimensions to baselines such as (generalized) FedAvg (Reddi et al., 2020), biased FedAvg , and the Local/Global approach of Liang et al.

Original: In biased
Revised: In biased

Original: FedAvg , we allow each client to learn a personalized bias vector b s of its output layer.
Revised: FedAvg , we allow each client to learn a personalized bias vector b s of its output layer.

Original: This will allow biased FedAvg to model the label skew at each client but, fundamentally, cannot model any other form of non-i.i.d-ness.
Revised: This will allow biased FedAvg to model the label skew at each client but, fundamentally, cannot model any other form of non-i.i.d-ness.

Original: Similarly, Liang et al.
Revised: Similarly, Liang et al.

Original: (2020) propose to split the model into local and global components by having local feature extractors and learning the upper layers of the neural network via FedAvg .
Revised: (2020) propose to split the model into local and global components by having local feature extractors and learning the upper layers of the neural network via FedAvg .

Original: We experimented with splitting LeNet-5 at every intermediate layer and report results with the best performing split: keeping the input layer local.
Revised: We experimented with splitting LeNet-5 at every intermediate layer and report results with the best performing split: keeping the input layer local.

Original: For ResNet-20, splitting after the ﬁrst block performed best.
Revised: For ResNet-20, splitting after the ﬁrst block performed best.

Original: First, we show that training with FedMix achieves higher global model accuracy.
Revised: First, we show that training with FedMix achieves higher global model accuracy.

Original: Second, we show that for a given communication budget, FedMix can be competitive against baselines.
Revised: Second, we show that for a given communication budget, FedMix can be competitive against baselines.

Original: Table 1 contains test set results for Cifar10, evaluated after 2 k communication steps and Femnist evaluated after 5 . 95 k .
Revised: Table 1 contains test set results for Cifar10, evaluated after 2 k communication steps and Femnist averaged over the last 2 k steps.

Original: Figures 4d, 4f reveal that the models have converged at that point.
Revised: Figures 4d, 4f reveal that the models have converged at that point.

Original: For Cifar100, we observe no overﬁtting and report test-set accuracies averaged over the last 10 evaluations.
Revised: For Cifar100, we observe no overﬁtting and report test-set accuracies averaged over the last 10 evaluations.

Original: Figure 4 shows learning curves for different settings of FedMix .
Revised: Figure 4 shows learning curves for different settings of FedMix .

Revised: All of these results were obtained by using the server version of the global parameters; the results with additional local ﬁne-tuning at the client level are discussed in Appendix G.

Original: For Cifar10 we can see that except for the initial training phase, FedMix improves over FedAvg for a given communication budget.
Revised: For Cifar10 we can see that except for the initial training phase, FedMix improves over FedAvg for a given communication budget.

Original: For the same communication costs per communication round,
Revised: For the same communication costs per communication round,

Original: FedMix offers better performance than FedAvg , as can be seen in comparing FedMix ( K = 2 ) to
Revised: FedMix offers better performance than FedAvg , as can be seen in comparing FedMix ( K = 2 ) to

Original: FedAvg LeNet-5 with 61 channels and FedMix ( K = 4 ) to FedAvg LeNet-5 with 113 channels.
Revised: FedAvg LeNet-5 with 61 channels and FedMix ( K = 4 ) to FedAvg LeNet-5 with 113 channels.

Original: The channel count has been chosen such that the model size approximates 2 and 4 separate standard
Revised: The channel count has been chosen such that the model size approximates 2 and 4 separate standard

Original: LeNet experts respectively.
Revised: LeNet experts respectively.

Original: Against biased FedAvg and Local/Global Liang et al.
Revised: Against biased FedAvg and Local/Global Liang et al.

Original: (2020), FedMix can improve signiﬁcantly in terms of inference on a new shard (Figure 8 in Appendix C).
Revised: (2020), FedMix can improve signiﬁcantly in terms of inference on a new shard (Figure 12 in Appendix F).

Original: When endowing FedMix with the ability to learn a single personalized bias vector b s across all experts, we can harness some of the advantages of biased FedAvg for inference on a new shard while retaining FedMix ’ advantage when performing inference on a new shard.
Revised: When endowing FedMix with the ability to learn a single personalized bias vector b s across all experts, we can harness some of the advantages of biased FedAvg for inference on a new shard while retaining FedMix ’ advantage when performing inference on a new shard.

Original: For Cifar100 , we can see that learning progresses much faster for FedMix in terms of communication rounds than FedAvg .
Revised: For Cifar100 , we can see that learning progresses much faster for FedMix in terms of communication rounds than FedAvg .

Original: Against FedAvg , we can see that FedMix improves with both K= and K=10 on a per round progression, and with K=10 and expert pruning, it eventually provides a better accuracy and communication trade-off than FedAvg .
Revised: Against FedAvg , we can see that FedMix improves with both K= and K=10 on a per round progression, and with K=10 and expert pruning, it eventually provides a better accuracy and communication trade-off than FedAvg .

Original: Biased FedAvg seems to be the better inductive bias in this task, although one should be weary that its usefulness is limited to label skew and does not generalize to other sources of non-i.i.d.ness (see section 4.2).
Revised: Biased FedAvg seems to be the better inductive bias in this task, although one should be weary that its usefulness is limited to label skew and does not generalize to other sources of non-i.i.d.ness (see section 4.2).

Original: This is in contrast to
Revised: This is in contrast to

Original: FedMix which is designed to handle any type of explicit non-i.i.d.ness.
Revised: FedMix which is designed to handle any type of explicit non-i.i.d.ness.

Original: For Femnist we compare against FedAvg and observe that generally FedMix is more prone to overﬁtting due to the extra amount of model parameters and expert specialization.
Revised: For Femnist we compare against FedAvg and observe that generally FedMix is more prone to overﬁtting due to the extra amount of model parameters and expert specialization.

Original: For this reason, we tie the parameters of  all  the layers of the LeNet-5 experts, except for those at the output layer. In this setting, the extra communication cost of FedMix compared to FedAvg is negligible and with
Revised: For this reason, we learn the parameters of the local gating networks via FedAvg across all clients. The gating network is still localized through the use of q ( z | s ) in pooling of h k ( x ) (equation 9).

Original: K = 4 experts we see that  FedMix reaches better performance.
Revised: We see that increasing K improves the performance of FedMix .

Original: Alternative choices for K , such as 2 , 10 , did not seem to improve upon FedAvg or K = 4 .
Revised: Since FedMix effectively requires the model to ﬁrst identify the right expert, followed by evaluating that expert’s prediction, compounding errors can lead to FedMix being outperformed by FedAvg , such as for K = 2 .

Original: Biased FedAvg performed identically to standard FedAvg since the non-i.i.d-ness does not predominantly stem from label-skew.
Revised: Biased FedAvg performs similar to standard FedAvg since the non-i.i.d-ness does not predominantly stem from label-skew.

Original: [Table]

Original: [Figure]
Revised: [Figure]

Original: 4. R OTATED MNIST
Revised: 4. R OTATED MNIST

Original: To show that FedMix is not limited to label-skew, we create a federated rotated MNIST dataset with 100 clients.
Revised: To show that FedMix is not limited to label-skew, we create a federated rotated MNIST dataset with 100 clients.

Original: Instead of label skew, each client randomly chooses a multiple of 45 degrees  from a different probability distribution over 8 possible rotation angles to rotate a digit with.
Revised: Instead of label skew, each client randomly chooses a multiple of 45 degrees [Table] from a different probability distribution over 8 possible rotation angles to rotate a digit with.

Original: Each client’s distribution is drawn from Dir( α = 1 . 0 ).
Revised: Each client’s distribution is drawn from Dir( α = 1 . 0 ).

Original: At test time, each data-point is randomly rotated according to the client’s distribution.
Revised: At test time, each data-point is randomly rotated according to the client’s distribution.

Original: Additionally, we create a dataset where instead of uniform sampling of labels, we replicate the non-i.i.d label skew described for Cifar10 above and combine it with the rotation non-i.i.d-ness.
Revised: Additionally, we create a dataset where instead of uniform sampling of labels, we replicate the non-i.i.d label skew described for Cifar10 above and combine it with the rotation non-i.i.d-ness.

Original: We compare FedMix where q is conditioned on y or on the degrees of rotation for a data-point against baselines.
Revised: We compare FedMix where q is conditioned on y or on the degrees of rotation for a data-point against baselines.

Original: Figures 5 and 6 show that biased FedAvg can improve in the presence of label skew but collapses to standard FedAvg for rotation non-i.i.d-ness.
Revised: Figures 5 and 6 show that biased FedAvg can improve in the presence of label skew but collapses to standard FedAvg for rotation non-i.i.d-ness.

Original: FedMix on the other hand has an advantage in both cases.
Revised: FedMix on the other hand has an advantage in both cases.

Original: Furthermore, in presence of both sources of noni.i.d-ness, FedMix with q ( z | y ) can outperform biased FedAvg .
Revised: Furthermore, in presence of both sources of noni.i.d-ness, FedMix with q ( z | y ) can outperform biased FedAvg .

Original: In the presence of only rotation (Figure 5), conditioning FedMix on the true source of non-i.i.d-ness is advantageous. 
Revised: In the presence of only rotation (Figure 5), conditioning FedMix on the true source of non-i.i.d-ness is advantageous. 4. L ABEL P ERMUTATIONS

Revised: Apart from non-i.i.d-ness in p ( y ) and p ( x ) , we can expect the mapping p ( y | x ) itself to be different between clients.

Original: {+We replicate the experimental setup of (Sattler et al., 2019) with 20 clients for Cifar10, C = 1 . 
+}
Revised: {+We replicate the experimental setup of (Sattler et al., 2019) with 20 clients for Cifar10, C = 1 . 
+}

Revised: and

Revised: E = 3 , albeit with LeNet-5.

Original: D ISCUSSION
Revised: Each client is randomly assigned one of 4 different permutations of its labels, determining the cluster assignment q ( z | s ) that FedMix has to learn. Although FedMix is designed to distinguish different regions of the input space, we show that it can perform user clustering. The gating function p ( z | x, s ) learns to correctly identify, for each datapoint, the expert corresponding to the permutation of s , thus recovering the original clustering; Figure 7 illustrates this effect (please note that the ordering of columns is arbitrary). It should be mentioned, that FedMix converges very quickly compared to the results in (Sattler et al., 2019), as this result was obtained after only 10 rounds.

Revised: After this point, FedMix essentially trains 4 independent models on the respective i.i.d subsets of the data that have the same label permutation.

Revised: Appendix B contains additional discussions for K (cid:54) = 4 and experimental details.

Original: With FedMix we have introduced a federated learning algorithm that explicitly takes the non-i.i.d characteristics of a federated dataset into account.
Revised: With FedMix we have introduced a federated learning algorithm that explicitly takes the non-i.i.d characteristics of a federated dataset into account.

Original: Clients can learn to align specialized experts on sub-regions of the data space and achieve higher performance compared to FedAvg , especially in situations where the source of the non-i.i.d nature is known.
Revised: Clients can learn to align specialized experts on sub-regions of the data space and achieve higher performance compared to FedAvg , especially in situations where the source of the non-i.i.d nature is known.

Original: This assumption is very strong in real-world federated scenarios and we expect a more ﬂexible alignment process than a global q ( z | y ) to be the most interesting avenue for future research.
Revised: This assumption is very strong in real-world federated scenarios and we expect a more ﬂexible alignment process than a global q ( z | y ) to be the most interesting avenue for future research.

Original: Nevertheless, we do showsome encouraging results on the rotated MNIST experiment, where even with the “wrong” conditioning we were able to improve upon FedAvg .
Revised: We show encouraging results on experiments with non-i.i.d-ness in p ( y | x ) (permutation) and p ( x ) (rotation).

Original: In the future, we will explore ways to perform automatic selection of K as well as automatic selection of architecture elements to share between experts, trading-off gradient alignment and communication budgets.
Revised: In the future, we will explore ways to perform automatic selection of K as well as automatic selection of architecture elements to share between experts, trading-off gradient alignment and communication budgets.

Original: A G RADIENT D IVERGENCE
Revised: A E XPERIMENTAL S ETUP

Revised: For Cifar10, we replicate the federated data split of Hsu et al.

Revised: The dataset is split across 100 clients, whose data-points are drawn according to their label from a Dir ( α = 1 . 0) distribution without replacement.

Revised: For the base model, we use a LeNet-5 architecture (LeCun et al., 1998).

Revised: We use a SGD optimizer with a learning rate of 0 .

Revised: 05 and a batch size B = 64 locally and the

Revised: Adam (Kingma & Ba, 2014) optimizer with its default hyperparameters at the server by interpreting the difference of the local from the global model as a gradient (Reddi et al., 2020).

Revised: We sample 10 clients without replacement on each round (but with replacement across rounds) and train for

Revised: E = 1 local epochs.

Revised: For Cifar100 we replicate the data split of Reddi et al.

Revised: The dataset is split into 500 clients by using a hierarchical model over the coarse and ﬁne labels, with the same hyperparameters as the ones provided by Reddi et al.

Revised: The other hyperparameters are the same as Cifar10 with the exception of the batch size, where we use B = 20 , as well as the architecture, where we use a ResNet-20 with group normalization (Wu & He, 2018) layers instead of batch normalization (Ioffe & Szegedy, 2015).

Revised: We augment the data by random cropping from apixel padded image and horizontal ﬂipping.

Revised: Finally, for the Femnist dataset, we similarly followed the setup of Reddi et al.

Revised: (2020) with the same LeNet-5 architecture and hyperparameters of Cifar with the exception of the batch size where we used B = 20 .

Revised: h k ( x ) is deﬁned as the input to an expert’s output layer.

Revised: B P ERMUTATION E XPERIMENTS

Revised: For the experiments in Section 4.3, we train FedMix with K = 4 and a base LeNet-5 model.

Revised: The optimization hyperparameters are chosen identically to what is described in A, except for the optimization of q ( z | s ) .

Revised: The local objective in this case is

Revised: [Equation]

Revised: In order to speed up convergence, we found it beneﬁcial to not perform gradient ascent and instead solve it directly using Lagrange multiplies, i.e. , ∇ π,λ L s ( D s , w , q π ) + λ

Revised: [Equation]

Revised: Kk =1 π k − 1) = 0 .

Revised: The solution for π k = q

Revised: ( z = k | s ) at client s takes the form of

Revised: [Equation]

Revised: For these permutation experiments, FedMix is tasked to recover a cluster-assignment, as opposed to a partitioning of the input space.

Revised: We have found it unnecessary to change β from its default value of 1 .

Revised: For an approximation of π k with a mini-batch of size M , we approximate log

Revised: [Equation]

Revised: and perform dampening π t

Revised: = α · π t − 1 + (1 − α ) · π .

Revised: We found our initial pick of α = 0 .

Revised: 75 to perform well.

Revised: B. M ODIFYING K

Revised: In Section 4.3, we choose K = 4 a-priori as the number of experts that ﬁts the data generating process with 4 clusters.

Revised: in Figure 8, we plot q ( z | y ) in case we miss-specify K and choose K =and K = 5 .

Revised: In this plot, we can consider K = 4 as the ground truth for the correct cluster assignment.

Revised: Please note that ordering of the columns is arbitrary and can be different across the three experiments.

Revised: Comparing K = 4 with K = 5 , we see that FedMix assigns two experts to the same cluster, while still correctly disentangling all clients.

Revised: For K = 3 , FedMix correctly identiﬁes two clusters while being forced to mix together the other two clusters into one expert.

Revised: C P RIVACY I MPLICATIONS

Revised: Privacy is one of the key motivations for research and deployment of Federated Learning.

Revised: Even though privacy is not a focus of this paper, we brieﬂy discuss some implications of making explicit use of q ( y | s ) in FedMix .

Revised: The update rule described in equation 8 requires access to the marginal q ( z | s ) =

Revised: (cid:80) y p ( y | s ) q φ ( z | y ) at the server.

Revised: At the same time, the server has access to the parameters φ that were used in computing p ( z | s ) before being sent to the server.

Revised: Therefore, in principle, it could solve q ( z | s ) =

Revised: (cid:80) y p ( y | s ) q φ ( z | y ) with respect to p ( y | s ) and thus obtain the marginal label distribution at the client.

Revised: In practice this is not as straightforward to do as (a) the probability matrix q φ ( z | y ) is not always invertible and (b) solutions that use the pseudo-inverse, empirically, are not very accurate in capturing the entire distribution.

Revised: Irrespective of FedMix , we want to shine light on the possibility to reconstruct p ( y | s ) at the serverside in standard FedAvg .

Revised: Assume a randomly initialized model being sent to a client s , where the client performs a single full data set update step on the output layer’s bias vector b s .

Revised: Assuming a softmax cross-entropy loss L s , the average gradient with respect to a the k-th

Revised: entry b k takes the form of

Revised: [Equation]

Revised: (14) where 1 is the indicator function and π i,k is the softmax probability of class k of datapoint i .

Revised: With a randomly initialized model, these softmax probabilities can be assumed to be uniform, leading to an average gradient of

Revised: [Equation]

Revised: where N c is the number of classes.

Revised: Upon sending the updated bias vector b s = b − α ∂L

Revised: s ∂b k to the server, it can easily reconstruct the marginal label distribution.

Revised: Figure 9 shows, for every client in our Cifar10 setup (Appendix A), in red the true marginal p ( y | s ) and in blue the reconstructed marginal based on (the same) randomly initialized model being sent to each client.

Revised: Clearly, we have high congruence.

Revised: In Figure 10, we investigate the more realistic setup of E = 1 with mini-batch stochastic gradient descent at the client level.

Revised: In order to avoid reconstructing the multi-step update, we simply normalize the difference ( b − b s ) and interpret it as marginal label distribution.

Revised: We see that multiple updates (on average: 8 ) reduce the congruence between the true and reconstructed marginal, however the information leakage is still remarkable.

Revised: D G RADIENT D IVERGENCE

Original: We aim to track the divergence of updates for a subset S (cid:48) of shards at the server at time step t for
Revised: We aim to track the divergence of updates for a subset S (cid:48) of shards at the server at time step t for

Original: FedMix and FedAvg .
Revised: FedMix and FedAvg .

Original: Therefore we deﬁne a metric inspired by Kim et al.
Revised: Therefore we deﬁne a metric inspired by Kim et al.

Original: (2020); Sattler et al.
Revised: (2020); Sattler et al.

Original: (2019) to deﬁne divergence of gradients
Revised: (2019) to deﬁne divergence of gradients

Original: ∆ tk,i = ω tk − ω t +1 k,i for some subset ω k of the parameters w k of expert k as
Revised: ∆ tk,i = ω tk − ω t +1 k,i for some subset ω k of the parameters w k of expert k as

Original: [Equation]
Revised: [Equation]

Original: For FedAvg , the above metric collapses to
Revised: For FedAvg , the above metric collapses to

Original: [Equation]
Revised: [Equation]

Original: In Figure 1 in the main text, we plot the sum of GD (∆ tk ) across all parameters ω k
Revised: In Figure 1 in the main text, we plot the sum of GD (∆ tk ) across all parameters ω k

Original: ( i.e. , convolutional kernels, weights and biases) of the LeNet-5 experts in comparison to ω for FedAvg .
Revised: ( i.e. , convolutional kernels, weights and biases) of the LeNet-

Original: B A BLATION STUDIES
Revised: E A BLATION STUDIES

Original: We investigate the characteristics of FedMix by varying the number of experts K for Cifar10, as well as by selecting different levels η for pruning experts.
Revised: We investigate the characteristics of FedMix by varying the number of experts K for Cifar10, as well as by selecting different levels η for pruning experts.

Original: Figures 7a and 7c show learning curves of several values of K .
Revised: Figures 11a and 11c show learning curves of several values of K .

Original: We can see that a higher number of experts has the potential to achieve higher accuracies at the cost of more required communication.
Revised: We can see that a higher number of experts has the potential to achieve higher accuracies at the cost of more required communication.

Original: With increasing K , the modeling task for a single expert becomes progressively easier and the challenge is moved towards identifying the right expert for a given data-point.
Revised: With increasing K , the modeling task for a single expert becomes progressively easier and the challenge is moved towards identifying the right expert for a given data-point.

Original: This could subsequently lead to overﬁtting in the federated setting if not for the shared feature extractor h ( x ) (Eq.
Revised: This could subsequently lead to overﬁtting in the federated setting if not for the shared feature extractor h ( x ) (Eq.

Original: equation 9).
Revised: equation 9).

Original: In case we choose K to be equal to the number of classes, we observe that FedMix assigns one cluster to each class and the p ( z | x , s ) is reminiscent of Liang et al.
Revised: In case we choose K to be equal to the number of classes, we observe that FedMix assigns one cluster to each class and the p ( z | x , s ) is reminiscent of Liang et al.

Original: In case we further increase K , we observe marginally improved performance by multiple experts being assigned to the same label.
Revised: In case we further increase K , we observe marginally improved performance by multiple experts being assigned to the same label.

Original: Figures 7b and 7d depict the inﬂuence of different pruning levels η on the communication efﬁciency and ﬁnal performance FedMix with K = 4 .
Revised: Figures 11b and 11d depict the inﬂuence of different pruning levels η on the communication efﬁciency and ﬁnal performance FedMix with K = 4 .

Original: We can see that higher levels of pruning do reduce communication signiﬁcantly, however they cause some loss of performance.
Revised: We can see that higher levels of pruning do reduce communication signiﬁcantly, however they cause some loss of performance.

Original: [Figure]
Revised: [Figure]

Original: C N EW SHARD INFERENCE
Revised: F N EW SHARD INFERENCE

Original: In the main text, we discussed performance of the individual algorithms when evaluating new data on a shard that took part in training and therefore has access to trained local model parameters.
Revised: In the main text, we discussed performance of the individual algorithms when evaluating new data on a shard that took part in training and therefore has access to trained local model parameters.

Original: This is the case for all considered methods except FedAvg , for which there exist no localized parameters.
Revised: This is the case for all considered methods except FedAvg , for which there exist no localized parameters.

Original: For inference on a new shard, we use the sum of all local validation sets as a proxy.
Revised: For inference on a new shard, we use the sum of all local validation sets as a proxy.

Original: Liang et al.
Revised: Liang et al.

Original: (2020) propose to ensemble the representations of the local feature extractors across clients before evaluating the global part of the network.
Revised: (2020) propose to ensemble the representations of the local feature extractors across clients before evaluating the global part of the network.

Original: We ﬁnd this approach to work quite poorly in practice.
Revised: We ﬁnd this approach to work quite poorly in practice.

Original: Presumably pretraining the local feature extractors with FedAvg ameliorates this behaviour.
Revised: Presumably pretraining the local feature extractors with FedAvg ameliorates this behaviour.

Original: In biased FedAvg , we ensemble the individual local biases across clients to receive a single global bias.
Revised: In biased FedAvg , we ensemble the individual local biases across clients to receive a single global bias.

Original: For FedMix , we marginalize the local gating predictions to achieve a global gating prediction p ( z | x ∗ ) =
Revised: For FedMix , we marginalize the local gating predictions to achieve a global gating prediction p ( z | x ∗ ) =

Original: (cid:80) S s =1 p θ s ( z | x ∗ , s ) p ( s ) .
Revised: (cid:80) Ss = 1 p θ s ( z | x ∗ , s ) p ( s ) .

Original: Predictions can then be made by marginalizing across experts using this global gating function: p ( y | x ∗ )
Revised: Predictions can then be made by marginalizing across experts using this global gating function: p ( y | x ∗ )

Original: = (cid:80) Kz
Revised: = (cid:80) Kz

Original: =1 p w k ( y | x ∗ , z ) p ( z | x ∗ ) .
Revised: =1 p w k ( y | x ∗ , z ) p ( z | x ∗ ) .

Original: Interestingly, Cifar10 and Cifar100 display different behaviours when comparing FedMix to the baselines.
Revised: Interestingly, Cifar10 and Cifar100 display different behaviours when comparing FedMix to the baselines.

Original: For Cifar10, biased FedAvg performs much worse compared to FedAvg while FedMix performs well.
Revised: For Cifar10, biased FedAvg performs much worse compared to FedAvg while FedMix performs well.

Original: For Cifar100, biased FedAvg is indistinguishable from FedAvg , outperforming
Revised: For Cifar100, biased FedAvg is indistinguishable from FedAvg , outperforming

Original: FedMix .
Revised: FedMix .

Original: We leave the investigation of these discrepeancies to future work.
Revised: We leave the investigation of these discrepeancies to future work.

Original: [Figure]
Revised: [Figure]

Original: As an alternative to marginalizing the local gate predictions using p ( s ) , investigating the graphical model in Figure 2 reveals the possibility for marginalization with p ( s | x ∗ ) :
Revised: As an alternative to marginalizing the local gate predictions using p ( s ) , investigating the graphical model in Figure 2 reveals the possibility for marginalization with p ( s | x ∗ ) :

Original: [Equation]
Revised: [Equation]

Original: It is for this third evaluation case that training local generative models p ( x | s ) for each client becomes interesting, as they allow to compute the responsibilities p ( s | x ∗ ) at test time.
Revised: It is for this third evaluation case that training local generative models p ( x | s ) for each client becomes interesting, as they allow to compute the responsibilities p ( s | x ∗ ) at test time.

Original: In practice, however, the success of this approach depends heavily on the correctness of p ( s | x ∗ ) , which in turn depends on the ability of the local generative models p ( x | s ) to assign high probability to data that resembles
Revised: In practice, however, the success of this approach depends heavily on the correctness of p ( s | x ∗ ) , which in turn depends on the ability of the local generative models p ( x | s ) to assign high probability to data that resembles

Original: D s and low probability to out-of-distribution data.
Revised: D s and low probability to out-of-distribution data.

Original: Training and calibrating generative models for this task is in itself an active area of research (Nalisnick et al., 2018) and investigating how clients in a federated setting might exchange information to facilitate this process is not yet explored.
Revised: Training and calibrating generative models for this task is in itself an active area of research (Nalisnick et al., 2018) and investigating how clients in a federated setting might exchange information to facilitate this process is not yet explored.

Original: We therefore leave a thorough evaluation of this case to future work and only present here a MNIST (LeCun et al., 2010) experiment.
Revised: We therefore leave a thorough evaluation of this case to future work and only present here a MNIST (LeCun et al., 2010) experiment.

Original: We train FedMix with K = 4 and experts of two hidden layer ReLU MLP with 200 hidden units on MNIST.
Revised: We train FedMix with K = 4 and experts of two hidden layer ReLU MLP with 200 hidden units on MNIST.

Original: We split the dataset into S = 100 clients according to the procedure described in Liang et al.
Revised: We split the dataset into S = 100 clients according to the procedure described in Liang et al.

Original: FedMix achieves 97 .
Revised: FedMix achieves 97 .

Original: 7% average validation accuracy compared to 97 .
Revised: 7% average validation accuracy compared to 97 .

Original: 0% with FedAvg after 600 communication steps.
Revised: 0% with FedAvg after 600 communication steps.

Original: Independently for each client, we train a small variational autoencoder with a 32 -dimensional latent space using a two-layer MLP with 512 andhidden units respectively as encoder and mirrored decoder structure.
Revised: Independently for each client, we train a small variational autoencoder with a 32 -dimensional latent space using a two-layer MLP with 512 andhidden units respectively as encoder and mirrored decoder structure.

Original: We optimize the VAEs using
Revised: We optimize the VAEs using

Original: Adam with standard hyper parameters, B = 10 and perform early stopping on the local validation sets after no improvement for three epochs.
Revised: Adam with standard hyper parameters, B = 10 and perform early stopping on the local validation sets after no improvement for three epochs.

Original: After training, each client communicates their VAE to the server, where we evaluate p ( s | x ∗ ) to marginalize over the local gating functions according to p ( z | x ∗ ) =
Revised: After training, each client communicates their VAE to the server, where we evaluate p ( s | x ∗ ) to marginalize over the local gating functions according to p ( z | x ∗ ) =

Original: (cid:80) Ss =1 p ( z | x ∗ , s ) p
Revised: (cid:80) Ss =1 p ( z | x ∗ , s ) p

Original: ( s | x ∗ ) .
Revised: ( s | x ∗ ) .

Original: With this procedure, FedMix achieves 95 .
Revised: With this procedure, FedMix achieves 95 .

Original: 9% test set accuracy, compared to FedAvg with 96 . 9% .
Revised: 9% test set accuracy, compared to FedAvg with 96 . 9% .

Original: Marginalizing with p ( s ) instead of p ( s | x ∗ ) achieves 96 .
Revised: Marginalizing with p ( s ) instead of p ( s | x ∗ ) achieves 96 .

Original: 63% , showing the limitation of the approach in that any error in p ( s | x ∗ ) propagates into the expert assignment.
Revised: 63% , showing the limitation of the approach in that any error in p ( s | x ∗ ) propagates into the expert assignment.

Original: Reliable out-of-distribution detection capabilities in the individual estimators for p ( x | s ) are therefore necessary.
Revised: Reliable out-of-distribution detection capabilities in the individual estimators for p ( x | s ) are therefore necessary.

Original: D S PECIALIZATION
Revised: G L OCAL F INE - TUNING

Revised: In the main text we have compared algorithms by evaluating local data-sets on the server-provided global parameters.

Revised: For biased FedAvg , Local/Global and FedMix those parameters were made complete with the client-speciﬁc local parameters.

Revised: In Appendix F we have discussed the performance of FedMix when evaluating on a new shard.

Revised: Here, we discuss the evaluation when using locally ﬁne-tuned models.

Revised: Figure 13 shows the learning curves of Cifar10 when evaluating the validation set on parameters just before communication to the server.

Revised: This corresponds to local ﬁne-tuning for E = 1 epochs on the local data-set.

Revised: This form of model personalization is an option for all federated learning methods.

Revised: Due to the relatively small size of the local data-sets, this form of personalization can cause signiﬁcant levels of overﬁtting, as is evident from Figure 13.

Revised: Since regularization with respect to this form of personalization adds another level of complexity to the experimental design and hyperparameter tuning, we leave a thorough evaluation to future work.

Original: In Figure 9 we see how q φ ( z | y ) changes over time.
Revised: In Figure 14a we see how q φ ( z | y ) changes over time.

Original: The entropy of q φ ( z | y = c ) , i.e. the weight with which each expert is assigned to a speciﬁc label c decreases over time until each label is assigned exactly one expert with probability 1 . 
Revised: The entropy of q φ ( z | y = c ) , i.e. the weight with which each expert is assigned to a speciﬁc label c decreases over time until each label is assigned exactly one expert with probability 1 . The original MoE formulation, however, quickly collapses to a single expert, as can be see in Figure 14b. In the federated learning scenario, this corresponds to performing standard FedAvg on the single surviving expert. [Figure]

Revised: [Figure]

Original: E F ED M IX WITH EXPERT PRUNING
Revised: I F ED M IX WITH EXPERT PRUNING

Original: The FedMix algorithm (Algorithm 1) can be extended to ignore the communication of experts between server and client if, for this given client, that particular expert is not updated during training.
Revised: The FedMix algorithm (Algorithm 1) can be extended to ignore the communication of experts between server and client if, for this given client, that particular expert is not updated during training.

Original: If q φ ( z = k | s )
Revised: If q φ ( z = k | s )

Original: =
Revised: =

Original: E p ( y | s )
Revised: E p ( y | s )

Original: [ q φ ( z = k | y )]
Revised: [ q φ ( z = k | y )]

Original: = 0 for an expert k on client s , then parameters of that expert observe no gradient and the expert is effectively pruned from the local library of available experts.
Revised: = 0 for an expert k on client s , then parameters of that expert observe no gradient and the expert is effectively pruned from the local library of available experts.

Original: Since the parameterization of q φ ( z | s ) does not allow for exact values of zero, we introduce η , such that if q φ ( z = k | s ) <
Revised: Since the parameterization of q φ ( z | s ) does not allow for exact values of zero, we introduce η , such that if q φ ( z = k | s ) <

Original: η/K , we consider expert k to be pruned.
Revised: η/K , we consider expert k to be pruned.

Original: Since the value of q φ ( z = k | s ) during local optimization at the client is subject to change, it is possible that q φ ( z = k | s ) brieﬂy lies above the threshold even though it had just been pruned away.
Revised: Since the value of q φ ( z = k | s ) during local optimization at the client is subject to change, it is possible that q φ ( z = k | s ) brieﬂy lies above the threshold even though it had just been pruned away.

Original: Therefore the server considers expert k pruned for a given client s (cid:48) only if q φ ( z = k | s = s (cid:48) )
Revised: Therefore the server considers expert k pruned for a given client s (cid:48) only if q φ ( z = k | s = s (cid:48) )

Original: < 0 .
Revised: < 0 .

Original: 9 · η/K .
Revised: 9 · η/K .

Original: Algorithm 2 details this approach.
Revised: Algorithm 2 details this approach.

Original: Note that we write p ( z | s ) =
Revised: Note that we write p ( z | s ) =

Original: E x ∼ D s
Revised: E x ∼ D s

Original: [ p θ s ( z | x , s )] , i.e. the true marginal, however in practice we make use of the cheap to compute marginal approximation q φ ( z | s ) =
Revised: [ p θ s ( z | x , s )] , i.e. the true marginal, however in practice we make use of the cheap to compute marginal approximation q φ ( z | s ) =

Original: E y ∼ D s [ q φ ( z | y )] .
Revised: E y ∼ D s [ q φ ( z | y )] .

Original: Algorithm function S ERVER SIDE
Revised: Algorithm function S ERVER SIDE

Original: Initialize φ and K vectors W = [ w 1 , . . .
Revised: Initialize φ and K vectors W = [ w 1 , . . .

Original: , w K ] for round t in 1 , . . .
Revised: , w K ] for round t in 1 , . . .

Original: T do
Revised: T do

Original: S (cid:48) ← random subset of the clients
Revised: S (cid:48) ← random subset of the clients

Original: Initialize ∆ t W = 0 , ∆ t φ =
Revised: Initialize ∆ t W = 0 , ∆ t φ =

Original: Initialize p ( z | s ) = 1 /K
Revised: Initialize p ( z | s ) = 1 /K

Original: ∀ s for s in S (cid:48) do
Revised: ∀ s for s in S (cid:48) do

Original: W (cid:48)
Revised: W (cid:48)

Original: ← [ w k | p ( z = k | s ) ≥ 0 .
Revised: ← [ w k | p ( z = k | s ) ≥ 0 .

Original: 9 · η/K ]
Revised: 9 · η/K ]

Original: W ts , φ ts , p ( z | s ) ← C LIENT SIDE
Revised: W ts , φ ts , p ( z | s ) ← C LIENT SIDE

Original: ( s, φ , W (cid:48) )
Revised: ( s, φ , W (cid:48) )

Original: Store p
Revised: Store p

Original: ( z | s ) end for p ( s | z ) ← p ( z | s ) p ( s ) / (cid:80) s ∈ S (cid:48) p ( z | s ) p ( s ) for s in S (cid:48) do ∆ t w k + = p
Revised: ( z | s ) end for p ( s | z ) ← p ( z | s ) p ( s ) / (cid:80) s ∈ S (cid:48) p ( z | s ) p ( s ) for s in S (cid:48) do ∆ t w k + = p

Original: ( s | z = k )( w t − 1 k − w ts,k ) ∀ k ∆ t φ + = N s N S (cid:48)
Revised: ( s | z = k )( w t − 1 k − w ts,k ) ∀ k ∆ t φ + = N s N S (cid:48)

Original: ( φ t − 1 − φ ts ) end for ∆ t φ − = ∇ φ H ( (cid:80)
Revised: ( φ t − 1 − φ ts ) end for ∆ t φ − = ∇ φ H ( (cid:80)

Original: c
Revised: c

Original: q φ ( z | y =
Revised: q φ ( z | y =

Original: c ) p ( y = c ))
Revised: c ) p ( y = c ))

Original: w t +11: K ← A DAM ( ∆ t w 1: K , β ) φ t +1 ← A DAM ( ∆ t φ , β ) end for end function function C LIENT SIDE ( s, φ , W )
Revised: w t +11: K ← A DAM ( ∆ t w 1: K , β ) φ t +1 ← A DAM ( ∆ t φ , β ) end for end function function C LIENT SIDE ( s, φ , W )

Original: Get local parameters θ s for epoch e in 1 , . .
Revised: Get local parameters θ s for epoch e in 1 , . .

Original: , E do for batch b ∈ B do q (cid:48) φ ( z | s ) , q φ ( z | s ) ← E y ∼D s [ q φ ( z | y )]
Revised: , E do for batch b ∈ B do q (cid:48) φ ( z | s ) , q φ ( z | s ) ← E y ∼D s [ q φ ( z | y )]

Original: q (cid:48) φ ( z = k | s ) = 

Revised: q (cid:48) φ ( z = k | s ) =

Original: if q φ ( z = k | s ) < η/K q (cid:48) φ ( z | s ) ← q (cid:48) φ ( z | s ) / (cid:80) k q (cid:48) φ ( z = k | s ) q (cid:48) φ ( z | y ) ← q φ ( z | y ) q (cid:48) φ ( z = k | y )
Revised: if q φ ( z = k | s ) < η/K q (cid:48) φ ( z | s ) ← q (cid:48) φ ( z | s ) / (cid:80) k q (cid:48) φ ( z = k | s ) q (cid:48) φ ( z | y ) ← q φ ( z | y ) q (cid:48) φ ( z = k | y )

Original: = 

Revised: =

Original: if q φ ( z = k | s ) < η/K q (cid:48) φ ( z | y ) ← q (cid:48) φ ( z | y ) / (cid:80) k q (cid:48) φ ( z = k | y )
Revised: if q φ ( z = k | s ) < η/K q (cid:48) φ ( z | y ) ← q (cid:48) φ ( z | y ) / (cid:80) k q (cid:48) φ ( z = k | y )

Original: L s ← E q (cid:48) φ ( z | y b ) [log p w z ( y b | x b , z )]
Revised: L s ← E q (cid:48) φ ( z | y b ) [log p w z ( y b | x b , z )]

Original: +
Revised: +

Original: E q φ ( z | y b ) [log p θ s ( z | x b , s )]
Revised: E q φ ( z | y b ) [log p θ s ( z | x b , s )]

Original: φ + = α ∇ φ L s
Revised: φ + = α ∇ φ L s

Original: W + = α ∇ W L s θ s + = α ∇ θ s L s end for end for q ( z | s ) ← E y
Revised: W + = α ∇ W L s θ s + = α ∇ θ s L s end for end for q ( z | s ) ← E y

Original: ∼D s [ q φ ( z | y )]
Revised: ∼D s [ q φ ( z | y )]

Original: W (cid:48)
Revised: W (cid:48)

Original: ← [ w k | q
Revised: ← [ w k | q

Original: ( z = k | s ) ≥ η/K ] q (cid:48)
Revised: ( z = k | s ) ≥ η/K ] q (cid:48)

Original: ( z | s ) ← q ( z | s ) q (cid:48) ( z = k | s ) =
Revised: ( z | s ) ← q ( z | s ) q (cid:48) ( z = k | s ) =

Original: if q ( z = k | s ) < η/K q (cid:48) ( z | s ) ← q (cid:48)
Revised: if q ( z = k | s ) < η/K q (cid:48) ( z | s ) ← q (cid:48)

Original: ( z | s ) / (cid:80) k q (cid:48)
Revised: ( z | s ) / (cid:80) k q (cid:48)

Original: ( z = k | s ) (cid:46)
Revised: ( z = k | s ) (cid:46)

Original: Renormalization return W (cid:48) ; φ ; q (cid:48)
Revised: Renormalization return W (cid:48) ; φ ; q (cid:48)

Original: ( z | s ) end function
Revised: ( z | s ) end function