Source-linked AI summary

Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillation

Sebastian Hofstätter, Sophia Althammer, Michael Schröder, Mete Sertkan, Allan Hanbury

arXiv:2010.02666v2cs.IR

TL;DR

Efficient neural ranking models offer lower query latency but face an effectiveness gap, and differing score ranges make direct cross-architecture distillation difficult. The paper adapts distillation with Margin-MSE, which transfers teacher margins to efficient architectures; it improves re-ranking and dense BERTDOT retrieval while preserving efficiency, with teacher ensembles usually helping further.

  • Problem

    Efficient ranking architectures reduce query-time cost but may sacrifice effectiveness, while differing architecture-specific score ranges complicate knowledge distillation.

  • Method

    The paper distills BERTCAT teacher margins over query, relevant-passage, and non-relevant-passage triples into efficient students using a model-agnostic Margin-MSE procedure.

  • Results

    Margin-MSE improves re-ranking effectiveness across evaluated efficient architectures without compromising efficiency, while also improving dense BERTDOT retrieval; teacher ensembles usually provide further gains.

  • Takeaways & Limitations

    Cross-architecture margin distillation can narrow the efficiency-effectiveness trade-off and produce competitive dense retrieval without specialized, more costly training methods.

  • Takeaways & Limitations

    The approach retains rare negative-margin samples that can reverse pair orderings, and ColBERT dense-retrieval experiments were limited by substantially higher index storage requirements.

Abstract

from arXiv · show

Retrieval and ranking models are the backbone of many applications such as web search, open domain QA, or text-based recommender systems. The latency of neural ranking models at query time is largely dependent on the architecture and deliberate choices by their designers to trade-off effectiveness for higher efficiency. This focus on low query latency of a rising number of efficient ranking architectures make them feasible for production deployment. In machine learning an increasingly common approach to close the effectiveness gap of more efficient models is to apply knowledge distillation from a large teacher model to a smaller student model. We find that different ranking architectures tend to produce output scores in different magnitudes. Based on this finding, we propose a cross-architecture training procedure with a margin focused loss (Margin-MSE), that adapts knowledge distillation to the varying score output distributions of different BERT and non-BERT passage ranking architectures. We apply the teachable information as additional fine-grained labels to existing training triples of the MSMARCO-Passage collection. We evaluate our procedure of distilling knowledge from state-of-the-art concatenated BERT models to four different efficient architectures (TK, ColBERT, PreTT, and a BERT CLS dot product model). We show that across our evaluated architectures our Margin-MSE knowledge distillation significantly improves re-ranking effectiveness without compromising their efficiency. Additionally, we show our general distillation method to improve nearest neighbor based index retrieval with the BERT dot product model, offering competitive results with specialized and much more costly training methods. To benefit the community, we publish the teacher-score training files in a ready-to-use package.

1 INTRODUCTION

Efficient neural ranking architectures reduce query-time cost but can sacrifice effectiveness, motivating cross-architecture knowledge distillation from BERTCAT teachers. The paper proposes Margin-MSE and reports improved effectiveness for efficient models, including smaller effectiveness gaps and competitive dense retrieval.

  • Motivation: Efficient neural ranking architectures transfer computation to indexing or use reduced interactions to lower query latency, trading off effectiveness.BERTCAT requires query-passage concatenation to pass through all Transformer layers at query time, while TK, ColBERT, PreTT, and BERTDOT make efficiency-oriented architectural choices.
  • Motivation: Different ranking architectures converge to distinct output-score ranges, complicating direct cross-architecture knowledge distillation.BERTCAT tends toward positive relevant-document scores and negative non-relevant scores, TK produces negative averages, and BERTDOT and ColBERT produce high dot-product scores.
  • Approach: Margin-MSE distills the teacher’s relevant-versus-non-relevant score margin rather than its raw score range, allowing students to retain architecture-specific scoring ranges.The method uses query, relevant-passage, and non-relevant-passage triples and optimizes the student margin toward the teacher margin with mean squared error.
  • Findings: Teacher ensembles generally improve student re-ranking effectiveness beyond single-teacher distillation across retrieval metrics.The ensemble averages scores from BERT-Base, whole-word-masked BERT-Large, and ALBERT-large BERTCAT teachers.
  • Findings: For dense nearest-neighbor retrieval, teacher distillation improves BERTDOT effectiveness and yields competitive results without dense-retrieval-specific training methods.The comparison is against more costly approaches such as index-based passage sampling and in-batch negatives.
  • Findings: Distilled efficient architectures outperform BERTCAT baselines on several metrics, while BERTDOT and TK have significantly smaller effectiveness gaps.The reported trade-off results indicate that PreTT and ColBERT no longer require the stated effectiveness compromise.

2 RETRIEVAL MODELS

The section compares neural ranking architectures that trade effectiveness, latency, computation, and storage differently. BERTCAT is the strongest but query-expensive model, while BERTDOT, ColBERT, PreTT, and TK shift computation or representation design to improve efficiency.

  • BERTCAT: BERTCAT concatenates each query and passage, processes them jointly through BERT, and scores the resulting CLS representation with a linear layer.Because concatenation requires the Transformer layers to process each query-passage pair at query time, BERTCAT has substantial query-time computation.
  • BERTDOT: BERTDOT independently encodes queries and passages, pre-computes passage representations, and scores query-passage pairs with a dot product.This design is highly efficient in storage and query latency but compresses information more strongly than BERTCAT, reducing effectiveness.
  • Comparison: Table 1 compares baseline nDCG@10, indexing capability, latency, memory, and model characteristics, while the evaluation measures latency for one query and 1,000 documents.The summary characterizes BERTCAT as prohibitive for efficient production use, BERTDOT as the most efficient BERT-based model, and TK as cheaper but less effective overall.
  • ColBERT: ColBERT delays query-document interaction until after BERT encoding, using max-pooling over query-term interactions and summing the resulting scores.Its aggregation requires n*m dot products and is roughly as efficient as BERTDOT, but pre-computed passage representations have much higher storage costs.
  • PreTT: PreTT lies between BERTCAT and ColBERT by separately processing query and passage layers before concatenating them and computing the remaining layers.The evaluated implementation omits optional representation compression for comparability and uses concatenation after three of six layers.
  • TK: TK independently contextualizes query and passage terms with shallow Transformers, cosine-similarity matching, Gaussian kernels, and kernel-pooling aggregation.Kernel-pooling acts as a soft histogram over similarity ranges before the final weighted score is computed.

3 CROSS-ARCHITECTURE KNOWLEDGE DISTILLATION

The paper addresses the mismatch between ranking architectures’ score ranges by distilling teacher information through score margins rather than absolute scores. Margin-MSE trains efficient students on teacher-derived margins from the same query, relevant passage, and non-relevant passage triples.

  • Motivation: Neural ranking triples provide only binary relevance labels, although re-ranking requires more nuanced distinctions among many candidate passages.The paper uses BERTCAT’s stronger generalization capabilities as a source of additional training information for other architectures.
  • Motivation: Different ranking architectures learn distinct output-score ranges, making absolute-score distillation unsuitable across architectures.BERTCAT tends toward positive relevant-document scores and negative non-relevant scores, while TK, BERTDOT, and ColBERT exhibit different average ranges.
  • Margin-MSE: Margin-MSE trains students to match the teacher’s score margin between a relevant and non-relevant passage for each query.The loss uses triples of queries, relevant passages, and non-relevant passages, with the teacher margin serving as the student’s label.
  • Margin-MSE: Margin-MSE discards the original binary relevance information because teacher margins can be negative and teacher pairwise ranking accuracy exceeds 98% during training.The authors therefore regard adding the binary information to the ranking loss as redundant.
  • Training procedure: The staged procedure trains BERTCAT with collection labels, stores its scores for all training data, and independently trains each student with Margin-MSE.Reusing the same training triples enables the teacher scores to be stored once for experimentation and sharing.

4 EXPERIMENT DESIGN

The experiments use MSMARCO-Passage and specified optimization, stopping, batching, and initialization choices to evaluate the ranking models. The implementation combines PyTorch and HuggingFace Transformers for ranking with Anserini for first-stage retrieval.

  • Implementation: PyTorch and HuggingFace Transformers support re-ranking training and inference, while Anserini supports first-stage indexing and retrieval.These tools define the software stack used for the reported experiments.
  • Datasets: The study uses MSMARCO-Passage, including 49,000 sparsely judged development queries and 43 densely judged TREC-DL’19 queries.The collection contains 8.8 million passages and a proposed training set of 40 million triples.
  • Datasets: TREC graded relevance labels are binarized at 2 for MRR and MAP, and teachers are evaluated on the full training set.The evaluation uses both the MSMARCO-DEV query set and the TREC-DL’19-derived query set.
  • Optimization: All models use Adam, with a learning rate of 7 *10^-6 for BERT layers and 10^-5 for TK, plus early stopping based on validation nDCG@10.Training batches contain 32 examples.
  • Model setup: Student language models start from six-layer DistilBERT, selected because it provides a close lower bound at half the runtime.The implementation repeats ColBERT’s query MASK augmentation eight times and concatenates PreTT sequences after three layers.

5 RESULTS

The experiments show that Margin-MSE improves cross-architecture distillation across efficient ranking models, with teacher ensembles generally outperforming single teachers. Distillation also improves dense retrieval and shifts the efficiency-effectiveness trade-off without changing model latency, although dense retrieval introduces practical trade-offs.

  • 5.1 Optimization Study: Margin-MSE outperforms pointwise MSE and weighted RankNet on all metrics across three ranking architectures, while distillation generally improves each model over its original baseline.The ablation uses a single BERT-BaseCAT teacher on MSMARCO-DEV.
  • 5.2 Knowledge Distillation Results: Teacher ensembles generally improve student effectiveness beyond single-teacher training, with the ensemble preferred overall despite exceptions for TK and PreTT on TREC-DL’19.The ensemble improves all models on sparse MSMARCO-DEV labels, while BERT-BaseDOT is an outlier under single-teacher training.
  • 5.2 Knowledge Distillation Results: DistilBERTDOT consistently gains from single- and ensemble-teacher training, whereas BERT-BaseDOT shows little improvement with the single teacher but substantial improvement with the ensemble.BERT-Base has a slight advantage over DistilBERT when trained on T2, but DistilBERTDOT exhibits more consistent gains.
  • 5.2 Knowledge Distillation Results: Dense retrieval removes BM25 reranking, making the pipeline more efficient and less complex but increasing false-positive risk and reducing interpretability.ColBERT dense retrieval was not extensively evaluated because its storage requirement rises from 2GB of plain text to a 2TB index.
  • 5.2 Knowledge Distillation Results: Dense BERTDOT retrieval with teacher ensembles is more effective than single-teacher or no-teacher training and remains competitive with specialized, more costly dense-retrieval methods.The comparison uses nearest-neighbor retrieval with a flat Faiss index; the approach is slightly disadvantaged on MSMARCO DEV MRR@10 but outperforms models that also report TREC-DL’19 results.
  • 5.3 Closing the Efficiency-Effectiveness Gap: The efficiency-effectiveness plots show that distilled DistilBERTCAT, PreTT, and ColBERT close the gap to BERT-BaseCAT and can surpass its single-instance results on TREC-DL’19.Latency is measured as median neural-model inference time without preprocessing or disk access, and teacher training changes weights rather than architecture.

6 TEACHER ANALYSIS

Teacher analysis examines score-margin distributions, per-query nDCG changes, and the use of teacher ensembles. The analysis finds diverse teacher margins and mixed but generally improving per-query effects from teacher training.

  • The same BERTCAT architecture produces three distinct score patterns across teachers, supporting diversity in the ensemble.
  • Negative margins can reverse pair orderings, but such swaps are rare, involve small score differences, and affect only a few percent of training data.
  • Figure 5 compares margins between relevant and non-relevant documents for three teacher models on MS MARCO-Passage training data.
  • Figure 6 compares per-query nDCG@10 changes between original training and T1 or T2 training on TREC-DL’19.
  • Across both models, T1 and T2 training improves roughly 60 % of queries and decreases results on 33 %; T2 changes are more extreme.

7 RELATED WORK

Related work covers efficiency improvements for neural relevance models, adapted training procedures, and knowledge distillation. Prior studies include smaller same-architecture models, weakly supervised student-teacher training, and distillation across several retrieval settings.

  • Efficient relevance models: Efficient relevance-model research includes kernel approximations with locality-sensitive hashing and static index pruning to reduce query-processing latency.
  • Adapted training procedures: Adapted training procedures address limited annotation through weak supervision, semi-supervised student-teacher learning, and multiple weak-supervision sources.
  • Knowledge distillation: Knowledge distillation transfers capabilities from large pretrained language models to smaller models while targeting little effectiveness loss and lower computational cost.
  • Knowledge distillation: TinyBERT and DistilBERT distill BERT while maintaining accuracy across a variety of natural language understanding tasks.
  • Knowledge distillation: Prior information-retrieval studies distill sequential recommenders, retrieval chatbots, different BERTCAT sizes, TinyBERT, and PARADE-style ranking models.

8 CONCLUSION

The paper applies cross-architecture distillation from BERTCAT to efficient passage-ranking models using margin-based supervision. Margin-MSE and teacher ensembles improve effectiveness while preserving efficiency, with competitive retrieval results beyond re-ranking.

  • Cross-architecture knowledge distillation improves query-latency-efficient neural passage ranking models taught by full-interaction BERTCAT.
  • Margin-MSE matches relevant–non-relevant passage margins instead of raw scores because architectures converge to different scoring ranges.
  • Margin-MSE outperforms a simple pointwise MSE loss, and teacher ensembles are beneficial in most passage-retrieval cases.
  • Teacher-ensemble training lets single efficient-model instances outperform their single-instance teacher models despite fewer parameters and less interaction capacity.
  • The method shifts the evaluated models’ effectiveness-efficiency trade-off toward greater effectiveness and produces competitive nearest-neighbor retrieval results.
Loading 2010.02666v2…