opensearch-neural-sparse-encoding-asym-v1

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.

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.

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 two towers live in subfolders — query in query_0_MLMTransformer/, document in document_0_MLMTransformer/ 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(f"{local}/document_0_MLMTransformer").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)

License

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

Copyright

Copyright OpenSearch Contributors. See NOTICE for details.