W3 & Q3

W3: This paper does not include an analysis of hyperparameter selection, which is important for SAM-based approaches.

Q3: A sensitivity analysis of key hyperparameters, such as the radius γ, would add valuable insight.

Thanks for your valuable advice. In the new revision, we provide Parameter analysis in Section 4.6.

Major points related to the selection of $\gamma$ is concluded as follows:

• The selection of $\gamma$ is based on the grid search: the optimal value of $\gamma$ varies across tasks, datasets, and models for SAM-based methods. As a result, we conducted a grid search for $\gamma$ on the PACS dataset to determine the appropriate value for both FedSAM and GFM. The local optimum ($\gamma=0.02$) exists for both methods.
• GFM exhibits a relatively flatter optimum compared to FedSAM: this could be attributed to the improved consistency of local models in GFM, which reduces the need for local flatness to achieve sufficient global flatness.

Q4

Q4: The work by Sun et al. [1] also explores global flatness in heterogeneous FL using an ADMM algorithm to address constrained optimization. Including a citation of this work and a comparison would enhance the paper.

In this revision, we cited [1] in the Introduction Section and compared to it in Table 1.

W1 & Q1

W1: The min-max optimization process in the local updates remains somewhat unclear. Further clarification would enhance understanding of the authors’ claim. Q1: Could the authors provide a more detailed analysis of the min-max process? For instance, analyzing the solutions of $\phi_i$ and $\theta_i$ would improve understanding of this component.

Thanks for your valuable advice. We provide a more detailed analysis of the min-max process in the new revision as shown in Appendix F, including:

• The optimal solution of $\phi_i$: we show that the augmentation model aims to generate samples that exhibit the largest prediction discrepancy with respect to the true label. In the subsequent minimization step, $\theta_i$ seeks to improve its performance on these challenging samples. Through this process, $\theta_i$ is guided to align its behavior with $\theta$.
• Existence of the desirable saddle point: we prove the existence of the saddle point in the min-max process in Proposition 1. It can be inferred that the saddle point solution in Proposition 1 is desirable because it achieves comparable global performance to the global model $\theta$, as demonstrated by $p(y|x;\theta_i^*) = s \cdot p(y|x;\theta)$. In this way, the local update can be effectively supplemented with global information as stated, leveraging both the global model and the augmentation model.

Proposition 1: A saddle point solution exists for the min-max problem in Eq. (14). Construct $\theta_i^*$ such that $p(y|x;\theta_i^*) = s \cdot p(y|x;\theta)$ for any $x$ in the support set with true label $y$, where $s = \frac{1}{\max_x p(y|x;\theta)}$. Then, there exists $\phi_i^*$ such that $\theta_i^*$ is the local minimum of $\mathcal{E}_{a(D_i;\phi_i^*)}(\theta_i^*)$. Consequently, the pair $(\theta_i^*, \phi_i^*)$ constitutes a saddle point solution, satisfying $\Delta \mathcal{E}(\theta_i^*,\phi_i)\leq\Delta \mathcal{E}(\theta_i^*,\phi_i^*)\leq \Delta \mathcal{E}(\theta_i,\phi_i^*)$.

W2 & Q2

W2: This paper lacks direct validation to demonstrate whether the augmented data provides a more effective surrogate for global data compared to local data. A more focused assessment on this point would strengthen the findings.

Q2: Regarding point 2 in Weaknesses: Does the augmented data act as a better surrogate for global data than local data? A more direct evaluation would be helpful.

To provide a more direct evaluation, we focus on the trained model after local updates. As shown in Proposition 1, the trained model should perform similar to the global model. Consequently, we evaluate the effect of augmented data and local data by measuring the mean forgetting rate between the trained model and the global model on the global dataset $\hat{D}$ as shown in Appendix C.2. There are two major conclusions:

• It is evident that GCA adopted in GFM effectively ameliorates the situation of forgetting.
• The amelioration is more evident in the later stages of training when the global model becomes stronger.

W1

Assumption 1 of the paper mentions that local models may outperform global models is COUNTER-INTUITIVE, and I think this hypothesis can be illustrative to a certain extent, but it does not cover so the federated learning situation. Because the local model is stronger than the global model in the federated learning chaotic situation is not a completely impossible situation.

This is a very good point. In certain chaotic federated learning scenarios, it is indeed possible for a local model to outperform the global model. For example, this could occur if a local client has a large amount of data with a distribution similar to the global data, while other clients have limited data or significantly differing distributions.

However, we mainly focus on the federated domain generalization (FedDG) setting, where each node contains data from a single domain with sufficient volume. In this context, it is almost impossible for a model trained on one domain to perform well across all domains. As such, we believe that Assumption 1 holds under the FedDG setting presented in the paper. For scenarios where the assumption might not hold, caution should be exercised when applying our method, as the global model might impose incorrect constraints in such cases.

We have revised Assumption 1 to include additional constraints, as follows:

Assumption 1. If (1) data distributions $\{D_i\}\_{i=1}^{M_S}$ across clients exhibit a non-trivial degree of heterogeneity, and (2) each client has access to a sufficiently large local dataset to estimate the data distribution. Then during the training phase, local models $\{\theta_i\}\_{i=1}^{M_S}$ and their aggregate $\sum_i p_i \theta_i$​, when weighted by coefficients specific to clients, satisfy the following inequality:

$$\hat{\mathcal{E}}\_D(\sum_i p_i \theta_i)\leq\sum_ip_i\hat{\mathcal{E}}\_D(\theta_i),$$

where $p_i$ represents the coefficient of client $i$.

Q1

It is not particularly clear to me what is the difference between the paper's approach and the client-side optimization of local models using SAM-like methods?

Thank you for your question. One key conclusion from Theorem 2 is that the performance on unseen domains is directly tied to the local flatness of seen clients on augmented samples. This suggests that instead of designing a specific global flatness strategy, we can effectively combine any existing local flatness minimizer with the proposed GCA.

As a result, our algorithm integrates a client-side flatness optimization method with the proposed GCA. The presence or absence of GCA is the critical distinction between GFM and purely client-side optimization methods. Moreover, GFM is designed to achieve global flat minima, whereas local flatness methods like SAM focus only on finding local flat minima.

Q2

As mentioned in the article, such methods require more power on the client side. The computational paper on how to reduce client-side computation doesn't go much further than that.

In the revised version, we further optimized training speed to reduce computational costs on the client side by implementing the following strategies:

• Parallel Augmentation: Previously, the basic augmentation in GFM was implemented serially. In the revised version, we have optimized this by implementing parallel augmentation.

• Disable Gradients for Classification Model Parameters During Augmentation Updates: Gradients for the classification model's parameters are now disabled when updating the augmentation model, reducing computational overhead.

As a result of these improvements, the augmentation cost has been reduced from 1.36V to 1.16V (with the total cost reduced from 3.06V to 2.86V), taking the cost of FedSAM (1.7V) as the baseline cost for the classification model. To further reduce the computational cost, we use subsampled low-dimension images as the inputs to the augmentation model. The resulting color and affine parameters are then applied to the original images to generate the augmented images, as discussed in Appendix C.5.

Fortunately, the final performance is not significantly impacted by the input image dimensions used in the augmentation model. By adopting this strategy, the augmentation cost reduces from 1.16V to 0.6 V, which only accounts for 35% of the classification cost. Moreover, as shown in Table 7, among the methods designed to achieve flatter minima for the global model (including FedGAMMA, FedSMOO, and GFM), our method imposes the lowest space constraints.

Q3

Is it right to infer that Theorem 1 is a restatement and not an original contribution?

Yes, it is right. We stated it in the Line 136 of the paper.

Q5

Q5.1: Can c>e? How do the authors choose the values of these hyperparameters? It is quite important to understand this especially since this is an adversarial optimization that is often unstable under the slightest change of hyperparameters.

$c$ refers to the update interval, and $e$ refers to the iterations of local updates. $c>e$ means during the whole local training, the augmentation model won't be updated, which we don't suggest. To understand the selection of $c$, we conduct parameter analysis in Section 4.6 of the revised version. We briefly introduce its conclusions as follows:

• The performance improves with more frequent updates of the augmentation model: the strategy of alternating one iteration of augmentation with one iteration of classification achieves the best performance.
• To balance performance and efficiency, we set $c = 10$ for related experiments.

Q5.2: How many models need to be maintained by a client to run this optimization?

Each client needs to maintain the received global model $\theta^{r-1}$, the local model $\theta^r_i$, and the augmentation model $\phi_i$.

Q5.3: Could the authors expand line 11 to include loop oversampling, gradient estimation, and update rule? Are the ϕis also aggregated?

Sure, and we apologize for the error in line 11.In practice, we update $\phi_i$ as shown on the right-hand side of Eq. (14), rather than Eq. (9). Additionally, line 11 does not include loop oversampling; instead, $\phi_i$ is updated solely based on the current batch $X_i$ by augmenting it and feeding it through the prediction model to compute gradients. In the revised version, line 11 has been expanded to lines 11–12 to provide more detail.

​ 11: Compute $g_{\phi_i}=\nabla_{\phi_i} \hat{\mathcal{E}}_{a(D_i;\phi_i)}(\theta_i) - \hat{\mathcal{E}}_{a(D_i;\phi_i)}(\theta)$ on $X_i$

​ 12: Update $\phi_i=\phi_i+\rho_\phi g_{\phi_i}$ Update augmentation model $\phi_i$

$\{\phi_i\}$ do not need to be aggregated, although doing so is feasible. Aggregating them would only affect the initialization of $\phi_i$ in each round.

Q6

Authors mention that the SAM optimizer uses $\gamma=0.02$. How do they arrive at this value? Could they talk about how to interpret this value or how it might impact the training?

The radius $\gamma$ is a critical hyperparameter in SAM-based methods, as it determines the range of model perturbation. The optimal value of $\gamma$ varies across tasks, datasets, and models. So we perform a grid search to determine the appropriate value of $\gamma$ as shown in Section 4.6.

Q1

Authors mention that GFM = FedAvg + SAM + GCA. However, FedSAM is FedAvg + SAM. Could the authors clarify their contribution with respect to FedSAM?

The focus of the paper is to seek flatter minima of the global model. As demonstrated in Theorem 2, we say the performance on the unseen domain is directly related to the local flatness of seen clients on augmented samples. This insight suggests that there is no need to design a dedicated global flatness strategy; instead, we can combine any existing local flatness optimizer with the proposed GCA, which we consider a key strength and contribution of our work. Therefore, to make the overall algorithm as simple as possible, we adopt the well-known SAM optimizer in local update.

To further demonstrate, we show in the following that our method could combine with another local flatness optimizer. Specifically, we use SWAD [5] to replace SAM in local training ($n_{converge}=5,n_{tolerance}=6$ and the ratio of tolerance as $0.2$). It is worth noting that, due to time constraints, we have not extensively tuned the hyperparameters, and the results may be suboptimal.

To summary, we clarify the contributions as follows:

• Though the SAM component is shared between GFM and FedSAM, the focus differs significantly. GFM emphasizes achieving flatness for the global model, whereas FedSAM is concerned only with local flatness. (It is also worth noting that SAM was not originally proposed in FedSAM.)
• Our work can be combined with other local flatness optimizer to take place of SAM, like SWAD.
• Finally, we consider the compatibility of reusing established components in addressing new challenges, with accompanying theoretical guarantees, as an important contribution of our work.

Reference

[5] Cha, Junbum, et al. "Swad: Domain generalization by seeking flat minima." Advances in Neural Information Processing Systems 34 (2021): 22405-22418.

Q2

Currently the authors are working in the restrictive setting that each client has one domain each. How do they see their method adapting if each client has access to more than one but not all the domains?

Theoretically, the applicability of our method primarily depends on whether Assumption 1 is satisfied. In the mentioned case, Assumption 1 remains valid because, in principle, a local model trained on a subset of domains cannot outperform the global model when evaluated across all domains. To further validate it, we conducted experiments on PACS dataset. In this setting, when testing on the photo domain, $A,C,S$ datasets are randomly split into $A_1,A_2,C_1,C_2,S_1,S_2$, which were then recombined as $(A_1, C_1)$, $(A_2, S_1)$, and $(C_2, S_2)$ to represent different clients.

The results above indicate that the mentioned scenario is analogous to the single-domain case.

W4

Apart from comparing the methods based on performance that authors should also present the comparison in terms of computational and space complexity, e.g., plot comparing FLOPs vs performance or number of parameters shared vs performance, etc.

Thank you for the valuable advice. In the original paper, we compared the running time and space cost of GFM and FedAvg. Following your suggestion, we further refined the table and included three metrics—running time, space cost of local updates (former element in the space column), and the parameters required for sharing (latter element in the space column) —to evaluate each method in Table 7 of the revised manuscript.

From these results, we can conclude the following:

• The additional parameters introduced by the augmentation model are relatively small.
• The actual running time increases to approximately 1.5 to 1.7 times that of the FedSAM baseline.
• Among methods (FedGAMMA, FedSMOO, and GFM) designed to achieve flatter minima for the global model, our method imposes the lowest space constraints.

W2

Also from table 2, it can be seen that FedAvg+GCA wins in two of the 4 cases bringing up the next question: What is the impact of data augmentation in general when compared to the adversarial data-augmentation that the author's propose.

We would like to clarify that in the following table. We has conducted comparisons with several data augmentation methods: RA (RandAugment [1]), AA (AutoAugment [2]), and Cutout [3]. These methods represent widely-used general data augmentation strategies for images. Notably, RA and AA consist of color and geometric transformations, similar to GCA.

These results clearly demonstrate the effectiveness of GCA in the context of federated domain generalization. For the PACS dataset, we further illustrate that our methods are compatible with these general augmentation strategies, as shown below:

Reference

[1] Cubuk, Ekin D., et al. "Randaugment: Practical automated data augmentation with a reduced search space." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops. 2020.

[2] Cubuk, Ekin D., et al. "Autoaugment: Learning augmentation strategies from data." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.

[3] DeVries, Terrance. "Improved Regularization of Convolutional Neural Networks with Cutout." arXiv preprint arXiv:1708.04552 (2017).

W3.1 & W3.2 & Q4

W3.1: The statement "aggregation helps generalization." in line 191 needs either proof or a citation from prior work that proves it. While there is a citation in the paper that advocates that generalization can be achieved via flatness, one cannot infer that flatness is the necessary condition for generalization".

Q4: Could the authors mention some citations that use a similar assumption as "assuming aggregation helps generalization" on lines 191-192?

The statement "aggregation helps generalization" refers to Assumption 1 in our paper (line 187). If I understand your concern correctly, you are asking whether flatness contributes to generalization. There is an insightful study [4] in the literature that discusses the relationship between flatness, generalization, and SAM across various architectures and data distributions. We briefly introduce its conclusions below:

• Flatness is neither a necessary nor sufficient condition for generalization: For instance, sharp minima can sometimes generalize well, while flat minima may generalize poorly.

• Situations fall into three regimes: (1) Flattest minimizers of training loss provably generalize and sharpness minimization algorithms find generalizable models. (2) There exists a flattest minimizer that generalizes poorly. Also, sharpness minimization algorithms fail to find generalizable models. (3) There exist flattest minimizers that do not generalize but the sharpness minimization algorithm still finds the generalizable flattest model empirically.

While theoretical results cannot guarantee that "flatness helps generalization" under all architectures and data distributions, empirical evidence demonstrates its effectiveness across a variety of tasks and model architectures. Therefore, pursuing flatness remains a reasonable strategy to achieve better generalization in the context of the FedDG task.

W3.2: Secondly the authors need to prove that this statement is further not affected by the relaxation in assumption 1.

Assumption 1 pertains to the relationship between models before and after aggregation in the FedDG setting and is not directly related to the concept of flatness. Therefore, we believe it is not relevant to the statement "flatness helps generalization".

Reference

[4] Wen, Kaiyue, Zhiyuan Li, and Tengyu Ma. "Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization." Advances in Neural Information Processing Systems 36 (2024).

W1 & W3.3

W1: The ablation study in Table 2 shows that the main benefits are achieved by using SAM and/or GCA. This conflicts the key contribution of this work, i.e., GFM. Particularly, the authors mention "combining them can complement each other’s strengths" in line 516, but looking at the table the same cannot be inferred.

W3.3: Overall, while it irrefutable that authors are able to generate improved performances, it is not convincing that the reason for improvements can be attributed totally to flatness of the minima of the global model.

Thank you for your careful review and constructive feedback. We understand your concerns and acknowledge that the ablation study in Table 2 may not fully illustrate the effectiveness of GFM. In GFM, global flatness is highly coupled with the proposed GCA component. Additionally, as an augmentation strategy, GCA can influence generalization performance, further complicating the analysis of global flatness. To address this, we have conducted extensive ablation studies on GFM across a wide range of datasets, providing deeper insights into global flatness, as detailed in Section 4.5 of the revised manuscript.

Combining the new table with the original one, we have the following observations to answer your question:

• The effectiveness of global flatness is evident in some cases: it's worth noting that, in the OfficeHome and TerraInc datasets, GCA alone does not enhance generalization performance, which underscores the importance and effectiveness of the claimed improved flatness of the global model.
• GCA or FedSAM alone can achieve the best performance in certain domains of specific datasets: This typically occurs when either flatness or augmentation negatively impacts generalization performance. In such cases, removing the problematic component yields the best results. While the underlying causes of these failure modes require further investigation, it is encouraging to note that GFM consistently improves the performance in the average meaning.

Regarding the statement "combining them can complement each other’s strengths," we intended to convey that when one component underperforms, the other compensates, resulting in an overall performance improvement compared to the baseline without both of them. However, to avoid potential confusion, we have removed this statement from the revised manuscript.

W3

The organization of this paper could be improved. The related works and preliminaries sections span two pages, while the experimental section only covers three pages. The experimental section should be expanded with more detail in the main text.

Thank you for your kind suggestion. We will revise the related works section by moving content related flatness in federated learning setting to the introduction Section, placing other parts in the Appendix A. In the meanwhile, we add more experiments in the experimental section.

W4

It would be better to add hyperlinks to the citations for the methods listed in all the tables. Currently, all methods are cited only in Section 5.2, which makes it difficult to link each method to its corresponding paper.

Thanks for your suggestion. We have added hyperlinks in Table 1.

W5

The definition of flatness in Equation 3 is based on the assumption that θ is close to a local minimum, meaning the first-order gradient is near zero. There is no such assumption around equation 3.

Thanks for your comment. We apologize for the imprecise expression. $\mathcal{E}^\gamma(\theta)$ is not a direct measurement when $\theta$ is not close to a local minimum but is a promising objective to seek flat minima, as used in [2]. So we adopt $\mathcal{E}^\gamma(\theta)$ as our objective in the paper. We have modified the corresponding expression in this revision.

[2] Foret, Pierre, et al. "Sharpness-aware minimization for efficiently improving generalization." arXiv preprint arXiv:2010.01412 (2020).

Q1

How is the loss landscape illustrated in Figure 2? What method is used to generate the loss landscape in this figure?

We use a visualization technique similar to that described in Garipov et al. (2018). Specifically, three local models are used to define the averaging plane, where each point represents a set of model parameters. The value at each point corresponds to the empirical risk of the associated model. By setting the global model and model $a$ as X-axis and model $b$ in the upper of X, the plot is generated.

Q3

Did the authors conduct experiments on CIFAR-10 or CIFAR-100, as some baselines reported results on these datasets?

Because the focus of our paper is federated domain generalization, we didn't conduct experiments on CIFAR-10 or CIFAR-100 which only involves a single domain. We introduced these baselines of heterogeneous FL for providing a broader range of references on FedDG problem.

Q4

The authors mentioned 'However, in federated learning, this becomes challenging due to the privacy concerns of the data.', which seems to be one of the motivations to derive the proposed GFM. What is the difference of the privacy protection between the GFM and FedSAM?

There is no difference of privacy protection between GFM and FedSAM and both methods adhere to the principle of data privacy protection. We mentioned 'However, in federated learning, this becomes challenging due to the privacy concerns of the data.' to indicate that it's challenging to seek global flatness in the federated learning setting. This statement applies equally to FedSAM. Thus, FedSAM seeks local flatness to avoid violating this requirement, while GFM tried to seek global flatness by virtue of GCA, which adhere the principle as well.

W1 & Q2

W1.1: The proposed method is very similar to the baseline, FedSAM; however, it has not been discussed in either the related works section or the introduction.

Since there are relatively few studies in the literature focusing on seeking flat minima under federated learning settings, we did not create a separate subsection in the related work section. Instead, we integrated this discussion into the latter part of the flat minima subsection, which might have made it less noticeable to readers. In the revised version, we have clarified this part and FedSAM in the third paragraph of the Introduction section to improve visibility and understanding.

W1.2: In comparing the proposed method to FedSAM, I do not see a clear improvement or significant difference between them.

Regarding the relationship with FedSAM, we do not view our method as merely a refinement of the specific FedSAM approach. As demonstrated in Theorem 2, the performance on the unseen domain is directly related to the local flatness of seen clients on augmented samples. This insight suggests that instead of designing a specific global flatness strategy, we can combine any existing local flatness optimizer with the proposed GCA, which we consider a key strength and contribution of our work. We selected SAM primarily due to its simplicity and widespread adoption. It is also worth noting that the SAM optimization scheme was not originally proposed by FedSAM, which explains why we did not specifically emphasize FedSAM in the paper.

We show in the following that our method could combine with another local flatness optimizer. Specifically, we use SWAD [1] to replace SAM in local training ($n_{converge}=5,n_{tolerance}=6$ and the ratio of tolerance as $0.2$).

To summary, the difference between two methods are as follows:

• While both GFM and FedSAM share the SAM component, their focuses are fundamentally different. GFM emphasizes achieving flatness for the global model, whereas FedSAM focuses solely on local flatness.
• Our proposed augmentation strategy (GCA) can be flexibly combined with other local flatness optimizers, such as SWAD, in place of SAM.

W1.3: The only notable difference appears to be the augmentation network discussed around Equation 9, which is directly adopted from Suzuki (2022).

About the augmentation network, though the network structure was proposed in Suzuki (2022), we only borrowed its network structure (similar to borrowing ResNet for classification tasks). The global model-constrained adversarial data augmentation objective (Eq. (14)), however, is specifically designed for the global flatness objective. Therefore, we believe this does not overshadow our contribution.

W1.4: Additionally, the ablation study does not evaluate the performance impact of removing this augmentation network.

Q2: What would the performance be if FedSAM were combined with the augmentation network? Additionally, is there an ablation study on the augmentation network's impact within GFM?

When the augmentation network is removed, our method reduces to FedAvg+SAM, which is essentially FedSAM as discussed in line 480. What's more, results of GCA in the ablation study reveal the augmentation network's impact within GFM.

Reference

[1] Cha, Junbum, et al. "Swad: Domain generalization by seeking flat minima." Advances in Neural Information Processing Systems 34 (2021): 22405-22418.

W2

What is the primary motivation of this paper? In the introduction, the authors discuss issues related to finding flat minima and the neglect of data privacy. However, there does not appear to be any specific design or consideration aimed at enhancing privacy protection in this paper.

The primary motivation is to seek flat minima of the global model to improve domain generalization ability. To avoid violating data protection principles, we adopted the GCA component to generate a surrogate for the desired global data distribution in each local clients.