Over-parameterized Student Model via Tensor Decomposition Boosted Knowledge Distillation

TL;DR

This work proposes a tensor decomposition strategy that effectively over-parameterizes the relatively small student model through an efficient and nearly lossless decomposition of its parameter matrices into higher-dimensional tensors, facilitating the transfer of knowledge of large models.

Abstract

Increased training parameters have enabled large pre-trained models to excel in various downstream tasks. Nevertheless, the extensive computational requirements associated with these models hinder their widespread adoption within the community. We focus on Knowledge Distillation (KD), where a compact student model is trained to mimic a larger teacher model, facilitating the transfer of knowledge of large models. In contrast to much of the previous work, we scale up the parameters of the student model during training, to benefit from overparameterization without increasing the inference latency. In particular, we propose a tensor decomposition strategy that effectively over-parameterizes the relatively small student model through an efficient and nearly lossless decomposition of its parameter matrices into higher-dimensional tensors. To ensure efficiency, we further introduce a tensor constraint loss to align the high-dimensional tensors between the student and teacher models. Comprehensive experiments validate the significant performance enhancement by our approach in various KD tasks, covering computer vision and natural language processing areas. Our code is available at https://github.com/intell-sci-comput/OPDF.

Figures (2)

• Figure 1: The overview of over-parameter distillation framework (OPDF) for knowledge ditillation. a, We use MPO decomposition to realize the over-parameter procedure for the student model. The auxiliary tensors of the student model are trained to imitate the auxiliary tensors of the teacher model closely. b, We present an illustrative example of MPO decomposition. A parameter matrix $\mathbf{W}_{I\times J}$ is decomposed into central tensor and auxiliary tensors.
• Figure 2: The impact of over-parameterization scale, learning rate, and various components of the OPDF on distillation model performance is explored. Figure \ref{['Scale_factor']} demonstrates the performance of the LGTM and DBKD model on the MRPC task following the implementation of the OPDF. Figure \ref{['learning_rate']} presents the performance of DBKD + OPDF with different over-parameterization scales on the MRPC task. Figure \ref{['ablation']} displays the performance of the theseus model across various tasks, utilizing different over-parameterization methods and integrating various components of the OPDF.