HyperAIHyperAI

Command Palette

Search for a command to run...

Distillation par scores uniquement pour la recherche dense compacte

Kirill Dubovikov Martin Takáč Salem Lahlou

Résumé

Les grands modèles d'embedding améliorent la qualité de la recherche, mais le déploiement de grands encodeurs en ligne est coûteux. Nous étudions si un récupérateur compact peut apprendre le comportement de classement d'un enseignant à partir de vecteurs de scores, sans accès aux états cachés de l'enseignant. L'étudiant s'entraîne sur des lignes construites à partir de positifs de référence et de candidats négatifs produits par notre pipeline de génération de données ; nous évaluons séparément, comme extension, l'extraction de négatifs difficiles entre étudiant et enseignant. Nous utilisons un objectif de vecteur de scores centré sur les lignes, une implémentation efficace en mémoire de la perte PairMSE uniforme sur toutes les paires. Sur un panel fixe de huit tâches d'évaluation, notre protocole de distillation récupère jusqu'à 50 % de l'écart entre le modèle de base et l'enseignant. L'étudiant distillé de 0,6 milliard de paramètres est 4,7 fois plus rapide pour l'encodage des requêtes et 9,7 fois plus rapide pour l'encodage des documents que la fusion séquentielle en ligne de l'enseignant. Les performances de transfert externe après distillation restent mitigées, donc nos résultats soutiennent la compression des classements de l'enseignant dans le cadre de protocoles de recherche appariés.

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.


Créer de l'IA avec l'IA

De l'idée au lancement — accélérez votre développement IA avec le co-codage IA gratuit, un environnement prêt à l'emploi et le meilleur prix pour les GPU.

Codage assisté par IA
GPU prêts à l’emploi
Tarifs les plus avantageux

HyperAI Newsletters

Abonnez-vous à nos dernières mises à jour
Nous vous enverrons les dernières mises à jour de la semaine dans votre boîte de réception à neuf heures chaque lundi matin
Propulsé par MailChimp