Venue: NAACL 2022

# Causal Distillation for Language Models

Zhengxuan Wu, Atticus Geiger, Joshua Rozner, Elisa Kreiss, Hanson Lu Thomas Icard, Christopher Potts, Noah D. Goodman

Equal contribution. \({}^{*}\)Correspondence authors.

###### 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).

Towards, Machine Learning, ICML

## 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 performance 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 objective (e.g., sequence classification). Sanh et al. (2019), Sun et al. (2019), and Jiao et al. (2019) 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_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 representations on output. To achieve this, we employ the _interchange intervention training_ (IIT) method of Geiger et al. (2021b). 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 at some \(\langle\texttt{MASK}\rangle\)." 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 \(\langle\texttt{MASK}\rangle\) solid."? DIITO then pushes the student model to output the same logits as the teacher, i.e., matching the teacher's output distribution under

the counterfactual setup.

To assess the contribution of distillation with DIITO, we begin with BERTBASE(Devlin et al., 2019) and distill it under various alignments between student and teacher while pretraining on the WikiText-103M corpus (Merity et al., 2016) 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).

## 2 Related Work

Distillation was first introduced in the context of computer vision (Hinton et al., 2015) and has since been widely explored for language models (Sun et al., 2019; Sanh et al., 2019; Jiao et al., 2019). For example, Sanh et al. (2019) propose to extract information not only from output probabilities of the last layer in the teacher model, but also from intermediate layers in the fine-tuning stage. Recently, Rotman et al. (2021) 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 (Feder et al., 2021; Pryzant et al., 2021; Vig et al., 2020; Elazar et al., 2021; Giulianelli et al., 2020; Geiger et al., 2020, 2021a). Geiger et al. (2021b) extend these methods to network optimization. We contribute to this existing research by adapting intervention-based optimization to the task of language model distillation.

## 3 Causal Distillation

Here, we define our distillation training procedure. See Algorithm 1 for a summary.

**GetVals**. The GetVals operator is an activation-value retriever for a neural model. Given a neural model \({\cal M}\) containing a set of neurons \(\mathbf{N}\)

Figure 1: An IIT update in the context of masked language modelling (MLM). The teacher network (top) has 6 layers and the student (bottom) has 3 layers, and we align layer 2 in the student with layers 3–4 in the teacher. Solid lines are feed-forward connections, red lines show the flow of backpropagation, and dashed lines indicate interchange interventions. 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.

 

(an internal representations) and an appropriate input \({\bf x}\), GetVals\((\mathcal{M},{\bf x},{\bf N})\) is the set of values that \({\bf N}\) takes on when processing \({\bf x}\). In the case that \({\bf N}\) represents the neurons corresponding to the final output, GetVals\((\mathcal{M},{\bf x},{\bf N})\) is the output of model \(\mathcal{M}\) when processing \({\bf x}\) (i.e., output from a standard forward call of a neural model).

**SetVals**. The SetVals operator is a function generator that defines a new neural model with a computation graph that specifies an intervention on the original model \(\mathcal{M}\)Pearl (2009); Spirtes et al. (2001). SetVals\((\mathcal{M},{\bf N},{\bf v})\) is the new neural model where the neurons \({\bf N}\) are set to constant values \({\bf v}\). Because we overwrite neurons with \({\bf v}\) in-place, gradients can back-propagate through \({\bf v}\).

```
Require: Student model \(\mathcal{S}\), teacher model \(\mathcal{T}\), student output neurons \({\bf N_{\mathcal{S}}^{y}}\), alignment \(\Pi\), shuffled training dataset \(\mathcal{D}\).
1:\(\mathcal{S}\).train()
2:\(\mathcal{T}\).eval()
3:\(\mathcal{D}^{\prime}\) = random.shuffle(\(\mathcal{D}\))
4:\({\bf N_{\mathcal{T}}^{y}}=\Pi({\bf N_{\mathcal{S}}^{y}})\)
5:while not converged do
6:for\(\{{\bf x}_{1},{\bf y}_{1}\}\), \(\{{\bf x}_{2},{\bf y}_{2}\}\)in iter(\(\mathcal{D}\), \(\mathcal{D}^{\prime}\))do
7:\({\bf N_{\mathcal{S}}}=\text{sample\_student\_neurons}()\)
8:\({\bf N_{\mathcal{T}}}\) = \(\Pi({\bf N_{\mathcal{S}}})\)
9:with no_grad:
10:\(\mathcal{T}_{a}=\textsc{SetVals}(\)
11:\(\mathcal{T},{\bf N_{\mathcal{T}}},\textsc{GetVals}(\mathcal{T},{\bf x}_{1},{ \bf N_{\mathcal{T}}}))\)
12:\(o_{\mathcal{T}}=\textsc{GetVals}(\mathcal{T}_{a},{\bf x}_{2},{\bf N_{\mathcal{ T}}^{y}})\)
13:\(\mathcal{S}_{a}=\textsc{SetVals}(\)
14:\(\mathcal{S},{\bf N_{\mathcal{S}}},\textsc{GetVals}(\mathcal{S},{\bf x}_{1},{ \bf N_{\mathcal{S}}}))\)
15:\(o_{\mathcal{S}}=\textsc{GetVals}(\mathcal{S}_{a},{\bf x}_{2},{\bf N_{\mathcal{ S}}^{y}})\)
16:\(\mathcal{L}^{\textsc{DIITO}}=\textbf{get\_loss}(o_{\mathcal{T}},o_{\mathcal{S}})\)
17: Calculate \(\mathcal{L}_{\textsc{MLM}}\), \(\mathcal{L}_{\textsc{CE}}\), \(\mathcal{L}_{\textsc{Cos}}\)
18:\(\mathcal{L}=\mathcal{L}_{\textsc{MLM}}+\mathcal{L}_{\textsc{CE}}+\mathcal{ L}_{\textsc{Cos}}+\mathcal{L}^{\textsc{DIITO}}\)
19:\(\mathcal{L}\).backward()
20:Step optimizer
21:end while
```

**Algorithm 1 Causal Distillation via Interchange Intervention Training**

Intersubly sample a pair of examples from a training dataset \(({\bf x}_{1},{\bf y}_{1}),({\bf x}_{2},{\bf y}_{2})\in\mathcal{D}\). Next, where \({\bf N}\) is the set of neurons that we are targeting for intervention, we define \(\mathcal{M}_{\mathbf{N}}^{x_{1}}\) to abbreviate the new neural model as follows:

\[\textsc{SetVals}\big{(}\mathcal{M},{\bf N},\textsc{GetVals}(\mathcal{M},{\bf x }_{1},{\bf N})\big{)} \tag{1}\]

This is the version of \(\mathcal{M}\) obtained from setting the values of \({\bf N}\) to be those we get from processing input \({\bf x}_{1}\). The interchange intervention targeting \({\bf N}\) with \({\bf x}_{1}\) as the source input and \({\bf x}_{2}\) as the base input is then defined as follows:

\[\textsc{IntInv}(\mathcal{M},{\bf N},{\bf x}_{1},{\bf x}_{2}) \stackrel{{\mathrm{def}}}{{=}}\\ \textsc{GetVals}(\mathcal{M}_{\mathbf{N}}^{{\bf x}_{1}},{\bf x}_{ 2},{\bf N^{y}}) \tag{2}\]

where \({\bf N^{y}}\) are the output neurons. In other words, IntInv\((\mathcal{M},{\bf N},{\bf x}_{1},{\bf x}_{2})\) is the output state we get from \(\mathcal{M}\) for input \({\bf x}_{2}\) but with the neurons \({\bf N}\) set to the values obtained when processing input \({\bf x}_{1}\).

**DIITO**. DIITO employs \(\mathcal{T}\) as the teacher model, \(\mathcal{S}\) as the student model, \(\mathcal{D}\) as the training inputs to both models, and \(\Pi\) as an alignment that maps sets of student neurons to sets of teacher neurons. For each set of student neurons \({\bf N}_{\mathcal{S}}\) in the domain of \(\Pi\), we define DIITO loss as:

\[\mathcal{L}_{\textsc{CE}}^{\textsc{DIITO}~{}\mathrm{def}} \!\! 

accumulations over batches as in Sanh et al. (2019) 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 Devlin et al. (2019) and DistilBERT Sanh et al. (2019), 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 Sanh et al. (2019), 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 feed-forward layer of a Transformer block, as in Figure 1. We define \(L_{\mathcal{T}}\) and \(L_{\mathcal{S}}\) to be the number of layers in the student and teacher, respectively. In addition, we define \(\mathcal{S}^{j}_{i}\) and \(\mathcal{T}^{j}_{i}\) to be the representations in the \(i\)th row and \(j\)th column in the student and teacher, respectively. An alignment \(\Pi\) is a partial function from student representations to sets of teacher representations. We test three alignments:

* \(\Pi\) is defined on all student representations: \(\Pi(\mathcal{S}^{j}_{i})=\{\mathcal{T}^{j}_{4i+k}:0\leq k<L_{\mathcal{T}}/L_{ \mathcal{S}}\}\)
* \(\Pi\) is defined for the row \(L_{\mathcal{S}}~{}/~{}2\): \(\Pi(\mathcal{S}^{j}_{L_{\mathcal{S}}/2})=\{\mathcal{T}^{j}_{L_{\mathcal{T}}/2}\}\)
* \(\Pi\) is defined on the student representations in the first and second rows: \(\Pi(\mathcal{S}^{j}_{1})=\{\mathcal{T}^{j}_{L_{\mathcal{T}}-2}\}\) and \(\Pi(\mathcal{S}^{j}_{2})=\{\mathcal{T}^{j}_{L_{\mathcal{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 (**DIITOFULL**), random tokens (**DIITOFULL**+Random**) and masked tokens (**DIITOFULL**+Masked**). We also include \(\mathcal{L}^{\text{DIITO}}_{\text{Cos}}\) to the FULL alignment (**DIITOFULL**+\(\mathcal{L}^{\text{DIRTO}}_{\text{Cos}}\)).

## 5 Results

**Language Modeling**. We first evaluate our models using perplexity on the held-out evaluation data from WikiText. As shown in Table 1, DIITO brings performance gains for all alignments. Our best result is from the FULL alignment with the \(\mathcal{L}_{\text{Cos}}\) (**DIITOFULL**+\(\mathcal{L}^{\text{DIRTO}}_{\text{Cos}}\)), which has \(-2.24\) perplexity compared to standard DistilBERT trained with the same amount of data.

**GLUE**. The GLUE benchmark Wang et al. (2018) covers different natural language understanding tasks. We report averaged GLUE scores on the

\begin{table}
\begin{tabular}{l c c c c c c c c} \hline \hline  & \multicolumn{2}{c}{Pretraining} & WikiText & GLUE & CoNLL-2003 & SQuAD v1.1 \\ Model & Layers & Tokens & Perplexity & Score & ace & F1 & EM & Fl \\ \hline BERTBASEDevlin et al. (2019) & 12 & 3.3B & 10.27 (–)† & 82.75 (–) & 96.40 (–) & 92.40 (–) & 80.80 (–) & 88.50 (–) \\ (Wikipedia+BookCorpus) & & & & & & & & \\ \hline DistilBERT WikiText & 3 & 0.1B & 29.51 (0.32) & 67.42 (1.10) & 97.88 (0.04) & 88.89 (0.29) & 26.04 (0.93) & 68.38 (0.77) \\ DIITOFULLER WikiText & 3 & 0.1B & 26.04 (0.93) & 69.30 (1.08) & 98.03 (0.04) & 89.69 (0.18) & 58.74 (0.69) & 70.23 (0.57) \\ DIITOFULLER WikiText & 3 & 0.1B & 25.97 (0.63) & 69.01 (1.69) & 98.03 (0.03) & 89.82 (0.18) & 58.75 (0.49) & 70.21 (0.41) \\ DIITOFULL WikiText & 3 & 0.1B & 24.85 (0.58) & 69.36 (0.87) & 98.02 (0.03) & 89.67 (0.16) & 58.72 (0.67) & 70.50 (0.56) \\ \hline DistilBERT WikiText & 6 & 0.1B & 15.69 (0.11) & 75.80 (0.42) & 98.48 (0.03) & 92.12 (0.23) & 70.23 (0.75) & 79.99 (0.55) \\ DIITOFULLER WikiText & 6 & 0.1B & 14.32 (0.12) & 76.71 (0.47) & 98.56 (0.04) & 92.47 (0.19) & 71.93 (0.31) & 81.32 (0.23) \\ DIITOFULLER WikiText & 6 & 0.1B & 14.93 (0.23) & 76.80 (0.34) & 98.51 (0.02) & 92.36 (0.27) & 71.47 (0.28) & 81.01 (0.23) \\ DIITOFULL WikiText & 6 & 0.1B & 13.59 (0.25) & 76.67 (0.21) & 98.53 (0.04) & 92.35 (0.24) & 71.96 (0.29) & 81.33 (0.25) \\ \hline DIITOFULL+Random (WikiText) & 6 & 0.1B & 13.95 (0.18) & 76.84 (0.29) & 98.54 (0.03) & 92.41 (0.24) & 71.90 (0.54) & 81.27 (0.39) \\ DIITOFULL Masked (WikiText) & 6 & 0.1B & 13.99 (0.16) & 76.80 (0.32) & 98.55 (0.03) & 92.45 (0.18) & 71.77 (0.59) & 81.09 (0.42) \\ DIITOFULL+\(\mathcal{L}^{\text{DIRTO}}_{\text{Cos}}\) (WikiText) & 6 & 0.1B & 13.45 (0.19) & 77.14 (0.37) & 98.54 (0.04) & 92.35 (0.24) & 71.94 (0.31) & 81.35 (0.23) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Performance on the development sets of the WikiText, GLUE benchmark, CoNLL-2003 corpus for the name-entity recognition task, and SQuAD v1.1 for the question answering task. The score is the averaged performance scores with standard deviation (SD) for all tasks across 15 distinct runs. †Numbers are imputed from released models on Hugging Face Wolf et al. (2020).

 

development sets by fine-tuning our distilled models in Table 1. Individual task performance scores for each GLUE task are included in Table 2 in the Appendix. The results suggest that distilled models with DIITO lead to consistent improvements over standard DistilBERT trained under the same setting, with our best result (\({\bf DIITO_{FULL}+\mathcal{L}^{DIITO}_{Cos}}\)) being +1.77% higher.

**Named Entity Recognition**. We also evaluate our models on the CoNLL-2003 Named Entity Recognition task Tjong Kim Sang and De Meulder (2003). We report accuracy and Macro-F1 scores on the development sets. We fine-tune our models for three epochs. Our best performing model (\({\bf DIITO_{HIDDE}}\)) numerically surpasses not only standard DistilBERT (+0.38% on F1 score) trained under the same setting, but also its teacher, BERTBASE (+0.05% on F1 score). Though these improvements are small, in this case distillation produces a smaller model with _better performance_.

**Question Answering**. Finally, we evaluate on a question answering task, SQuAD v1.1 Rajpurkar et al. (2016). We report Exact Match and Macro-F1 on the development sets as our evaluation metrics. We fine-tune our models for two epochs. DIITO again yields marked improvements (Table 1). Our best result is from the vanilla \({\bf Full}\) alignment (\({\bf DIITO_{FULL}}\)), with +2.46% on standard DistilBERT trained under the same setting.

**Low-Resource Model Distillation** We experiment with an extreme case in a low-resource setting where we only distill with 15% of WikiText, keeping other experimental details constant. Our results suggest that DIITO training is also beneficial in extremely low-resource settings (Figure 2).

**Layer-wise Ablation** We further study the effect of DIITO training with respect to the size of the student model through a layer-wise ablation experiment. As shown in Figure 3, we compare GLUE performance for models trained with standard distillation pipeline and with DIITO training (\({\bf DIITO_{FULL}}\)). Specifically, we compute the averaged GLUE scores following the same procedure described in Section A.4. Our results suggest that DIITO training brings consistent improvements over GLUE tasks with smaller models booking the greatest gains.

## 6 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 Sanh et al. (2019) while consuming 97% less training data. These findings suggest that DIITO is a promising tool for effective model distillation.

## References

* S. Beckers, F. Eberhardt, and J. Y. Halpern (2020)Approximate causal abstractions. In Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, Vol. 115, Tel Aviv, Israel, pp. 606-615. External Links: Link Cited by: SS1.
* S. Beckers and J. Y. Halpern (2019)Abstracting causal models. Proceedings of the AAAI Conference on Artificial Intelligence33 (01), pp. 2678-2685. External Links: Document, ISSN 0037-649X Cited by: SS1.
* J. Devlin, M. Chang, K. Lee, and K. Toutanova (2019)BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for

Figure 3: GLUE score distribution across 15 distinct runs of students in different sizes. Following the evaluation for BERT Devlin et al. (2019). we exclude WNLI for evaluation.

Figure 2: Perplexity score distribution for the development set of WikiText of models trained in a low-resource setting. The best model is the one with the richest alignment structure.

 

Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_, pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
* Elazar et al. (2021) Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav Goldberg. 2021. Ammeseic probing: Behavioral explanation with amnesic counterfactuals. _Transactions of the Association for Computational Linguistics_, 9:160-175.
* Feder et al. (2021) Amir Feder, Nadav Oved, Uri Shalit, and Roi Reichart. 2021. CausaLM: Causal Model Explanation Through Counterfactual Language Models. _Computational Linguistics_, 47(2):333-386.
* Geiger et al. (2021a) Atticus Geiger, Hanson Lu, Thomas Icard, and Christopher Potts. 2021a. Causal abstractions of neural networks. In _Advances in Neural Information Processing Systems_, volume 34, pages 9574-9586. Curran Associates, Inc.
* Geiger et al. (2020) Atticus Geiger, Kyle Richardson, and Christopher Potts. 2020. Neural natural language inference models partially embed theories of lexical entailment and negation. In _Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP_, pages 163-173, Online. Association for Computational Linguistics.
* Geiger et al. (2021b) Atticus Geiger, Zhengxuan Wu, Hanson Lu, Josh Rozner, Elisa Kreiss, Thomas Icard, Noah D. Goodman, and Christopher Potts. 2021b. Inducing causal structure for interpretable neural networks. ArXiv:2112.00826.
* Giulianelli et al. (2020) Mario Giulianelli, Marco Del Tredici, and Raquel Fernandez. 2020. Analysing lexical semantic change with contextualised word representations. In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_, pages 3960-3973, Online. Association for Computational Linguistics.
* Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. _arXiv preprint arXiv:1503.02531_.
* Jiao et al. (2019) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2019. TinyBERT: Distilling BERT for natural language understanding. _arXiv:1909.10351_.
* Merity et al. (2016) Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture models. _arXiv:1609.07843_.
* Pearl (2009) Judea Pearl. 2009. _Causality: Models, Reasoning and Inference_, 2nd edition. Cambridge University Press, USA.
* Pryzant et al. (2021) Reid Pryzant, Dallas Card, Dan Jurafsky, Victor Veitch, and Dhanya Sridhar. 2021. Causal effects of linguistic properties. In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_, pages 4095-4109, Online. Association for Computational Linguistics.
* Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. SQuAD: 100,000+ questions for machine comprehension of text. In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing_, pages 2383-2392, Austin, Texas. Association for Computational Linguistics.
* Rotman et al. (2021) Guy Rotman, Amir Feder, and Roi Reichart. 2021. Model compression for domain adaptation through causal effect estimation. _arXiv:2101.07086_.
* Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. _arXiv:1910.01108_.
* Spirtes et al. (2001) Peter Spirtes, Clark N Glymour, and Richard Scheines. 2001. _Causation, Prediction, and Search_, 2nd edition. MIT Press.
* Sun et al. (2019) Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for BERT model compression. In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_, pages 4323-4332, Hong Kong, China. Association for Computational Linguistics.
* Sang and Meulder (2003) Erik F. Tjong Kim Sang and Fien De Meulder. 2003. Introduction to the CoNLL-2003 shared task: Language-independent named entity recognition. In _Proceedings of the Seventh Conference on Natural Language Learning at HLT-NAACL 2003_, pages 142-147.
* Vig et al. (2020) Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shiheber. 2020. Causal mediation analysis for interpreting neural NLP: The case of gender bias. _arXiv:2004.12265_.
* Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_, pages 353-355.
* Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Fantowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Phu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020. Transformers: State-of-the-art natural language processing. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_, pages 38-45, Online. Association for Computational Linguistics.

 

Appendix

### Standard Distillation Objectives

In our setting, our teacher model \({\mathcal{T}}\) is a BERT model, and our student model \({\mathcal{S}}\) is a shallower BERT model with fewer layers.

Assume that we randomly draw a training example \((\mathbf{x}_{1},\mathbf{y}_{1})\in{\mathcal{D}}\), where \(\mathbf{x}_{1}\) is the input to our models and \(\mathbf{y}_{1}\) is the corresponding ground truth (the token prediction at each masked position). We denote the model predictions (output logits) as \({\mathcal{T}}(\mathbf{x}_{1})\) and \({\mathcal{S}}(\mathbf{x}_{1})\). Additionally, we denote the contextualized representation for tokens for \(\mathbf{x}_{1}\) at the last layer as \(\text{BERT}_{\mathcal{T}}(\mathbf{x}_{1})\) and \(\text{BERT}_{\mathcal{S}}(\mathbf{x}_{1})\).

We adopt the three standard distillation objectives of Sanh et al. (2019):

* The masked language modeling loss of the student model calculated over all examples using the cross-entropy loss as follows: \[\sum_{\{\mathbf{x}_{1},\mathbf{y}_{1}\}\in{\mathcal{D}}}\text{CE}({\mathcal{ S}}(\mathbf{x}_{1}),\mathbf{y}_{1})\] (4)
* Following Hinton et al. (2015), the smoothed cross-entropy loss measuring the divergence between the student and teacher outputs as follows: \[\sum_{\mathbf{x}_{1}\in{\mathcal{D}}}\text{CE}_{\mathcal{S}}({\mathcal{S}}( \mathbf{x}_{1}),{\mathcal{T}}(\mathbf{x}_{1}))\] (5)
* The cosine embedding loss defined in terms of the final hidden states of the teacher and the student as follows: \[\sum_{\mathbf{x}_{1}\in{\mathcal{D}}}\text{Cos}(\text{BERt}_{\mathcal{S}}( \mathbf{x}_{1}),\text{BERt}_{\mathcal{T}}(\mathbf{x}_{1}))\] (6)

As a result, comparing to standard DistilBERT, DIITO essentially adds a new type of objective by pushing the student model to become a _causal abstraction_ of the teacher model.

### Causal Distillation Objectives

In addition to our causal loss \(\mathcal{L}^{\text{DIITO}}_{\text{CE}}\), we also propose a new loss \(\mathcal{L}^{\text{DIITO}}_{\text{Co}}\) which is identical to \(\mathcal{L}_{\text{Co}}\) with interchange interventions. In this section, we provide a formal definition for \(\mathcal{L}^{\text{DIRTO}}_{\text{Co}}\).

We denote our teacher and student models as \({\mathcal{T}}\) and \({\mathcal{S}}\) respectively. Using the notational conventions from Section3, we use \(\mathbf{N}^{\mathcal{V}}_{\mathcal{T}}\) and \(\mathbf{N}^{\mathcal{W}}_{\mathcal{S}}\) to represent the neurons corresponding to the final output for each model. Likewise, we use \(\mathbf{N}^{L_{\mathcal{T}}}_{\mathcal{T}}\) and \(\mathbf{N}^{L_{\mathcal{S}}}_{\mathcal{S}}\) to represent the neurons representing contextualized representation for each token after the final BERT layer.

Assuming we randomly sample a pair of examples from a training dataset \((\mathbf{x}_{1},\mathbf{y}_{1}),(\mathbf{x}_{2},\mathbf{y}_{2})\in{\mathcal{D}}\), we can then rewrite our causal loss \(\mathcal{L}^{\text{DIITO}}_{\text{CE}}\) by rearranging Eqn.2 and Eqn.3 as follows:

\[\sum_{\mathbf{x}_{1},\mathbf{x}_{2}\in{\mathcal{D}}}\text{CE}_{\mathcal{S}} \Big{(}\textsc{GetVals}({\mathcal{M}}^{\mathbf{x}_{1}}_{\mathbf{S}},\mathbf{ x}_{2},\mathbf{N}^{\mathcal{Y}}_{\mathcal{S}}), \tag{7}\]

where \({\mathcal{M}}^{\mathbf{x}_{i}}_{\mathbf{S}}\) and \({\mathcal{M}}^{\mathbf{x}_{i}}_{\mathbf{T}}\) are derived as in Eqn.1 for each model respectively. Crucially, Eqn.7 can be regarded as the _causal_ form of the standard smoothed cross-entropy loss with interchange intervention. Likewise, we can further define the \(\mathcal{L}^{\text{DIITO}}_{\text{Co}}\) as:

\[\sum_{\mathbf{x}_{1},\mathbf{x}_{2}\in{\mathcal{D}}}\text{Cos}\Big{(}\textsc{ GetVals}({\mathcal{M}}^{\mathbf{x}_{1}}_{\mathbf{S}},\mathbf{x}_{2}, \mathbf{N}^{L_{\mathcal{S}}}_{\mathcal{S}}), \tag{8}\]

with adjusted interchange alignments for \(\mathbf{N}^{L_{\mathcal{T}}}_{\mathcal{T}}\) and \(\mathbf{N}^{L_{\mathcal{S}}}_{\mathcal{S}}\).

### Evaluation Set-up

GlueWe fine-tune for 25 epochs for the smaller datasets (RTE and CoLA) and 3 epochs for the others. Following Devlin et al. (2019) and Sanh et al. (2019), we use Matthew's Correlation for CoLA, F1 for MRPC and QQP, Spearman correlation for STSWB, and accuracy for all the other tasks in GLUE.

### Reproducibility

To foster reproducibility and provide a fair comparison between methods, we distill BERT for each condition with three distinct random seeds. We then fine-tune each model with five distinct random seeds. Consequently, we report results aggregated from three distinct runs for the language modeling task, and 15 distinct runs for others.

Named Entity RecognitionWe follow the experimental set-up in the Hugging Face Wolf et al. (2020) repository for evaluation for the CoNLL-2003 Named Entity Recognition task Tjong Kim Sang and De Meulder (2003). For fine-tuning, we set the learning rate to \(5e^{-5}\) with an effective batch size of 32 for three epochs.2 

Question AnsweringWe use the experimental set-up of Sanh et al. (2019) for evaluation on SQuAD v1.1 Rajpurkar et al., 2016). For fine-tuning, we set the learning rate to \(3e^{-5}\) with an effective batch size of 48 for two epochs. We set the stride to 128.

\begin{table}
\begin{tabular}{l c c c c c c c c c c} \hline \hline  & \multicolumn{3}{c}{Pretraining} & \multicolumn{3}{c}{General Language Understanding Evaluation (GLUE)} \\ Model & Layers & Tokens & CoLA & MNLI & MRPC & QNLI & QQP & RTE & SST-2 & STS-B \\ \hline BERT\({}_{\text{BASE}}\)Devlin et al. (2019) & 12 & 3.3B & 56.30 & 84.70 & 88.60 & 91.80 & 89.60 & 69.30 & 92.70 & 89.00 \\ (Wikipedia+BookCorpus) & & & & & & & & & & \\ DistilBERT Sanh et al. (2019) & 6 & 3.3B & 51.30 & 82.10 & 87.50 & 89.20 & 88.50 & 59.90 & 91.30 & 86.90 \\ (Wikipedia+BookCorpus) & & & & & & & & & & \\ \hline \hline DIIBIBERT (WikiText) & 3 & 0.1B & 22.78 & 71.55 & 82.51 & 82.12 & 82.16 & 55.43 & 86.47 & 56.33 \\ DIITO\({}_{\text{MEDLE}}\) (WikiText) & 3 & 0.1B & 23.21 & 72.97 & 82.81 & 83.15 & 82.83 & 55.98 & 86.52 & 66.93 \\ DIITO\({}_{\text{ALTE}}\) (WikiText) & 3 & 0.1B & 24.12 & 72.80 & 82.16 & 82.88 & 82.85 & 57.29 & 87.31 & 62.65 \\ DIITO\({}_{\text{FULL}}\) (WikiText) & 3 & 0.1B & 25.01 & 72.85 & 82.71 & 83.05 & 82.85 & 55.37 & 86.92 & 66.15 \\ \hline DistilBERT (WikiText) & 6 & 0.1B & 40.43 & 78.95 & 87.45 & 84.76 & 84.96 & 60.10 & 89.38 & 80.40 \\ DIITO\({}_{\text{MDLE}}\) (WikiText) & 6 & 0.1B & 43.97 & 79.47 & 87.57 & 85.45 & 85.21 & 60.72 & 89.97 & 81.33 \\ DIITO\({}_{\text{DATE}}\) (WikiText) & 6 & 0.1B & 43.93 & 79.49 & 87.70 & 85.79 & 85.22 & 60.14 & 90.31 & 81.79 \\ DIITO\({}_{\text{FULL}}\) (WikiText) & 6 & 0.1B & 43.43 & 79.66 & 88.17 & 85.57 & 85.28 & 59.95 & 90.01 & 81.26 \\ \hline DIITO\({}_{\text{FULL}}\)+Random (WikiText) & 6 & 0.1B & 44.27 & 79.70 & 88.06 & 85.63 & 85.34 & 60.89 & 89.76 & 81.08 \\ DIITO\({}_{\text{FULL}}\)+Masked (WikiText) & 6 & 0.1B & 43.39 & 79.63 & 87.88 & 85.61 & 85.30 & 61.06 & 89.97 & 81.58 \\ DIITO\({}_{\text{FULL}}\)+\({}_{\text{Co}}^{\text{DITO}}\) (WikiText) & 6 & 0.1B & 45.17 & 79.68 & 88.18 & 85.83 & 85.31 & 60.94 & 90.32 & 81.69 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Model performance results on the development sets of the GLUE benchmark. The GLUE score is the averaged performance scores across 15 distinct runs with precision aligned for a fair comparison. Following the evaluation for BERT Devlin et al. (2019), we exclude WNLI for evaluation.