How to train a sparse encoder with Sentence Transformers

Sparse retrieval preserves exact vocabulary matches while learning expansion terms that connect a query with relevant wording absent from the query itself. A Sentence Transformers SparseEncoder can adapt that behavior to the language in a search corpus instead of relying only on a general-purpose checkpoint.

The current SPLADE training path combines paired query and document text with SparseMultipleNegativesRankingLoss. SpladeLoss wraps that ranking objective with separate query and document regularizers, which control how many vocabulary dimensions remain active in the learned representations.

The sample uses a tiny public checkpoint, eight support-search pairs, and two CPU training steps so the full save-and-reload path can be tested quickly. Production training needs a representative dataset, an evaluator, tuned regularizer weights, and enough steps to compare retrieval quality with sparsity before the model is used for indexing.

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

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

   The training extra installs datasets, accelerate, and the current trainer dependencies.
   Related: How to install Sentence Transformers with pip

2. Create train_sparse_encoder.py with the imports and output paths.

   train_sparse_encoder.py

   import os
   from pathlib import Path
    
   import torch
   from datasets import Dataset
   from sentence_transformers import (
       SparseEncoder,
       SparseEncoderTrainer,
       SparseEncoderTrainingArguments,
   )
   from sentence_transformers.sparse_encoder.losses import (
       SparseMultipleNegativesRankingLoss,
       SpladeLoss,
   )
    
   os.environ["TOKENIZERS_PARALLELISM"] = "false"
   os.environ["WANDB_DISABLED"] = "true"
    
   output_dir = Path("training-output/support-sparse-encoder")
   model_dir = Path("models/support-sparse-encoder")

   Distinct output paths prevent a new experiment from mixing trainer state or model files with an earlier run.

3. Append the paired support-search dataset to train_sparse_encoder.py.

   train_sparse_encoder.py

   train_dataset = Dataset.from_dict(
       {
           "anchor": [
               "reset a locked account",
               "reset a user password",
               "enable two-factor authentication",
               "rotate an API token",
               "restore a deleted project",
               "invite a new support agent",
               "export audit logs",
               "change the billing contact",
           ],
           "positive": [
               "Unlock the account from the admin users page.",
               "Open the user profile and send a password reset email.",
               "Open account security and enroll an authenticator app.",
               "Revoke the old API token and create a replacement token.",
               "Open deleted projects and restore the selected project.",
               "Open team settings and send an invitation email.",
               "Open compliance reports and export the audit log CSV.",
               "Open billing settings and update the primary contact.",
           ],
       }
   )

   Each anchor and positive value at the same row index forms a relevant query-document pair. The published pairs are smoke-test data rather than production training data.

4. Append the sparse model and regularized ranking loss to train_sparse_encoder.py.

   train_sparse_encoder.py

   model = SparseEncoder("sparse-encoder-testing/splade-bert-tiny-nq")
   ranking_loss = SparseMultipleNegativesRankingLoss(model=model)
   loss = SpladeLoss(
       model=model,
       loss=ranking_loss,
       document_regularizer_weight=3e-5,
       query_regularizer_weight=5e-5,
   )

   The tiny checkpoint keeps this CPU smoke test short. Production regularizer weights require retrieval and sparsity evaluation because stronger regularization can reduce index cost while also removing useful expansion terms.

5. Append the training arguments and trainer to train_sparse_encoder.py.

   train_sparse_encoder.py

   args = SparseEncoderTrainingArguments(
       output_dir=str(output_dir),
       max_steps=2,
       per_device_train_batch_size=2,
       learning_rate=2e-5,
       warmup_steps=0,
       use_cpu=True,
       save_strategy="no",
       eval_strategy="no",
       logging_strategy="no",
       disable_tqdm=True,
       report_to=[],
       seed=7,
   )
    
   trainer = SparseEncoderTrainer(
       model=model,
       args=args,
       train_dataset=train_dataset,
       loss=loss,
   )

   Larger batches provide more in-batch negatives to the ranking loss. Two steps exercise the API path but do not replace an evaluated epoch or step schedule.

6. Append the training, persistence, and retrieval smoke-test section to train_sparse_encoder.py.

   train_sparse_encoder.py

   train_result = trainer.train()
   model.save_pretrained(model_dir)
    
   trained_model = SparseEncoder(str(model_dir))
   query_embedding = trained_model.encode_query(
       ["How do I replace an expired API token?"],
       convert_to_tensor=True,
   ).coalesce()
   document_embeddings = trained_model.encode_document(
       [
           "Revoke the old API token and create a replacement token.",
           "Open billing settings and update the saved credit card.",
           "Open team settings and send an invitation email.",
       ],
       convert_to_tensor=True,
   ).coalesce()
    
   scores = torch.mm(
       query_embedding.to_dense(),
       document_embeddings.to_dense().T,
   )
   best_index = int(scores.argmax().item())
    
   print(f"Training rows: {len(train_dataset)}")
   print(f"Loss wrapper: {loss.__class__.__name__}")
   print(f"Training steps: {train_result.global_step}")
   print(f"Saved model: {model_dir}")
   print(f"Query active dimensions: {query_embedding._nnz()}")
   print(f"Document active dimensions: {document_embeddings._nnz()}")
   print(f"Top smoke-test document: d{best_index + 1}")
   print(f"Top smoke-test score: {scores[0, best_index]:.4f}")
    
   if train_result.global_step != 2:
       raise SystemExit(f"unexpected training steps: {train_result.global_step}")
   if query_embedding._nnz() == 0 or document_embeddings._nnz() == 0:
       raise SystemExit("sparse encoder returned no active dimensions")
   if not torch.isfinite(scores).all():
       raise SystemExit("non-finite sparse similarity score returned")
   if best_index != 0:
       raise SystemExit("the API token query did not rank the matching document first")

   Reloading from models/support-sparse-encoder tests the saved artifact rather than the in-memory training object. The assertions fail when training does not finish, either embedding has no active dimensions, a score is non-finite, or the matching document is not ranked first.

7. Run the completed sparse encoder training script.

   $ python train_sparse_encoder.py
   ##### snipped #####
   Training rows: 8
   Loss wrapper: SpladeLoss
   Training steps: 2
   Saved model: models/support-sparse-encoder
   Query active dimensions: 121
   Document active dimensions: 519
   Top smoke-test document: d1
   Top smoke-test score: 137.5335

   The trainer prints runtime metrics before the summarized lines. A d1 result means the reloaded sparse encoder ranked the API-token document above the billing and team documents.

8. Delete the sample script without removing the saved sparse encoder.

   $ rm train_sparse_encoder.py

   The models/support-sparse-encoder directory remains available for evaluation or indexing.
   Related: How to evaluate a sparse encoder with Sentence Transformers
   Related: How to build sparse semantic search with Sentence Transformers

9. Reload the saved sparse encoder after cleanup to confirm that it still ranks the matching document first.

   $ python - <<'PY'
   import torch
   from sentence_transformers import SparseEncoder
   
   model = SparseEncoder("models/support-sparse-encoder")
   query = model.encode_query(
       ["How do I replace an expired API token?"],
       convert_to_tensor=True,
   ).coalesce()
   documents = model.encode_document(
       [
           "Revoke the old API token and create a replacement token.",
           "Open billing settings and update the saved credit card.",
           "Open team settings and send an invitation email.",
       ],
       convert_to_tensor=True,
   ).coalesce()
   scores = torch.mm(query.to_dense(), documents.to_dense().T)
   top_document = int(scores.argmax().item()) + 1
   
   if top_document != 1:
       raise SystemExit(f"unexpected top document: d{top_document}")
   
   print(f"Reloaded top document: d{top_document}")
   PY
   ##### snipped #####
   Reloaded top document: d1

   Loading from models/support-sparse-encoder fails if the retained model is missing or unreadable, and a result other than d1 exits with an error.