DiPair: Fast and Accurate Distillation for Trillion-Scale
Text Matching and Pair Modeling

Figure 2 provides an overview of the proposed DiPair model. First, a transformer-based dual-encoder model is applied to the input pair; the output of an encoder is a sequence of token embeddings, which has the same sequence length as the tokenized input text. We then truncate the output sequences by only taking the first $N$ and $M$ token embeddings from the left and right inputs, respectively; the next step is to project those selected token embeddings into lower dimensions and merge them to form the new input sequence. The merged input sequence is then fed into the transformer (or an FFNN) head, and the first token embedding of the output sequence of the head is used as the representation of the initial input pair.

Note that, the dual-encoder will process the full-length input sequences. At the same time, the head only consumes a sequence of length $(N+M)$, which is typically much smaller than the length of the input sequences and ensures efficient execution of the head.

To create the training data for our proposed model, we use an expensive teacher model (e.g., a 12-layer BERT fine-tuned with human-rated data) to annotate a set of unlabeled text pairs (a.k.a. distillation set). The dual-encoder part of our model is initialized from the first few layers of a pre-trained BERT, and a novel two stage training strategy (see Sec. 3.7) is applied to boost the performance further. We defer more details of data specific model distillation to Sec. 4.3.

We now discuss each component of the proposed architecture in detail.

A dual-encoder is the key component of our proposed architecture, and we initialize our dual-encoder from pre-trained BERT (or tinyBERT, ALBERT, etc.). Our basic assumption is that the number of pairs is much larger than the set of unique inputs to the left or right encoders, and the bottleneck of serving our model is to run inference on the pairs with the head. Our proposed architecture, therefore, has an important benefit: increasing the model capacity does not increase the inference time as we can keep the head the same but use more expensive encoders. Figure 3(b) shows that increasing the number of layers of the encoders will often lead to better model performance.

This is the key step to speed up the model serving. Recall that the running time of a transformer-based model quadratically depends on the input sequence length. One of the most effective ways to reduce the running time is to reduce the input sequence length. However, as Table 4 reveals, blindly truncating the input to a BERT model will lead to a quick performance drop. Our key intuition is that, by using a dual-encoder + head architecture, we can focus on reducing the inference time of the head, instead of speeding up the encoders.

Therefore, we still use the full-length input sequences in our encoders, but aggressively reduce the input sequence length to the head. To be more concrete, before merging the outputted sequences from the two encoders, we take the first $N$ and $M$ token embeddings from the left and the right sequences, respectively; This truncation technique has several benefits:

• •

  It significantly speeds up the inference with the head, as the time complexity of transformer layers is quadratic w.r.t. the input sequence length.

• •

  It significantly reduces the amount of data we need to cache. Only the first few token embeddings need be stored as the output of the encoders.

• •

  $N$ and $M$ can be tuned to reflect the desired effectiveness and efficiency trade-off for a particular problem domain.

It is important to note that due to the end-to-end architecture of our model, even though we only use $(N+M)$ token embeddings from the output of the dual-encoder, the model learns to push the information of the input text to the first $(N+M)$ embeddings (thanks to the transformer layers, those selected token embeddings can interact with other token embeddings, and can be viewed as a summary of the full-length input sequences).

For each encoder, we add a projection layer to project each token embedding to a lower dimension. A projection layer is shared within an encoder, but different encoders may use different projection layers. There are two purposes of adding the projection layers:

• •

  Reduce storage. To run the inference with the proposed architecture, we need to cache all the outputs from the encoders.

• •

  Speed up the inference with the head. The time complexity of a transformer linearly depends on the embedding dimension.

In Table 5, we show that by choosing a proper projection layer, we can significantly reduce the embedding dimension with almost no quality drops.

After the projection layer, we merge the $N+M$ projected token embeddings into one sequence and feed it into the head. Like the BERT model, we also add position embeddings and segment embeddings to help the transformer head better aggregate the input sequence. The first token embedding (i.e., CLS embedding) of the transformer head is used as the final representation of the input pair.

Another advantage of using a head is that the head is tokenization-free: the input to the head is purely float tensors, and we do not need to preprocess/ tokenize the input text. This may lead to an additional speedup.

It is worth mentioning that a feedforward neural network (FFNN) can also be used as a head. An FFNN is faster than a transformer-head and often gives reasonable performance (though worse than a transformer head). See the experimental section (Sec. 4.6) for more discussion on these trade-offs.

In the standard dual-encoder model and the recent Poly-Encoders Humeau et al. (2020) work, the dot product between the embeddings is a scalar, which is not suited for tasks beyond regression/binary classification. On the other hand, our proposed architecture outputs a representation of the input pair and is therefore compatible with a wide range of loss functions.

It turns out that directly training the proposed models often leads to sub-optimal results (see Sec. 4.7 for more evidence). This is primarily because adding non-trivial layers on top of a well pre-trained dual-encoder during training may corrupt the knowledge that has been preserved in the dual-encoder. To address this issue, we propose to use a two-stage training strategy: we first freeze the dual-encoder part and only train the newly added parameters until convergence; we then unfreeze the dual-encoder and further train the entire model. A similar training strategy can be found in e.g., Wang et al. (2019).

Unlike the models proposed in the recent works MacAvaney et al. (2020); Humeau et al. (2020) where only pairs can be supported, our proposed architecture trivially extends to the scenario where we have $n$-ary tuple of textual objects $(a_{1},a_{2},\ldots,a_{n})$ as the model input, as we can simply replace the dual-encoder model with an $n$-encoder model. This feature is useful in many applications, such as QA tasks with context, or query to document scoring tasks with personalized information.

We evaluate our proposed methods on two datasets (Table 1 provides an overview):

• •

  Q2P-Mat is a binary classification task derived from the MSMARCO Passage Ranking data²²2https://microsoft.github.io/MSMARCO-Passage-Ranking/. Given a (query, passage) pair, the goal is to predict whether the passage contains the answer for the query. We measure the model performance using AUC-ROC. Appendix B lists more details.

• •

  P2T-Rel  is a regression task on a real-world ecommerce dataset. Given a (product, term) pair, the goal is to predict the relevance of the term to the product. We measure the model performance using Pearson correlation with the human judgments. Title and description are used as the product features. Appendix A provides several examples of (product, term) pairs.

There exist many knowledge distillation (see Sec. 2 for more details) works, but none of them has been optimized for pairwise input. We choose to compare our DiPair approach with the fastest BERT-based student model Turc et al. (2019) we are aware of, and our model is at least 8x faster (see Table 3(b)). We also compare our proposed approach with several other strong baselines:

• •

  BERT-Tiny: the fastest version of BERT released in Turc et al. (2019). This model has $2$ layers with $128$D word embeddings and $2$-head transformer. It is claimed to be $52$x faster than BERT-base (on TPU), and to the best of our knowledge, this is faster than any other BERT-based student models in the literature.

• •

  DE-Cos: BERT-based Dual-Encoder model. Cosine between left/right CLS embeddings is used as the similarity score.

• •

  DE-FFNN: BERT-based Dual-Encoder model. FFNN (Feedforward Neural Networks) is used to aggregate the left/right CLS embeddings into a similarity score. Unless otherwise stated, we fix the FFNN to be 2-Layer with dimensions x128x128. The input to the FFNN has dimension $768+768=1536$.

• •

  DiPairTSF: our proposed model, BERT-based Dual-Encoder, with a transformer-based head. $N$ and $M$ refer to the output sequence lengths (see Figure 2). In all experiments, we fix our head to be 2-Layer, 1-Head, 1024D intermediate size. The value of hidden_size (i.e., the dimension of the input token embeddings) is decided by the output of the projection layer.

• •

  DiPairFFNN: this is similar to DiPairTSF; the only difference is that the transformer-based head is replaced with an FFNN. The input to the FFNN has dimension $(N+M)$ * hidden_size (N, M defined in Figure 2 ). We use 2-Layer FFNN with dimensions x128x128 unless otherwise stated.

In all the aforementioned models (except BERT-Tiny), the dual-encoder is initialized from the first $K$ layers of the pre-trained BERT model as well as the token embedding matrix. Unless otherwise stated, we fix K=1 for P2T-Rel and K=4 for Q2P-Mat. The Left encoder and the right encoder will share parameters. For models with a projection layer, we use $D$ to represent the dimension of the projected result.

Our code is implemented with TensorFlow ⁴⁴4https://www.tensorflow.org/ and we use TPUv3 in all of our experiments. We use AdamW optimizer following the public BERT code. The warmup step is fixed to be 50k. Other parameters of the optimizer are identical to the default values set in the public BERT code ( weight_decay_rate=0.01, $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=1e^{-6}$).

We tune some other key hyper-parameters using the validation sets. We try multiple (learning rate, batch size) combinations and choose the best ones. In the two-stage training, the models are less sensitive to learning rates in the first stage, and we set the learning rate as 5e-5; we then train the models until they converge. In the second stage of training, the learning rate is set to be 5e-5 in DiPairTSF, DE-Cos, DiPairFFNN; we use batch size 512 and 4x4 TPU topology. For BERT-Tiny, we use batch size 128, learning rate 2e-6, and 2x2 TPU topology. All other hyperparameters related to model architecture are specified in Sec. 4.2.

Table 2 and Table 3 present the experimental results on P2T-Rel and Q2P-Mat datasets, respectively. Among all the student models with dual-encoder architecture, DiPairTSF consistently achieves the best performance. For the Q2P-Mat dataset, DiPairTSF  achieves similar AUC_ROC to BERT-Tiny; however, it achieves a 8x speedup.

Among all the student models, DE-Cos is the fastest one as it only requires dot product during inference. However, it has the worst performance, indicating that using Cosine function alone does not allow enough interaction between the input sequences embeddings.

To verify the importance of using a transformer-based head, we vary #params in the heads of DiPairTSF, DE-FFNN and DiPairFFNN. Table 4 presents the experimental results.

Comparing rows 1 and 2 in Table 4, the model quality of DiPairTSF  can be improved by increasing the head input sequences lengths ($N$ and $M$), although at the cost of longer inference time. On the other hand, rows 3-5 show that increasing #Params in FFNN head (e.g., using larger dimensions, more layers) does not lead to significant quality improvement for DiPairFFNN; even when the #Params of the FFNN head is 4x more than the transformer head, the model quality of DiPairTSF  is still considerably superior to that of DiPairFFNN(cf. rows 2 and 5). A similar conclusion can be made for DE-FFNN (rows 6-8).

Another interesting observation is that even with more input information and more parameters, DiPairFFNN  does not generate higher AUC_ROC than DE-FFNN. This might suggest that FFNN is not powerful enough to aggregate the input information effectively.

Overall, Table 4 illustrates the importance of using a transformer head if we want to achieve high model quality: Unlike FFNN-based heads, where we could not further improve the model via increasing #Params, a transformer-based head has more headroom to reduce the distillation gap further, and the desired quality-speed trade-off can be easily achieved by adjusting the values of $N$ and $M$.

Figure 3(a) shows that two-stage training, which is discussed in Section 3.7 has positive effects on all the methods we test. When the head is transformer-based, the two-stage training plays an important role: the AUC_ROC improves from $0.891$ to $0.930$.

On the other hand, the gain introduced by using two-stage training is less significant in other approaches such as DE-FFNN  and DiPairFFNN. This might be because FFNN is generally easier to train than transformer-based models, and thus initialization choices play a lesser role.