GAQAT: Gradient-Adaptive Quantization-Aware Training for Domain Generalization

Authors: Jiacheng Jiang, Chen Tang, Yuan Meng, Han Yu, Qun Li, Zhi Wang, Wenwu Zhu
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Model Quantization; Domain Generation; Sharpness-Aware Minimization

Abstract:
Research on loss surface geometry, such as Sharpness-Aware Minimization (SAM), shows that flatter minima improve generalization. Recent studies further reveal that flatter minima can also reduce the domain generalization (DG) gap. However, existing flatness-based DG techniques predominantly operate within a full-precision training process, which is impractical for deployment on resource-constrained edge devices that typically rely on lower bit-width representations (e.g., 4 bits, 3 bits). Consequently, low-precision quantization-aware training is critical for optimizing these techniques in real-world applications.
In this paper, we observe a significant degradation in performance when applying state-of-the-art DG-SAM methods to quantized models, suggesting that current approaches fail to preserve generalizability during the low-precision training process. To address this limitation, we propose a novel Gradient-Adaptive Quantization-Aware Training (GAQAT) framework for DG. 
Our approach begins by identifying the scale-gradient conflict problem in low-precision quantization, where the task loss and smoothness loss induce conflicting gradients for the scaling factors of quantizers, with certain layers exhibiting opposing gradient directions. This conflict renders the optimization of quantized weights highly unstable. To mitigate this, we further introduce a mechanism to quantify gradient inconsistencies and selectively freeze the gradients of scaling factors, thereby stabilizing the training process and enhancing out-of-domain generalization.
Extensive experiments validate the effectiveness of the proposed GAQAT framework. On PACS, both 3-bit and 4-bit exceed directly integrating DG and QAT by up to 4.5\%. On DomainNet, our 4-bit results deliver nearly lossless performance compared to the full-precision model, while achieving improvements of up to 1.39\% and 1.06\% over the SOTA QAT baseline for 4-bit and 3-bit quantized models, respectively.

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

Reviews and discussion:
<thread>
creation date:
1730013719188

thread_type:
official_review

<thread_object>
creator: reviewer_twza

<content>
summary:
This paper combines quantization-aware training with the domain generalization model. At first, the authors identify the scale factors (used in quantization) suffer from a gradient conflict problem, where the task loss and smoothness loss present conflicting gradients. As a result, GAQAT, which dynamically quantifies gradient conflicting and selectively freezes the gradients of scaling factors, is proposed to stabilize the training process and enhance domain generalization performance.

strengths:
1 This paper provides detailed implementation settings to help re-production.
2 The experimental results are comprehensive, which helps prove the effectiveness of GAQAT.

weaknesses:
I have some concerns about the writing and experiment comparison of this paper.

Writing:

1 In the main text, Figure 3 should be referred before Figure 4.

2 The gradient disorder is measured with two sequences according to Equation 4. However, in Line 254, the author says 'the gradient disorder of the  g_{task} decreases significantly as training progresses.' How do measure the gradient disorder of the g_{task} alone? 

3 Line 259 - 263. Why lower layers with lower gradient disorder are more likely to exhibit the phenomenon of two gradients in suboptimal equilibrium state? How the Figure 2 and Figure 4 prove this? What is the suboptimal equilibrium state?

4 Totally lost when reading Line 287-294

5 Line 297, 'Continuously **freezing** the g_{task} without **unfreezing**'? How can you freeze g_{task} while unfreezing it?


Experiment comparison:

1 Why only 4/4 and 3/3 bit-width?
2 Maybe the author should add some experiments that quantized the trained model with PTQ methods.
3 Does the quantized model is initialized from the pre-trained full-precision model?

questions:
See Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730022656700

thread_type:
official_review

<thread_object>
creator: reviewer_2jwq

<content>
summary:
The paper investigate the domain generalization gap under quantization regime. The propose a gradient adaptive quantization aware training to be able to make the optimization of low precision and smoothness together. They found the reason for the failure of previous methods is the gradient approximation use in low precision training, and suggest a method to resolve it.

strengths:
I think the main strength of the paper is the general idea that low precision training have different properties than the full precision counterpart, and we need to confront with them separately.

weaknesses:
I think that the paper require a better examination of the domain generalization under different quantization datatypes, group sizes, gradient approximation before it get can a general conclusion and try to solve it.. I think it is a must to be at a level expected from an ICLR paper.
Showing results only for INT quantization with learnable scale is not enough in my opinion.

questions:
1) Do you tried to use scaling factor that is not learnable and only is measured. It is a well used methods, not only for INT quantization also for FP quantization. It is interesting if in that cases we still see the gap between both losses - which is the main motivations for the paper.
2)  Interesting to see how to proposed methods work for other quantization method, such as FP. For example the MXFP4 datatype is a good candidate. Moreover, what about other gradients approximation that are not STE, like PWL or MAD (https://arxiv.org/pdf/2206.06501)
3) I would like to see also experiments in language models.

================================================

Dear authors,

Thank you for taking the time to try to address the questions and concerns raised.

Unfortunately, my two main concerns remain unresolved:

I believe the paper requires more in-depth research into different quantization techniques to achieve more accurate and comprehensive conclusions. In particular, the use of FP quantization, which is the preferred method for most AI accelerators, should be explored further.

The inclusion of language models is essential, as they face unique quantization challenges—such as handling outliers—that do not typically arise in models like ResNet.

I have decided to maintain my score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730534623142

thread_type:
official_review

<thread_object>
creator: reviewer_mcap

<content>
summary:
The paper introduces the Gradient-Adaptive Quantization-Aware Training (GAQAT) framework for enhancing domain generalization in quantized models, specifically targeting low-bit-width settings often used in resource-constrained devices. Addressing the limitations of current domain generalization methods under quantization, the authors identify gradient conflicts that arise during training due to discrepancies between task loss and smoothness loss. To manage this, GAQAT uses a gradient disorder metric to detect these conflicts and applies selective gradient freezing, stabilizing training and improving generalization. Extensive experiments demonstrate GAQAT's effectiveness, showing significant accuracy improvements on PACS and DomainNet datasets compared to baseline methods​.

strengths:
1. The paper identifies and addresses gradient conflicts between g_task and g_smooth, a key insight that stabilizes training and improves convergence. This reviewer fully agrees with this idea.

2. By dynamically freezing scaling factors based on gradient disorder, GAQAT mitigates gradient conflicts, leading to enhanced model performance in low-bit quantization.

3. Improved performance on quantized models in domain generalization scenario

weaknesses:
First, I would like to clarify that I am not highly experienced in the specific domain of Domain Generalization (DG) using ResNet-50. My review reflects this perspective (and if other reviewers find the contributions significant, I am open to adjusting my evaluation). Given the field's recent emphasis on Transformer-based research and development, it has been a few years since I worked with or optimized models like ResNet.

1. Applicability of Quantization to ResNet-50 in DG Scenarios: I find it challenging to fully appreciate the motivation for applying quantization to ResNet-50 in the context of Domain Generalization, especially with a focus on deployment to resource-constrained devices. While the paper assumes such deployment scenarios, it seems unlikely that practitioners would choose to deploy a large CNN model like ResNet-50 on an edge device. It appears this choice may stem from the use of benchmarks typically associated with DG research, but for the edge-device assumption, studying alternative, more compact models seems essential.

2. The identified gradient conflict between g_{task}​  and g_{smooth} recalls the classic challenge of balancing overfitting on domain-specific data with generalization (or regularization) to achieve a smoother loss landscape. This insight may not be limited solely to DG + QAT settings but could have implications in broader contexts. I would be interested in seeing if this theoretical framework extends to general quantization-aware training (QAT) or even Vision Transformers, for instance. This paper could potentially expand its contributions beyond a localized scope, making the findings more universally relevant.

questions:
included in weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730556046738

thread_type:
official_review

<thread_object>
creator: reviewer_l5zt

<content>
summary:
This paper introduces the Gradient-Adaptive Quantization-Aware Training (GAQAT) framework to achieve efficient domain generalization for low-precision quantized model. GAQAT tries to address scale-gradient conflict between gradients from task loss and smoothness loss( used in SAM) by selectively freezing the gradients of scaling factors. The proposed approach is tested on two Domainbed datasets and shows some improvement.

strengths:
1. GAQAT addresses domain generalization issue for low-precision quantized models, which is very practical. This is interesting and under-explored topic.
2. This work proposed an approach to tackle scale-gradient conflict by quantifying gradient inconsistencies and selectively freeze the gradients of scaling factors.
3. The paper is well written and easy to understand.

weaknesses:
1. Incorporating QAT from scratch is not a good idea. QAT is generally performed on a pretrained model and literature shows that it can improve results significantly. You can significantly improve results if you perform quantization after some iterations. 
2. The approach does not show any significant improvement in comparison to baseline methods. On DomainNet dataset, there is no difference between all three (Yours, LSQ, and LSQ +SAGM). On PACS you have a difference of less than 2% between LSQ and your method and you reported a difference of 4.5% in the abstract. 
3. It's a norm in the literature[1,2,3,4,5] to show results on all DomainBed[1] datasets with the mean and standard error over multiple runs. I would recommend authors to add those results as the experiments are currently insufficient.

[1] In Search of Lost Domain Generalization, Gulrajani et al., ICLR 2021.\
[2] Ensemble of Averages: Improving Model Selection and Boosting Performance in Domain Generalization, Arpit et al., NeurIPS 2022.\
[3] SWAD: Domain generalization by seeking flat minima. Cha et al., NeurIPS 2021.\
[4] Diverse weight averaging for out-of-distribution generalization. Rame et al., NeurIPS 2022.\
[5] Cross contrasting feature perturbation for domain generalization. Li et al., ICCV 2023.\

questions:
1. Why in Table 2 and Table 4, "ours" result for 4 bits are different?
2. Have you looked at the experiments with higher bit-precision? Does gradient conflict problem exist with higher bit-precision as well e.g 6,7 bit?
3. I would recommend to explore the generalization aspect of the QAT as well. I can clearly see that combining QAT with SAGM does not help, especially in the results for 3 and 4 bits on PACS dataset. 
4. Having an empirical demonstration of flat minima like SWAD[1] would be good as there is no gaurantee that your proposed approach will result in a flatter minima as compared to LSQ+SAGM.

Minor recommendation:\
Please swap the colors in Table 1, green color can be used for improvement as red is generally considered negative.

[1] SWAD: Domain generalization by seeking flat minima. Cha et al., NeurIPS 2021.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730566229839

thread_type:
official_review

<thread_object>
creator: reviewer_efhm

<content>
summary:
This paper addresses the degradation of domain generalization performance when conventional flatness based domain generalization are applied to quantized models. Existing flatness-based domain generalization techniques are typically used in full-precision training, which is impractical for deployment on resource-constrained edge devices. To address this limitation, the paper proposes a novel Gradient-Adaptive Quantization-Aware Training (GAQAT) framework for domain generalization. The authors first illustrate the scale gradient problem introduced by domain generalization in low precision quantization. In it the task loss and smoothness loss become contradictory therefore making the optimization highly unstable. To mitigate this problem, the proposed framework introduces a mechanism to quantify gradient inconsistencies and selectively freeze gradients, stabilizing the training process and enhancing out-of-domain generalization.

strengths:
The paper is very well structured in clearly identifying the problem statement, identifying the possible cause of the problem and then providing a solution.
The authors have clearly demonstrated the problem of opposing gradients incurred during training the quantized model in Figure 2 therefore establishing the problem statement.
The authors introduce a smoothing parameter in the quantizer to jointly optimize for loss curvature and task loss.
The idea of introducing gradient disorder by looking at the number of direction flips of gradient for smoothing parameters over n training steps is novel and simple to implement.
The idea of dynamically freezing only the task loss gradients for cases with low gradient disorder is innovative.
The final results looking very promising with almost all tasks showing best results with the proposed method.
The authors have also addressed some obvious questions in the ablations like impact of number of steps used for gradient disorder calculation and effect of freezing strategy.

weaknesses:
The experiments section seems a bit weak and might need a clearer explanation.

questions:
I would be curious to see 2 ends of the quantization results. 1/2 bit and 8 bit quantized models using the same method compared to DG-SAM + LSQ.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730608554382

thread_type:
official_review

<thread_object>
creator: reviewer_zrms

<content>
summary:
This submission proposed to conduct Quantization-aware Training (QAT) on Domain Generation (DG) task. Specially, it firstly observed and demonstrated empricially that objective of QAT is not aligned with GD: difference of accumlated gradient and perturbation of scaling factor in QAT. Based on the observation, it proposed a metric (gradient disorder) to measure the consistency of training. The metric is used to determine the optimization procedure of QAT and DG: During training, it calculates the disorder for QAT every K steps, if the metric is below a threshold, DG's optimization will be paused.

strengths:
- The quantitative measurement on disalignment between QAT and DG (Sec.3.2) is interesting and convincing. Though it may be noticed in other works that QAT affects the training of main task, seldom empirical works have been conducted for demonstration.
- The method to circumvent effect of QAT to main task (alternative update with one fixed) is straight forward and promissing.

weaknesses:
- Though objective disalignment (gradient inconsistency) between QAT and main task is observed in DG, the method proposed (Line 296-303) is not coupled with DG: it does not used any properties from DG and it can be applied to other problems that are involved with QAT (such as image classification, LLM with QAT). The inconsistency (method proposed is not coupled with the scene where motivation is discovered) harm the contribution. Author can either improve the method using DG's related parts or prove the method (and observation) is universal in QAT (for example conducting more experiments on other fields).

- The tuning of hyper-parameter is sophisticated: author should explain how step K and threshold $\tao$ are discovered and prove that these parameters are insensitive.

questions:
- Grident disorderd is formulated as the dicrepancy between two gradient sequence in Eq.2, how is it applied to a solo $g_\text{task}$ or $\delta_{t, S_i}$? It is calcualted between evaluation between successive steps?

- How the update rule (Line 296-303) is determined? In other words, it is possilbe to use $g_\text{task}$ (if disorder of $g_\text{task}$ is under certain value, scaling factor in quantizer is not updated). How about simply alternatively update of QAT and DG every K steps, which should be listed as a baseline.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732367055504

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely appreciate the constructive feedback provided by reviewers `zrms`, `eFhm`, `L5zt`, `mcAp`, `2JwQ`, and `TWZa`, as well as their acknowledgment of the strengths of our paper.

In this overall response, we emphasize the key strengths of our work as highlighted by the reviewers and address the main common concern. We have provided detailed responses to each reviewer’s individual questions and look forward to further discussions.

**Strengths**

- The combination of DG and QAT is meaningful and unexplored, and our work provides valuable insights into the conflict between the two gradients, as noted by `eFhm`, `L5zt`, and `mcAp`.
- The experimental approach is clear and straightforward, with existing results showing the effectiveness of our method, as highlighted by `zrms`, `eFhm`, and `TWZa`.
- The paper is well-structured and well-written, praised by `eFhm` and `L5zt`.

**Questions**

- The experiment requires more validation for different bit settings and visualizations of smoother results, as raised by `eFhm`, `L5zt`, and `TWZa`.
- Clearer explanation of the gradient disorder is needed, as requested by `mcAp` and `TWZa`.
- A clearer explanation of the update rule is necessary, as pointed out by `zrms` and `TWZa`.

In response to these points, we have updated the manuscript with new visualizations to demonstrate the experimental results and have corrected relevant descriptions and figure errors. For additional bit experiments, we have provided verification results for 5-bit and 8-bit in a separate response and will include more comprehensive experiments in the final version.

**New Results and Submissions Edits**

*All changes have been highlighted in blue*

In response to the reviewers' feedback, we have submitted an updated version incorporating the following changes:

- *New Experiments*: The newly added Figure 6 visualizes the loss curves for the four PACS domains under 4-bit quantization. As shown, the proposed method achieves significantly smoother surfaces.
- *Optimized Text and Figures*: We have revised the text based on the suggestions of reviewers `L5zt` and `TWZa`, and optimized the order of Figure 2 and Figure 3.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733228821903

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response - II

comment:
As the review process is nearing completion, we would like to summarize the reviewers' concerns and how we addressed them.

We responded to reviewer `zrms`, who pointed out that we had resolved issues related to hyperparameter selection and sensitivity **(W2)**, explained the concept of gradient disorder **(Q1)**, and clarified the origin of the update rules. We also experimentally analyzed the two other cases suggested by the reviewer, demonstrating their inapplicability **(Q2)**. However, regarding the relevance of our method to DG **(W1)**, the reviewer still had some concerns. We provided two rounds of updates in an effort to address these concerns.

We responded to reviewer `eFhm`, who pointed out his concerns about bit selection in experiments **(W1 & Q1)** had been explained, but further modifications are needed in the experimental section.  We provided preliminary results for higher bit rates and included visualizations of the loss surface to further demonstrate the effectiveness of our method. In the final version, we will include all results for 5-bit and 8-bit configurations and add loss surface comparisons for the IID scenario.



We responded to reviewer `L5zt`, addressing his doubts about the experimental setup. We resolved his concerns regarding the table results **(Q1)**, observations related to high bit rates **(Q2)**, the explanation of the significance of DG-SAM in our study **(Q3)**, and the visual evidence of the effectiveness of our method **(Q4)**. We clarified that our process is consistent with theirs and that our pre-trained model is more reasonable, as supported by recent papers **(W1)**. We also explained the experimental  improvements and demonstrated the stability of our method through experiments with multiple bit rates **(W2)**. In the final version, we will include results from multiple rounds of experiments **(W3)**.



We responded to reviewer `mcAp`, who mentioned not being familiar with the field of the paper. We focused on explaining the rationale behind the experimental setup **(W1)** he questioned and addressed the reviewers’ positive recognition of our contributions **(W2)**, hoping to help him understand our work.

We responded to reviewer `2JwQ` by citing classic surveys and the latest SOTA papers to clarify the reasonableness of our experimental setup and the invalidity of the additional experiments he suggested **(Q1&Q2&Q3)**, particularly on language models **(Q3)**.

We responded to reviewer `TWZa`, who acknowledged that most issues have been resolved. These include concerns related to the writing **(Writing 1-5)**, such as the freezing strategy, definitions related to gradient disorder, and experimental settings **(Experiment comparison 1-3)**. However, he still had some concerns about the writing. We will focus on revising these areas in the final version.

Regarding the two key points, we would like to clarify the following:

1. In the experimental section, the added 5-bit and 8-bit experiments, along with loss surface visualizations, demonstrated the effectiveness of our method. In the final version, we will include the full 5-bit and 8-bit results and more visualizations, and will perform multiple rounds of runs to average the results.
2. Regarding the writing, we have addressed the concerns raised by TWZa and highlighted these changes in blue. We’ve also adjusted the text for clarity. As we cannot make changes to the paper now, in the final version, we will provide clearer definitions, such as for gradient disorder, the freezing strategy, and the experimental setup, with examples and improved wording to ensure better understanding.

We appreciate the hard work of all the reviewers, as well as the AC, SAC, and PC, and hope our paper will inspire more contributions to the field.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work performs the domain generalization training in the quantized parameter space, which is quite different from existing practice, i.e., domain generalization first, and quantization later. In particular, DG-SAM, the prevalent method for domain generalization, does not work in their setting and therefore, this work designs a method to mitigate the conflict between the task loss and the smoothness loss in the quantized space. However, most reviews have concerns on the novelty and the writing.