Less is More: Task-aware Layer-wise Distillation for Language Model Compression
This work addresses model compression for language models, which is an incremental improvement in distillation techniques for efficiency and deployment.
The paper tackles the problem of compressing large language models into smaller ones via layer-wise distillation, which often leads to under-fitting due to capacity mismatch and redundant teacher information, by proposing Task-aware layEr-wise Distillation (TED) that uses task-aware filters to select useful knowledge, resulting in significant and consistent improvements over existing methods in continual pre-training and fine-tuning scenarios.
Layer-wise distillation is a powerful tool to compress large models (i.e. teacher models) into small ones (i.e., student models). The student distills knowledge from the teacher by mimicking the hidden representations of the teacher at every intermediate layer. However, layer-wise distillation is difficult. Since the student has a smaller model capacity than the teacher, it is often under-fitted. Furthermore, the hidden representations of the teacher contain redundant information that the student does not necessarily need for the target task's learning. To address these challenges, we propose a novel Task-aware layEr-wise Distillation (TED). TED designs task-aware filters to align the hidden representations of the student and the teacher at each layer. The filters select the knowledge that is useful for the target task from the hidden representations. As such, TED reduces the knowledge gap between the two models and helps the student to fit better on the target task. We evaluate TED in two scenarios: continual pre-training and fine-tuning. TED demonstrates significant and consistent improvements over existing distillation methods in both scenarios. Code is available at https://github.com/cliang1453/task-aware-distillation.