Medium system-design machine-learning

Should you train your own text encoder, fine-tune one, or call an embedding API?

The 30 second answer

A store catalog is far too small to pre-train a large language encoder, so I start from a pre-trained one. I size it to where it runs: catalog items are embedded in batch, so a larger model is affordable, while queries are embedded per request, so latency pushes toward smaller. I fine-tune contrastively on my own pairs, or use LoRA adapters, and weigh hosted APIs against cost and deprecation risk.

Do not pre-train. Data needs grow with model size, and one store's text will not even teach which rare words are related. Start from an existing encoder. Sizes vary widely: the all-MiniLM-L6-v2 sentence-transformer has about 22.7 million parameters, image-text CLIP-style models are larger, and recent LLM-based embedding models are larger again.

Size by where it runs. Bigger encoders usually produce better vectors, but you pay on every embedding. Catalog items are embedded once in a batch job, so a big model can be affordable. Queries are embedded on every request, so latency and total compute push the query side smaller.

Fine-tune on your objective. Reuse the contrastive setup: carts and the product that completed them, trained with an in-batch-negative loss such as Sentence Transformers' MultipleNegativesRankingLoss. With a modest dataset, freeze most layers and train the top ones to limit catastrophic forgetting.

LoRA. Freeze all original weights and add a pair of small low-rank matrices beside selected weight matrices; their product is added to the frozen layer's output, and only they train. The LoRA paper reports up to 10,000 times fewer trainable parameters than full fine-tuning of GPT-3 175B. Many adapters can share one frozen base model, one per use case.

Hosted APIs. Good for prototypes and smaller systems. At large scale teams often self-host for cost, fine-tuning and control over deprecation, because a retired embedding model forces a full re-embed; OpenAI's deprecation page, for example, lists first-generation embedding models shut down on 4 January 2024 with text-embedding-3-small as the replacement.

Sources (5)
  1. Hugging Face: sentence-transformers/all-MiniLM-L6-v2 model card
  2. Sentence Transformers docs: losses (MultipleNegativesRankingLoss)
  3. Hu et al., LoRA: Low-Rank Adaptation of Large Language Models
  4. Hugging Face PEFT docs: LoRA
  5. OpenAI API docs: Deprecations