LADER: Log-Augmented DEnse Retrieval for Biomedical Literature Search

Abstract

Queries with similar information needs tend to have similar document clicks, especially in biomedical literature search engines where queries are generally short and top documents account for most of the total clicks. Motivated by this, we present a novel architecture for biomedical literature search, namely Log-Augmented DEnse Retrieval (LADER), which is a simple plug-in module that augments a dense retriever with the click logs retrieved from similar training queries. Specifically, LADER finds both similar documents and queries to the given query by a dense retriever. Then, LADER scores relevant (clicked) documents of similar queries weighted by their similarity to the input query. The final document scores by LADER are the average of (1) the document similarity scores from the dense retriever and (2) the aggregated document scores from the click logs of similar queries. Despite its simplicity, LADER achieves new state-of-the-art (SOTA) performance on TripClick, a recently released benchmark for biomedical literature retrieval. On the frequent (“HEAD”) queries, LADER largely outperforms the best retrieval model by 39% relative NDCG@10 (0.338 v.s. 0.243). LADER also achieves better performance on the less frequent (“TORSO”) queries with 11% relative NDCG@10 improvement over the previous SOTA (0.303 v.s. 0.272). On the rare (“TAIL”) queries where similar queries are scarce, LADER still compares favorably to the previous SOTA method (NDCG@10: 0.310 v.s. 0.295). On all queries, LADER can improve the performance of a dense retriever by 24%–37% relative NDCG@10 while not requiring additional training, and further performance improvement is expected from more logs. Our regression analysis has shown that queries that are more frequent, have higher entropy of query similarity and lower entropy of document similarity, tend to benefit more from log augmentation.

1. INTRODUCTION

Biomedical literature search is an essential step for knowledge discovery and clinical decision support [8, 12, 17]. It has several distinct properties from other information retrieval (IR) tasks: (1) Most queries are short. The average length of a query is about 3.5 tokens in PubMed¹, a widely used biomedical literature search engine [7, 9, 10]; (2) Large-scale relevant query-document pairs can be easily collected from the user click logs. TripClick [29], a recently released benchmark for biomedical literature retrieval, contains 1.3 million query-document relevance signals collected from the Trip Database² logs; (3) Users mostly browse the documents on the first page [7], and the top 31% most clicked documents account for 80% clicks in the released Trip Database logs. These unique characteristics motivate the augmentation of biomedical literature search by directly retrieving from the click logs of similar queries, since queries with similar information needs tend to have similar document clicks [1, 38, 40]. This approach is essentially similar to recent retrieval augmentation methods [11, 20, 39]. Fig 1 shows an example query and clicked documents from its similar queries in the training set of TripClick. The objective of this work is to augment biomedical literature search with such documents.

Figure 1:

An example of using click logs of retrieved similar queries to augment biomedical literature search.

One model can simply return clicked documents from the logs of similar queries as the retrieval results, which is analogous to memory-based collaborative filtering in recommendation systems [33]. However, this might miss relevant documents that have not been clicked before. Therefore, clicked documents from logs are better used to augment an existing retriever, and they are theoretically retriever-agnostic. We choose to use dense retrievers, where queries and documents are encoded and matched in low-dimensional embedding space, because recent work has shown that dense retrievers based on pre-trained language models such as BERT [6] outperform traditional sparse retrievers on various tasks [19, 21, 37].

In this paper, we introduce Log-Agumented DEnse Retrieval (LADER), which is a novel and simple plug-in module that augments dense retrievers by interpolating the document scores aggregated from logs of similar queries and the original scores from a dense retriever. We conduct experiments on the recently introduced TripClick benchmark [29] and show that: (1) LADER outperforms previous state-of-the-art (SOTA) models on queries of all frequency groups, with close to 40% relative NDCG@10 improvement on the frequent (“HEAD”) queries; (2) LADER improves the backbone dense retriever by 24%–37% relative NDCG@10 while not requiring additional training, and it is expected to have further performance improvements with more logs. Our regression analysis shows that queries with higher frequency, lower entropy of similar query scores, and higher entropy of similar document scores, tend to benefit more from log augmentation.

2. METHODS

We describe the training of our backbone dense retriever in §2.1, and introduce how to do inference with the LADER module in §2.2.

2.1. Dense Retriever Training

We train a dense retriever that contains a query encoder (QEnc) and a document encoder (DEnc), both are 12-layer transformer (Trm) encoders [34]. The encoders are initialized with PubMedBERT-base [13], which is a biomedical domain-specific BERT model. During training, each instance contains a triple of a query $q$, a relevant document $d^{+}$ for the query, and a non-relevant document $d^{-}$ for the query. They are fed to their respective encoders and the last [CLS] hidden states are used as their embeddings:

$$\begin{gathered} E(q)=\mathrm{Q}\mathrm{E}\mathrm{n}\mathrm{c}(q)=\mathrm{T}\mathrm{r}\mathrm{m}([\mathrm{C}\mathrm{L}\mathrm{S}]q[\mathrm{S}\mathrm{E}\mathrm{P}]{)}_{[\mathrm{C}\mathrm{L}\mathrm{S}]} \\ E(d)=\mathrm{D}\mathrm{E}\mathrm{n}\mathrm{c}(d)=\mathrm{T}\mathrm{r}\mathrm{m}{\left([\mathrm{C}\mathrm{L}\mathrm{S}]d_{\text{title}}[\mathrm{S}\mathrm{E}\mathrm{P}]d_{\mathrm{a}\mathrm{b}\mathrm{s}\mathrm{t}\mathrm{r}\mathrm{a}\mathrm{c}\mathrm{t}}[\mathrm{S}\mathrm{E}\mathrm{P}]\right)}_{[\mathrm{C}\mathrm{L}\mathrm{S}]} \end{gathered}$$

where [CLS] and [SEP] are special tokens used in BERT.

We optimize the model with a combination of two loss functions: an in-batch negative log-likelihood loss ${ℒ}^i$ [19] and a triplet contrastive loss ${ℒ}^t$ [32]. The in-batch negative log-likelihood loss helps the dense retriever distinguish between positive documents and all other documents in the mini-batch:

$${ℒ}^i=-\mathrm{l}\mathrm{o}\mathrm{g}\frac{\mathrm{e}\mathrm{x}\mathrm{p}(E(q{)}^{T}E\left(d^{+}\right))}{\sum _{i\in B}\mathrm{e}\mathrm{x}\mathrm{p}(E(q{)}^{T}E\left(d_i\right))}$$

where $B$ denotes all documents from the mini-batch. The triplet loss further contrasts positive documents and hard negative documents:

$${ℒ}^t=\mathrm{m}\mathrm{a}\mathrm{x}\left(0,\mathrm{D}\mathrm{i}\mathrm{s}\mathrm{t}\left(E(q),E\left(d^{+}\right)\right)-\mathrm{D}\mathrm{i}\mathrm{s}\mathrm{t}\left(E(q),E\left(d^{-}\right)+\alpha \right)\right.$$

where Dist denotes a distance metric and $\alpha$ is the target margin between the positive and negative pairs.

The final loss is a weighted sum of the two loss functions:

$$ℒ=\beta {ℒ}^i+(1-\beta ){ℒ}^t$$

where $\beta$ is a hyper-parameter for loss weighting. We train parameters in QEnc and DEnc end-to-end by gradient-based optimizers.

2.2. LADER Inference

The overall architecture of LADER is shown in Figure 2. We use the trained QEnc to encode all queries in the training dataset and use the trained DEnc to encode all documents in the collection, getting $𝒬\in {\mathbb{R}}^{N_q\times 768}$ and $𝒟\in {\mathbb{R}}^{N_d\times 768}$, respectively. $N_q$ denotes the number of training queries and $N_d$ denotes the number of documents in the collection.

Figure 2:

Overall architecture of the LADER model.

During inference, we first encode a new query $q$ by QEnc. Then, we conduct a maximum inner-product search (MIPS) to get the top-$m$ most similar queries $\left(q_1,q_2,\dots ,q_m\right)$ and their inner-product similarities $\left(s_q^1,s_q^2,\dots ,s_q^m\right)$ from the training set, as well as the top-$n$ most similar documents $\left(d_1,d_2,\dots ,d_n\right)$ and their similarities $\left(s_d^1,s_d^2,\dots ,s_d^n\right)$ from the collection:

$$\begin{gathered} \left(q_1,s_q^1\right),\left(q_2,s_q^2\right),\dots ,\left(q_m,s_q^m\right)=\mathrm{M}\mathrm{I}\mathrm{P}\mathrm{S}(\mathrm{Q}\mathrm{E}\mathrm{n}\mathrm{c}(q),𝒬) \\ \left(d_1,s_d^1\right),\left(d_2,s_d^2\right),\dots ,\left(d_n,s_d^n\right)=\mathrm{M}\mathrm{I}\mathrm{P}\mathrm{S}(\mathrm{Q}\mathrm{E}\mathrm{n}\mathrm{c}(q),𝒟) \end{gathered}$$

The similarity scores are further normalized by a softmax function:

$$\begin{gathered} {\tilde{s}}_q^1,{\tilde{s}}_q^2,\dots ,{\tilde{s}}_q^m=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(s_q^1,s_q^2,\dots ,s_q^m\right) \\ {\tilde{s}}_d^1,{\tilde{s}}_d^2,\dots ,{\tilde{s}}_d^n=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(s_d^1,s_d^2,\dots ,s_d^n\right) \end{gathered}$$

We denote the mapping from a query to its relevant documents in the training set as Rel : $q\mapsto \left\{d\right\}$. The final score of a document $d$ in the candidate set $\left\{d_1,d_2,\dots ,d_n\right\}\cup \left\{\mathrm{R}\mathrm{e}\mathrm{l}\left(q_1\right),\mathrm{R}\mathrm{e}\mathrm{l}\left(q_2\right),\dots ,\mathrm{R}\mathrm{e}\mathrm{l}\left(q_m\right)\right\}$ is a weighted sum computed as follows:

$$\mathrm{S}\mathrm{c}\mathrm{o}\mathrm{r}\mathrm{e}(d)=\sum _{j\in [1,\dots ,n]}{\tilde{s}}_d^j\mathbf{1}\left(d=d_j\right)+\sum _{i\in [1,\dots ,m]}\lambda {\tilde{s}}_q^i\mathbf{1}\left(d\in \mathrm{R}\mathrm{e}\mathrm{l}\left(q_i\right)\right)$$

where the former part is the dense retrieval score and the latter part is the log-augmentation score from similar queries, 1 is the indicator function, and $\lambda$ is a hyper-parameter to control the extent of log augmentation. We return the candidate documents ranked by their final scores to the input query.

3. EXPERIMENTS

3.1. Settings

Dataset.

We evaluate our LADER method on the TripClick benchmark [29], which contains 692k unique queries and 1.5M documents (PubMed abstracts). Based on their frequencies, the queries are divided into HEAD (>44), TORSO (6–44), and TAIL (<6) subsets. Both the validation and test sets contain 1,175 queries for each HEAD, TORSO and TAIL subset. Following [14, 29], we use two sets of relevance scores: The “RAW” relevance is used to judge TORSO and TAIL queries, where clicked documents have a score of 1 and other documents have a score of 0. The Graded Document Click-Through Rate [3, 4] (“DCTR”) relevance is used to judge the HEAD queries, where the document score is defined as the number of clicks divided by the number of exposure. To train the dense retriever (§2.1), we use 10M triples released by [14], where each triple contains a TripClick training query, a clicked document, and a non-relevant document sampled from BM25 negatives.

Configuration.

We implement LADER with PyTorch [28] and HuggingFace’s libraries [35]. For training the dense retriever, we use the AdamW optimizer [22,25] for 20k steps with the learning rate of 2e-5 and weight decay of 1e-2, batch size of 256, 10k warmup steps, the cosine learning rate decay schedule, $\alpha =5.0$, $\beta =0.9$, and the Euclidean distance for the triplet loss. For inference, we implement MIPS with FAISS’s FlatIP index [18], $m=n=1000,N_q=685,649$, and $N_d=1,523,871$. We use $\lambda =0.5$ for the HEAD and TORSO queries, and $\lambda =0.2$ for the TAIL queries. We also experiment with two ablations: LADER w/o Log Augmentation (LA) where LA scores are set to 0, and LADER w/o Dense Retrieval (DR) where DR scores are set to 0.

Baseline methods for comparison.

We compare LADER with various baselines in [29] on all queries, including BM25 [30], RM3 Pseudo Relevance Feedback (PRF) [23, 26], Position Aware Convolutional Recurrent Relevance Matching (PACRR) [16], Match Pyramid (MP) [27], Kernel-based Neural Ranking Model (KNRM) [36], Convolutional KNRM (ConvKNRM) [5], and Transformer-Kernel (TK) [15]. For the HEAD queries, we further compare with current SOTA methods [14] based on pre-trained language models, including biencoders initialized by different BERT models [2, 13, 31]. Following their respective metrics, we compare with benchmark baselines [29] using NDCG@10, MRR, and Recall@10, and compare with HEAD query SOTA [14] using NDCG@10, MRR@10, and also Recall@1k.

3.2. Main Results

Comparison with benchmark baselines.

Table 1 shows comprehensive comparisons of LADER with a variety of benchmark baselines [29] on all queries. LADER outperforms all benchmark baselines on each query subset, and on all metrics except the Recall@10 for the TAIL queries. On the most frequent HEAD queries, LADER outperforms the best benchmark baseline method (TK) by large margins, showing 63% (0.338 v.s. 0.208), 53% (0.664 v.s. 0.434), and 61% (0.304 v.s. 0.189) relative gains on NDCG@10, MRR, and Recall@10, respectively. On the less frequent TORSO queries, the performance improvements over the previous SOTA are decent with 10% to 12% relative gains on all metrics. On the rare TAIL queries, LADER still performs favorably than previous SOTA on NDCG@10 (0.310 v.s. 0.295) and MRR (0.306 v.s. 0.280), but the performance is slightly lower on Recall@10 (0.449 v.s. 0.459).

Table 1:

Model performance on the TripClick test sets.

Model	HEAD (DCTR relevance)			TORSO (RAW relevance)			TAIL (RAW relevance)
	NDCG@10	MRR	Recall@10	NDCG@10	MRR	Recall@10	NDCG@10	MRR	Recall@10
BM25 [30]	0.140	0.290	0.138	0.206	0.283	0.262	0.267	0.258	0.409
RM3 PRF [23, 26]	0.141	0.300	0.136	0.194	0.261	0.254	0.242	0.227	0.384
PACRR [16]	0.175	0.356	0.162	0.212	0.302	0.262	0.267	0.261	0.409
MP [27]	0.183	0.372	0.173	0.243	0.347	0.297	0.281	0.280	0.409
KNRM [36]	0.191	0.393	0.173	0.235	0.338	0.283	0.272	0.265	0.409
ConvKNRM [5]	0.198	0.420	0.178	0.243	0.358	0.288	0.271	0.265	0.409
TK [15]	0.208	0.434	0.189	0.272	0.381	0.321	0.295	0.279	0.459
LADER (ours)	0.338	0.664	0.304	0.303	0.427	0.353	0.310	0.306	0.449
LADER w/o LA	0.247	0.532	0.237	0.241	0.350	0.293	0.260	0.257	0.394
LADER w/o DR	0.324	0.649	0.284	0.266	0.396	0.298	0.232	0.236	0.330

Baseline results are from [29]. LA: log augmentation; DR: dense retrieval.

Comparison with BERT DOT.

In Table 2, we compare LADER with more recent BERT-based retrievers (BERT DOT) [14], whose results are only available on the HEAD queries. Compared with the BERT DOT dense retriever, LADER has 39% (0.338 v.s. 0.243), 24% (0.659 v.s. 0.530), and 7% (0.893 v.s. 0.828) relative improvement on NDCG@10, MRR@10, and Recall@1k, respectively. This shows the effectiveness of log augmentation since our backbone dense retriever (LADER w/o LA) performs similarly to their counterparts.

Table 2:

LADER results on the HEAD queries compared to BERT DOT [14]

Model	NDCG@10	MRR@10	R@1k
BM25 [30]	0.140	0.276	0.834
BERT DOT [14]
w/ DistillBERT [31]	0.236	0.512	0.813
w/ SciBERT [2]	0.243	0.530	0.793
w/ PubMedBERT [13]	0.235	0.509	0.828
LADER (ours)	0.338	0.659	0.893
LADER w/o LA	0.247	0.526	0.878
LADER w/o DR	0.324	0.644	0.670
Log-Augmented (LA−)
Sparse retriever (BM25)	0.312	0.598	0.889
Raw PubMedBERT [13]	0.137	0.288	0.453

LA: log augmentation; DR: dense retrieval.

Improvement over the backbone dense retriever.

On all queries, LADER improves the performance of the backbone dense retriever by 24%–37% relative NDCG@10 (LADER v.s. LADER w/o LA) while not requiring additional training. The performance improvement is more significant on the HEAD queries than on the TORSO or TAIL queries, which will be further analyzed in the next section.

3.3. Analysis

Importance of trained dense retrievers.

In Table 2, we show the results of log augmentation with different backbone retrievers. Log-augmented raw PubMedBERT performs much worse than LADER (0.137 v.s. 0.338 NDCG@10), suggesting the importance of the backbone dense retriever to be trained on retrieval tasks (§2.1). We also replace the dense retriever in LADER with BM25, which we denote as LABM25 and implement with Pyserini [24]. LABM25 greatly improves the original BM25 baseline by 122% (0.312 v.s. 0.140 NDCG@10), which proves that the effectiveness of log augmentation is retriever-agnostic. However, LADER still outperforms LABM25 (0.338 v.s. 0.312 NDCG@10), indicating that the potential of log augmentation is better harnessed by dense retrievers.

Effects of log size.

In Figure 3, we show the performance of LADER using different proportions of queries sampled from the training set for the log augmentation. For all queries, the performance improves with the number of queries to retrieve, and since these curves have not saturated yet, more performance gains are expected with more logs. We also find that on lower frequency queries, more training queries are required to retrieve from for LADER to outperform current SOTA methods.

Figure 3:

NDCG@10 of LADER with different proportions of training queries to retrieve in log-augmentation.

What queries gain more from log-augmentation?

We collect 9 features for each query, including query length (QL), query group (HEAD, TORSO or TAIL), ENT(Q): the entropy of $\left[{\tilde{s}}_q^1,{\tilde{s}}_q^2,\dots ,{\tilde{s}}_q^m\right]$, ENT(D): the entropy of $\left[{\tilde{s}}_d^1,{\tilde{s}}_d^2,\dots ,{\tilde{s}}_d^n\right],$ E(REL): the expectation of relevant document number, and the average number of relevant documents in the top-1 and top-5 similar queries (REL1 and REL5). We fit a linear regression model to predict the gain of NDCG@10 by log-augmentation (LADER v.s. LADER w/o LA) using min-max normalized query features. The feature coefficients shown in Figure 4 indicate that: (1) Query frequency is the most important feature: more frequent (HEAD group) queries benefit more than average from log augmentation, while rare (TAIL group) queries benefit less than average; (2) Queries without very similar documents (higher ENT(D)), benefit more than queries with very similar documents; (3) Queries with very similar queries in the log (lower ENT(Q)), benefit more from log augmentation than queries without.

Figure 4:

Feature coefficients in the regression analysis. Positive values indicate more gains from log augmentation.

4. CONCLUSIONS AND LIMITATIONS

We present LADER, a simple and novel plug-in module that uses search logs to augment dense retrievers. Our results show that LADER achieves new SOTA on TripClick, and can largely improve the performance of its backbone retriever without additional training. We also provide thorough analyses of its characteristics.

One limitation of LADER is that it increases the search latency by about $N_q/N_d$ (45% for TripClick) due to the additional query-to-query retrieval step. Another drawback of this study is that we only use data from one search engine. It remains interesting to test the generalizability of LADER to a different search engine, which will be beneficial for cold-starting new literature search initiatives.