How to train a Matryoshka embedding model with Sentence Transformers

Dense retrieval systems often trade search quality against vector storage and comparison cost. Matryoshka training keeps several usable prefix lengths inside one embedding model, allowing the same saved model to serve indexes with different vector dimensions.

The MatryoshkaLoss wrapper applies a base embedding loss at every configured prefix, including the model's native output dimension. Truncation reduces downstream vector storage and comparison work; it does not make model training or embedding generation faster.

A compact support-ticket dataset is enough to exercise the trainer and reload path locally. A production model still needs representative training data and held-out retrieval evaluation at every dimension that will be deployed.

1. Install the Sentence Transformers training dependencies in the active Python environment.

   $ python -m pip install --upgrade "sentence-transformers[train]"

   A virtual environment keeps the trainer stack separate from system Python packages.

2. Create train_matryoshka.py with the support-ticket pairs and trainer imports.

   train_matryoshka.py

   from pathlib import Path
    
   from datasets import Dataset
   from sentence_transformers import (
       SentenceTransformer,
       SentenceTransformerTrainer,
       SentenceTransformerTrainingArguments,
   )
   from sentence_transformers.sentence_transformer.losses import CoSENTLoss, MatryoshkaLoss
    
    
   output_dir = Path("models/support-matryoshka")
   pairs = [
       ("reset a forgotten admin password", "help an administrator regain account access", 0.95),
       ("restore a deleted customer record", "recover a customer profile that was removed", 0.94),
       ("create a read-only database user", "add a database user that can only read", 0.92),
       ("rotate an expired API token", "replace an API token that has expired", 0.91),
       ("download the latest invoice", "retrieve the newest billing invoice", 0.90),
       ("configure daily database backups", "schedule a database backup every day", 0.89),
       ("invite a new support agent", "add another agent to the support team", 0.88),
       ("archive a completed support case", "close and archive a resolved support case", 0.86),
       ("reset a forgotten admin password", "download the latest invoice", 0.08),
       ("restore a deleted customer record", "invite a new support agent", 0.07),
       ("create a read-only database user", "archive a completed support case", 0.06),
       ("rotate an expired API token", "configure daily database backups", 0.05),
   ]
   train_dataset = Dataset.from_dict(
       {
           "sentence1": [pair[0] for pair in pairs],
           "sentence2": [pair[1] for pair in pairs],
           "score": [pair[2] for pair in pairs],
       }
   )

3. Append the base model and Matryoshka loss definition to train_matryoshka.py.

   model = SentenceTransformer(
       "sentence-transformers/paraphrase-MiniLM-L3-v2",
       model_kwargs={"torch_dtype": "float32"},
   )
    
   base_loss = CoSENTLoss(model=model)
   loss = MatryoshkaLoss(
       model=model,
       loss=base_loss,
       matryoshka_dims=[384, 128, 64, 32],
   )

   The loss list starts with the model's native 384-dimension output and contains each smaller prefix intended for downstream indexes.

4. Append the trainer configuration and model-saving call to train_matryoshka.py.

   args = SentenceTransformerTrainingArguments(
       output_dir=str(output_dir),
       num_train_epochs=1,
       per_device_train_batch_size=4,
       learning_rate=2e-5,
       warmup_steps=0,
       save_strategy="no",
       logging_strategy="no",
       report_to="none",
       seed=42,
   )
    
   trainer = SentenceTransformerTrainer(
       model=model,
       args=args,
       train_dataset=train_dataset,
       loss=loss,
   )
   trainer.train()
   model.save_pretrained(output_dir / "final")
   print(f"saved model: {output_dir / 'final'}")

5. Run train_matryoshka.py to train the Matryoshka model.

   $ python train_matryoshka.py
   ##### snipped #####
   {'train_runtime': '37.64', 'train_samples_per_second': '0.319', 'train_steps_per_second': '0.08', 'train_loss': '21.18', 'epoch': '1'}
   saved model: models/support-matryoshka/final

   Training time and loss vary with the processor, model, dataset, and package versions. The saved model path is the required state for the independent reload check.

6. Create verify_matryoshka.py for an independent ranking check at every trained dimension.

   verify_matryoshka.py

   from sentence_transformers import SentenceTransformer
    
    
   model = SentenceTransformer("models/support-matryoshka/final")
   texts = [
       "reset admin login",
       "help an administrator regain account access",
       "retrieve the newest billing invoice",
   ]
    
   for dimension in [384, 128, 64, 32]:
       embeddings = model.encode(
           texts,
           normalize_embeddings=True,
           truncate_dim=dimension,
       )
       scores = embeddings[1:] @ embeddings[0]
       assert embeddings.shape == (3, dimension)
       assert scores[0] > scores[1]
       print(
           f"{dimension:>3} dimensions: "
           f"related={scores[0]:.4f} unrelated={scores[1]:.4f}"
       )

7. Run verify_matryoshka.py to confirm the reloaded model preserves the expected ranking at each dimension.

   $ python verify_matryoshka.py
   384 dimensions: related=0.5981 unrelated=0.1180
   128 dimensions: related=0.6312 unrelated=0.0028
    64 dimensions: related=0.7110 unrelated=-0.0304
    32 dimensions: related=0.5644 unrelated=-0.0082

   Stored document vectors and query vectors must use the same width, so an index dimension change requires a rebuilt document index.