DeText: A Deep Text Ranking Framework with BERT

There are two categories of deep NLP based ranking models: representation based and interaction based models. Representation based models learn independent embeddings for the query and the document. DSSM (Huang et al., 2013) averages the word embeddings as the query/document embeddings. Following this work, CLSM/LSTM-RNN (Shen et al., 2014; Palangi et al., 2016) encodes word order information using CNN(LeCun and Bengio, 1995)/LSTM(Hochreiter and Schmidhuber, 1997), respectively. All these three works assume that there is only one field on the document side, and the document score is the cosine similarity score of the query/document embedding. NRM-F (Zamani et al., 2018) adds more fields in the document side and achieves better performance. One major weakness of representation based networks is the failure to capture local lexical matching, since the text embedding, e.g., a 128 dimensional vector, cannot summarize all the information in the original text.

To overcome the issue, interaction based models compare each part of the query with each part of the document. In DRMM (Guo et al., 2016), a cosine similarity is computed for each word embedding in the query and each word embedding in the document. The final document score is computed based on the pairwise word similarity score histogram. K-NRM (Xiong et al., 2017) and Conv-KNRM (Dai et al., 2018) extended DRMM by kernel pooling and pairwise ngram similarity, respectively. Recently, BERT (Devlin et al., 2019) has shown superior performance (Dai and Callan, 2019; Nogueira and Cho, 2019; Qiao et al., 2019) in ranking. It is considered an interaction based model, since the query string and document string are concatenated as one sentence, where transformer layer (Vaswani et al., 2017) compares every word pair in that sentence.

In experiments of previous works, interaction based methods usually produce better relevance results than representation based methods, at the cost of longer computation time introduced by the pairwise word comparison.

Commercial search engines have a strict requirement on the serving latency. Despite better relevance performance, the interaction based ranking approaches are not scalable due to the heavy interaction computation. Therefore, to our best knowledge, the representation based approaches are generally used for production.

With representation based approaches, existing work uses embedding pre-computing, either for documents (Ramanath et al., 2018; Yin et al., 2016) or for member profiles (personalization) (Grbovic and Cheng, 2018). It requires a huge amount of hard disk space to store the embedding, as well as a sophisticated system design to refresh the embeddings when there are any document/profile changes.

As illustrated in Figure 2, given multiple source (queries, user profiles) / target (documents) texts and traditional features, the DeText framework contains 6 components: Token Embedding Layer, Text Embedding Layer, Interaction Layer, Traditional Feature Processing, Multilayer-Perceptron Layer, and Learning-to-rank Layer. Specifically, DeText-BERT model uses BERT as the text embedding layer. In the rest of this section, we will illustrate the details of each component.

Input Text Data. The input text data is generalized as source and target texts. The source texts could be queries or user profiles. The target text could be documents. Both source and target could have multiple fields, which is different from most previous work (Huang et al., 2013; Shen et al., 2014; Palangi et al., 2016), where only two fields (query and document) are available. There are several advantages of using multiple fields: 1). enable personalization with text fields from user profiles, and 2). achieve better and more robust results.

Token Embedding Layer. The sequence of text tokens is transformed into an embedding matrix $E$. For text with $m$ tokens, the matrix has a size of $d\times m$, where $d$ is the number of token embedding dimensions. Depending on the text encoding methods, different token granularities are used: in CNN/LSTM, the tokens are words; in BERT, the tokens are subwords (Sennrich et al., 2016).

Text Embedding Layer. Under the representation based model structure, embedding is extracted independently for each text field. The embedding can be generated through various neural network components for deep natural language processing, such as BERT, CNN, LSTM, etc. The outcome of this layer is a $d$-dimensional embedding vector. More details are discussed in Section 4.3 and 5.2.

Interaction Layer. The interaction between source and target only happens after the text embedding is generated, which is the key difference of representation based methods from interaction based methods. Table 2 summarizes the different interaction methods, where $u_{s}$/$u_{t}$ is the source/target field embedding, respectively. Note that for every source and target pair, cosine similarity generates one feature, while the Hadamard product/concatenation generates many features (a vector).

Traditional Feature Processing. The existing hand-crafted features, such as personalization features, social networks features, user behavior features, are usually informative for ranking. To integrate them with deep NLP features, we use standard normalization and elementwise rescaling (Aksoy and Haralick, 2001) to better process the features:

where mean $\mu$ and standard deviation $\sigma$ are pre-computed from training data, and $w$ and $b$ are learned in the DeText-BERT model.

MLP Layer. Deep features, as the output of the interaction layer, are concatenated with the traditional features as the final features, followed by a Multilayer-Perceptron (MLP) (Pal and Mitra, 1992) layer to compute the final document score. The hidden layer in MLP is able to extract the non-linear correlations of deep features and traditional features.

LTR Layer. The last layer is the learning-to-rank layer that takes multiple target scores as input. DeText provides the flexibility of pointwise, pairwise or listwise LTR (Burges, 2010), as well as Lambda rank (Burges et al., 2007). Binary classification loss (pointwise learning-to-rank) can be used for systems where click probability is important to model, while pairwise/listwise LTR can be used for systems where only relative position matters.

The DeText framework enables model flexibility to adapt to demands of different productions, in terms of input data layer (multiple source/target fields), text embedding layer (CNN vs BERT), interaction layer (cosine/hadamard/concat), LTR (pointwise/pairwise/listwise), etc.

By enhancing the model flexibility, we can optimize the model effectiveness while maintaining efficiency. Firstly, representation based methods are used to bound the time complexity. Secondly, the flexibility of input data/interaction layer, together with traditional feature handling, enable us to experiment and develop scalable neural network models with strong relevance performance.

To use BERT in ranking model, we follow the approach of fine-tuning on pretrained BERT model (Devlin et al., 2019): The BERT model is firstly pretrained on unsupervised data, and then fine-tuned in ranking framework with supervised clickthrough data. To extract the text embedding, we use the embedding of a special token ”[CLS]”. The source/target embedding will later go through the interaction layer to generate deep features.

Previous work (Qiao et al., 2019) shows that directly training a representation based BERT ranking model does not yield good results. This is because the BERT fine-tuning requires a small learning rate (around 1e-5). Therefore, two optimizers are used in DeText, each with a different learning rate responsible for a different part of the model. For example, in our experiments (Section 6.1.4), we set 1e-5 for BERT components, and 1e-3 for other components. Using this dual learning rates strategy, a successful representation based ranking model with BERT can be trained.

In order to reduce the online serving latency and capture domain-specific semantics, we also pretrained a compact BERT model on LinkedIn’s in-domain data, named as LiBERT. More detailed can be found in Section 6.1.5.

DeText framework can support CNN for deep natural language processing in the text embedding layer. It is worth noting that we use word tokens instead of triletters as in prior work (Huang et al., 2013; Shen et al., 2014), since the latter lifts the computation (by an order of the character length of words). We follow the previous work (Kim, 2014) to generate the text embedding from word embedding matrix $E$. Specifically, it uses a one-dimensional CNN along the sentence length dimension. After max-pooling, the resulting text embedding vector has $f$ elements, where $f$ is the number of filters.

The multiple transformer layers in BERT is computationally time-consuming. Since DeText uses the representation based method, we are able to adopt the document embedding pre-computing approach for search ranking, as shown in the boxed section (doc pre-computing) in Figure 3. For offline, document embeddings are pre-computed with BERT and saved in an embedding store, which is a fast key-value store where key is the document id, and value is the pre-computed document embedding vectors. The store is refreshed on a regular basis (e.g., daily). For online, after document candidates are retrieved from search index, the corresponding document ids are used to fetch the document embeddings from pre-computed embedding store. The benefit of this approach is to have the heavy BERT online computation only happen on the queries. It can significantly save online computation time, since the document texts are much larger than the query texts. In the setting of 10 documents for one query, the online latency can be reduced from hundreds of milliseconds to tens of milliseconds with this deployment strategy.

DeText-CNN can be adopted with a different online integration strategy: real-time inference for both sources and targets. Compared to document embedding pre-computing, real-time inference can simplify online system design without the need of pre-computing or refreshing document embeddings. Hence the real-time inference could be a lightweight solution readily applied for many search engines. In this paper, we show that real-time inference can be achieved by (1) choosing a compact DeText structure without hurting relevance performance too much, and (2) two pass ranking that reduces 99 percentile (P99) latency.

We find that a compact CNN structure with small dimensions can perform well in our experiments (Table 7). This is mainly because traditional handcrafted features from the production systems already contain valuable information for ranking, so that the CNN model can be focusing on the signals that are missing in the traditional features.

Even with a simple network, the CNN computation time grows linearly with the number of retrieved documents. Therefore, we use a two pass ranking schema to bound the latency (Figure 3, two pass ranking box). The first ranker is a MLP (Pal and Mitra, 1992) with one hidden layer without the deep features, which is fast. After ranking, only the top ranked hundreds of documents are sent to the DeText-CNN model.²²2Note that the top hundreds of documents are in one worker, while the online ranking is distributed to many workers. Each worker is responsible for retrieving and ranking the documents on its own shard. This two pass ranking framework has several benefits: 1). easy to implement and deploy; 2). the MLP ranker can filter out a large amount of irrelevant documents, which provides a relatively small candidate set with high recall for CNN ranker; and 3). the latency is bounded since CNN is applied to a small set of candidates.

In this section, we show the great potential of applying DeText to applications beyond search ranking. We conducted extra experiments on two additional ranking tasks: job recommendation and query auto completion from job search.

For job recommendation, we model the job application probability. The input is a tuple of user id, job post id, and whether the user applied for the job or not. The source fields are from user profiles, including headline, job title, company, and skill. The target fields are from job posts, including job title, job company, job skill, and job country. We used logistic regression as a baseline that is close to production setting, and evaluated with AUC (Fawcett, 2006) metrics. For fair comparison, point-wise ranking (binary classification) is used with no hidden layer of MLP in DeText. Traditional features are kept the same as in the baseline model.

For query auto completion, the source fields are from member profiles, including headline, job title, and company; the target is the completed query. The baseline model is XGBoost with traditional hand-crafted features. We used the same set of traditional features in DeText with listwise LTR, and evaluated with MRR@10 (Bar-Yossef and Kraus, 2011), which is the reciprocal of the rank position of the correct answer.

Table 13 shows the offline results. DeText outperforms the baseline models in both tasks by a large margin, indicating that DeText is flexible enough to be applied in other ranking tasks.