Train and fine-tune sentence-transformer embedding models

Train Sentence Transformers helps adapt three model families: SentenceTransformer bi-encoders for reusable dense embeddings, CrossEncoder models that jointly score query-passage pairs, and SparseEncoder models for learned sparse retrieval. Dense embeddings support similarity, clustering, classification, paraphrase mining and deduplication. Cross-encoders rerank retrieved candidates; SPLADE produces vocabulary-weighted vectors for inverted indexes such as Elasticsearch, OpenSearch and Lucene.

The workflow selects a model family, aligns the dataset shape with its loss, establishes a baseline, trains within an approved budget and compares the selected checkpoint against held-out evaluation. It includes separate training templates and references for architecture, loss functions, evaluators, data preparation, memory management and remote execution. Variant recipes cover Matryoshka dimensions, adaptive layers, multi-dataset training, LoRA, distillation, multilingual adaptation, static embeddings and hard-negative mining. Referenced companion recipes must be available before using their specialized paths.

Use it for domain adaptation when a general retriever or reranker misses task-specific relevance, language or terminology. Choose dense retrieval when documents need independently reusable embeddings, joint scoring when a candidate pool already exists, and sparse retrieval when lexical index integration is a requirement. Compare suitable pretrained models before paying for a fresh training run. Discovery and leaderboard results change; model suitability, language coverage and licenses need individual review.

A training plan does not authorize computation, downloads, uploads, telemetry, paid hardware or publication. Establish data rights, privacy boundaries, approved hardware, time and cost ceilings first. Treat dataset content, model cards and downloaded artifacts as untrusted input. Review dependencies and pinned revisions; avoid arbitrary remote code and unsafe serialized inputs. Numerical examples for quality, memory, speed and dataset volume are illustrative planning hypotheses, not guaranteed outcomes or universal minimums.

Inputs include the task, approved base and teacher models, reviewed datasets, a lockfile, a loss-compatible column layout, evaluation splits and an execution budget. Recognized label columns are label, labels, score or scores; the relative order of remaining input columns determines their role. Remove unrelated columns unless a custom collator explicitly needs them. Convert string labels to the intended numeric type. Cosine regression needs labels on the intended cosine scale, while ranking objectives do not require that same absolute normalization.

Training templates require their sibling training_safety.py helper. Local execution requires ST_ALLOW_LOCAL_TRAINING=1, ST_MAX_STEPS from 1 to 1000, an existing symlink-free ST_RUN_PARENT and ST_REVISIONS_JSON mapping every requested model and dataset to an audited full 40-character commit. Entry points force Hub and dataset offline mode before SDK imports, create a private run directory, disable remote reporting and save locally. Reviewed caches must be provisioned separately; absent revisions or cached assets must fail rather than fall back online. These templates are single-process. Distributed execution requires a separately reviewed adaptation with rank-safe persistence.

The ST_MAX_ROWS default is 200, capped at 20,000 per explicit streaming dataset load. That limit does not bound bytes, cache size, memory, evaluation or elapsed time. NanoBEIR can load its own cached data and perform substantial work. A genuinely bounded smoke test needs tiny prepared data and a tiny evaluator; max_steps=1 alone does not bound preprocessing. Capped identifier-based datasets may lack referenced documents, so prepare self-contained text instead of inventing matches or silently enlarging the load. Offline flags are SDK conventions rather than a network sandbox; strict isolation needs process or operating-system controls.

Outputs include local checkpoints, final saved models, logs, baseline and post-training metrics, and an experiment record. Evaluate before training; split by query, document, entity or time as appropriate, grouping related query rows before expanding labeled pairs. Select checkpoints on development data and evaluate once on an untouched test set. Inspect actual emitted metric names and finite values. SMOKE_ONLY is pipeline evidence, not quality acceptance; a fixed 0.005 gain is descriptive, not statistical significance. Missing sparse activity metrics yield SPARSITY_UNVERIFIED. Approve sparse gains against task-specific latency, index-size and sparsity budgets as well as ranking quality.

For dense models, MultipleNegativesRankingLoss uses positive pairs or tuples with negatives and needs duplicate-free batches. Cached variants enlarge the contrastive pool through GradCache; gradient accumulation alone does not create equivalent in-batch negatives. GIST variants use a guide to filter false negatives. Explicit triplet objectives need trustworthy negatives, while batch-triplet objectives need multiple examples per class through GROUP_BY_LABEL. CoSENT, AnglE and cosine regression serve labeled similarity data; contrastive and softmax classification losses serve different label structures. Distillation can match teacher embeddings, score differences or candidate distributions. Denoising and ContrastiveTension support unsupervised paths.

Matryoshka objectives train multiple usable embedding dimensions; adaptive-layer variants support earlier exits. Matryoshka2d already combines dimension and layer adaptation, and GOR is a separate regularizer rather than a wrapper. Encoder models commonly use mean pooling, causal decoders need last-token pooling, and static embeddings trade contextual modeling for token lookup. Their float32 lookup-table size is vocabulary size multiplied by dimension and four bytes, so vocabulary choice materially affects the footprint. Multimodal backbones need the relevant preprocessing extras. Router models start with unaligned modality spaces, require alignment training and need projection when dimensions differ. Long-context defaults can admit unexpectedly large batches, so set a deliberate sequence limit.

For rerankers, binary cross-entropy uses one output and positive/negative labels, with class weighting derived from actual data. Multiclass cross-entropy needs multiple outputs. Other regression, distillation and ranking objectives require the documented activation setup, including Identity where a sigmoid would distort scores. LambdaLoss, ListNet, ListMLE, PListMLE, RankNet and related listwise paths need correctly shaped candidate arrays. Large candidate lists can create quadratic memory pressure: reduce loss mini-batches, then device batches, and finally candidate count with an explicit experiment change. Tiny normalized losses do not establish ranking quality. Custom heads must emit the required scores and remain importable when reloaded.

SPLADE requires a masked-language-model-compatible backbone and SpladePooling. Wrap its retrieval loss with SpladeLoss so FLOPS regularization actually encourages sparse outputs. The scheduler gradually introduces regularization; full strength at initialization can collapse outputs to zero. Query and document regularization weights usually differ, with stronger query regularization for sparser queries. CachedSpladeLoss is an alternative memory strategy. Monitor active query and document dimensions alongside dot-product retrieval metrics; dense outputs defeat the sparse deployment goal. CSR architectures require a different setup.

Evaluators cover retrieval, similarity correlation, binary classification, triplets, reranking, distillation and translation. Full-corpus retrieval is better reserved for milestones when expensive. SequentialEvaluator combines checks, while ReciprocalRankFusionEvaluator measures a deployed hybrid retriever. Cross-encoder evaluation must distinguish reranking an artificially positive-containing pool from end-to-end retrieval where the positive may be missing. Inspect primary_metric after baseline evaluation rather than guessing a version-independent key.

Training arguments coordinate evaluation and checkpoint saving for best-model selection. Early stopping is especially useful for rerankers that peak before their final step. Keep weights and optimizer state in full precision and use supported bf16 or fp16 autocast; do not enable both. Evaluations outside the trainer need their own autocast, particularly with FlashAttention 2. Cached losses conflict with gradient checkpointing. Multi-GPU contrastive negatives remain local unless supported cross-device gathering is enabled. FSDP has evaluator and layer-wrapping caveats; validate the exact stack before distributing a working single-process setup.

Hard-negative mining supports triplets, tuples, labeled pairs and candidate lists. Filter likely false negatives and preserve the intended training shape. Resume logic needs special care for streaming dataset order. Hyperparameter search belongs after a working manually tuned run, and random seeds do not guarantee identical results across hardware or software versions.

The example scripts target Python 3.10 or later with reviewed compatible training dependencies. One exact installation example is:

pip install "sentence-transformers[train]>=5.0"

Hugging Face Jobs requires separate account, region, data, dependency, hardware, maximum-cost, timeout and artifact-destination approval. A local script does not automatically mount local data or include its sibling helper. Pin the complete bundle, use least-privilege secret injection, authorize any remote tracking, inspect the returned job ID and verify exit status and artifact checksums. A timeout does not prove output persistence, and successful training alone does not prove publication, validation or final cost.

The code license is Apache-2.0; base models, teachers and datasets have separate terms.