QT-DoG: Quantization-Aware Training for Domain Generalization

We appreciate your recognition of our results' competitiveness and the efficiency of combining quantized models in ensembling. We have carefully considered your comments and have provided detailed responses to your queries below.

W1: How does your approach of replacing standard ERM training with quantization and using an ensemble of quantized models offer a novel contribution to domain generalization, given that quantization and model-soup are already established techniques?

While it is true that quantization has been applied in ERM in prior work, our key contribution lies in repurposing quantization as an implicit regularizer specifically for domain generalization (DG). Unlike traditional uses of quantization for compression, we show that quantization noise helps guide optimization toward flatter minima, which enhances OOD generalization. For more clarity, we reworded our first contribution as ``We are the first to demonstrate that quantization-aware training, traditionally used for model compression, can serve as an implicit regularizer, with quantization noise enhancing domain generalization."

Moreover, our ensemble of quantized models builds on model-soup but introduces a key innovation by combining quantization with ensembling in a way that reduces resource overhead while maintaining high performance. This fusion not only improves DG but also ensures scalability across diverse benchmarks. The combination of these techniques within the context of domain generalization, supported by extensive empirical validation, represents a solid and impactful contribution, which we believe to be of interest to the community.

W2: How does the manuscript define flat minima, and what theoretical evidence supports the connection between quantization and flatter minima?

We appreciate the reviewer’s feedback and acknowledge that the definition of flat minima needs to be more clearly presented. In the revised manuscript, we have clarified the discussion of flat minima. Flat minima refer to regions in the loss landscape where the loss remains relatively stable under small perturbations of the model parameters.

The theoretical connection is drawn in Equation 4, which relates the Hessian of the loss to the impact of quantization. Quantization noise acts as a perturbation, and the optimization process biases the model toward flatter regions where the Hessian eigenvalues are smaller, ensuring stability under these perturbations. We have revised the text to clarify this.

W3: Equation 5 does not provide theoretical evidence to support $\mathcal{F_\gamma(W^q)<F_\gamma(W)}$. How do you justify its inclusion, given the variety of quantization methods and settings?

Equation 5 is not intended to provide a theoretical proof or establish a formal connection but rather to serve as an empirical demonstration, akin to the approach used in SWAD [1]. It illustrates that the loss landscape becomes flatter with the application of quantization. This observation is supported by the plots in Figure 2, where it is evident that quantized models consistently achieve flatter loss landscapes compared to their non-quantized counterparts. While the diversity of quantization methods and settings makes a universal theoretical framework challenging, our empirical findings consistently validate this behavior across the datasets and quantization methods we evaluated.

[1] SWAD: Domain generalization by seeking flat minima. Cha et al., NeurIPS 2021.

W4: Not all quantization methods lead to flat minima. Why does the adopted method succeed where others fail?

This observation aligns with our findings. We demonstrate that quantization-aware training (QAT) encourages a flatter minima, which is not guaranteed with post-training quantization (PTQ). In QAT, flatter minima is achieved by incorporating quantization noise during training, acting as an implicit regularizer and smoothing the loss landscape (see Equation 4). In contrast, PTQ primarily involves clipping the network weights without any subsequent retraining., so it does not bring the same effect. This difference is reflected in Table 2, where PTQ performs worse on out-of-domain data compared to QAT. We provide a detailed discussion of this distinction in Section 4.3.4. If further clarification is needed, we are happy to revise this explanation.

Thank you for appreciating the clarity of our writing, the novelty of leveraging weight quantization for domain generalization, and the strength of our experimental results. We address your major questions and concerns below.

W1: Equation 4.

Thank you for your comment. We acknowledge that our discussion of Equation 4 might have been unclear. In short, and as updated in the paper, Equation 4 shows that the noise induced by the quantization process will facilitate the optimization to escape from sharp minima and instead settle in flatter ones. There is however no guarantees that these flatter minima have a lower loss value than others or are global minima. We hope the revised text for Equation 4 clarifies what we meant.

W2: How does large noise dominate optimization in Equation 4 and result in sub-optimal convergence?

The quantization-induced noise (\Delta\) in Equation 4 affects the loss function as follows:

$L(w + \Delta) \approx L(w) + \nabla L(w)^T \Delta + \frac{1}{2} \Delta^T \mathcal{H} \Delta,$

where $\mathcal{H}$ is the Hessian matrix. The second and third terms in this approximation introduce a regularization effect: they penalize sharp minima (large eigenvalues of $\mathcal{H}$) and encourage convergence to flatter regions of the loss surface. However, if the noise (\Delta\) becomes too large, it introduces over-regularization. This excessive noise can overly restrict the search space, preventing the model from reaching a good solution. Instead, the optimization process may focus on minimizing the loss in a way that avoids sharp regions, but sacrifices the ability to find true minimum of the loss function.

Table 8 evidences this trade-off: Moderate noise improves out-of-distribution (OOD) generalization, but excessive noise harms both in-domain and OOD performance. The noise perturbation introduces stochasticity into the optimization process, which biases the trajectory toward flatter minima, but too much noise destabilizes the process, leading to suboptimal convergence.

We have revised the text in the paper to clarify this.

W3: How does quantization affect model size, training costs, and inference costs?

Quantization reduces both the memory footprint and latency. For example, a ResNet-50 model running on an AMD EPYC 7302 processor achieves a latency of 34.28ms in full precision and 21.02ms with 8-bit quantization. While quantization can theoretically reduce training costs by enabling dynamic bit precision switching, our setup does not support this, so training time remains almost unchanged. However, inference time and model size are significantly reduced, as detailed in the results section of the paper.

Q1: How are ensemble quantization models trained—differently from the beginning or only at the quantization step?

The models are trained independently from initialization, using different random seeds to ensure diversity. We have clarified it in the updated manuscript.

Q2: How does the quantization step affect the model performance? How do the author select such hyperparameter?

The quantization steps were empirically determined. We added the ablation studies in the appendix to discuss the performance trends. Here is the ablation study performed on the PACS dataset:

Thanks for your responses. Most of my concerns have been addressed and I would like to maintain my original rating of 6.

Thank you for acknowledging the novelty of using quantization to achieve flatter minima and the robustness of our experimental results. We have addressed your questions and concerns below:

W1/Q2: Why each channel in a layer has a different scaling factor? How is the value of $s$ determined?

It is common practice[1,2,3] to use a learnable, channel-wise scaling factor $s$ because channels within a layer often exhibit varying activation and weight distributions. To account for these differences, channel-wise scaling factors are applied to normalize the perturbation $\Delta$ per channel. However, to explore the impact of this choice, we conducted an ablation study where we set $s$ at the layer level, rather than on a per-channel basis. We see that Channelwise $s$ can lead to 1.5% accuracy as compared to layerwise $s$. The results of this experiment on the PACS dataset with 7 bit quantization are shown below:

[1] Mixed-Precision Neural Network Quantization via Learned Layer-wise Importance, ECCV 2022.
[2] Learnable Companding Quantization for Accurate Low-bit Neural Networks, CVPR 2021.
[3] Learned Step Size Quantization, ICLR 2020.
[4] LSQ+: Improving low-bit quantization through learnable offsets and better initialization, CVPRW 2020.

W2: Measuring sharpness by equation 5 needs more discussion. When the model weight under measurement has different scales, perturbing weight with the same does not serve as a fair way to get around the local area of w'. Relative flatness (Definition 3 in [1]) could be considered to address the dependency of sharpness measurement over scale of model weight. [1] https://arxiv.org/pdf/2001.00939

Thank you for sugggesting this. We find this idea compelling and would like to conduct a proof of concept to explore its potential in our analysis. We have contacted the authors of [1] and will try to incorporate this in our framework to further extend our empirical evaluation.

[1] Relative Flatness and Generalization, NeurIPS 202.

Q1: How much does $w_q$ differ from $w$ in training?

Initially, the difference is relatively large when the weights gets clipped to quantization levels, but it decreases as training progresses and then remains almost constant throughout the rest of the training. Overall, the difference is relatively small and stays almost constant during training.

Q3: Given a fixed quantization bit $b$, how would $s$ influence the performance?

We are not sure to understand the reviewer's question. To clarify, $b$ is pre-determined by the user and the channel-wise $s$ values are learned during training; as such, they can automatically adapt to the given $b$ value.

Q4: What is the scale of w (e.g. norm) across different algorithms in section 3.3?

Following the reviewer's suggestion, we measured the norm of the different $w$s in the network and observed no significant difference between the norms of quantized and non-quantized parameters, even across different datasets.

We sincerely appreciate your recognition of the intriguing role of quantization noise in enhancing generalization, as well as your positive remarks on the organization of our work. Below, we address your major concerns in detail:

W1: The methods for quantization and ensemble are quite standard and not specific to domain generalization, limited novelty

Our approach is the first to use quantization to address domain generalization challenges. We present novel analysis to show that quantization-aware training introduces an implicit regularization effect, guiding the model toward flatter minima and thus enhancing out-of-distribution (OOD) generalization. We appreciate that the reviewer acknowledges this as an ``interesting observation", which further highlights that it brings novelty to the community.

Additionally, we introduce the Ensemble of Quantization (EoQ), which combines quantization and ensembling to amplify the benefits of both techniques. EoQ achieves state-of-the-art OOD performance while requiring minimal resource overhead, demonstrating its practical and methodological significance. Our work not only highlights a novel application of quantization noise but also provides empirical validation of its effectiveness across diverse datasets, setting it apart from standard techniques. We believe this to be of interest to the community.

W2:How does your work reconcile the flatness-generalization relationship with recent critiques[1,2]?

Thank you for highlighting this important discussion!

To better understand the conclusions of [1], we reached out to its authors, and Andriushchenko (the primary author) directed us to his post-ICML discussions (https://x.com/maksym_andr/status/1687395919442948096), where he acknowledges that sharpness (or its inverse, flatness) can still be a valuable measure of generalization. He states that their work ``doesn't imply that sharpness is useless, particularly since the empirical success of SAM is undeniable." This perspective aligns well with our findings. We do not claim that flatter minima are universally better, but instead demonstrate their utility in reducing sensitivity to domain shifts, as supported by experiments across multiple benchmarks.

Furthermore, the work in [2] is orthogonal to our findings. It shows that, on certain datasets, applying SAM to normalization layers only can improve performance over applying it across the entire model. We added discussion around [1,2] in the revised manuscript to address the nuanced relationship between flatness and generalization.

[1] Andriushchenko, Maksym, et al. "A modern look at the relationship between sharpness and generalization." ICML, 2023.
[2] Mueller, Maximilian, et al. "Normalization layers are all that sharpness-aware minimization needs." NeurIPS, 2023.

W3: How does quantization noise compare to uniform weight noise and other weight perturbation schemes?

Quantization noise and uniform weight noise share similarities in that both introduce perturbations to the model's parameters. However, quantization noise specifically arises from the discretization of the weights, which can lead to a more structured form of regularization due to the rounding or truncation during the quantization process. In contrast, uniform weight noise typically adds random perturbations with a uniform distribution, which may not exhibit the same structured regularization properties.

Below, we provide the results of our ablation study on the PACS dataset with uniform noise with different minimum and maximum value:

We hope that our response have provided clarity and effectively addressed your concerns. We would greatly appreciate it if you could acknowledge this. If there are any remaining questions or unresolved concerns, we would be more than happy to provide further clarification. Thank you sincerely for your time and valuable feedback—it is greatly appreciated.