BERT tiny multi-vector encoder trained on MS MARCO

This is a Multi-Vector Encoder model finetuned in two stages from prajjwal1/bert-tiny 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: prajjwal1/bert-tiny
• Maximum Sequence Length: 512 tokens
  • Maximum Query Length: 32 tokens
  • Maximum Document Length: 256 tokens
• Output Dimensionality: 128 dimensions
• Similarity Function: MaxSim
• Supported Modality: Text
• Training Dataset:
  • msmarco-bm25
• Language: en
• License: mit

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', 'query_length': 32, 'document_length': 256, 'query_expansion': {'strategy': 'min', 'attend': False, 'token': None, 'length': 32}, 'architecture': 'BertModel'})
  (1): Dense({'in_features': 128, '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': [], 'skiplist_tasks': ['document'], '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("multi-vector-encoder-testing/bert-tiny-multi-vector")
# Run inference: each input becomes a sequence of per-token vectors (variable length).
queries = [
    'calories in kirkland ravioli',
]
documents = [
    'There are 140 calories in a 4 pieces serving of Kirkland Signature Four Cheese Ravioli. Calorie breakdown: 45% fat, 31% carbs, 24% protein.',
    'There are 140 calories in a 4 pieces serving of Kirkland Signature Four Cheese Ravioli. Calorie breakdown: 45% fat, 31% carbs, 24% protein.',
    'Current Local Time: Cleveland, Ohio is in the Eastern Time Zone: The Current Time in Cleveland, Ohio is: Thursday 1/18/2018 10:41 PM EST Cleveland, Ohio is in the Eastern Time Zone',
]
query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings[0].shape, document_embeddings[0].shape)
# (32, 128) (39, 128)

# Get the MaxSim similarity scores
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[25.4110, 25.4110,  8.3390]])

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: 501,907 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: 4 tokens mean: 8.81 tokens max: 16 tokens	min: 25 tokens mean: 78.08 tokens max: 174 tokens	min: 22 tokens mean: 76.59 tokens max: 180 tokens

• Samples:

  query	positive	negative
  sociopath define	Updated September 08, 2016. Both psychopaths and sociopaths are defined as someone who is suffering from Antisocial Personality Disorder. Both groups show a pervasive pattern of disregard for the rights and feelings of others. There are, however, subtle differences between the two groups.	Define sociopath. sociopath synonyms, sociopath pronunciation, sociopath translation, English dictionary definition of sociopath. n. A psychopath or a person with antisocial personality disorder. so′ci·o·path′ic adj. so′ci·op′a·thy n. n psychiatry another name for psychopath...
  what county is tarrytown ny in	ABOUT US. The Music Hall, an 1885 landmark in Tarrytown, NY is Westchester County's oldest theater and one of the region's busiest music venues, welcoming 85,000 visitors every year, including tens of thousands of children. Please complete all required fields!	Tarrytown, NY Other Information. 1 Located in WESTCHESTER County, New York. 2 Tarrytown, NY is also known as: 3 N TARRYTOWN, NY. NORTH TARRYTOWN, 1 NY. PHILIPSE MANOR, 2 NY. POCANTICO HILLS, 3 NY. SLEEPY HOLLOW, NY. SLEEPY HOLLOW MANOR, NY.
  what temperature do you grill a t-bone at	Step 2. Move your T-bones to the medium heat side of your grill and continue grilling. If you like your T-bone medium rare, grill for four to five minutes on each side or until a meat thermometer reads 130 to 140 degrees Fahrenheit.For a medium steak, grill six to seven minutes per side or until a meat thermometer reads 140 to 150 degrees.iming and technique are the keys to grilling a 1-inch-thick T-bone steak to perfection. This thicker cut requires different treatment from a T-bone, for example, 1/2- to 3/4-inch thick.	Preheat your grill using two temperature settings. If you are using a gas grill, set one side to high and the other to a medium setting, then close the lid for 10 to 15 minutes.iming and technique are the keys to grilling a 1-inch-thick T-bone steak to perfection. This thicker cut requires different treatment from a T-bone, for example, 1/2- to 3/4-inch thick.

• Loss: MultiVectorMultipleNegativesRankingLoss with these parameters:

  {
      "scale": 1.0,
      "similarity_fct": "colbert_scores",
      "mini_batch_size": null,
      "score_mini_batch_size": null,
      "gather_across_devices": false
  }
  

Evaluation Dataset

msmarco-bm25

• Dataset: msmarco-bm25 at ce8a493
• Size: 1,024 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 tokens mean: 8.65 tokens max: 15 tokens	min: 21 tokens mean: 80.76 tokens max: 174 tokens	min: 16 tokens mean: 79.88 tokens max: 224 tokens

• Samples:

  query	positive	negative
  what is the bottom lip piercing called	1 Standard Lip Piercing – This is a single piercing done off centered on the lower lip. 2 A captive bead ring (CBR) or labret stud can be worn. 3 Monroe Piercing – This is a single piercing done on the left side of the upper lip and is named for the mole on Marilyn Monroe’s lip. 4 Usually a labret stud is worn in this piercing.	Also known as lower-lip piercing or bottom lip piercing .The labret piercing is placed at the labrum (below bottom lip, above chin). Popular among men and women, this style looks super cool and trendy.
  what are the mind and body	This is known as dualism. Dualism is the view that the mind and body both exist. There are two basic types of dualism: o Descartes dualism: The view that the mind and body function separately, without interchange. o Cartesian dualism argues that there is a two-way interaction between mental and physical substances. Dualism is in contrast to monism that states the mind and body are the same thing.	Quotes About What Matters In Life. “What is in your mind position or disposition your mind, body and spirit in the best or worst way. What you are yet to accept into your mind exposes your mind to and keep your mind on what you are yet to accept and what has not yet come into your mind least controls your mind, body and spirit. Browse By Tag.
  what breed of dogs have green eyes	Best Answer: There are many dog breeds that CAN have green eyes, such as Australian Shepherds, Border Collies, Siberian Huskies, and others, but it is an uncommon occurrence. However it won't be a bright green like a cat's eye. It'll be a somewhat subdued shade of blueish-grey with green overtones.	What are some dog breeds that have or can have green eyes? What breed is this dog? What is the breed of a dog, which looks like a fox and has light blue eyes, called? Rohit Akut, love and respect animals be friendly have had lots of different pets..

• Loss: MultiVectorMultipleNegativesRankingLoss with these parameters:

  {
      "scale": 1.0,
      "similarity_fct": "colbert_scores",
      "mini_batch_size": null,
      "score_mini_batch_size": null,
      "gather_across_devices": false
  }
  

Training Hyperparameters

Non-Default Hyperparameters

• per_device_train_batch_size: 128
• max_steps: 10000
• learning_rate: 1e-05
• warmup_steps: 0.05
• weight_decay: 0.01
• bf16: True
• disable_tqdm: True
• per_device_eval_batch_size: 32
• load_best_model_at_end: True
• seed: 12
• batch_sampler: no_duplicates

All Hyperparameters

Click to expand

• per_device_train_batch_size: 128
• num_train_epochs: 3.0
• max_steps: 10000
• learning_rate: 1e-05
• lr_scheduler_type: linear
• lr_scheduler_kwargs: None
• warmup_steps: 0.05
• optim: adamw_torch_fused
• 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: 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: 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: 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: 12
• 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
• 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: no_duplicates
• multi_dataset_batch_sampler: proportional
• router_mapping: {}
• learning_rate_mapping: {}
• warmup_ratio: None
• max_length: None

Training Logs

Click to expand

• The bold row denotes the saved checkpoint.

Training Time

• Training: 15.5 minutes
• Evaluation: 19.4 seconds
• Total: 15.9 minutes

Framework Versions

• Python: 3.11.6
• Sentence Transformers: 6.1.0.dev0
• Transformers: 5.16.1
• PyTorch: 2.10.0+cu128
• Accelerate: 1.14.0
• Datasets: 4.8.4
• Tokenizers: 0.23.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",
}

MultiVectorMultipleNegativesRankingLoss

@misc{henderson2017efficient,
    title={Efficient Natural Language Response Suggestion for Smart Reply},
    author={Matthew Henderson and Rami Al-Rfou and Brian Strope and Yun-hsuan Sung and Laszlo Lukacs and Ruiqi Guo and Sanjiv Kumar and Balint Miklos and Ray Kurzweil},
    year={2017},
    eprint={1705.00652},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}