Research map / Label-efficient and robust learning / Knowledge distillation
USTAD: Unified Single-model Training Achieving Diverse Scores for Information Retrieval
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains a single BERT-base Transformer to serve as cross-encoder reranker, dual-encoder retriever and late-interaction scorer, and proposes EmbedDistill (embedding-matching distillation with asymmetric student DE that inherits the teacher's frozen document encoder).
Problem
Modern IR pipelines use separate Transformer models for retrieval (dual-encoder) and reranking (cross-encoder), increasing complexity; cross-encoders also give poor embeddings and standard score-based distillation ignores teacher embeddings, leaving a teacher-student generalization gap.
Approach
A single Transformer is trained jointly in cross-encoder mode (concatenated query-document input with dual pooling) and dual-encoder mode (separate inputs with empty padding), with score-matching between modes; the teacher's embeddings are then used for embedding-matching distillation to students, including asymmetric students that freeze the teacher's document encoder and train only a small query
Key results
- Asymmetric student with 11.3M trainable parameters (BERT-mini query encoder, inherited teacher document embeddings) reaches about 95-97% of teacher performance, per the paper's claim.
- On NQ test relaxed recall, USTAD CE mode reaches R@5 79.9 and R@20 87.3, competitive with YONO (79.1/86.7) and AR2-d (81.5 R@5).
- On BEIR, a 6-layer DistilBERT student trained with EmbedDistill gets average nDCG@10 44.0, against 44.5 for Jina and 46.6 for Contriever (12 layers).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add score-based distillation to direct student training (ablation) | DistilBERT student · NQ dev reranking | R@1: 39.5 -> 42.4 | Helped |
| inherit teacher document embeddings (ablation) | DistilBERT student with score-based distillation · NQ dev reranking | R@1: 42.4 -> 42.8 | Helped |
| add query embedding matching (ablation) | DistilBERT student with inherited document embeddings · NQ dev reranking | R@1: 42.8 -> 43.4 | Helped |
| add generated queries (ablation) | DistilBERT student with query embedding matching · NQ dev reranking | R@1: 43.4 -> 43.8 | Helped |
| add score-based distillation to direct student training (ablation) | DistilBERT student · MSMARCO dev reranking | MRR@10: 27.0 -> 33.2 | Helped |
| inherit teacher document embeddings (ablation) | DistilBERT student with score-based distillation · MSMARCO dev reranking | MRR@10: 33.2 -> 35.4 | Helped |
| add query embedding matching (ablation) | DistilBERT student with inherited document embeddings · MSMARCO dev reranking | MRR@10: 35.4 -> 36.1 | Helped |
| add generated queries (ablation) | DistilBERT student with query embedding matching · MSMARCO dev reranking | MRR@10: 36.1 -> 36.3 | Helped |
| add score-based distillation to direct student training (ablation) | DistilBERT student · NQ dev retrieval | R@5: 36.2 -> 65.3 | Helped |
| inherit teacher document embeddings (ablation) | DistilBERT student with score-based distillation · NQ dev retrieval | R@5: 65.3 -> 69.9 | Helped |
| add query embedding matching (ablation) | DistilBERT student with inherited document embeddings · NQ dev retrieval | R@5: 69.9 -> 72.7 | Helped |
| add generated queries (ablation) | DistilBERT student with query embedding matching · NQ dev retrieval | R@5: 72.7 -> 73.4 | Helped |
Surprising findings
- Standard CE models produce degenerate query/document embeddings, so a CE model cannot directly act as a retriever without the proposed input design and dual-pooling.
Limitations
- Late-interaction mode is only partially evaluated (Table 1 ablation) and not fully integrated.
- Teacher and student results largely reported on NQ and MS MARCO dev sets; MS MARCO has no public test set.
- Gains from query generation are inconsistent across settings.
Experimental setup
Datasets: Natural Questions MS MARCO BEIR
Benchmarks: Natural Questions MS MARCO BEIR
Metrics: MRR@10 nDCG@10 recall@k strict recall relaxed recall
Baselines: AR2-g AR2-d YONO R2-D2 RocketQA ACNE DPR SentenceBERT Contriever TAS-B GenQ ANCE
Builds on: ColBERT DPR EmbedDistill-style score distillation ListMLE
Closest papers by meaning
- Hierarchical Retrieval: The Geometry and a Pretrain-Finetune Recipe (NeurIPS 2025 · similarity 0.87)
- U-MARVEL: Unveiling Key Factors for Universal Multimodal Retrieval via Embedding Learning with MLLMs (ICLR 2026 · similarity 0.86)
- A Statistical Framework for Data-dependent Retrieval-Augmented Models (ICML 2024 · similarity 0.86)
- Adaptive Retrieval and Scalable Indexing for k-NN Search with Cross-Encoders (ICLR 2024 · similarity 0.85)
- NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models (ICLR 2025 · similarity 0.85)
- MUVERA: Multi-Vector Retrieval via Fixed Dimensional Encoding (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.