SlowGuess/NAACL_2022_OCR huggingface dataset:
Zhengxuan Wu\*, Atticus Geiger\*, Joshua Rozner, Elisa Kreiss, Hanson Lu
Thomas Icard, Christopher Potts, Noah D. Goodman
{wuzhengx, atticusg}@stanford.edu
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 CoNLL2003 (named entity recognition).
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 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
*Equal contribution. †Correspondence authors.
Human Language Technologies, pages 4288 - 4295
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.
the counterfactual setup.
To assess the contribution of distillation with DIITO, we begin with  $\mathrm{BERT}_{\mathrm{BASE}}$  (Devlin et al., 2019) and distill it under various alignments between student and teacher while pretraining on the WikiText-103M corpus (Merit 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).
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 in-
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; Giulianielli 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  $\mathcal{M}$  containing a set of neurons  $\mathbf{N}$
Algorithm 1 Causal Distillation via Interchange Intervention Training
Require: Student model  $S$ , teacher model  $\mathcal{T}$ , student output neurons  $\mathbf{N}_{S}^{\mathbf{y}}$ , alignment  $\Pi$ , shuffled training dataset  $\mathcal{D}$ .
1:  $S$  .train()
2:  $\mathcal{T}$  .eval()
3:  $\mathcal{D}' =$  randomshuffle  $(\mathcal{D})$
4:  $\mathbf{N}_{\mathcal{T}}^{\mathbf{y}} = \Pi (\mathbf{N}_{\mathcal{S}}^{\mathbf{y}})$
5: while not converged do
6: for  $\{\mathbf{x}_1,\mathbf{y}_1\} ,\{\mathbf{x}_2,\mathbf{y}_2\}$  in iter  $(\mathcal{D},\mathcal{D}^{\prime})$  do
7:  $\mathbf{N}_S =$  sample/student_neurons()
8:  $\mathbf{N}_{\mathcal{T}} = \Pi (\mathbf{N}_{\mathcal{S}})$
9: with no_grad:
10:  $\mathcal{T}_a = \mathrm{SETVALS}($
11:  $\mathcal{T},\mathbf{N}_{\mathcal{T}},\mathrm{GETVALS}(\mathcal{T},\mathbf{x}_1,\mathbf{N}_{\mathcal{T}}))$
12:  $o_{\mathcal{T}} = \mathrm{GETVALS}(\mathcal{T}_a,\mathbf{x}_2,\mathbf{N}_{\mathcal{T}}^{\mathbf{y}})$
13:  $S_{a} = \mathsf{SETVALS}$
14:  $\mathcal{S},\mathbf{N}_{\mathcal{S}},\mathbf{GETVALS}(\mathcal{S},\mathbf{x}_1,\mathbf{N}_{\mathcal{S}})$
15:  $o_{\mathcal{S}} = \mathrm{GETVALS}(S_a,\mathbf{x}_2,\mathbf{N}_{\mathcal{S}}^{\mathbf{y}})$
16:  $\mathcal{L}^{\mathrm{DIITO}} = \mathbf{get\_loss}(o_{\mathcal{T}}, o_{\mathcal{S}})$
17: Calculate  $\mathcal{L}_{\mathrm{MLM}},\mathcal{L}_{\mathrm{CE}},\mathcal{L}_{\mathrm{Cos}}$
18:  $\mathcal{L} = \mathcal{L}_{\mathrm{MLM}} + \mathcal{L}_{\mathrm{CE}} + \mathcal{L}_{\mathrm{Cos}} + \mathcal{L}^{\mathrm{DIITO}}$
19: L.backup()
20: Step optimizer
21: end while
(an internal representations) and an appropriate input  $\mathbf{x}$ , GETVALS  $(\mathcal{M},\mathbf{x},\mathbf{N})$  is the set of values that  $\mathbf{N}$  takes on when processing  $\mathbf{x}$ . In the case that  $\mathbf{N}$  represents the neurons corresponding to the final output, GETVALS  $(\mathcal{M},\mathbf{x},\mathbf{N})$  is the output of model  $\mathcal{M}$  when processing  $\mathbf{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},\mathbf{N},\mathbf{v})$  is the new neural model where the neurons  $\mathbf{N}$  are set to constant values  $\mathbf{v}$ . Because we overwrite neurons with  $\mathbf{v}$  in-place, gradients can back-propagate through  $\mathbf{v}$ .
Interchange Intervention. An interchange intervention combines GETVALS and SETVALS operations. First, 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}$ . Next, where  $\mathbf{N}$  is the set of neurons that we are targeting for intervention, we define  $\mathcal{M}_{\mathbf{N}}^{\mathbf{x}_1}$  to abbreviate the new neural model as follows:
$$
\operatorname {S E T V A L S} \left(\mathcal {M}, \mathbf {N}, \operatorname {G E T V A L S} \left(\mathcal {M}, \mathbf {x} _ {1}, \mathbf {N}\right)\right) \tag {1}
$$
This is the version of  $\mathcal{M}$  obtained from setting the values of  $\mathbf{N}$  to be those we get from processing input  $\mathbf{x}_1$ . The interchange intervention targeting  $\mathbf{N}$  with  $\mathbf{x}_1$  as the source input and  $\mathbf{x}_2$  as the base input is then defined as follows:
$$
\begin{array}{l} \operatorname {I N T I N V} \left(\mathcal {M}, \mathbf {N}, \mathbf {x} _ {1}, \mathbf {x} _ {2}\right) \stackrel {\text {d e f}} {=} \\ \operatorname {G E T V A L S} \left(\mathcal {M} _ {\mathbf {N}} ^ {\mathbf {x} _ {1}}, \mathbf {x} _ {2}, \mathbf {N} ^ {\mathbf {y}}\right) \tag {2} \\ \end{array}
$$
where  $\mathbf{N}^{\mathbf{y}}$  are the output neurons. In other words, INTINV  $(\mathcal{M},\mathbf{N},\mathbf{x}_1,\mathbf{x}_2)$  is the output state we get from  $\mathcal{M}$  for input  $\mathbf{x}_2$  but with the neurons  $\mathbf{N}$  set to the values obtained when processing input  $\mathbf{x}_1$ .
DIITO. DIITO employs  $\mathcal{T}$  as the teacher model,  $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  $\mathbf{N}_S$  in the domain of  $\Pi$ , we define DIITO loss as:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {C E}} ^ {\mathrm {D I I T O} \stackrel {\mathrm {d e f}} {=}} \\ \sum_ {\mathbf {x} _ {1}, \mathbf {x} _ {2} \in \mathcal {D}} \operatorname {C E} _ {\mathrm {S}} \left( \begin{array}{l} \operatorname {I N T I N V} (\mathcal {S}, \mathbf {N} _ {\mathcal {S}}, \mathbf {x} _ {1}, \mathbf {x} _ {2}), \\ \operatorname {I N T I N V} (\mathcal {T}, \Pi (\mathbf {N} _ {\mathcal {S}}), \mathbf {x} _ {1}, \mathbf {x} _ {2}) \end{array} \right) \tag {3} \\ \end{array}
$$
where  $\mathrm{CE_S}$  is the smoothed cross-entropy loss measuring the divergences of predictions, under interchange, between the teacher and the student model.
Distillation Objectives. We adopt the standard distillation objectives from DistilBERT (Sanh et al., 2019) (defined formally in Appendix A.1):  $\mathcal{L}_{\mathrm{MLM}}$  for the task-specific loss for the student model,  $\mathcal{L}_{\mathrm{CE}}$  for the loss measuring the divergence between the student and teacher outputs on masked tokens, and  $\mathcal{L}_{\mathrm{Cos}}$  for the loss measuring the divergence between the student and teacher contextualized representations on masked tokens in the last layer. Our final training objective for the student is a linear combination of the four training objectives reviewed above:  $\mathcal{L}_{\mathrm{MLM}}, \mathcal{L}_{\mathrm{CE}}, \mathcal{L}_{\mathrm{Cos}},$  and  $\mathcal{L}_{\mathrm{CE}}^{\mathrm{DIITO}}$ . In a further experiment, we introduce a fifth objective  $\mathcal{L}_{\mathrm{Cos}}^{\mathrm{DIITO}}$  which is identical to  $\mathcal{L}_{\mathrm{Cos}}$ , except the teacher and student are undergoing interchange interventions (see Appendix A.2 for details).
We adapt the open-source Hugging Face implementation for model distillation (Wolf et al., 2020).<sup>1</sup> We distill our models on the MLM pretraining task (Devlin et al., 2019). We use large gradient
1https://github.com/huggingface/transformers
Model	Layers	Pretraining Tokens	WikiText Perplexity	GLUE Score	CoNLL-2003 acc	F1	SQuAD v1.1 EM	F1
BERTBASE (Devlin et al., 2019) (Wikipedia+BookCorpus)	12	3.3B	10.27 (−)†	82.75 (−)	96.40 (−)	92.40 (−)	80.80 (−)	88.50 (−)
DistilBERT (Sanh et al., 2019) (Wikipedia+BookCorpus)	6	3.3B	17.48 (−)†	79.59 (−)	98.39 (−)†	93.10 (−)†	77.70 (−)	85.80 (−)
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)
DIITOMIDDLE (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)
DIITOLATE (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)
DistilBERT (WikiText)	6	0.1B	15.69 (1.51)	75.80 (0.42)	98.48 (0.03)	92.12 (0.23)	70.23 (0.75)	79.99 (0.55)
DIITOMIDDLE (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)
DIITOLATE (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)
DIITOTFULL (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)
DIITOTFULL+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)
DIITOTFULL+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)
DIITOTFULL+LCOS+DIITO( 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)
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).
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 feedforward 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}_i^j$  and  $\mathcal{T}_i^j$  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:
FULL II is defined on all student representations:
$$
\Pi (\mathcal {S} _ {i} ^ {j}) = \{\mathcal {T} _ {4 i + k} ^ {j}: 0 \leq k <   L _ {\mathcal {T}} / L _ {\mathcal {S}} \}
$$
MIDDLE  $\Pi$  is defined for the row  $L_{\mathcal{S}} / / 2$  ..  $\Pi (\mathcal{S}_{L_{\mathcal{S} / / 2}}^{j}) = \{\mathcal{T}_{L_{\mathcal{T} / / 2}}^{j}\}$
LATE II is defined on the student representations in the first and second rows:
$$
\Pi (\mathcal {S} _ {1} ^ {j}) = \{\mathcal {T} _ {L \mathcal {T} - 2} ^ {j} \} \text {a n d} \Pi (\mathcal {S} _ {2} ^ {j}) = \{\mathcal {T} _ {L \mathcal {T} - 1} ^ {j} \}
$$
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  $(\mathsf{DIITO}_{\mathsf{FULL}})$ , random tokens  $(\mathsf{DIITO}_{\mathsf{FULL}} + \mathsf{Random})$  and masked tokens  $(\mathsf{DIITO}_{\mathsf{FULL}} + \mathsf{Masked})$ . We also include  $\mathcal{L}_{\mathrm{Cos}}^{\mathrm{DIITO}}$  to the FULL alignment  $(\mathsf{DIITO}_{\mathsf{FULL}} + \mathcal{L}_{\mathrm{Cos}}^{\mathsf{DIITO}})$ .
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}_{\mathrm{Cos}}$  ( $\mathbf{DIITO}_{\mathrm{FULL}} + \mathcal{L}_{\mathrm{Cos}}^{\mathrm{DIITO}}$ ), 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
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.
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  $(\mathbf{D}\mathbf{I}\mathbf{T}\mathbf{O}_{\mathrm{FULL}} + \mathcal{L}_{\mathrm{Cos}}^{\mathbf{D}\mathbf{I}\mathbf{T}\mathbf{O}})$  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  $(\mathbf{D}\mathbf{I}\mathbf{T}\mathbf{O}_{\mathbf{M}\mathbf{i}\mathbf{D}\mathbf{L}\mathbf{E}})$  numerically surpasses not only standard DistilBERT  $(+0.38\%$  on F1 score) trained under the same setting, but also its teacher,  $\mathrm{BERT_{BASE}}$ $(+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 MacroF1 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 FULL alignment  $(\mathbf{D}\mathbf{I}\mathbf{T}\mathbf{O}_{\mathbf{F}\mathbf{U}\mathbf{L}})$ , 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 stan
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.
dard distillation pipeline and with DIITO training  $(\mathbf{D}\mathbf{I}\mathbf{T}\mathbf{O}_{\mathbf{F U L L}})$ . 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.
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.
Sander Beckers, Frederick Eberhardt, and Joseph Y. Halpern. 2020. Approximate causal abstractions. In Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, pages 606-615, Tel Aviv, Israel. PMLR.
Sander Beckers and Joseph Y. Halpern. 2019. Abstracting causal models. Proceedings of the AAAI Conference on Artificial Intelligence, 33(01):2678-2685.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina 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