HyperAIHyperAI

Command Palette

Search for a command to run...

5 days ago
Model Training

Score-Only Distillation for Compact Dense Retrieval

Kirill Dubovikov Martin Takáč Salem Lahlou

Abstract

Large embedding models improve retrieval quality, but serving large encoders online is expensive. We study whether a compact retriever can learn teacher ranking behavior from score vectors without access to teacher hidden states. The student trains on rows built from ground-truth positives and negative candidates produced by our data generation pipeline; we evaluate student-teacher hardnegative mining separately as an extension. We use a row-centered score-vector objective, a memory-efficient implementation of uniform all-pairs PairMSE loss. On a fixed eight-task evaluation panel, our distillation protocol recovers up to 50% of the base-to-teacher gap. The distilled 0.6B student is 4.7× faster for query encoding and 9.7× faster for document encoding than sequential online teacher fusion. External-transfer performance after distillation remains mixed, so our evidence supports compression of teacher rankings under matched retrieval protocols.

One-sentence Summary

Researchers from Mohamed bin Zayed University of Artificial Intelligence propose a score-only distillation method that trains a compact dense retriever using a row-centered score-vector objective (uniform all-pairs PairMSE loss) without teacher hidden states, recovering up to 50%50\%50% of the base-to-teacher gap and yielding a 0.6B0.6\text{B}0.6B student 4.7×4.7\times4.7× faster for query encoding and 9.7×9.7\times9.7× faster for document encoding, compressing teacher rankings.

Key Contributions

  • The paper defines a black-box score-vector distillation protocol that trains compact retrieval students using only teacher scores, without access to teacher hidden states or representations.
  • This work proposes a row-centered implementation of uniform all-pairs PairMSE, a memory-efficient objective, and compares it against label-only and positive-negative MarginMSE controls under identical rows, student, and evaluation settings.
  • Results characterize the scope of the protocol: score-vector distillation improves compact students under matched retrieval protocols, and the paper provides empirical analysis of extensions including teacher fusion, active learning, and hard negative mining.

Introduction

Dense retrieval pipelines increasingly depend on large embedding models, but their size inflates serving cost, index refresh cost, and operational complexity, motivating the compression of these models into compact retrievers. Prior distillation methods for retrieval typically require access to teacher embeddings, hidden states, or logits, or they assume cross-encoder teachers that jointly encode query-document pairs. These approaches are not applicable when the teacher is a black-box embedding model whose only output is query-document score vectors. The authors propose a score-vector distillation protocol that trains a compact bi-encoder student solely from teacher-supplied scores over candidate rows. They introduce a row-centered reformulation of uniform all-pairs PairMSE loss that avoids explicit pairwise margin tensors and demonstrate that the protocol improves compact students across two architectures and an eight-task retrieval panel.

Method

The authors formulate score-vector distillation as a row-wise training procedure. For each query qqq, a candidate row contains a set of documents d0,,dk1d_0, \dots, d_{k-1}d0,,dk1, where d0d_0d0 is the labeled positive and the remaining k1k-1k1 documents are negatives. The student retriever assigns scores si=fθ(q,di)s_i = f_\theta(q, d_i)si=fθ(q,di), while a teacher model provides target scores tit_iti for the same query-document pairs. The teacher targets are normalized cosine scores, and the student and teacher need not share a common embedding space.

When multiple teachers are available, the authors apply equal fusion: the teacher score vectors are averaged coordinate-wise to produce a single target vector ttt. This fusion changes only the target scores, leaving the candidate rows, student architecture, and downstream evaluation unchanged.

The core training objective is the row-centered score vector distillation loss, designed to align student scores with teacher scores while discarding row-level additive offsets. Let r=stRkr = s - t \in \mathbb{R}^kr=stRk be the residual vector. The loss is defined as

Lcenter(s,t)=1krmean(r)1k22.\mathcal{L}_{\text{center}}(s, t) = \frac{1}{k} \| r - \operatorname{mean}(r) \mathbf{1}_k \|_2^2.Lcenter(s,t)=k1rmean(r)1k22.

This loss is invariant to any constant shift added to all scores in a row, focusing learning on the relative score differences that determine document ranking.

Using the variance identity i,j(rirj)2=2ki(rimean(r))2\sum_{i,j} (r_i - r_j)^2 = 2k \sum_i (r_i - \operatorname{mean}(r))^2i,j(rirj)2=2ki(rimean(r))2 and noting that (sisj)(titj)=rirj(s_i - s_j) - (t_i - t_j) = r_i - r_j(sisj)(titj)=rirj, the centered loss is equivalent to a uniform all-pairs PairMSE over ordered document pairs:

Lallpairs(s,t)=1k2i=0k1j=0k1[(sisj)(titj)]2=2Lcenter(s,t).\mathcal{L}_{\text{allpairs}}(s, t) = \frac{1}{k^2} \sum_{i=0}^{k-1} \sum_{j=0}^{k-1} \left[ (s_i - s_j) - (t_i - t_j) \right]^2 = 2 \mathcal{L}_{\text{center}}(s, t).Lallpairs(s,t)=k21i=0k1j=0k1[(sisj)(titj)]2=2Lcenter(s,t).

Thus, optimizing the centered objective implicitly enforces the same pairwise teacher preferences as all-pairs PairMSE, but avoids constructing explicit pair tensors. The two losses differ only by a constant factor for a fixed row width.

The authors compare this formulation against two control objectives. The first, Positive-Negative MarginMSE, matches only the margins between the labeled positive and each negative:

Lmargin=1k1j=1k1[(s0sj)(t0tj)]2.\mathcal{L}_{\text{margin}} = \frac{1}{k-1} \sum_{j=1}^{k-1} \left[ (s_0 - s_j) - (t_0 - t_j) \right]^2.Lmargin=k11j=1k1[(s0sj)(t0tj)]2.

The second is a label-only DPR-style contrastive objective that ignores all teacher scores. It optimizes the labeled positive against the other documents in the row using a fixed logit scale τ=0.05\tau = 0.05τ=0.05:

LCE=logexp(s0/τ)i=0k1exp(si/τ).\mathcal{L}_{\mathrm{CE}} = -\log \frac{\exp(s_0 / \tau)}{\sum_{i=0}^{k-1} \exp(s_i / \tau)}.LCE=logi=0k1exp(si/τ)exp(s0/τ).

This control tests whether the candidate rows and labeled positives alone are sufficient without teacher score supervision. The entire distillation pipeline operates on fixed candidate rows, with the loss function driving the student to reproduce the teacher's relative scoring patterns.

Experiment

The evaluation uses BEIR snapshots with training rows from three datasets and full-corpus retrieval on both row-source and held-out datasets, comparing student models trained with teacher score vectors against label-only and margin-based baselines. Score-vector distillation consistently improves compact retrievers beyond label-only training, and matching all candidate pairs in the row captures more teacher signal than margin-based losses. However, gains are domain-specific: the student recovers about one-quarter to one-half of the teacher gap and transfers poorly to unrelated corpora, while offering significant inference speed and memory improvements. Tested extensions including hard negative mining and active learning did not produce reliable further benefits.

Distilling teacher score vectors into a compact student retriever improves full-corpus retrieval quality over the frozen base, with the E5 student gaining more than the Qwen student. The label-only control on the same training rows performs worse than the base, confirming that teacher-provided score information is essential. Equal fusion of teacher scores does not consistently outperform single-teacher targets for Qwen, while it achieves the strongest result for E5. Score-vector distillation lifts both student models above their bases; the E5 student improves by roughly 0.07 NDCG@10 and the Qwen student by about 0.02, while the label-only control falls below the frozen base. For Qwen, equal fusion of teacher scores yields an NDCG@10 of 0.472 with a confidence interval that overlaps single-teacher estimates, so no clear advantage emerges; for E5, equal fusion is reported as the strongest setting.

Ablation experiments demonstrate that matching the full centered teacher-score vector (centered MSE) yields stronger retrieval quality than using only positive-negative margin loss, unit-normalized direction matching, or rank-weighted pairs. Removing teacher scores entirely with label-only cross-entropy drops performance sharply. Hard negative mining that selects negatives where teacher and student disagree yields a slight single-seed improvement, but its multi-seed interval overlaps the equal-fusion baseline, so a robust advantage is not established. Centered MSE outperforms positive-negative MarginMSE, unit-centered MSE, and rank-weighted PairMSE, indicating that the relative ordering among all negatives in the row carries useful teacher signal. Label-only cross-entropy without teacher scores reduces full-corpus retrieval quality dramatically compared to using teacher score vectors. Hard negative mining with a compact 6-candidate row achieves the best single-seed score, but its three-seed interval overlaps the equal-fusion 32-candidate baseline, so no stable ordering emerges.

Distilling teacher score vectors into a compact student retriever improves full-corpus retrieval over a frozen base, with the E5 model gaining more than Qwen, while a label-only control performs worse than the base, confirming that teacher-provided relative ordering of negatives is essential. Ablation shows that matching the full centered teacher-score vector (centered MSE) yields stronger retrieval quality than margin-based, direction-based, or rank-weighted losses, and dropping teacher scores entirely causes a sharp performance drop. Equal fusion of multiple teacher scores achieves the strongest result for E5 but offers no consistent advantage for Qwen, and hard negative mining does not provide a robust improvement over the baseline.


Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing

HyperAI Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp