Causal Distillation for Language Models (2022)

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 dynamics of the teacher through a distillation interchange intervention training objective (DIITO). DIITO pushes the student model to become a causal abstraction of the teacher model – a faithful model with simpler causal structure. DIITO is fully differentiable, easily implemented, and combines flexibly with other objectives. Compared against standard distillation with the same setting, DIITO results in lower perplexity on the WikiText-103M corpus (masked language modeling) and marked improvements on the GLUE benchmark (natural language understanding), SQuAD (question answering), and CoNLL-2003 (named entity recognition).

## 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 performance by training a simpler student model from a larger teacher model @xcite @xcite . @xcite 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 objective (e.g., sequence classification). @xcite , @xcite , and jiao et al. (2019) * equal contribution. ¶ correspondence authors. adapt this method, strengthening it with additional supervision to align internal representations between 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 @xcite @xcite of the teacher model: the simpler student will faithfully model the causal effect of teacher representations on output. to achieve this, we employ the interchange intervention training (iit) method of @xcite . the distillation interchange intervention training objective (diito) aligns a high-level student model with a low-level teacher model and performs interchange interventions (swapping of aligned internal states); during training the high-level model is pushed to conform to the causal dynamics of the low-level model. figure 1 shows a schematic example of this process. here, hidden layer 2 of the student model (bottom) is aligned with layers 3 and 4 of the teacher model. the figure depicts a single interchange intervention replacing aligned states in the left-hand models with those from the right-hand models. this results in a new network evolution that is shaped both by the original input and the interchanged hidden states. it can be interpreted as a certain kind of counterfactual as shown in figure 1 : what would the output be for the sentence "i ate some ⟨mask⟩." if the activation values for the second token at the middle two layers were set to the values they have for the input "the water ⟨mask⟩ solid."? diito then pushes the student model to output the same logits as the teacher, i.e., matching the teacher's output distribution under logits salad i ate @xmath0. logits froze the @xmath1. pizza logits i ate @xmath2. froze logits the @xmath3. in this case, the student originally predicted the token "salad" under the interchange intervention, while the teacher predicted the token "pizza" under an aligned interchange intervention. diito trains the student to minimize the divergence between the student logits and the teacher logits under the interchange intervention. this updates the student to conform to causal dynamics of the teacher. to assess the contribution of distillation with diito, we begin with bert base @xcite and distill it under various alignments between student and teacher while pretraining on the wikitext-103m corpus @xcite achieving -2.24 perplexity on the mlm task compared to standard distilbert trained on the same data. we then fine-tune the best performing distilled models and find consistent performance improvements compared to standard distilbert trained with the same setting on the glue benchmark (+1.77%), conll-2003 name-entity recognition (+0.38% on f1 score), and squad v1.1 (+2.46% on em score). ## related work distillation was first introduced in the context of computer vision @xcite and has since been widely explored for language models @xcite @xcite . for example, @xcite propose to extract information not only from output probabilities of the last layer in the teacher model, but also from in-termediate layers in the fine-tuning stage. recently, @xcite adapt causal analysis methods to estimate the effects of inputs on predictions to compress models for better domain adaptation. in contrast, we focus on imbuing the student with the causal structure of the teacher. interventions on neural networks were originally used as a structural analysis method aimed at illuminating neural representations and their role in network behavior @xcite @xcite @xcite @xcite . @xcite extend these methods to network optimization. we contribute to this existing research by adapting intervention-based optimization to the task of language model distillation. ## causal distillation here, we define our distillation training procedure. see algorithm 1 for a summary. ## experimental set-up we adapt the open-source hugging face implementation for model distillation @xcite ). 1 we distill our models on the mlm pretraining task @xcite . we use large gradient accumulations over batches as in @xcite for better performance. specifically, we distill all models for three epochs for an effective batch size of 240. in contrast to the setting of 4k per batch in bert @xcite and distilbert @xcite , we found that small effective batch size works better for smaller dataset. we weight all objectives equally for all experiments. with our new objectives, the distillation takes approximately 9 hours on 4 nvidia a100 gpus. student and teacher models. our two students have the standard bert architecture, with 12 heads with a hidden dimension of 768. the larger student has 6 layers, the smaller 3 layers. our pretrained teacher has the same architecture, except with 12 layers. following practices introduced by @xcite , we initialize our student model with weights from skipped layers (one out of four layers) in the teacher model. we use wikitext for distillation to simulate a practical situation with a limited computation budget. we leave the exploration of our method on larger datasets for future research. alignment. our teacher and student bert models create columns of neural representations above each token with each row created by the feedforward layer of a transformer block, as in figure 1 . we define l t and l s to be the number of layers in the student and teacher, respectively. in addition, we define s j i and t j i to be the representations in the ith row and jth column in the student and teacher, respectively. an alignment π is a partial function from student representations to sets of teacher representations. we test three alignments: in the first and second rows: π(s j 1 ) = {t j l t -2 } and π(s j 2 ) = {t j l t -1 } for each training iteration, we randomly select one aligned student layer to perform the interchange intervention, and we randomly select 30% of token embeddings for alignment for each sequence. we experiment with three conditions with the full alignment: consecutive tokens (diito full ), random tokens (diito full +random) and masked tokens (diito full +masked). we also include l diito cos to the full alignment (diito full +l diito cos ). ## conclusion in this paper, we explored distilling a teacher by training a student to capture the causal dynamics of its computations. across a wide range of nlp tasks, we find that diito leads to improvements, with the largest gains coming from the models that use the richest alignment between student and teacher. our results also demonstrate that diito performs on-par (maintaining 97% of performance on glue tasks) with standard distilbert @xcite while consuming 97% less training data. these findings suggest that diito is a promising tool for effective model distillation.