Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup
Paper • 2101.06983 • Published • 3
EmbeddingGemma-2 · AllNLI LoRA proof of concept
A ~60-second proof-of-concept fine-tune of google/embeddinggemma-2
(744M params, text tower used) on the pair subset of
sentence-transformers/all-nli
with a contrastive objective. Only the text backbone's attention + MLP projections carry adapters
(≈5.0M trainable parameters, r=16).
Training recipe
| Setting | Value |
|---|---|
| Loss | CachedMultipleNegativesRankingLoss (mini-batch 8) |
| Batch size | 64 (in-batch negatives, no duplicates sampler) |
| Steps | 9 (~60 s on an Nvidia T4, fp16) |
| Learning rate | 2e-5, 10% warmup |
| Sequence length | 128 |
| LoRA | r=16, alpha=32, dropout 0.1 on q/k/v/o/gate/up/down_proj of the text tower |
| Prompts | Model's built-in query / document prefixes |
The recipe follows the one used for sentence-transformers/embeddinggemma-300m-medical (full-scale training with this loss lifted nDCG@10 from 0.847 to 0.903; this POC runs 9 steps of it).
Results (what 60 seconds gets you)
- Train loss: 0.72 → ~0.68 mean over 9 steps (noisy at this scale — this is ~0.18% of one epoch)
- Eval loss (1,000 AllNLI dev pairs): 0.981 after 5 steps → 0.959 after 9 steps
- Retrieval sanity check after training: cos(anchor, positive) = 0.80 / 0.88 vs cos(anchor, negative) = 0.61 / 0.56 — anchors rank their true positives above unrelated sentences.
This is a pipeline proof of concept, not a trained-for-quality checkpoint: the run never finishes a warmup-then-converge cycle at 9 steps.
Usage
from peft import PeftModel
from sentence_transformers import SentenceTransformer
base = SentenceTransformer("google/embeddinggemma-2")
model = PeftModel.from_pretrained(base[0].auto_model, "ML-Intern-lab/embeddinggemma2-nli-lora-poc")
base[0].auto_model = model
embeddings = base.encode(["hello world", "a second sentence"])
The repo contains the standard adapter package only: adapter_config.json + adapter_model.safetensors (~19 MB, extracted from the run's full-model save; all 168 lora_B tensors verified non-zero, i.e. the trained weights).
References
- MultipleNegativesRankingLoss: Henderson et al., Efficient Natural Language Response Suggestion for Smart Reply (2017); Reimers & Gurevych, Sentence-BERT (2019)
- CachedMNRL: Gao et al., Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup (2021)
- Downloads last month
- 5
Model tree for ML-Intern-lab/embeddinggemma2-nli-lora-poc
Base model
google/embeddinggemma-2