GAQAT: gradient-adaptive quantization-aware training for domain generalization

TL;DR

GAQAT tackles domain generalization under low-precision quantization by uncovering gradient conflicts between task-driven and flatness-driven objectives in quantized training. It introduces gradient disorder as a metric and a dynamic selective-freezing mechanism to stabilize scale gradients in quantizers, enabling joint optimization of quantization and smoothing. By integrating a smoothing term into the quantizer and applying domain-aware quantization-aware training, GAQAT achieves state-of-the-art DG performance on PACS and DomainNet at 3- and 4-bit precision, closely matching or surpassing full-precision baselines. This work enables efficient, deployment-ready DG models for edge devices with practical performance gains.

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, our 3-bit and 4-bit models outperform direct DG-QAT integration by up to 4.5%. On DomainNet, the 4-bit model achieves near-lossless performance compared to full precision, with improvements of 1.39% (4-bit) and 1.06% (3-bit) over the SOTA QAT baseline.

Figures (7)

• Figure 1: Illustration of GAQAT. Compared to full-precision weight gradients, the tensor-wise scale gradients have only two directions: positive and negative. For the newly introduced task-related scale gradients, we apply the GAQAT method for selective freezing. We calculate the disorder of each scale's task gradient $\mathbf{g}_{{\text{task}}}$ and freeze those with disorder below a certain threshold to improve the model's generalization ability.
• Figure 2: Results of cumulative gradients every 350 steps in the 4-bit test on the PACS ART domain, revealing conflicts in the scaling factors.
• Figure 3: Results of task and smoothness gradient disorder of scaling factors over 350 steps in the 4-bit test on the PACS ART domain, revealing in some layers, the gradient disorder of the $\mathbf{g}_{{\text{task}}}$ decreases significantly as training progresses.
• Figure 4: Results of freezing over 350 steps in the 4-bit test on the PACS ART domain, using gradient disorder as an indicator, with no unfreezing. The findings suggest that instability in gradient fluctuations is partly caused by interference between scaling factors during training. Moreover, the gradient disorder indicator proves to be a useful metric for determining when to freeze.
• Figure 5: Results of cumulative gradients every 2111 steps in the 3-bit test on the DoaminNet Clipart and Infograph domains, revealing fewer anomalous gradients compared to 4-bit, with $\mathbf{g}_{{\text{task}}}$ dominating.

Theorems & Definitions (1)

• Definition 3.1