Source-linked AI summary

D2LLM: Decomposed and Distilled Large Language Models for Semantic Search

Zihan Liao, Hang Yu, Jianguo Li, Jun Wang, Wei Zhang

arXiv:2406.17262v1cs.CL

TL;DR

Semantic search needs accurate query–passage matching without sacrificing efficient retrieval. D2LLM decomposes an LLM cross-encoder into a bi-encoder with attention pooling and an interaction module, and distills teacher knowledge into the student. It outperforms five leading methods across three benchmark tasks, while retaining the operational efficiency associated with bi-encoders.

  • Problem

    Bi-encoders enable efficient precomputed retrieval but can miss query–passage nuances, whereas LLM cross-encoders capture them but require fresh inference for each pairing.

    Source: §1 · pp. 1–2
  • Method

    D2LLM decomposes an LLM cross-encoder into a bi-encoder with Pooling by Multihead Attention and an Interaction Emulation Module, then distills teacher knowledge through contrastive, rank, and feature imitation.

    Source: §1 · p. 2
  • Results

    D2LLM outperforms five leading methods across three benchmark tasks, including a 6.45% lead over BGE on NLI.

    Source: §1 · p. 2
  • Takeaways & Limitations

    D2LLM combines the high accuracy associated with cross-encoders and the operational efficiency of bi-encoders for semantic search.

    Source: §5 · p. 9

    Training D2LLM demands substantial computational resources because it relies on large language models, which may be prohibitive for those with limited compute.

    Source: §7 · p. 9

Abstract

from arXiv · show

The key challenge in semantic search is to create models that are both accurate and efficient in pinpointing relevant sentences for queries. While BERT-style bi-encoders excel in efficiency with pre-computed embeddings, they often miss subtle nuances in search tasks. Conversely, GPT-style LLMs with cross-encoder designs capture these nuances but are computationally intensive, hindering real-time applications. In this paper, we present D2LLMs-Decomposed and Distilled LLMs for semantic search-that combines the best of both worlds. We decompose a cross-encoder into an efficient bi-encoder integrated with Pooling by Multihead Attention and an Interaction Emulation Module, achieving nuanced understanding and pre-computability. Knowledge from the LLM is distilled into this model using contrastive, rank, and feature imitation techniques. Our experiments show that D2LLM surpasses five leading baselines in terms of all metrics across three tasks, particularly improving NLI task performance by at least 6.45%. The source code is available at https://github.com/codefuse-ai/D2LLM.

1 Introduction

Semantic search must balance accurate query–passage matching with the efficiency needed for large-scale use. D2LLM decomposes an LLM cross-encoder into a precomputable student model and distills the teacher’s knowledge; the paper reports gains over five leading methods across three tasks.

1 Introduction

  • Semantic search aims to find passages matching a query by semantic links, resolving lexical mismatches across applications such as retrieval, question answering, and dialogue.The passage also names recommendation, fact checking, and retrieval-augmented generation.Source: §1 · p. 1
  • Bi-encoders enable precomputed passage vectors and efficient querying, but their constrained interaction can miss nuances and their training is data-intensive.Their performance may also generalize poorly to unseen domains.Source: §1 · p. 1
  • LLM cross-encoders jointly process queries and passages for granular understanding and zero-shot adaptation, but require fresh inference for each query–passage pair.This prevents caching passage vectors and can hinder real-time, high-volume processing.Source: §1 · p. 2
  • D2LLM decomposes an LLM cross-encoder into a bi-encoder with PMA and an Interaction Emulation Module, then distills teacher knowledge through contrastive, rank, and feature imitation.The bi-encoder separately embeds queries and passages so passage vectors can be stored in advance.Source: §1 · p. 2
  • D2LLM outperforms five leading methods across three benchmark tasks, with a 6.45% lead over the heavily finetuned BGE model on NLI.The introduction also reports a 14.39% improvement over the second-best LLaRA.Source: §1 · p. 2

2 Related Works

Prior semantic-search work includes standalone bi-encoders, cross-encoders, and methods distilling cross-encoder effectiveness into bi-encoders. D2LLM uses a static teacher, multiple imitation losses, and an interaction module to address limitations described for earlier approaches.

2 Related Works

  • Bi-encoders offer compact retrieval but can struggle with complex query–passage relationships, while cross-encoders model them more effectively at a computational cost that limits real-time use.The related work also notes that finetuning alone may not ensure adaptability and generalization.Source: §2 · p. 2
  • Earlier cross-encoder-to-bi-encoder distillation used pseudo-labeling or KL divergence, while AR2 and RocketQAv2 explored adversarial training and dynamic distillation.The paper identifies scalability, incomplete capture of LLM knowledge, and insufficient interaction modeling as challenges for applying these approaches to LLMs.Source: §2 · pp. 2–3
  • D2LLM refines the student using a static, task-specific teacher, contrastive, rank, and feature imitation losses, and an Interaction Emulation Module for query–passage interactions.The related-work passage frames this design as addressing scalability and interaction-modeling challenges in prior approaches.Source: §2 · p. 3
  • LLM bi-encoder methods commonly use contrastive training with a special passage token, an adaptation that may underuse LLM capabilities and underperform specialized bidirectional encoders.The passage distinguishes finetuning-only methods from approaches combining continued pretraining and finetuning.Source: §2 · p. 3

3 D2LLM

D2LLM decomposes an LLM cross-encoder into a bi-encoder student with PMA and IEM, then distills teacher knowledge through contrastive, rank, and feature imitation. The design retains precomputable passage vectors while modeling query–passage interactions.

3 D2LLM

  • D2LLM decomposes an LLM cross-encoder into a bi-encoder, PMA, and IEM student intended to retain efficiency while emulating teacher accuracy.The bi-encoder precomputes passage vectors while IEM models query–passage interplay.Source: §3 · p. 3, §3.1.2 · p. 4

3.1.1 Teacher Model T

  • The teacher uses prompted joint query–passage inputs and a classification-token embedding to score symmetric or asymmetric search matches.Prompts guide the LLM to produce a yes-or-no response for the pair.Source: §3.1.1 · p. 3, Figure 1 · p. 4

3.1.2 Student Model S

  • PMA applies attention pooling to query and passage token embeddings, using a learnable anchor to dynamically weight token information.The paper contrasts this flexible weighting with fixed weights in mean, min, or max pooling.Source: §3.1.2 · p. 4
  • IEM concatenates query and passage embeddings and uses separate MLP branches to encode prompt-specific interactions for symmetric and asymmetric search.The lightweight MLP also serves as a flexible similarity metric beyond cosine similarity.Source: §3.1.2 · p. 5

3.2 Distillation

  • Distillation trains the student with contrastive, rank, and feature imitation objectives to transfer different aspects of the teacher’s knowledge.The student is trained with these objectives alongside the PMA, IEM, linear layer, and bi-encoder.Source: §3.2 · p. 5, §3.2.4 · p. 6

3.2.1 Contrastive Imitation

  • Contrastive imitation uses teacher relevance scores to weight samples, emphasizing higher-scoring positives and easy negatives while remaining robust to latent positives among negatives.The authors contrast this with standard contrastive loss.Source: §3.2.1 · p. 5

3.2.2 Rank Imitation

  • Rank imitation aligns teacher and student rankings of positive and hard-negative samples, while a second objective distinguishes hard negatives from in-batch easy negatives.The alignment uses Pearson correlation; in-batch negatives are excluded from the positive-versus-hard-negative comparison.Source: §3.2.2 · p. 5

3.2.3 Feature Imitation

  • Feature imitation matches teacher and student similarity matrices across positive and hard-negative query–passage pairs to transfer relational patterns.The method compares relationships between pair representations rather than directly aligning classification-token embeddings.Source: §3.2.3 · p. 6

4 Experiments

Across three Chinese-focused tasks, the experiments compare D2LLM with five baselines and report its strongest gains on NLI; further analyses examine STS and IR, runtime, and component contributions.

4 Experiments

  • The evaluation covers NLI and STS with 0.3M training samples each and IR with 0.9M, comparing D2LLM with five state-of-the-art baselines.All baselines have at least as many tunable parameters as D2LLM except SGPT.Source: §4 · p. 6

4.2 Results on STS and IR

  • On STS, D2LLM beats same-dataset methods in 10 of 14 cases, while BGE-ft sometimes and original BGE consistently outperform it; detailed IR results are deferred to Appendix E.Finetuning the teacher on 0.3M STS samples improves its average by 7.17%; D2LLM-ft improves 1.69% over D2LLM and beats other same-data methods by at least 17.42% on average.Source: Table 2 · p. 7, §4.2 · p. 8

4.3 Runtime Analysis

  • D2LLM uses precomputed passage vectors and an MLP for scoring, with only marginally higher relevance-scoring time than bi-encoders and markedly better efficiency than the cross-encoder teacher.The runtime comparison on T2Retriever covers query vectorization, relevance scoring, and passage ranking; passage-embedding preprocessing is excluded for bi-encoders.Source: §4.3 · p. 8, Figure 2 · p. 8

4.4 Ablation Studies

  • The ablations identify contrastive, rank, and feature imitation losses, PMA, and IEM as important; larger teacher LLMs also yield more capable students.Mean pooling or [EOS] cannot aptly replace PMA, while cosine similarity alone does not capture query-passage dynamics adequately.Source: §4.4 · pp. 8–9, Table 3 · p. 8

5 Conclusion

D2LLM distills knowledge from an LLM teacher into an efficient semantic-search student, combining cross-encoder accuracy with bi-encoder efficiency.

5 Conclusion

  • D2LLM uses specialized modules and losses to distill an LLM teacher’s knowledge into a compact student that combines cross-encoder accuracy with bi-encoder efficiency.Together, the experiments support this synthesis of accuracy and operational efficiency.Source: §5 · p. 9

6 Ethical Considerations

The authors characterize the research as fundamental and report using exclusively open-source datasets, while acknowledging responsibilities around semantic-search applications.

6 Ethical Considerations

  • The authors report using exclusively open-source datasets and encourage continued ethical consideration of semantic search’s influence on information access and dissemination.They do not anticipate significant social implications from this fundamental research.Source: §6 · p. 9

7 Limitations

The paper identifies three limitations: the IEM may miss details captured by the teacher, tuning the combined-loss weights is complex, and training requires substantial computational resources.

7 Limitations

  • The IEM may omit intricate details captured by the teacher LLM, while increasing its complexity may improve performance at the expense of efficiency.The combined loss also has three weight parameters that are complex to tune, and training may be prohibitive for users with limited computational means.Source: §7 · p. 9

A Datasets

The paper evaluates semantic search on NLI, STS, and IR using task-specific training and testing data, with testing data drawn from CMTEB.

A Datasets

  • NLI tests entailment between sentence pairs on OCNLI and CMNLI, using 0.3M training examples and Accuracy, AP, Precision, and Recall.The NLI training data comprises SNLIzh, NLI-zh, ATEC, BQ, LCQMC, and PAWSX.Source: A Datasets · p. 12
  • STS predicts sentence-pair similarity on seven test datasets, using the same training data as NLI and Pearson and Spearman correlations for evaluation.The test datasets are ATEC, BQ, LCQMC, PAWSX, STSB, AFQMC, and QBQTC.Source: A Datasets · p. 12
  • IR ranks relevant documents for queries, training on a 0.9M sample and testing on six retrieval datasets with MRR@10 and Recall@10.The training sample combines T2Ranking, DuReader, cMedQA2, and mMARCO; testing includes T2Retrieval, DuRetrieval, CovidRetrieval, CmedqaRetrieval, MedicalRetrieval, and MMarcoRetrieval.Source: A Datasets · p. 12

B Implementation Details

Implementation uses a shared LLM foundation for D2LLM and comparable baselines, while task-specific teacher finetuning supplies binary relevance judgments for distillation.

B Implementation Details

  • D2LLM uses Qwen-7B-Chat, 89M tunable parameters, a 32-head PMA, and a two-MLP IEM; teacher outputs are stored before student training.Training takes about 10 hours for symmetric tasks and 22 hours for IR, versus about 9 and 20 hours for LLaRA.Source: B Implementation Details · pp. 12–13

C Baselines

  • The five baselines include BGE, RocketQAv2, SGPT, Udever, and LLaRA, spanning BERT-style and GPT-style bi-encoders with differing training and pooling strategies.LoRA ranks for Udever and LLaRA are set to 40 to match D2LLM’s trainable parameter count.Source: C Baselines · p. 13

D Teacher Model Finetuning

  • Teacher finetuning uses equal numbers of positive pairs and hard negatives, task-appropriate prompts, and binary “yes” or “no” responses.LoRA rank 32 aligns the teacher’s learnable parameter count with comparable methods.Source: D Teacher Model Finetuning · p. 13

E Results on IR

The IR results favor D2LLM across most reported instances, while ablations examine its losses and modules and compare shared versus specialized IEM branches.

F Ablation Studies

  • The ablation studies assess the contribution of the proposed losses and modules through targeted modifications to the model.The corresponding results are presented in Table 3.Source: F Ablation Studies · p. 13, §F.1 · p. 14

E Results on IR

  • D2LLM outperformed methods trained on the same data in 8 of 12 IR instances, although BGE achieved the highest performance.The authors note that the LLM cross-encoder teacher was impractically slow for real-world retrieval.Source: E Results on IR · p. 13, Table 4 · p. 14

F.1 Impact of Losses and Modules

  • Replacing CI with standard contrastive loss reduced performance by 3.59%, which the authors attribute to sensitivity to positives concealed among hard negatives.The modified setup replaces the proposed CI loss with standard contrastive loss.Source: §F.1 · p. 14
  • Ablation reductions were 6.57% for −RI_PH, 4.92% for −RI_HI, 2.17% for −FI, 1.71% for −PMA+mean, and 1.51% for −PMA+[EOS].Replacing IEM with cosine similarity reduced performance by 3.83%, and reducing the teacher from 7B to 1.8B reduced it by 14.76%.Source: §F.1 · p. 14

F.2 Impact of Single and Dual Branches in the IEM

  • D2LLM-dual matched specialized models on their corresponding symmetric or asymmetric tests, while each specialized model struggled on the other relation type.The comparison covers OCNLI and CMNLI for symmetric relations and T2Retrieval for asymmetric relations.Source: Table 5 · p. 14, §F.2 · p. 15

F.3 Weight Sensitivity Analysis

Across NLI and IR tasks, performance is stable over a range of loss weights, with a similar optimal combination across tasks; the IEM emulates query–passage interaction, while quantization is proposed to improve real-world search efficiency.

F.3 Weight Sensitivity Analysis

  • The authors choose loss weights sequentially rather than by grid search, conserving resources but not guaranteeing globally optimal parameters.They first tune α, then β, then γ near previously selected values.Source: §F.3 · p. 15
  • α = 1 is optimal: lower values correlate with suboptimal performance, while higher values bring a decline.The authors associate low α with inadequate use of the teacher’s ranking capacity.Source: §F.3 · p. 15
  • β = 0.3 yields the most significant performance augmentation in the reported experiments.β determines attention to the LRI loss term.Source: §F.3 · p. 15
  • Increasing γ improves performance up to a point, but excessive feature-imitation weight disrupts the balance among losses and impairs overall performance.The passage does not specify the point at which the decline begins.Source: §F.3 · pp. 15–16
  • A similar optimal weight combination across tasks and stable performance over a parameter range mean D2LLM generally needs no laborious task-specific tuning.The experiments use α = 1, β = 0.3, and γ = 0.1 unless otherwise specified.Source: §F.3 · p. 16

G Discussion on IEM

  • The IEM takes concatenated query and passage representations; its f1 layer emulates teacher cross-attention, and f2 routes symmetric or asymmetric relations to the corresponding MLP branch.The branch produces a “yes” or “no” representation to handle both search types.Source: G Discussion on IEM · p. 16
  • D2LLM-dual outperforms D2LLM-cos-mixed on symmetric and asymmetric tasks; the IEM still trails cross-encoders because interaction follows separate query and passage embedding derivation.The authors attribute the gap to missing the continuous cross-attention modeling of cross-encoders.Source: G Discussion on IEM · p. 16

H Efficiency Enhancement in Real-world Applications

  • Quantizing PMA embeddings and IEM parameters is proposed to reduce storage and accelerate search, with a slight accuracy trade-off.The passage also suggests that faster search can contribute to better ranking performance in subsequent procedures.Source: H Efficiency Enhancement in Real-world… · p. 16

I Bad Case Study

The bad-case study presents NLI and IR examples where student predictions or rankings miss positive or negative labels; the authors suggest teacher misjudgments may contribute to such errors.

I Bad Case Study

  • In two NLI cases, the student predicts 0.5959 for a contradiction labeled 0 and 0.3541 for an entailment labeled 1; teacher logits are 0.9141 and 0.4688, respectively.The authors propose that teacher misjudgments may hinder student learning.Source: I Bad Case Study · p. 17
  • For two positive IR examples, the student ranks the motorcycle passage 32 and the phone-help passage 46.The reported teacher logits are 0.081 and 0.135, respectively.Source: I Bad Case Study · p. 17
  • The authors suggest that teacher misjudgments may hinder student learning and propose adding error examples through in-context learning or fine-tuning the teacher while avoiding catastrophic forgetting.They present these as potential strategies, not demonstrated remedies.Source: I Bad Case Study · p. 17
Loading 2406.17262v1…