PaperlayerView exact revision on arXiv

Source-linked AI summary

LexMAE: Lexicon-Bottlenecked Pretraining for Large-Scale Retrieval

Tao Shen, Xiubo Geng, Chongyang Tao, Can Xu, Xiaolong Huang, Binxing Jiao, Linjun Yang, Daxin Jiang

arXiv:2208.14754v2cs.IR

TL;DR

Language modeling favors low-entropy words, while lexicon-weighting retrieval needs importance assigned to high-entropy words relevant to query and document meaning. LexMAE uses an unsupervised CBoW bottleneck between an encoder and weakened decoder to learn importance-aware lexicon representations. Its fine-tuned retriever improves passage retrieval, delivers state-of-the-art document results, and transfers strongly across 12 BEIR datasets.

  • Problem

    Language modeling favors low-entropy words, whereas relevance-oriented lexicon weighting favors high-entropy words, creating an objective mismatch for retrieval.

    Source: §1 · p. 2

  • Method

    LexMAE learns document-specific lexicon-importance distributions through an unsupervised CBoW bottleneck between a language-modeling encoder and weakened masking-style decoder.

    Source: §1 · p. 2

  • Results

    LexMAE improves previous state-of-the-art passage retrieval by 1.5% MRR@10 with ∼13× speed-up, achieves 44.4% MRR@100 on document retrieval, and is state of the art on BEIR’s 12 datasets.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The conclusion reports state-of-the-art retrieval quality with high efficiency and state-of-the-art zero-shot transfer on BEIR after fine-tuning LexMAE.

    Source: §5 · p. 9

    The encoder retains an MLM objective despite its conflict with lexicon-importance learning, using it as regularization and to support robust learning from token-level noise.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

In large-scale retrieval, the lexicon-weighting paradigm, learning weighted sparse representations in vocabulary space, has shown promising results with high quality and low latency. Despite it deeply exploiting the lexicon-representing capability of pre-trained language models, a crucial gap remains between language modeling and lexicon-weighting retrieval -- the former preferring certain or low-entropy words whereas the latter favoring pivot or high-entropy words -- becoming the main barrier to lexicon-weighting performance for large-scale retrieval. To bridge this gap, we propose a brand-new pre-training framework, lexicon-bottlenecked masked autoencoder (LexMAE), to learn importance-aware lexicon representations. Essentially, we present a lexicon-bottlenecked module between a normal language modeling encoder and a weakened decoder, where a continuous bag-of-words bottleneck is constructed to learn a lexicon-importance distribution in an unsupervised fashion. The pre-trained LexMAE is readily transferred to the lexicon-weighting retrieval via fine-tuning. On the ad-hoc retrieval benchmark, MS-Marco, it achieves 42.6% MRR@10 with 45.8 QPS for the passage dataset and 44.4% MRR@100 with 134.8 QPS for the document dataset, by a CPU machine. And LexMAE shows state-of-the-art zero-shot transfer capability on BEIR benchmark with 12 datasets.

1 INTRODUCTION

Lexicon-weighting retrieval uses weighted sparse vocabulary-space representations, but its relevance-oriented importance weights conflict with language modeling’s preference for low-entropy words. LexMAE addresses this mismatch with an unsupervised CBoW bottleneck, and its fine-tuned retriever improves passage retrieval while achieving strong document and BEIR results.

1 INTRODUCTION

  • Large-scale retrieval fetches top query-relevant documents from huge collections, making efficiency a fundamental prerequisite.Query-independent document indexing and lightweight relevance metrics are common practices, often implemented with a two-tower structure.Source: §1 · p. 1
  • Lexicon-weighting retrieval represents queries and documents with weighted vocabulary-space terms, drawing on PLM coordinate knowledge to mitigate lexicon mismatch.It shares a high-level inspiration with BM25 but learns dynamic compression and expansion lexicons and their weights end to end.Source: §1 · p. 1
  • MLM favors low-entropy words that are often function words or collocations, whereas relevance-oriented lexicon weighting emphasizes high-entropy semantic words.The paper links this mismatch to cases where DistilBERT outperforms BERT-base and RoBERTa fails to converge under lexicon-weighting fine-tuning.Source: §1 · p. 2
  • LexMAE pre-trains a language-modeling encoder to produce document-specific lexicon-importance distributions, using a CBoW bottleneck to learn them without supervision.Its architecture combines the encoder, lexicon-bottlenecked module, and weakened masking-style decoder.Source: §1 · p. 2
  • LexMAE improves previous state-of-the-art passage retrieval by 1.5% MRR@10 with ∼13× speed-up, achieves 44.4% MRR@100 on document retrieval, and is state of the art on BEIR’s 12 datasets.The BEIR examples include Natural Questions, HotpotQA, and FEVER.Source: §1 · p. 2

2 RELATED WORK

Related work describes dense-vector retrieval as encoding queries and documents in a shared low-dimensional semantic space, while lexicon-weighting methods operate on learned term weights. The cited approaches differ in how they derive sparse representations, and prior methods directly fine-tune language models despite the objective mismatch.

PLM-based Dense-vector Retrieval

  • Dense-vector retrieval encodes queries and documents in a shared low-dimensional semantic space and scores relevance by spatial distance.The passage notes an objective gap between language-model pretraining and dense-vector document-compressing fine-tuning; dense representations also entail large indexes and high latency.Source: §2 · pp. 2–3

Lexicon-weighing Retrieval

  • Lexicon-weighting methods use neural-model term weights; cited approaches expand sparse representations through document-query concurrence or MLM coordinate terms.The passage says these methods directly fine-tune pretrained language models despite the mismatch between general language modeling and relevance-oriented lexicon weighting.Source: §2 · p. 3

3 LEXMAE: LEXICON-BOTTLENECKED MASKED AUTOENCODER

LexMAE uses a lexicon-importance distribution and a continuous bag-of-words bottleneck to connect language-modeling representations to reconstruction. A weakened masking-style decoder trains the bottleneck, and the pre-training objective combines encoder and decoder losses.

3.1 LANGUAGE MODELING ENCODER

  • The encoder uses masked language modeling to predict masked tokens from corrupted input, with 80% replaced by [MASK], 10% by random tokens, and 10% left unchanged.Its word embeddings, Transformer, and language-modeling head map input tokens to vocabulary-space logits.Source: §3.1 · p. 3

3.2 LEXICON-BOTTLENECKED MODULE

  • The lexicon-importance distribution weights a continuous bag-of-words bottleneck, which the decoder uses to reconstruct the original document.The distribution is derived by max-pooling encoder logits across sequence positions and normalizing them; the resulting bottleneck is a weighted sum of learnable word embeddings.Source: §3.2 · p. 4

3.3 WEAKENED MASKING-STYLE DECODER

  • The weakened decoder receives the bottleneck and an additionally masked input, then reconstructs tokens absent from the encoder input.Its weakened design uses aggressive masking and shallow Transformer layers, described as two layers.Source: §3.3 · p. 4

3.4 PRE-TRAINING OBJECTIVE & FINE-TUNING FOR LEXICON-WEIGHTING RETRIEVER

  • LexMAE’s pre-training loss adds the encoder and decoder losses, while tied embedding matrices are trained with gradients cut off through the bottleneck’s word-embedding matrix.The gradient cutoff focuses training on the lexicon-importance distribution rather than that matrix.Source: §3.4 · p. 5
  • For retrieval fine-tuning, the decoder is removed and the language-modeling encoder supplies lexicon-weighting representations for query-document contrastive learning.The retrieval loss includes a FLOPS regularizer for representation sparsity.Source: §3.4 · p. 5
  • Top-K sparsifying keeps only the highest-weighted document lexicons at inference, offering a targeted efficacy-efficiency trade-off with almost no extra overhead.It avoids tuning sparsity during fine-tuning by zeroing weights outside the top K.Source: §3.4 · p. 5

4 EXPERIMENT

Experiments evaluate LexMAE on MS-Marco, TREC Deep Learning, and BEIR, alongside retrieval efficiency, storage, transfer, and model-choice analyses. LexMAE achieves strong retrieval results, and ablations report that its performance is stable across several design changes.

4 EXPERIMENT

  • The experiments use official MS-Marco queries, evaluate passage retrieval on MS-Marco Dev and TREC DL’19/DL’20, and assess zero-shot transfer on 12 BEIR datasets.They also evaluate MS-Marco document retrieval; pretraining uses BERTbase initialization and takes 14 hours on 8 A100 GPUs.Source: §4 · pp. 5–6

4.1 MAIN EVALUATION

  • LexMAE improves on SimLM by +1.5% MRR@10 on MS-Marco passage retrieval and achieves new state-of-the-art performance on TREC Deep Learning 2019 and 2020.It also exceeds Co-Self-Disil by +5.1% MRR@10.Source: §4.1 · pp. 6–7, Table 1 · p. 6
  • On BEIR zero-shot transfer, LexMAE is best on 7 of 12 out-of-domain datasets and has the best overall ‘BEST ON’ and ‘AVERAGE’ metrics.The reported metrics exclude in-domain results.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 EFFICIENCY ANALYSIS AND COMPARISON

  • LexMAE has the best efficacy-efficiency trade-off among dense-vector, quantized-dense, and lexicon-based methods, with 1.5% higher MRR@10 and 14.1× acceleration than SimLM.With top-L sparsifying, it can achieve competitive performance with SimLM at 100+ QPS; with four tokens per passage, it scores 24.0% MRR@10 versus BM25’s 18.5%. Lexicon-based methods, including LexMAE, also require less index and per-document representation storage than dense-vector methods.Source: §4.2 · p. 7, Figure 2 · p. 8, Table 4 · p. 8

4.3 FURTHER ANALYSIS

  • A lexicon-first pipeline that reranks its top-K candidates with SimLM improves the previous state-of-the-art hybrid method by 2.4% MRR@10 on MS-Marco Dev.The alternative ensemble combines retrieval scores but incurs overhead from running both large-scale retrieval systems.Source: §4.3 · p. 8, Table 5 · p. 8
  • Without fine-tuning, LexMAE’s 110M-parameter lexicon-importance embedding beats BM25 and is competitive with 6.1B-parameter SGPT-CE on MS-Marco.Across all three fine-tuning stages, LexMAE achieves the best or competitive results.Source: §4.3 · pp. 8–9, Figure 3 · p. 8, Table 6 · p. 8

4.4 MODEL CHOICES & ABLATION STUDY

  • Softmax-norm CBoW bottlenecks outperform [CLS] and no-bottleneck variants in lexicon-weighting retrieval, although [CLS] has better decoding loss.Saturated-norm performs well early in fine-tuning, while softmax-norm performs better later; other masking strategies have minor downstream effects, and the authors describe LexMAE as stable to various changes.Source: §4.4 · p. 9, Table 7 · p. 9
  • Back-propagating to the word embedding matrix creates a learning shortcut and worsens fine-tuning results; sharing encoder and decoder LM heads or adding a bottleneck-specific head causes a minor drop.Replacing BERT with DistilBERT still outperforms a number of competitors.Source: §4.4 · p. 9

5 CONCLUSION

LexMAE pretraining is presented as addressing the mismatch between masked language modeling and relevance-oriented lexicon importance. The authors report strong zero-shot results and state-of-the-art retrieval quality and efficiency after fine-tuning.

5 CONCLUSION

  • LexMAE pretraining is proposed to alleviate the mismatch between masked language modeling and relevance-oriented lexicon importance.Source: §5 · p. 9
  • After fine-tuning on a large-scale retrieval benchmark, LexMAE achieves state-of-the-art retrieval quality with high efficiency.Source: §5 · p. 9
  • LexMAE shows strong zero-shot performance after pretraining and state-of-the-art zero-shot transfer on BEIR.The authors also report that efficacy-efficiency analyses of retrieval latency and storage memory verify the superiority of fine-tuned LexMAE.Source: §5 · p. 9

A.1 FINE-TUNING PIPELINE

The fine-tuning pipeline adapts a dense-vector retrieval pipeline and replaces its on-the-fly reranker with an off-the-shelf reranker. Training proceeds through BM25 negatives, retriever-generated hard negatives, and reranker distillation.

A.1 FINE-TUNING PIPELINE

  • An off-the-shelf reranker from Zhou et al. (2022b) replaces the retriever-specific reranker trained on the fly in Wang et al. (2022), avoiding its high computation overheads.The fine-tuned reranker is directly adopted from Zhou et al. (2022b) to avoid expensive reranker training.Source: §A.1 · p. 15, Figure 6 · p. 16
  • Stage 1 samples BM25 top-K1 negatives for each query and trains the retriever with a contrastive loss and FLOPS regularization.The stage-1 loss is given in Equation (11).Source: §A.1 · p. 15
  • Stage 2 samples hard negatives from the top-K2 candidates scored by the stage-1 retriever and applies contrastive training with FLOPS regularization.The stage-2 loss is given in Equation (12).Source: §A.1 · p. 15
  • Stage 3 samples top-K3 hard negatives from the stage-2 retriever and distills a well-trained reranker alongside the contrastive objective.The reranker is an expensive cross-encoder whose parameters are frozen during KL-divergence distillation.Source: §A.1 · p. 16

A.2 FINE-TUNING SETUPS

Fine-tuning uses shared training settings across its three stages, with different batch and negative counts for stages 1–2 versus stage 3. The loss weights are tuned, while the random seed is fixed.

A.2 FINE-TUNING SETUPS

  • The three stages share a 2 × 10^-5 learning rate, three epochs, LexMAE initialization, 144-token document length, and 32-token query length.Fine-tuning is limited to one A100 GPU, unlike the four GPUs used by Wang et al. (2022).Source: §A.2 · p. 16
  • Stages 1 and 2 use batches of 24 queries with 1 positive and 15 negatives; stage 3 uses 16 queries with 1 positive and 23 negatives.The smaller stage-3 batch allows more negatives while fitting one GPU's memory.Source: §A.2 · p. 16
  • Negative-sampling depths are K1 = 1000 for BM25 and K2 = K3 = 200 for the other negatives; tuned loss weights are λ1 = 0.002 and λ2 = λ3 = 0.008.The random seed is fixed at 42 without tuning.Source: §A.2 · p. 16

B LEXICON-WEIGHTING INFERENCE FOR LARGE-SCALE RETRIEVAL

Lexicon-weighting retrieval uses real-valued sparse representations for relevance scoring, then converts them into lexicons and virtual frequencies for term-based retrieval. Anserini builds an inverted index that is queried with similarly converted query representations.

B LEXICON-WEIGHTING INFERENCE FOR LARGE-SCALE RETRIEVAL

  • The relevance metric is the dot product between real-valued sparse lexicon-weighting representations, which support gradient back-propagation and end-to-end learning.Source: B LEXICON-WEIGHTING INFERENCE FOR… · p. 16
  • For term-based retrieval, nonzero vector elements become lexicons and each virtual frequency is quantized as ⌊100 × v⌋.The passage notes that directly using real-valued sparse representations is inefficient and infeasible for open-source term-based systems such as LUCENE and Anserini.Source: B LEXICON-WEIGHTING INFERENCE FOR… · p. 16
  • The inference workflow converts document vectors to lexicons and frequencies, builds an Anserini inverted index, converts each test query similarly, and retrieves top document candidates.Source: B LEXICON-WEIGHTING INFERENCE FOR… · p. 16

C EXPLANATION OF DIFFERENT RECALL METRICS

The section distinguishes official all-positive-macro Recall@N from DPR’s one-positive-enough Recall@N and reports that LexMAE remains the best retriever among its competitors under the latter calculation.

C EXPLANATION OF DIFFERENT RECALL METRICS

  • Official Recall@N follows the evaluation definitions used by trec_eval and Anserini.The text calls this calculation all-positive-macro Recall@N.Source: C EXPLANATION OF DIFFERENT RECALL… · p. 17
  • Official all-positive-macro Recall@N is usually lower than DPR’s one-positive-enough Recall@N, especially at smaller N.The all-positive calculation accounts for multiple positive documents, whereas one-positive-enough Recall@N follows DPR.Source: C EXPLANATION OF DIFFERENT RECALL… · p. 17
  • LexMAE remains the best retriever among its competitors under one-positive-enough Recall@N.The section presents this calculation separately for more precise comparisons.Source: C EXPLANATION OF DIFFERENT RECALL… · p. 17

D SPARSIFYING LEXICON REPRESENTATIONS

The section introduces inference-time top-K sparsification to control document lexicon representations without repeated fine-tuning, and specifies how the thresholded representation is built and used.

D SPARSIFYING LEXICON REPRESENTATIONS

  • Lexicon-weighting representations are described as intrinsically efficient and compatible with term-based retrieval systems such as BM25.The passage contrasts them with dense-vector methods using product quantization and inverted files to trade performance for memory and time efficiency.Source: D SPARSIFYING LEXICON REPRESENTATIONS · p. 17
  • Tuning λ for sparsity requires repeated fine-tuning and does not reliably predict the resulting sparsity.Earlier methods that constrain activated lexicons also require extra fine-tuning efforts.Source: D SPARSIFYING LEXICON REPRESENTATIONS · p. 17
  • Top-K sparsification keeps only the highest-weighted lexicons in each document representation and sets the rest to zero.The threshold t is the K-th largest value in the representation, and the mask selects values in its top K.Source: D SPARSIFYING LEXICON REPRESENTATIONS · p. 17
  • If K is below the number of already activated lexicons, top-K sparsification leaves the representation unchanged.The resulting sparsified representations are used to build the inverted index.Source: D SPARSIFYING LEXICON REPRESENTATIONS · p. 17
  • The inference-time method offers controllable sparsification with almost zero extra overhead.Multiple K values can be derived from the same original representation, so fine-tuning and document embedding run only once.Source: D SPARSIFYING LEXICON REPRESENTATIONS · p. 17

E MORE EXPERIMENTS

The experiments report LexMAE’s advantage over prior pre-training objectives and retrieval-and-rerank pipelines. The background contrasts dense-vector and lexicon-weighting representations and their relevance-matching approaches.

E MORE EXPERIMENTS

  • LexMAE can outperform many state-of-the-art retrieval-and-rerank pipelines.The compared rerankers are costly because they process each query-document concatenation, unlike query-agnostic bi-encoder representations.Source: E MORE EXPERIMENTS · p. 18
  • LexMAE improves the previous best by 1.3% MRR@10 on MS-Marco Dev.The comparison fine-tunes the pre-training objectives on MS-Marco BM25 negatives.Source: E MORE EXPERIMENTS · p. 18

F BACKGROUND: DENSE-VECTOR AND LEXICON-WEIGHTING RETRIEVAL

  • Recent PLMs serve as neural encoders that are fine-tuned on downstream retrieval datasets.The background groups PLM retrieval methods into dense-vector and lexicon-weighting paradigms.Source: F BACKGROUND: DENSE-VECTOR AND… · p. 18
  • Dense-vector methods encode queries and documents as low-dimensional real-valued vectors and score them with measures such as dot product or cosine similarity.The passage notes limitations in index size, retrieval latency, and preservation of lexicon overlap.Source: F BACKGROUND: DENSE-VECTOR AND… · p. 18
  • Lexicon-weighting methods encode text as weighted sparse vocabulary vectors and compute relevance with lexical matching, such as sparse dot product or BM25.The passage describes contextual weighting and aggregation of vocabulary lexicons for documents and queries.Source: F BACKGROUND: DENSE-VECTOR AND… · p. 19
Loading 2208.14754v2…