In-Distribution Consistency Regularization Improves the Generalization of Quantization-Aware Training

2024-02-21

AI Summary

Quantization-aware training (QAT) suffers from poor generalization and instability—exhibiting sensitivity to minor perturbations in inputs or weights. To address this, this paper introduces in-distribution consistency regularization (CR) into QAT for the first time, jointly enhancing model robustness against such perturbations during both training and inference. Unlike prior approaches, our method requires no knowledge distillation and simultaneously mitigates two fundamental causes of QAT failure: training instability and insufficient test-time generalization. Evaluated on CIFAR-10, our approach improves ResNet-18 and MobileNet-V2 by 3.79% and 3.84%, respectively, over standard QAT baselines—surpassing existing state-of-the-art QAT methods and even outperforming their full-precision counterparts. These results demonstrate superior generalization capability and broad architectural compatibility across diverse network families.

Technology Category

• Computer Vision: Adversarial Attacks & Robustness
• Machine Learning: Calibration & Uncertainty Quantification
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Although existing Quantization-Aware Training (QAT) methods intensively depend on knowledge distillation to guarantee performance, QAT still suffers from severe performance drop. The experiments have shown that vanilla quantization is sensitive to the perturbation from both the input and weights. Therefore, we assume that the generalization ability of QAT is predominantly caused by both the intrinsic instability (training time) and the limited generalization ability (testing time). In this paper, we address both issues from a new perspective by leveraging Consistency Regularization (CR) to improve the generalization ability of QAT. Empirical results and theoretical analysis verify that CR would bring a good generalization ability to different network architectures and various QAT methods. Extensive experiments demonstrate that our approach significantly outperforms current state-of-the-art QAT methods and even the FP counterparts. On CIFAR-10, the proposed method improves by 3.79% compared to the baseline method using ResNet18, and improves by 3.84% compared to the baseline method using the lightweight model MobileNet.

Problem

• Quantized Aware Training
• Performance Instability
• Limitation Adaptability

Innovation

• Consistency Regularization
• Quantization-Aware Training
• Performance Enhancement