Causal Distillation for Language Models

TL;DR

This paper introduces Causal Distillation with interchange intervention training (IIT), a novel method that enhances language model distillation by encouraging the student to imitate the teacher's causal computation process, leading to improved performance.

Contribution

It proposes IIT as a new objective for distillation that promotes causal abstraction, improving efficiency and performance of language models over standard methods.

Findings

Lower perplexity on Wikipedia masked language modeling

Improved results on GLUE benchmark

Better performance on SQuAD and CoNLL-2003

Abstract

Distillation efforts have led to language models that are more compact and efficient without serious drops in performance. The standard approach to distillation trains a student model against two objectives: a task-specific objective (e.g., language modeling) and an imitation objective that encourages the hidden states of the student model to be similar to those of the larger teacher model. In this paper, we show that it is beneficial to augment distillation with a third objective that encourages the student to imitate the causal computation process of the teacher through interchange intervention training(IIT). IIT pushes the student model to become a causal abstraction of the teacher model - a simpler model with the same causal structure. IIT is fully differentiable, easily implemented, and combines flexibly with other objectives. Compared with standard distillation of BERT, distillation via IIT results in lower perplexity on Wikipedia (masked language modeling) and marked improvements on the GLUE benchmark (natural language understanding), SQuAD (question answering), and CoNLL-2003 (named entity recognition).