QT-DoG: Quantization-Aware Training for Domain Generalization

Saqib Javed, Hieu Le, Mathieu Salzmann

TL;DR: QT-DoG leverages weight quantization to promote flatter minima, enhancing generalization across unseen domains while reducing model size and computational costs.

Abstract: A key challenge in Domain Generalization (DG) is preventing overfitting to source domains, which can be mitigated by finding flatter minima in the loss landscape. In this work, we propose Quantization-aware Training for Domain Generalization (QT-DoG) and demonstrate that weight quantization effectively leads to flatter minima in the loss landscape, thereby enhancing domain generalization. Unlike traditional quantization methods focused on model compression, QT-DoG exploits quantization as an implicit regularizer by inducing noise in model weights, guiding the optimization process toward flatter minima that are less sensitive to perturbations and overfitting. We provide both an analytical perspective and empirical evidence demonstrating that quantization inherently encourages flatter minima, leading to better generalization across domains. Moreover, with the benefit of reducing the model size through quantization, we demonstrate that an ensemble of multiple quantized models further yields superior accuracy than the state-of-the-art DG approaches with no computational or memory overheads. Code is released at: https://saqibjaved1.github.io/QT_DoG/.

Lay Summary:

Deep learning models often struggle when tested in conditions that differ from those on which they were trained. For example, a model trained to recognize objects in daytime images may not perform well at night. It is a key obstacle to making AI systems reliable in the real world, where conditions frequently change. Domain Generalization (DG) addresses this problem and aims to learn models that perform well not only on the training (source) domains but also in new, unseen (target) data distributions.

In our research, we introduce Quantization-Aware Training for Domain Generalization (QT-DoG) and demonstrate that weight quantization helps deep learning models become more reliable in unfamiliar settings. We demonstrate that quantization-aware training, a technique traditionally employed to reduce model size and improve inference speed, can also serve as a form of regularization. By introducing small structured noise during training, we guide the model to avoid overly sensitive solutions and instead learn more stable patterns that enhance generalization.

We also combine several quantized models into an Ensemble of Quantization (EoQ) and show that EoQ outperforms many larger models and more complex training algorithms, without requiring significant additional computing power.

Link To Code: https://saqibjaved1.github.io/QT_DoG/

Primary Area: General Machine Learning->Transfer, Multitask and Meta-learning

Keywords: Domain Generalization, Quantization, Ensemble, Network Compression, Flat Minima, Regularization

Submission Number: 1828