Building Embedding Models for Large-Scale Real-World Applications

Sahil Dua, Tech Lead at Google for Gemini embedding models, delivered a comprehensive technical presentation on building and deploying embedding models for large-scale real-world applications. The talk covered the complete lifecycle from architecture design to production deployment, with particular emphasis on practical considerations for search and RAG (Retrieval-Augmented Generation) applications.

The Critical Role of Embedding Models

Embedding models serve as the backbone for modern search and retrieval systems, converting various inputs—text, images, or videos—into numerical representations called embeddings. These embeddings capture the semantic meaning of inputs and enable efficient similarity search across massive datasets. Dua illustrated this with a simple example: when users search for "cute dogs," embedding models sift through billions of images to find the most relevant results by comparing vector representations.

The applications are broad and impactful:

• Document and media retrieval: Finding relevant content across vast corpora
• Personalized recommendations: Capturing user preferences to suggest relevant products or content
• RAG applications: Augmenting large language model responses with retrieved context to improve accuracy and reduce hallucinations
• Data deduplication: Identifying and removing redundant data during LLM training

Architecture Deep Dive

The presentation provided a detailed breakdown of embedding model architecture, which consists of several key components:

Tokenizer: Converts input strings into token IDs by breaking text into smaller units. This is the first step in processing any text input.

Embedding Projection: Maps token IDs to their corresponding vector representations using a large vocabulary table. This creates the initial numerical representation of each token.

Transformer: The core component that enriches token-level embeddings by incorporating context from surrounding tokens. This bidirectional attention mechanism allows the model to understand the full sequence.

Pooler: Combines token-level embeddings into a single representation for the entire sequence. Mean pooling is the most common approach, averaging all token embeddings to create a unified vector.

Output Projection Layer: An optional linear layer that transforms the pooled embedding to a specific dimension. This allows control over the final embedding size.

For multimodal applications, the architecture adapts: vision encoders replace tokenizers for images (breaking images into patches), and video processing treats videos as sequences of frames.

Training Techniques

The presentation emphasized contrastive learning as the primary training approach. The core principle is ensuring similar inputs have close embeddings while dissimilar inputs have distant embeddings. The training process involves:

Next Sentence Prediction: For supervised learning, splitting text into query-document pairs where the document is the subsequent sentence.

Span Corruption: For unsupervised learning, masking portions of text and training the model to recognize that corrupted versions of the same text should have similar embeddings.

Hard Negative Mining: Adding challenging negative examples that are semantically similar but contextually different (e.g., an Italian restaurant in New York vs. London) to improve model discrimination.

Distilling Large Models for Production

Given the computational constraints of production environments, the talk detailed techniques for distilling large language models into smaller, efficient embedding models:

Scoring Distillation: Using the teacher model's similarity scores to train the student model, ensuring the smaller model predicts similar relevance scores.

Embedding Distillation: Training the student model to produce embeddings similar to the teacher's embeddings, preserving semantic understanding.

Combined Approach: Using both scoring and embedding distillation for optimal results.

When designing student architectures, several considerations are crucial:

• Model depth and width: Balancing quality (deeper, wider models) against latency requirements
• Attention mechanism: Using multi-head attention for better quality despite higher memory usage, as memory constraints are less critical for smaller models
• Output dimension: Choosing appropriate embedding sizes (typically 64-256 dimensions) to balance quality, storage, and search costs

Evaluation Strategies

Evaluating embedding models presents unique challenges, particularly when golden labels are unavailable. The presentation outlined a comprehensive evaluation framework:

Standard Retrieval Evaluation: When golden labels exist, the process involves generating embeddings for queries and documents, computing top-K nearest neighbors, and measuring metrics like recall, NDCG, and mean reciprocal rank.

Auto-Rater Approach: When golden labels are unavailable, an auto-rater model (typically a large language model) scores the relevance of retrieved results. The position-weighted average score is particularly valuable as it accounts for the importance of ranking order.

Production Deployment Challenges

Serving embedding models at scale involves addressing several critical challenges:

Query Latency: Since query embedding generation occurs on the critical serving path, optimization techniques include:

• Server-side dynamic batching to improve GPU utilization
• Model quantization to reduce precision and improve speed
• Using smaller query models while maintaining larger document models for offline indexing

Document Indexing Cost: Processing billions of documents requires:

• Large batch sizes for efficient inference
• Parallel processing across multiple GPUs/TPUs
• Smaller embedding dimensions to reduce storage and search costs

Nearest Neighbor Search Latency: Finding relevant documents quickly involves:

• Using approximate algorithms instead of exact matching for scalability
• Choosing appropriate vector databases (e.g., Spanner k-NN on Google Cloud, OpenSearch on AWS)
• Incremental index updates to handle dynamic content

Practical Considerations for Model Selection

For teams without resources to train custom models, the presentation provided guidance on selecting off-the-shelf models:

• Intended use case: Ensure the model is trained for your specific domain (e.g., shopping, legal)
• Language coverage: Verify the model supports all required languages
• Training data: Check for domain-specific knowledge relevant to your application
• Model size and efficiency: Match the model's performance characteristics to your serving requirements
• Output dimension: Ensure the model provides embeddings of the appropriate size
• Licensing: Verify commercial usage rights
• Community support: Assess documentation and community resources
• Benchmark performance: Review quality metrics on relevant benchmarks

Key Takeaways

The presentation concluded with several actionable insights:

1. Embedding models are fundamental to modern search and RAG applications
2. Evaluation can be complex without golden labels, often requiring auto-rater models
3. Large models must be distilled for production efficiency
4. Production deployment requires managing both real-time serving and offline indexing
5. Model selection involves balancing multiple factors including quality, efficiency, and licensing

The comprehensive coverage of embedding model development and deployment provides a practical roadmap for teams looking to implement these technologies at scale, with particular emphasis on the real-world challenges and trade-offs involved in production systems.