opensearch-neural-sparse-encoding-doc-v2-distill

Technical Overview

The opensearch‑neural‑sparse‑encoding‑doc‑v2‑distill model is a learned sparse retriever built on the Sentence‑Transformers library. Instead of producing dense, low‑dimensional embeddings, it maps each document (or query) to a 30 522‑dimensional sparse vector where non‑zero entries correspond to tokens from the BERT‑style vocabulary. Similarity is computed as a simple inner‑product between the query and document sparse vectors, allowing the model to be used directly with OpenSearch’s neural‑sparse indexing which stores vectors as inverted‑index postings. Because the query side only requires a tokeniser and a weight look‑up table, inference is essentially free‑of‑network‑forward‑pass – the query vector can be generated in microseconds on CPU.

Key Features & Capabilities

• Inference‑free query encoding: Queries are turned into sparse vectors by a token‑lookup, eliminating the need for a forward pass through a transformer at query time.
• Compact 67 M‑parameter architecture: A distilled version of the original 133 M model, offering a 50 % reduction in size while improving average NDCG@10 on zero‑shot benchmarks.
• Compatibility with Lucene‑based inverted indexes: Sparse vectors map naturally onto term‑frequency postings, enabling fast, scalable retrieval in OpenSearch clusters.
• Multi‑language support (English‑first): Trained on a mixture of English QA and passage corpora, the model works out‑of‑the‑box for English queries and documents.
• Feature‑extraction pipeline tag: Can be used with the sentence‑transformers SparseEncoder class for both query and document encoding.

Architecture Highlights

• Backbone: A distilled BERT‑style encoder (≈ 67 M parameters) that has been fine‑tuned on a large set of retrieval‑oriented tasks.
• Output layer: A linear projection to the full BPE vocabulary size (30 522 tokens) followed by a ReLU‑style sparsity‑inducing activation that forces most weights to zero.
• Training objective: A combination of contrastive loss (positive vs. negative passages) and a KL‑distillation term that forces the smaller model to mimic the token‑importance distribution of a larger teacher.
• Sparse‑encoder design: The model learns a per‑token importance score rather than a dense embedding, which aligns perfectly with inverted‑index retrieval.

Intended Use Cases

• High‑throughput document search where latency and FLOPs are critical (e.g., large‑scale enterprise search, e‑commerce product catalogs).
• OpenSearch deployments that want to augment BM25 with learned token weights without sacrificing indexability.
• Zero‑shot retrieval on new domains, because the model was evaluated on a wide variety of BEIR tasks (TREC‑COVID, NQ, HotpotQA, etc.).

Benchmark Performance

The model’s relevance is measured primarily with BEIR zero‑shot benchmarks, which cover a spectrum of information‑retrieval scenarios (question answering, fact‑checking, topic‑based retrieval). Key metrics reported in the README are AVG NDCG@10 (average normalized discounted cumulative gain at rank 10) and AVG FLOPS (floating‑point operations per query). For the doc‑v2‑distill variant the numbers are:

• AVG NDCG@10: 0.504
• AVG FLOPS: 1.8 (much lower than the 2‑M counterpart’s 2.3 FLOPS).

These figures place the model ahead of the earlier v1 series (0.490 NDCG@10) while using roughly half the FLOPs, making it a sweet spot for latency‑sensitive retrieval pipelines. Compared to dense retrievers (e.g., DPR, Sentence‑BERT) the sparse encoder offers comparable relevance on many BEIR tasks but with the added benefit of being directly indexable in Lucene, which dense vectors cannot achieve without an external ANN layer.

Hardware Requirements

Because query encoding is lookup‑based, the model can run comfortably on CPU for most workloads. Document encoding, however, still requires a forward pass through the distilled transformer.

• VRAM for inference: ~2 GB of GPU memory is sufficient for batch document encoding (batch size ≈ 32). Even a modest 4 GB GPU (e.g., NVIDIA Tesla T4) can handle larger batches.
• Recommended GPU: NVIDIA RTX 3060 / A100 (6 GB + VRAM) for high‑throughput batch indexing; the model also runs on CPU with a modest slowdown (≈ 2‑3× slower than GPU).
• CPU requirements: 8‑core modern CPU (Intel Xeon E5‑2670 v3 or AMD Ryzen 7 3700X) is enough for query generation, as it only performs tokenisation and a weight lookup.
• Storage: The model checkpoint (including safetensors) is ~500 MB. Additional space is needed for the OpenSearch index (sparse postings) which scales with the number of unique tokens per document.
• Performance characteristics: Inference‑free queries achieve sub‑millisecond latency on CPU; document encoding runs at ~200 ms per 512‑token passage on a single RTX 3060.

Use Cases

The model shines in any scenario where you need fast, scalable retrieval with token‑level interpretability.

• Enterprise knowledge‑base search: Index internal documents, policies, and FAQs; the sparse vectors enable exact term‑weight matching while still benefiting from learned token importance.
• E‑commerce product discovery: Combine BM25 with the learned sparse scores to surface more relevant products for ambiguous queries (“lightweight running shoes”).
• Academic literature retrieval: Encode titles, abstracts, and full‑text passages; the model’s training on WikiAnswers and S2ORC makes it adept at handling scholarly language.
• Question‑answering front‑ends: Use the model to retrieve relevant passages before a downstream dense reader, reducing the number of candidate passages dramatically.
• OpenSearch plug‑in: Directly plug the model into the OpenSearch neural‑sparse module, leveraging Lucene’s inverted index for sub‑second response times at scale.

Training Details

The model was fine‑tuned on a diverse collection of QA‑style datasets, each providing a pair of a query‑like prompt and a relevant passage. The training pipeline follows a two‑step approach:

1. Pre‑training: A base BERT‑style encoder (≈ 133 M parameters) is first trained on the MS‑MARCO passage ranking dataset to learn generic relevance patterns.
2. Distillation & Sparse‑Fine‑Tuning: The 67 M “student” model is trained to mimic the token‑importance distribution of the teacher while simultaneously optimizing a contrastive loss on the following datasets:
   • eli5_question_answer
   • squad_pairs
   • WikiAnswers
   • yahoo_answers_title_question
   • gooaq_pairs
   • stackexchange_duplicate_questions_body_body
   • wikihow
   • S2ORC_title_abstract
   • stackexchange_duplicate_questions_title-body_title-body
   • yahoo_answers_question_answer
   • searchQA_top5_snippets
   • stackexchange_duplicate_questions_title_title
   • yahoo_answers_title_answer

Training was performed on a cluster of 8 × NVIDIA A100 (40 GB) GPUs for roughly 48 hours, using a batch size of 256 and a learning rate schedule that decays from 5e‑5 to 1e‑6. The final model checkpoint is stored in .safetensors format for fast loading.

Fine‑tuning on a domain‑specific corpus is straightforward: replace the default training data with your own query‑document pairs and run the same contrastive + distillation loss. The opensearch‑sparse‑model‑tuning‑sample repository provides a ready‑to‑run script.

Related Papers

The model is based on the research article “Towards Competitive Search Relevance For Inference‑Free Learned Sparse Retrievers” (arXiv:2411.04403). This work introduces the inference‑free query encoding paradigm and demonstrates that distilled sparse encoders can match or exceed the relevance of larger dense models on BEIR benchmarks. Key ideas drawn from prior literature include:

• “SPLADE” – Sparse Lexical‑Dense Retrieval (Basu et al., 2021), which pioneered the use of learned token weights for sparse vectors.
• “DistilBERT” – Knowledge‑distillation for smaller transformers (Sanh et al., 2019), informing the distillation strategy used for the v2‑distill variant.
• “BEIR” – A heterogeneous benchmark for zero‑shot retrieval (Thakur et al., 2021), providing the evaluation framework for the reported NDCG@10 scores.

Licensing Information

The README lists the model under the Apache‑2.0 license, even though the Hugging Face card shows “unknown”. Assuming the Apache‑2.0 license applies, the following holds:

• Commercial use: Permitted without fee. You may embed the model in proprietary search products or SaaS offerings.
• Modification & redistribution: You may modify the model weights or code and redistribute them, provided you retain the original copyright notice and include a copy of the Apache‑2.0 license.
• Patents: The license grants a patent‑grant for contributions made by the licensor.
• Attribution: Required. Cite the model name and the original paper (see section 6) in any public documentation or academic work.
• Trademark: The “OpenSearch” brand is owned by Amazon; you must not imply endorsement unless you have explicit permission.