Paper: Domain-Inspired Sharpness-Aware Minimization Under Domain Shifts (ICLR-2024)

Reviewer comment (Theory):
More discussions on $\lambda$ in eq. (7) are needed. This is a critical parameter that considers the variance/domain shifts in DISAM. However, this $\lambda$ does not appear in Theorem 1. Can the authors illustrate more on this point? And how does the choice of $\lambda$ influence convergence and generalization?

Author response:
Thank you very much for the reviewer's advice. We have elaborated more discussion about the role of $\lambda$ in DISAM in Appendix B.2 (pages 18-21) in the revised submission. In the proof of Theorem 1, specifically Eq. (15) on page 20, **$\lambda$ is integrated into $\beta$, serving as a hyperparameter that regulates the weight adjustment in DISAM**. It functions by modulating the degree of correction for domain shifts:
$$
\beta^i_t = \alpha_i - \frac{2\lambda}{M} \left (\mathcal{L}^i(w_t) - \frac{1}{M} \sum_{j=1}^M \mathcal{L}^j(w_t) \right)
$$

**The influence of $\lambda$:** The choice of $\lambda$ influences how aggressively DISAM responds to variance or domain shifts, with **a higher $\lambda$ leading to more pronounced adjustments in $\beta$**. Our experimental analysis in Figures 5(c) and (d) on page 9, reveals that DISAM's performance remains relatively stable across a wide range of $\lambda$ values. However, **choosing too large $\lambda$ can result in overly aggressive early training adjustments, yielding the negative impact on the convergence process and leading to the increased variance in repeated experiments**. Consequently, we adopted a default $\lambda$ value of 0.1 in all experiments.

Reviewer comment (Novelty):
Relation with a recent work (https://arxiv.org/abs/2309.15639). The paper above also proposes approaches to reduce variance for finding perturbations, although not designed for the domain generalization setting. How does this work relate with the proposed DISAM?

Author response:
Thank you for recommending this excellent work on variance suppression (VaSSO)[1]. We have added this method into the revised submission with the proper discussion. Generally, while DISAM and VaSSO both enhance the perturbation direction generation in SAM, they target different aspects:

- **VaSSO Approach:** VaSSO aims to **reduce noise from mini-batch sampling** by using averaged previous perturbation directions. Its primary focus is on **stabilizing perturbation generation within the same domain**.

- **DISAM's Unique Focus:** In contrast, DISAM incorporates domain information to **specifically address domain-level convergence inconsistencies**, a challenge prevalent in domain shift scenarios. DISAM's approach is to **impose a variance minimization constraint on domain loss** during the perturbation generation process, thereby enabling a more representative perturbation location and enhancing generalization.

We are running the experiments to compare/combine DISAM with VaSSO. Once the experiments are finished, we will report the results here and include the results in the submission.

**Reference:**

[1]. Enhancing Sharpness-Aware Optimization Through Variance Suppression, arXiv2023.

Reviewer comment (Theory):
Theorem 1 illustrates that the *convergence* of DISAM benefits from $\Gamma$. Can the authors explain more on the discussion of 
> as DISAM enjoys a smaller $\Gamma$ than SAM, DISAM can permit the potential larger $\rho$ than that in SAM, thus yielding a better generalization  
In particular, how does the convergence rate link with generalization?

Author response:
Thank you for the advice on explanation about convergence and generalization, and we have added more discussion and analysis in the Appendix B.3 on page 21. In the following, we provide some clarification on these points.

- **Generalization Theorem of SAM:** In the SAM framework, the parameter $\rho$ plays a crucial role in determining generalizability. As established in [1], there exists an upper bound on the generalization error for SAM, suggesting that a larger $\rho$ could potentially enhance generalization, provided that convergence is not impeded. Here is the relevant generalization bound from [1]:
> For any $\rho > 0$ and any distribution $\mathcal{D}$, with probability $1- \delta$ over the choice of the training set $S\sim \mathcal{D}$,
> $$\mathcal{L} _{\mathcal{D}} (w) \leq \max _{\| \epsilon\| _2 \leq \rho} \mathcal{L} _{S}(w+\epsilon) + \sqrt{\frac{k \log \left(  1 + \frac{\| w \| _2^2}{\rho^2} (1+\sqrt{\frac{\log (n)}{k}})^2 \right) + 4 \log \frac{n}{\delta} + \tilde{O}(1)}{n-1}} $$
> where $n = |S|$, $k$ is the number of parameters and we assumed $\mathcal{L} _{\mathcal{D}}(w) \leq \mathbb{E} _{\epsilon_i \approx \mathcal{N}(0, \rho)} [\mathcal{L} _{\mathcal{D}}(w+\epsilon)]$. DISAM leverages a smaller $\Gamma$ than SAM, as shown in Theorem 1 in page 5. This allows DISAM to employ a potentially larger $\rho$, enhancing generalizability.

- **Practical Implications:** Combining the above theorem with the convergence theorem (Theorem 1 on page 5), there is a trade-off with respect to $\rho$. A larger $\rho$ might theoretically enhance generalization but poses greater challenges for convergence. This reflects the intuitive notion that searching for flatter minima across a broader range is inherently more complex, which can potentially affect training efficiency. However, **if $\mathcal{L}_{S} (w + \epsilon)$ can be converged with a sufficiently small value, a larger $\rho$ corresponds to better generalization**.  DISAM, with a smaller $\Gamma$ compared to SAM, converges faster, which means that under the same convergence speed, a larger $\rho$ can be used to achieve better generalization. 


- **Empirical Validation:** Our experiments, as illustrated in Figures 3(c) and (d) on page 6, demonstrate that DISAM effectively employs a larger $\rho$ compared to traditional SAM. DISAM's ability to handle a larger $\rho$ results in both consistent convergence and improved generalization compared to SAM, demonstrating its superiority in domain shift scenarios.

We appreciate the reviewer's comments on the theoretical parts and have adjusted the corresponding parts to make these points more clear.

**Reference:**

[1]. Sharpness-aware minimization for efficiently improving generalization, ICLR2021.

Reviewer comment (Experiments):
SAM-based optimization incurs twice the computational overhead and additional storage overhead in comparison to the commonly used SGD. While DISAM, the method proposed in this paper, demonstrates faster convergence under domain shift conditions when compared to SAM, it does not include a comparison with optimizers such as SGD or Adam.

Author response:
Thank you for highlighting the point of the computational cost. We have to admit that although SAM-based methods help  improve generalization, it is ususally at the expense of approximately double the computational effort per iteration compared to standard SGD, due to the sharpness-aware perturbation. DISAM shares the similar cost as other SAM-based methods, since we target to address the domain-level inconsistency during the sharpeness estimation, instead of the training accelaration. In the revision version, we **highlight this point with the empirical verification** in Figure 8 on page 25, pointing out that compared DISAM with ERM, revealing that although DISAM converges faster than SAM, it does not outpace ERM in terms of convergence speed. **The potential extension combined with the acceleration methods like [1,2] can be explored** in the future works. 

**Reference:** 

[1]. Efficient sharpness-aware minimization for improved training of neural networks, ICLR2022.

[2]. Sharpness-aware training for free, NeurIPS2022.

Reviewer comment (Experiments):
This paper employs multiple benchmarks to evaluate the performance of multi-source domain generalization. The article highlights the need for advancements in the domain shift perspective of the SAM method and suggests conducting comparisons between DISAM and the state-of-the-art (SOTA) method to further validate the effectiveness of the proposed approach.

Author response:
We would like to kindly clarify that in Table 1, we has conducted comparison with some SOTA methods [1-2] like the latest SAM-based method SAGM [1] in combination with CORAL [2]. To further address the reviewer's concern, we here add two new SOTA methods "Decompose, Adjust, Compose"(DAC-SC) [3] and Fishr [4] (following reviewer d1EH's recommendation). The table below presents the performance of different methods across multiple domains:
| **Method**            | **PACS** | **VLCS** | **OfficeHome** | **TerraInc** | **DomainNet** | **Avg.** |
| --------------------- |:--------:|:--------:|:--------------:|:------------:|:-------------:|:--------:|
| ERM                   |   85.5   |   77.5   |      66.5      |     46.1     |     43.8      |   63.9   |
| SAM                   |   85.8   |   79.4   |      69.6      |     43.3     |     44.3      |   64.5   |
| **DISAM**             |   87.3   |   80.1   |      70.7      |     47.9     |     45.8      | **66.4** |
| Fishr (ICML2022) [4]  |   86.9   |   78.2   |      68.2      |     53.6     |     41.8      |   65.7   |
| **Fishr+DISAM**       |   87.5   |   79.2   |      70.7      |     54.8     |     43.9      | **67.2** |
| DAC-SC (CVPR2023) [3] |   87.5   |   78.7   |      70.3      |     46.5     |     44.9      |   65.6   |
| **DAC-SC+DISAM**      |   88.7   |   79.1   |      70.6      |     47.4     |     45.6      | **66.3** |
| SAGM (CVPR2023) [1]   |   86.6   |   80.0   |      70.1      |     48.8     |     45.0      |   66.1   |
| **SAGM+DISAM**        |   87.5   |   80.7   |      71.0      |     50.0     |     46.0      | **67.0** |

These results clearly illustrate the consistent improvement of DISAM on a range of methods in the field of multi-source domain generalization.

**Reference:** 

[1]. Sharpness-aware gradient matching for domain generalization, CVPR2023.

[2]. Deep coral: Correlation alignment for deep domain adaptation, ECCV2016.

[3]. Decompose, Adjust, Compose: Effective Normalization by Playing with Frequency for Domain Generalization, CVPR2023.

[4]. Fishr: Invariant gradient variances for out-of-distribution generalization, ICML2022.

Reviewer comment (Experiments):
The value of $\rho$ in DISAM significantly influences both the convergence speed and generalizability. And it needs more discussion on how to effectively determine the value to maximize the benefits of proposed method.

Author response:
Thank you for the suggestion. Following the advice, we have enriched the discussion about the value of $\rho$ in the Appendix B.3 (page 21) of the revised version. Regarding the experiments, we actually have conducted the corresponding experiments to verify this aspect. We kindly refer the reviewer to Appendix B.3 for more details. We summarize the parts for the reviewer's questions as follows.


- **Generalization Theorem of SAM:** In the SAM framework, the parameter $\rho$ plays a crucial role in determining generalizability. As established in [1], there exists an upper bound on the generalization error for SAM, suggesting that **a larger $\rho$ could potentially enhance generalization, provided that convergence is not impeded**. Here is the relevant generalization bound from [1]:
> For any $\rho > 0$ and any distribution $\mathcal{D}$, with probability $1- \delta$ over the choice of the training set $S\sim \mathcal{D}$,
> $$\mathcal{L} _{\mathcal{D}} (w) \leq \max _{\| \epsilon\| _2 \leq \rho} \mathcal{L} _{S}(w+\epsilon) + \sqrt{\frac{k \log \left(  1 + \frac{\| w \| _2^2}{\rho^2} (1+\sqrt{\frac{\log (n)}{k}})^2 \right) + 4 \log \frac{n}{\delta} + \tilde{O}(1)}{n-1}} $$
> where $n = |S|$, $k$ is the number of parameters and we assumed $\mathcal{L} _{\mathcal{D}}(w) \leq \mathbb{E} _{\epsilon_i \approx \mathcal{N}(0, \rho)} [\mathcal{L} _{\mathcal{D}}(w+\epsilon)]$. This theorem's proof focuses solely on the magnitude of $\rho$, thus affirming the applicability of this theoretical framework to DISAM.

- **Practical Implications:** When considering the convergence Theorem 1 on page 5 alongside the above generalization theorem, a critical trade-off emerges with respect to $\rho$. **A larger $\rho$ might theoretically enhance generalization but poses greater challenges for convergence**. This reflects the intuitive notion that searching for flatter minima across a broader range is inherently more complex, which can potentially affect training efficiency.

- **Empirical Validation:** DISAM, with its accelerated convergence, can utilize a larger $\rho$ while still maintaining an acceptable convergence. This advantage is empirically showcased in Figures 3(c) and (d) on page 6, where we demonstrate that DISAM effectively employs a larger $\rho$ compared to traditional SAM. This ensures both convergence and enhanced generalization. Such a capability to balance between convergence efficiency and generalization is a distinguishing feature of DISAM over conventional SAM methods.

**Reference:** 

[1]. Invariant risk minimization, arXiv2019.

Reviewer comment (Experiments):
The article presents a theoretical analysis suggesting that larger values of parameter $\rho$ should lead to improved generalization, given that convergence is guaranteed. It is important to reflect this aspect in the experiments to provide stronger evidence and validation.

Author response:
Thank you for the suggestion. Following the advice, we have enriched the discussion about the value of $\rho$ in the Appendix B.3 (page 21) of the revised version. Regarding the experiments, we actually have conducted the corresponding experiments to verify this aspect. We kindly refer the reviewer to Appendix B.3 for more details. We summarize the parts for the reviewer's questions as follows.


- **Generalization Theorem of SAM:** In the SAM framework, the parameter $\rho$ plays a crucial role in determining generalizability. As established in [1], there exists an upper bound on the generalization error for SAM, suggesting that **a larger $\rho$ could potentially enhance generalization, provided that convergence is not impeded**. Here is the relevant generalization bound from [1]:
> For any $\rho > 0$ and any distribution $\mathcal{D}$, with probability $1- \delta$ over the choice of the training set $S\sim \mathcal{D}$,
> $$\mathcal{L} _{\mathcal{D}} (w) \leq \max _{\| \epsilon\| _2 \leq \rho} \mathcal{L} _{S}(w+\epsilon) + \sqrt{\frac{k \log \left(  1 + \frac{\| w \| _2^2}{\rho^2} (1+\sqrt{\frac{\log (n)}{k}})^2 \right) + 4 \log \frac{n}{\delta} + \tilde{O}(1)}{n-1}} $$
> where $n = |S|$, $k$ is the number of parameters and we assumed $\mathcal{L} _{\mathcal{D}}(w) \leq \mathbb{E} _{\epsilon_i \approx \mathcal{N}(0, \rho)} [\mathcal{L} _{\mathcal{D}}(w+\epsilon)]$. This theorem's proof focuses solely on the magnitude of $\rho$, thus affirming the applicability of this theoretical framework to DISAM.

- **Practical Implications:** When considering the convergence Theorem 1 on page 5 alongside the above generalization theorem, a critical trade-off emerges with respect to $\rho$. **A larger $\rho$ might theoretically enhance generalization but poses greater challenges for convergence**. This reflects the intuitive notion that searching for flatter minima across a broader range is inherently more complex, which can potentially affect training efficiency.

- **Empirical Validation:** DISAM, with its accelerated convergence, can utilize a larger $\rho$ while still maintaining an acceptable convergence. This advantage is empirically showcased in Figures 3(c) and (d) on page 6, where we demonstrate that DISAM effectively employs a larger $\rho$ compared to traditional SAM. This ensures both convergence and enhanced generalization. Such a capability to balance between convergence efficiency and generalization is a distinguishing feature of DISAM over conventional SAM methods.

**Reference:** 

[1]. Invariant risk minimization, arXiv2019.

Reviewer comment (Evaluation):
Regarding the open class generalization of the clip-based model, further experimental analysis should be conducted to elucidate the reasons behind the superior performance of DISAM.

Author response:
Thank you for the suggestion. To clarify the reasons, we complemented more analysis in Appendix C.6 on page 25 in the revised version. Generally, according to the results in Table 3 (we select some results in the following table for reference) and Figure 4, we can see: 

- **ERM tends to overfit to training data classes:** As shown in the table below, although CoOp and CLIPOOD perform better on base classes than zero-shot, their performance on new classes is worse than zero-shot. This suggests that the fine-tuned parameters **overfit to the existing training data distribution from both the domain and class perspectives**. Figure 4 visualizes the change in performance trends during the training process, and we observe a trend where ERM initially performs well on base classes but then exhibits a decline on new classes, suggesting a collapse of the feature space onto the training data classes.

- **DISAM improves generalization on new classes:** Although SAM offers some relief from overfitting, its performance on new classes does not match zero-shot levels. In contrast, DISAM, by minimizing sharpness more effectively, shows improved performance on new classes, especially in domain shift scenarios.


| **Method** | **Results on Base** | **Results on New** |
| ---------- |:-------------------:|:------------------:|
| Zero-shot  |        72.6         |        67.4        |
| CoOp       |        74.4         |        66.3        |
| **CoOp+DISAM** |      **75.3**       |      **69.6**      |
| CLIPOOD    |        76.0         |        66.9        |
| **CLIPOOD+DISAM** |      **77.0**       |      **69.7**      |

Reviewer comment (Novelty):
This paper considers the domain-level convergence consistency in SAM for multiple domains, and proposes to adopt the domain loss variance in training loss. The convergence consistency is a general issue, and the solution is normal, thus the novelty is not so clear for publication in ICLR.

Author response:
**Differences from general convergence consistency issue:**
- **Distinct focus:** DISAM focuses on the issue where **SAM-based methods are unable to accurately estimate sharpness in domain shift scenarios**, leading to the ineffective sharpness minimization and reduction in generalization performance.
- **Enhancing on top of general methods:** While traditional solutions[1,2,3] aim at convergence consistency in parameter optimization, DISAM's methodology is distinct and orthogonal. It builds upon methods like V-REx[2] and Fishr[3], but **goes further in enhancing out-of-domain generalization through better sharpness minimization**. This is evident in our experiments, where combining DISAM with Fishr results in significant performance gains (as shown in the table below).

| **Method**      | **PACS** | **VLCS** | **OfficeHome** | **TerraInc** | **DomainNet** | **Avg.** |
| --------------- |:--------:|:--------:|:--------------:|:------------:|:-------------:|:--------:|
| V-REx[2]        |   84.9   |   78.3   |      66.4      |     46.4     |     33.6      |   61.9   |
| **V-REx+DISAM** |   85.8   |   78.4   |      70.5      |     45.9     |     42.3      |   64.6   |
| Fishr[3]        |   86.9   |   78.2   |      68.2      |     53.6     |     41.8      |   65.7   |
| **Fishr+DISAM** |   87.5   |   79.2   |      70.7      |     54.8     |     43.9      | **67.2** |


**Novelty and contributions:**
- We first identify that the straightforward application of SAM has **a detrimental impact on training under domain shifts** (as shown Figure 1 and Table 1). Specifically, we observed that the way SAM generates perturbation directions amplifies the inconsistency in convergence between domains, leading to inaccurate sharpness estimation and making sharpness minimization less effective.
- DISAM handle the above challenge by imposing a variance minimization constraint on domain loss **during the sharpness estimation process**, thereby enabling a more representative perturbation location and enhancing generalization. 
- The **significant improvements in extensive experimental results** (as shown in Table 1-3) validate DISAM's novelty and practical relevance.

**Reference:**

[1]. Invariant risk minimization, arXiv2019.

[2]. Out-of-distribution generalization via risk extrapolation (rex), ICML2021.

[3]. Fishr: Invariant gradient variances for out-of-distribution generalization, ICML2022.

Reviewer comment (Theory):
In the definition of the variance between different domain losses, the values of loss between different domains are restricted. Which one is more important? The value of losses in different domains, or the minimization speed of loss in different domains?

Author response:
Firstly, we are sorry that one typo in Eq. (7) has mislead your understanding. In the revised manuscript, we have addressed such a typo issue in the description of Eq.(7) as follows.

$$
\min _{w} \mathbb{E}  _{\xi \in \mathcal{S}} [\mathcal{L} _{DISAM}(w;\xi)] \triangleq \min _{w}  \max _{ \| 
 \epsilon \| _2 \leq \rho} \left [ \sum _{i=1} ^M \alpha _i \mathcal{L} _i (w + \epsilon) -  \lambda \text{Var}\{\mathcal{L} _i(\hat{w} + \epsilon)\} _{i=1} ^M \right]
$$

Here **$\hat{w}$ is $w$ without derivative taken during optimization**, and it only makes effect in the $\max_{\| \epsilon\|_2 \leq \rho}$ loop without affecting the first term.

- **Variance Minimization Focus:** DISAM primarily focuses on minimizing variance during the generation of perturbation directions $\epsilon$. The outer optimization w.r.t. $w$ does not involve a trade-off between the empirical loss term and the variance term as we enforce $\hat{w}$ (assigned by $w$) without derivative taken.
- **Crucial Role of Minimizing Variance:** Minimizing domain-level variance in the perturbation generation loop is critical. Our experiments, illustrated in Figures 5(c) and 5(d) on page 9, show a marked decrease in generalization performance when $\lambda=0$, confirming its essential effectiveness in DISAM. Furthermore, DISAM exhibits a robust performance across a broad range of $\lambda$ values.

Reviewer comment (Theory):
In the learning of multiple domains, there is Multi-Objective Optimization, so the domain-level convergence consistency is a general issue under domain shifts? Or the convergence consistency is a general issue in Multi-Objective Optimization?

Author response:
**Differences from multi-objective optimization:**
- The research topic is intrinsically different from multi-objective optimization. Specially, the goal of DISAM has a single ultimate objective, improving the generalization performance of the model trained under multiple domains. This emphasizes both in-domain generalization and out-of-domain generalization, while multi-objective optimization usually refers to improving the multiple objectives of all collaborated tasks.
- Methodologically, we still use the SAM optimization framework, and **do not involve multi-objective optimization process during training**. Specifically, $\{\alpha_i=n_i/N\}_{i=1}^M$ are constant in Eq.(7) unlike the task-specific variables to be learnt in multi-objective optimization. We observed the negative impact of domain-level convergence inconsistency on SAM-based methods during the perturbation direction generation process. DISAM achieves better perturbation directions for out-of-domain generalization by minimizing the variance of the domain losses.

Reviewer comment (Novelty):
This paper considers the domain-level convergence consistency in SAM for multiple domains, and proposes to adopt the domain loss variance in training loss. The convergence consistency is a general issue, and the solution is normal, thus the novelty is not so clear.

Author response:
**Differences from general convergence consistency issue:**
- **Distinct focus:** DISAM focuses on the issue where **SAM-based methods are unable to accurately estimate sharpness in domain shift scenarios**, leading to the ineffective sharpness minimization and reduction in generalization performance.
- **Enhancing on top of general methods:** While traditional solutions[1,2,3] aim at convergence consistency in parameter optimization, DISAM's methodology is distinct and orthogonal. It builds upon methods like V-REx[2] and Fishr[3], but **goes further in enhancing out-of-domain generalization through better sharpness minimization**. This is evident in our experiments, where combining DISAM with Fishr results in significant performance gains (as shown in the table below).

| **Method**      | **PACS** | **VLCS** | **OfficeHome** | **TerraInc** | **DomainNet** | **Avg.** |
| --------------- |:--------:|:--------:|:--------------:|:------------:|:-------------:|:--------:|
| V-REx[2]        |   84.9   |   78.3   |      66.4      |     46.4     |     33.6      |   61.9   |
| **V-REx+DISAM** |   85.8   |   78.4   |      70.5      |     45.9     |     42.3      |   64.6   |
| Fishr[3]        |   86.9   |   78.2   |      68.2      |     53.6     |     41.8      |   65.7   |
| **Fishr+DISAM** |   87.5   |   79.2   |      70.7      |     54.8     |     43.9      | **67.2** |


**Novelty and contributions:**
- We first identify that the straightforward application of SAM has **a detrimental impact on training under domain shifts** (as shown Figure 1 and Table 1). Specifically, we observed that the way SAM generates perturbation directions amplifies the inconsistency in convergence between domains, leading to inaccurate sharpness estimation and making sharpness minimization less effective.
- DISAM handle the above challenge by imposing a variance minimization constraint on domain loss **during the sharpness estimation process**, thereby enabling a more representative perturbation location and enhancing generalization. 
- The **significant improvements in extensive experimental results** (as shown in Table 1-3) validate DISAM's novelty and practical relevance.

**Reference:**

[1]. Invariant risk minimization, arXiv2019.

[2]. Out-of-distribution generalization via risk extrapolation (rex), ICML2021.

[3]. Fishr: Invariant gradient variances for out-of-distribution generalization, ICML2022.