Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup
Paper • 2101.06983 • Published • 3
How to use ML-Intern-lab/embeddinggemma-2-allnli-triplet with sentence-transformers:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("ML-Intern-lab/embeddinggemma-2-allnli-triplet")
sentences = [
"Emeralds?",
"I am not sure.",
"Emeralds?",
"Diamonds?"
]
embeddings = model.encode(sentences)
similarities = model.similarity(embeddings, embeddings)
print(similarities.shape)
# [4, 4]This is a sentence-transformers model finetuned from google/embeddinggemma-2 on the sentence-transformers/all-nli dataset. It maps inputs to a 768-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, classification, clustering, and more.
SentenceTransformer(
(0): Transformer({'transformer_task': 'feature-extraction', 'modality_config': {'text': {'method': 'forward', 'method_output_name': 'last_hidden_state'}, 'image': {'method': 'forward', 'method_output_name': 'last_hidden_state'}, 'audio': {'method': 'forward', 'method_output_name': 'last_hidden_state'}, 'video': {'method': 'forward', 'method_output_name': 'last_hidden_state'}, 'message': {'method': 'forward', 'method_output_name': 'last_hidden_state', 'format': 'structured'}}, 'module_output_name': 'token_embeddings', 'architecture': 'EmbeddingGemma2Model'})
(1): Pooling({'embedding_dimension': 768, 'pooling_mode': 'mean', 'include_prompt': True})
(2): Normalize({'module_input_name': 'sentence_embedding', 'module_output_name': 'sentence_embedding'})
)
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("ML-Intern-lab/embeddinggemma-2-allnli-triplet")
# Run inference
queries = [
'Two women are embracing while holding to go packages.',
]
documents = [
'Two woman are holding packages.',
'Two kids in numbered jerseys wash their hands.',
'A man selling donuts to a customer.',
]
query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings.shape)
# [1, 768] [3, 768]
# Get the similarity scores for the embeddings
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[0.8584, 0.3772, 0.3133]])
all-nli-devTripletEvaluator| Metric | Value |
|---|---|
| cosine_accuracy | 0.9459 |
anchor, positive, and negative| anchor | positive | negative | |
|---|---|---|---|
| type | string | string | string |
| modality | text | text | text |
| details |
|
|
|
| anchor | positive | negative |
|---|---|---|
A person on a horse jumps over a broken down airplane. |
A person is outdoors, on a horse. |
A person is at a diner, ordering an omelette. |
Children smiling and waving at camera |
There are children present |
The kids are frowning |
A boy is jumping on skateboard in the middle of a red bridge. |
The boy does a skateboarding trick. |
The boy skates down the sidewalk. |
CachedMultipleNegativesRankingLoss with these parameters:{
"scale": 20.0,
"similarity_fct": "cos_sim",
"mini_batch_size": 16,
"mini_batch_num_tokens": null,
"gather_across_devices": false,
"directions": [
"query_to_doc"
],
"partition_mode": "joint",
"hardness_mode": null,
"hardness_strength": 0.0
}
anchor, positive, and negative| anchor | positive | negative | |
|---|---|---|---|
| type | string | string | string |
| modality | text | text | text |
| details |
|
|
|
| anchor | positive | negative |
|---|---|---|
Two women are embracing while holding to go packages. |
Two woman are holding packages. |
The men are fighting outside a deli. |
Two young children in blue jerseys, one with the number 9 and one with the number 2 are standing on wooden steps in a bathroom and washing their hands in a sink. |
Two kids in numbered jerseys wash their hands. |
Two kids in jackets walk to school. |
A man selling donuts to a customer during a world exhibition event held in the city of Angeles |
A man selling donuts to a customer. |
A woman drinks her coffee in a small cafe. |
CachedMultipleNegativesRankingLoss with these parameters:{
"scale": 20.0,
"similarity_fct": "cos_sim",
"mini_batch_size": 16,
"mini_batch_num_tokens": null,
"gather_across_devices": false,
"directions": [
"query_to_doc"
],
"partition_mode": "joint",
"hardness_mode": null,
"hardness_strength": 0.0
}
per_device_train_batch_size: 128num_train_epochs: 1max_steps: 1500learning_rate: 2e-05warmup_steps: 0.1bf16: Trueper_device_eval_batch_size: 128batch_sampler: no_duplicatesper_device_train_batch_size: 128num_train_epochs: 1max_steps: 1500learning_rate: 2e-05lr_scheduler_type: linearlr_scheduler_kwargs: Nonewarmup_steps: 0.1optim: adamw_torch_fusedoptim_args: Noneweight_decay: 0.0adam_beta1: 0.9adam_beta2: 0.999adam_epsilon: 1e-08optim_target_modules: Nonegradient_accumulation_steps: 1average_tokens_across_devices: Truemax_grad_norm: 1.0label_smoothing_factor: 0.0bf16: Truefp16: Falsebf16_full_eval: Falsefp16_full_eval: Falsetf32: Nonegradient_checkpointing: Falsegradient_checkpointing_kwargs: Nonetorch_compile: Falsetorch_compile_backend: Nonetorch_compile_mode: Noneuse_liger_kernel: Falseliger_kernel_config: Noneuse_cache: Falseneftune_noise_alpha: Nonetorch_empty_cache_steps: Noneauto_find_batch_size: Falselog_on_each_node: Truelogging_nan_inf_filter: Trueinclude_num_input_tokens_seen: nolog_level: passivelog_level_replica: warningdisable_tqdm: Falseproject: huggingfacetrackio_space_id: Nonetrackio_bucket_id: Nonetrackio_static_space_id: Noneper_device_eval_batch_size: 128prediction_loss_only: Trueeval_on_start: Falseeval_do_concat_batches: Trueeval_use_gather_object: Falseeval_accumulation_steps: Noneinclude_for_metrics: []batch_eval_metrics: Falsesave_only_model: Falsesave_on_each_node: Falseenable_jit_checkpoint: Falsepush_to_hub: Falsehub_private_repo: Nonehub_model_id: Nonehub_strategy: every_savehub_always_push: Falsehub_revision: Noneload_best_model_at_end: Falseignore_data_skip: Falserestore_callback_states_from_checkpoint: Falsefull_determinism: Falseseed: 42data_seed: Noneuse_cpu: Falseaccelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}parallelism_config: Nonedataloader_drop_last: Falsedataloader_num_workers: 0dataloader_pin_memory: Truedataloader_persistent_workers: Falsedataloader_prefetch_factor: Nonedataloader_multiprocessing_context: Nonedataloader_in_order: Trueremove_unused_columns: Truelabel_names: Nonetrain_sampling_strategy: randomlength_column_name: lengthddp_find_unused_parameters: Noneddp_bucket_cap_mb: Noneddp_broadcast_buffers: Falseddp_static_graph: Noneddp_backend: Noneddp_timeout: 1800fsdp: Nonefsdp_config: Nonedeepspeed: Nonedebug: []skip_memory_metrics: Truedo_predict: Falseresume_from_checkpoint: Nonelocal_rank: -1prompts: Nonebatch_sampler: no_duplicatesmulti_dataset_batch_sampler: proportionalrouter_mapping: {}learning_rate_mapping: {}warmup_ratio: None| Epoch | Step | Training Loss | Validation Loss | all-nli-dev_cosine_accuracy |
|---|---|---|---|---|
| -1 | -1 | - | - | 0.9036 |
| 0.0046 | 20 | 1.6320 | - | - |
| 0.0092 | 40 | 1.5520 | - | - |
| 0.0138 | 60 | 1.4255 | - | - |
| 0.0184 | 80 | 1.2147 | - | - |
| 0.0229 | 100 | 1.0517 | - | - |
| 0.0275 | 120 | 0.8570 | - | - |
| 0.0321 | 140 | 0.8000 | - | - |
| 0.0367 | 160 | 0.7947 | - | - |
| 0.0413 | 180 | 0.7525 | - | - |
| 0.0459 | 200 | 0.6815 | - | - |
| 0.0505 | 220 | 0.6914 | - | - |
| 0.0551 | 240 | 0.6813 | - | - |
| 0.0574 | 250 | - | 0.3524 | 0.9406 |
| 0.0596 | 260 | 0.6740 | - | - |
| 0.0642 | 280 | 0.6508 | - | - |
| 0.0688 | 300 | 0.6537 | - | - |
| 0.0734 | 320 | 0.6400 | - | - |
| 0.0780 | 340 | 0.6726 | - | - |
| 0.0826 | 360 | 0.6459 | - | - |
| 0.0872 | 380 | 0.6518 | - | - |
| 0.0918 | 400 | 0.6561 | - | - |
| 0.0964 | 420 | 0.6366 | - | - |
| 0.1009 | 440 | 0.6344 | - | - |
| 0.1055 | 460 | 0.6060 | - | - |
| 0.1101 | 480 | 0.6290 | - | - |
| 0.1147 | 500 | 0.7045 | 0.3279 | 0.9443 |
| 0.1193 | 520 | 0.6654 | - | - |
| 0.1239 | 540 | 0.6329 | - | - |
| 0.1285 | 560 | 0.6248 | - | - |
| 0.1331 | 580 | 0.6289 | - | - |
| 0.1376 | 600 | 0.6293 | - | - |
| 0.1422 | 620 | 0.6200 | - | - |
| 0.1468 | 640 | 0.6097 | - | - |
| 0.1514 | 660 | 0.6558 | - | - |
| 0.1560 | 680 | 0.6488 | - | - |
| 0.1606 | 700 | 0.6108 | - | - |
| 0.1652 | 720 | 0.6311 | - | - |
| 0.1698 | 740 | 0.6408 | - | - |
| 0.1721 | 750 | - | 0.3174 | 0.9449 |
| 0.1744 | 760 | 0.6349 | - | - |
| 0.1789 | 780 | 0.6083 | - | - |
| 0.1835 | 800 | 0.6153 | - | - |
| 0.1881 | 820 | 0.6040 | - | - |
| 0.1927 | 840 | 0.6180 | - | - |
| 0.1973 | 860 | 0.5697 | - | - |
| 0.2019 | 880 | 0.6652 | - | - |
| 0.2065 | 900 | 0.6085 | - | - |
| 0.2111 | 920 | 0.6082 | - | - |
| 0.2156 | 940 | 0.6263 | - | - |
| 0.2202 | 960 | 0.5992 | - | - |
| 0.2248 | 980 | 0.5950 | - | - |
| 0.2294 | 1000 | 0.5778 | 0.3149 | 0.9453 |
| 0.2340 | 1020 | 0.5890 | - | - |
| 0.2386 | 1040 | 0.6056 | - | - |
| 0.2432 | 1060 | 0.6154 | - | - |
| 0.2478 | 1080 | 0.6102 | - | - |
| 0.2524 | 1100 | 0.5502 | - | - |
| 0.2569 | 1120 | 0.6006 | - | - |
| 0.2615 | 1140 | 0.6149 | - | - |
| 0.2661 | 1160 | 0.5941 | - | - |
| 0.2707 | 1180 | 0.6153 | - | - |
| 0.2753 | 1200 | 0.5927 | - | - |
| 0.2799 | 1220 | 0.5759 | - | - |
| 0.2845 | 1240 | 0.5894 | - | - |
| 0.2868 | 1250 | - | 0.3132 | 0.9456 |
| 0.2891 | 1260 | 0.5852 | - | - |
| 0.2936 | 1280 | 0.5578 | - | - |
| 0.2982 | 1300 | 0.5786 | - | - |
| 0.3028 | 1320 | 0.5501 | - | - |
| 0.3074 | 1340 | 0.6244 | - | - |
| 0.3120 | 1360 | 0.6009 | - | - |
| 0.3166 | 1380 | 0.5955 | - | - |
| 0.3212 | 1400 | 0.5901 | - | - |
| 0.3258 | 1420 | 0.6017 | - | - |
| 0.3304 | 1440 | 0.6022 | - | - |
| 0.3349 | 1460 | 0.5904 | - | - |
| 0.3395 | 1480 | 0.6046 | - | - |
| 0.3441 | 1500 | 0.5862 | 0.3131 | 0.9459 |
@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",
}
@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}
}
@misc{oord2019representationlearningcontrastivepredictive,
title={Representation Learning with Contrastive Predictive Coding},
author={Aaron van den Oord and Yazhe Li and Oriol Vinyals},
year={2019},
eprint={1807.03748},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/1807.03748},
}
Base model
google/embeddinggemma-2