Data-Augmented Quantization-Aware Knowledge Distillation

📄 arXiv: 2509.03850v1 📥 PDF

作者: Justin Kur, Kaiqi Zhao

分类: cs.LG, cs.CV

发布日期: 2025-09-04

备注: 10 pages, 2 figures

💡 一句话要点

提出数据增强感知的量化知识蒸馏方法，提升低比特模型精度

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 量化感知训练 知识蒸馏 数据增强 上下文互信息 低比特模型

📋 核心要点

  1. 现有量化感知训练和知识蒸馏方法忽略了数据增强对低比特模型性能的影响。
  2. 提出一种基于上下文互信息的指标来评估和选择数据增强策略，无需额外训练。
  3. 实验表明，该方法能显著提升现有量化感知训练和知识蒸馏算法在多种模型和数据集上的性能。

📝 摘要（中文）

本文结合量化感知训练(QAT)和知识蒸馏(KD)，旨在创建具有竞争力的低比特深度学习模型。现有KD和QAT工作主要从网络输出的角度，通过设计更好的KD损失函数或优化QAT的前向和反向传播来提高量化模型的准确性。然而，对于输入转换（如数据增强(DA)）的影响关注不足。量化感知KD和DA之间的关系仍未被探索。本文旨在解决以下问题：如何在量化感知KD中选择好的DA，特别是对于低精度模型？我们提出了一种新颖的指标，该指标根据DA最大化上下文互信息（与图像标签非直接相关的信息）的能力来评估DA，同时确保每个类别的预测平均接近真实标签。该方法自动对DA进行排序和选择，所需训练开销极小，并且与任何KD或QAT算法兼容。广泛的评估表明，使用我们的指标选择DA策略可以显著提高各种模型架构和数据集上的最先进的QAT和KD性能。

🔬 方法详解

问题定义：现有量化感知训练(QAT)和知识蒸馏(KD)方法在设计低比特模型时，主要关注网络输出和训练过程的优化，忽略了数据增强(DA)对模型性能的影响。尤其是在低精度量化模型中，合适的数据增强策略的选择至关重要，但现有方法缺乏有效的DA选择机制。

核心思路：论文的核心思路是利用数据增强来提升模型的泛化能力，尤其是在量化模型的上下文中。通过最大化上下文互信息，即图像中与标签无关的信息，可以使模型学习到更多鲁棒的特征表示。同时，需要保证增强后的数据仍然能够正确分类，因此需要约束模型的预测结果接近真实标签。

技术框架：该方法主要包含以下几个阶段：1) 定义上下文互信息(Contextual Mutual Information)的计算方式，用于评估不同数据增强策略的有效性。2) 设计一个评估指标，该指标同时考虑上下文互信息的最大化和预测结果与真实标签的接近程度。3) 使用该指标对不同的数据增强策略进行排序。4) 选择排名靠前的数据增强策略，并将其应用于量化感知训练和知识蒸馏过程中。

关键创新：该方法最重要的技术创新点在于提出了一种新的数据增强策略评估指标，该指标能够有效地衡量数据增强策略在量化感知训练和知识蒸馏中的有效性。与现有方法相比，该方法能够自动选择合适的数据增强策略，而无需手动调整或进行大量的实验。

关键设计：关键设计包括：1) 上下文互信息的具体计算方法，需要选择合适的特征提取器和互信息估计方法。2) 评估指标的构建，需要平衡上下文互信息的最大化和预测结果的准确性。3) 数据增强策略的选择范围，需要考虑不同的数据增强方法及其组合。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，使用该方法选择的数据增强策略能够显著提高现有QAT和KD算法的性能。例如，在ImageNet数据集上，使用ResNet-18模型进行4比特量化时，该方法能够将模型的精度提高1-2个百分点，超过了现有的最先进方法。

🎯 应用场景

该研究成果可应用于资源受限的边缘设备，例如移动设备、嵌入式系统和物联网设备。通过选择合适的数据增强策略，可以显著提高低比特模型的精度，从而在这些设备上部署更复杂的深度学习模型，实现更智能的应用，例如智能监控、自动驾驶和医疗诊断。

📄 摘要（原文）

Quantization-aware training (QAT) and Knowledge Distillation (KD) are combined to achieve competitive performance in creating low-bit deep learning models. Existing KD and QAT works focus on improving the accuracy of quantized models from the network output perspective by designing better KD loss functions or optimizing QAT's forward and backward propagation. However, limited attention has been given to understanding the impact of input transformations, such as data augmentation (DA). The relationship between quantization-aware KD and DA remains unexplored. In this paper, we address the question: how to select a good DA in quantization-aware KD, especially for the models with low precisions? We propose a novel metric which evaluates DAs according to their capacity to maximize the Contextual Mutual Information--the information not directly related to an image's label--while also ensuring the predictions for each class are close to the ground truth labels on average. The proposed method automatically ranks and selects DAs, requiring minimal training overhead, and it is compatible with any KD or QAT algorithm. Extensive evaluations demonstrate that selecting DA strategies using our metric significantly improves state-of-the-art QAT and KD works across various model architectures and datasets.