Efficient Federated Domain Translation 
Authors: Zeyu Zhou
Sheikh Shams Azam
Christopher Brinton
David I. Inouye
Venue: ICLR 2023 poster (2023)

Abstract: A central theme in federated learning (FL) is the fact that client data distributions are often not independent and identically distributed (IID), which has strong implications on the training process. While most existing FL algorithms focus on the conventional non-IID setting of class imbalance or missing classes across clients, in practice, the distribution differences could be more complex, e.g., changes in class conditional (domain) distributions. In this paper, we consider this complex case in FL wherein each client has access to only one domain distribution. For tasks such as domain generalization, most existing learning algorithms require access to data from multiple clients (i.e., from multiple domains) during training, which is prohibitive in FL. To address this challenge, we propose a federated domain translation method that generates pseudodata for each client which could be useful for multiple downstream learning tasks. We empirically demonstrate that our translation model is more resource-efficient (in terms of both communication and computation) and easier to train in an FL setting than standard domain translation methods. Furthermore, we demonstrate that the learned translation model enables use of state-of-the-art domain generalization methods in a federated setting, which enhances accuracy and robustness to increases in the synchronization period compared to existing methodology.

Published as a conference paper at ICLR 2023
EFFICIENT FEDERATED DOMAIN TRANSLATION
Zeyu Zhou, Sheikh Shams Azam, Christopher Brinton, David I. Inouye Elmore Family School of ECE, Purdue University (zhou1059, azaml, cgb, dinouye @purdue. edu
ABSTRACT
A central theme in federated learning (FI.) is the fact that client data distributions are often not independent and identically distributed (IID), which has strong impli- cations on the training process. While most existing FL algorithms focus the conventional non-III setting of class imbalance or missing classes across clients, in practice, the distribution differences could be more complex, e.g., changes in class conditional (domain) distributions. In this paper, we consider this complex case in FL wherein each client has access only one domain distribution. For tasks such as domain generalization, most existing learning algorithms require access data from multiple clients (i.e., from multiple domains) during training, which is prohibitive in FL. To address this challenge, we propose federated domain translation method that generates pseudodata for each client which could be useful for multiple downstream learning tasks. We empirically demonstrate that our translation model more resource-efficient (in terms of both communication and computation) and easier to train in FL setting than standard domain trans- lation methods. Furthermore, we demonstrate that the learned translation model enables use of state-oftte tefffthe-tttttt domain generalization methods in a federated setting, which enhances accuracy and robustness to increases in the synchronization period compared to existing methodology.
INTRODUCTION
Distribution shift across clients a well-known challenge in the Federated Learning (FL) community (Huang et al.) 2021). Most existing works have considered this from the perspective of class imbalance or missing classes (i.e., shift marginal distribution of classes) across clients, form of non independent and identically distributed (IID) datasets (Zhao et al. 2018). In particular, these works typically assume implicitly that the class conditional distribution of data is the same. In practice, however, the conditional distributions across different clients could be very different, e.g.. in computer vision, there is a shift in the data distribution (specifically, illumination) of images captured during the day versus night irrespective of the class label (Lengyel et al. 2021). This can lead to significant model generalization errors even we solve the issue of class shifts. Translating between datasets is one promising strategy for mitigating the more general shift across distributions of different clients. Moreover, it could solve the problem of Domain Generalization (DG) which requires model to generalize unseen domains (Nguyen et 2021)
domain translation model is one that can translate data from different distributions, typically attempting to align the conditional shift across distributions. In centralized settings, many translation methods have been proposed, such as StarGAN (Choi et al. 2018). However, in FL, domain translation models can be difficult to train because most existing methods require access data across all domains. Prior literature does not consider this natural setting of federated domain translation where domain datasets are distributed across clients.
In this paper. we empirically demonstrate that a naive implementation of state-of- state-of-the-an the--hherrt art (SOTA) translation models the Fl. context indeed performs poorly given communication limitations between the server and clients that often exist in practice (Azam et al. 2022a). Then, we propose leveraging an iterative translation model, Iterative Naive Barycenter (INB) (Zhou et al., (2022). which is much more amenable to FL training in terms of communication efficiency and data privacy considerations. We
"Currently Apple.

---

Published as a conference paper at ICLR 2023
empirically demonstrate that this modification obtains far superior performance to standard translation methods in the FL setting, and that it can aid in solving the challenge of DG in FL settings.
Our main contributions are summarized as follows:
We develop a federated domain translation methodology based on the recent iterative approach INB, which is more amenable to the FL setting than standard translation methods. We analytically show the equivalence between our federated algorithm and original INB which important for enabling usage of INB the federated setting.
We further propose several FL-motivated improvements to INB, including the use of variable- bin-width histograms, which significantly reduce communication costs.
We empirically demonstrate that our FedINB approach performs significantly better than standard translation models under the practical limited communication setting.
As one application, we demonstrate the feasibility of leveraging our federated translation model to aid in federated domain generalization. We also show that our federated DG method provides substantial improvements in robustness an increasing synchronization period, allowing reductions in communication overhead.
1.1 BACKGROUND: UNPAIRED TRANSLATION METHODS
Unpaired domain translation is the task of learning to translate between every pair of domains using only umpaired samples from each domain (Zhu et al. 2017). Formally, let M be the number of domains and Pm (z) denote the true m-th domain distribution. Let Xm, = {am ~ Pm Jimi denote the training dataset from the m-th domain distribution, where xm Rd is the number of samples per domain, and d is the number of dimensions. Also, let fm- denote the translation model from the m-th domain to the m'-th domain. Given this notation, the translation problem is usually formulated as minimizing a distribution divergence D (e.g., Jensen-Shannan Distance (JSD) for adversarial learning) between the translated and true distributions with some regularization term Rt
1
M
{f min glagg
Pere) + A AR(fm-mm')
1 THU 7m
where Ps..... is the distribution of the samples translated from the m-th domain to the m' -th domain, i.., the distribution of mm-rmm (FEm) where F' ~ Pm-
Standard GAN-based Translation Methods Zhu al. 2017) proposes CycleGAN, which esti-
mates unpaired translation models between two domains, using adversarial loss to approximate the divergence term and cycle consistency loss for the regularization term. StarGAN (Choi et al. 2018) extends CycleGAN by proposing a unified model for domain translation between multiple domains using a single translation model that takes the source and target domain labels as input. key issue with most existing translation models is that the computation of their objective requires access to data from all domains in the training. which is prohibited in an FL setting. For example, in StarGAN, to compute the domain classification loss for fake data, need discriminator trained on other domains. While the issue could be mitigated by federated algorithms such as FedAvg (McMahan al. 2017). this requires frequent global synchronization across domains and be hard to train as we show in Section[] While more advanced unified translation models exist (e.g.. StarGANv2 (Choi al. (0200), they are trained in similar ways to StarGAN and will suffer from the same drawbacks. Besides, many existing translation models learn pairwise translation (Zhu et al.. 2017 Park et al.) 2020) which would require an excessive computation and communication effort as the number of clients in an FL. setting increases. Thus, we focus on StarGAN in our experiments as an archetype model of standard translation methods.
Iterative Naive Barycenter (INB) In contrast standard translation approaches, the Iterative Naive Barycenter (INB) method (Zhou et al. 2022) builds up a deep translation model by solving sequence of much simpler problems that are highly amenable to the FL setting (as will be described in the next
section). INB learns deep invertible transformations T,, (L.) o...o (1) (where L is the number of layers) that map each domain distribution to shared latent distribution. Given these invertible
2

---

Published as a conference paper at ICLR 2023
(1)
0 T1
91
91
Backtraoking Search
8 122
0 02
92
bary
e AM
9M-
9M
Clients
Server
Clients
Server
Figure 1: High-level visualization of our Fed-multi max- K-SW methodology with VW histograms. At each iteration, there are transmissions: (1) histograms of data from each domain, (2) histograms of the empirical barycenter, (3) gradients Im, and (4) updated projection O. These transmissions the main sources of communication cost of FedINB which we minimize in this paper.
transformations, domain translation model can then be expressed as fm-mm & Tm o Tms i.e., translate the shared latent distribution and then inverse translate to any other domain. Importantly, each invertible layer ((9) is fit greedily by solving simple problems: (1) finding the projection matrix that maximizes the multi-distribution K-sliced Wasserstein divergence (Zhou et al., 2022). and (2) solving ID Wasserstein barycenter problems along the projection directions. Because ID Wasserstein barycenter problems are known to have closed-form solutions in terms of the domain ID CDFs, the key computational challenge is solving (1), which is formally defined as:
M
K
7L
ΣΣΣ i=
100 am ) D/q
max 0:0 0=1T MKn
where 2n E _dxn is the sample data matrix for the m-th domain, 0 10....., 0k), (0 xm)) signify the samples from the m-th domain distribution projected along the direction O₂ sorted in ascending order, D.j,k A M7 Σm=1 xmm)(g) is the empirical barycenter along direction Ox. is the dimension of the data, K < d is the number of projection directions, Ix E IR k is the identity matrix, is vector index. and is the sorted index. The corresponding algorithm is called multi-max-RR K-SW (Zhou et al. 2022 Algorithm 3). More details of INB can be found in Appendix
2
FEDERATED INB
2.1 PROBLEM STATEMENT: FEDERATED TRANST.TIOON WITH CLASS-CONDITIONA SHIFTS
We will make natural assumption that there is a one-to-one mapping between clients and domains, ie., each client only has access to data from its own domain. We will also extend the unpaired translation task to consider both a data instance x and its corresponding class label y, where Pm (z, y) will denote the joint distribution of and for the m-th domain. For this extension, conditional translation model fm (xly) which aims to translate from Pm (x: y) to Pm (xly) given class label y. Ultimately, our goal is to learn domain translation models across all clients such that they can be applied for other downstream federated tasks. Most existing non-lID FL works assume that the client distributions only exhibit class imbalance, i.e., the marginal class distributions are different (pm (v) * Pow (y)), but the class-cou flasscconiiionall distributions are equal (pm (x/y) = Pm (2]y)). In contrast, we focus on the case where the class-conditional distributions differ across clients.
seek class-
2.2 FEDERATED MAXIMUM MUL.T-DIITTRIUTTION K-SLICED WASSERSTEIN DIVERGENCE
In Algorithm|] I tion is done and when communication is needed in FL. For example, (Server) means the computation is done the central server and {Server Clients} means the following data will need to be
adapt the original multi-max: max---xxx-------- K-SW algorithm to show explicitly where computa-
3

---

Published as a conference paper at ICLR 2023
transmitted from the server all clients. A major problem with implementing INB in the federated setting is the loss of gradient informtation during transmission. In particular, the empirical barycenter xm tal becomes a constant (because it is transmitted and the gradient is not tracked) instead of a function of 0. However, we show that treating y constant will actually return the same gradient value. Formally, we make the following remark (proof in Appendix B).
Vii],k A 1
Mr Σm= (0
Remark 1. Fed-multi-max -aa-tiit-ia----wwww------------- and Moultimaax---See compute equivalent results despite Fed- multi-max-K treating y as a constart rather than a function of
Proof sketch. Each client computes biased gradient Vødm Vi dm + bm where bm is the bias of the gradient estimate. However, for the special case where the cost function is c(x, y) = ||æ y 2 these biases cancel out each other so that the sum of the biased client gradients is equal the true gradient, i.e., Em Vadm Em Vite dm Em bm = Em V dm because Σm b, = O.
Algorithm Fed-multt max- K-SW
Input: Samples from the M class distributions (æ) 22,...; BM), number of directions K, max number of iterations J
Output: Estimated projection matrix 0
(Server) Randomly initialize 0 E dx satisfying 0 TK. 0 = 00, 0x]
[Server -> Clients) 0
for j 21, J] do
[Clients} Compute projections and sort samples along each ID projection (0 am )
(Clients -> Server) (0 am)mm) )
[Server] Compute empirical barycenter Diii = 1
Em= (0 x
30
|Server Clients] V((]A
[Clients} Compute local objective dmm. MKn E Σ n 1=
10 3m 11
[Clients] Compute gradient 9m Vodm
{ (Clients + Server) 9m
[Server] Aggregate gradients g = Em Im. u [g. 0) v= 00 -gl
[Server] Update with backtracking line search 0 = 0 -u(12x + 3 uuuu)-a -2-=
|Server - Clients] O
if converge then Stop end if
end for
2.3 FEDERATED ID BARYCENTER
The next step of INB to solve the ID Wasserstein barycenter problems along each direction inde- pendently and estimate the corresponding invertible transportation map between marginal distribution and barycenter. The procedure can be found in Algorithm We first estimate histogram densities at each client and send those histogram bins the server. Then compute the inverse CDF of the barycenter at the server and send the histogram bins of the inverse barycenter back to each client. We summarize the full algorithm of FedINB in Algorithm of Appendix C
2.4 VARIABLE-BIa WIDTH HISTOGRAMS FOR MAX-K-SW K-SWD
Intuitively, the sorting operation on the clients in Algorithm| will improve data privacy because the samples are sorted independently for each projection direction. The joint dependencies between samples are broken and thus joint samples cannot be reconstructed from these sorted samples. From a distribution perspective, after sorting each direction independently, the transmitted data only contains information about the marginals of the domain distributions along the projection directions rather than information about the joint domain distributions. However, in practice, there may still be a concern that real data is being transmitted, as is often prohibited in FL. Additionally, since we are sending all data along those projection, the number of parameters being transmitted will increase linearly with the number of samples, which could lead high communication cost. Hence, propose to use histogram density estimators to approximate the empirical slice Wasserstein distance objective in multi-max- --S.. Specifically, instead of sending data directly, we first fit histogram
4

---

Published as a conference paper at ICLR 2023
densities at each client and send the histogram bins to the server. To avoid high computational cost, we use Variable- Bin- Width histograms (VW histograms) as defined next.
Definition (Variable- Biiiabble-i----ee Width histograms). Given a number of bins B, the probability of each bin is assumed to be equal to 1/B while the bin edges b, are selected as the uniformly spaced quantiles of the empirical distribution, i.e., bo (0), b, = (i (i/ B) Vi E (1..... B), where F-1 is pseudo-inverse of the empirical CDF of the samples and B is the number of bin edges.
By this definition, the samples in each bin the same (assuming the number of samples is exactly divisible by the number of bins). While a standard histogram has uniformly spaced bin edges where the probability of each bin depends on the data, VW histograms have uniform spaced cumulative probabilities where the bin edges depend on the data. A benefit of using VW histograms is that the computation of ID barycenter histograms is merely the average of histogram bin edges because the barycenter is based on the inverse CDF (in contrast, mixture model is based on mixing the CDF). After computing the barycenter histograms, we can readily employ the inverse CDF with uniformly spaced quantiles to estimate pseudo empirical barycenter samples y V/I.a A F. (i/n):
dm
m Ek=1 Plp EE (0T 1e
The visualization of Algorithm with VW histograms can be found in Figure Beyond communica- tion cost, VW histograms also provides improved data privacy as it sends a smoothed and compressed version of the marginal distributions rather than sending the sorted real samples. From an information theory perspective, VW histograms send less information because the histogram estimation is not one-to-one mapping. A more comprehensive discussion on privacy can be found in Appendix D
(more As the number of bins increases, the algorithm approaches the original Fed-multi-max---33 details in Appendix (.2). We empirically show that using limited number of bins is enough to lead to good optimization result and may even be preferred over a large number of bins because it may regularize the problem. To summarize, using VW histograms in FedINB improves communication costs and privacy while also possibly improving translation performance due to its regularizing effect.
2.5 PRACTICAL IMPROVEMENTS
We briefly introduce two practical improvements that can reduce communication costs further. The first simplifies the line search procedure of the original multi-max- --SW algorithm and is discussed in Appendix C.3 The second improvement uses autoencoders to reduce dimensionality because the computational and communication cost of the algorithm scale linearly with the data dimension d (as detailed in Section 2.6). Also, reduced dimensionality may improve the performance of INB because INB does not natively scale well to high-dimensional data.
Algorithm 2 Fed- 1ee-1DD-aaa 1D-BBaryeettr
Input: Samples from the M domain distributions (21, 22, 2M,, weight vector Output: Estimated invertible alignment maps (f1, f22," ... .ts)
[Clients] Estimate the ID CDF of Zmi F. HisoogamDeemiyyyssiaaaoooo (zm. [Clients - Server) Fm
(Server) Estimate the inverse CDF of barycenter Frury Σ m Wen F-
(Server -> Clients} Foery
[Clients) Construct invertible alignment map f. Poory o Fm
return (t1, (2, ... (tas)
2.6 COMMUNICATION COST ANALYSIS OF FEDINB
We will now discuss the communication complexity of FedINB in terms of the number of parameters to be transmitted from each client to the server. Let B be the number of histogram bins for the CDF of local slices and the inverse CDF of the barycenter in Algorithm 2 and let V as the number of histogram bins multi-max- K-SW requires us to transmit slices of samples, empirical barycenter samples, projection matrix and gradient. Thus, the number of parameters needed to be transmitted is O(K (n d))
for VW histograms in Algorithm [ At each iteration at each layer of INB, Fed-
5

---

Published as a conference paper at ICLR 2023
per iteration of Fed-multi-max-Ke K-S.. At each layer of INB, the total number of parameters being transmitted for red-multi-max K-SW is OJJK(n + d)), where we assume that the number of backtracking iterations is at most small constant more than the maximum number of iterations = (J). For Fed-Bary (applying Algorithm 2 to K dimensions simultaneously), we need to communicate histogram bins between the client and the server. The number of parameters being transmitted is (K After applying VW histograms and the simplification for backtracking line search, the total communication cost of each layer of INB reduces to OVJK(V + d)) per layer (assuming B & JV, which is true in practice). Thus, the communication is linear in terms of all the relevant parameters, and reducing the dimensionality d with an autoencoder will immediately provide communication benefits for high-dimensional data. More details can be found in Appendix
3 RELATED WORK
Beyond the translation models introduced in Section 1.1 we discuss other related works here.
Konevený et al. (016) originally proposed the FL framework for promoting data privacy (Shokri Shmatikov 2015 Azam et al. 2022b) and enhancing communication efficiency (Wang al. 2020 Recent years have seen increasing interest in the adaptation of FL frameworks to settings of non-lII data with respect to the class labels. For example, Zhao et al. (2018) propose sharing subset of client data to create global data pool that reduces the gradient divergence across different clients with non-lID data. Another line of work (Lin et al., (021)) proposes to tackle client data heterogeneity by incorporating distributed average model consensus within a semi-decentralized architecture. However, these works do not consider the larger set of complex data heterogeneity settings that include a shift in input data distribution (or domain distribution) across clients.
Federated learning
Domain generalization Domain generalization is an application of domain translation that targets training models that can generalize to unseen distributions. While most existing methods require access to data from all domains in the training, there are several methods that can be used as regularization for training at each local client. For example, Carlucei al. (2019) regularize the model's 's training by solving solving jigsaw puzzles on each image. Zhang et al. (2020) apply deelply stacked transformation on the source domains to simulate the possible domain shift. Another main branch of domain generalization is to learn domain-invarian features. DIRT (Nguyen et al. (021)
explicitly employs a pretrained domain translation model by minimizing the following objective:
MM Fp... ( (().(())) + 1 M E. M
where g(z) is encoder trying to find the bomainn-nnvrrian feature and e is the classification loss given the representation g(x) in the encoded space. In their paper, they use StarGAN as f, and as discussed in Section 1 propose use FedINB for federated domain generalization.
Understandably, there are few works that consider the problem I in federated setting owing to its high complexity as discussed in Section Liu et al. (02I) solve the lack of domain knowledge across clients by creating a common bank of amplitude spectrum of data which is aggregated by accumulating the amplitude spectrums obtained by application of Fast-Fourier Transforms (FFTs) on the client data. While this trick leads to significant performance gains, it significantly reduces privacy because the clients essentially sharing half of their datasets (i.e., their amplitude spectrums) with all other clients. A more detailed discussion on privacy can be found in Appendix D Besides, the proposed method by Liu al. (202l) is specifically developed for image segmentation tasks and it is not clear how can the same be extended to more general task and other data modalities such as text, tabular or other modalities.
Federated domain generalization
of domain generalization!
Note that this is different from domain adaptation. Suppose we have source domains and target domains. In DG, each source (client) has access to both data instance æ and class label y. and the goal is to generalize to target domains (new clients) where neither are observed in the training. In domain adaptation, except for and y in the source domain, we also observe the data instance z of the target domain in the training.

---

Published as a conference paper at ICLR 2023
4 EXPERIMENTS
Before presenting our experimental results, we first introduce a naming convention to describe many different setups of FedINB. For simplicity, we will remove the "Fed" prefixes and call it INB since all experiments are done in a federated setting. With the autvencoder, the model is called AEINB if using shared autoencoder for all domains, and IndAEINB if using autoencoders trained on each domain separately. When using VW histograms to replace the transmission of real data in Sid-multi-max-K-SSy K-SW, call them HistINB, HistAEINB, HistIndAEINB separately. The number of layers L, number of dimensions K and max number of iterations (in Ted-multi-max-K-SW J are three key hypeparameters
tune when building INB model. We will represent them in the format L-K-J. As an example, "HistInAAEINB-LIOS" MO-NB--LO-KKO-JOOO 10-J100" refers to an INB model with L 10, K 10, J 100, and it incorporates an autoencoder trained locally and uses VW histograms when optimizing for 0.
Datasets Following the setup in Zhou et al. (2022), we test FedINB on Rotated MNIST and FashionMNIST GGhifary al. 2015). For the federated domain translation experiment, there are clients participating in the training and each has data from one of domains 0, 15, 30, 45, 60, where the number represents the degree of clocker-e rotation and means the original images. For the federated domain generalization task, we use the translation model trained on domains 0, 15,30, 30, 45, 60. We test the classification model's generalizing ability to domains and 90.
Metries For the federated domain translation experiment, we use the empirical Wasserstein Distance (WD) and FID score (Heusel et al. 2017) between the original samples and translated samples as eval- uation metric. Wasserstein Distance computed WD
W (D fm (frm) where each WD is computed with the Sinkhorn algorithm (Cuturi 2013). The FID score is com- puted as FID 1 Σm Lmn FID fmm (2m)). The WD and FID metrics reported are the M
E.m
average of 10 classes of digits since we are conducting class-wise translation. All FID score results can be found in Appendix [F] To compare communication cost, we consider the number parameters needed be transmitted between one client and the server. We also measure the performance of our federated translation models by using it for downstream federated DG tasks.
4.1 PRACTICAL IMPROVEMENT OF FEDINB
Due to space limitations, only show figure for each investigation with Rotated MNIST. More figures with different setups of INB can be found Appendices [. and F22 Figures with Rotated
FashionMNIST. qualitative results and investigation of the FedINB optimization can be found in Appendices F3 and F.6 respectively.
Autoencoder. We investigate how adding autoencoders part of the translation model can help with translation performance and communication cost. As can be seen in Figure 2a both AEINB and IndAEINB achieve lower WD and lower communication cost than INB while other parameters kept the same. Additionally, we show that by using an independent autoencoder for cach domain separately, we can obtain lower WD for the same communication cost while simultaneously reducing the work of training federated shared autoencoder.
Insensitivity J. In the original INB paper, the authors set 200 for MNIST (LeCun & Cortes 2010) and FahsionMNSIT (Xiao et al. 2017). This could lead to high communication cost due to frequent transmission of data (even if the data size is small). We empirically show that even when the optimization algorithm does not converge and we stop early, it will not affect the final performance much. One explanation for this is that the projection matrix found suffers less from overfitting. As shown in Figure 2b after 10 layers of INB, the model with J 50, 100, 200 converge to similar point. Even the model with J = 30 almost begins to converge at the same point; only when / is decreased to a quite small value like 10 will the convergence of INB require more layers.
VW histograms, As shown in Figures 2c 2d and 2e when using VW histograms and decreasing V. the communication cost is significantly reduced. In comparison to the number of samples being transmitted by INB and IndAEINB (n = 10000), V = 500 can tolerate the transmission of substantially less samples while achieving similar performance.
7

---

Published as a conference paper at ICLR 2023
AE: L10-K10-J200
IndAEINB: L=10, K=10
VW hist(V=500): L10 1100-000
Avg. Wasserstein-
35.0
INB
Avg. Wassersteins
27.8
J-200
400
32.5
AEINB
J-100
IndAEINB
IndAEINB
27.4
1-50
Wasse day
35
HistINB
1-30
HistinAEINN
27.0
1=10
27.5
30
26.6
25.0
0.0
1.0 Communication Cost 1e9 18 params sent)
Communication Cost leg
Communication Cost 1eg
[# params sent)
(# params sent)
(a) Different autoencoders
(b) Numer of iterations
(c) VW histograms
HistINB: L10-K10-J100
HistIndAEINB:
L10-K10-J100
FedStarGAN
Aug. Wassersteing
45
V-2000
5=500
Avg wasserstain
40
V=200
V-500
AVg. Wassersteinz
200
FedStarGAN-4
HistndAEINB
40
FedStarGAN 16
35
FedStarGAN-122
35
100
30
30
2.5
Communication Cost 1c8
Communicatiin Cost 1e8 (# params sent)
Communication Cost 1c8 (# params sent!
(# params sent)
(d) Influence of
(e) Influence of V
(I) Convergence of FedStarGAN
Figure 2: Wasserstein Distance (lower is better) as communication cost increases for Rotated MNIST.
(a) Different autoencoders used for INB. (b) Different number of inner iterations used for IndAEINB. (c) Using Atriale-BBinWidiee (VW) histograms for INB. (d) HistINB: different number of histogram bins used for VW histograms. (e) HistIndAEINB: different number of histogram bins used for VW histograms. (f) Failure of FedStarGAN. The INB used here is HistIndAEINB-LIOS H0--100-K1O-J 10-J0000 110 with
V 500. The starting points of HistAEINB and FedStarGAN differ because we begin tracking the WD after applying autoencoders to INB (details can be found in Appendix E
4.2
COMMUNICATION EAPENIICATTOONEEEICEENT DOMAIN TRANSLATION
We compare with FedStarGAN to demonstrate that our translation model is much more efficient in terms of both communication and computational cost. We empirically show that the training of StarGAN fails in the federated setting. We adapt StarGAN to the federated setting by training individual StarGAN models on each client (using the local data, i.e., single domain) and synchronizing the client generator and discriminator after a certain number of mini-batches (controlled using the sync_stepp hyperparameter) using FedAvg. We try various sync_step sizes to fairly evaluate the federated extension of StarGAN (FedStarGANN, As note, FedStarGAN- 128 means using FedAvg to train StarGAN and synchronizing after each 128 mini-batches (sync_stee=1222
As shown in Figure 3 a fixed communication cost, even though FedStarGAN can generate images with good quality, it fails to translate it to the target domain, i.e., can only reconstruct the input but not translate a rotated version of the image. This is in line with our observation during training that while FedStarGAN can optimize for reconstruction loss, it does not converge over the domain adversarial loss. Moreover, in Figure 21] we observe that FedStarGAN- fastest) converges to its starting point, justifying our argument that it can only reconstruct images. Additionally, is important to note that INB takes around minutes to finish training on a single RTX A5000 GPU while FeNStarGAN takes around 2.5 hours to generate good samples on a single Tesla P100 GPU (notice it can only reconstruct samples rather than translate samples).
(which converges
4.3 FEDERATED DOMAIN GENERALIZATION VIA PSUEDO TRANSLATED DOMAIN DATA
Once the translation models have been trained in a federated way, clients can exchange their models so that each can translate to every other domain. This enables the training of DIRT (Nguyen et al. 202T) for DG using these translation models (that were also learned via FL). Specifically, we combine FedAvg and DIRT named FedDIRT where each client m. minimizes the following local objective:
Ep.,, (o))))))) g(x)) + 1 M EM= dis(g(z), fm (x)))
(5)
StarGAN is trained with the standard FedAvg algorithm (MMMahan et al. 2017). To improve performance, we synchronize quite frequently- after certain mini-batches rather than certain epochs.
8

---

Published as a conference paper at ICLR 2023
O
2
3
4
0
/
2 3 4
0
/
2 7 2 7
4 9 4
3
5
7
8 3
9
S
6
7
8 9
5
6
4
#
0
/
3 8
5
11
1
11
1
5
6
9
(a) INB
(b) FAN-16 (c) FeNStarGAN-128
Figure 3: Qualitative results for federated domain translation using INB and FedStarGAN when around 2 x 10t parameters transmitted. The first two rows are original samples from domain O and the last two rows are samples translated to domain 60. (a) Our BistnndAENNB-LJOO 10-KK0-J1100
algorithm: It translates the digits to the expected domain. (b) FedStarGAN after 160 iterations and Fedavg every 16 steps: It fails to generate legible digits. (c) FedStarGAN after 1280 iterations and
FedAvg every 128 steps: The generated samples are well reconstructed but are not translated.
Table 1: Classification accuracy unseen domains for Rotated MNIST. For each target domain, consider aggregating the model after each mini-batch and 100 mini batches. The mean and standard deviation are taken over 4 runs. For FedDIRT, use HistIndAEINB with J 100 and V 500.
Model
(I-bcc)))
75 (10-bath)
75 (100-bath))
(1-bhh))))
(10-hatch)
90 (100-
FedDIRIL22-K10)
92.2 1
89.3 I. 4.2
91.9 0.9
69.8 1 0.7
70.4 2.2
69.9 0.3
FedDIRTYL2-K200
92.2 1 1.3
90.8 1. 1.2
91.4 0.6
69.4 1. 0.8
71.1 1.3
69.8 1.5
FedAvg
85.2 1 0.7
85.1 1 0.5
80.1 2.3
63.8 2.1
63.6 1 0.8
55.6 1 2.0
Table 2: Classification accuracy on unseen domains for Rotated FashionMNIST.
Model
75 (1-batch)
75 (10-batch)
75 (100-batch)
(1-batc)))
(10-batch)
90 (100-batch)
FedDIRITL20-K10
65.8 1.3
63.6 I 0,9
63.0 1
18.2 18.2.1.0 0.9
18.1 1 0.4.
18.8 0.6
FedDIRT((22-K3R
65.5 +
64.5 + 1.5
63.5 +
19.0 0.3
18.5 1.0
18.9 0.4
FodAvg
51.9 + 2.9
50.4 + 1.2
40.4 + 4.1
13.7 + 1.6
14.6 + 1.8
13.1 + 0.9
and the server aggregates the model after certain mini-batches. Note that this local objective only uses data from the m-th client but leverages the shared FedINB translation models for the DIRT regularization. We run FedAvg without any domain translation regularization loss term as our baseline. For fairesss, we use the same CNN structure for both models. Table I and Table 2 give the results for training the domain translation model on domains 0,15,30,45,66 and using it to train a model that can generalize to unseen domains 75 and 90. More results can be found in Appendix F We that regularization from domain translation model significantly improves the model's 's ability to generalize unseen domains. Additionally, (i.e., decrease the FL aggregation frequency), which is an important objective in communication- constrained FL systems (Lin et al.) 2021), FedDIRT achieves similar performance as before while the performance of FedAvg drops significantly. The regularization helps make the federated training more stable. This can also justified by the higher standard deviation of most FedAvg runs.
observe
notice that we increase the synchronization period
5 DISCUSSION AND CONCLUSION
Despite the theoretical and empirical advantages of FedINB, some limitations should be noted for future work. For example, it is difficult for FedINB to translate high dimensional images: possible solution is to further reduce the dimension via a (pretrained) deep neural network in the federated domain generalization task. Also, we have not provided a theoretical privacy guarantee for FedINB, which could possibly be accomplished by utilizing DP-SGD (Abadi et al. 2016) for Fed-multi-max--o S-SW. More discussion on these points can be found in Appendix C.5
In this paper, we proposed a federated domain translation approach (FedINB) that can mitigate non-IID issues in federated learning tasks. Our model can handle the harder task conditional shift in comparison to most existing FL methods. We then proposed several improvements to FedINB that decrease communication costs for practical resource-constrainee FL systems. We empirically showed that our translation model performs substantially better than FL versions of standard translation models (StarGAN). Finally, as an application, we demonstrated that combining FedINB with SOTA domain generalization methods leads strong performance in federated domain generalization.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
Z.Z. and D.I. acknowledge support from NSF (ISS-2212097) and ARL (V9112N--000--21,, S.A. and C.B. acknowledge support from NSF (CNS-2146171) and ONR (N000142212305),
REFERENCES
Martín Abadi, Andy Chu, lan J. Goodfellow, H. Brendan McMahan, lyya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Edgar R. Weippl, Stefan Katzenbeisser, Christopher Kruegel, Andrew C. Myers, and Shai Halevi (eds.), Proceedings of the 2016 ACM SIGSAC Conference on Computer and Communications Security, Vienna, Austria, October 24-28, 2016, pp. 308-318. ACM, 2016, doi: 10.114522976449.2778311 URL httpp:://dott org/
10.1145/2976749.1 2997733
Sheikh Shams Azam, Seyyedali Hosseinalipour, Qiang Qiu, and Christopher Brinton. Recycling model updates in federated learning: Are gradient subspaces low-rank? In International Conference on Learning Representations, 2022a.
Sheikh Shams Azam, Taejin Kim, Seyyedali Hosseinalipour, Carlee Joe- foe-Wong, Saurabh Bagchi, and Christopher Brinton. Can we generalize and distribute private representation learning? In International Conference on Artificial Intelligence and Statistics, pp. 11320-11340, PMLR, 2022b.
Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 950-959. PMLR, 2020. URL http: /ttt:://rroe http:://ppccceednnn... mlr. wwt.pppeessssbeeeeeeeeee html
Fabio Maria Carlucci, Antonio D' D'Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi.
Domain generalization by solving jigsaw puzzles. In IEEE Conference on Computer Vision and Putlern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pp. 2029-2222322 2238 Computer Vision Foundation IEEE, 2019. doi: 10. 1099/CVPR.2019.000
URL http: com content CVPR 2019 htt099/hmmmm///m Sattmn/CCaruccee Domain.
7/openaccess.
thecvf. tereenaccess.theef.....
Generalization by Solving Saaiiigg_iggsaww Puzzles CVPR 2019
Yunjey Choi, Min-Je Choi, Munyoung Kim, Jung-Woo Ha, Sunghun Kim, and Jaegul Choo. Stargan: Unified generative adversarial networks for multi-domain image-to- image translation. In 2018 IEEE Conference on Computer Vision and Pattern Recog-
nition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pp. 8789-87977 Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10.1 10.1109/CVPR. 2018.00916. URL Choi StarGAN Unified Generative CVPR 2018 paper. html
thecvf. com/ www.co//ccnceeeeee content cvpr 2018/html/
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pp. 8185-8194, Computer Vision Foundation / IEBE, 2020. doi: 10. N0.11090/2 200.206000022000002222 theevf. . www/ content CVPR 2020/htmll httphhhmmhhhhhha StarGAN v2 Diverse Image Synthesis for Sor_MMutipllee Domains CVP__ 2020 paper. html
URL littps://ppeaccess
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Christopher C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger (eds.), Advances in Neural Information Processing Systems 26: 27th Annual Conference Neural Information Processing Systems 2013. Proceedings of a meeting held December 5-8, 2013, Lake Tahoe, Nevada, United States, pp. 2292-2300, 2013. URL httpppppp/pooeee.nggggg nwwrips.......m ccc cc/paper/ html
Biwei Dai and Uros Seljak. Sliced iterative normalizing flows. In Proceedings of the 38th Interna- fional Conference on Machine Learning, ICML., 2021.
10

---

Published as a conference paper at ICLR 2023
Cynthia Dwork. Differential privacy. In Michele Bugliesi, Bart Preneel, Vladimiro Sassone, and Ingo Wegener (eds.), Automata, Languages and Programming. 33rd International Colloquium, ICALP 2006, Venice, Italy, July 10- 14. 2006, Proceedings, Part II, volume 4052 of Lecture Notes in Computer Science, pp. 1-12. Springer, 2006. doi: 10.1007/1 1787006 77doi
URL https
1007 787666
Cynthia Dwork and Aaron Roth. The algorithmic foundations of differential privacy. Found. Trends Theor Comput. Sei., 9(344)211--00) 2014. doi: 10. 00.156/0440000042,, URL https
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. Imagenet-trainee cnns are biased towards texture; increasing shape bias improves accuracy and robustness. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL |attpa://openzeviii
net (eet/Eoumm: -endiddd-yyn 3ygh J09KX
Muhammad Ghifary, W. Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 2551-2559. IEEE Computer Society, 2015, doi: 10.100//CC.....5.2.22. URL httph: or0.0....0.0000000....00.0...... ICCV. 2015.293
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. In Isabelle Guyon, Ulrike von Luxburg. Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 6626-6637, 2017. URL httpppts://poccceeeiiiiis
www.diinnnnnnnnnuueueeeeeee neur iisi..pppp.. ps. CC wttppppppppppp..pcppccccccccepee 2017 201///211///////// / 8a (a11533777rab0erreee111116666722
Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangehuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning on non-iid data. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
Jakub Konevený, H. Brendan McMahan, Felix X. Yu, Peter Richtarik, Ananda Theertha Suresh, and Dave Bacon, Federated learning: Strategies for improving communication efficiency. In NIPS Workshop Private Multi-Pani Marti-Party Machine Learning, 2016. URL http:/axxxx.o../aaaa/a/ abs/
1610. 05492
Yann LeCun and Corinna Cortes.
MNIST handwritten digit database.
2010. URL httth://aaann.... http://aannnecu.....mm http://cc//////////////c/a/ cttp///b//x//i////
Attila Lengyel, Sourav Garg, Michael Milford, and Jan C van Gemert. Zero-shot day-night domain adaptation with a physics prior. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4399 4409, 2021.
Ninghui Li. Wahbeh H. Qardaji, Dong Su, Yi Wu, and Weining Yang. Membership privacy: a unifying framework for privacy definitions. In Ahmad-Reza Sadeghi, Virgil D. Gligor, and Moti Yung (eds.), 2013 ACM SIGSAC Conference on Computer and Communications Security, CCS' 13, Berlin, Germany, November 4-8, 2013, pp. 889-900. ACM, 2013. doi: 10.1145/2508855/25516668 URL /ttp:://doo..ooo 01t////....001101111511//333333 00199.25660 251225566666 6686
Frank Po-Chen Lin, Seyyedali Hosseinalipour, Sheikh Shams Azam, Christopher G Brinton, and Nicolo Michelusi. Semi-decentralized federated learning with cooperative d2d local model aggregations. IEEE Journal on Selected Areas in Communications, 39(12):3851-386, 2021.
Quande Liu, Cheng Chen, Jing Qin, Qi Dou, and Pheng-Amn Heng. Feddg: Federated domain generalization medical image segmentation via episodic learning in continuous frequency space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1013-1023, 2021.
11

---

Published as a conference paper at ICLR 2023
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera Arcas. Communication efficieen learning of deep networks from decentralized data. In Aarti Singh and Xiaojin (Jerry) Zhu (eds.), Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FI,, USA, volume 54 of Proceedings of Machine Learning Research, pp. 1273-1282. PMLR, 2017. URL http:
https/////mmmmmaanmmmmmmm //proceedings.o mir prwss v54 wwm.mannia..cmccmmmmmmmmmmm....a aa html
A. Tuan Nguyen, Toan Tran. Yarin Gal, and Atilim Gunes Baydin. Domain invariant rep- resentation learning with domain density transformations. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Ad- vances Neural Information Processing Systems 34: Annual Conference on Neural In- formation Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 5264 5275, 2021. URL httpstttp//tooeediiiiggg 2 08271795575 8b 2ac.155555882222222cc11333333333
neurips...... wttppppp///2/..///// 2021/has// html
Taesung Park, Alexei A. Efros, Richard Zhang, and Jun- Yan Zhu. Contrastive learning for unpaired image-to-image translation. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm (eds.), Computer Vision ECCV 2020 16th European Conference, Glasgow; UK, August 23-28, 2020, Proceedings, Part IX, volume 12354 of Lecture Notes in Computer Science, pp. 319 345. Springer, 2020. doi: 10. 1007/977 100-00305055----- 10. 15.1007/911
URL https://doi. org/
030 -30-0055500 58545-7-199
Reza Shokri and Vitaly Shmatikov. Privacy-preserving deep learning. In Proceedings of the 22nd ACM SIGSAC conference on computer and communications security, pp. 1310-1321, 2015.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:761 1-7623, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms, 2017.
Ling Zhang, Xiaosong Wang. Dong Yang. Thomas Sanford, Stephanie A. Harmon, Baris Turkbey, Bradford J. Wood, Holger Roth, Andriy Myronenko, Daguang Xu, and Ziyue Xu. Generalizing deep learning for medical image segmentation to unseen domains via deep stacked transformation. IEEE Trans. Medical Imaging. 39(7):2531-254)) 2020. doi: 21110//MM..000.0990
URL
https://doi www.ddo..ogg/.... 1tt.//10.1.09MM1M1 2tt.://M....00....9333
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint / 00582, 2018
Zeyu Zhou, Ziyu Gong, Pradeep Ravikumar, and David Inouye Iterative alignment flows. In Proceedings of The 25th International Conference Artificial Intelligence and Statistics, 2022.
Man Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 20/7 pp. 2242-22511111 IEEE Computer Society, 2017. doi: 10.1 10.110//CC.22017.24444 URL / 1109/ ICW///CCC.... 2017.241
12

---

Published as a conference paper at ICLR 2023
A OVERVIEW
We have organized our appendix as follows:
Appendix includes the proof of Remark !
Appendix C includes more details about the INB algorithm, its implementation, limitations and future works.
Appendix includes more discussion on privacy.
Appendix includes experiment and implementation details.
Appendix includes more experiment results.
B PROOF OF REMARK
Before starting the proof of Remark 1 let us first introduce the following definition and lemma to deal with the sorting operation in the objective.
Definition 2. The sorting permutation matrix Pz E ka associated with vector Z E RR is defined as ZPz = sorted( Z) such that (ZP2); S (ZP2)i+1, =0, O....,
Lemma Let A Ixdd X E gdxx, and denote the sorting permutation matrix of AX as PAX. We have:
0(AXPAx) (XPAx) 0A
Proof. Note that AXPAX = (Ax)(AX) where sort the sorting function. With this, we first note that the Jacobian of sort with respect it's inputs is equal to the permutation matrix that sorts that input.
Json(3) Osort(2)
(2)
PT
Oz
where v = P. is sorted, i.e., 2/2 $ 2/4,, Vi < n Blondel et al.
2020
Now we want the following:
@sort(AX) @sort(A. (AX) 0AX 0A 0(AX)
pm x = (XP)T.
0A
Now prove Remark
Proof of Remark [ This remark in essence states that when compute the gradient of sum of the loss among all clients with respect to the projection matrix, the empirical barycenter could be regarded constant (zero gradient). Intuitively, the proof relies the fact that using squared distance, and the empirical barycenter can be seen as a mean value over domains.
Let A E Rixd X2 Rdxn
1....M M and define the loss term
M Σ
M
L
AXPPA
Σ (X
PAX,)
M
Note that L is equivalent to Σ M
(AX.) 3M (AX,)) Then define the loss I as
L
AXX 2
Blondel, Mathieu, et al. "Fast differentiable sorting and ranking." International Conference Machine Learning. PMLR, 2020.
13

---

Published as a conference paper at ICLR 2023
where C 3/ Σド= AX5Pax, but will be regarded as a constant when computing the gradient of L And for simplicity, we denote X; X,PAx, We have:
OL
2 Σ (AX,PAX.
Σ AAX )
M
PAX,
1
AX,,
0A
M j=1
DA
M
2 (AX; AX)
1
Σ X
M
j=1
j=1
Σ AX 4
M Σ Σ
A
ΣΣ AX!XT
M 3
= 2 Σ AXXXT
Σ Σ AX;X;"
M
where (1) uses Lemma Additionally,
0LC
= E(AXPAX - -0) (AX)
0A
-2 (AX -OX
M D (AX; Σ AX M
j=1 = 2 Σ AXXXT
M ΣΣ AXX 1
AL
0d
C MORE DISCUSSION OF THE FEDINB ALGORITHM
C.I ORIGINAL INB AND FEDINB
Zhou et al. (2022) propose the Iterative Naïve Barycenter (INB) algorithm to align high dimensional distributions via decomposing the problem into simpler ID alignment problems with closed-form solutions. Specifically, they first reduce the dimension via a orthogonal projection matrix found by minimizing maximum K-sliced Wasserstein divergence as shown in Equation (2). Informally, this objective finds orthogonal directions where the empirical Wasserstein-2 distance to the empirical Wasserstein barycenter is maxinized-where where both the empirical Wasserstein-2 distance and barycen- can be solved in closed-form only using 1D sorting. They adopt the optimization approach of (Dai & Seljak 202T) to optimize over the Stiefel manifold of orthonormal matrices via projected gradient descent with backtracking line search. Then, along each dimension, they solve the ID Wasserstein Barycenter and find the corresponding mapping to the Barycenter. By iteratively finding the projection matrix and ID mapping the Barycenter, they construct deep translation model between distributions. More details about INB can be found in (Zhou et al. 2022).
In our paper, we propose adapting INB for federated domain translation. The main FedINB algorithm employing Algorithms 1&2 is given in Algorithm 3
C.2 VW HISTOGRAMS
One key improvement made for FedINB is to use VW histograms to approximate the empirical slice Wasserstein distance objective in Algorithm It noteworthy that if the number of histogram bins the same as number of samples, the FedINB with VW histograms becomes equivalent to INB. In this case, the quantiles sent would just be the samples and each of them has uniform density. As
14

---

Published as a conference paper at ICLR 2023
Algorithm 3 Federated Iterative Naive Barycenter
Input: Samples from M clients 21, 22, ................................ MM. number of directions K, number of itera- tions/layers
Output: Estimated domain translation models (71, .e..... TM)
[Client] T - id. Vm M for é
L} do
]Client] Vm, 2m + Tm (Im)
|Server, Client] x- Fed-multi-mad max- max-n-Sw ((M) (Algorithm]
for k k= ( K} do
(Client) Vm, zm 0% 2m
|Server, Client] t1, tM. = ID-Barycenter<2(1.... 2/1)) [Algorithm] 2
end for
|Client] Vmn, Im f- [m..1...,, tm.k]
[Client) Vm, Tm (ax) t- ml (0 Tm (2)) end for
return (Ti, T2. TM)
show empirically in Section 4 and Appendix F even if the number of bins is smaller than the number of samples, this approximation will not hurt the final performance and can reduce the communication cost.
C.3 COMMUNICATION SIMPLIHCATION FOR BACKTRACKING LINE SEARCH
When optimizing the projection matrix 0, the original line search which requires recomputing the loss with an updated 0 where only the learning rate is changed. Thus, given fixed maximum number of iterations J, we may need to run extra iterations J'. At each backtracking iteration, updated & to recompute the loss so that we can determine whether to accept this update. We propose locally record the change of so that if backtracking line search determines a reduction in the learning rate, we can reverse-engineer the new updated 0 without explicit synchronization of the new 0.
algorithm does backtracking
need to compute slices and the empirical barycenter with
C.4 COMPUTATIONAL COMPLEXITY AND COMMUNICATION COST
Computational complexity. The complexity of the maximization (Algorithm is O((J J')nM l) + K2d + K))), where J, J,',"" M. n, d, K are the number of iterations for Red-multi-max-K-SW extra iterations for backtracking line search, domains, samples per domain, dimensions, and latent dimensions, respectively. If we use VW histograms, the complexity reduces log V) K2d + K))). The complexity of the minimization (applying Algorithm for each latent dimension independently) is O).MK..
[(J + JV
Communicution cost. We gave an approximated communication cost in the main paper. Here, we provide more exact estimation. At each iteration of INB, for each client, the communication cost of the maximization (Algorithm [ is O(( + J') (n + d)) and O(J + KV + d)) for FedINB and FedINB with VW histograms respectively. The communication cost of minimization (applying Algorithm 2 to dimensions simultaneously) is O (BB, + B2)) where B, and B, are the number of histogram bins for the CDF of local slices and inverse CDF of the Barycenter.
C.5 LIMITATION AND FUTURE WORKS
Similar to the original INB, a limitation of FedINB is that it does not perform well for high dimensional image datasets. One of the reasons could be that FedINB is greedy algortihm. On one hand, this leads to benefits in computation and communication. On the other hand, this does not guarantee we find the global optimum the alignment problem. From a theoretical perspective, as pointed out in Zhou et al. (2022), the best can guarantee is that at each iteration, given the current O, the local alignment map f. optimal. Thus, in optimizing 9, especially considering the addition of VW histogram and reducing the maximum number of iterations J due to communication concerns, the
15

---

Published as a conference paper at ICLR 2023
empirical maximization of multi-mat matiiimax----- K-SW may not find the true maximum. But, empirically show in the experiment section, neither leads to drop in performance unless J or the maximum number of histogram bins V is really small. To further improve this fundamentally, we could investigate more on the optimization on the manifold of orthonormal matrices.
To overcome this, we provide practical improvements such using autoencoders to improve the performance. Also, FedINB could employ some specific tricks for images such as patch-based based hierarchical structure in SINF (Dai & Seljak| 202T) (though we do not leverage these image-specific tricks in this paper). Instead, we focus on a more general domain-agnostic method that does not assume any image-based structure in the model design. Since we focus on the more general case, Wasserstein Distance is a more reasonable metric as it is Furthermore, for the application of federated DG on high dimensional image datasets, we could share pretrained deep neural networks, which is known to be good at extracting meaningful features. Then we train the federated translation model in the intermediate representation. need this preliminiary transformation to be invertible since we don't need to get the exact translated images.
and not specific to images.
this case, we don't t
Another limitation of FedINB is that we need to transmit the translation model if we want to use it for other downstream tasks such as federated domain generalization. While this is more desirable than directly transmitting data in comparison models such as ELCFS (Liu et al. 202T). we propose possible method to mitigate this issue altogether. For federated domain generalization, we can utilize the shared space learned by FedINB for applying FedDIRT. Specifically, each client will minimize the following local objective:
(n.. g(x)) - dis(g()), (e
The key difference is that the local objective on the m-th client would only need its own translation model to the shared space Tm rather than needing Tm-mw', which would require the translation models from other clients. The intuition is that even though the shared space learned by FedINB is the actual Wasserstein Barycenter, may still keep the geometric structure of each domain and it may share the same invariant representation as the distribution of each domain. Hence, the second term may still serve as good regularizer for DIRT. This has the potential to further improve privacy preservation (because the client-specifie translation models would not be shared) and communication cost (because the translation models would not have to be transmitted). However, we leave a formal investigation of FedDIRT with the shared space of FedINB for future work.
Regarding privacy, we discussed about why sorting and VW histograms could improve privacy. More detailed discussion and limitation are in the next section.
D DISCUSSION ON PRIVACY
In this section, we will discuss how our FedINB with VW histograms is better than ELCFS (Liu et al., )02T) and original FedINB in terms of privacy. However, we want clarify upfront that we do not claim our model satisfies any strict privacy metric from the machine learning and federated learning areas such as differential privacy (DP) (Dwork, 2006). Fl Rather, we only seek demonstrate better privacy in comparison to prior and naïve approaches. Regarding DP, we provide discussion at the end of this section and will leave more formal analysis and careful investigation for future works.
In what follows, we will compare methods from the perspective of membership inference attacks. Positive membership privacy concerned about an adversary's 's ability to infer that an entity partici- pated the training as defined in (Li et al. 2013 Definition 3.1 3. As a running example for clarity, we will assume the hospital scenario where each client is hospital and each observation is a patient's data (e.g., cell images), but the following arguments could be made in other privacy scenarios as well. In our case, we make two assumptions: 1) The adversary has the original data of one patient (which may may not be one of the training examples one of the hospitals) and 2) the adversary can view the information sent to and from the central server. The membership inference question is: Given these assumptions, can the adversary determine whether the patient's 's data was included in the training data or not?
Cynthia, Dwork. "Differential privacy." Automata, languages and programming (2006): 1-12.
$ii, Ninghui, et al. "Membership privacy: unifying framework for privacy definitions." ACM SIGSAC Conference on Computer Communications Security. 2013.
16

---

Published as a conference paper at ICLR 2023
We first discuss membership attacks for ELCFS and then discuss them with respect to the privacy improvements of FedINB. ELCFS first does a fast Fourier transform of all data in each client to get the amplitude spectrum and phase spectrum. Then they transmit the low-frequency component of the amplitude spectrum to the server to form a distribution bank. In the training, each client acquires those low-frequency spectrums from other clients to create pseudodata from different domains. We believe this operation significantly violates the positive membership privacy. An adversary can compute the fast Fourier transform of the original data (assumption 1) and check whether its low-frequency component of amplitude spectrum is in the distribution bank, which is shared to each client in the training (assumption 2)
The sorting of slices of data before sending to the server in [ makes membership attacks more difficult than if the sorting was done on the server. Intuitively this is because sorting destroys joint sample information. First, consider the case where sorting is done on the server. We argue that an membership inference attack would be almost trivial in this case. The adversary could project the original data (assumption 1) using the current projection matrix (assumption 2) and match this projected data with a whole row in the projected training examples that are transmitted during training (assumption 2). If there is not whole-row match, then the patient is not in the training data. If there is a whole-row match, then the patient is more likely to be in the training data, where K determines the amount of precision-i.e., if K d and all data instances are unique, then this proves that the patient is in the training data, but if K = 1, there may be multiple patients that have the same projected value. On the other hand, if sorting is applied the clients before sending, then the adversary must attempt to determine a match by matching on each feature independently because the sorting indices are hidden. Thus, the adversary cannot match on whole rows anymore. The amount of actual privacy will depend on the non-uniqueness of individual values Furthermore, sending quantiles instead of real samples improves the privacy with respect to membership attacks because directly sending projected samples may hurt privacy. With VW histograms, most samples' exact information would be impossible to reconstruct (as histograms are lossy compressions of data samples) and trivial membership inference could be prevented. In some cases, such as the min and max values, membership attacks could still be possible but they are significantly more difficult for VW histograms than merely sending samples. Again, we do not claim DP guarantees but diseuss below how these could be added on top of our framework.
Finally, provide several possibilities ensure DP on top of the proposed FedINB. First, our method could be extended to leverage SP- (Abadi et al. 2016 for the optimization of projection matrix by merely clipping the gradient values and adding noise before locally updating the parameters. Regarding transmission of VW histograms for the barycenter calculations, our method could also be extended to compute differentially private histograms by adding noise the histogram estimates as is suggested in Example 3.2 of Dwork & Roth 2014) Finally. we could use standard DP techniques for composing multiple DP mechanisms (in this case DP-SGD and and DP histograms) into a joint DP mechanism that has DP guarantees (Dwork & Roth 2014). We leave a more careful investigation of DP methods in this context for future work.
E EXPERIMENTAL DETAILS
E.1 FEDERATED DOMAIN TRANSLATION
Metrics As discussed in the main paper, we choose WD and FID as our two met- rics, which gives a fair comparison against baselines. They are computed as WD I Em Emv W (mF J'm -m (Fm²) )) and FIDD = NIZ Em Em FID( I'm (m (fixed
M
'For common Deep Learning datasets such as high-dimensional images, the low-frequency component of amplitude spectrum is very likely to be unique given the high signal-to-now ratio of real images.
If the feature values are discrete, then there is likely to be many overlapped values and membership attacks will be difficult. If the feature values are continuous, then each value will be theoretically unique but adding some small I Laplace noise, as is standard in diffcrential privacy, would ensure that values are no longer unique.
$ Abadi, Martin, al. "Deep leaming with differential privacy." " ACM SIGSAC Conference on Computer and Communications Security. 2016.
Dwork, Cynthia, and Aaron Roth. "The algorithmic foundations of differential privacy." Found. Trends Theor. Comput. Sei. 9.3-4 (2014): 21 1-407.
17

---

Published as a conference paper at ICLR 2023
typo in the main paper). WD is computed with the Sinkhorn algorithm with maximum iterations set to 100 and c = 0.0001. At the starting point, f is the identity function.
Dataset For training, we use 10,000 samples from the MNIST and FashionMNIST training set as the dataset of domain O, where each class has 1000 samples. Then we use all samples to generate 10,000 samples for all other training domains (15.30,45,60) So, the total size of training data is 50,000. For evaluation of FedINB, we also use 10,000 samples from the MNIST and FashionMNIST test set, and create other samples in the same way. So the total size of test data is also 50,000. For evaluation of FedStarGAN, due to high computational cost, we only use 2,000 samples create test dataset so the total size of test data is 10,000. We find that the difference of evaluation score caused by different sizes of test dataset is negligible especially, for the WD score, and will not affect our conclusion in the paper.
FedINB with autoencoders In both AEINB and IndAEINB, we use the same autoencoder structure. For MNIST, the encoder is composed of [nn. Conv2d (1, 16, 3, padding-l), nn MaxPool2dd (2), nn. Conv2d 8, 3,
nn ReLU (inplace= padding=1) nn ReL (inplace-True) torch. nn in PyTorch. The decoder is composed of [nn.Conv2d 8, 3, padding=1), Rn. (implac====e) Cn. 16, 3, padding=1), nn. ReIU (tggg (aataaaf nn. ww..pmmmlliggiii..
nn Max where nn represents
nn. wwwp:/ppnnppppiiippeeeeemmm
2d (scale_factore)
nn. Conv2d (16, 1, 3,
padding-1), nn. Sin.iigmoii () where the last Sigmoid activation is used to map the output to the range of 1]. For FashionMNIST, the encoder is composed of [nn. Con.ddd stride-2, padding-1), nn. ReLU (True), nn. Convd (8, 16, 3, stride-2, padding=1), nn. mn.BatchNorm2d (16), nn. ReLU (True), nn. conv2d(16, 32, 3, stride=2, padding=0), nn RnLU (True) of Lnn.ConnTTrassssee22... (32, 16, 3, stride=2, battput_padding=0) nn nn.BaaccNNrm22d111) nn. ReL. nn (166 8, 3, stride=2, padding=1, output outpuut_padding-), nn (a2cchNorm22(8)) nn ReI htw.:/nnCnnneamommmmmmm Con.Transposee22 , nn. Siggigm
8, 3,
The decoder is composed
1, 3, stride-2, padding=1, As pointed out in the main paper, AEINB uses shared autoencoder that trained on data from all training domains. In practice, we should use some federated algorithm to train a shared autoencoder while each client has access to data from one domain. However, as a test, we show that even in this case, IndAEINB achieves competitive or even better result and in the federated case, the performance of AEINB cannot be better. Hence, we only report results of AEINB using this autoencoder.
For IndAEINB, each client has its autoencoder that is trained the same training dataset used to train INB.
We want to clarify that when reporting scores of AEINB and IndAEINB, the first number (at number of parameters sent equal to 0) is after applying autoencoders. This is why the starting point of them is different from that of other models such as INB.
Implementation of VW histograms In Algorithm when computing the objective, we use VW histograms to approximate the empirical Wasserstein Barycenter. At each client, we find the quantiles of sliced data (the number of quantiles is determined by V/) and send the quantiles to the server and average them. Then we have the VW histogram density of the empirical Wasserstein Barycenter using the averaged quantiles. Finally we create quantile of [0, 1] and use the inverseCDF of the histogram density estimator to create pseudodata of the empirical Wasserstein Barycenter.
Federated implementation of INB To emulate a federated version of INB, we need detach the gradient of the empirical Wasserstein Barycenter in Algorithm Other than that, there is not any practical difference whether we physically separate the data across different clients and simulate the transmission of data. In terms of communication cost, we compute the number of parameters being transmitted and report them.
FedStarGAN In order to implement FedStarGAN we start with the centralized implementation of StarGAN on RMNIST made available by the authors of DIRT. We next extend this to federated
18

---

Published as a conference paper at ICLR 2023
setting by defining 5 disjoint clients (1 for each domain of RMNIST) such that each client has its own copy of generator and discriminator. The 5 source domains are distributed among the 5 clients such that each client has data from only 1 of the 5 domains. The generator and discriminator undergo periodic FedAvg. We control the aggregation frequency by define a tunable parameter s step_st ze that determines the number of local update steps that client models must undergo before FedAvg aggregation and synchronization. For aggregation, we use the uniformly weighted FedAvg. Our experiments show that even 10 rounds of aggregation steps (using different values of sync_step) consumes more communication bandwidth than the HistInd/EIINB model but underperforms in terms of WD score (Fig. d)).
E.2 FEDERATED DOMAIN GENERALIZATION
Dataset For training, we use 1000 samples from MNIST and FashionMNIST training set as the dataset of domain 0 where each class has 1000 samples. Then we use all samples to generate 1000 samples for all other training domains (15.30,45660) So, the total size of training data is 5000. We decrease the size of the dataset to make this a harder domain generalization problem, and we retrain INB using less data. For evaluation, we use 1000 samples from MNIST and FashionMNIST training set to generate test data new domains 75 and 90.
Metries We use the accuracy of the trained classifier when employed to a new domain as the metric. Note that the INB and autoencoders are not needed after the training of classifier. As finding the most appropriate validation domain is another hyperparameter to tune, we choose to use the simplest way to report accuracy we report the test accuracy at a fixed point. For faimess, we report the accuracy of FedDIRT and FedAvg at the same point and check that both have already converged before recording (this could be validated by Figure [4]. For Rotated MNIST and all test domains, we run FedDIRT//edAAg (1 (1-batch) for 2000 iterations, FedDIRT/FedAvg (10-batch) for 2500 iterations and FedDIRTT/FeAAvg (100-batch) for 3000 iterations. For Rotated Fashion and all synchro- nization steps,
FedDIRT/FedAvg (Domain 75) for 1500 iterations and FedDIRT/FedAvg (Domain 90) for 3000 iterations.
FedDIRT For the network structure (encoder and classifier) and training hyperparameters, we modify based on the default setup in the repository of DIRT (Nguyen al. [021), which can be found at their public repository
The only difference is that we change Batch Normalization to Instance Normalization. For FedINB, we choose a few IndAEINB with VW histograms (i.e., HistIndAEINB) since they are optimal in terms of overall translation performance and communication cost. Note this case, autocoders and INB should be considered together as single translation model f. At each batch in the training, we randomly assign target domain for each sample which different from the default setup of DIRT (they assign the same target domain for whole batch in the training). For MNIST, we use 64 as batch size and 0.001 as learning rate. We set the regularization weight of DIRT 10 batch size and 0.0001 learning rate. We set the regularization weight of DIRT to be 10.
to be 2. For MNSHi we use 128 as
FedAvg The FedAvg model structure is the same as FedDIRT. We use the uniformly weighted averaging of model updates similar to the original implementation McMahan et al.] (017). We consider three different configurations of FedAvg as baselines: FedAvg with saync_stesss__ FedAvg with sync_sttttteeee cenee_sses=l0,..__11 and FedAvg with sync_st iinc_sttps=100. Learning rate and batch size are the same as those used for corresponding FedDIRT.
F ADDITIONAL RESULTS
F.I FID SCORE OF RESULTS IN SECTION
In Figure 4 we include results of FID. One observation is that for FID, the influence of the blurriness caused by applying autoencoders is more significant than the domain shifts. This is unsurprising because FID is based on a modern CNN and CNNs are known to focus texture fine- fine-grained
The weight before domain invariant feature term.
19

---

Published as a conference paper at ICLR 2023
Autoencoder: L10 110-K00--200
IndAEINB: L= 10, K=10
106
100
J-200
J-100
Average FID
Average FID
104
J=50
80
INB
J-30
AEINB
102
J-10
IndAEINB
100
40
0.5
1.0
15
20
25
3.0
100
2.5
Communication Cost (# params sent)
leg
Communication Cost (# params sent)
1eg
(a) Different autoencoders
(b) Number of iterations
VW hist(V=500): K10-K10-1100
FedStarGAN
Average Waserstein Distance
500
100
400
HistIndAEINE
INB
HistINB
IndAEINB HistINB
Average
FedStarGAN.4
60
HistindAEING
200
FedStarGAN-16
FedStarGAN-129
100
0.00
0.5
0.5
0.75
1.00
1:25
150
175
0.0
05
10
1.5
2.0
Communication Cost (# params sent)
100
Communication Cost (# params sent)
1e0
(c) VW histograms
(d) Convergence of FedStarGAN
Figure 4: FID (lower is better) communication cost increases for Rotated MNIST. (a) Different autoencoders used for INB. (b) Different J used for IndAEINB. (c) Using VW histograms for INB. (d) Failure of FedStarGAN. The INB used here is HistndAAEINB-LLO--IOO--- BBsINB-LIO-K1OJJ10 with V = 500.
with 500 and
details instead of high-level semanties (Geirhos et 2019 Because proposing more general approach to federated translation that not specific to image data, we suggest that the domain-agnostic metric of WD is better metrie for our context rather than the image-specific FID metric. Besides, as shown in Figure [2] AEINB and IndAEINB generate reasonable rotations of the original samples (albeit somewhat blurry because of the AE), and keeping the semanties of the domain translation (i.e., correct rotation) is more important than preserving the fine details for downstream tasks like DG.
F.2 DIFFERENT SETUPS OF INB WITH ROTATED MNIST
In this section, we include results with more different setups of INB.
As shown in Figure with different setups of INB (different L, K, J, whether to use VW histograms and different r), we observe that terms of WD, IndAEINB better than using shared autoconder not using autoencoders.
Autoencoder.
we include more results with different INB setups to study the influence of decreasing J. We can conclude that using higher J does not lead to better performance after convergence.
Max iterations Figure
Number of slices K. As shown in Figure 7 communication cost becomes times higher but the overall performance improved. Hence, use lower K in following experiments to reduce communication cost.
we can observe that as we increase K to 30, the
"eiiroo,, Robert, et al. "ImageNet-trainee" CNNs are biased towards texture; increasing shape hias improves accuracy and robustness." International Conference Learning Representations. 2018.
20

---

Published as a conference paper at ICLR 2023
Autoencoder: L10-K30-J200
Autoencoder: L10-K30 110-K30-J200
Waserstein Distance
100
AEINB
32
IndAEINB
Average FID
INB
30
AEINB
60
IndAEINB
28
26
Communication Cost
Communication Cost (# params sent)
(# params sent)
(a) WD: LI0-K30-J200
(b) FID: 111---30-2220
Autoencoder: L10-K10-J100
Autoencoder: L10- 11-K100-1100
Average Waserstein Distance
34
INB
100
AEINB
IndAEINB
Average FID
INB
30
IndAEINB
28
60
0.00
0.25
0.50
0.75
1.00
1225
1.50
1.75
0.00
0.25
0.00
0.75
1.00
1.25
150
175
Communication Cost (# params sent)
169
Communication Cost (# params sent)
169
(c) WD: 110-0100
(d) FID: L10-L10----00 K10-J100
Autoencoder: 10--100-1100 (V=500)
Autoencoder: L10-K10 -10-K1000 (V=500)
Average Waserstein Distance
34
HistiNB
HistAEINB
100
32
HistIndAEINB
Average FID
80
HistiNB
30
HIstAEINB
28
10
1.5
2.0
2.5
0.0
0.
1.0
1.5
2.0
25
3.0
Communication Cost
10B
Communication Cost params sent)
lca
params
WD: LI0-K10-J1000-550 L10-K10-J100 (V=2000)
(f) 110-KK0-J100--55
Autoencoder:
Autoencoder: L10-K10-J100 (V=2000)
Average Waserstein Distance
34
HistiNB
100
HIstAEINB
32
HistIndAEINS
Average FID
HistiBB
30
HistAEINB
HistindAEINB
28
26
2
8
Communication Cost
1eB
Communication Cost
168
params sent)
(# params
(g) WD: LI0-KKIOOOJJ0000
(h) FID: L10-K100-1000-V2000
Figure 5: Influence of autoencoders with different setups of INB for Rotated MNIST.
VW histograms. As shown in Figure g as we decrease V. both HistIBB and HistInAAEINB can achieve good performance (note the actual number of samples each client is 10,000).
21

---

Published as a conference paper at ICLR 2023
INB: L=10, K=10
INB: L=10, K=10
Average Waserstein Distance
35
J-200
J-200
34
J-100
J-100
J-50
Average FID
-550
33
-=30
J=30
I=10
J=10
32
39
31
38
30
37
Communication Cost (# params sent]
Communication Cost (# params sent)
(a) WD: INB
(b) FID: INB
INB: L=10, K=30
INB: L=10, K=30
Average Waserstein Distance
33
J-200
42
J-200
34
1-100
1=100
5=50
Average FID
5=50
3=30
3=30
32
1=10
1=
38
31
30
29
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
044
0
0.8
1.0
Communication Cost
1610
Communication Cost (# params sent)
1e10
(# params sent)
(c) WD: INB
(d) FID: INB
IndAEINB: L=10, K=10
IndAEINB: L=10, K=10
102
Average Waserstein Distance
28.00
J-200
J-200
27.75
--000
101
J-100
27.50
-50
-30
Average FID
J-30
J=50
27.25
1-10
100
J-10
99
26.75
98
0.0
0.5
1.0
1.5
2.0
2.5
1.5 Communication Cost (# params sent)
Communication Cost
109
(# params sent)
(e) WD: IndAEINB
FID: IndAEINB
IndAEINB: L=10, K=30
IndAEINB: L=10, K=30
102
Average Waserstein Distance
28.00
J-200
J-200
27.75
J-100
J-100
27.50
5=50
FID
100
J=50
J-30
27.25
J-10
J-10
27.00
96
26.50
0.6 Communication Cost (# params sent)
00
0.5
0.8
1.
Communication Cost (# params sent)
1e10
(g) WD: IndAEINB
(h) FID: IndAEINB
Figure 6: Inffuence of J for Rotated MNIST.
F.3 DIFFERENT SETUPS OF INB WITH ROTATED FASHIONMNIST
In this section, we include quantitative federated domain translation results with Rotated FashionM- NIST, The experiment setup here not exactly the same that with Rotated MNIST. We skip some
22

---

Published as a conference paper at ICLR 2023
INB: L=10, J=200
Average Waserstein Distance
INB: L=10, J2200
K-10
47.5
K-30
K-30
45.0
Average FID
40
42.5
35
40.0
37.5
30
0.6
Communication Cost (# params sent)
Communication Cost (# params sent)
1e20
(a) WD: INB
(b) FID: INB
Average Waserstein Distance
IndAEINB: J=200
IndAEINB: L=10, I=
J=200
K=10
K=10
K=30
104
K-30
Average FID
102
35
100
0.0
0.2
0.4
0.6
0.8
100 1e10
0.0
0.2
0.4
0.6
0.8
10
Communication Cost
Communication Cost (# params sent)
1+10
(# params sent)
(c) WD: IndAEINB
(d) FID: IndAEINB
Figure Influence of for Rotated MNIST.
HistINB: L10-K10-J100
HistINB: L10-K10-J100
Average Waserstein Distance
45
V-500 V-2000
V-500
V-2000
Average FID
40
35
40
38
30
Communication Cost (# params sent)
Communication Cost (# params sent)
(a) WD: HistINB
(b) FID: HistINB
HistIndAEINB: L10-K10-J100
HistIndAEINB: L10-KIN K10-J100
Average Waserstein Distance
40
V-500
106
V=500
V=2000
V-2000
Average FID
104
35
102
30
100
Communication Cost (# params sent]
le8
Communication Cost params sent)
160
(c) WD: HistIndAEINB
(d) FID: HistIndAEINB
Figure 8: Influence of V for Rotated MNIST.
23

---

Published as a conference paper at ICLR 2023
Waserstein Distance
34
HistindAEINB: L=20, K=10, V=500
J-200
Average Waserstein Distance
HistindAEINB: L=20, K=10, V=500
J-200
32
J=100
150
]=100
J-50
145
J=50
J=30
J-30
140
28
135
26
130
125
Communication Cost
Communication Cost (# params sent)
(# params sent)
(a) WD: L20-K10-V500
(b) FID: 120-K10-V500
Figure 9: Influence of for Rotated PashionMNIS'.
Average Wazerstein Distance
34
HistindAEINB: L= 10, 1=10,, V=500
K-10
Average Waserstein Distanc
140
HistIndAEINB: L=10, J=100, V=500
32
K-20
139
K=30
28
137
K=10
26
136
K-20
K-30
135
Communication Cost (# params sent)
Communication Cost
(# params sent)
(a) WD: L10-J100-V55
(b) FID: L10-J100-V500
Figure 10: Influence of K for Rotated MNSHTOnMNIITT.
preliminary tests such as investigation of autoencoders and we test with HistIndAEINB instead of IndAEINB for J and K.
Max iterations J. As shown in Figure 9 observe that HistIndAEINB with J 200 and 100 achieve similar performance in terms of WD while J 50 and 30 seem to perform worse. On one hand, we send histograms instead of all samples hand, it is inherently harder task to find the best projection for Rotated FashionMNIST such that more number of iterations is as expected.
that in Figure with Rotated MNIST). On other
Number slices K. As shown in Figure 10 as we increase K to 30, the communication cost becomes 3 times higher but the overall performance is not improved in terms of WD.
Number of histogram bins V. As shown in Figure 11 10000 and FashionMNIST is much harder task, the performance of HistIndAEINB does not drop until V as small as 50.
even if the total number of samples is
F.4 MORE QUALITATIVE TRANSLATION RESULTS
In Figure 12 we include more qualitative results with MNIST. We observe that in comparison to INB, IndAEINB and HistIndAEINB generate smooth samples. Moreover, samples translated by IndAEINB and HistIndAEINB are very similar to each other. Figure [ we include qualitative results with MNSS..
24

---

Published as a conference paper at ICLR 2023
HistindAEINB: L=10, K=10, J=100
]=100
Average Waserstein Distance
HistindAEINB: L=10, K=10, J=100
34
V-2000
Average Waserstein Distance
150.0
32
V-500
V-100
147.5
V-2000
V=500
V=50
145.0
V=100
V-50
28
142.5
26
137.5
135.0
Communication Cost (# params sent)
Communication Cost (# params sent)
168
(a) WD: LI0-K10-J100
(b) FID: LI0-K10-J100
Average Waserstein Distance
HistindAEINB: L=20, K=20.J=100
Average Waserstein Distance
150
HistindAEINB: L=20, 2=20.
34
V=2000
V=2000
32
V=100
V=500
145
V-500
V=100
30
V=50
V-50
140
26
135
24
130
0.0
0.5
100
3.5
2.0
2.5
0.0
0.5
1.0
1.5
2.0
25
Communication Cost (# params sent)
169
Communication Cost (# params sent)
149
(c) WD: 220-020-J100
(d) FID: 1220-K2-J100
Figure 11: Influence of V for Rotated FaShionMNIST.
F.5 MORE RESULTS OF FEDERATED DOMAIN GENERALIZATION
In Table 3 and Table we include results with more INB setups. We can observe that HiSIInAAEINB achieves similar performance, and they are all better than FedAvg. In Figure [4 we plot the test accuracy of FedDIRT and FedAvg to demonstrate that the training already converges when we record the test accuracy.
F.6 INVESTIGATION OF THE OPTIMIZATION OF FEDINB
In this section, further compare the optimization of FedINB and original INB, we investigate the optimization of the multi-max K-SW subproblem throughout the full optimization process with respect to the AE and VW histograms improvement. Note that the convergence of the whole algortihm has been investigated in Figure 2a Figure 2c] Figure [] and Figure 8 which indicates that autoencoders lead a better alignment and VW histograms do not hurt the performance unless V really small. So focus their impact multi-max. max---mxx------------- here.
As reminder, the objective of multi max-max----------- is MKn E EK In [5] we track the change of maximum K-sliced Wasserstein Divergence through the training of a 10-layer (L 10) FedINB for different digits where constrain the maximum number of iterations to be 100 (J = 100). We can observe that using VW histograms barely leads to any difference in the optimization: the point they converge and the speed they converge at are very similar. For IndAEINB, we cannot directly compare the numbers as they are essentially computed difference spaces. But we can observe that IndAEINB actually converges faster especially in the last layers, and the max-K max-K-SWW is much closer to O in the last few layers. This could result from that we are running multi-max- K-SW in space in much lower dimension and the task becomes easier. Besides, we observe that in most cases, multi-max-K- number of iterations even if there is barely any improvement. This also explains why we could reduce J without hurting the performance (as explored in Section 4.I.. In Table 5 also track the total number of iterations with different setups. We can observe that INB and HistINB requires similar iterations while AEINB takes relatively less iterations.
mm) (0T V().
keeps running till it reaches the maximum
25

---

Published as a conference paper at ICLR 2023
Table Classification accuracy at unseen domains for Rotated MNIST. For each target domain, we try to aggregate model after each 1, 10, 100 mini-batches. The mean and standard deviation are taken over runs. For FedDIRT, we use HistIndAEINB with 100 and V 500 (with different L and
K)
Model
(1-btch))
(10-batth)
75 (100-batch)
90 (10-batth)) (1-batc))
(100-batch)
(cDDIRT(L10-K20)
92.4 + 0.3
91.7 + 1.2
90.2 + 1.0
70.1 0.3
68.5 + 3.8
65.7 2.9
FedDIRT(I.22 (0))
92.2 + 0.3
89.3 4.2
91.9 + 0.9
69.8 0.7
70.4 2.2
69.9 0.3
FedDIRTYLIO-II
92.4 0.4
92.2 + 0.7
90.9 + 0.8
70.5 0.5
70.8 0.8
65.9
FedDIRT((20--200)
92.2 = 1.3
90.8
1.2
91.4 0.6
+..4
71.1 1.3
69.8 1.5
FodAvg
85.2
0.5
80.1 2.3
63.8
63.6 0.8
35.6
Table 4: Classification accuracy at unseen domains for Rotated FashionMNIST. For each target domain, we try to aggregate model after each I, 10, 100 mini-batches. The mean and standard deviation are taken over runs. For FedDIRT, we use HistIndAEINB with J = 100 and V = 500 (with different L and K)
Model
75 (1-batch
(10-butc)
75 (100-batch)
90 (1-batch))
90 (10-batch)
90 (100-butch)
FeaDIRTRLOO 10-KK0)
65.8 1.3
6.91
62.9 1
0.5
0.6
19.0 0.4
FedDIRTYL20-K4 (10)
65.8 = 1.3
63.6 0.9
63.0 + 1.0
1.9
18.1 15.2044 0.4
18.8 + 0.6
PEDDIRTLLO-KIIN
65.1
1 1.8
63,5 1 0.6
62.6 0.6
1.5
18.3 0.3
18.8 1 0.6
FodDIRT(L20-330
65.5 + 1.8
64.5 + 1.5
63.5 + 1.3
19.0 0.3
18.5 + 1.0
18.9 + 0.4
FodAvg
51.9
50.4
40.4 +
13.7
14.6 1.8
13.1 0.9
Table 5: Number of iterations of multi-max-K-SVV with L 10, K
10, 100 for RotatedM-
NIST. The numbers are averaged over 4 runs. Note that the extra steps (I. = 10 x 100 = 1000)
come from the backtracking line search explained in Appendix C.3
INB
HistINB
IndAEINB
Digit 0
1,605.25
1,605.00
1,586.50
Digit 1
1,609.00
1,612.75
1,580.50
Digit 2
1,601.75
1,603.50
1,584.75
Digit 3
1,599.50
1,600.00
1,583.25
Digit 4
1,597.75
1,594.50
1,381.25
Digit 5
1,605.00
1,605.25
1,587.75
Digit 6
1,601.00
1,601.00
1,582.75
Digit 7
1,600.75
1,601.25
1,581.25
Digit 8
1,601.50
1,600.75
1,583.25
Digit 9
1,598.75
1,599.75
1,580.00
Average
1,602.03
1,602.38
1,583.13
In Figure 16 we increase the maximum number of iterations to be 300 and observe similar results. In Table we observe that total number of iterations of HistINB is slightly larger than INB. but in practice, we don't t actually need this many iterations.
we also observe that in the last few layers, the maximum K-sliced Wasserstein Diver- gence found by the algorithm is very small which is consistent with the observation of convergence of FedINB Figure 2c
From Figure 15|
26

---

Published as a conference paper at ICLR 2023
Table 6: Number of iterations of multi-max-K-SW with L = 10, K = 10, J NIST. The numbers are averaged over 4 runs.
300 for RotatedM-
INB
HistINB
IndAEINB
Digit 0
4,767.67
4,800.67
4,756.50
Digit 1
4,780.00
4,782.67
4,756.50
Digit 2
4,778.00
4,777.67
4,756.00
Digit 3
4,779.00
4,779.33
4,753.50
Digit 4
4,778.67
4,781.67
4,754.50
Digit 5
4,778.67
4,784.33
4,758.50
Digit 6
4.743.33
4.788.00
4,753.50
Digit 7
4,776.67
4,786.00
4,754.50
Digit 8
4,779.00
4,778.67
4,754.50
Digit 9
4,781.67
4,789.00
4,754.00
Average
4.774.27
4,784.80
4,755.20
27

---

Published as a conference paper at ICLR 2023
4
0 1
0
7
8
6
8 8650571
1
2
6
4
6
066289233
3
(a) INB: 10= K=10, I=200
0
7
7
8
6
8
1655771
2
6
6
3
2
8
2
0
(b) IndAEINB: L=10,, K=10, J=200
4
0
0
1
7
8
6
8
6
2
6
&
8
9
2
2
3
0
(c) HistnndAEINB: L=10, K=10, J=200, V=500
Figure 12: More qualitative results for Rotated MNIST. The first three rows are original samples from domain O. The last three rows are translated samples (from domain 0 to 4).
28

---

Published as a conference paper at ICLR 2023
(a) HistIndAEINB: L=10, K=10, J=100, V=500
Figure 13: Qualitative results for Rotated MashionMNIST. The first three rows are original samples from domain O. The three rows are translated samples (from domain to 4).
Domain 75 (1-batch)
Domain (1-batch)
0.6
Test Accuracy
0.8
0.5
Test
0.4
0.7
Fed2IRT(L20-
FedDIRT[L20-K20)
FedAug
FedAvg
0.3
Rer serations
Iterations
(a) Rotated MNIST
(b) Rotated FashionMNIST
Figure 14: Test accuracy in the federated domain generalization experiments.
29

---

Published as a conference paper at ICLR 2023
Convergence of max-K-SW (label=0)
Convergence of max-k-SW (label=0)
HistiNB
200
IndAEINB
max: KWW objective
INB
max max-k-sw K-SW Objective
INB
150
150
100
100
50
50
01
iterations
Iterations
INB and HistlNB (V 500) for digit
(h) INB and IndAEINB for digit
Convergence of max-K-SW SW (label=1)
Convergence
max-K-SW (label=1)
HistINB
IndAEINB
max-k K.5 Objective
INB
max-k-SW Objective
INB
200
200
100
100
0
200
400
600
800
3000
:200
1400
1600
200 0
400
600
800
2000
1200 1400 2600
Iterations
Itteratoons
(c) INB HistINB (V 500) for digit
(d) INB and IndAEINB for digit 1.
Convergence of max-K- SW (label=2)
Convergence of max-K-SW (label=2)
250
HistINB
IndAEINB
max SW objective
SW Objective
200
INB
200
INB
150
SW
100
100
50
50
400
600
1900
1200
1C00
1661
6.0 300
1200 1000 1800
Iterations
Iterations
(e) INB and HistINB (V 500) for digit 2.
(f) INB and IndAEINB for digit 2.
Figure 15: Convergence multi-max- K-SW with 10, K = 10, J 100 for RotatedMNIST.
The curves describe the change of loss through 10 layers of FedINB with different setups. The sudden change of loss because we apply ID-Ban D--Barycenter there to align the projected distribution. We want to note that the value of objective for INB and IndAEINB is not directly comparable because for IndA EINB we are applying FedINB in a space in much smaller dimension.
30

---

Published as a conference paper at ICLR 2023
Convergence of max-k-SW (label=0)
Convergence of max-K-SW (label-0)
max 5W Objective
200
INB
HistINB
max-k-WW Objective
200
IndAEINB
150
150
SW
100
50
0
1000
2000
3000
4200
1000
2000
5000
4000
5000
Iterations
Iterations
(a) INB and HistNNB (V = 500) for digit O.
(b) INB and IndAEINB for digit O.
Convergence of max-K-SW (label=1)
Convergence of max-k-SW (label=1)
HistlNB
IndAEINB
max-K-SW ob Dbfeciiee
300
mox-K-SW Objective
300
INB
200
200
100
100
0-
1000
2000
2000
4000
5000
0
1000
2000
3000
4000
5000
Iterations
Iterations
(c) INB and HistINB (V = 500) for digit
(d) INB and IndAEINB for digit 1.
Convergence of max-K-SW (label=2)
Convergence of max-K-SW (label=2)
HistiNB
IndAEINB
max-K-5W Objective
max-K-5W Objective
200
INB
200
INB
100
100
50
0
1000
2008
3030
4000
5000
0
1000
2000
3000
4000
5000
Iterations
Iterations
(e) INB HistINB (V 500) for digit
(f) INB and IndAEINB for digit
Figure 16: Convergence of Pilti-max-K-SWW with L 10, K 10, J 300 for RotatedMNIST.
31


Reviews:
[{"review_id": "CPc1AujFBmt", "reviewer": "Reviewer_zb7q", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper starts with pointing out that the prevalent perspective regarding modeling heterogeneous FL clients by looking at the label distribution is limiting, since, in practice, the data distribution may differ between clients, too. The authors propose to mitigate this concern using a variant of an iterative domain translation model that they adapted to the FL scenario. They evaluate on a one-domain-per-client setting using Rotated MNIST and FashionMNIST.", "summary_of_the_review": "To a reviewer without a background in federated learning, this looks like a strong submission.", "strength_and_weaknesses": "I am no expert in federated learning, so I cannot comment on how novel the contributions are and how relevant this work is to practical federated learning. But on a high level, the paper seems to be well motivated. I could easily follow the overall story. I did not check 100% of the formalisms, but I could not find any vagueness/clarity issues in the ones I did check. The claims regarding communication efficiency are clearly demonstrated in Fig. 2.\n\nTables 1 and 2 are also clearly demonstrating the benefits regarding domain generalization, but I'm wondering why they do not include a baseline from Liu et al. (2021b) (see Sec. 3). Also, if the advantage over Liu et al. (2021b) is that this work is amenable to other data modalities such as text, it would be nice to evaluate on a text-based task.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "4xS5JZc1-BR", "reviewer": "Reviewer_CNpn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on a challenging setting in federated learning (FL) where class conditional distribution varies amongst clients. To address this issue, the authors develop a federated domain translation methodology based on an advanced iterative translation model, namely Iterative Naive Barycenter(INB).\nINB is more viable to the FT setting than standard translation methods, e.g. StarGAN. Furthermore, the authors propose several FL-motivated improvements to INB, which significantly reduce communication costs. Empirical results show that the proposed federated translation model not only performs better than FL versions of baseline translation model, but also improves robustness", "summary_of_the_review": "This is a solid paper which applies INB in a Federated Learning setting with practical improvement to reduce the communication cost.", "strength_and_weaknesses": "Strength:\n\n1. Novelty: The authors propose to apply INB in FL setting for the problem where different clients have their own class distribution. In addition, the authors propose several FL-motivated improvements to INB, including the use of variable-bin-width histograms and autoencoder.\n2. Performance: Empirical results on Rotated MNIST and FashionMNIST show that the proposed method outperforms FL versions of standard translation models under practical settings of limited client-server communication capability.\n\nWeakness and Suggestions:\n\nWhile it does propose practical improvement, the work seems incremental in applying INB in federated learning.\n\nThe key point in INB is to learn the transformation which maps each domain distribution to a shared latent distribution. While it could be reasonable or practical for some domains, there could be cases where the mapping is difficult to learn. This may need to be discussed as potential weaknesses.\n\nI am wondering how much the difference is between a INB and a FedINB, where there be a balance/trade-off between communication and the performance?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-wVV-dz8lA", "reviewer": "Reviewer_uckp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a federated domain translation approach that can mitigate conditional shift issues in federated learning tasks. The proposed translation model is empirically shown that it performs better than the state of the art FedStarGAN. The paper also shows that by combining FedINB with DIRT, regularization from domain translation model significantly improves the model’s ability to generalize to unseen domains. I am not particularly familiar with the topic to have indepth comments. However I think the work is solid and should be accepted to ICLR.", "summary_of_the_review": "An well-studied federated domain translation approach that can mitigate conditional shift issues in federated learning tasks.", "strength_and_weaknesses": "* Strength:\n\nThe proposed framework sounds to me. I also found the experiments solid, showing the proposed model is more efficient in terms of both communication and computational cost than FedStarGAN. Experiments also show that practical improvements proposed in the paper seem helpful to reduce communication costs.\n\n* Con: \n\nI found some parts of the writing are not clear. However, I also must say that I am not the particularly familiar with the topic, so it also contributes to some difficulty in following the paper to me (see Clarity section).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster