Knowledge Distillation with Perturbed Loss: From a Vanilla Teacher to a Proxy Teacher

Rongzhi Zhang, Jiaming Shen, Tianqi Liu, Jialu Liu, Michael Bendersky, Marc Najork, Chao Zhang

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: knowledge distillation, language model, NLP

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We proposed a perturbed loss for knowledge distillation, complemented by a systematic search method to determine the perturbation coefficients, and we validated its effectiveness both theoretically and empirically.

Abstract: Knowledge distillation is a popular technique to transfer knowledge from large teacher models to a small student model. Typically, knowledge distillation employs teacher-forcing learning, where the student learns to imitate the teacher by minimizing the KL divergence of its output distribution with the teacher’s output distribution. In this work, we argue that such a learning objective is sub-optimal because there exists a discrepancy between the teacher’s output distribution and the ground truth label distribution. Therefore, forcing the student to blindly imitate the unreliable teacher output distribution leads to inferior performance. To this end, we propose a novel knowledge distillation objective PTLoss by first representing the vanilla KL-based distillation loss function via a Maclaurin series and then perturbing the leading-order terms in this series. This perturbed loss implicitly transforms the original teacher into a proxy teacher with a distribution closer to the ground truth distribution. We establish the theoretical connection between this “distribution closeness” and the student model generalizability, which enables us to select the PTLoss’s perturbation coefficients in a principled way. Extensive experiments on five datasets demonstrate PTLoss can significantly improve the distillation effectiveness for teachers of various scales.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 9063