ColBERT ModernBERT-base trained on MS MARCO triplets with GradCache

This is a Multi-Vector Encoder model finetuned from answerdotai/ModernBERT-base on the msmarco-bm25 dataset using the sentence-transformers library. It maps inputs to sequences of 128-dimensional token-level vectors and scores them with late interaction (MaxSim), useful for semantic search with late interaction.

Model Details

Model Description

• Model Type: Multi-Vector Encoder
• Base model: answerdotai/ModernBERT-base
• Maximum Sequence Length: 8192 tokens
• Output Dimensionality: 128 dimensions
• Similarity Function: maxsim
• Supported Modality: Text
• Training Dataset:
  • msmarco-bm25
• Language: en
• License: apache-2.0

Model Sources

• Documentation: Sentence Transformers Documentation
• Documentation: Multi-Vector Encoder Documentation
• Repository: Sentence Transformers on GitHub
• Hugging Face: Multi-Vector Encoders on Hugging Face

Full Model Architecture

MultiVectorEncoder(
  (0): Transformer({'transformer_task': 'feature-extraction', 'modality_config': {'text': {'method': 'forward', 'method_output_name': 'last_hidden_state'}}, 'module_output_name': 'token_embeddings', 'architecture': 'ModernBertModel'})
  (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'module_input_name': 'token_embeddings', 'module_output_name': 'token_embeddings'})
  (2): MultiVectorMask({'skiplist_words': [], 'keep_only_token_ids': None})
  (3): Normalize({'module_input_name': 'token_embeddings', 'module_output_name': 'token_embeddings'})
)

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 MultiVectorEncoder

# Download from the 🤗 Hub
model = MultiVectorEncoder("tomaarsen/multivector-ModernBERT-base-msmarco-cached-contrastive")
# Run inference: each input becomes a sequence of per-token vectors (variable length).
queries = [
    'what is territorial sovereignty',
]
documents = [
    'territorial sovereignty. Exclusive right of a state to exercise its powers within the boundaries of its territory.',
    'Territorial preservation, as Agnew explains is merely one aspect of a states territorial integrity (2005). The lack of territorial sovereignty is often a key characteristic of so-called failed states where effective monopoly over the internal means of violence is lost.',
    '1 Active Transportï\x82§ Active Transport requires the cell to use energy, usually in the form of ATP.ï\x82§ Active Transport creates a charge gradient in the cell membrane. 2  For example in the mitochondrion, hydrogen ion pumps pump hydrogen ions into the intermembrane space of the organelle as part of making ATP. 3  15.',
]
query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings[0].shape, document_embeddings[0].shape)
# (6, 128) (24, 128)

# Get the MaxSim similarity scores
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[5.6691, 5.5729, 1.3079]])

Evaluation

Metrics

Multi Vector Information Retrieval

• Datasets: NanoMSMARCO, NanoNQ, NanoFiQA2018, NanoClimateFEVER, NanoDBPedia, NanoFEVER, NanoFiQA2018, NanoHotpotQA, NanoMSMARCO, NanoNFCorpus, NanoNQ, NanoQuoraRetrieval, NanoSCIDOCS, NanoArguAna, NanoSciFact and NanoTouche2020
• Evaluated with MultiVectorInformationRetrievalEvaluator

Multi Vector Nano BEIR

• Dataset: NanoBEIR_mean
• Evaluated with MultiVectorNanoBEIREvaluator with these parameters:

  {
      "dataset_names": [
          "msmarco",
          "nq",
          "fiqa2018"
      ],
      "dataset_id": "sentence-transformers/NanoBEIR-en"
  }
  

Multi Vector Nano BEIR

• Dataset: NanoBEIR_mean
• Evaluated with MultiVectorNanoBEIREvaluator with these parameters:

  {
      "dataset_names": [
          "climatefever",
          "dbpedia",
          "fever",
          "fiqa2018",
          "hotpotqa",
          "msmarco",
          "nfcorpus",
          "nq",
          "quoraretrieval",
          "scidocs",
          "arguana",
          "scifact",
          "touche2020"
      ],
      "dataset_id": "sentence-transformers/NanoBEIR-en"
  }
  

Training Details

Training Dataset

msmarco-bm25

• Dataset: msmarco-bm25 at ce8a493
• Size: 99,000 training samples
• Columns: query, positive, and negative
• Approximate statistics based on the first 100 samples:

  	query	positive	negative
  type	string	string	string
  modality	text	text	text
  details	min: 5 tokensmean: 9.4 tokensmax: 18 tokens	min: 28 tokensmean: 83.3 tokensmax: 197 tokens	min: 24 tokensmean: 76.0 tokensmax: 226 tokens

• Samples:

  query	positive	negative
  what is an agate made of	Agate is the name given to a group of silicate minerals that are made up primarily of chalcedony. Chalcedony is a member of the quartz family of minerals.	What is an agate, though, and what are the properties and feng shui meaning of agate? Let's find out. WHAT IS THE MEANING OF AGATE? As a form of chalcedony (type of quartz) agate exhibits a variety of colours , shapes, as well as an often present gentle iridescence.
  what is the socratic method?	The Socratic Learning Method (SLM) is a constructivist learning approach consisting of four key. steps: eliciting relevant preconceptions, clarifying preconceptions, testing ones own. hypotheses or encountered propositions, and deciding whether to accept the hypotheses or. propositions.	The Socratic Learning Method and the Inquiry-Based Learning Method. Since the Socratic dialogues are among the earliest documented instances of learning through. inquiry, it is reasonable to argue that what is now known as inquiry-based learning can trace its. origin to the Socratic Learning Method.
  what ditto means	• DITTO (noun). The noun DITTO has 1 sense: 1. a mark used to indicate the word above it should be repeated. Familiarity information: DITTO used as a noun is very rare. • DITTO (verb). The verb DITTO has 1 sense: 1. repeat an action or statement. Familiarity information: DITTO used as a verb is very rare.	Valerie Hill 0 I had texted one of my Auntie's and she sent the word Ditto back to me I was like what do that mean. Now I know!!! Acronym. DITTO in text means the same, or me too, or I agree.

• Loss: CachedMultiVectorMultipleNegativesRankingLoss with these parameters:

  {
      "score_metric": "colbert_scores",
      "mini_batch_size": 32,
      "mini_batch_num_tokens": null,
      "score_mini_batch_size": 32,
      "scale": 1.0,
      "size_average": true,
      "gather_across_devices": false
  }
  

Evaluation Dataset

msmarco-bm25

• Dataset: msmarco-bm25 at ce8a493
• Size: 1,000 evaluation samples
• Columns: query, positive, and negative
• Approximate statistics based on the first 100 samples:

  	query	positive	negative
  type	string	string	string
  modality	text	text	text
  details	min: 4 tokensmean: 9.18 tokensmax: 21 tokens	min: 22 tokensmean: 76.76 tokensmax: 210 tokens	min: 30 tokensmean: 84.36 tokensmax: 210 tokens

• Samples:

  query	positive	negative
  expected return is a function of	security market line Security market line (SML) is the representation of the capital asset pricing model. It displays the expected rate of return of an individual security as a function of systematic, non-diversifiable risk (its beta). beta Average sensitivity of a security's price to overall securities market prices.	A portfolio's expected return is the sum of the weighted average of each asset's expected return. Calculate a portfolio's expected return. To calculate the expected return of a portfolio, you need to know the expected return and weight of each asset in a portfolio.
  what smoker temperature for barbecue chicken	Chickens smoked hot and fast over indirect heat on the grill. But for pulled chicken, I wanted a slightly more intense smokiness to balance the barbecue sauce. I decided to do a side-by-side comparison, cooking one on the grill over indirect heat at around 375°F, and one in the smoker at 225°F.	Cook for an hour and 15 minutes. Check your smoked chicken breasts to make sure the temperature is still holding at about 250 degrees. Flip the chicken breasts and close the lid again. After 30 minutes, check the internal temperature of the chicken with a meat thermometer. You are looking for a temperature of 160 degrees before you can pull them off the smoker. Serve your smoked chicken breasts with a side of barbecue sauce for dipping.
  in classification of matter what is an element	Matter can be in the same phase or in two different phases for this separation to take place. Key Terms. mixture: Something that consists of diverse, non-bonded elements or molecules. element: A chemical substance that is made up of a particular kind of atom and cannot be broken down or transformed by a chemical reaction.	2. What four elements make up 96% of all living matter? The four elements that make up 96% of all living matter are oxygen, carbon, hydrogen and nitrogen. 3. What is the difference between an essential element and a trace element? An essential element is an element that an organism needs to live a healthy life and reproduce. A trace element is required by an organism in only minute quantities. Section 2 4. Sketch a model of an atom of helium, showing the electrons, protons, neutrons, and atomic nucleus. Neutrons Protons Electrons 5.

• Loss: CachedMultiVectorMultipleNegativesRankingLoss with these parameters:

  {
      "score_metric": "colbert_scores",
      "mini_batch_size": 32,
      "mini_batch_num_tokens": null,
      "score_mini_batch_size": 32,
      "scale": 1.0,
      "size_average": true,
      "gather_across_devices": false
  }
  

Training Hyperparameters

Non-Default Hyperparameters

• per_device_train_batch_size: 256
• num_train_epochs: 1
• learning_rate: 3e-05
• warmup_steps: 0.05
• bf16: True
• per_device_eval_batch_size: 32
• load_best_model_at_end: True
• batch_sampler: no_duplicates

All Hyperparameters

• per_device_train_batch_size: 256
• 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_fused
• optim_args: None
• weight_decay: 0.0
• adam_beta1: 0.9
• adam_beta2: 0.999
• adam_epsilon: 1e-08
• optim_target_modules: None
• gradient_accumulation_steps: 1
• 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: None
• 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: True
• include_num_input_tokens_seen: no
• log_level: passive
• log_level_replica: warning
• disable_tqdm: False
• 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: False
• hub_private_repo: None
• hub_model_id: None
• 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: 42
• data_seed: None
• 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: 0
• dataloader_pin_memory: True
• dataloader_persistent_workers: False
• dataloader_prefetch_factor: None
• 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
• warmup_ratio: None
• local_rank: -1
• prompts: None
• batch_sampler: no_duplicates
• multi_dataset_batch_sampler: proportional
• router_mapping: {}
• learning_rate_mapping: {}
• max_length: None

Training Logs

• The bold row denotes the saved checkpoint.

Training Time

• Training: 41.5 minutes
• Evaluation: 7.2 minutes
• Total: 48.7 minutes

Framework Versions

• Python: 3.11.6
• Sentence Transformers: 5.7.0.dev0
• Transformers: 5.13.1
• PyTorch: 2.10.0+cu128
• Accelerate: 1.14.0
• Datasets: 4.8.4
• Tokenizers: 0.22.2

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",
}

CachedMultiVectorMultipleNegativesRankingLoss

@misc{gao2021scaling,
    title={Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup},
    author={Luyu Gao and Yunyi Zhang and Jiawei Han and Jamie Callan},
    year={2021},
    eprint={2101.06983},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}