Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Contrastive Distillation on Intermediate Representations for Language Model Compression

Siqi Sun, Zhe Gan, Yu Cheng, Yuwei Fang, Shuohang Wang, Jingjing Liu

TL;DR

Large language models are costly to deploy; traditional distillation relying on L2 losses on hidden states may miss structural information in intermediate representations. CoDIR introduces a contrastive distillation framework that aligns teacher and student intermediate representations via an InfoNCE objective on mean-pooled per-layer embeddings, complemented by standard KD on logits and a memory-bank mechanism for efficient negative sampling. Evaluations on GLUE show that a 6-layer student compressed from RoBERTa-base can achieve competitive results with substantial inference speedups, outperforming prior task-agnostic compression methods. The approach demonstrates that leveraging intermediate representations with contrastive learning yields stronger regularization and practical gains for language-model compression across pre-training and finetuning regimes.

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.

Contrastive Distillation on Intermediate Representations for Language Model Compression

TL;DR

Large language models are costly to deploy; traditional distillation relying on L2 losses on hidden states may miss structural information in intermediate representations. CoDIR introduces a contrastive distillation framework that aligns teacher and student intermediate representations via an InfoNCE objective on mean-pooled per-layer embeddings, complemented by standard KD on logits and a memory-bank mechanism for efficient negative sampling. Evaluations on GLUE show that a 6-layer student compressed from RoBERTa-base can achieve competitive results with substantial inference speedups, outperforming prior task-agnostic compression methods. The approach demonstrates that leveraging intermediate representations with contrastive learning yields stronger regularization and practical gains for language-model compression across pre-training and finetuning regimes.

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.

Paper Structure

This paper contains 29 sections, 5 equations, 1 figure, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Language Model Compression
  4. Contrastive Representation Learning
  5. CoDIR for Model Compression
  6. Framework Overview
  7. Contrastive Distillation
  8. Intermediate Representation
  9. Contrastive Objective
  10. Pre-training and Finetuning Adaptation
  11. Memory Bank
  12. Finetuning
  13. Pre-training
  14. Experiments
  15. Datasets
  16. ...and 14 more sections

Figures (1)

  • Figure 1: Overview of the proposed CoDIR framework for language model compression in both pre-training and finetuning stages. "Trm" represents a Transformer block, $X$ are input tokens, $f^t$ and $f^s$ are teacher and student models, and $X_0, \{X_i\}_{i=1}^K$ represent one positive sample and a set of negative samples, respectively. The difference between CoDIR pre-training and finetuning mainly lies in the negative example sampling method.