Scaling and Stabilizing Large-Scale Embedding-Based Retrieval

2026-07-10

AI Summary

This work addresses the challenges of insufficient training signal quality and performance instability in large-scale embedding retrieval, particularly the degradation caused by the discrepancy between candidate set sizes during training and inference. To mitigate these issues, the authors propose a unified training framework that integrates hybrid hard negative mining—combining cross-batch online sampling with offline cross-encoder reranking enhanced by meta-heuristic filtering—and a warm-start knowledge distillation strategy tailored for legacy models. This approach ensures behavioral consistency while enabling smooth model evolution. Empirically, the system successfully upgraded its encoder from DistilBERT to GTE-base, achieving a 7.34% improvement in offline NDCG@5 and a 0.50% increase in online total revenue, as validated through large-scale experiments and A/B testing.

Abstract

Embedding-based retrieval (EBR) is foundational to large-scale e-commerce search, yet its effectiveness is often constrained by the quality of training signals and the representational capacity of the encoder. Standard dual-encoders suffer from a training-inference gap: they are optimized on narrow candidate pools but must discriminate against hundreds of millions of items during inference. Furthermore, while transitioning to higher-capacity backbones can mitigate this gap, simply replacing a mature model can lead to inconsistent retrieval behavior and a loss of the domain-specific knowledge established in previous iterations. In this paper, we present a unified pipeline deployed at Walmart that addresses both signal quality and model evolution. Our contributions are two-fold: (1) Hybrid Hard Negative Mining: We integrate Online Cross-Batch Sampling to increase negative diversity by an order of magnitude and Hybrid Offline Mining, which combines cross-encoder predictions with metadata heuristics to identify nuanced mismatches. (2) Legacy-Aware Distillation: We transition from DistilBERT to a higher-capacity GTE-base encoder. To ensure a smooth and superior transition, we introduce a Warm-Start Distillation technique that transfers domain-specific expertise from the legacy model to the new backbone. Validated through extensive offline experiments and online A/B testing, the proposed pipeline is deployed in live production, delivering a +7.34% improvement in NDCG@5 and a +0.50% lift in gross revenue.

Problem

• embedding-based retrieval
• training-inference gap
• model evolution
• domain-specific knowledge
• large-scale retrieval

Innovation

• Hybrid Hard Negative Mining
• Legacy-Aware Distillation
• Embedding-Based Retrieval
• Warm-Start Distillation
• Cross-Batch Sampling