Improving Knowledge Distillation in Transfer Learning with Layer-wise Learning Rates

TL;DR

This work tackles the deteriorating performance of knowledge distillation as task complexity grows by introducing a layer-wise learning-rate optimization that aligns per-layer representations between teacher and student. It applies this layer-wise LR scheme to attention-map, Jacobian-map, and Hessian-map based transfer learning, using per-layer losses and Jensen-Shannon divergence to continuously adapt learning rates at crucial layers. Across CIFAR-10, CIFAR-100, and CoCo benchmarks, the approach yields consistent improvements, with larger gains observed on harder tasks, especially for Jacobian and Hessian-based methods. The method is simple to integrate with existing distillation pipelines and demonstrates practical impact in improving stability and accuracy of compact student models.

Abstract

Transfer learning methods start performing poorly when the complexity of the learning task is increased. Most of these methods calculate the cumulative differences of all the matched features and then use them to back-propagate that loss through all the layers. Contrary to these methods, in this work, we propose a novel layer-wise learning scheme that adjusts learning parameters per layer as a function of the differences in the Jacobian/Attention/Hessian of the output activations w.r.t. the network parameters. We applied this novel scheme for attention map-based and derivative-based (first and second order) transfer learning methods. We received improved learning performance and stability against a wide range of datasets. From extensive experimental evaluation, we observed that the performance boost achieved by our method becomes more significant with the increasing difficulty of the learning task.

Figures (2)

• Figure 1: Schematic representation of the layer-wise learning rate optimization. On the left, $TL_N$ represents teacher layers and $SL_N$ represents the corresponding layers in the student model that dimensionally match with the teacher; the corresponding loss between them is represented by $\it{l}$. The teacher model has more layers than the student model. On the right are the steps of updating the learning rate ($\alpha_j^{(t-1)}$) based on the JSD loss (between the student and the teacher layer at the second identical layer) and prior momentum ($\eta_j^{(t-1)}$) at an interval of 25 epochs.
• Figure 2: Learning curves (epoch vs. accuracy) for with and without layerwise learning in attention-map, jacobian, and hessian-based knowledge distillation. (a) Learning curves along with some samples from the CIFAR-10 dataset (b) Learning curves along with some samples from the CIFAR-100 dataset (c) Learning curves along with some samples from the CoCo dataset. For CIFAR experiments, 'Direct' refers to 'Constant Update' in Table 1, whereas in the CoCo experiment, 'Direct' refers to None in Table 2