Seeking Global Flat Minima in Federated Domain Generalization via Constrained Adversarial Augmentation

Authors: Chenyang Wang, Junjun Jiang, Xingyu Hu, Xianming Liu
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Federated Domain Generalization, Flat Minima, Data Augmentation

Abstract:
Federated domain generalization (FedDG) aims at equipping the federally trained model with the domain generalization ability when the model meets new clients with domain shifts. Among factors that possibly indicate generalization, the loss landscape flatness of the trained model is an intuitive, viable, and widely studied one. However, pursuing the flatness of the global model in the FedDG setting is not trivial due to the restriction to preserve data privacy. To address this issue, we propose GFM, a novel algorithm designed to seek Global Flat Minima of the global model. Specifically, GFM leverages a global model-constrained adversarial data augmentation strategy, creating a surrogate for global data within each local client, which allows for split sharpness-aware minimization to approach global flat minima. GFM is compatible with federated learning without compromising data privacy restrictions, and theoretical analysis further supports its rationality by demonstrating that the objective of GFM serves as an upper bound on the robust risk of the global model on global data distribution. Extensive experiments on multiple FedDG benchmarks demonstrate that GFM consistently outperforms previous FedDG and federated learning approaches.

Decision: reject
Reviewer scores: 3, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730271878739

thread_type:
official_review

<thread_object>
creator: reviewer_itls

<content>
summary:
The paper addresses the challenge of federated domain generalization (FedDG), where a federated model must generalize effectively to new clients facing domain shifts. Authors introduce GFM, an algorithm designed to achieve global flat minima during the federated training of global model. GFM uses a globally-constrained adversarial data augmentation technique, which generates surrogate global data within each local client, improving the model’s flatness while preserving data privacy. Supported by theoretical analysis, this approach ensures that GFM effectively finds global flat minima. Extensive experiments further highlight GFM’s superior performance over previous FL and FedDG methods across multiple benchmarks.

strengths:
1.	The exploration of the global loss landscape within FedDG research is novel, and the proposed approach of enhancing global flatness through data augmentation represents an innovative and promising direction.
2.	Theoretical analysis in this paper is rigorous, skillfully incorporating global flatness into local updates. This theoretical foundation effectively guides the algorithm design, resulting in a well-grounded and interpretable solution.
3.	Claims made in the methodology, including Assumption 1 and those related to model flatness, are well-validated and substantiated by experimental results.

weaknesses:
1.	The min-max optimization process in the local updates remains somewhat unclear. Further clarification would enhance understanding of the authors’ claim.
2.	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.
3.	This paper does not include an analysis of hyperparameter selection, which is important for SAM-based approaches.

questions:
1.	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.
2.	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.
3.	A sensitivity analysis of key hyperparameters, such as the radius $\gamma$, would add valuable insight.
4.	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.

[1] Sun, Yan, et al. “Dynamic Regularized Sharpness Aware Minimization in Federated Learning: Approaching Global Consistency and Smooth Landscape.” International Conference on Machine Learning. PMLR, 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730548145660

thread_type:
official_review

<thread_object>
creator: reviewer_hj24

<content>
summary:
The paper proposes a method GFM that flattens the global model to a minimum in the case of federated learning. the paper proposes that the method of finding perspectives on global data is based on a solid theoretical foundation that improves the generalization of the distributional bias of a new client.

strengths:
1. the paper's provides a sound and complete theoretical foundation for their method.
2. the paper provides sufficient experiments to show that their method is better and more suitable than FedSAM to be applied to the federated learning method of divisional offsets.
3. the paper provides sufficient literature, citations to illustrate the relevant issues.

weaknesses:
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.

questions:
1. 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?

2. 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
</content>
</thread_object>

</thread>
<thread>
creation date:
1731206639021

thread_type:
official_review

<thread_object>
creator: reviewer_e99p

<content>
summary:
This paper proposes a federated domain generalization method, GFM, which incorporates techniques from sharpness-aware learning. The authors aim to find a flat minimum by coordinating clients during federated training. In addition to sharpness-aware learning, a global model-constrained adversarial data augmentation strategy is employed to better regulate the sharpness of the converged models. The authors claim that existing methods overlook data privacy concerns, motivating the design of the GFM training algorithm. A series of experiments is conducted to verify its effectiveness.

strengths:
1. The paper provides rigorous preliminaries and thorough theoretical support, with detailed proofs included in the appendix. 

2. A series of experiments are conducted here to verify the effectiveness of the proposed method.

3. The proposed method is tailored for domain generalization tasks in both its theoretical derivation and algorithm design.

weaknesses:
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. In comparing the proposed method to FedSAM, I do not see a clear improvement or significant difference between them. The only notable difference appears to be the augmentation network discussed around Equation 9, which is directly adopted from Suzuki (2022). Additionally, the ablation study does not evaluate the performance impact of removing this augmentation network.

2. 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.

3. 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.

4. 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.

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

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

2. 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?

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

4. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731208108219

thread_type:
official_review

<thread_object>
creator: reviewer_njhx

<content>
summary:
The authors adapt the work that ties flat minima to domain generalization in the context of Federated Learning. They show that the prior work in the centralized case does not naturally extend to federated learning. They then move on to build an intuitive case for its extension and propose a new method that utilizes sharpness-aware minimization and data augmentation to simulate global data distribution during client optimization.

strengths:
The paper is easy to read, and the authors explain the concepts in a linear fashion, making it intuitive to understand the decisions they make while formulating their proposed solution. The authors also clearly identify the shortcomings of existing works and the reasons why they would fail to extend to FL, thus proving the necessity of this work.

weaknesses:
While the paper writing is good, the problem is relevant, and the proposed solution is headed in a promising direction, several weaknesses and gaps in the study need to be addressed. The main concerns are listed below:
1. 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. 
2. 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. 
3. The study in section 5.4 is very thorough and quite informative. The main question that the paper does not address is two-fold 
	- 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"
	- secondly the authors need to prove that this statement is further not affected by the relaxation in assumption 1.
	- 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.
4. 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.

questions:
1. Authors mention that GFM = FedAvg + SAM + GCA. However, FedSAM is FedAvg + SAM. Could the authors clarify their contribution with respect to FedSAM?
2. 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?
3. Is it right to infer that Theorem 1 is a restatement and not an original contribution?
4. Could the authors mention some citations that use a similar assumption as "assuming aggregation helps generalization" on lines 191-192?
5. In Algorithm 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.
	- How many models need to be maintained by a client to run this optimization?
	- Could the authors expand line 11 to include loop oversampling, gradient estimation, and update rule? Are the $\phi_i$s also aggregated?
6. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732296743040

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response by Authors

comment:
# Summary

We thank reviewers for their valuable feedback , and appreciate the great efforts made by all reviewers, ACs, SACs and PCs.

We appreciate that the reviewers have multiple positive impressions of our work, including: (1) focused problem is **of necessity and novelty** (Njhx, iTLS); (2) the proposed solution **is headed in a promising direction** (Njhx, iTLS); (3) **thorough theoretical support** (E99p, hj24, iTLS); (4) **sufficient experiments** (E99p, hj24, iTLS);

We provide a summary of our updates, and for detailed responses, please refer to the feedback of each comment/question point-by-point.

- **Relationship with Related Work:** We clarified the relationship between our work and related studies (e.g., FedSAM) and added an introduction to closely related works in the **Introduction** section. Content from the Related Works section was reorganized, and additional citations were included.
- **Improved Analysis:** We enriched the theoretical analysis of the min-max optimization process and provided detailed explanations in Appendix F. The ablation study is modified for better illustration in Section 4.5. A sensitivity analysis of the SAM radius parameter ($\gamma$) and update interval $c$ was added in Section 4.6. Results on surrogate data effectiveness and mean forgetting rates were added in Appendix C.2. 
- **Computational Cost:** A computational efficiency study was included, showing reduced augmentation costs through parallelization and subsampling strategies (see Appendix C.5).
- **Revised Presentation:** Key expressions and equations were refined for clarity.

The above updates in the revised draft (including the regular pages and the Appendix) are highlighted in blue color. 

We once again express our gratitude to all reviewers for their time and efforts devoted to evaluating our work. We eagerly anticipate your further responses and are hopeful for a favorable consideration of our revised manuscript.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper designs GFM that seeks flat global minima in federated domain generalization. The reviewers agreed that it is well written but there is a major lingering concern regarding the overall implications of the work after the discussion phase. I agree with these concerns and hence, assess that the paper does not cross the bar for publication at ICLR. In particular, GFM can be thought of as FedSAM plus an augmentation network called global model-constrained adversarial data augmentation (GCA). Reviewers Njhx and E99p feel that it remains unclear whether the empirical improvements are entirely attributed to GCA. In other words, it remains unconvincing that the reason for improvements can be attributed totally to flatness of the minima of the global model, but perhaps it is due to GCA.