Ensemble D istillation for S tructured P rediction: C alibrated, A ccurate, F ast— C hoose T hree (2020)

Abstract: Modern neural networks do not always produce well-calibrated predictions, even when trained with a proper scoring function such as cross-entropy. In classification settings, simple methods such as isotonic regression or temperature scaling may be used in conjunction with a held-out dataset to calibrate model outputs. However, extending these methods to structured prediction is not always straightforward or effective; furthermore, a held-out calibration set may not always be available. In this paper, we study ensemble distillation as a general framework for producing well-calibrated structured prediction models while avoiding the prohibitive inference-time cost of ensembles. We validate this framework on two tasks: named-entity recognition and machine translation. We find that, across both tasks, ensemble distillation produces models which retain much of, and occasionally improve upon, the performance and calibration benefits of ensembles, while only requiring a single model during test-time.

## introduction for a calibrated model, an event with a forecast confidence p occurs in held-out data with probability p. calibrated probabilities enable meaningful decision making, either by machines such as downstream probabilistic systems @xcite , or by end-users who must interpret and trust system outputs @xcite . the calibration of modern neural models has recently received increased attention in both the natural language processing and machine learning communities. a major finding is that modern neural networks do not always produce wellcalibrated predictions. as a result, much recent work has focused on improving model calibration, predominantly with post-hoc calibration methods @xcite . however, post-hoc calibration methods have primarily been developed in the context of classification tasks. thus, it is unclear how these methods will affect the performance of sequence-level structured prediction tasks @xcite . additionally, post-hoc calibration methods require a held out calibration dataset, which may not be available in all circumstances. to improve calibration, an alternate approach is model ensembling, which is closely related to approximating the intractable posterior distribution over model parameters @xcite @xcite . although computationally expensive, both at training and inference time, ensembling does not require a separate calibration set. furthermore, ensembles have been found to be competitive or even outperform other calibration methods, particularly in more challenging settings such as dataset shift @xcite . in this paper, we study ensemble distillation as a means of achieving calibrated and accurate structured models while avoiding the high cost of naive ensembles at inference time @xcite . ensemble distillation consists of two stages: in the first stage, we select a base model for the task, such as a recurrent neural network or transformer, and then train an ensemble of k such models, ensuring diversity either via sub-sampling ( §4) or with different random seeds ( §5). in the second stage, the ensemble of k teacher models is distilled into a single student model. prior work has examined the effects of ensemble distillation on measures of uncertainty in vision tasks @xcite . to our knowledge, this is the first systematic study of the effect of ensemble distillation on the calibration of structured prediction models-we consider ner and nmt-which we find poses distinct challenges both in terms of measuring caliensembles vastly improve both performance and calibration over individual models, and ensemble distillation is able to retain much of this improvement in a single model. further, we find that even small ensembles, e.g. of size 3, are enough to see significant improvements over single models. experimental details are described in §5. to this end, our contributions may be summarized as follows: our findings suggest that ensemble distillation has potential to become a standard training recipe in settings where calibration is important. ## calibration given an arbitrary observation x and a model with parameters θ, we are interested in the predictive uncertainty, p θ (y | x), of an event y . our objective is to compute the predictive uncertainty of p θ over of a finite sample of held-out data, {(x (i) , y (i) )} n @xmath0. we then say that p θ is calibrated if the predictive uncertainty agrees with held-out observations; that is, if the model predicts an event with confidence p, then that event prediction is correct p% of the time. calibrated models can be useful for downstream systems which benefit from accurate estimates of uncertainty @xcite . recently, it has been noted that a large portion of modern neural networks are not well calibrated after training @xcite @xcite , although it has been found that pre-training can help with this in natural language processing (desai and durrett, 2020). ## ensemble distillation for ner we evaluate the calibration and performance effects of ensemble distillation on ner models. in these experiments, we examine teacher ensembles that use either strong independence assumptions (subsequently referred to as "iid"), or 1st order markov assumptions. these settings allow us to examine the effects of distilling globally-marginalized versus locally-marginalized structured distributions into student models. we experiment on the 2003 conll dataset @xcite , which contains datasets in english and german, and consider both languages in our experiments. our ner models use representations from pretrained masked language models: bert for english and multilingual-bert (mbert) for german @xcite . given an input sequence @xmath1, . . . , x t } bert outputs representations for each x t . 4 we consider two separate models in our experiments: the 'iid' model makes predictions based solely off of the token-level logits output from a feed-forward layer applied to the bert representations, making each prediction ŷt independently from all others. the 'crf' model instead passes the representations through a bidirectional lstm layer, and the result into a conditional random field (crf) with learned transition parameters @xcite . all models are trained using the adafactor optimizer; we use a learning rate of 1e-4 for training the ensembles, and 5e-5 during distillation. for each dataset, we trained @xmath2. to encourage diversity, each model in a given framework uses a different 1/10 split foot_4 of the training set for its early stopping criterion, in addition to using a different random seed. we then consider ensembles of sizes @xmath3, 6, 9 models, where for @xmath4, 6 the individual models are chosen randomly, but in such a way that the ensemble of 3 is always a subset of the ensemble of 6. during inference time, each ensemble's per-token distribution q k (y t | x) is averaged uniformly to create the ensemble's distribution q(y t | x). in iid ensembles, q k (y t | x) is taken directly from the logits at timestep t. in the crf model, the forward-backward algorithm is used to compute distributions for each timestep which are globally normalized over all possible output sequences y. predictions are then made for each token based on the maximum likelihood prediction from the ensembled distribution, q. for comparison, we also train a collection of 9 models, each with a different random initialization, on the entire training dataset and report their average performance and standard error. note that this setup disadvantages the ensemble, as each individual model in an ensemble has access to strictly less training data than the individual models. ## ensemble distillation for nmt in this section, we evaluate ensemble distillation for nmt models. 9 state-of-the-art nmt models such as the transformer are auto-regressive, meaning that the probability of a given target y t is a function of all previous targets @xmath5. thus, distilling teacher information in this scenario is different from what is done in ner; the structure level knowledge which is being distilled is inherently greedy (the teacher distributions do not take into account future sequences) and the distributions are built off of the gold labeled sequences up to that point (making it difficult to distill the global behavior of the ensemble). all experiments are run on the wmt16 @xmath6, using the vanilla transformer-base architecture from @xcite . 10 we use a vocabulary of 32k symbols based on a joint source and target byte pair encoding @xcite . unlike in our ner setting, all models are trained on the full training set, with variation being instilled only through different random initializations and data order. all models considered use stochastic weight averaging (swa; @xcite . additionally, to evaluate the effect of label smoothing @xcite on calibration we consider 2 variations of nmt experiments: models trained on standard cross-entropy loss (ce-swa), and models trained on cross-entropy loss with a smoothing factor of λ = 0.1 (ls-swa). models are added to ensembles based on order of random seed. 11 during ensemble inference, the next output token is taken from the argmax of the averaged tokendistribution across all models in the ensemble. ## conclusion summary of contributions. we present a systematic study of the effect of ensembles on the calibra-