RikkaBotan/stable-static-embedding-fast-retrieval-mrl-en-v2

[image]

If you would like to know more details:

[SSE Technical Article](https://huggingface.co/blog/RikkaBotan/stable-static-embedding-technical-report)

[SSE v2 Technical Article](https://huggingface.co/blog/RikkaBotan/stable-static-embedding-v2-technical-report)

[image]

(a) Retrieval performance (nDCG@10) across NanoBEIR English tasks. (b) Mean nDCG@10 vs. inference speed (QPS: queries per second) measured on TREC-COVID and Quora using an Intel® Core™ Ultra 7 265K (3.90 GHz) with batch size 32.

🩵 SSE: Stable Static Embedding for Retrieval MRL v2 🩵

A lightweight, faster and powerful embedding model

Performance Snapshot The released SSE checkpoint reaches NanoBEIR mean nDCG@10 = 0.5159 at 512 dimensions. The prior static reference, static-retrieval-mrl-en-v1, scores 0.5030 at its native 1024-dimensional width. SSE uses a 15.6M-parameter checkpoint versus 31.3M for the reference and stores half as many values per full-width embedding.

SSE v2 keypoint

The released checkpoint is evaluated as a complete training recipe combining contrastive ranking, Matryoshka prefix supervision, a 15-source data mixture, large-batch one-epoch optimization, and a lightweight post-pooling transform.

At 256 dimensions, the Matryoshka prefix reaches a NanoBEIR mean nDCG@10 of 0.5036, matching the prior static reference's native 1024-dimensional score of 0.5030.

This provides the same retrieval quality with one quarter of the vector storage and one quarter of the coordinate operations for dense similarity scoring.

• —Matryoshka truncation

[image]

• —PCA Analysis

[image]

🩵 Why Choose SSE Retrieval MRL? 🩵

✅ Highest NanoBEIR nDCG@10 among the measured throughput-comparable static models ✅ 15.6M parameters — approximately half the 31.3M parameters of the prior static reference ✅ 512D native output — half the reference width; the 256D prefix matches the reference's native 1024D score ✅ Matryoshka-ready — smoothly truncate to 256D/128D/64D/32D with graceful degradation ✅ Apache 2.0 licensed — free for commercial & personal use ✅ CPU-optimized — runs faster on edge devices & modest hardware

🩵 Model Details 🩵

SentenceTransformer(
  (0): SSE(
    (embedding): EmbeddingBag(30522, 512, mode='mean')
    (dyt): SeparableDyT()
  )
)

[image]

🩵 Mathematical formulations 🩵

Dynamic Tanh Normalization (DyT) enables magnitude-adaptive gradient flow for static embeddings. For input dimension x, DyT computes $$ yk = ck \tanh(ak xk + b_k) $$ with learnable parameters. The gradient of x is:

$$ \frac{\partial yk}{\partial xk} = ck ak \, \mathrm{sech}^2(ak xk + b_k). $$

For saturated dimensions |x| > 1 $$ |ai xi + bi| \gg 1 $$ yields exponential decay $$ \mathrm{sech}^2(z) \sim 4e^{-2|z|} $$ suppressing gradients as $$ \partial yi / \partial xi \to 0 $$ For non-saturated dimensions |x| << 1 , $$ \mathrm{sech}^2(z) \approx 1 $$ preserves near-constant gradients $$ \partial yj / \partial xj \approx cj a_j $$ This magnitude-dependent gating attenuates learning signals from noisy, large-magnitude dimensions while maintaining full gradient flow for stable, informative dimensions—providing implicit regularization that enhances generalization without explicit hyperparameters.

🩵 Evaluation Results (NanoBEIR) 🩵

The values below were re-measured with the same NanoBEIR evaluator used in the paper.

Top performance on community-based retrieval (Quora) and scientific fact verification!

🩵 How to use? 🩵

import torch
from sentence_transformers import SentenceTransformer

# load (remote code enabled)
model = SentenceTransformer(
    "RikkaBotan/stable-static-embedding-fast-retrieval-mrl-en-v2",
    trust_remote_code=True,
    device="cuda" if torch.cuda.is_available() else "cpu",
)

# inference
sentences = [
    "Stable Static embedding is interesting.",
    "SSE works without attention."
]

with torch.no_grad():
    embeddings = model.encode(
        sentences,
        convert_to_tensor=True,
        normalize_embeddings=True,
        batch_size=32
    )

# cosine similarity
# cosine_sim = embeddings[0] @ embeddings[1].T
cosine_sim = model.similarity(embeddings, embeddings)

print("embeddings shape:", embeddings.shape)
print("cosine similarity matrix:")
print(cosine_sim)

🩵 Retrieval usage 🩵

import torch
from sentence_transformers import SentenceTransformer

# load (remote code enabled)
model = SentenceTransformer(
    "RikkaBotan/stable-static-embedding-fast-retrieval-mrl-en-v2",
    trust_remote_code=True,
    device="cuda" if torch.cuda.is_available() else "cpu",
)

# inference
query = "What is Stable Static Embedding?"
sentences = [
    "SSE: Stable Static embedding works without attention.",
    "Stable Static Embedding is a fast embedding method designed for retrieval tasks.",
    "Static embeddings are often compared with transformer-based sentence encoders.",
    "I cooked pasta last night while listening to jazz music.",
    "Large language models are commonly trained using next-token prediction objectives.",
    "Instruction tuning improves the ability of LLMs to follow human-written prompts.",
]


with torch.no_grad():
    embeddings = model.encode(
        [query] + sentences,
        convert_to_tensor=True,
        normalize_embeddings=True,
        batch_size=32
    )

print("embeddings shape:", embeddings.shape)

# cosine similarity
similarities = model.similarity(embeddings[0], embeddings[1:])
for i, similarity in enumerate(similarities[0].tolist()):
    print(f"{similarity:.05f}: {sentences[i]}")

🩵 Training Hyperparameters 🩵

Non-Default Hyperparameters

• —eval_strategy: steps
• —per_device_train_batch_size: 2048
• —gradient_accumulation_steps: 4
• —learning_rate: 0.1
• —adam_beta2: 0.9999
• —adam_epsilon: 1e-10
• —num_train_epochs: 1
• —lr_scheduler_type: cosine
• —warmup_ratio: 0.1
• —bf16: True
• —dataloader_num_workers: 8
• —batch_sampler: no_duplicates

🩵 Training Datasets 🩵

We learned from 15 datasets:

All trained with MatryoshkaLoss**

🩵 Training results 🩵

• —train loss

[image]

• —NanoBEIR mean NDCG@10

[image]

🩵 About me 🩵

Japanese independent researcher having shy and pampered personality. Twin-tail hair is a charm point. Interested in nlp. Usually using python and C.

X(Twitter): https://twitter.com/peony__snow

[image]

🩵 Acknowledgements 🩵

The author acknowledge the support of Saldra, Witness and Lumina Logic Minds for providing computational resources used in this work.

I thank the developers of sentence-transformers, python and pytorch.

I thank all the researchers for their efforts to date.

I thank Japan's high standard of education.

And most of all, thank you for your interest in this repository.

🩵 Citation 🩵

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

MatryoshkaLoss

@misc{kusupati2024matryoshka,
    title={Matryoshka Representation Learning},
    author={Aditya Kusupati and Gantavya Bhatt and Aniket Rege and Matthew Wallingford and Aditya Sinha and Vivek Ramanujan and William Howard-Snyder and Kaifeng Chen and Sham Kakade and Prateek Jain and Ali Farhadi},
    year={2024},
    eprint={2205.13147},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

MultipleNegativesRankingLoss

@misc{henderson2017efficient,
    title={Efficient Natural Language Response Suggestion for Smart Reply},
    author={Matthew Henderson and Rami Al-Rfou and Brian Strope and Yun-hsuan Sung and Laszlo Lukacs and Ruiqi Guo and Sanjiv Kumar and Balint Miklos and Ray Kurzweil},
    year={2017},
    eprint={1705.00652},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}