QT-DoG: Quantization-aware Training for Domain Generalization

2024-10-08

arXiv.org

AI Summary

Domain generalization (DG) suffers from source-domain overfitting, while flat minima of the loss landscape enhance cross-domain generalization. To address this, we propose a quantization-aware training framework that introduces weight quantization as an implicit regularizer—marking the first application of quantization for DG. Theoretically and empirically, we show that quantization noise inherently promotes flatter minima. Building upon this insight, we further design a zero-overhead ensemble strategy leveraging multiple quantized models, requiring no additional inference cost. Our method consistently outperforms state-of-the-art approaches across diverse DG benchmarks, network architectures, and quantization schemes, simultaneously improving accuracy and enabling model compression. It is also fully compatible with mainstream DG techniques. The core contribution lies in uncovering the intrinsic link among quantization, loss landscape flatness, and generalization, thereby establishing the first quantization-driven flat-optimization paradigm for DG.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Domain Generalization (DG) aims to train models that perform well not only on the training (source) domains but also on novel, unseen target data distributions. A key challenge in 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 theoretical insights 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. Our extensive experiments demonstrate that QT-DoG generalizes across various datasets, architectures, and quantization algorithms, and can be combined with other DG methods, establishing its versatility and robustness.

Problem

• Prevent overfitting to source domains in Domain Generalization
• Use weight quantization to find flatter loss minima
• Enhance generalization across domains with quantized models

Innovation

• Uses weight quantization for flatter minima
• Quantization acts as implicit regularizer
• Ensemble of quantized models boosts accuracy