Lost in Distillation: A Case Study in Toxicity Modeling (2022)

Abstract: In an era of increasingly large pre-trained language models, knowledge distillation is a powerful tool for transferring information from a large model to a smaller one. In particular, distillation is of tremendous benefit when it comes to real-world constraints such as serving latency or serving at scale. However, a loss of robustness in language understanding may be hidden in the process and not immediately revealed when looking at high-level evaluation metrics. In this work, we investigate the hidden costs: what is “lost in distillation”, especially in regards to identity-based model bias using the case study of toxicity modeling. With reproducible models using open source training sets, we investigate models distilled from a BERT teacher baseline. Using both open source and proprietary big data models, we investigate these hidden performance costs.

## introduction the revolution in natural language processing brought on by transformers, which have now been employed in virtually all major text processing applications, also brought substantially higher computational costs. the typical bert model @xcite has over 100m parameters and 12 layers. the prospect of using these models in production settings without special purpose hardware quickly led practitioners to seek techniques to reduce the computational costs. an approach widely advocated is to employ the technique of knowledge distillation to improve the performance of a simpler student model by training on additional unsupervised data that has been labeled by the larger teacher model @xcite . the ability to draw upon the wellspring of nearly unlimited unsupervised data and to leverage the higher performance of a much larger model, while maintaining the lower serving costs of a smaller model, has led to rapid adoption of this practice. however, closer analysis of the performance of distilled models reveals that while they may be able to erect a facade of high accuracy, they fail to capture important aspects of the knowledge represented in the teacher models. we present a particular method of using distillation that we used to improve the performance of our models through pseudo-labeling of unsupervised data, while retaining the model architecture and number of parameters. while, for some metrics we saw nearly asymptotic performance to the teacher model, using other metrics we discovered important differences. while we do not know if this problem will manifest across all differences in architecture and parameterization -we want to caution researchers who are exploring distillation as a potential quick fix. models with unintended biases has received considerable attention with multiple survey papers both generally @xcite and for natural language in particular @xcite . two popular implementations of the distillation paradigm of creating a vast training set using large models to label unsupervised data are presented in @xcite and @xcite . the primary goal of this work is producing a model with similar performance characteristics on the target task, but with lower a resource footprint. @xcite suggests pre-training and fine-tuning compact models as an alternative to traditional distillation. however, the effects on model bias were not reported in these studies. several other works explore this idea in modes similar to the work we present here, although often with a different array of model architectures. @xcite and mangalwedhekar (2021) both include cnns as one of the target models. @xcite ; @xcite all present additional studies regarding distillation and the performance of the models in terms of fidelity to the teacher model. specifically regarding bias in the distillation or model compression setting, @xcite report reduction in bias in contrast to our findings, although in a generation application. however, @xcite makes clear that biases from the training data can also be preserved or exacerbated in a similar distillation setting. @xcite raises several risks of large language models overall, including identity-based bias. we show that these risks can be magnified with the use of distillation, and that high-level accuracy metrics can hide nuances in performance, especially when large models are built to address a wide range of use cases. ## toxicity modeling we have chosen to use the problem of "toxic" comment classification to illustrate the difficulty that we observed in distillation. this is due to the ready availability of training resources for this task, the practical real-world need to address this problem, and the clear risks @xcite of identity term bias and other modeling pitfalls. several diagnostic frameworks that were proposed to highlight the limitations of classification systems in general can also be used to highlight the problems with distillation in particular. our primary framework is the method of measuring classifier unintended bias associated with neutral or ambiguous identity terms. this framework was introduced in @xcite and expanded in @xcite along with the civil comments dataset that is our primary source of supervised training data. in addition we use the diagnostic hatecheck test set @xcite . recently works that study implicitly abusive language @xcite , where careful attention to the context and implication of the comments is required. we include these evaluation challenges for our models. ## models we found the bias effects of distillation to be remarkably persistent from a small to a very large scale. we created smaller, reproducible models entirely from publicly available resources, and duplicated the same findings on a very large model to show the generality of these findings. ## evaluating performance and bias experiments are run on a variety of evaluation sets to assess the classification performance of the teacher, baseline and distilled models. in assessing both the civil comments based models and the big data models, we compare the distilled student and baseline models performance against the teacher models. results are shown in ## conclusion the experimental section illustrates for both our more robust big data toxicity models and smaller reproducible versions that we are able to distill large transformer teacher models into smaller student models with very similar metrics on the evaluation datasets drawn from the same distribution. however, upon closer inspection, the distilled models consistently under-perform and even demonstrate serious weakness when examined on larger and more difficult suites of test sets. in particular, identity-based bias for the toxicity models is noticeably worse in the distilled model versions, even with the addition of significant quantities biasmitigating data. in our own work, we became interested in distillation because it allowed us to maintain our existing architecture and serving costs, but allowed us to improve our models to what seemed like performance parity with the promising new bert models. we quickly noticed that distilled models performed worse, consistently, in our bias metrics. while the technique of data augmentation has helped us mitigate these biases, that technique has proven to be less effective in distillation settings. in trying to tackle biases, whether caused by sampling methods, the annotators, or the models themselves, there are always other potential biases that we are not yet measuring. for these reasons we have concluded that there may be subtle and intangible benefits to using large models. importantly for us, data augmentation techniques for bias mitigation perform better with transformer models, at least to the limits of our ability to measure. while distillation seemingly lifts student model performance to new heights of accuracy, it may be a pale imitation of the often profound context sensitive classifications that are produced by the teacher models. we hope that this caution and advice with help other practitioners who face similar choices.