Accurate Knowledge Distillation via n-best Reranking (2024)

Abstract: We propose utilizing n-best reranking to enhance Sequence-Level Knowledge Distillation (Kim and Rush, 2016) where we extract pseudo-labels for student model’s training data from top n-best hypotheses and leverage a diverse set of models with different inductive biases, objective functions or architectures, including some publicly-available large language models, to pick the highest-quality hypotheses as labels. The effectiveness of our proposal is validated through experiments on the WMT’21 German ↔ English and Chinese ↔ English translation tasks. Our results demonstrate that utilizing pseudo-labels generated by our n-best reranker leads to a significantly more accurate student model. In fact, our best student model achieves comparable accuracy to a large translation model from (Tran et al., 2021) with 4.7 billion parameters, while having two orders of magnitude fewer parameters.

## introduction knowledge distillation (kd) @xcite plays a pivotal role in many machine learning tasks including neural machine translation (nmt). this is evident in recent translation evaluations @xcite @xcite , where the majority of submissions incorporate kd into their training pipelines. central to this paper is kd's primary strength, which lies in its ability to utilize a larger teacher model to train a smaller student model effectively. consequently, the accuracy of the teacher model correlates with that of the student model. one well-known yet simple approach for enhancing the accuracy of machine learning models involves model ensembling @xcite , as also has been applied by @xcite for kd. in this approach, the underlying models are typically trained on the same datasets but with varying random initializations. many existing works applying kd to nmt have adopted this approach, including its dominant variant, namely sequencelevel kd @xcite , which deploys an ensemble of teacher models to generate the pseudolabels of the student models' training data. to maintain the simplicity of the inference process when generating the labels, the sequence-level kd often imposes the requirements that the underlying models share the same vocabulary and network architecture. these factors, unfortunately, restrict the types of models that can be included in the ensemble, thereby limiting the avenues for improving the teacher accuracy. we introduce an n-best reranking approach to extend the vanilla sequence-level kd to allow kd to benefit from a more diverse type of models. nbest reranking is a well-known and powerful approach for boosting translation accuracy, as evident in @xcite @xcite just to name a few. unfortunately, the associated inference cost can be too high for this approach to be deployed in online environments with tight latency constraint. by applying n-best reranking for distillation, the elevated computational cost is shifted to training stage without affecting the latency of the deployed student model, thus sharing a similar motivation as @xcite . our proposed approach involves a two-step process. the first step involves generating a highquality n-best list from each source sentence. our initial study indicates a potential gain of almost 10 bleu on our validation set if we consider hypotheses beyond top-1. the final decision on which hypothesis to be selected is deferred to the second step where we take advantage of models with various architectures, inductive biases, sources of training data and objective functions to rerank the n-best list. to increase model diversity further, we also incorporate open-source large pretrained models in our experiments, tapping into their availability and increased translation capabilities. this flexibility stands in contrast to the rigid type of teacher models deployed by the vanilla sequence-level kd. we showcase our reranker's effectiveness in two scenarios, namely the traditional and iterative kds. in the former, pseudo-labels are directly utilized for training student models, while the latter, often dubbed as self-training, involves an extra step of iteratively retraining the teacher models using the pseudo-labels @xcite . in the first scenario, we thoroughly examine the accuracy of student models trained with pseudo-labels generated by our reranker, comparing them to models trained with pseudo-labels from the vanilla sequence-level kd. in the second scenario, we investigate whether our reranker can yield a valuable cascading effect by progressively enhancing teacher models using the pseudo labels from teacher models from the previous iteration. these improved teacher models, in turn, may generate superior pseudo-labels for training more accurate student models. we also explore efforts that are directed at scaling up our method to distill large-scale training data more efficiently. this includes methods such as model selection, where we choose a smaller set of models for distillation that results in only a minimal accuracy drop, and transfer set reduction, where we decrease the volume of data to be distilled to only include in-domain samples. more concretely, we conduct extensive experiments on wmt'21 german ↔ english and chinese ↔ english translation tasks. our final model is as accurate as a large multilingual model with 4.7 billion parameters, despite having only 68 million parameters. ## n-best reranking for distillation our n-best reranker formulates q θ (t|s) as a weighted log-linear model, which is parameterized by a collection of m models, m(s, t) ∈ r m , and their associated weights λ ∈ r m . each model, m i (s, t) ∈ r, assigns a real-valued score that indicates the plausibility of a hypothesis t being the translation of s according to the model, hence we refer to m(s, t) as scoring models. the score of each scoring model m i (s, t) is weighted by λ i and then combined with all other models to produce the final score. while the models are considered static, their associated weights are trainable parameters, learned via discriminative training. we discuss the process in section 3.1. to generate pseudo-labels, our reranker applies the following arg max formula: ## experimental results to showcase the efficacy of our proposal, we conduct large-scale experiments on wmt21 german ↔ english and chinese ↔ english translation tasks. our baseline is the vanilla sequence-level kd @xcite ) that employs the aforementioned lr models as its teachers. we constrained the student model's capacity to approximately 68 million parameters, in line with the transformer base architecture. we use fairseq @xcite for training and inference of our in-house models. more details about these models can be found in appendix a, including other experimental setup including the bitext used mainly for teacher model training and the monolingual data primarily used for student model training. we use the wmt19 set to learn λ weights for our reranker, the wmt20 set as our validation set and the wmt21 set as our blind test set. for these sets, we use the maximum number of references provided. to report accuracy, we use sacre-bleu @xcite with the following signature nrefs:k|case:mixed|eff:no|tok:13a|smooth :exp where k is the number of reference(s). for our main results, we additionally report chrf @xcite with this signature nrefs:k|case:mixed|eff:yes|nc:6|nw:0|spa ce:no and comet22 @xcite using wmt22-comet-da model. for generating n-best list and student model's hypothesis, we set beam size to 8 and 5 respectively. in section 4.1, we first focus on intrinsic evaluation, comparing the accuracy of the n-best reranker with that of the sequence-level kd's teacher models on validation sets. in section 4.2 and section 4.3, we then shift to primary evaluations where we assess the utility of the pseudo-labels generated by our n-best reranker for training student model and retraining teacher models. we mainly focus on the @xmath0. ## related work our proposal intersects with many works in various ways. the idea of utilizing n-best reranking to im-prove accuracy has been extensively investigated as far back as the era of statistical machine translation if not earlier, for example in @xcite @xcite and more recently in @xcite @xcite . in these recent work, n-best reranking incurs significantly higher inference time from running multiple models over the n-best list, thus may not be practical for real-world systems. in contrast, our work makes a practical trade-off by shifting the heavy computational cost of n-best reranking to training data preprocessing without affecting the latency of the deployed model. our work shares the same motivation as @xcite , but we consider a larger and more diverse set of models. the idea of looking at n-best hypotheses for knowledge distillation has been also investigated in the original sequence-level kd paper @xcite , namely sequence-level knowledge interpolation which we consider as one of our baseline where the authors propose to approximate the mode with the hypothesis that scores the highest according some translation metrics. however, since this approach requires the ground truth, the application of this variant is limited to distilling parallel data. in contrast, since our n-best reranker is trained on a tune set, our approach is applicable for distilling unlabelled monolingual data. our n-best reranker incorporates various models as reranking models. some of these models have been applied to knowledge distillation. for example, @xcite deploys nearest neighbor machine translation models. meanwhile, @xcite combines direct models with channel and language models. @xcite trains domain-specific teacher models to distilled in-domain training data for training multi-domain student model. on the other hand, our n-best reranker incorporates significantly larger number of models, including the aforementioned. also, we deploy these models to score hypotheses, rather than to generate them, which is significantly faster. self-training has also been frequently investigated for machine translation in statistical and neural era @xcite . recently, it is often dubbed as iterative knowledge distillation and can be found as a winning formula in many evaluation campaigns @xcite . in this work, we apply self-training using high-quality pseudo-labels from n-best reranker which produces accurate results. ## summary and future work we enhance the sequence-level knowledge distillation @xcite by incorporating nbest reranking. thus, rather than improving the accuracy of the teacher models by following neural scaling laws alone, our proposed method do so by leveraging a multitude of models with different inductive biases, objective functions or architectures to collaboratively rescore n-best hypotheses and identify the best pseudo-labels. furthermore, we observed a relatively strong cascading effect, where teacher models finetuned using pseudolabeled data are more accurate, leading to the generation of more accurate pseudo-labels for the next iteration and resulting in an even more accurate student model. we also explore efficienty efforts to scale up n-best reranking via model selections and transfer set reductions, resulting in a reduction in distillation time. our final student model demonstrates up to 4.0 bleu point improvement over baseline systems and is on par with a strong large translation model on german↔english and chinese↔english translation tasks, despite having only 1/70 th the parameters. for future work, we intend to improve the efficacy of our approach by incorporating more powerful large language models that are finetuned towards translation tasks as well as models that more explicitly capture fine-grained phenomena such as gender or number agreements. to improve the efficiency, we also plan to investigate methods to automatically identify transfer sets at fine-grained sentence level as well as ways to speed up the scoring process further, for instance by utilizing only unnormalized probabilty score like in @xcite .