Optimizing Retrieval Components for a Shared Backbone via Component-Wise Multi-Stage Training

TL;DR

The paper addresses optimizing a shared dense retrieval backbone for multiple legal information retrieval tasks by introducing a unified multi-stage training framework with progressively harder supervision (Stage 1 semantic coverage, Stage 2 hard negative mining, Stage 3 robustness calibration) applied to both embeddings and rerankers. It demonstrates that embedding and reranking components exhibit stage-dependent trade-offs, and that a mixed-stage end-to-end configuration can outperform a single shared baseline in RAG scenarios. An offline A/B test and production deployment validate practical gains with only modest latency overhead, supporting deployment in a shared service. The work highlights the value of component-wise, stage-aware model selection to balance coverage, precision, and system latency in real-world legal retrieval systems.

Abstract

Recent advances in embedding-based retrieval have enabled dense retrievers to serve as core infrastructure in many industrial systems, where a single retrieval backbone is often shared across multiple downstream applications. In such settings, retrieval quality directly constrains system performance and extensibility, while coupling model selection, deployment, and rollback decisions across applications. In this paper, we present empirical findings and a system-level solution for optimizing retrieval components deployed as a shared backbone in production legal retrieval systems. We adopt a multi-stage optimization framework for dense retrievers and rerankers, and show that different retrieval components exhibit stage-dependent trade-offs. These observations motivate a component-wise, mixed-stage configuration rather than relying on a single uniformly optimal checkpoint. The resulting backbone is validated through end-to-end evaluation and deployed as a shared retrieval service supporting multiple industrial applications.

Figures (3)

• Figure 1: Overview of a shared retrieval infrastructure in legal AI systems, highlighting the retrieval backbone optimized in this work.
• Figure 2: Multi-stage training pipeline for exploring progressive refinement of a unified retrieval model.
• Figure 3: Embedding recall as a function of retrieval budget across different training stages.