Ensemble Distillation Approaches for Grammatical Error Correction

For sequence models, distillation can be performed in multiple ways [14, 24]. The approach adopted in this work is token-level knowledge distillation, and is one of the simplest methods. The teacher and student $\bm{\phi}$ use the same reference back-history ${\bm{y}}_{<l}$ (teacher-forcing) and input $\bm{x}$. The KL-divergence between the ensemble and student token-level categorical distributions is then minimised:

$\displaystyle\mathcal{L}(\bm{\phi})=\frac{1}{L}\sum_{l=1}^{L}{\tt KL}\bigg{(}\left(\frac{1}{M}\sum_{m=1}^{M}{\bm{\pi}}_{l}^{(m)}\right)\;\bigg{\|}\;{\tt P}\big{(}y_{l}|\bm{y}_{<l},\bm{x};\bm{\phi}\big{)}\bigg{)}$

(2)

Extending distribution distillation described in the previous section, this section introduces distribution distillation for sequence models. For token-level distribution distillation, the student with parameters $\bm{\lambda}$ predicts the Dirichlet distribution with parameters $\bm{\alpha}_{l}$ for the $l$-th token:

$\displaystyle{\tt P}(\bm{\pi}|\bm{y}_{<l},\bm{x};\bm{\lambda})={\tt Dir}(\bm{\pi};\bm{\alpha}_{l}),\medspace\bm{\alpha}_{l}=\bm{f}(\bm{y}_{<l},\bm{x};\bm{\lambda})$

Given $\{\bm{\pi}^{(1:M)}_{l}\}_{l=1}^{L}$ from the sequence ensemble, the distribution distilled model can be trained using negative log-likelihood (NLL):

$\displaystyle\mathcal{L}(\bm{\lambda})=-\frac{1}{ML}\sum_{l=1}^{L}\sum_{m=1}^{M}\ln\bigg{(}{\tt Dir}(\bm{\pi}_{l}^{(m)};\bm{\alpha}_{l})\bigg{)}$

(3)

Eq. (3) optimizes the parameters of the distribution distilled model directly from the ensemble predictions for each time instance. Though this yields an appropriate criterion to find ${\bm{\lambda}}$, predicting the distribution over the ensemble members for all back-histories, it may be very challenging to optimise the network parameters. To simplify this optimization, a two stage approach may be adopted. First, the predictions of the ensemble for each back-history is modelled by a Dirichlet distribution with parameters ${\tilde{\bm{\alpha}}}_{l}$ where

$\displaystyle\bm{\tilde{\alpha}}_{l}=\operatorname*{arg\,max}_{\bm{\alpha}}\bigg{\{}\frac{1}{M}\sum_{m=1}^{M}\ln\Big{(}{\tt Dir}(\bm{\pi}_{l}^{(m)};\bm{\alpha})\Big{)}\bigg{\}}$

The distribution distillation parameters are now trained to minimize the KL-divergence between this ensemble Dirichlet and the predicted distilled Dirichlet:

$\displaystyle\mathcal{L}(\bm{\lambda})=\frac{1}{L}\sum_{l=1}^{L}{\tt KL}\Big{(}{\tt Dir}(\bm{\pi};\bm{\tilde{\alpha}}_{l})\;\Big{\|}\;{\tt Dir}(\bm{\pi};\bm{\alpha}_{l})\Big{)}$

(4)

This has the same general form as eq. (2) but is based on the KL-divergence between distributions of distributions, rather than just the expected posterior distribution based on the ensemble.

Once a distribution distilled sequence model has been obtained, the probability of predicting class $c$ is then:

$\displaystyle{\tt P}(y=\omega_{c}|\bm{y}_{<l},\bm{x};\bm{\lambda})=\mathbb{E}_{{\tt P}(\bm{\pi}|\bm{y}_{<l},\bm{x};\bm{\lambda})}\Big{[}{\tt P}(y=\omega_{c}|\bm{\pi})\Big{]}=\frac{\alpha_{l,c}}{\alpha_{l,0}}$

as would be expected when the student parametrises a Dirichlet. This shows that achieving sequence based distribution distillation is a straightforward generalization of both ensemble distribution and sequence distillation.

Both eq. (3) and (4) result in viable objectives for obtaining a distribution distilled sequence model $\bm{\lambda}$, and from which uncertainty metrics can be derived. Instead of calculating entropy over the classes as in eq. (1), uncertainties for sequence models have to enumerate all possible sequences $\bm{y}\in\mathcal{Y}$. Let $\bm{\Pi}=\{\bm{\pi}_{l}\}_{l=1}^{L}$ be the sequence of categorical distributions from which the sequence $\bm{y}$ is generated:

$\displaystyle{\tt P}(\bm{\Pi}|\bm{x};\bm{\lambda})=\prod_{l=1}^{L}{\tt Dir}(\bm{\pi}_{l};\bm{\alpha}_{l}),\medspace{\tt P}(\bm{y}|\bm{\Pi})=\prod_{l=1}^{L}{\tt P}(y_{l}|\bm{\pi}_{l})$

The sequence uncertainties for the student $\bm{\lambda}$ can be expressed as:

$\displaystyle\underbrace{\mathcal{I}[\bm{y},\bm{\Pi}|\bm{x};\bm{\lambda}]}_{\begin{subarray}{c}Knowledge\\ Uncertainty\end{subarray}}=\underbrace{\mathcal{H}\big{[}{\tt P}(\bm{y}|\bm{x};\bm{\lambda})\big{]}}_{\begin{subarray}{c}Total\\ Uncertainty\end{subarray}}-\underbrace{\mathbb{E}_{{\tt P}(\bm{\Pi}|\bm{x};\bm{\lambda})}\big{[}\mathcal{H}[{\tt P}(\bm{y}|\bm{\Pi})]\big{]}}_{\begin{subarray}{c}Expected\ Data\\ Uncertainty\end{subarray}}$

(5)

However, as noted in [15], calculating the uncertainties require the intractable computation of enumerating all $\bm{y}\in\mathcal{Y}$. Instead the same approximations made in [15] will be utilized here. Given $S$ sampled sequences: $\forall\bm{y}_{<l}^{(s)}\subset\bm{y}^{(s)}\sim{\tt P}(\bm{y}|\bm{x},\bm{\lambda})$, the following Monte-Carlo approximations can be made:

	$\displaystyle\mathcal{H}\big{[}{\tt P}(\bm{y}|\bm{x};\bm{\lambda})\big{]}\approx$	$\displaystyle\frac{1}{S}\sum_{s,l}\mathcal{H}\big{[}{\tt P}(y_{l}|\bm{y}_{<l}^{(s)},\bm{x};\bm{\lambda})\big{]}$		(6)
	$\displaystyle\mathbb{E}_{{\tt P}(\bm{\Pi}|\bm{x};\bm{\lambda})}\big{[}\mathcal{H}[{\tt P}(\bm{y}|\bm{\Pi})]\big{]}\approx$	$\displaystyle\frac{1}{S}\sum_{s,l}\mathbb{E}_{{\tt P}(\bm{\pi}|\bm{y}_{<l}^{(s)},\bm{x};\bm{\lambda})}\big{[}\mathcal{H}[{\tt P}(y_{l}|\bm{\pi})]\big{]}$		(7)

The sequence-level uncertainties, or rates, in [15] may be obtained by normalising the quantities in eqs. (6) and (7) by the sequence length.

Training distribution distilled models can be significantly harder than training to standard distillation. Hence, a two model approach will also be explored. This is based on the observation that the distribution distilled model tends to have high Spearman’s rank correlation with the ensemble predicted uncertainties in teacher-forcing mode used in training. In contrast, evaluating the model in free-run decoding, so predictions are based on ${\hat{\bm{y}}}_{<l}$, the same consistency between the model and the ensemble was not observed. To address this, the distribution distilled model was fed the back-history from the distilled model, as the distilled model was found to be less sensitive to the teacher-forcing and free-running mismatch.

Assuming we have a distilled model $\bm{\phi}$ obtained from eq. (2), and a distribution distilled model $\bm{\lambda}$ obtained from either eq. (3) or (4), one can then perform free-run decoding according to:

	$\displaystyle\bm{\hat{\pi}}_{l}=\bm{f}(\bm{\hat{y}}_{<l},\bm{x};\bm{\phi});\>\>\>\>\hat{y}_{l}\sim{\tt P}(y_{l}|{\hat{\bm{y}}}_{<l},{\bm{x}};{\bm{\phi}})$
	$\displaystyle\hat{\bm{\alpha}_{l}}=\bm{f}(\bm{\hat{y}}_{<l},\bm{x};\bm{\lambda})$

The distilled model $\bm{\phi}$ is used to predict the output sequence, which then guides the second model $\bm{\lambda}$ to return higher quality uncertainties that can be derived from $\{\hat{\bm{\alpha}_{l}}\}_{l=1}^{L}$. Although this method (referred to as guided uncertainty approach; GUA) does not yield the same efficiency as standard distribution distillation, it ensures that the best attributes of $\bm{\phi}$ and $\bm{\lambda}$ are maintained in testing.

The data used in this paper is the same as those described in [25]. This data has been manually annotated with grammatical errors and corrections. The training set and FCE (a specific test set) have been taken from the Cambridge Learner Corpus [26] and includes written exams of candidates with different L1 languages. The FCE test set is a held-out subset of the corpus and therefore, in-domain (ID) with the training data. NICT-JLE [27] is a public speech corpus based on non-native speech. Only manual transcriptions of these interviews involving Japanese learners are available, along with the grammatical corrections. No audio is available. BULATS [28, 29] is a corpus based on a free speaking business English test, where candidates were prompted for responses up to 1 minute. Both manual (BLT) and ASR (BLT${}_{\text{asr}}$) transcriptions were used, the average ASR WER for this data was 19.5% (see [30] for details). The candidates are drawn from 6 L1s and span the full range of proficiencies in English. These sets are out-of-domain (OOD) and will be used to test the performance of any uncertainties derived from distribution distilled models and ensembles, see Table (1).

All models used in this work are transformers based on [31], with the default parameters of 6 layers, $d_{model}=512$, $d_{ff}=2048$, 8 heads, and 10% dropout. Input words were mapped to randomly initialized embeddings of dimension 512. Random seeds were used to generate an ensemble of 5 models $M=5$ [10]. As 5 models are used in the ensemble, this increases the memory requirements by a factor 5 over a single model, and approximately 5 for the decoding cost depending on the form of beam-search being used. It is possible to utilize more advanced ensemble generation methods, but is not key for this work [19, 20, 32]. Additionally, sequence models can make predictions in different ways: expectation of products, and products of expectations [15]. Since it has been found in [15] that products of expectations performs better, it will be adopted in this work when evaluating ensembles. Beam-search with a beam of one will be used, and uncertainty metrics will be evaluated on this single output.

As noted in [12], when training, the target categorical is often concentrated in one of the classes, while the Dirichlet predicted by the student often has a distribution spread over many classes. This implies that the common support between the teacher and student is poor. Optimizing KL-divergence when there is a small non-zero common support between target and prediction can lead to instability, and therefore, temperature annealing will be used to alleviate this [2]. In this work, when training a distribution distilled system, the targets from the ensemble are annealed with a temperature $T$, following a schedule from $T=10.0$ to $3.0$. Further reducing the temperature down to $T=1.0$ resulted in instability during training. To remain consistent, all results concerning ensembles will also be based on the final temperature $T=3.0$.

The GEC performance metric used in this work is GLEU [33], and uncertainty metrics will be based at the token level. Furthermore, the distribution distilled models and ensembles will be compared when a fraction of samples are rejected (based on some metric) and passed on to an oracle. In these cases, the rejection will be based on the highest sentence level uncertainty metrics, and will be compared to simply rejecting sentences based on their length. An additional metric based on knowing the true target will be ${\tt L}\cdot(1-{\tt GLEU})$ (referred to as manual), where ${\tt L}$ is the true sequence length; the inclusion of length is important as sentence level GLEU does not take length into account even though it has an effect when used at a dataset level [34]. System performance will be evaluated using relative area under the curve ${\tt AUC_{RR}}$ as defined by:

$\displaystyle{\tt AUC_{RR}}=\frac{{\tt AUC}-{\tt AUC}_{\tt random}}{{\tt AUC}_{\tt manual}-{\tt AUC}_{\tt random}}$

(8)

following the same reasoning as in [35] to simplify comparison between metrics; ${\tt AUC}_{\tt random}$ refers to fully random rejection. Finally GLEU performance will also be reported at 10% rejection.

Table 2 shows the GLEU performance of a range of different models. As expected the ensemble performs best, showing performance gains over the individual ensemble members. Distilling the ensemble, using eq. (2), again yielded performance gains over the individual ensemble members, though not to the same level as the original ensemble but at reduced memory and computational cost.

Secondly, the table also shows the performance of distribution distilled models. The simplified training using KL-divergence, eq. (4), outperforms NLL, eq. (3), though neither matches the average performance of the ensemble members. This illustrates the challenges in training sequence ensemble distribution distillation models. As the KL model performed better, it will be used as the uncertainty model in GUA, though both the NLL and KL based models yielded high Spearman rank correlation coefficient with the ensemble uncertainty measures. The same trend can be seen for both the in-domain data (FCE) and the out-of-domain data.

Table 3 shows the average word-level uncertainties predicted for two datasets, FCE and the most mismatched BLT${}_{\text{asr}}$.

The guided model approach behaves similarly to the ensemble. High knowledge uncertainty (KU) is, according to theory, an indication of a sample being OOD and in this case as expected KU is higher for BLT${}_{\text{asr}}$—an out of-of-domain test set. Furthermore, it can be seen that all uncertainties for BLT${}_{\text{asr}}$ are higher than FCE across both models.

Table 4 shows relative AUC (AUC_RR) for FCE and BLT${}_{\text{asr}}$, as well as mix of the two which assess whether the uncertainty measures can detect the more challenging OOD speech data. It is interesting that as a baseline simply using the length of the of the output sequence (Length) is a good baseline, as longer sentences will tend to have more complex grammatical structure and opportunity for mistakes. KU performs best in rejecting challenging inputs. for both the ensemble and GUA, outperforming simple length selection for all conditions.

Table 5 shows the performance when a fixed percentage, 10%, is ”rejected” and manually corrected. These results can be compared to the baseline numbers in Table 2, show significant gains from rejecting 10% of the data. Similar trends to that shown in AUC_RR can be seen, KU generally yields the best performance, outperforming the baseline length approach in all conditions. GUA yields similar performance to the original ensemble, but at reduced memory and computational costs.