opensearch-neural-sparse-encoding-asym-v1

opensearch-neural-sparse-encoding-asym-v1: search relevance, retrieval FLOPS, and query inference time

Select the model

OpenSearch neural sparse encoding models produce sparse term vectors over the BERT WordPiece vocabulary and are served with an inverted index. asym-v1 moves neural sparse search from inference-free to inference-light: it is an asymmetric (two-tower) model that pairs a small query encoder (bert-medium, 41 M) with a large document encoder (GTE, 137 M), which runs offline at index time. The query encoder keeps the full semantics of the query at an affordable cost: it encodes a query in less than 30 ms on a single CPU thread. asym-v1 achieves the strongest BEIR-13 score among OpenSearch sparse models and ranks #2 among sparse models under 256M parameters on the MTEB BEIR benchmark (53.48 average NDCG@10 over 15 tasks), with low retrieval FLOPS and query processing cost.

Model Inference-free for Retrieval Model Parameters AVG NDCG@10 AVG FLOPS
opensearch-neural-sparse-encoding-v1 137M 0.524 11.4
opensearch-neural-sparse-encoding-v2-distill 67M 0.528 8.3
opensearch-neural-sparse-encoding-asym-v1 41M (query) + 137M (doc) 0.555 4.5
opensearch-neural-sparse-encoding-doc-v1 βœ”οΈ 137M 0.490 2.3
opensearch-neural-sparse-encoding-doc-v2-distill βœ”οΈ 67M 0.504 1.8
opensearch-neural-sparse-encoding-doc-v2-mini βœ”οΈ 23M 0.497 1.7
opensearch-neural-sparse-encoding-doc-v3-distill βœ”οΈ 67M 0.517 1.8
opensearch-neural-sparse-encoding-doc-v3-gte βœ”οΈ 137M 0.546 1.7

Search relevance and query cost

Comparison with other sparse retrievers. The query encoder is the model that runs at search time; inference-free models only tokenize the query.

Model Query encoder BEIR-13 NDCG@10 ↑ Retrieval FLOPS ↓ Query encoding time, p50 ↓
granite-embedding-30m-sparse 30M 0.506 10.06 26.1 ms
opensearch-neural-sparse-encoding-v2-distill 67M 0.528 5.73 38.9 ms
opensearch-neural-sparse-encoding-doc-v3-gte tokenizer only 0.546 0.92 2.0 ms
linkup-sparseup-embed-v1 149M 0.564 7.02 78.8 ms
opensearch-neural-sparse-encoding-asym-v1 41M 0.555 4.42 26.9 ms

asym-v1 has the lowest retrieval FLOPS among the models that run a query encoder; only inference-free models are lower.

Overview

  • Paradigm: asymmetric learned-sparse (SPLADE) retrieval. Both towers map text to a sparse vector over the 30,522-entry BERT WordPiece vocabulary; relevance = sparse dot product, served by an inverted index.
  • Query encoder: Started from prajjwal1/bert-medium (41.4 M, 512-hidden / 8-layer). It runs efficiently on CPU: on a single thread of an Intel Xeon 8275CL it encodes a query in 26.9 ms (p50), so search-time latency stays low even without a GPU.
  • Document encoder: Warm-started from opensearch-neural-sparse-encoding-doc-v3-gte; runs offline at index time.

Usage (Sentence Transformers)

First install the Sentence Transformers library:

pip install -U "sentence-transformers>=5.0" "transformers>=5.22.1"

It requires sentence-transformers >= 5.0 (sparse Router / MLMTransformer / SpladePooling modules) and transformers >= 5.22.1 for native GTE support. Then you can load this model and run inference.

from sentence_transformers import SparseEncoder

model = SparseEncoder("opensearch-project/opensearch-neural-sparse-encoding-asym-v1")

query_embeddings    = model.encode_query(["What's the weather in New York now?"])
document_embeddings = model.encode_document(["Currently New York is rainy."])

print(model.similarity(query_embeddings, document_embeddings))   # tensor([[18.19]])

# Inspect the sparse expansion terms (token, weight):
print(model.decode(query_embeddings[0], top_k=8))
# [('york', 2.93), ('now', 2.43), ('weather', 2.39), ('nyc', 1.97), ('ny', 1.86), ('current', 1.82), ('sunny', 1.62), ('today', 1.45)]

Usage (HuggingFace Transformers)

The query tower lives in query_0_MLMTransformer/ and the document tower in the repository root. Both load with transformers >= 5.22.1 and share the BERT WordPiece tokenizer.

import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForMaskedLM, AutoTokenizer

local = snapshot_download("opensearch-project/opensearch-neural-sparse-encoding-asym-v1")
tokenizer   = AutoTokenizer.from_pretrained(f"{local}/query_0_MLMTransformer")
query_model = AutoModelForMaskedLM.from_pretrained(f"{local}/query_0_MLMTransformer").eval()      # bert-medium query tower
doc_model   = AutoModelForMaskedLM.from_pretrained(local).eval()                              # GTE document tower

def encode(model, texts, max_length):
    feat = tokenizer(texts, return_tensors="pt", truncation=True, max_length=max_length, padding=True)
    feat.pop("token_type_ids", None)
    with torch.no_grad():
        logits = model(**feat).logits                                   # [batch, seq_len, vocab]
    # SPLADE: max-pool log(1 + relu(logits)) over (masked) tokens -> [batch, vocab] sparse term weights
    return (torch.log1p(torch.relu(logits)) * feat["attention_mask"].unsqueeze(-1)).max(dim=1).values

q_vec = encode(query_model, ["What's the weather in New York now?"], max_length=256)   # query tower (search-time)
d_vec = encode(doc_model,   ["Currently New York is rainy."],        max_length=512)   # document tower (index-time)
print((q_vec * d_vec).sum(dim=-1).item())                                              # ~18.19

Usage (ONNX Runtime, query encoder)

The query encoder is also available as an ONNX model in onnx/query_encoder.onnx, so queries can be encoded on a CPU with ONNX Runtime and no PyTorch. The graph includes SPLADE pooling: it takes input_ids and attention_mask (int64, [batch, sequence]) and returns sparse_embedding (float32, [batch, 30522]), the term weights over the BERT WordPiece vocabulary. The tokenizer is not part of the graph; use query_0_MLMTransformer/tokenizer.json.

pip install onnxruntime tokenizers huggingface_hub numpy
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

repo = "opensearch-project/opensearch-neural-sparse-encoding-asym-v1"
tokenizer = Tokenizer.from_file(hf_hub_download(repo, "query_0_MLMTransformer/tokenizer.json"))
tokenizer.enable_truncation(256)
tokenizer.enable_padding(pad_id=0, pad_token="[PAD]")
session = ort.InferenceSession(hf_hub_download(repo, "onnx/query_encoder.onnx"), providers=["CPUExecutionProvider"])

queries = ["What's the weather in New York now?"]
enc = tokenizer.encode_batch(queries)
(query_embeddings,) = session.run(None, {
    "input_ids": np.array([e.ids for e in enc], dtype=np.int64),
    "attention_mask": np.array([e.attention_mask for e in enc], dtype=np.int64),
})  # [batch, 30522] term weights

top = np.argsort(-query_embeddings[0])[:8]
print([(tokenizer.id_to_token(int(i)), round(float(query_embeddings[0, i]), 2)) for i in top])
# [('york', 2.93), ('now', 2.43), ('weather', 2.39), ('nyc', 1.97), ('ny', 1.86), ('current', 1.82), ('sunny', 1.62), ('today', 1.45)]

Detailed Search Relevance (BEIR-13, NDCG@10)

Model Average Trec Covid NFCorpus NQ HotpotQA FiQA ArguAna Touche DBPedia SCIDOCS FEVER Climate FEVER SciFact Quora
v1 0.524 0.771 0.360 0.553 0.697 0.376 0.508 0.278 0.447 0.164 0.821 0.263 0.723 0.856
v2-distill 0.528 0.775 0.347 0.561 0.685 0.374 0.551 0.278 0.435 0.173 0.849 0.249 0.722 0.863
doc-v1 0.490 0.707 0.352 0.521 0.677 0.344 0.461 0.294 0.412 0.154 0.743 0.202 0.716 0.788
doc-v2-distill 0.504 0.690 0.343 0.528 0.675 0.357 0.496 0.287 0.418 0.166 0.818 0.224 0.715 0.841
doc-v2-mini 0.497 0.709 0.336 0.510 0.666 0.338 0.480 0.285 0.407 0.164 0.812 0.216 0.699 0.837
doc-v3-distill 0.517 0.724 0.345 0.544 0.694 0.356 0.520 0.294 0.424 0.163 0.845 0.239 0.708 0.863
doc-v3-gte 0.546 0.734 0.360 0.582 0.716 0.407 0.520 0.389 0.455 0.167 0.860 0.312 0.725 0.873
asym-v1 (this) 0.555 0.750 0.358 0.586 0.742 0.405 0.621 0.333 0.438 0.179 0.874 0.327 0.741 0.866

License

This project is licensed under the Apache v2.0 License.

Copyright

Copyright OpenSearch Contributors. See NOTICE for details.

Downloads last month
7
Safetensors
Model size
0.1B params
Tensor type
F32
Β·
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support