Contrastive Distillation on Intermediate Representations for Language Model Compression

TL;DR

This paper introduces CoDIR, a contrastive distillation framework that improves language model compression by capturing structural knowledge in intermediate representations, leading to better performance on benchmarks.

Contribution

The paper proposes a novel contrastive distillation method for intermediate layers, surpassing existing L2-based approaches in language model compression.

Findings

Outperforms state-of-the-art methods on GLUE benchmark

Effective in both pre-training and fine-tuning stages

Leverages contrastive learning to capture richer intermediate representations

Abstract

Existing language model compression methods mostly use a simple L2 loss to distill knowledge in the intermediate representations of a large BERT model to a smaller one. Although widely used, this objective by design assumes that all the dimensions of hidden representations are independent, failing to capture important structural knowledge in the intermediate layers of the teacher network. To achieve better distillation efficacy, we propose Contrastive Distillation on Intermediate Representations (CoDIR), a principled knowledge distillation framework where the student is trained to distill knowledge through intermediate layers of the teacher via a contrastive objective. By learning to distinguish positive sample from a large set of negative samples, CoDIR facilitates the student's exploitation of rich information in teacher's hidden layers. CoDIR can be readily applied to compress large-scale language models in both pre-training and finetuning stages, and achieves superb performance on the GLUE benchmark, outperforming state-of-the-art compression methods.