Distillation of encoder-decoder transformers for sequence labelling (2023)

Abstract: Driven by encouraging results on a wide range of tasks, the field of NLP is experiencing an accelerated race to develop bigger language models. This race for bigger models has also underscored the need to continue the pursuit of practical distillation approaches that can leverage the knowledge acquired by these big models in a compute-efficient manner. Having this goal in mind, we build on recent work to propose a hallucination-free framework for sequence tagging that is especially suited for distillation. We show empirical results of new state-of-the-art performance across multiple sequence labelling datasets and validate the usefulness of this framework for distilling a large model in a few-shot learning scenario.

## introduction sequence labelling (sl) can be defined as the task of assigning a label to a span in the input text. some examples of sl tasks are: i) named entity recognition (ner), where these labelled spans refer to people, places, or organizations, and ii) slot-filling, where these spans or slots of interest refer to attributes relevant to complete a user command, such as song name and playlist in a dialogue system. in general, these spans vary semantically depending on the domain of the task. despite the strong trend in nlp to explore the use of large language models (llms) there is still limited work evaluating prompting and decoding mechanisms for sl tasks. in this paper we propose and evaluate a new inference approach for sl that addresses two practical constraints: data scarcity leads us to consider highperforming encoder-decoder based llms. we address deployment concerns by considering distillation of such models into much smaller sl architectures, for instance bi-directional long short memory (bilstm) @xcite units, through the use of both labelled and unlabelled data. a standard distillation approach, knowledge distillation (kd) @xcite , requires access to the probability that the teacher network assigns to each of the possible output tags. this probability distribution is typically unavailable at inference time for llms; thus, distillation of encoderdecoder models needs to resort to pseudo-labels: foot_0 the student is trained on the one-hot labels that the teacher assigns to examples in an unlabelled dataset. this prevents the student model from learning those relationships among the probabilities of the incorrect classes that the teacher has learned. similar arguments apply to decoder-only models. in this paper, we propose sent ′ , a simple modification of the simplified inside sentinel+tag (sent) format by @xcite . we combine our target sequence format with a scoring mechanism for decoding, which we collectively call sentscore. this combination results in an effective framework that allows us to employ a language model to perform sequence labelling and knowledge distillation. we show that sentscore is an hallucination-free decoding scheme, and that even with smaller models it outperforms the original sent format across a variety of standard sl datasets. our proposed sentscore method defines a sequence of scores over the output tags that can be aligned with those generated by the sequence tagging student network, making kd possible. we find an advantage in terms of performance in using kd as opposed to just pseudo-labels as a distillation objective, especially for smaller distillation datasets. in sum, our contributions are: using llms to perform sequence tagging is discussed by athiwaratkun et al. (2020); yan et al. (2021); paolini et al. (2021); qin and joty (2021); xue et al. (2022) and @xcite . while these previous works have minor differences in the prompting format of the models, all but the last one include input tokens as part of the target sequence. different from our work, all previous models are prone to hallucinate. distillation refers to training a small student model from scratch using supervision from a large pretrained model @xcite . distillation of transformer-based models for different nlp tasks is typically discussed in the context of encoder-only models (e.g. @xcite @xcite , with a few exceptions looking at distillation of decoder-only models (e.g. @xcite . in this paper we will discuss two approaches to distillation: pseudo-labels and knowledge distillation (kd). in the first approach the student model is trained on the hard labels generated by the teacher on some (unlabelled) dataset. in the second approach additional soft information provided by the teacher is used: typically the probability distribution the teacher assigns to the labels. in the context of sequence labelling, using pseudo-labels allows us to perform distillation on any teacher-student architecture pair. kd, on the other hand, requires access to the teacher's probability distribution over the output tags. these are not usually available in language models for which the output distribution is over the whole vocabulary of tokens. we are not aware of other works which modify the decoder inference algorithm to generate such probabilities. however, there is recent work distilling internal representations of the teacher model, with the most closely related work to us being @xcite . in that work the authors distill a multilingual encoderonly model into a bilstm architecture using a two-stage training process. this two-stage process, however, assumes a large unlabelled set for distilling internal model representations, embedding space, and teacher logits, and another significant amount of labelled data for directly training the student model using cross-entropy loss. ## datasets we select seven english datasets that have been used in recent work on slot labelling: atis @xcite , snips @xcite , mit corpora (movie, movietrivia, and restaurant) 2 , and the english parts of mtop @xcite and of mtod @xcite . some statistics about the datasets are shown in we use the original training, development, and test sets of the snips, mtop, and mtod datasets. for the atis dataset we use the splits established in the literature by @xcite , in which a part of the original training set is used as the dev set. similarly, we follow @xcite 3 to obtain a dev set out of the original training set for each of the mit datasets. we notice that all datasets, with the exception of movietrivia, contain some duplicates. among these, all apart from restaurant contain examples in the test set that are also duplicated in the train and dev sets. this happens for fewer than 30 instances, with the exception of mtod, where more than 20% of the test set examples are also found in the train and dev sets. how these duplicates are handled varies across the literature; we do not remove duplicates from the datasets used in our main results. however, for mtod, we also obtained results on a version of the dataset that was deduplicated as follows: if an example is duplicated, we retain it in the highest priority (defined below) split and removed from the others. to ensure the test set is as close as possible to the original test set, we order the splits in ascending order of priority as follows: test, dev, and train. we found that the f1 scores on the deduped mtod dataset are within 0.5 points those on the original mtod dataset across all experiments; as such, we do not report the deduped results in the following sections. in addition to covering different domains, there are noticeable differences across the datasets in terms of the number of tags and the number of labelled examples for evaluation and testing, as can be seen in ## score-based sequence labelling using llms for sequence tagging requires reframing the problem as a sequence-to-sequence task. in raman et al. ( 2022 ), the strategy that proved the most effective, at least when applied to the mt5 encoder-decoder architecture, was the simplified inside sentinel+tag (sent in this paper). in this format (see given the set t of string labels to be used to annotate a span of text, the scheme used to associate tags across tokens is a modification of the standard bio scheme: we use @xmath0, a single tag i for each token that continues a labelled span, and o to tag tokens that do not belong to labelled spans. we refer to this scheme as simplified inside bio (sbio), and we indicate with t ≡ t ∪{i, o} the tag set associated to it. @xcite argue that the success of sent can be attributed to two factors: 1) on the one hand, the use of sentinel tokens mimics the denoising objective that is used to pretrain mt5; 2) on the other hand, when compared to other decoding strategies, sent does not require the decoder to copy parts of the input sentences and also produces shorter outputs. both these facts supposedly make the task easier to learn and reduce the number of errors from the decoder (hallucinations, as they are often referred to in the literature). we remark however that any output format among those described in the literature can be made completely free of hallucinations by constraining decoding (either greedy or beam search based) through a finite state machine enforcing the desired output format (see for instance @xcite . in what follows we describe our proposed decoding approach that builds on this previous work. ## experimental settings we evaluate the models by computing the f1 score on the test set of each dataset. f1 is calculated following the conll convention as defined by @xcite , where an entity is considered correct iff the entity is predicted exactly as it appears in the gold data. we show micro-averaged f1 scores. the first set of experiments we performed are intended to investigate whether our proposed senttscore approach is competitive with respect to recent results on the same datasets ( this discrepancy is resolved by the invariance of the softmax under constant shifts of its arguments. sentt ′ : the pretrained model is the same as that used for senttscore. the goal of this baseline is to assess improvements attributed to our proposed decoding mechanism. this model is also the closest model to prior sota. the main difference between our results and those in @xcite is the pretrained model. they used a multilingual t5 model @xcite with 580m parameters, whereas we use a smaller monolingual version @xcite . all the above models were trained with the adamw optimizer @xcite . the best checkpoint for each training job was selected based on highest micro-f1 score on the validation set. all pretrained transformer models are downloaded from huggingface. ## results the comparisons between baselines, sent ′ , and sentscore are shown in ## limitations and future work a reasonable critique to our focus on real-world constraints is the simple fact the datasets we are using are not real-world ones. from noise to tokenization choices, many issues arise when considering datasets outside of the academic domain. however, we believe our methods are simple enough to be applicable to real-world scenarios and our results to be independent of these various subtleties. some issues that could be addressed in future work have to do with the exploration of even larger models and different architectures such as decoderonly ones @xcite @xcite @xcite . we note, however, that in all our experiments we finetune all the weights of the pretrained models we use. when using extremely large models this becomes impractical. recent work @xcite suggests that kd with compact encoder-only student models, such as bert, is a promising avenue for further research. exploring the pure few-shot scenario, or only finetuning a subnetwork, for instance by using adapters à la @xcite , would be also interesting. ## conclusion real-time systems need to find a trade-off between performances and computing resources, the latter constraint coming either from budget or some other service requirement. such trade-offs become particularly evident with large pretrained transformer models, which achieve sota results on many nlp tasks at the cost of being extremely hard and expensive to deploy in a real-world setting. the standard solution for this is distillation. in this paper we have revisited these issues for the sl task, which is often the first crucial step in many real-world nlp pipelines. we propose a new inference algorithm, sentscore, that allows us to leverage the performance of arbitrarily large encoder-decoder transformer architectures by distilling them into simpler sequence taggers using kd as opposed to just pseudo-labelling.