D i P air: Fast and Accurate Distillation for Trillion-Scale Text Matching and Pair Modeling (2020)

Abstract: Pre-trained models like BERT ((Devlin et al., 2018) have dominated NLP / IR applications such as single sentence classification, text pair classification, and question answering. However, deploying these models in real systems is highly non-trivial due to their exorbitant computational costs. A common remedy to this is knowledge distillation (Hinton et al., 2015), leading to faster inference. However – as we show here – existing works are not optimized for dealing with pairs (or tuples) of texts. Consequently, they are either not scalable or demonstrate subpar performance. In this work, we propose DiPair — a novel framework for distilling fast and accurate models on text pair tasks. Coupled with an end-to-end training strategy, DiPair is both highly scalable and offers improved quality-speed tradeoffs. Empirical studies conducted on both academic and real-world e-commerce benchmarks demonstrate the efficacy of the proposed approach with speedups of over 350x and minimal quality drop relative to the cross-attention teacher BERT model.

## introduction modeling the relationship between textual objects is critical to numerous nlp and information retrieval (ir) applications @xcite . this subsumes a number of different problems such as textual entailment, semantic text matching, paraphrase identification, plagiarism detection, and relevance modeling. for example, modeling the relationship between queries and documents / ad keywords is central to search engines / digital advertisement systems @xcite . recently neural network-based models have demonstrated large gains in this space @xcite . in particular, the transformer / bert family of models @xcite correspondence to chenjiecao@google.com 2018; @xcite @xcite have set a new bar for these semantic text matching problems. however, the computational costs of these models have proven to be prohibitively expensive, thus limiting their use in real-world applications @xcite . for example, on the e-commerce relevance-scoring task (p2t-rel dataset) discussed in sec. 4.1, scoring the (trillion+) text pairs would take years. one popular remedy is to distill these expensive teacher models @xcite into lightweight student models. training these students using examples labeled by the teacher has been shown to maintain quality while enabling faster inference. the key to the effectiveness of distillation techniques is a good trade-off between student quality and inference speed. however, as we show here, existing knowledge distillation techniques @xcite @xcite ) fall short on the quality-speed trade-off when dealing with pairs of texts. on one hand, approaches that model the texts jointly (i.e., using cross-attention) even one as highly optimized as bert-tiny @xcite are still orders of magnitude too slow. on the other hand, techniques that model the texts independently such as the dual-encoder models foot_0 @xcite @xcite @xcite are able to run efficient inference on large-scale text pairs. by exploiting the independence of the texts, these techniques can significantly speed up inference by caching/indexing embeddings of individual texts. however, this speedup comes at a significant costwith sharply reduced scoring quality. the key drawback here is that these independent models lack the ability to mimic the cross-attention enabled teachers and model the joint nuances and facets of the texts. as a motivating example, consider the ecommerce term relevance-scoring task. for the product "black sport nike shoes for boys size wide", terms such as "black", "wide footwear" and "nike shoes" are all relevant. however, enforcing similarity between the independently modeled term and product will lead to the embeddings of "black" and "nike shoes" being incorrectly considered similar. motivated by this, we propose dipair for fast and accurate distillation of large-scale text matching and pair modeling. dipair aims to combine the best of both worlds: like dual-encoder models, it leverages common pre-computation, while at the same time modeling the text jointly -with cross-attention -using multiple contextual embeddings for each text. in particular, we extract a small fraction of the output token embeddings from each text, and then jointly model this smaller "sequence" using a transformer head (we use the term head to refer to the component that consumes the outputs of a dual-encoder model, see figure 2 ). we demonstrate that a two-stage, end-to-end training allows the proposed dipair model to learn richer multifaceted semantic representations of the text pairs. the resulting dipair model is 350x+ faster with minimal quality drop relative to the teacher on academic and real-world e-commerce datasets. in summary, our main contributions include: ## experiments in this section, we conduct experimental studies. we aim to answer the following questions through our experiments: ## open questions dipair has been discussed in the context of knowledge distillation in this work, but it can be trivially extended to more scenarios, as we can train it directly. the proposed framework raises several research questions. learning dynamics of our model recall that, in our framework, each encoder outputs its first few token embeddings as the input to the head, and we end to end to train the model to force the encoder to push the information of the input text into those outputted embeddings. however, it is unclear to us what those outputted embeddings actually learn. it would be interesting to understand the learning dynamics of our model. ## conclusion and future work in this work, we reveal the importance of customizing models for problems with pairwise/n-ary input and propose a new framework, dipair, as an effective solution. this framework is flexible, and we can easily achieve more than 350x speedup over a bert-based teacher model with no significant quality drop.