wilsonmarciliojr/matryoshka-embed-knn-b16

SentenceTransformer based on distilbert/distilroberta-base

This is a sentence-transformers model finetuned from distilbert/distilroberta-base on the all-nli-knn-hard-negatives dataset. It maps sentences & paragraphs to a 768-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more.

Model Details

Model Description

• Model Type: Sentence Transformer
• Base model: distilbert/distilroberta-base
• Maximum Sequence Length: 512 tokens
• Output Dimensionality: 768 dimensions
• Similarity Function: Cosine Similarity
• Training Dataset:
  • all-nli-knn-hard-negatives

Model Sources

• Documentation: Sentence Transformers Documentation
• Repository: Sentence Transformers on GitHub
• Hugging Face: Sentence Transformers on Hugging Face

Full Model Architecture

SentenceTransformer(
  (0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'RobertaModel'})
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
)

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 SentenceTransformer

# Download from the 🤗 Hub
model = SentenceTransformer("wilsonmarciliojr/matryoshka-embed-knn-b16")
# Run inference
sentences = [
    'A baby at the end of a slip and slide at a party',
    'A man is playing with a baby on a deck.',
    'A baby in a bib is making funny faces at the camera.',
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 768]

# Get the similarity scores for the embeddings
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[1.0000, 0.3685, 0.3925],
#         [0.3685, 1.0000, 0.3452],
#         [0.3925, 0.3452, 1.0000]])

Evaluation

Metrics

Semantic Similarity

• Dataset: sts-test-768
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 768
  }
  

Metric	Value
pearson_cosine	0.7119
spearman_cosine	0.6905

Semantic Similarity

• Dataset: sts-test-512
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 512
  }
  

Metric	Value
pearson_cosine	0.7104
spearman_cosine	0.6889

Semantic Similarity

• Dataset: sts-test-256
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 256
  }
  

Metric	Value
pearson_cosine	0.7092
spearman_cosine	0.6883

Semantic Similarity

• Dataset: sts-test-64
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 64
  }
  

Metric	Value
pearson_cosine	0.697
spearman_cosine	0.677

Semantic Similarity

• Dataset: sts-test-2
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 2
  }
  

Metric	Value
pearson_cosine	0.2065
spearman_cosine	0.3003

Training Details

Training Dataset

all-nli-knn-hard-negatives

• Dataset: all-nli-knn-hard-negatives at c7814a7
• Size: 3,204,256 training samples
• Columns: anchor, positive, negative_1, negative_2, negative_3, negative_4, and negative_5
• Approximate statistics based on the first 1000 samples:

  	anchor	positive	negative_1	negative_2	negative_3	negative_4	negative_5
  type	string	string	string	string	string	string	string
  details	min: 8 tokens mean: 16.58 tokens max: 50 tokens	min: 6 tokens mean: 14.56 tokens max: 59 tokens	min: 5 tokens mean: 9.62 tokens max: 16 tokens	min: 5 tokens mean: 9.24 tokens max: 16 tokens	min: 5 tokens mean: 9.16 tokens max: 18 tokens	min: 6 tokens mean: 9.43 tokens max: 17 tokens	min: 5 tokens mean: 9.41 tokens max: 18 tokens

• Samples:

  anchor	positive	negative_1	negative_2	negative_3	negative_4	negative_5
  An older man is drinking orange juice at a restaurant.	An elderly man is drinking orange juice at a cafe.	An elderly gentleman eats.	A man has many oranges in his baskets.	An elderly person is being served food	A man works at a restaurant	There is a older man.
  An older man is drinking orange juice at a restaurant.	A man drinking orange juice while walking.	An elderly gentleman eats.	A man has many oranges in his baskets.	An elderly person is being served food	A man works at a restaurant	There is a older man.
  An older man is drinking orange juice at a restaurant.	A man drinks orange juice and walks outside.	An elderly gentleman eats.	A man has many oranges in his baskets.	An elderly person is being served food	A man works at a restaurant	There is a older man.

• Loss: MatryoshkaLoss with these parameters:

  {
      "loss": "MultipleNegativesRankingLoss",
      "matryoshka_dims": [
          768,
          512,
          256,
          64,
          2
      ],
      "matryoshka_weights": [
          1,
          1,
          1,
          1,
          1
      ],
      "n_dims_per_step": -1
  }
  

Evaluation Dataset

all-nli-knn-hard-negatives

• Dataset: all-nli-knn-hard-negatives at c7814a7
• Size: 103,904 evaluation samples
• Columns: anchor, positive, negative_1, negative_2, negative_3, negative_4, and negative_5
• Approximate statistics based on the first 1000 samples:

  	anchor	positive	negative_1	negative_2	negative_3	negative_4	negative_5
  type	string	string	string	string	string	string	string
  details	min: 9 tokens mean: 17.34 tokens max: 36 tokens	min: 6 tokens mean: 17.12 tokens max: 52 tokens	min: 5 tokens mean: 9.27 tokens max: 21 tokens	min: 6 tokens mean: 9.98 tokens max: 21 tokens	min: 5 tokens mean: 9.35 tokens max: 21 tokens	min: 5 tokens mean: 9.12 tokens max: 16 tokens	min: 6 tokens mean: 9.47 tokens max: 23 tokens

• Samples:

  anchor	positive	negative_1	negative_2	negative_3	negative_4	negative_5
  Two women are embracing while holding to go packages.	Two women in a embrace of greetings, one of them is holding flowers and they are greeting each other of a kiss.	Two women are in the city.	The women each have one head.	Two women are drinking wine and having a conversation.	women carry food on plates	Two people are kissing each other.
  Two women are embracing while holding to go packages.	Two women wearing boots and holding bags are talking to each other.	Two women are in the city.	The women each have one head.	Two women are drinking wine and having a conversation.	women carry food on plates	Two people are kissing each other.
  Two women are embracing while holding to go packages.	Two women are wet while holding hands with a long building and buses in the background.	Two women are in the city.	The women each have one head.	Two women are drinking wine and having a conversation.	women carry food on plates	Two people are kissing each other.

• Loss: MatryoshkaLoss with these parameters:

  {
      "loss": "MultipleNegativesRankingLoss",
      "matryoshka_dims": [
          768,
          512,
          256,
          64,
          2
      ],
      "matryoshka_weights": [
          1,
          1,
          1,
          1,
          1
      ],
      "n_dims_per_step": -1
  }
  

Training Hyperparameters

Non-Default Hyperparameters

• eval_strategy: steps
• per_device_train_batch_size: 16
• per_device_eval_batch_size: 16
• num_train_epochs: 1
• warmup_ratio: 0.1
• fp16: True
• batch_sampler: no_duplicates

All Hyperparameters

Click to expand

• overwrite_output_dir: False
• do_predict: False
• eval_strategy: steps
• prediction_loss_only: True
• per_device_train_batch_size: 16
• per_device_eval_batch_size: 16
• per_gpu_train_batch_size: None
• per_gpu_eval_batch_size: None
• gradient_accumulation_steps: 1
• eval_accumulation_steps: None
• torch_empty_cache_steps: None
• learning_rate: 5e-05
• weight_decay: 0.0
• adam_beta1: 0.9
• adam_beta2: 0.999
• adam_epsilon: 1e-08
• max_grad_norm: 1.0
• num_train_epochs: 1
• max_steps: -1
• lr_scheduler_type: linear
• lr_scheduler_kwargs: {}
• warmup_ratio: 0.1
• warmup_steps: 0
• log_level: passive
• log_level_replica: warning
• log_on_each_node: True
• logging_nan_inf_filter: True
• save_safetensors: True
• save_on_each_node: False
• save_only_model: False
• restore_callback_states_from_checkpoint: False
• no_cuda: False
• use_cpu: False
• use_mps_device: False
• seed: 42
• data_seed: None
• jit_mode_eval: False
• use_ipex: False
• bf16: False
• fp16: True
• fp16_opt_level: O1
• half_precision_backend: auto
• bf16_full_eval: False
• fp16_full_eval: False
• tf32: None
• local_rank: 0
• ddp_backend: None
• tpu_num_cores: None
• tpu_metrics_debug: False
• debug: []
• dataloader_drop_last: False
• dataloader_num_workers: 0
• dataloader_prefetch_factor: None
• past_index: -1
• disable_tqdm: False
• remove_unused_columns: True
• label_names: None
• load_best_model_at_end: False
• ignore_data_skip: False
• fsdp: []
• fsdp_min_num_params: 0
• fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}
• fsdp_transformer_layer_cls_to_wrap: None
• accelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}
• deepspeed: None
• label_smoothing_factor: 0.0
• optim: adamw_torch
• optim_args: None
• adafactor: False
• group_by_length: False
• length_column_name: length
• ddp_find_unused_parameters: None
• ddp_bucket_cap_mb: None
• ddp_broadcast_buffers: False
• dataloader_pin_memory: True
• dataloader_persistent_workers: False
• skip_memory_metrics: True
• use_legacy_prediction_loop: False
• push_to_hub: False
• resume_from_checkpoint: None
• hub_model_id: None
• hub_strategy: every_save
• hub_private_repo: None
• hub_always_push: False
• hub_revision: None
• gradient_checkpointing: False
• gradient_checkpointing_kwargs: None
• include_inputs_for_metrics: False
• include_for_metrics: []
• eval_do_concat_batches: True
• fp16_backend: auto
• push_to_hub_model_id: None
• push_to_hub_organization: None
• mp_parameters:
• auto_find_batch_size: False
• full_determinism: False
• torchdynamo: None
• ray_scope: last
• ddp_timeout: 1800
• torch_compile: False
• torch_compile_backend: None
• torch_compile_mode: None
• include_tokens_per_second: False
• include_num_input_tokens_seen: False
• neftune_noise_alpha: None
• optim_target_modules: None
• batch_eval_metrics: False
• eval_on_start: False
• use_liger_kernel: False
• liger_kernel_config: None
• eval_use_gather_object: False
• average_tokens_across_devices: False
• prompts: None
• batch_sampler: no_duplicates
• multi_dataset_batch_sampler: proportional
• router_mapping: {}
• learning_rate_mapping: {}

Training Logs

Epoch	Step	sts-test-768_spearman_cosine	sts-test-512_spearman_cosine	sts-test-256_spearman_cosine	sts-test-64_spearman_cosine	sts-test-2_spearman_cosine
-1	-1	0.6905	0.6889	0.6883	0.6770	0.3003

Framework Versions

• Python: 3.10.12
• Sentence Transformers: 5.0.0
• Transformers: 4.53.1
• PyTorch: 2.7.1+cu126
• Accelerate: 1.8.1
• Datasets: 3.6.0
• Tokenizers: 0.21.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",
}

MatryoshkaLoss

@misc{kusupati2024matryoshka,
    title={Matryoshka Representation Learning},
    author={Aditya Kusupati and Gantavya Bhatt and Aniket Rege and Matthew Wallingford and Aditya Sinha and Vivek Ramanujan and William Howard-Snyder and Kaifeng Chen and Sham Kakade and Prateek Jain and Ali Farhadi},
    year={2024},
    eprint={2205.13147},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

MultipleNegativesRankingLoss

@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}
}