Ensemble Transformer for Efficient and Accurate Ranking Tasks: an Application to Question Answering Systems (2022)

Abstract: Large transformer models can highly improve Answer Sentence Selection (AS2) tasks, but their high computational costs prevent their use in many real-world applications. In this paper, we explore the following research question: How can we make the AS2 models more accurate without significantly increasing their model complexity? To address the question, we propose a Multiple Heads Student architecture (named CERBERUS), an efficient neural network designed to distill an ensemble of large transformers into a single smaller model. CERBERUS consists of two components: a stack of transformer layers that is used to encode inputs, and a set of ranking heads; unlike traditional distillation technique, each of them is trained by distilling a different large transformer architecture in a way that preserves the diversity of the ensemble members. The resulting model captures the knowledge of heterogeneous transformer models by using just a few extra parameters. We show the effectiveness of CERBERUS on three English datasets for AS2; our proposed approach outperforms all single-model distillations we consider, rivaling the state-of-the-art large AS2 models that have 2.7× more parameters and run 2.5× slower. Code for our model is available at https://github.com/amazon-research/wqa-cerberus.

## introduction answer sentence selection (as2) is a core task for designing efficient retrieval-based web qa systems: given a question and a set of answer sentence candidates (e.g., retrieved by a search engine), as2 models select the sentence that correctly answers the question with the highest probability. as2 research originated from the trec competitions @xcite , which targeted large amounts of unstructured text. as2 models are very efficient, and can enable web-powered question answering systems of real-world virtual assistants such as alexa, google home, siri, and others. as most research areas in text processing and retrieval, as2 has been dominated by the use of ever larger transformer model architectures @xcite . these models are typically pre-trained using language modeling tasks on large amounts of text @xcite @xcite , and then fine-tuned on specific downstream tasks @xcite @xcite . @xcite achieved an impressive accuracy by fine-tuning pre-trained transformers to the as2 task on the target datasets. they established the new state of the art performance for as2 using a roberta large model. unfortunately, larger transformer models come at a cost: they require large computing resources, consume a lot of energy (critically impacting the environment @xcite ), and may have unacceptable latency and/or memory usage. these downsides are critical for as2 applications, where, for any given query, a model is required to score hundreds or thousands of candidates to select the top-k answers. therefore, in this work, we investigate how as2 models can be made more accurate without significantly increasing their complexity. previous work has addressed the general problem of high computational cost of transformer models by developing techniques for reducing their overall size while maintaining most of their performance @xcite @xcite . in particular, knowledge distillation (kd) techniques have been shown to be particularly effective @xcite @xcite @xcite . kd techniques use a larger model, known as a teacher, to obtain a smaller and thus more efficient model, known as a student @xcite . the student is trained to mimic the output of the teacher. however, we empirically show that, at least for as2, base models trained through distillation are still significantly behind the state of the art, i.e., models based on large transformers. in this paper, we introduce a new transformer model for as2 that matches the state of the art while being dramatically more efficient. our main idea is based on the following considerations: first, in recent years, several transformer model families have been introduced, each pretrained using different datasets and modeling techniques @xcite . second, ensembling several diverse models has shown to be an effective way to improve performance in many question answering and ranking tasks @xcite @xcite . our contribution lies in a new approach to approximate a computationally expensive ranking ensemble into a single efficient architecture for as2 tasks. more specifically, our investigation proceeds as follows. first, we optimize ranking architectures for as2 by training k student models to replicate k unique teacher architectures. when ensembled, we show that they achieve better performance than any standalone models at the cost of increased computational burden. then, to preserve the accuracy of this ensemble while achieving lower complexity, we propose a new multiple heads student architecture, which we refer to as cerberus. as shown in fig. 1 , cerberus is composed of a shared encoder body and multiple ranking heads. the encoder body is designed to derive a shared representation of input sequences, which gets fed to ranking heads. we show that if each ranking head is trained to mimic a unique teacher distribution, it is possible to achieve the desirable diversity through ensemble model while being significantly more efficient. we train a cerberus model using three different teachers: roberta @xcite , elec-tra @xcite , and albert @xcite . we conduct experiments on three as2 datasets: asnq @xcite , wikiqa @xcite , and an internal corpus (ias2). our results show that cerberus consistently improves over all models trained with single teachers, rivaling performance of much larger models including multiple variants of ensemble models; further, cer-berus matches current state-of-the-art as2 models (tanda by @xcite al. (2020) ), while saving 64% and 60% in model size and latency, respectively. in summary, our contribution is four-fold: (i) we propose cerberus, an efficient architecture specifically designed to distill an ensemble of heterogeneous transformer models into a single transformer model for as2 tasks while preserving ensemble diversity. ## methodology we build up to introducing cerberus by first formalizing the as2 task (section 3.1), and then summarizing typical transformer distillation and ensembling techniques (section 3.2). finally, details of the cerberus approach are explained in section 3.3. ## results here we present our main experimental findings. in section 5.1, we compare cerberus to stateof-the-art models and other distillation techniques using three datasets (ias2, asnq, wikiqa). in sections 5.2 -5.4, we motivate our design and hyperparameter choices for cerberus by empirically validating them. finally, in section 5.5, we discuss inference latency of cerberus comparing to other transformer models. ## conclusions and future work in this work, we introduce a technique for obtaining a single efficient as2 model from an ensemble of heterogeneous transformer models. this efficient approach, which we call cerberus, consists of a sequence of transformer blocks, followed by multiple ranking heads; each head is trained with a unique teacher, ensuring proper distillation of the ensemble. results show that the proposed model outperforms traditional, single teacher techniques, rivaling state-of-the-art as2 models while saving 64% and 60% in model size and latency, respectively. cerberus enables large-like as2 accuracy while maintaining base-like efficiency. further analysis demonstrates that reported improvements in as2 performance are due to to three key factors: (i) multiple ranking heads, (ii) multiple teachers, and (iii) heterogeneity in teacher models. future work would focus on two key aspects: how cerberus performs on non-ranking tasks, and whether it could achieve similar improvements in ranking tasks outside qa. for the former, we remark that, while the core idea of cerberus can be extended to tasks such as those in the glue benchmark @xcite , further investigation is necessary in establishing the best set of trade-offs for different objectives and metrics. a similar concern exists in the case of extending cerberus to ranking tasks, such as ad-hoc retrieval.