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

If you would like to know more details:

SSE Technical Article

SSE v2 Technical Article

(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

• PCA Analysis

---

🩵 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()
  )
)

---

🩵 Mathematical formulations 🩵

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

$$ \frac{\partial y_k}{\partial x_k} = c_k a_k , \mathrm{sech}^2(a_k x_k + b_k). $$

For saturated dimensions |x| > 1 $$ |a_i x_i + b_i| \gg 1 $$ yields exponential decay $$ \mathrm{sech}^2(z) \sim 4e^{-2|z|} $$ suppressing gradients as $$ \partial y_i / \partial x_i \to 0 $$ For non-saturated dimensions |x| << 1 , $$ \mathrm{sech}^2(z) \approx 1 $$ preserves near-constant gradients $$ \partial y_j / \partial x_j \approx c_j 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

• NanoBEIR mean NDCG@10

🩵 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

🩵 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}
}