wilsonmarciliojr/matryoshka-embed-knn

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")
# 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.4520, 0.5252],
#         [0.4520, 1.0000, 0.2916],
#         [0.5252, 0.2916, 1.0000]])

Evaluation

Metrics

Semantic Similarity

• Datasets: sts-dev-768 and sts-test-768
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 768
  }
  

Metric	sts-dev-768	sts-test-768
pearson_cosine	0.7937	0.7534
spearman_cosine	0.7849	0.7348

Semantic Similarity

• Datasets: sts-dev-512 and sts-test-512
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 512
  }
  

Metric	sts-dev-512	sts-test-512
pearson_cosine	0.7938	0.7538
spearman_cosine	0.7856	0.7355

Semantic Similarity

• Datasets: sts-dev-256 and sts-test-256
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 256
  }
  

Metric	sts-dev-256	sts-test-256
pearson_cosine	0.7909	0.7523
spearman_cosine	0.7838	0.7354

Semantic Similarity

• Datasets: sts-dev-64 and sts-test-64
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 64
  }
  

Metric	sts-dev-64	sts-test-64
pearson_cosine	0.78	0.7396
spearman_cosine	0.7781	0.7274

Semantic Similarity

• Datasets: sts-dev-2 and sts-test-2
• Evaluated with EmbeddingSimilarityEvaluator with these parameters:

  {
      "truncate_dim": 2
  }
  

Metric	sts-dev-2	sts-test-2
pearson_cosine	0.2275	0.2158
spearman_cosine	0.3273	0.323

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: 324
• per_device_eval_batch_size: 324
• 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: 324
• per_device_eval_batch_size: 324
• 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	Training Loss	Validation Loss	sts-dev-768_spearman_cosine	sts-dev-512_spearman_cosine	sts-dev-256_spearman_cosine	sts-dev-64_spearman_cosine	sts-dev-2_spearman_cosine	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
0.1011	1000	15.5729	11.0151	0.8103	0.8099	0.8066	0.7934	0.3373	-	-	-	-	-
0.2022	2000	11.3861	10.5141	0.8084	0.8084	0.8063	0.7939	0.3636	-	-	-	-	-
0.3033	3000	10.5075	10.5011	0.8061	0.8068	0.8051	0.7960	0.3659	-	-	-	-	-
0.4044	4000	9.9884	10.7506	0.7900	0.7906	0.7875	0.7770	0.3387	-	-	-	-	-
0.5056	5000	9.5581	10.7602	0.7997	0.7997	0.7974	0.7882	0.3496	-	-	-	-	-
0.6067	6000	9.2037	10.6260	0.7930	0.7930	0.7910	0.7846	0.3549	-	-	-	-	-
0.7078	7000	8.9519	10.5886	0.7910	0.7921	0.7903	0.7840	0.3342	-	-	-	-	-
0.8089	8000	8.7682	10.6864	0.7896	0.7903	0.7881	0.7819	0.3311	-	-	-	-	-
0.9100	9000	8.6166	10.6835	0.7849	0.7856	0.7838	0.7781	0.3273	-	-	-	-	-
-1	-1	-	-	-	-	-	-	-	0.7348	0.7355	0.7354	0.7274	0.3230

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