Train 400x faster Static Embedding Models with Sentence Transformers

TL;DR

Hugging Face released a training recipe for static embedding models that run 100x–400x faster on CPU while retaining at least 85% of the quality of standard transformer‑based embeddings, and published two concrete models (static‑retrieval‑mrl‑en‑v1 for English retrieval and static‑similarity‑mrl‑multilingual‑v1 for multilingual similarity) together with training scripts, evaluation results, and Weights & Biases logs.

Method

The approach combines contrastive learning with MultipleNegativesRankingLoss and optional Matryoshka Representation Learning to train static encoders that perform a simple token‑embedding lookup instead of attention‑based encoding, yielding speedups of several orders of magnitude with minimal quality loss.

Training Details

Requirements

Training uses the Sentence Transformers library with components: dataset, loss function, training arguments, evaluator, and trainer.

Model Inspiration

Two models were targeted: an English‑only retrieval model and a multilingual general similarity model, both initialized with a StaticEmbedding module wrapping a BERT‑based tokenizer (bert‑base‑uncased or bert‑base‑multilingual‑uncased) and an embedding dimension of 1024.

Dataset Selection

For English retrieval, 30 datasets were selected, including gooaq, msmarco (triplet), squad, s2orc (title‑abstract‑pair), allnli (triplet), paq, trivia_qa, msmarco_10m, swim_ir (en), pubmedqa (triplet‑20), miracl (en‑triplet‑all), mldr (en‑triplet‑all), and mr_tydi (en‑triplet‑all). For multilingual similarity, datasets with parallel sentences (wikititles, tatoeba, talks, europarl, global_voices, jw300, muse, wikimatrix, opensubtitles) and positive‑pair datasets (stackexchange‑duplicates, quora‑duplicates, wikianswers‑duplicates, allnli, simple_wiki, altlex, flickr30k_captions, coco_captions, nli_for_simcse, negation) were combined, totaling 30 training datasets.

Loss Function

MultipleNegativesRankingLoss was chosen because the batch size of 2048 fits on an RTX 3090, avoiding the overhead of CachedMultipleNegativesRankingLoss and the slowdown of a guide model in GISTEmbedLoss. MatryoshkaLoss was applied on top, with dimensions [32, 64, 128, 256, 512, 1024].

Training Arguments

Both models used: num_train_epochs=1, per_device_train_batch_size=2048, per_device_eval_batch_size=2048, learning_rate=2e-1, warmup_ratio=0.1, bf16=True, batch_sampler=BatchSamplers.NO_DUPLICATES, multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL, eval_strategy=steps, eval_steps=250 (English) or 1000 (multilingual), save_strategy=steps, save_steps matching eval_steps, save_total_limit=2, logging_steps matching eval_steps, logging_first_step=True, and a run‑specific output directory.

Evaluator

The English retrieval model used NanoBEIREvaluator for zero‑shot retrieval assessment; the multilingual model relied on MTEB tasks (STS, Classification, Pair Classification) for evaluation.

Hardware

Training was performed on an RTX 3090 GPU, an i7‑13700K CPU, and 32 GB RAM.

Overall Training Scripts

The provided scripts load the datasets, instantiate the StaticEmbedding model, compose the loss with MatryoshkaLoss, set the training arguments, optionally run an evaluator, train with SentenceTransformerTrainer, and save the final model. The English retrieval script took 17.8 hours, consuming 2.6 kWh and emitting 1 kg CO₂; the multilingual script took 3.1 hours, consuming 0.5 kWh and emitting 0.2 kg CO₂.

Usage

Both models are loaded via SentenceTransformer with the model name and device="cpu). Inference is identical to standard Sentence Transformers: model.encode returns embeddings, and model.similarity computes cosine similarity. The truncate_dim argument enables Matryoshka‑style dimensionality reduction (e.g., truncate_dim=256). The models work out‑of‑the‑box with LangChain, LlamaIndex, Haystack, and txtai.

Performance

English Retrieval

On NanoBEIR, static‑retrieval‑mrl‑en‑v1 achieves an NDCG@10 of 0.5032, which is 87.4 % of the score of all‑mpnet‑base‑v2 (0.5757). On CPU it processes 107 419.51 sentences per second, 397× faster than all‑mpnet‑base‑v2 (270.40 sentences/s). On GPU it processes 97 171.47 sentences per second, 24× faster than all‑mpnet‑base‑v2 (4043.13 sentences/s). Matryoshka evaluation shows that halving the dimension to 512 reduces NDCG@10 by only 1.47 % (0.5032 → 0.4957).

Multilingual Similarity

Relative to multilingual‑e5‑small, static‑similarity‑mrl‑multilingual‑v1 scores 92.3 % on STS, 95.52 % on Pair Classification, and 86.52 % on Classification. It is approximately 125× faster on CPU and 10× faster on GPU than multilingual‑e5‑small. Matryoshka evaluation indicates that reducing the dimension to 256 (4× smaller) incurs only a 0.56 % drop in English STS performance.

Conclusion

Static embedding models trained with the presented recipe deliver 100×–400× CPU speedups and 10×–25× GPU speedups while preserving at least 85 % of the quality of common transformer‑based embeddings. The released models enable efficient on‑device, in‑browser, and edge‑computing use cases with minimal accuracy loss.

Next Steps

Users can replace existing Sentence Transformer models with static‑retrieval‑mrl‑en‑v1 or static‑similarity‑mrl‑multilingual‑v1, or train their own static embeddings on task‑specific data. Potential improvements include hard‑negative mining, model souping, curriculum learning, guided false‑in‑batch negatives filtering, seed‑optimized random initialization, tokenizer retraining, gradient caching via CachedMultipleNegativesRankingLoss, and model distillation from larger encoders.