Compute Optimal Embedding Eval

compute-optimal-embedding-eval

Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe — Ziarko et al. (2024) (arXiv:2406.04165, 2024)

What this evaluates

Evaluates the compute-optimal fine-tuning recipe for repurposing decoder-only LLMs into text embedding models. It measures how different computational budgets and fine-tuning methods affect both training contrastive loss and downstream retrieval/similarity performance.

Datasets

• BAAI BGE — total 200000000; splits: train (200000000)
• MTEB — total ?; splits: test (-1)

Metrics

• contrastive loss (primary) — range: other
  • The final training loss from the contrastive objective computed over (query, value) pairs.
• MTEB downstream score — range: [0, 1]
  • Average performance across 8 selected tasks (one per category) from the MTEB benchmark, using standard task-specific metrics.

Input / output format

Input: Semantically related (query, value) text pairs.

Output: Dense vector embeddings for each text, used to compute contrastive loss and downstream task scores.

Scoring recipe

model = load_pythia_model(size)
optimizer = AdamW(model)
scheduler = CosineLR(optimizer, warmup_frac=0.1)
for batch in BAAI_BGE_loader(batch_size=1024, ctx_len=75):
    loss = contrastive_loss(model(batch.query), model(batch.value))
    loss.backward()
    optimizer.step()
final_loss = loss.item()
mteb_scores = []
for task in MTEB_subset[8_tasks]:
    preds = model.encode(task.test_queries)
    gold = task.test_labels
    score = compute_mteb_metric(preds, gold)
    mteb_scores.append(score)
avg_mteb = sum(mteb_scores) / len(mteb_scores)

Common pitfalls

• Training for more than one epoch on BAAI BGE causes diminishing returns; the protocol explicitly uses <1 epoch.
• LoRA rank selection is highly sensitive; ranks of 8 or 2048 are suboptimal, while 32 or 128 yield the best loss.
• In block freezing, token-embedding parameters are always frozen, unlike standard transformer block freezing.

Evidence (verbatim from paper)

In addition to controlling the final training contrastive loss achieved by the models, we also measure downstream performance by evaluating the models on a representative subset of the MTEB benchmark*(Muennighoff et al., [2023b])*. The benchmark defines eight categories of tasks in total (e.g., retrieval, semantic text similarity). We select one task from each category by determining which ones are the most correlated with the performance for the whole category (see Appendix[B] for details).

Citation

@misc{ziarko2024repurposing,
  title={Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe},
  author={Ziarko et al. (2024)},
  year={2024},
  note={arXiv:2406.04165}
}

• arXiv: 2406.04165