Causal Distillation for Language Models

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 4288 - 4295 July 10-15, 2022 ©2022 Association for Computational Linguistics Causal Distillation for Language Models Zhengxuan Wu∗¶, Atticus Geiger∗¶, Joshua Rozner, Elisa Kreiss, Hanson Lu Thomas Icard, Christopher Potts, Noah D. Goodman Stanford University {wuzhengx, atticusg}@stanford.edu Abstract Distillation efforts have led to language mod- els 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 objec- tive (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 dynamics of the teacher through a distillation interchange intervention training objective (DIITO). DIITO pushes the student model to become a causal abstrac- tion of the teacher model – a faithful model with simpler causal structure. DIITO is fully differentiable, easily implemented, and com- bines flexibly with other objectives. Compared against standard distillation with the same set- ting, DIITO results in lower perplexity on the WikiText-103M corpus (masked language mod- eling) and marked improvements on the GLUE benchmark (natural language understanding), SQuAD (question answering), and CoNLL- 2003 (named entity recognition). 1 Introduction Large pretrained language models have improved performance across a wide range of NLP tasks, but can be costly due to their large size. Distillation seeks to reduce these costs while maintaining per- formance by training a simpler student model from a larger teacher model (Hinton et al., 2015; Sun et al., 2019; Sanh et al., 2019; Jiao et al., 2019). Hinton et al. (2015) propose model distillation with an objective that encourages the student to produce output logits similar to those of the teacher while also supervising with a task-specific objec- tive (e.g., sequence classification). Sanh et al. (2019), Sun et al. (2019), and Jiao et al. (2019) ∗Equal contribution. ¶Correspondence authors. adapt this method, strengthening it with additional supervision to align internal representations be- tween the two models. However, these approaches may push the student model to match all aspects of internal states of the teacher model irrespective of their causal role in the network’s computation. This motivates us to develop a method that focuses on aligning the causal role of representations in the student and teacher models. We propose augmenting standard distillation with a new objective that pushes the student to become a causal abstraction (Beckers and Halpern, 2019; Beckers et al., 2020; Geiger et al., 2021a) of the teacher model: the simpler student will faithfully model the causal effect of teacher rep- resentations on output. To achieve this, we employ the interchange intervention training (IIT) method of Geiger et al. (2021b). The distillation inter- change intervention training objective (DIITO) aligns a high-level student model with a low-level teacher model and performs interchange interven- …

Content truncated to 500 words