Adversarial Moment-Matching Distillation of Large Language Models
Chen Jia
TL;DR
This paper reframes knowledge distillation for large language models as imitation learning by matching action-value moments rather than directly distilling probability distributions. It introduces an adversarial, two-player minimax training algorithm that jointly optimizes on-policy and off-policy moment-matching distances via learnable $Q$-value functions, guided by policy-gradient updates. Empirical results on instruction-following and task-specific tasks show state-of-the-art performance, demonstrating the effectiveness of moment-matching over traditional distribution distances. The approach offers a principled way to capture long-horizon knowledge transfer and could enhance practicality for deploying smaller, efficient LLMs without large sacrifices in performance.
Abstract
Knowledge distillation (KD) has been shown to be highly effective in guiding a student model with a larger teacher model and achieving practical benefits in improving the computational and memory efficiency for large language models (LLMs). State-of-the-art KD methods for LLMs mostly rely on minimizing explicit distribution distance between teacher and student probability predictions. Instead of optimizing these mandatory behaviour cloning objectives, we explore an imitation learning strategy for KD of LLMs. In particular, we minimize the imitation gap by matching the action-value moments of the teacher's behavior from both on- and off-policy perspectives. To achieve this action-value moment-matching goal, we propose an adversarial training algorithm to jointly estimate the moment-matching distance and optimize the student policy to minimize it. Results from both task-agnostic instruction-following experiments and task-specific experiments demonstrate the effectiveness of our method and achieve new state-of-the-art performance.