Knowledge Distillation with Training Wheels

TL;DR

This paper introduces a generalized framework for knowledge distillation that allows models to learn from teachers during training and selectively seek help at test-time, improving performance and flexibility in language tasks.

Contribution

It formulates knowledge distillation as an entropy-regularized optimization problem and develops a new algorithm using Path Consistency Learning and constrained reinforcement learning for test-time assistance.

Findings

Improved translation and summarization accuracy.

Enhanced control over teacher assistance during inference.

Unlocks new operating points beyond existing decoding methods.

Abstract

Knowledge distillation is used, in generative language modeling, to train a smaller student model using the help of a larger teacher model, resulting in improved capabilities for the student model. In this paper, we formulate a more general framework for knowledge distillation where the student learns from the teacher during training, and also learns to ask for the teacher's help at test-time following rules specifying test-time restrictions. Towards this, we first formulate knowledge distillation as an entropy-regularized value optimization problem. Adopting Path Consistency Learning to solve this, leads to a new knowledge distillation algorithm using on-policy and off-policy demonstrations. We extend this using constrained reinforcement learning to a framework that incorporates the use of the teacher model as a test-time reference, within constraints. In this situation, akin to a human learner, the model needs to learn not only the learning material, but also the relative difficulty of different sections to prioritize for seeking teacher help. We examine the efficacy of our method through experiments in translation and summarization tasks, observing trends in accuracy and teacher use, noting that our approach unlocks operating points not available to the popular Speculative Decoding approach.