Importance Analysis for Dynamic Control of Balancing Parameter in a Simple Knowledge Distillation Setting

TL;DR

This work tackles the problem of balancing distillation and downstream-task losses in knowledge distillation by arguing for a dynamic balancing parameter $\lambda$. It provides a theoretical analysis showing that the per-step loss reduction $\Delta \mathcal{L}^{i+1}$ depends on $\lambda$ via a first-order Taylor expansion and the geometry of the gradient pair $(\nabla L_{dist}, \nabla L_{cls})$, including their magnitudes and the angle $\phi$ between them. The key result is that $\Delta \mathcal{L}^{i+1}$ is a quadratic function of $\lambda$, implying that fixing $\lambda$ during training can hinder optimization and that adapting $\lambda$ to the training state can enhance convergence. This work lays groundwork for adaptive KD strategies that adjust $\lambda$ in response to gradient signals, potentially improving real-time model compression performance in resource-constrained settings.

Abstract

Although deep learning models owe their remarkable success to deep and complex architectures, this very complexity typically comes at the expense of real-time performance. To address this issue, a variety of model compression techniques have been proposed, among which knowledge distillation (KD) stands out for its strong empirical performance. The KD contains two concurrent processes: (i) matching the outputs of a large, pre-trained teacher network and a lightweight student network, and (ii) training the student to solve its designated downstream task. The associated loss functions are termed the distillation loss and the downsteam-task loss, respectively. Numerous prior studies report that KD is most effective when the influence of the distillation loss outweighs that of the downstream-task loss. The influence(or importance) is typically regulated by a balancing parameter. This paper provides a mathematical rationale showing that in a simple KD setting when the loss is decreasing, the balancing parameter should be dynamically adjusted

Figures (2)

• Figure 1: This figure depicts the simple knowledge distillation setting of this paper.
• Figure 2: The magnitudes of $\nabla_{\theta_s^{i}}\mathcal{L}_{\text{dist}}$ and $\nabla_{\theta_s^{i}}\mathcal{L}_{\text{cls}}$ were independently sampled from a uniform distribution over $[10^{-5}, 10^{-1}]$. The red curve depicts cases in which the angle between the two gradient vectors is acute, whereas the green curve corresponds to an obtuse angle. For each trial, that angle was also sampled uniformly at random within the appropriate range.