Training and Finetuning Reranker Models with Sentence Transformers

Hugging Face has detailed a framework for training and finetuning reranker models (Cross Encoders) using the Sentence Transformers library. By training on domain-specific data, developers can create small, efficient rerankers that outperform significantly larger general-purpose models in specialized retrieval tasks.

Reranker Models vs. Embedding Models

Cross Encoder reranker models evaluate the relevance between pairs of texts (such as a query and a document) by processing them together through a shared neural network to produce a single output score. This differs from bi-encoders (embedding models), which embed texts independently into vectors and compute similarity via distance metrics.

While Cross Encoders are more computationally expensive because they must process every possible pair, they generally provide higher accuracy. Consequently, the industry standard for high-performance search is a two-stage "retrieve and rerank" pipeline: using a fast embedding model for initial retrieval and a Cross Encoder to refine the top-k results.

Core Training Components

Training a reranker model involves five primary components integrated through the CrossEncoderTrainer:

1. Dataset

The CrossEncoderTrainer supports datasets.Dataset or datasets.DatasetDict objects. Data can be sourced from the Hugging Face Hub or local files (CSV, JSON, Parquet, Arrow, SQL). To ensure compatibility with the chosen loss function, the dataset must have a column named "label", "labels", "score", or "scores" for labels, and the remaining columns must match the required number of inputs for the loss function.

2. Hard Negative Mining

Model performance depends heavily on the quality of negatives. While "soft negatives" are completely unrelated, "hard negatives" are passages that appear relevant but are not. Sentence Transformers provides the mine_hard_negatives function to identify these challenging examples, which forces the model to be more precise in its distinctions.

3. Loss Functions

Loss functions guide the optimization process based on the available data. While learning-to-rank losses like LambdaLoss or ListNetLoss exist, BinaryCrossEntropyLoss remains a highly effective and simpler option for labeled pairs.

4. Training Arguments

Customization is handled via CrossEncoderTrainingArguments, allowing developers to adjust learning rates, batch sizes, warmup ratios, and precision settings (FP16/BF16). The batch_sampler=BatchSamplers.NO_DUPLICATES setting is specifically recommended for losses using in-batch negatives.

5. Evaluators

To track performance beyond simple loss, Sentence Transformers offers several built-in evaluators:

  • CrossEncoderClassificationEvaluator: For binary or multiclass labels.
  • CrossEncoderCorrelationEvaluator: For similarity scores (e.g., using the STSb dataset).
  • CrossEncoderRerankingEvaluator: For evaluating reranking performance using queries, positives, and negatives.
  • CrossEncoderNanoBEIREvaluator: A lightweight evaluator for English reranking.

Multi-Dataset Training

The CrossEncoderTrainer supports training on multiple datasets simultaneously, even if they have different formats or require different loss functions. This is managed via a dictionary of datasets and an optional dictionary of loss functions. Sampling strategies include:

  • ROUND_ROBIN: Equal sampling from each dataset until one is exhausted.
  • PROPORTIONAL: Sampling proportional to the size of each dataset, ensuring all samples are used.

Performance Evaluation and Results

In a practical application, a reranker based on ModernBERT-base was finetuned on 99k query-answer pairs from the GooAQ dataset. Using BinaryCrossEntropyLoss and hard negative mining, the resulting model (tomaarsen/reranker-ModernBERT-base-gooaq-bce) outperformed 13 common open-source rerankers, including models up to 4x its size.

GooAQ NDCG@10 Results (Top 30 Reranking)

Model Parameters Realistic NDCG@10 Evaluation NDCG@10
Retriever only (No reranking) - 59.12 59.12
BAAI/bge-reranker-large 560M 73.20 77.46
mixedbread-ai/mxbai-rerank-large-v2 1.54B 75.40 80.04
ModernBERT-base-gooaq-bce 150M 77.14 83.51
ModernBERT-large-gooaq-bce 396M 79.42 85.81

Technical Training Tips

  • Prevent Overfitting: Cross Encoders overfit quickly. Use an evaluator with load_best_model_at_end and metric_for_best_model to capture the peak performance model.
  • Balance Negatives: Relying exclusively on hard negatives can degrade performance on easier tasks. Mixing random negatives with hard negatives can mitigate this issue.
  • Efficiency: Finetuning a small reranker on domain-specific data can simultaneously improve search accuracy and reduce inference latency compared to using a massive general-purpose model.

Sources