UltraBERT-150M reranker MS MARCO pilot

This is a Cross Encoder model finetuned from dmis-lab/UltraBERT-150M on the parquet dataset using the sentence-transformers library. It computes scores for pairs of texts, which can be used for text reranking and semantic search.

Model Details

Model Description

• Model Type: Cross Encoder
• Base model: dmis-lab/UltraBERT-150M
• Maximum Sequence Length: 291 tokens
• Number of Output Labels: 1 label
• Supported Modality: Text
• Training Dataset:
  • parquet
• Language: en
• License: apache-2.0

Model Sources

• Documentation: Sentence Transformers Documentation
• Documentation: Cross Encoder Documentation
• Repository: Sentence Transformers on GitHub
• Hugging Face: Cross Encoders on Hugging Face

Full Model Architecture

CrossEncoder(
  (0): Transformer({'transformer_task': 'sequence-classification', 'modality_config': {'text': {'method': 'forward', 'method_output_name': 'logits'}}, 'module_output_name': 'scores', 'architecture': 'UltrabertForSequenceClassification'})
)

Usage

Direct Usage (Sentence Transformers)

First install the Sentence Transformers library:

pip install -U sentence-transformers

Then you can load this model and run inference.

from sentence_transformers import CrossEncoder

# Download from the 🤗 Hub
model = CrossEncoder("tomaarsen/ultrabert-150m-reranker-msmarco524k-3e-5-20261014")
# Get scores for pairs of inputs
pairs = [
    ['how much does david cameron earn', "OK-So definitively David Cameron who is the current British Prime Minister earns £142,500 per year. Recently the Civil Service disclosed the salary's of all of the Cabinet members who earn over £150,000. David Cameron's £142,500 was used as a bench mark for this."],
    ['how much does david cameron earn', 'Celebrities wealth, income, and salary in 2016, 2017. Cameron Dallas. Handsome Vine star Cameron Dallas earned his first million before he was 20 years old. His current net worth is approximately $1.2 million as of 2017 and although he never admitted having any official girlfriend, we know the names of last two girls he was dating.'],
    ['distance from the island of palms to downtown charleston, sc', 'Distance, Gas Consumption and Emission Notes. Distance from Isle of Palms, SC to Charleston, SC is 18Miles or 29 Km. You can get this distance about 39 mins. If you want to planning travel with plane for 13 Miles or 20 Km, You can get this distance about 31 mins . A car with an average MPG will needs 0.83 gallons of gas to get the route between these points. The estimated cost of gas to get between Isle of Palms, SC and Charleston, SC is $1.88. During the route, an average car will release 16.26 pounds of CO 2 to the atmosphere. Your carbon footprint is 0.9 pounds of CO 2 per mile.'],
    ['distance from the island of palms to downtown charleston, sc', 'Once youâ\x80\x99ve done your obligatory sunbathing and beachcombing at the Isle of Palms County Park Beach, head to nearby Charleston, a 30-minute drive away, which has historic sites, landmarks, shopping districts, golf courses and family attractions like The South Carolina Aquarium and Fountain Walk, a new dining, entertainment and shopping complex nearby.'],
    ['how much does nicu cost per day', 'According to a 2010 article in Managed Care Magazine, the average cost for infants in NICU is around $3,000 per day. This average does not include surgeries or helicopter transports.'],
]
scores = model.predict(pairs)
print(scores)
# [ 6.9062 -3.8906  5.9688 -2.3438  4.4375]

# Or rank different texts based on similarity to a single text
ranks = model.rank(
    'how much does david cameron earn',
    [
        "OK-So definitively David Cameron who is the current British Prime Minister earns £142,500 per year. Recently the Civil Service disclosed the salary's of all of the Cabinet members who earn over £150,000. David Cameron's £142,500 was used as a bench mark for this.",
        'Celebrities wealth, income, and salary in 2016, 2017. Cameron Dallas. Handsome Vine star Cameron Dallas earned his first million before he was 20 years old. His current net worth is approximately $1.2 million as of 2017 and although he never admitted having any official girlfriend, we know the names of last two girls he was dating.',
        'Distance, Gas Consumption and Emission Notes. Distance from Isle of Palms, SC to Charleston, SC is 18Miles or 29 Km. You can get this distance about 39 mins. If you want to planning travel with plane for 13 Miles or 20 Km, You can get this distance about 31 mins . A car with an average MPG will needs 0.83 gallons of gas to get the route between these points. The estimated cost of gas to get between Isle of Palms, SC and Charleston, SC is $1.88. During the route, an average car will release 16.26 pounds of CO 2 to the atmosphere. Your carbon footprint is 0.9 pounds of CO 2 per mile.',
        'Once youâ\x80\x99ve done your obligatory sunbathing and beachcombing at the Isle of Palms County Park Beach, head to nearby Charleston, a 30-minute drive away, which has historic sites, landmarks, shopping districts, golf courses and family attractions like The South Carolina Aquarium and Fountain Walk, a new dining, entertainment and shopping complex nearby.',
        'According to a 2010 article in Managed Care Magazine, the average cost for infants in NICU is around $3,000 per day. This average does not include surgeries or helicopter transports.',
    ]
)
# [{'corpus_id': ..., 'score': ...}, {'corpus_id': ..., 'score': ...}, ...]

Training Details

Training Dataset

parquet

• Dataset: parquet
• Size: 1,048,576 training samples
• Columns: query, passage, and label
• Approximate statistics based on the first 100 samples:

  	query	passage	label
  type	string	string	float
  modality	text	text
  details	min: 4 tokensmean: 8.23 tokensmax: 15 tokens	min: 30 tokensmean: 88.86 tokensmax: 206 tokens	min: 0.0mean: 0.5max: 1.0

• Samples:

  query	passage	label
  how much does david cameron earn	OK-So definitively David Cameron who is the current British Prime Minister earns £142,500 per year. Recently the Civil Service disclosed the salary's of all of the Cabinet members who earn over £150,000. David Cameron's £142,500 was used as a bench mark for this.	1.0
  how much does david cameron earn	Celebrities wealth, income, and salary in 2016, 2017. Cameron Dallas. Handsome Vine star Cameron Dallas earned his first million before he was 20 years old. His current net worth is approximately $1.2 million as of 2017 and although he never admitted having any official girlfriend, we know the names of last two girls he was dating.	0.0
  distance from the island of palms to downtown charleston, sc	Distance, Gas Consumption and Emission Notes. Distance from Isle of Palms, SC to Charleston, SC is 18Miles or 29 Km. You can get this distance about 39 mins. If you want to planning travel with plane for 13 Miles or 20 Km, You can get this distance about 31 mins . A car with an average MPG will needs 0.83 gallons of gas to get the route between these points. The estimated cost of gas to get between Isle of Palms, SC and Charleston, SC is $1.88. During the route, an average car will release 16.26 pounds of CO 2 to the atmosphere. Your carbon footprint is 0.9 pounds of CO 2 per mile.	1.0

• Loss: BinaryCrossEntropyLoss with these parameters:

  {
      "activation_fn": "torch.nn.modules.linear.Identity",
      "pos_weight": 1.0
  }
  

Training Hyperparameters

Non-Default Hyperparameters

• per_device_train_batch_size: 32
• num_train_epochs: 1
• learning_rate: 3e-05
• warmup_steps: 0.05
• optim: adamw_torch
• weight_decay: 0.01
• gradient_accumulation_steps: 4
• bf16: True
• tf32: True
• logging_nan_inf_filter: False
• disable_tqdm: True
• per_device_eval_batch_size: 32
• push_to_hub: True
• hub_private_repo: False
• hub_model_id: tomaarsen/ultrabert-150m-reranker-msmarco524k-3e-5-20261014
• load_best_model_at_end: True
• seed: 12
• data_seed: 12
• dataloader_num_workers: 2

All Hyperparameters

• per_device_train_batch_size: 32
• num_train_epochs: 1
• max_steps: -1
• learning_rate: 3e-05
• lr_scheduler_type: linear
• lr_scheduler_kwargs: None
• warmup_steps: 0.05
• optim: adamw_torch
• optim_args: None
• weight_decay: 0.01
• adam_beta1: 0.9
• adam_beta2: 0.999
• adam_epsilon: 1e-08
• optim_target_modules: None
• gradient_accumulation_steps: 4
• average_tokens_across_devices: True
• max_grad_norm: 1.0
• label_smoothing_factor: 0.0
• bf16: True
• fp16: False
• bf16_full_eval: False
• fp16_full_eval: False
• tf32: True
• gradient_checkpointing: False
• gradient_checkpointing_kwargs: None
• torch_compile: False
• torch_compile_backend: None
• torch_compile_mode: None
• use_liger_kernel: False
• liger_kernel_config: None
• use_cache: False
• neftune_noise_alpha: None
• torch_empty_cache_steps: None
• auto_find_batch_size: False
• log_on_each_node: True
• logging_nan_inf_filter: False
• include_num_input_tokens_seen: no
• log_level: passive
• log_level_replica: warning
• disable_tqdm: True
• project: huggingface
• trackio_space_id: None
• trackio_bucket_id: None
• trackio_static_space_id: None
• per_device_eval_batch_size: 32
• prediction_loss_only: True
• eval_on_start: False
• eval_do_concat_batches: True
• eval_use_gather_object: False
• eval_accumulation_steps: None
• include_for_metrics: []
• batch_eval_metrics: False
• save_only_model: False
• save_on_each_node: False
• enable_jit_checkpoint: False
• push_to_hub: True
• hub_private_repo: False
• hub_model_id: tomaarsen/ultrabert-150m-reranker-msmarco524k-3e-5-20261014
• hub_strategy: every_save
• hub_always_push: False
• hub_revision: None
• load_best_model_at_end: True
• ignore_data_skip: False
• restore_callback_states_from_checkpoint: False
• full_determinism: False
• seed: 12
• data_seed: 12
• use_cpu: False
• accelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}
• parallelism_config: None
• dataloader_drop_last: False
• dataloader_num_workers: 2
• dataloader_pin_memory: True
• dataloader_persistent_workers: False
• dataloader_prefetch_factor: None
• dataloader_multiprocessing_context: None
• dataloader_in_order: True
• remove_unused_columns: True
• label_names: None
• train_sampling_strategy: random
• length_column_name: length
• ddp_find_unused_parameters: None
• ddp_bucket_cap_mb: None
• ddp_broadcast_buffers: False
• ddp_static_graph: None
• ddp_backend: None
• ddp_timeout: 1800
• fsdp: None
• fsdp_config: None
• deepspeed: None
• debug: []
• skip_memory_metrics: True
• do_predict: False
• resume_from_checkpoint: None
• local_rank: -1
• prompts: None
• batch_sampler: batch_sampler
• multi_dataset_batch_sampler: proportional
• router_mapping: {}
• learning_rate_mapping: {}
• warmup_ratio: None

Training Logs

• The bold row denotes the saved checkpoint.

Training Time

• Training: 57.6 minutes
• Evaluation: 3.8 minutes
• Total: 1.0 hours

Framework Versions

• Python: 3.12.12
• Sentence Transformers: 6.2.0.dev0
• Transformers: 5.16.1
• PyTorch: 2.10.0+cu128
• Accelerate: 1.15.0
• Datasets: 5.1.0
• Tokenizers: 0.23.3

Additional Resources

• Training and Finetuning Reranker Models with Sentence Transformers: the end-to-end guide for training or finetuning Cross Encoder (reranker) models.
• Multimodal Embedding & Reranker Models with Sentence Transformers: use text, image, audio, and video reranker models through the same API.
• Training and Finetuning Multimodal Embedding & Reranker Models with Sentence Transformers: training multimodal Cross Encoders.

Citation

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

Retrieval scaling experiment

One epoch on 524,288 shared MS MARCO triplets, seed 12, LR 3e-05. Query content is clipped to 32 tokens, passage content to 256, combined cap 291. Checkpoint selected by the original 1,024-query validation set. Test scores do not select checkpoints. Experiment artifacts.