Federated Domain Generalization with Data-free On-server Matching Gradient
Authors: Trong Binh Nguyen
Duong Minh Nguyen
Jinsun Park
Viet Quoc Pham
Won-Joo Hwang
Venue: ICLR 2025 Poster (2025)
Keywords: Federated Learning
Domain Generalization
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: The paper propose a on-server approach to aggregate clients' informations and statistics without the data transmissions from clients to solve the federated domain generalization

Abstract: Domain Generalization (DG) aims to learn from multiple known source domains a model that can generalize well to unknown target domains. One of the key approaches in DG is training an encoder which generates domain-invariant representations. However, this approach is not applicable in Federated Domain Generalization (FDG), where data from various domains are distributed across different clients. In this paper, we introduce a novel approach, dubbed Federated Learning via On-server Matching Gradient (FedOMG), which can efficiently leverage domain information from distributed domains. Specifically, we utilize the local gradients as information about the distributed models to find an invariant gradient direction across all domains through gradient inner product maximization. The advantages are two-fold: 1) FedOMG can aggregate the characteristics of distributed models on the centralized server without incurring any additional communication cost, and 2) FedOMG is orthogonal to many existing FL/FDG methods, allowing for additional performance improvements by being seamlessly integrated with them. Extensive experimental evaluations on various settings demonstrate the robustness of FedOMG compared to other FL/FDG baselines. Our method outperforms recent SOTA baselines on four FL benchmark datasets (MNIST, EMNIST, CIFAR-10, and CIFAR-100), and three FDG benchmark datasets (PACS, VLCS, and OfficeHome). The reproducible code is publicly available~\footnote[1]{\url{https://github.com/skydvn/fedomg}}.

Published as a conference paper at ICLR 2025
FEDERATED DOMAIN GENERALIZATION WITH DATA-FRED ON-SERVER MATCHING GRADIENT
Trong-Binh Nguyen'', Minh-Duong "guuen'",,Jnnee Park't, Quoe- Viet Pham', Won Joo Hwang"! 1 Pusan National University, Republic of Korea, 2 Trinity College Dublin, Ireland
* Equal contribution, I Corresponding author
( Doinngguyentrong, duongnm, jspark, hwangwi) @pusan ac kr viet. ie
ABSTRACT
Domain Generalization (DG) aims to learn from multiple known source domains model that can generalize well to unknown target domains. One of the key approaches in DG is training an encoder which generates Nomain-innvriian repre- sentations. However, this approach is not applicable in Federated Domain General- ization (FDG), where data from various domains are distributed across different clients. In this paper, we introduce novel approach, dubbed Federated Learning via On-server Matching Gradient (FedOMG), which can efficiently leverage domain information from distributed domains. Specifically, we utilize the local gradients as information about the distributed models to find an invariant gradient direction across all domains through gradient inner product maximization. The advantages are two-fold: 1) FedOMG can aggregate the characteristics of distributed models on the centralized server without incurring any additional communication cost, and 2) FedOMG is orthogonal to many existing FL/FDG methods, allowing for additional performance improvements by being seamlessly integrated with them. Extensive experimental evaluations various settings demonstrate the robustness of FedOMG compared to other FL/FDG baselines. Our method outperforms recent SOTA baselines four FL benchmark datasets (MNIST, EMNIST, CIFAR-IO, and CIFAR- and three FDG benchmark datasets (PACS, VLCS, and OfficeHome). The reproducible code is publicly available
INTRODUCTION
Federated Learning (FL) has gained widespread recognition due to its ability to train models collabo- ratively across multiple clients while keeping their individual data secure. However, one practical challenge, how to ensure that models trained on sites with heterogeneous distributions generalize to target clients with unknown distributions, known as Federated Domain Generalization (FDG), remains under-explored. While label distribution shift has been considered in traditional FL, FDG focuses on the feature shift among clients and considers each client as an individual domain.
FDG shares a similar goal as standard Domain Generalization (DG) (Nguyen et al. 2022b
al. 2018), i.e., generalizing from multi-source domains to unseen domains. However, unlike DG. where knowledge from various domains can be jointly utilized to develop an efficient algorithm, FDG prohibits direct data sharing among clients, which makes most existing DG methods hardly applicable. To address the challenges inherent to FDG, recent methods circumvent these difficulties by adopting alternative strategies, e.g., unbiased local training within each isolated domain, local data
preprocessing (Huang al.) 2022 2024), knowledge sharing among users (Liu et al. 2021 Chen
ct al. 2023). As a consequence, an open question remains for FDG:
How can knowledge across domains be effectively leveraged design FDG algorithms that achieve performance comparable to DG, while avoiding additional communication overhead?
https:///githubb..om///y/

---

Published as a conference paper at ICLR 2025
Addressing the aforementioned question, we inherit the gradient matching rationale (Shi et al. 2022 Rame et al. 2022). The rationale is straightforward: the model is domain-invariane invariant, the gradients induced by e across different domains should exhibit high correlation with each other. To this end, we propose a novel on-server invariant gradient aggregation approach for FDG, dubbed Federated learning via On-server Matching Gradient (FedOMG). Specifically, leverage the available local gradients at the server to design a gradient-based approach for finding domain-invariant models. To design an on-server optimization strategy, we draw inspiration from the Gradient Inner Product (GIP) optimization problem introduced in Fish (Shi et al. 2022) to implement gradient matching, However, the current GIP approach has two significant limitations for FDG: (1) directly minimizing GIP, as described in (Shi et al., 2022 Alg. 2), incurs substantial computational overhead due to the need to compute second-order derivatives of model parameters related to the GIP term; and (2) the surrogate approach for Fish, as outlined in (Shi et al. 2022] Alg. 1), requires continuous transmission of client models, leading to excessive communication overhead. To address these issues, we propose an indirect optimization method that leverages surrogate optimization variables instead of model parameters, thereby eliminating the need for second-order derivatives in GIP. Furthermore, to enhance computational efficiency, we introduce efficient convex optimization formulation streamline the on-server optimization process. The advantages of our design are as follows:
1. FedOMG can leverage local client knowledge to train global model, similar to conventional DG. Consequently, our global model has better knowledge of all clients, thus, leading to high generalization when applied to Out-Of-Distributin (OOD) data.
2. Due to the design of on-server gradient matching, our method is orthogonal to many classic FDG methods, where current FDG methods focus on the local training (Huang al. 2022 2024). As consequence, the integration of FedOMG with other FDG methods can achieve improved performance.
3. Due the utilization of available local gradients server, our method does not require data transmission from distributed devices to the server, or sharing data among users (Liu et al.) [021] Chen et 2023). Consequently, can maintain the efficiency of communication and privacy in FL settings.
Contribution. Our contributions can be summarized follows. 1) We propose FedOMG, communication-cffee method for learning local gradients to facilitate the on-server gradient matching. 2) We conduct theoretical analysis to reveal that FedOMG is robust to the reduction of the generalization gap of the unseen target domains. 3) We conduct experimental evaluations under FL settings and FDG settings. In FL settings, consider four FL datasets (MNIST, EMNIST, CIFAR-10, and CIFAR-100), In FDG settings, we consider the DG benchmark datasets (PACS, VLCS, and OfficeHomee) Experimental results show that FedOMG significantly outperforms existing FL methods in terms of IID, non-IID dataset, and OOD generalization.
models in FDG, which leverages the
PRELIMINARIES
Notations: By default, (:-)) and || denote the Euclidean inner product and Euclidean norm, respectively. In our research, we use d.)(s-) to represent any distance metric M between two distributions. We denote O as Hadamard product. We denote IP(O; D) the output distribution of model given input as dataset DD We denote 05"), offr) d(f)) E RRXM as the global, local, trained "f'),(())
local model parameters used in the FL system round r, respectively. Let g(o("=)) ;D,) E RIxM be local gradient of client u at epoch of round r, g(0(r); (D.) E RRIM be gradient of client u at round (ne). g(((r) = =
r. Without loss of generality, we abuse g(f(f)), ;,.) = gu
Problem Setup: We consider an FL. system comprising set of source clients Us {ulu 1, ........ .U))............. and target clients U4 {ulu = Us + Us 2,...,Us + UT). Each client gains access to its data D. = {x,, y J with N₂ data samples, where u E (Us: Source and target clients are assigned to the source datasets Ds (Dulu E Us} and target datasets Dy = (Dulu e 3, respectively. Each client is trained for E epochs every round. The system objectives are two-fold.
2

---

Published as a conference paper at ICLR 2025
Objective 1 (guarantee on the conventional FL setting) The method has perform well the source clients set Ug, i.e., 0 = arg ming Es = Eueus (wE Du), where (0, Du) is the local empirical risk, Tuu represents the weight on client u, satisfying Eueti 1
Objective 2 (guarantee on the FDG setting (Zhang al. 2023b) Let H. be a hypothesis space of VC-dimension M, HHHH( (Du, Dz) is the domain divergence. For any E (0, 1), the generalization gap on an unseen domain Dy is bounded by the following with a probability of least 6:
log M + log
C(0,DT)
Σ Yu E(0;Dw)+ wELss
dHuH((uN + VELLT
E
+ (1)
2N
where 5 is the optimal combined risk on Ds and Dr. One crucial challenge in FDG the existence of a domain shift between the sets of source clients' data D and target clients' data Dr, respectively.
primary goals of FDG. While Objective easily be achieved by training on source clients Us. Objective Objective(rrequiress requires the FDG system to minimize the discrepancy between data from source clients Us and unseen target clients UT. However, Inequality requires the availability of both source and target data (i.e., D, E Ds, D. E DD, Vu, E Uy) Consequently, one of the most significant approaches to achieving generalization is to design model that is domain- invarian (Li et al. 2018 Zhao et al. 2019) by learning through the objective:
Objectives [anghhighhlghh
0" mig E10 Ds) + Σ Σ (MP((O.D...P(::D..))
(2)
WEMs wElls
where À is the scaling coefficient of domain- invariian regularization. The principle of Nomain-invariann learning is to learn model that remains invariant across source domains, rather than reducing the divergence between the source and unavailable target domains. As a consequence, learning efficiently reduces the generalization gap between source and target domains (Li et al. 2018). However, due to the decentralization of the data, each client does not have access to other clients' data, ie., Du n D. = Ø, Vu, v E Us u e U. Consequently, it is challenging to find the domain-invariant invariant features within the FDG framework.
3 MOTIVATION BEHIND ON-SERVER OPTIMIZATION FEDERATED LEARNING
3.1 FROM FEDERATED LEARNING TO ON-SERVER OPTIMIZATION FEDERATED LEARNING
To efficiently aggregate knowledge from all clients, propose an on-server optimization framework designed to identify an optimal aggregated gradient for achieving domain invariance, thus, eliminating the need for direct access to client data. This approach leverages meta-learning decompose the FDG optimization problem into two meta-learning steps. Specifically, we reformulate the FDG optimization problem from Eq. (2) into two sequential stages: local update and meta update. Drawing on the meta-learning principle (Hospedales et 2022). the problem is formulated as follows:
0(r) -0( 0(r-1) -g
Σ uells
VE(O()) D.) + E
Vd Vdl.(F(:.D.,PPPO.D...
(3a)
EEMs
s.t. (() = ((r mg Dw):
(3b)
where yu and ny are the local and global learning rate of FL system, respectively. The local update in Eq. (3b) is designated for on-device training, while the meta update in Eq. (3a) is utilized for on-server update. The on-server update in 3ap consists of two terms: the first term VE(a!)) D.) refers to the FL update, and the second term (P Dw), I((; D&)) represents the update of domain divergence reduction between any pair of clients from the source client set Us. However, as aforementioned in Section 2 the integration of Eq. (2) into Eq. (3) is infcasible due to the requirement of accessibility to all source clients' data.
3

---

Published as a conference paper at ICLR 2025
3.2 MOTIVATION OF DOMAIN GRADIENT MATCHING
Addressing the demand for local data access in conventional DG approach, we utilize Invariant Gradient Direction (IGD):
Definition (Invariant Gradient Direction (Si et al. 2022)) Considering a model 0 with task of finding Comain-invariant features, the features generated by the model domainn-nvariiie if the two gradients g(0; Du). and g(0; Do) point to similar direction, i.e., g(0; D.) g(0; D.) 0. Otherwise, the invariance cannot be guaranteed if g(0; Dw) g(0; D.) S.
Given Definition D we can leverage the local gradients as data for on-server optimization. Recently, many researches have been carried out to find the invariant gradient direction among domains. regularization of the joint
However, most of the researches consider gradient minimization objective function. For instance, Fishr (Rame et al. 2022 Eq. (4)) considers
N Σ S No
Epists = Us
Σ E(0()) wells
Σ NElls
11, s.t. afr)
9 ) (4)
Us
where gu = N. Ei Pu.i. and atr) = Ug Eveut du are the client mean gradient and global mean gradient variance, respectively. Another approach is to use GIP as regularization. For instance, Fish (Shi et al. [2022] Eq. (4)) considers
Eiis
1 Us
Σ E(0M) wElls
X Us(Us
2
g ( 1867)
BElls vElls
On-server data required if we leverage Eq. (4) Eq.
5 the objective function for on-server optimization. Acknowledge this shortcoming, we design a meta-learning based approach and apply the gradient matching solely during the meta-update stage. Therefore, the FDG update can be represented as follows:
B(r) g(r-1) ng (o(cb:
(6a)
s.t. GGD arg max
Σ WElds
(g(0: ((0::).g(og) (clo:)).(o():DD)) T(4) -000 mg(0(r-1); D.))
f(r-1) -
(6b)
Here, we abuse ghe g(olr); ; D.) as the local gradient of user u on round r after E epochs. g(0;-) -) 0 0 is the gradient computed by learnable parameter 0, thus, we do not consider dataset. Note that Eq. (6) has a high similarity with (Shi et al. 2022 Eq. (4)). Apart from (Shi
al. 2022), we use Eneiks (o(O; )) g(ol) :-)) instead of Eveus Eneus (9l" (sk", .g!"") because (r)
(o.") 90 ) is double-edged bounded by (o(0; -), g!)) and (g(0;)), "!)")) Vu, v. explained in the following.
Lemma (Triangle inequality for cosine similarity (Schubert 2021 Let Da (r) gg .0(0;)) be three vectors in a M -ypprrplaann, then the following bounds hold:
( g (0(:)))))) - v(1 (ole) (1) (g(0) )) ),gg) )12)))))))) < (o("))
< (ol)) g(0; )) (g(0;)), V (ol))) .(0;; )1 (9(0; )), (.))
By proposing Eq. (6. we establish the meta objective function Eq. 6ay. which leverages only the clients' gradient as training data. The remaining issues Eq. (6) are two-folds. Firstly, the optimization problem over the model f may not achieve good generalization due to the lack of training data at every optimization step (which not applicable to model with large parameters). Secondly, as noted by (Sii et al. (2022), minimizing the GIP can result in significantly high computational overhead, primarily due to the requirement of computing second-ord -oddrr derivatives of the model parameters associated with the GIP term. In our proposed method, we will focus making this objective function feasible while not require extensive computational resources due to the second-orde -oddr derivatives of the model parameters.
4

---

Published as a conference paper at ICLR 2025
4 FEDOMG: DATA FREE ON-SERVER MATCHING GRADIENT
4.1 OVERALL SYSTEM ARCHITECTURE
We propose a novel method FedOMG, which efficiently aggregates the clients' knowledge learn generalized global model 0g. The main differences between our approach and other current FL and FDG approaches are two-fold. Firstly, to guarantee the communication efficiency of FL settings, we leverage the available local gradients as training data for our on-server optimization.
g/ 0(r 0(r-0) where 0(r,0) 0(r-1).
(8)
Secondly, we design novel on-server optimization approach. The goal of this on-server optimization to find encoder with the capability of generating domain-invariant features. To this end, leverage clients' gradients in the global model 8g that achieves the invariant gradient direction on all domains.
4.2 INTER-CLIENS GRADIENT MATCHING
To make the objective function 6 feasible, we target to do the following: 1) limit the searching space of the meta-update, 2) leverage Pareto optimality to reduce the computation overheads, and 3)
indirectly find invariant gradient direction.
Indirect Search of Invariant Gradient Direction. As mentioned in Section 3.2 the utilization of GIP may induce significant computation overheads due to the on-server training, the optimization over model parameter requires the second order derivative. To this end, instead of finding a gradient
solution with M -ddimenniionll optimization variable ofc) arg max, Eneus (g(0; 1), g(0l;; -)),
consider the invariant gradient direction 9GGD convex combination of local gradients go
(r)
FIGD
lg =
Yugh"
(r) where, I = { ns g = {o , ................................ 9is
(9)
uEUs
By doing so, the FDG update in Eq. 6 is reduced to
g(r)
11
NggioD
f(r-1)
"7 -ggiicgg(7)
(10a)
g
s.t. TIGD = arg max
Σ uells
[gg (r) g( d(r) = 0(r-1)) (
(((M)()) (w))
(10b)
By indirectly optimizing the joint gradients over the auxiliary optimization set r, the need for computing second-order derivatives with respect to the model parameters 0 eliminated. Additionally, the optimization variable dimensionality is reduced from M to Us (M > Us), which further improve the computational efficiency.
Searching Space Limitation.
The straightforward optimization of GIP, as formulated in Eq. (1O), may introduce an optimization bias toward gradients with the dominating magnitudes. This bias can result in a loss of generalization, potentially causing the optimization process to overlook clients that contribute less significantly during a given communication round. An alternative approach for GIP involves leveraging cosine similarity. By doing the influence of gradient magnitudes among local gradients is minimized, allowing the focus to shift to optimizing the angles between gradients. However, this method is computationally intensive and challenging simplify into more practical formulation. To this end, we limit the searching space into the M -bll. For instance,
FIGD arg max
Σ (Ig "g!"")
(ITg
(r)
x rellot?11?
(11)
Nells
Gradicnt matching
Scarching space limitation
where (2) represents the gradients of the referenced FL methods (e.g., FedAvg (McMahan al.
2017), where of EuEus N.ggu /Enetts. N..) Beside the aforementioned first term, Eq. 11
consists of searching space limitation. Specifically, searching space limitation activated by con- straining the searching radius, preventing it from diverging too far from the specific range (i.e.,
(r) BM (g 11))). where Bar ( x (12) is M-ball centered at 9FL and having JIGD
(r)
5

---

Published as a conference paper at ICLR 2025
radius =||9FL"11 II.) Note that, the searching space limitation term reduces arg max, ki Kil || (after the relaxation in Theorem D which corresponds to the denominator in the application of cosine similarity. By constraining the search space remain close to the reference FL gradient, the computational overhead of on-server optimization can be significantly reduced. This is because the optimization process requires fewer iterations to converge to optimal results. Additionally, this constraint mitigates the risk of the optimization process converging to suboptimal solutions.
(r)
(*)
Pareto-based Optimization. At every iteration, Eq. consists of loop over Ug user, which requires (Us) computation complexity. reduce the objective's complexity, we consider Pareto front, which provides trade-off among the different objectives. We have the following definitions:
Definition 2 (Pareto dominance (Zitzler & Thiele 1999)) Let 0 A E RM be two points, 6 is said dominate 6 (0t > (h) and only if & (0%) S E0))), Vu U and (00) (0"),,, 3u E U.
Definition (Pareto optimality (Zitzler Thiele 1999) is a Pureto optimal point and ((0) a Pareto optimal objective if it does not exist E RM satisfying > -- The set of all Pareto optimal points is called the Pareto set. The projection of the Pareto set onto the loss space is called the Pareto front.
Leveraging Definitions
2 and we have the following lemma:
Lemma The average cosine similarity between given gradient vector 9IGD and the domain- speccffc by the worst case cosine similarity as follows:
gradient is
Σ aEUs
min acMis
(g
:IGGD
(r)
U
Lemma Zaalooss us to realize that Eq. (1I) can be reduced to maximizing the worst-case scenario. Thus, have the surrogate FDG update as follows:
8 g(r-1) ny9iGD (r)
0gr-1) Ng (oiigg
(12a)
s.t. FIGD argmax max min
[ (Ig'"),9''" 7
- -/1 -1
x 11997? ]
(12b)
r
wells
However, solving Eq. 12 is complex due to the min-max problems with two variables. Thus, we
simplify the optimization problem Eq. (12) as follows:
Theorem (FedOMG solution) Given r = (- E Us Eneus 70 1} is the set of learn-
(r) able coefficients at each round Invariant gradient direction DiGD is characterized as follows:
9IGD 9FL
(r) (r)
K Kl1grz
1ng (r) s.t. - arg min
(r) 9FL + k) K1l9 x1199221119)
(13)
Proof of Theorem is detailed in Appendix F.5 Theorem T provides an alternative solution for aggregating local gradients the server. By relaxing Eq. into Theorem [I] the optimization process becomes simplified. Instead of requiring iterative computations over all participating users in arg maxp minwetss (rg("))g"")) the optimization can be performed through straightforward computation, c.g., arg min: (rg(() 9 The detailed algorithm is demonstrated as in Alg. |
Integratability. 9LL represents the gradients of referenced FL algorithm, while the local gradients ga
(r) depend that specific FL algorithm. Consequently, by using Theorem |
can integrate on-server optimization into other FL algorithms. Furthermore, Theorem] [ also indicates that the IGD can be reduced to the chosen FL algorithms if we choose appropriate ppper-parameters:
Corollary When the radius of the r-hypersphere reduces to 0, the IGD is reduced to the referenced (r)
FL algorithm. For instance, lim JIGD
9FL

---

Published as a conference paper at ICLR 2025
5 THEORETICAL ANALYSIS
To prove the generalization capability of FedOMG, we have the following lemma:
Lemma For any 0 E ④ the domain divergence (A, B) bounded by the expectation of gradient divergence between domain A and domain B.
- 16
dncn B)
dgooo (A B),
(14)
where dgoo(A) B) is the gradient divergence of model 0 when training in two domains A and B.
Proof of Lemma is detailed in Appendix F.3 Lemma 3 establishes the connection between domain
divergence and the gradient divergence it induces. Consequently, the domain shift dH H (P.,, D.) from Eq. ( can be interpreted as gradient divergence, which serves as the primary minimization objective our work. We then prove that the generalization gap on unseen domain Dy by the optimal solution on Ds is upper-bounded by the generalization gap in the source domains in Theorem 2 From Theorem the domain generalization gap on an unseen domain Dy is bounded by X and the domain divergence dHAH (D., D.).
Let 0 denote the global model after R rounds FL, Or and 07 mean the local optimal for each client and the unseen target domain, respectively. The local objectives follow the μ strongly convex from Assumption 2 For any 8 (0, 1), the domain generalization gap for the unseen domain Dy can be bounded by the following equation with a probability of at least 1-2
Theorem 2
ED (0R) ED,(Ob,)
dgoel (Go,(D.. + dHnnH(Ds DT)
log M
V log
Ms.
Σ Tu (0)+ WElls
Σ vells
S*
11
V2N
where Du, D,, are the sampled counterparts from the domain D,, D., respectively. dgon (D.., Do) denotes the gradient divergence of model 0 when training on two different domains D., and D...
Proof of Theorem 2 is detailed in Appendix F.4 The generalization gap at round R on target domain Ep, (0 Ep, b affected by the divergence among domains' gradients (D.. D.). These domain gradients are affected by the FedOMG algorithm, where our goal is minimize the gradient divergence at every round. By combining Theorem and Lemma we obtain the upper-boundary on the FedOMG.
Corollary 2 Assume FedOMG solution given by Jigp arg max; Eveus Eveue deeo (Ge D.). Generalization gap of FedOMG is bounded by the followings with a probability of at least
(r)
(A") Ed, (0 5
Σ Yu E wells
(0)
max G
Σ
dgoe (Du, Du,DD D.)
P
wells
log
UsM + log
+ dHs
+
V2N,
Consequently, the gradient matching can reduce the FedOMG generalization gap every round, thus reducing the generalization gap on the unseen target data domain Dy.
EXPERIMENTAL EVALUATIONS
ILLUSTRATIVE TOY TASKS
Fig. presents the performance of FedOMG on Rect-4, revealing several key findings. Firstly, significant gradient conflict is observed when training in systems with heterogeneous users (e.g.. feature divergence), while in the IID setting, the gradient is more aligned toward one direction.
7

---

Published as a conference paper at ICLR 2025
Secondly, in FedAvg, users 1 and 2 tend to steer the global model away from the ideal ERM (where data is randomly sampled and not affected by heterogeneity). This issue is exacerbated in the FDG setting, causing the gradient to shift upward. meaning the bias increases (ideally, the bias should be close to zero). Thirdly, FedOMG demonstrates effective gradient matching, aligning closely with ERM gradients in the FL setting moving towards the optimal point in the FDG setting. The details of experimental settings are demonstrated in Appendix C
3D Pat Model Para maters Across Federated Leaming Rounds
3D Plot of Model Farameters Across Federated Learning Rounds
ttxe
cMn tier
Uie'2 Clkbe
Feshig Gozd
Usee
Uvet
EMM Joar4
OMG Geea
100 User
DMC Chibul
Fechi User
ONG 3
Uver
User
3
Reight
Weight
Wilght
Figure I: Illustrative toy task on two settings. 1) FL settings, where all users are participating in the training (left figure), 2) FDG setting: one user is excluded in the training (right figure).
6.2 EXPERIMENTAL MAIN RESULTS
FL performance comparison with baselines. Tab. ! presents the results for our method in conventional FL setting. The best result is highlighted in pink, the second-best result is in bold, and our proposed method highlighted in yellow. We observe that our method consistently achieves the best or second-best performance across all datasets compared to other methods. Notably, our proposed algorithm demonstrates robustness, as integrates effectively with other FL methods, resulting in superior performance. Full results are shown in Appendix E.1
Discussion on orthogonality of FedOMG with FL baselines. Performances of the integration of FedOMG into three other FL algorithms (i.e., PerAvg, FedRod, FedBabu) are shown in Tab. T The integration of advanced methods consistently outperforms standard FL algorithms. Notably, the combination of FedRod and FedOMG yields significant enhancement in performance. This improvement can be attributed to the complementary nature of these two techniques. FedRod introduces layer disentanglement, which enhances the personalization capabilities of FL clients. In contrast, FedOMG focuses on on-server gradient matching, which improves the overall generalization of the FL model. Together, these orthogonal approaches create synergistic effect that boosts both personalization and generalization in FL.
Table Comparison of methods across different datasets and non-IID scenarios. The setups are 100 users at 20% user participation (i.c. a E {0.1, 0.5}). Results are averaged after 5 times.
Setting
Non-IID (cr = 0.1)
Non-IID (o = 0.5)
Method
MNIST
CIFARIO
CIFARI00
EMNIST
MNIST
CIFARIO
CIFARIO0
EMNIST
FedAvg
93.77 + 0.33
60.89
0.12
28.78 0.
86.45 + 0.30
91.47
0.26
59.12 0.21
0..22
84.23 0.188
PerAvg
94.44 = 0.26
64.76 2 0.13
06.3720.32
87.87
2 0.17
93.52
0.32
64.46
0.02
27.29
0.220
84.87 0.01
FedRod
98.09 0.27
88.90=000
44.33
97.50 + 0.24
96.85
= 0.11
70.52
0.31
0..12
96.46
0.22
FedPac
96.90
87.81
48.83
97.74
0.32
94.63
0.22
73.02
0.22
94.10
FedBabu
96.40 0: 0.08
88.19
0.19
49.18 11 0.09
96.76
2 0.31
95.
= 0.24
70.91
= 0.26
28.33
0.23
93.11
0.02
FedAS
97.91 0.22
89.15
0.06
50.37
0.18
97.71
0.15
96.78
0.33
75.75
0.21
32.57
0,31
94.53 0.12
FedOMG
98.75 0.00
90.40
0.03
48.76
+ 0.15
98.91
0.02
98.58
0.09
72.68
0.02
30.64
0.32
98.41
0.02
Per
95.04
0.26
84.46
0.13
44.52
0.32
98.18
0.32
94.52
0.32
70.76
0.02
28.29
0.20
90.68
0.01
RedRod+OMG
99.63
0.08
95.88
0.04
0.16
99.62
0.03
0.08
76.80
0.12
36.78
0.04
99.10
0.63
FedBBabuxOOG
98.12 + 0.04
92.19
0.09
$1.27
0.09
98.62
0.20
97.28
0.09
75.41
0.13
31,82 + 0.16
97.25 0.02
FDG performance. In Tab. 2 we report the results for domain generalization performance. Our
FedOMG consistently outperforms other methods, i.e., 5% 6% gain in target domain accuracy.
8

---

Published as a conference paper at ICLR 2025
Discussion orthogonality of FedOMG with FDG baselines. Performances of the integration of FedOMG into three other FDG algorithms (i.c., PerAvg, FedRod, FedBabu) are shown in Tab. [ The integration of advanced methods consistently outperforms standard FDG algorithms. However, only FedSAM+OMG outperforms the standalone FedOMG model. This is because FedSAM and FedOMG address orthogonal and fundamental aspects of FDG (i.e., FedSAM reduces the sharpness of the loss landscape, while FedOMG focuses on mitigating gradient divergences across domains). In contrast, the regularization techniques used in FedIIR and FedSR primarily aim to minimize domain divergence on the client side, contribution that overlaps to some extent with FedOMG's approach.
Table 2: Domain generalization performance on VLCS, PACS, and OfficeHome Datasets
Algoriti
VLCS
PACS
OfficHoome
Agg
Ang
Apg
VedAkg
61.1
916
654
91.7
06
72
4
77.9
0.
$10-1:04 0.4
82.
11.7
483-4.01
128
0.
25-3
67.5
PedGz
7.
569
10
06
734 93.9
03
812
7:
12. 0:1
835
56.
543
01
753
0.
74.7
652
FwiSAM
14..1.03
929
96.1
89
733
816
935
643
FedSit
02..1.03
623 :
938 0.
74.4
066
755
94.
066
0.
01.7::0.1
934
57.9
02
80.3
06
0.1
2.
643
7.6.1.01
991. 0.7 98.1 002
70.5
75-3
04.
011
R50
111
7631
02
1.
54.3
51.1
87.9
05
12.7
106
75.1
650
IdOOMG
K2.3 0.5
9.4 W1.
79.1
98.8
W1.
0.4
81.4
58.4
0.4
77.8
TU0
DMBIRGOMG
64.0-00 0.3 90.74-01 72.6
77.5
91.7
110-411
903
093-4.0.
85.1
63.0
0.
53.8
0.5
74.3
76.9 4
66.5
YodSAM--NIG
K2.7 67 694
8.9 W13-5000 9:3 78.5
NO.5 98.3
0.1
1809 2
53.7
6:33
6.3
55.5
65.8
569
11.
7569
9.5
79.3
0.7 7099
PaiSB-OG
0.3 948 1.02
768
793
3.3
84.8
CIFAR10
CIFAR100
1.0
18st Accuracy
Test Accuracy
Test Accuracy
Test Accuracy
Rd-
TeacMG
fedio0
0.7
Pardeg
FeSRCD-CMM
FedDASU
fadieg
Farkig
Redh00--NG
Feshig
Fod24BU
Pasheg
FADROO-OR
2042
FudiikU
Communication Rounds
20
Communication Rounds
Communication Rounds
Communication Rounds
Figure 2: Performance comparison without pretrained models for a = 1, = 20.
6.3 ABLATION STUDIES
Effects of training FL without pretrained models: Fig. 2 demonstrates the FL training without pretrained models. FedOMG results in the more stable convergence than most other FL baselines along with FedPAC and FedRod. Notably, the combination of FedOMG and FedRod shows substantial improvement in performance compared to other baselines. In summary, FedOMG, both individually and in combination with other FL methods, offers enhanced robustness when training FL models from scratch without relying on pretrained models.
Global learning rate yp In Fig. 3 we present FedOMG results with varied 77. and fixed K to 0.5, and each result is the
20-
90
2 m
Tost Accuracy
average of three runs. In EMNIST and CIFAR-10, the gradient matching task is straightforward. Therefore, higher accelerates convergence to the optimal state. However, more challenging datasets with domain shift issues, i.c., PACS, VLCS, and OfficeHome, a lower learning rate proves to be more effective.
ac.
VLCS Datanet
VLCS Dataset
PACS Dusaset
PACS Datsac:
Officetoone Dulase:
Offichhome Datast
CIAALL Datast
CFVB1 Daranst
Global Learning Rate
Hypersphere Radius
Figure 3: Global LR r).
Figure 4: Searching ratio k.
Searching radius ratio In this evaluation, we test our method with four different values of across seven benchmark datasets. Each experiment is run three times, and the results are averaged. As shown in Fig. 4 FedOMG performs optimally when the radius is set to zs 0.5. When - 0. the performance leans towards the FedAvg, while setting a high K, the performance degrades gradually as the searching space become large, thus, is struggled to find the optimal solution.

---

Published as a conference paper at ICLR 2025
0.9
FwdDING Domain
0,14
0.75
Fedores Domain
FoddmiG Domain Bedomis_Demain3 5
0.12
Cosine Similarity
Cosine Similarity
0.70
6.10
6.7
Fwddug Domair
Ceneralizatoon
0.08
Feddug Domain
Fediug Domais 5
6.08
0.04
FedSAM
0.5
0.59
0.02
0.00
FoOMMG
GAo
Step
Step
Communication Rounds
(a) Gradient cosine of FedAvg
(b) Gradient cosine of FedOMG
(c) Generalization gap.
The evaluations gradient invariance. Algorithms with stronger invariance properties result in smaller gaps in cosine similarity between each domain's s and the global gradients. The algorithms evaluated on PACS dataset, where source domains are A, C, S, and target domain is P.
Figure 5:
Evaluation on gradient invariance: To assess the invariant property of FedOMG, we calculate the cosine similarity between the gradients of each source domain and the global gradient, and visualize in Figs. Sa 5b As it can be seen from the figures, the significance of FedOMG compared to FedAvg two-fold. First of all, the cosine similarities of all domains FedOMG hold low variance among domains. This mean that all of the domains share the same divergence compared with the global gradients. Secondly, the overwhelmingly high cosine value in domain S of FedAvg may results from the overfitting to domain S. As consequence, the FedAvg significantly lacks generalization capability compared to that of the FedOMG (see Fig. 5c
Discussion on impact of local epochs and local learning rate: Fig. Gjillustrate the fine-tuning of FedOMG under varying numbers of local train- ing epochs and learning rates. The results indi- cate that FedOMG reaches optimal performance when the number of local epochs is set to 5. This suggests that after epochs, the local gradients adequately represent the users' gradient direction during each round. Extending the local training duration beyond this point incurs higher compu- tational overhead while reducing the generaliza- tion of on-server gradient matching, potentially due to overfitting to individual user datasets. Fur- thermore, the optimal performance is observed when the local learning rate is set 0.001.
0.8
Less
0.6
0.3
1 Local Training Rounds ac00
0.1
0,005 0.01 0.05
Leaming Rate
Figure 6: Evaluations local training epochs and learning rate.
Additional results. Detailed results when training with different target domain VLCS, PACS. OfficeHome are demonstrated in Appendix E.3 Results on computation cost are demonstrated in Appendix E.4 Additional results conducted on FDG benchmark (Bai al. 2024) demonstrated in Appendix E,5
CONCLUSION
In this paper, we present FedOMG, novel and effective federated domain generalization method aimed addressing the challenge of heterogeneous client data distributions by learning an invariant gradient. Our method introduces two key innovations: (1) utilizing the local gradient as critical factor for server-side optimization, and (2) devising training strategy that identifies optimal coeffi- cients to approximate the invariant gradient across client source domains. We provide theoretical proof to ensure the global invariant gradient solution and conduct extensive experiments demon- strating FedOMG's significant improvements in both standard FL and domain generalization (FDG) settings. This work establishes an optimal approach for server optimization, leading to enhanced FL performance while maintaining privacy constraints.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by Institute of Information communications Technology Planning & Evaluation (ITPP) under the Artificial Intelligence Convergence Innovation Human Resources grant funded by the Korea governmen((MIITT) This re- search was supported by the MSIT(Ministry of Science and ICT), Korea, under the ITRCCInformation Technology Research Center) support program(IITP-2025 20L-TTP-222-----222222 2023-00260098) supervised by the ITP(Instiutee for Information & Communications Technology Planning & Evaluation). This work was supported by the National Research Foundation of Korea(NRF) grant funded by the Korea The work of Quoc- Viet Pham is supported in part by the European Union under the ENSURE-6G project (Grant Agreement No. 101 101182933).
Development 2--2--2--2223-002-20
government(MSSTTRRS (SMSSTTTTRRS-222200033666000
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. Int. Conf. Learn. Represent., May. 2021.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint
Jul. 2020.
Ruqi Bai, Saurabh Bagchi, and David I. Inouye. Benchmarking algorithms for federated domain generalization. In Int. Conf. Learn. Represent., May 2024.
Hong-You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. Int. Conf. Learn. Represent., May. 2022.
Junming Chen, Meirui Jiang, Qi Dou, and Qifeng Chen. Federated domain generalization for image recognition via cross-client style transfer. In Wint. Conf. on App. of Comput. Vis, Jan. 2023.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and André van Schaik. EMNIST: an extension of MNIST to handwritten letters. Proc. IEEE, Feb. 2017.
Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. Personalized federated learning with theoretical guarantees: model-agnostic meta-learning approach. In Adv. Neural Inform. Process. Syst., Dec. 2020.
Ziqing Fan, Shengchao Hu, Jiangehao Yao, Gang Niu, Ya Zhang, Masashi Sugiyama, and Yanfeng Wang. Locally estimated global perturbations are better than local perturbations for federated Sharpness-award aware minimization. In Int. Conf. Mach. Learn., Jul. 2024.
Yaming Guo, Kai Guo, Xiaofeng Cao, Tieru Wu, and Yi Chang. Out-of catoof-distribution generalization of federated learning via implicit invariant relationships. In Int. Conf. Mach. Learn., Jul. 2023.
Timothy Hospedales, Antreas Antoniou, Paul Micaclli, and Amos Storkey. Meta-Learning in Neural Networks: A Survey IEEE Trans. Pattern Analysis. Machine. Intell, Sep. 2022.
W. Huang, M. Ye, Z. Shi, and B. Du. Generalizable heterogeneous federated cross-correlation and instance similarity learning, IEEE Trans. Pattern Analysis. Machine. Intell, Feb. 2024.
Wenke Huang, Mang Ye, and Bo Du. Learn from others and be yourself in heterogeneous federated learning. In IEEE Conf. Comput. Vis. Pattern Recog., Apr. 2022.
Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedexp: Speeding up federated averaging via extrapolation. In Int. Conf. Learn. Represent., May. 2023.
Enyi Jiang, Yibo Jacky Zhang, and Sanmi Koyejo. Principled federated domain adaptation: Gradient projection and auto-weighting. In Int. Conf. Learn. Represent. May. 2024.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Int. Conf. Mach. Learn., pp. 5132-5143, 2020.
11

---

Published as a conference paper at ICLR 2025
Alex Krizhevsky. Learning multiple layers of features from tiny images. University of Toronto, May. 2012.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proc. IEEE, Aug. 1998.
Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain generalization. Int. Conf. Comput. Vis., Oct. 2017.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization for heterogeneous networks. In Int. Conf. Mach. Learn, Apr. 2020.
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Eur. Conf. Comput. Vis., Sep. 2018
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Adv. Neural Inform. Process. Syst., Dec. 2020.
Quande Liu, Cheng Chen, Jing Qin, Qi Dou, and Pheng-Ann Heng. FedDG: Federated domain generalization on medical image segmentation via episodic learning in continuous frequency space. In JEEE Conf. Comput. Vis. Pattern Recog., Jun. 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--fffccient learning of deep networks from decentralized data. In Proc. of Int. Conf. on Arti. Intell. and Stat., Apr. 2017.
Tuan Nguyen, Philip Torr, and Ser Nam Lim. Fedsr: A simple and effective domain generalization method for federated learning, Adv. Neural Inform. Process. Syst., 35:38831-388-4 Dec. 2022a.
A. Tuan Nguyen, Toan Tran, Yarin Gal, Philip Torr, and Atilim Gunes Baydin. KL guided domain adaptation. In Int. Conf. Learn. Represent., May. 2022b.
Jaehoon Oh, SangMook Kim, and Se- Young Yun. FedBABU: Toward enhanced representation for federated image classification. In Int. Conf. Learn. Represent. May. 2022.
Kunjal Panchal, Sunav Choudhary, Subrata Mitra, Koyel Mukherjee, Somdeb Sarkhel, Saayan Mitra, and Hui Guan. Flash: Concept drift adaptation in federated learning. In Int. Conf. Mach. Learn., Jul. 2023.
Jun-Gyu Park, Dong-Jun Han, Jinho Kim, Shiqiang Wang, Christopher G. Brinton, and Jackyun Moon. StableFDG: Style and attention based learning for federated domain generalization. In Ady. Neural Inform. Process. Syst., Nov. 2023.
Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In Int. Conf. Mach. Learn., Jul. 2022.
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of distribbiint generalization. In Int. Conf. Mach. Learn., volume 162, pp. 18347-18377, 17-23 Jul. 2022.
Erich Schubert. A Triangle Inequality for Cosine Similarity. Springer International Publishing, Jul. 2021.
Yuge Shi, Jeffrey Secly, Philip Torr, Siddharth N, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. Int. Conf. Learn. Represent., May. 2022.
Yan Sun, Li Shen, Shixiang Chen, Liang Ding, and Dacheng Tao. Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. In Int. Conf. Mach. Learn., Jul. 2023.
Zhenheng Tang, Yonggang Zhang, Shaohuai Shi, Xinmei Tian, Tongliang Liu, Bo Han, and Xiaowen Chu. Fedimpro: Measuring and improving client update in federated learning. Int. Conf. Learn. Represent., May. 2024.
12

---

Published as a conference paper at ICLR 2025
Antonio Torralba and Alexei A. Efros. Unbiased look dataset bias. IEEE Conf. Comput. Vis. Pattern Recog., Jun. 2011.
Hemanth Venkateswara, José Eusébio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. IEEE Conf. Comput. Vis. Pattern Recog.. Jul. 2017.
Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. Int. Conf. Learn. Represent., May. 2023.
Xiyuan Yang. Wenke Huang, and Mang Ye. FedAS: Bridging inconsistency in personalized federated learning. In IEEE Conf. Comput. Vis. Pattern Recog-, pp. 11986-11995, Jun. 2024.
Liling Zhang, Xinyu Lei, Yichun Shi, Hongyu Huang, and Chao Chen. Federated learning for iot devices with domain generalization. IEEE Internet of Things Journ., Oct. 2023a.
Ruipeng Zhang, Qinwei Xu, Jiangehao Yao, Ya Zhang. Qi Tian, and Yanfeng Wang. Federated domain generalization with generalization adjustment. In IEEE Conf. Comput. Vis. Pultern Recog., Apr. 2023b.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. Ins. Conf. Mach. Learn., Jun. 2019.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogencous federated learning. In Int. Conf. Mach. Learn., Jul. 2021.
E. Zitzler and I.. Thiele. Multi-objective evolutionary algorithms: a comparative case study and the strength pareto approach. IEEE Trans. on Evo. Comput, Aug. 1999.
13

---

Published as a conference paper at ICLR 2025
Appendix
The appendix is organised as follows:
Appendix provides the literature reviews.
Appendix B
provides details of proposed algorithm.
Appendix provides toy dataset description.
Appendix D provides extensive experiment settings and hyper-parameten
Appendix provides additional results for standard FI. and ablation test for generalization performance.
Appendix provides proof on theorems.
A
RELATED WORKS
Federated Learning:
To contextualize our paper's contribution to on-server optimization, FL can
be divided into two primary categories: 1) local training and 2) on-server aggregation. In the first
category, substantial efforts have heen dedicated to improving the performance of FI., particularly in
scenarios involving non-IID data. FedRod (Chen & Chao 2022) proposes using balanced softmax
for training generic model while utilizing vanilla softmax for personalized heads. FedBABU (Oh
et al. 2022) keeps the global classifier fixed during feature representation learning and performing local adaptation through fine-tuning FedPAC (Xu et al. 2023) adds regularizer local training to
learn task-invariant representations. Similar to FedPAC, FedAS (Yang et al.) (2024) finds the shared parameters which are task-invariant by leveraging the parameter alignment technique. Recently, researchers have focused on on-server aggregation to improve the performance of FL. This approach can be implemented by utilizing generative AI on the server with pseudo data (Zhu et al. (021), or
leveraging local gradients to adjust the learning rate (Jhunjhunwala et al., 2023 Panchal et al. 2023).
Federated Domain Generalization and Adaptation: Huang al. (2022) and Huang et (0024)
leverage unlabeled public data to reduce bias toward local data distributions. Nguyen et al. (002aa)
regularizes local training via conditional mutual information. Guo al. (2023) adds regu larization on the local training to implicitly align gradients between the local global models. Qu et al. (2022), (un (2023) and Fan et al. (2024) leverage sharpness aware minimization to smooth the loss landscape. Tang al. (2024) utilizes the shared local hidden features to reduce the gradient dissimilarity. Zhang ct al. (2023b) introduces server-side weight aggregation adjustments leveraging the auxiliary local generalization gap. Diang et al. 2024 utilizes the divergence between source and target data gradients to formulate a joint server aggregation rule. However, requires access to target domain data, which may not be feasible in domain generalization scenarios, PPrk al. (2023) introduced style-based strategy to enhance the diversity of data on source client. Additionally, they integrated attention into the network to highlight common and important features across different domains.
B DETAILED ALGORIHTM
In Algorithm I we present the detailed training procedure of our FedOMG approach. The primary contribution of our method lies the server-side optimization, while the client-side follows the standard FedAvg algorithm (McMahan et al. (017), ensuring no additional communication or com- putation overhead. During each communication round, clients receive the global model parameters and perform local training using SGD for E epochs. The locally updated parameters are then sent back to the centralized server for optimization. On the server, local gradients are computed and used as key components in solving Eq. [15). The resulting gradient closely approximates the invariant gradient direction, as demonstrated in Section 4 Subsequently. the global gradient is computed as shown in Eq. (16), and the global model is updated according Eq. (17).
14

---

Published as a conference paper at ICLR 2025
Algorithm 1: Federated Learning via On-server Matching Gradient
Input: set of source clients Us, number of communication rounds R. local learning rate 7), global learning rate Ilg, searching space hyper-parameter K.
Output: 0g40
Clients Update:
for client u E Us do
Receive global model 0!-) ofr)
for local epoch e E do
Sample
from local data Du;
Calculate gradient (E(()))) .));
Update client's 's model: g(re+1)_ = 0lrr)) - -V ,));
end for
Upload client's model (r.EE) server;
end for
Server Optimization:
for round 0,..., R do
Clients Updates:
Calculate gu 0 0(1)) .g [9_ Iu E Us}:
Calculate gEL (e.g.. gPL 1 E Iu the FedAvg update);
(r)
(r)
(r)
Solve for Nx:
arg min (g gn. + x xllsn (r) (g
(15)
Update the model:
D (r)
K|
g
(16)
|| g
(()) ((r-1)
7g9IGD
(17)
end for
15

---

Published as a conference paper at ICLR 2025
C TOY DATASET DESCRIPTION
Dataset and Settings Descriptions. Based on Zhao et al. (2019), we design a toy dataset, coined Rect-4, a synthetic binary classification dataset with 4 domains, according 4 different users. each domain, a two-dimensional key point Iı = (44.1,Id..) is randomly selected in the two-dimensional space with varying region distributions. To visualize the gradient of the toy dataset, we design 1 I-layer, 2-parameters network. To this end, we can visualize the gradients in 3-D space, consisting of parameters and weight.
We visualize the training data in Figs. and 8 domains. To this end, we design the data where point are distributed into rectangular with different size and shape. The rationale of designing the data distribution is as follows:
In the FDG setting, the users are from different
The global dataset consists of two classes from two rectangular regions, which has the classification boundary equal y O.
Each domain-wise dataset has different classification boundary (e.g., x -6 for domain 1). We add the noisy data on every domains so that the user assign to each domain will tend to learn the local boundary instead of the global boundary. Thus, divergence more clearly, as the global boundary is not the optimal solution when learn on local dataset.
observe the gradient
All of the local classification boundary is orthogonal from the global classification boundary, thus, we can make the learning more challenging despite the simplicity of the toy dataset.
Domain
Domain
Domain
Clan
Clasi
Cles
Clasi
Feature 2
Har Harzona decalos Bouncay
Feature 2
Resture 2
Feature
Hot Deriaaaaaaaiiieee dear worreay
Glas
Clsi
Halcoical beciloe Bourdry
Heriecna DDccioo Deardayy
Feature
Feature
Festure
Feature
Figure 7: Illustration of users with different domains.
Domain
Domain 2
Domain
Domain
Chams
Feature 2
Feature 2
Feature 2
Feature 2
Hundey
Feature
Feature
Feature
Feature
Figure 8: Illustration of users with same domains.
16

---

Published as a conference paper at ICLR 2025
D EXTENSIVE EXPERIMENTAL SETTINGS
D.1 DATASETS
MNIST [ecun et al.
1998) consists of 70,000 grayscale images of handwritten digits and ten classes
EMNIST (Cohen et al.
2017) extension version of original MNIST dataset consists of 70,000
samples and 10 classes
CIFARIO Krizhevsky "Bird", "Cat", "Deer",
2012
consists 60,000 images across 10 classes "Airplane", "Automobile", "Frog", "Horse". "Ship" and "Truck"
Dog
CIFAR-100 Krizhevsky
2012 is a more challenging extension of CIFAR- 10, consisting 60,000
images distributed across 100 distinct classes.
VLCS Torralba & Efros
2011) includes 10,729 images from four domains: "VOC2007", "LabelMe", "SUN00 "SUN09" and five classes ("bird', 'car', 'chair', 'dog" and 'person").
"Caltech101",
PACS (Li et al. 2017) includes 9,991 images from four domains: "Photos", "Art", "Cartoons", and "Sketches" and seven classes ('dog', 'elephant', 'giraffe', 'guitar", 'horse', 'house", and 'person").
OfficeHome (Venkateswara et al. 2017) includes four domains: "Art", "Clipart", "Product", and "Real". The dataset contains 13,588 samples and sixty five classes.
D.2 BASELINES
We consider the following baselines: FedAvg (McMahan et al. 2017). PerAvg Fallah et al. 2020). FedRod (Chen & Chao 2022). FedBABU (Oh al. (022), FedPac (Xu et al. 2023
FL Baselines.
and FedAS (Yang et al. 2024) comparing with our proposed method FedOMG.
FDG Baselines. We implement the following methods: FedAvg (McMahan al.
2017), FedGA
iZhang et al., 2023b),, FedIIR (Guo et al. 2023), FedSam (Qu et al. 2022), FedSR (gguyen et al.
2022ab and StableFDG (Park et al., 2023
D.3 EVALUATION METRIC
We report accuracy as the standard metric, defined as the ratio of correctly paired samples to the total number of samples, with Top-1 accuracy being used. FL scenarios, accuracy is reported as the ratio of the total number of correct samples across all clients to the total number of samples. In FDG scenarios, accuracy is reported for each test domain individually, and the average accuracy across all domains is provided.
D.4 MODEL
For the standard FL task, we use a CNN architecture consisting of two convolutional layers (32 and 64 channels) followed by max pooling layers. It includes two fully connected layers with 512 and number output classes units, respectively, and ends with softmax output for class probabilities.
For the domain generalization task, employ pretrained ResNet-18 backbone model ZZhng ctt al. 2023b] Park et al. 22023] Guo et al. (023), which is downloaded from Py PoTcch. The detailed architecture of the ResNet- model used is outlined in Tab. 3
Table 3: Summary of model architecture of ResNet-18 model for FDG.
Layer Name
Number of Layers
Parameters (M)
ConvI
0.009
Conv2 (Residual Blocks)
4
0.073
Conv3 (Residual Blocks)
4
0.231
Conv4 (Residual Blocks)
4
0.919
Conv5_ (Residual Blocks)
4
3.673
Fully Connected
0.513
Total
18
11.69M
17

---

Published as a conference paper at ICLR 2025
D.5 HYPER-PARAMERR PARAMETERS
In this section, we present the hyper-parameters chosen for FedOMG across different datasets.
Table 4: Hyper-paramett Summary PACS, VLCS, and OfficeHome
Method
parameee
PACS
VLCS
OfficeHome
Global
5 10-
5 x 10-2
5 x 10-2
Training Ir Iteration
25
25
25
FedOMG
Momentum
0.5
0.5
0.5
Searching radius
0.5
0.5
0.5
Global Ir
5 x 10-2
x 10-2
x 10-2
Training Ir Iteration
25
25
25
21
21
21
FedSR+OMG
Momentum
0.5
0.5
0.5
Searching radius
0.5
0.5
0.5
L2 Regularizer
0.01
0.01
0.01
Cmi Regularizer
0.001
0.001
0.001
Syper-parameter
1 10
1 10
1 10
Global Ir
5 10
5 10
5 10
Training Ir
25
25
25
FedIIR+OMG
Iteration
21
21
21
Momentum
0.5
0.5
0.5
Searching radius
0.5
0.5
0.5
Global Ir
5 x 10-2
5 x 10
5 x 10-2
Training Ir
25
25
25
Iteration
21
21
21
FedSAM+OMG
Momentum
0.5
0.5
0.5
Searching radius
0.5
0.5
0.5
Perturbation control
0.6
0.6
0.6
momentum parameter 3
0.6
0.6
0.6
D.6 IMPLEMENTATION DETAILS
To evaluate the performance of our proposed FDG method, we adopt the methodologies described in Lin et al. (2020); Acar et al. (202T) to simulate non-lID data. We utilize the Dirichlet distribution with two levels of data heterogeneity, specifically
0.1 and comprises three configurations of 100 clients, with join rate ratios of 1, 0.6, and 0.4, respectively, and global communication rounds set at 800. Each client undergoes 5 local training rounds using local learning rate of 0.005, an SGD optimizer, and a batch size of 16. To ensure a fair comparison, all methods are evaluated with the
0.5. Our experimental setup
network architecture and settings.
For FDG scenarios, we evaluate model performance using the conventional leave-one-domaine method (Guo et al.) 2023 Fan et al. 2024), where domain designated the test domain and all other domains are used for training. Number of clients set to match the number of source domains.
In line with Zhang al. (223bb;; Park al. (2023): ; Guo et al. (2023). we use pretrained ResNet- backbone model. Our experiments involve 100 global communication rounds with local training rounds, utilizing the SGD optimizer with learning rate of 0.001 and training batch size of 16. Each experiment is conducted three times, and we report the average performance of the global model on the test domain, with each domain serving as the test domain once.
E ADDITIONAL RESULTS AND ABLATION TESTS
In this section, provide additional results and conduct ablation test, including (1) FL accuracy performance across different datasets, (2) convergence analysis, (3) the effect of using a pretrained model, and (4) ablation test various global learning rates and hyper-sphere radius.
18

---

Published as a conference paper at ICLR 2025
Table Comparison of methods across different datasets and non-IID scenarios.
Problem
Nou-HID
Non-lID
Dataset
Method
MNIST
CIFARIO
CIFAR100
EMNIST
MNIST
CIFARIO
CIFARIO0
EMNIST
Number of elients U - 100 and participation ratio - 0.4
FedOMG
(Ours)
FedOMG+RRd
99.34
0.04
86.74 0.05
51.70
0.23
99.31
0.01
98.90
0.11
75.08
0.12
33.97 1: 0.14
99.06 1: 0.01
PerAvg
0.12
87.04
0.04
83.53
FedRod
97.73
0.17
82.09
0.12
40.18
95.25
0.20
94.88
0.29
68.03
4: 0.32
27.33
0.16
92.46 1 0.20
FedPac
95.05 0.28
81.74
0.15
40.78
0.18
92.92
0.09
93.36
0.20
70.87
0.25
28.56
0.26
91.44 - 0.05
Fedibabu
FedAS
95.98
+: 0.24
83.35
0.06
46.20
0.76
93.71
0.15
98
0.33
72.29
0.21
31.21
0.31
92.53 + 0.12
FedAvg
88.04
0.22
55.75
0.01
25.28
0.05
85.83
0.07
85.76
0.17
49.44
0.27
22.73
0.19
82.57
0.26
Number of clients U 100 and participation ratio 0.6
FedOMG
(Ours)
0.10
0.19
FeddMG-Rad
99.55 0.15
93.72
0.23
54.69 0.12
99.61
0.03
98.93
0.11
75.84
0.12
34.51
0.10
99.08 1:18::00 0.01
PerAvg
62.199400.
0.09
87.34
0.09
89.44 0.24
60.74
0.32
25.53
0.26
83.74 1724--0..22
FedRod
97.92 97.92.10..2 0.27
86.42 1 0.99
0.22
95.51
0.09
95.82
0 0.32
68.14
0.32
28.56
= 0.09
92.92
0.28
FedPac
96.64 0.30
86.35
0.07
42.08
0.21
93.92 0.07
94.24 0.30
72.23
0.19
28.98
0.04
91.55 0.10
FedBabu
86.36
FedAS
97.12 0.12
87.35
0.13
48.46
0.21
93.91
0.15
98
0.13
0.31 53,42
73.75
31,70
0.11
93.53
0.12
FedAvg
90.17 0.22
56.35
0.06
27.46
0.14
85.94
0.15
87.90
0.33
0.21
24.07
0.31
83.71
0.19
Number of clients
100 and participation ratio
FedOMG
(Ours)
0,02
FedOMG+Rod
99.63 0.08
95.81
0.04
$5.39
0.16
99.62
0.03
99.55 0.08
76,80
0.12
36.78
0.04
99.10 1:100:1:0 0.63
PerAvg
94.44 0.266
64.76 0.13 0.13
36.27
0.332
87.87
0.17
93.52 0.32
64.46
+ 0.02
27.29
0.200
84,87 0.01
FedRod
98.09 0.27
88.90
44.33
0.26
97.50
0.24
96.85
70.52
0.31
28.17
0.12
96,46
FedPac
96.90 0.03
87.81
0.17
48.83 0.04
97.74
+ 0.32
94.63
0.22
73.02
+-1 0.16
29.94
+ 0.24
94,10 + 0.09
FedBabu
88.19
0.09
95
0.24
0.26
FedAS
97.91 1: 0.22
89.15
1 0.06
50.37
0.18
97.71
0.15
96.78
0.33
75.75
2 0.21
32,57 1: 0.31
94.53
0.12
FedAvg
93.77
0.33
60.89
0.12
28.78
0.10
86.45
= 0.30
91.47
0.26
59.12
25.55
84,23 + 0.18
E. ADDITIONAL RESULTS
We show the additional performance comparison between our method FedOMG and FL baselines under different data heterogeneity level {0.1, 0.5) in Table 5 FedOMG+ROD and FedOMG achieve the best and second-best performance across all dataset.
E.2 PERFORMANCE OF FEDOMG AND BASELINES WITHOUT PRETRAINED MODELS
MNIST
EMNIST
CIFARIO
CIFAR100
0.3
Test Accuracy
Pirtig
FAGRDGHOPG
Pativp
TedRoD+ONG
Faskao
TedRcO
FuSDABU
FedBiNUU
Feditl
Fackig
Test
0.2-
Teckeg
CCMG
Endivg
Feckop
FedRCO
Communication Rounds
Communication Rounds
Communication Rounds
Communication Rounds
Figure 9: Performance comparison of FL. algorithms without pretrained models for 0.1,U 20.
MNIST
EMNIST
CIFAR10
CIFAR100
Test Accuraty
0.2
Test accuraay cuuuray....
FeRRDO-OMG
TEdRCD+ONG Pariing
05
Parkig
Festoo
0.3
TedicO
Pedtiii Pariig
Nettc
Tedlay
fuding
Petive FedRc0
Fodlvo
TecROD
0.4
Communication Rounds
Communication Rounds
Communication Rounds
Communication Rounds
Figure 10: Performance comparison of FL algorithms without pretrained models for 0.1,U=
19

---

Published as a conference paper at ICLR 2025
EMNIST
CIFAR10
0.25
Test Accurary
Rest Accuracy
0.20
Pettar
Parlig
Aveitla iock
Feat788
reston
PudRCD
Irakon
Fediary
20
Communication Rounds
130
Communication Rounds
Communication Rounds
Communication Rounds
Figure 11: Performance comparison of FL algorithms without pretrained models for a 1, 40.
MNIST
EMNIST
CIFAR10
CIFAR100
Accurac
Test Acrurayy
1.-
Accuracy
2.8
Test Accuracy
3.4
2.6
FedRICC-ONG
Pactip
Panhi
Fective
Perdeg
tedleg
Fathyg
Festig
Teanop
Pediiy
Communication Rounds
Communication Rounds
Communication Rounds
Communication Rounds
Figure 12: Performance comparison of FL algorithms without pretrained models for = 60.
CIFAR1O
CIFAR100
Accuras
Test Accuracy
Test Accuracy
0.25
FueOHG
FuoncD-CMK
82
0.6
Fa22MC FesnCD-CW CMC
TeSBAZU
PupCMIC
Puth00 DMS
FestCC
FestADU
Fx5t00-300 DMC
TedFAC Fudb-bu
Test
Fardeg
Tusleg
Felkg
Feshig
Pastag
Poskig
10c
200
30
301 Communication Rounds
120
233
100
200
30
100
30
365 230 120 485 vin nin
Communication Rounds
Communication Rounds
Communication Rounds
Figure 13: Performance comparison FL algorithms without pretrained models for a 1,U 60.
20

---

Published as a conference paper at ICLR 2025
E.3 DOMAIN PERFORMANCE
Table 6: Performance comparison on VLCS dataset.
Methods
Backbone
S
Avg
Centralized Methods
ERM Fishr Fish
ResNet-18
74.0
67.9 = 0.7
97.6 + 0.3
70.9 0.2
77.6
ResNet-18
75.7 0.3
67.3 0.5
97.6 + 0.7
72.2 0.9
78.2
ResNet-18
74.4
66.5 0.3
97.6 + 0.6
72.7 = 0.6
77.8
Federated Methods
FedSR
ResNet-18
72.8 + 0.3
62.3
93.8 + 0.5
74.4 0.6
75.8
FedGA
ResNet-18
74.4
56.9 = 1.0
94.3
68.9
73.4
StableFDG
ResNet- 18
59.2 0.1
9.21
70.2 1.1
75.3
FedOMG
ResNet-11
82.3 - 0.5
67.5 = 0.4
99.3 + 0.1
79.1 = 0.5
82.0
Table 7: Performance comparison on PACS dataset.
Methods
Backbone
P
A
C
S
Avg
Centralized Methods
ERM Fish
ResNet-18
96.2 + 0.3
86.5 1.0
81.3 + 0.6
82.7 1.1
86.7
ResNet-18
97.9 + 0.4
87.9 0.6
80.8 + 0.5
81.1 - 0.8
86.9
Fishr
ResNet-18
95.6 + 0.3
84.7 1.2
81.1 + 0.2
80.6 - 0.5
85.5
Federated Methods
FedSR
ResNet-18
94.0 0.6
82.8 = 1.5
75.2 + 0.5
81.7 = 0.8
83.4
FedGA
ResNet-18
93.9 + 0.2
81.2 0.7
76.7 + 0.4
82.5
0.1
83.5
StableFDG
ResNet-18
94.8 + 0.1
83.0 1.1
79.3
+ 0.2
79.7
+ 0.8
84.2
FedOMG
ResNet-18
98.0 + 0.2
89.7
0.4
81.4 + 0.8 84.3
0.5
88.4
Table 8: Performance comparison on OfficeHome dataset.
Methods
Backbone
C
P
Avg
Centralized Methods
ERM Fish
ResNet-18
61.7 + 0.7 $3 0.3
74.1 = 0.1
76.2 + 0.3
66.4
ResNet-18 ResNet- 18
63.4 = 0.8
54.2 0.3
76.4 0.3
78.5 - 0.2
68.2
Fishr
60.2 = 0.5
52.2 + 0.6
75.0 0.2
76.5 + 0.3
66.0
Federated Methods
FedSR
ResNet-18
57.9 0.2
50.3 + 0.6
73.3 0.1
0.1
64.3
FedGA
ResNet-
58.5 = 0.4
54.3
73.3 = 0.8
74.7 + 1.0
65.2
StableFDG
ResNet-
57.1 0.3
0.5
72.7 0.6
72.1 0.8
65.0
FedOMG
ResNet-18
65.4 = 0.4
58.1 0.3
77.5 0.4
78.9 - 0.5
70.0
21

---

Published as a conference paper at ICLR 2025
E.4 EVALUATIONS ON COMPUTATION OVERHEAD
Fig. dedetntttrasss the computation time of FedOMG, FedROD, and FedPAC. The test are evaluated on same number of communication rounds (i.e., R 800). FedOMG shows the most computation efficient on most of settings (e.g., MNIST, EMNIST, CIFAR-10), On CIFAR-100, FedOMG compu- tation time is lower than FedROD. FedPAC tends consume significantly higher computation time than that of FedROD and FedOMG (up to 10 times).
Algorithm
160
Computation Time (hours)
FedOMG
140
FedPAC
100
80
0
MNIST
EMNIST
CIFAR10 CIFAR100
Figure 14: Computation time comparison across different datasets.
E.5 EVALUATIONS ON FEDERATED DOMAIN GENERALIZATION BENCHMARK
As highlighted in (Bai et 2024), the performance of FDG is observed to decline when influenced by the combined effects of domain shift and non-independent and identically distributed (non-IID) data. To investigate this, we conducted experimental evaluations using the Celeb-A dataset, employing
benchmark with two distinct non-lID settings (a 0.1, 1). We compare FedOMG with FedADG
(Zhang et al.) 2023a), FedSR (Nguyen al. 2022a), Scaffold (Karimireddy et al., 2020). FedAvg
(McMahan et al. 2017) and FedProx (Li et al 2020) The results of these evaluations are presented
Figure 15
Test Accuracy Worst Group
Test Accuracy Worst Group
0.6
FL_DG DG FedADG
FL Scaffold
FL DE, RdAADG
DL FedAvg
FedAvve
FL
DG FedProx
FL_DG FadOMG
RL DG RedPrax
FL, DE, RdOOG
DL ARL
Accuracy
Accuracy
FL_DG_Scarde Scaffold
0.4
0.2
0.2
12
18
8
12
18
Step
Step
(a) Gradient cosine of FedAvg
(b) Gradient cosine of FedOMG
The evaluations of FDG and FL algorithms CelebA on two non-lID settings. Fig. 15a uses a = 0.1, and Fig. 15b uses a 1.0). We use the worst group test accuracy as a metric to evaluate the algorithms Celeb-A.
Figure 15:
As shown in Figs. 15a and 15b FedOMG exhibits slower convergence during the initial stages. However, unlike other algorithms that tend to reach saturation quickly, FedOMG demonstrates gradual and consistent improvement over time. This behavior be attributed FedOMG's tendency to prioritize identifying invariant gradient direction, effectively balancing rapid convergence with stable and invariant updates. This approach enhances long-term generalization performance.
22

---

Published as a conference paper at ICLR 2025
F PROOF ON THEOREMS
F.1 TECHNICAL ASSUMPTIONS
Assumption (I-smoothness) Each local function is Lipschitz smooth, that
[VE(x;D.)) VE Pw)ll s - Do)ll Vu Us.
(18)
Assumption (j-strongly convex) Each local objective function is Lipschitz. smooth, that is, [[V[(;; VE(u; P.)Il 2 pl] D.) Eyy: D.) Yu E Us
(19)
Assumption 3 (Domain triangle inequality (Zhao et al., 2019)) For any hypothesis space H. it can readily verified that dx((ss) satisfies the triangular inequality:
dRAH(D, D") dHAH (D,D') dHAU(D D
(20)
F.2 TECHNICAL LEMMAS
Lemma 4 we have (0) Evelis PuE then for any unseen domain Dy, we have:
dDAH(Ds, = Σ (a Dr).
(21)
wells
Proof. From the definition of dHAH(::-) in AArjovsyy et al. 2020), we can get
dHAH((ss Dr)
sup ACAMAn
(ro(()
PrD, (A)l=2 =2
sup AEAHAN
Σ TaPr((A)
PrD, (A)
uells
Σ uells
Yu (r)(A) -rro (4)
sup AEANAR
< 2 sup
Σ
(A) - Proy (A)
AEANAR wEMss
2 Σ
Yu sup wells AEANAR
(AD(() PrD, (A)l
Σ
VadHAH Tad DT).
DT)
(22)
wells
Lemma For any E θ the expectation risk gap between domain and domain is bounded by the domain divergence dHAH B).
[Ea(0) - EB(()) < ().
(23)
Proof. By the definition of dHAH(---) in (Arjovsky et al. 2020), we have: dHAn(A,B)
sup Pro=alf(2;0)) 0) * f(z; 0') Pre f(z;6))). 0.0' ge
(24)
where fxx; 0) means the prediction function data æ with model parameter 0. We chose o' as parameter of the label function, then f(x; 0) + (x; 8') means the loss function L(x; 0), so have:
dHAM B) = sup Pr Pra-a [C(x 0)] Pr PP [[((x; 0)] 2 (0) - (g))
(25)
6Ee
Here, (a) holds due Assumption
23

---

Published as a conference paper at ICLR 2025
Lemma 6 (Guarantee of source domain invariance) we have Eg(() =Eneus Vul then for any domain Ds we have:
Σ
uu, Ds) S
Σ Σ
(Dd D.).
(26)
wells vElds
Proof. From the definition of dHAH((:)-) in Arjovsky et al. 2020), we can get
Σ BEuls
TodHAH Ds) =2 Σ The
sup AEANoH
Pro. (A) Prrs (A)
uElls
2
Σ
Yu AEANAH
sup
PrD.
Σ
M.Pro,(A)
NElls
wElds
Σ
Yu
sup AEANAH
Σ
YoPro, (A)
Σ
YoPrD. (A)l
Nells
vells
vells
S
Σ
Tu
Σ
To AEAHAN
sup
PrD, (A) Pro, (A)l
NElls
NEMs
Σ
Σ
Yu Yudnan (D., D.).
(27)
"EMys wElls
F.3 PROOF ON LEMMA
Proof. By the definition of duau(:;)) :) in (Arjovsky et al. 2020), we have:
dHan(A, B) = 2 sup 0,8' ee
(f + f(;; 0' Pla= + f(z; (2;00)1
(28)
where f(x; 8) means the prediction function on data æ with model parameter 6. We chose o' as parameter of the label function, then f(x; 0) + f(;; 0') means the loss function C(x; 0), have:
dHAH(A, B) sup Pr (e>xA(((x;0)) Prz~B(((x;;0) 0)]
0-
sup (e) - Eb(O)] sup (0) VEB(()) S - dgoo(A, B). (29) eg@
Here, dgo0( (A, B) the gradient divergence, given the model and (a) holds due to Assumption 2
F.4 PROOF ON THEOREM 2
Let D. be the sampled counterpart from the domain D.,. we have E is an empirical risk of Du, i.e., Ep. = 1/N.. ΣN L([(1) 0), v(). We also have expected risk &D. defined as &D, = B(a.yED.)) ((teEE)))(((((;;;))) y)]. For a given 0 E O, with the definition of generalization bound, the following inequality holds with at most Us for each domain D., (Us is the number of users, which is also the number of domains).
E._ (0) Ep.) (0) log M + log Us/8
(30)
2N.
Moreover, from Lemma we have 1Ep, (0) ED, (0) S AdNAH (Dw, DT) for each domain. Then let us consider Eq. (30). we can obtain the following inequalities with the probability at least greater than 8 U.
log M + log Us/8
min Ep, (0)) E.. (0) + d'
2N,
log M + log Us/8
< (0)
(8) + 2
dHaM (D., Dr) +
(31)
2N,
24

---

Published as a conference paper at ICLR 2025
We denote the local optimal each client of source set u, Us as 8.. If we choose specific parameter 7 = ming E (0) which is the local optimal on the unseen domain T. the above third rewrite the above inequalities into:
inequality still holds. Then,
E, (0:) & ED, (6,) dnan(Du, (T) log M + log U
(32)
2.N.
Considering on each domain, Eq. (2 holds. By similar derivation process, can obtain the inequality between T and D with the probability at least greater than 1 - 5.
Σ E. (07) (05) +
E Th
(dnnn(Dnn DT) +
log M + log Us/ô
(33)
2N,
HElls
MELls
Combining the Eq. FL. For instance,
33 and Objective
we have Theorem 2 with the global model after R rounds
ED, (B EDr(07m
&D. (0) - -Ep((0)) + dHAH DT) + log M log 3
log M + log Us 0
< Σ
+
V2N₂
22N
uEldg
Σ Yu Ep_ (8) dHAH((, DT) + log 0
Usa v log
6
+ +
(34)
V2N.
wElls
To further analyze the convergence bound, we consider the Assumption
3 For instance,
ED, (BR) EDy (0
Σ Yu E (0) + +dHAH((D, Dr) + V
log M 0
v llg UsM
+ 5
(35)
v2N₂
wElds
M V Vlog
5 + log
Ug 8
Σ Yu E (0) + dHAH((D, Ds) + dHAH(Ds, DT)
+ +
V2Nu
uElls
(b)
dHAH(Du. D.)
M log o
+ /log
M 8
Σ Tu E (0) + Σ
+ dHAAHDDs
+
V2N,
wells
vElls
We have (b) holds due to Lemma Applying Lemma 3 we have:
Ed, (BR) EDy (0b,))
dgoo dgoo(Du.
(D.. De)
V
log M 8
+ v log UsM
Σ Tu ED, (0) + Σ
+ (Ds DT)
+
1
V2NNa
WEMls
vells
25

---

Published as a conference paper at ICLR 2025
F.5 PROOF ON THEOREM 1
To find the optimal solution of invariant gradient direction DiGD in Eq. 11 we consider x VIGD and
(r)
consider the maximization problem with optimization variable. Denote ø x 199. 2 Note
(r)
that min, (gu 9FL min,
Eu Valk (r))) The Lagrangian of the objective is
A (r) 2 1lof - x11² + 2 s.t. A
max min Aiy
Σ 1a x Es
(36)
Since the problem a convex programming and and Slater's condition is satisfied for x > (meanwhile, if к 0, it can be easily verified that all results hold trivially), the strong duality holds. Consequently, the order of the min and max operations can be interchanged. For instance.
min max A,
Σ Tull uEuls
2 Ilooi =112
2 Φ s.t. A20 0.
(37)
Taking A₁ into consideration. If we consider X, y as constant, x the variable, x achieves the optimal solution when 0A1/0x- = 0. Another speaking, we have A(x - ((4)) - Σv= Vuh or specifically,
DEL
(r)
Tulh(()) /2.
(38)
Therefore, we have the followings:
A =
) T
9FL +
Σ Tuhl 12)
11 (9fl +
(E Tahl /2)
112
2
u=1
)uh(?)) T
9r..
(r)
> yuhl!") (1)
/x)
\ 1 Σ Yah
A
2 Φ
(39)
Substituting 91 (r) Eum Yaha (r) Consider the optimization problem of Eq.
33
we have:
A1 = = (r)T
9FL + 91
/2)
|lof
/A12
+
2
= =66 9F1.
(T) (r)
+ 9
1 (r)
1
|lf
A 2
21
Dr DAL
TT) (r)
+ loo
1
11 +
z0.
X
(40)
2X
Therefore, have Eq. 37 is equivalent
(r)T gr
(r) 112
min Ary MLL
+
+ 2
(41)
2x
Next, we consider À as variable to find the optimal value. Subsequently, optimization problem Eq. 41 is equivalent to the following relationship:
0
A2 =
llo 112 -0 0. 2xa
(42)
0x
Therefore, the equation achieves the optimality as A 11 1/2 Combining with Eq. 41 and Eq. 38 have the followings:
SIGD
DFL
K
(r) s.t. | = arg min (r) D1
gf DFL + k x
(1)
(43)
1
This solves the problem.
26


Reviews:
[{"review_id": "LaVpRMvLBj", "reviewer": "Reviewer_XXzG", "summary": "The paper proposes a gradient-matching strategy on the server for federated domain generalization (FDG) using a meta-learning approach.\nThis essentially boils down to a convex combination update rule instead of a simple average as in FedAvg.\nIn practice, the optimal combination of local updates is found by bounding the optimization near a simpler update rule like FedAvg update rule.\nThe paper provides some theoretic results and compares to some baseline methods in several datasets showing performance improvement empirically and that the method can be combined with other methods.", "questions": "- Why do you formulate this as a bi-level optimization problem? Why is this necessary? It is important to lead the reader up to this point rather than just stating it. Furthermore, this is not even a bi-level problem in 3a and 3b. This is just 2 equations. Are you minimiizing 3a subject to 3b? If so, it's not even clear that 3b is minimization problem, it's just a constraint perhaps? This is incorrectly formalized.\n\n- What is the intuition between the difference between yours and [Shi et al., 2022]? Is it something like the some of inner pairwise products is bounded by the sum of inner products between a mean vector and each vector? Does this have relationship to standard sum of squares ideas?\n\n- How is the update rule for your method different than FedAvg? It seems like a different aggregation method but is difficult to understand how it is different than simple FedAvg. Could you provide an more explicit comparison and discussion on how it differs? When and why would the weights be different than FedAvg and when would the be equivalent?\n\n\n\n\n- Other than computational, is there a reason to limit the search space to a convex combination of local gradients? This doesn't seem necessary.\n\n- Why is a search space limitation needed? This is not well-motivated but just stated as fact. It seems that constraining to a convex combination + constraining to be near a simpler method like FedAvg strongly regularizes the method. But it is not clear why this is necessary or justified, perhaps other than an empirical argument.\n\n- Eq 11 is not a multi-objective optimization problem as it is written. Essentially it has been reduced to a single objective via scalarization $\\gamma$ parameter. Thus, it is not clear why multi-objective optimization is needed or required. \n\n- Eq 11 - Why is $\\kappa$ needed? It seems that this term does not depend on $\\Gamma$ and thus can be ignored.\n\n- There is no lead up to Theorem 1 and this theorem is not well-explained. Why is this an easier problem? Why did you need to go through the min-max problem setup? This whole derivation seems very convoluted and does not lead logically to the next step. A major revision and explanation is needed.\n\n- The theoretic analysis seems to use the same basic tools and techniques from prior works. Could you briefly explain the similarity of each main lemma, theorem or corollary w.r.t. its prior (probably non-FL) counterparts? Are there any new theoretic techniques used?\n\n*Summary of Review*\nThe basic idea of doing gradient matching on the server seems natural and the paper proposes one practical way to implement this. Furthermore, the empirical results are fairly strong either by itself or in combination with other methods. However, the method is not explained well and the theoretic analysis seems very similar to prior work (useful but not itself much of a contribution).  The lack of comparison to a Federated DG benchmark paper from last year's ICLR also calls the results into question.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Proposes an new aggregation method to improve federated domain generalization based on gradient matching. \n\n- Provides some theoretic analysis on the proposed method.\n\n- The empirical results show a significant improvement over baselines used and demonstrate that FedOMG can be combined with other methods for a good performance boost (caveat with not comparing to a known Federated DG benchmark paper, see Weaknesses).", "weaknesses": "- The paper does not compare to the federated domain generalization benchmark [Bai et al. 2024] that exactly fits this setting. This benchmark paper enables comparison of your method to multiple prior DG methods adapted from central methods and federated DG methods. In particular, Bai et al. [2024] noticed that simple federated averaging or convergence-specific FL methods could beat most prior Federated DG methods. Bai et al. [2024] also control for hyperparameter tuning costs. I would like to see this method applied within this framework to compare to the methods in that benchmark. This will place FedOMG on the same playing field as prior methods and would only require wrapping the method so that it is compatible with the framework. (Additionally, Bai et al. [2024] noticed that the PACS dataset is quite a bit different from other DG datasets so seeing on other datasets in the benchmark would be helpful.)\n\n- The explanation of the method and the key insights are not written well. Several parts are simply incorrect though I can almost guess at what is meant. Others may be correct but are not explained well or justified appropriately. Please see questions below.\n\n[Bai et al., 2024] Bai, R., Bagchi, S., & Inouye, D. I. (2024). Benchmarking algorithms for federated domain generalization. ICLR.", "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": "", "extra_scores": {}}, {"review_id": "wQwFfP8XVb", "reviewer": "Reviewer_9XVR", "summary": "This study introduces Federated Learning via On-server Matching Gradient (FedOMG) for Federated Domain Generalization (FDG). Unlike traditional methods, FedOMG leverages local gradients to identify an invariant gradient direction across domains, enabling efficient feature aggregation on a central server without extra communication costs. It is also compatible with existing FL/FDG methods for enhanced performance.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\tThis paper propose a method that is effective and highly compatible with existing FL algorithms and a large number of valid experiments have been conducted on the empirical side.\n2.\tAlthough I'm not an expert in this area (FL combined with DG), I think it's a very interesting work, benefiting from the theoretical explanation of how the authors derive the final method step by step, and the theoretical analysis seems solid enough.", "weaknesses": "1. What confuses me is the explanation of the motivation for this new approach (Section 3.2), and I would appreciate it if the authors could explain this part in more detail.\n2. In lines 203-215, the authors seem to consider an alternative gradient for solving the $M$ -dimensional optimization, and I would like to know if this approach is a rough estimate, and if so, can you discuss the implications in detail \n3. a formulation “any FL algorithm” was used in line 267 - 269, I would have expected the authors to mention a great deal of relevant work here to demonstrate this overly certain claim", "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": "", "extra_scores": {}}, {"review_id": "duMHDcllw2", "reviewer": "Reviewer_HBia", "summary": "This paper introduces FedOMG to address the challenge of Federated Domain Generalization. The core idea behind FedOMG is to leverage local gradients from distributed models as domain-specific information. FedOMG maximizes the inner product of gradients to ensure that the model finds a gradient direction that is invariant across all domains. Extensive experiments conducted to evaluate the effectiness of the proposed model under federated setup.", "questions": "1. How does on-server gradient alignment directly contribute to reducing domain-specific biases without compromising the generalization capabilities of the global model?\n\n2. Could you clarify how using gradients alone (instead of data) ensures privacy in FDG? Are there specific privacy-preserving guarantees that FedOMG provides through gradient-only use? What limitations, if any, exist in relying solely on gradient matching for privacy preservation, especially with highly heterogeneous client data?\n\n3. What motivated the choice of complex optimization techniques like convex optimization and Pareto optimality in this context? How do these methods specifically enhance FedOMG’s performance over simpler approaches? Could simpler optimization approaches potentially achieve similar outcomes, or are the proposed techniques essential to achieving the method's goals?\n\n4. Could you provide a more intuitive explanation of how gradient direction alignment helps achieve domain invariance across clients?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Theoretical analysis propoded which seems to be true.", "weaknesses": "1. The approach to invariant gradient direction through convex optimization and Pareto optimality increases the computational complexity.\n\n2. The method’s effectiveness relies on the assumption that domain gradients will align under the invariant gradient matching, which may not hold well in highly heterogeneous data settings.\n\n3. The connection between the stated motivations (privacy, communication efficiency, and domain generalization) and the chosen methodology (on-server gradient matching and invariant gradient direction) is not clear.", "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": "", "extra_scores": {}}, {"review_id": "aITo08Usqr", "reviewer": "Reviewer_cQAL", "summary": "This paper introduces FedOMG, a method that efficiently leverages domain information from distributed domains to prove the performance of the federated domain generalization (FDG) problem. The authors propose an approach that finds an invariant gradient direction across all domains through gradient inner product maximization to achieve this. The authors take the FDG problem as a multi-objective optimization problem and optimize it by finding the Pareto front. The extensive experiment results show the strengths and effectiveness of FedOMG, which is both superior in performance and computationally efficient.", "questions": "- Could the authors provide more intuitions behind designing the search space (why search in a ball and use L2 norm)? Also, why are you using FedAvg for g_FL here? Would other federated aggregation rules for non-iid data work better here?\n- Also, the authors use a convex combination of local gradients to find the invariant gradient direction. Could you also provide more explanations and intuitions for the design behind that?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea is presented with solid theory, which is generally easy to follow.\n-  Framing the FDG problem as a multi-objective optimization problem using gradient matching is interesting and intuitive. Also, it aligns with the nature of Federated learning with optimization goals.\n- The results are strong across many datasets and baselines, and FedOMG is compatible with other related baselines. Many ablation studies are presented to show the effectiveness of the proposed method.", "weaknesses": "- **The presentation of the Introduction and Related work**: The authors mention some previous work and I think it would be good to discuss their differences, limitations, and connections, compared with FedOMG. Also, in the introduction, the authors can discuss more on the novelty and design intuitions/details of FedOMG. Now it reads only mentioning gradient matching, which is too general from my perspective. Presentation-wise, this paper can be improved greatly in my opinion.\n\n- **About theory**: In Objective 2, Theorem 2, and Corollary 3, some variables are not defined, such as µ, M, and δ. Also, it would be good if the authors could provide more explanations on how to derive Corollary 2.\n\n- **Toy tasks**: The figure shows too many trajectories and is a little bit confusing. The authors can think of ways to present it conveying the conclusions clearer and more straightforward.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)