Instance Temperature Knowledge Distillation

2024-06-27

arXiv.org

AI Summary

Existing knowledge distillation methods lack a long-term perspective in tuning the temperature parameter, optimizing only the immediate loss at each stage and thus failing to adapt to the dynamically varying difficulty across learning phases. This work formulates instance-level temperature selection as a sequential decision-making problem with delayed rewards and solves it via proximal policy optimization (PPO). We propose three key innovations: (1) an instance-aware dynamic state representation encoding per-sample difficulty and confidence; (2) a reward calibration mechanism to mitigate bias induced by delayed feedback; and (3) an efficient exploration strategy ensuring training stability. The resulting plug-and-play framework consistently improves performance across diverse distillation methods—e.g., KD, RKD, CRD—on both image classification and object detection benchmarks. It exhibits strong generalizability, modular design for easy integration, and requires no architectural modifications to teacher or student models.

Technology Category

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

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Knowledge distillation (KD) enhances the performance of a student network by allowing it to learn the knowledge transferred from a teacher network incrementally. Existing methods dynamically adjust the temperature to enable the student network to adapt to the varying learning difficulties at different learning stages of KD. KD is a continuous process, but when adjusting the temperature, these methods consider only the immediate benefits of the operation in the current learning phase and fail to take into account its future returns. To address this issue, we formulate the adjustment of temperature as a sequential decision-making task and propose a method based on reinforcement learning, termed RLKD. Importantly, we design a novel state representation to enable the agent to make more informed action (i.e. instance temperature adjustment). To handle the problem of delayed rewards in our method due to the KD setting, we explore an instance reward calibration approach. In addition,we devise an efficient exploration strategy that enables the agent to learn valuable instance temperature adjustment policy more efficiently. Our framework can serve as a plug-and-play technique to be inserted into various KD methods easily, and we validate its effectiveness on both image classification and object detection tasks. Our project is at https://www.zayx.me/ITKD.github.io/.

Problem

• Dynamic temperature adjustment in knowledge distillation
• Sequential decision-making for temperature control
• Delayed rewards handling in reinforcement learning

Innovation

• Reinforcement learning for temperature adjustment
• Instance reward calibration for delayed rewards
• Efficient exploration strategy for policy learning