Large neural models (such as Transformers) achieve state-of-the-art performance for information retrieval (IR). In this paper, we aim to improve distillation methods that pave the way for the resource-efficient deployment of such models in practice. Inspired by our theoretical analysis of the teacher-student generalization gap for IR models, we propose a novel distillation approach that leverages the relative geometry among queries and documents learned by the large teacher model. Unlike existing teacher score-based distillation methods, our proposed approach employs embedding matching tasks to provide a stronger signal to align the representations of the teacher and student models. In addition, it utilizes query generation to explore the data manifold to reduce the discrepancies between the student and the teacher where training data is sparse. Furthermore, our analysis also motivates novel asymmetric architectures for student models which realizes better embedding alignment without increasing online inference cost. On standard benchmarks like MSMARCO, we show that our approach successfully distills from both dual-encoder (DE) and cross-encoder (CE) teacher models to 1/10th size asymmetric students that can retain 95-97% of the teacher performance.
Paper
Similar papers
Peer review
Summary
This paper presents a distillation technique called EmbedDistill for Information Retrieval, aiming to align the pooled representations of various IR models. The focus of this study revolves around two distillation paradigms: DE-to-DE distillation and CE-to-DE distillation. The former involves straightforward embedding matching, while the latter utilizes special tokens for alignment.
Strengths
1. Embedding matching is a practical method for distillation. This work pioneers the application of feature mapping to the distillation of Information Retrieval (IR) models, contributing to the existing body of knowledge in this area. 2. The experimental findings presented in this study reveal certain positive outcomes when applying EmbedDistill to IR models.
Weaknesses
1. The concept of embedding matching lacks novelty and has been extensively explored in prior studies [1, 2, 3]. Pooling techniques are also commonly employed to handle the matching of intermediate features. 2. The theoretical proof, while complex, does not significantly contribute to understanding the proposed method, as it is essentially a standard feature mapping approach already described in previous works. 3. There appears to be no direct correlation between the term "Geometric" and the proposed method. If a straightforward matching technique can be considered a geometric algorithm, it implies that most feature-based KD methods could also be classified as geometric. [1] Jiao, Xiaoqi, et al. "Tinybert: Distilling bert for natural language understanding." arXiv preprint arXiv:1909.10351 (2019). [2] Romero, Adriana, et al. "Fitnets: Hints for thin deep nets." arXiv preprint arXiv:1412.6550 (2014). [3] Hofstätter, Sebastian, et al. "Improving efficient neural ranking models with cross-architecture knowledge distillation." arXiv preprint arXiv:2010.02666 (2020).
Questions
My main concern lies in the novelty of this work. Please refer to the weaknesses.
Rating
3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
3 good
Presentation
2 fair
Contribution
1 poor
Limitations
N/A
Summary
This paper proposes EmbedDistill, a novel distillation approach for IR that directly aligns the embeddings from the student and teacher. The authors conduct a theoretical analysis of the teacher-student generalization gap, strengthening the importance of embedding alignment and suggesting the use of asymmetric encoders. The architecture can be directly applied in DE to DE distillation. For CE to DE distillation, the authors propose dual-pooling to obtain document and query embeddings from CE.
Strengths
- The authors show theoretically that embedding matching can close the teacher-student generalization gap in IR setting, and support the use of asymmetric configurations.
Weaknesses
- The authors claim that EmbedDistill can be used for CE to DE distillation, but the empirical results show poor performance. It only improves the score distillation method from 33.3 to 33.7 (+0.4), while the teacher score is 37.0. This raises doubts about the effectiveness of EmbedDistill in CE to DE distillation. EmbedDistill appears to only work in DE to DE distillation, which is an easier scenario where embedding matching is trivial. This significantly weakens the contribution of this work, as EmbedDistill can’t learn from CE, the stronger architecture. - One possible reason for this is that the embeddings extracted from a dual-pooled teacher may not be good. Since the query and document embeddings are concatenated in the re-ranker, the model can simply make the representations of the two close to each other when the query and document are relevant, without properly distributing the representations as required in the DE retrieval scenario. This poses a major challenge for embedding matching in CE to DE distillation, and the authors do not seem to address it adequately. - There are errors and inconsistencies in the formulas presented in the paper and the appendix. For instance, in line 130, $R\left(s^{\mathrm{s}}, s^{\mathrm{t}} ; \mathcal{S}_n\right)$ is not defined (should it be the empirical risk?). In the proof of Lemma C.3, the authors use $K_Q$ and $K_D$ in Eq. (25) but use $K$ elsewhere. The last line of Eq. (26) is incorrect; it should be $=\mathbb{E}\left[\left(1-y-\left(1-\sigma\left(F(q)^{\top} G(d)\right)\right)\right) f(q)^{\top} g(d)\right]$. The authors should thoroughly check all formulas and correct any errors. - The document needs improvement in terms of writing quality. - Section 5, which covers the result analysis, is difficult to follow. - The authors use too many parentheses to add additional explanations, which interrupts the reading flow. - The captions of the tables contain too much information. - The teacher performance is presented in the content of the table of NQ test and BEIR, but is described in the caption of the table of NQ dev and MS MARCO dev, which can confuse the reader. - The author uses too much italics in the introduction, which obscures the key points of the paper.
Questions
Please explain Figure 4 in Appendix G.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
2 fair
Presentation
2 fair
Contribution
2 fair
Limitations
As pointed out in the Weaknesses, EmbedDistill does not seem to work in CE to DE distillation, which limits the contribution.
Summary
This paper proposes to learn dense neural IR models by distilling not only the teacher scores but also by aligning the learned representations. The paper has a theoretic part that motivate the actual embedding distillation with two propositions. The first one shows that the teacher-student gap is bounded by an expression where some terms are the difference between teacher/student embeddings. The second bounds one of term of the expression by leveraging the $u$-covering number of function classes for teacher and student. Experiments show that using a regularization based on embeddings improve results in two settings: dense to dense and cross-encoder to dense distillation. In the latter case, the authors propose a simple but original approach to get a document and query representation. The experiments are conducted on the Natural Questions and MS-Marco datasets.
Strengths
This modification of the distillation procedure can easily be applied to all dense model, so the impact of the paper might be important in changing the mainstream procedure. There is also a nice result showing that just using embedding alignment (without the other part of the distillation loss) words pretty well. The mathematics justifying the approach are also interesting, by integrating the distance between teacher and student embeddings in the picture, but the bound do not seem to be of practical interest (apart from motivating the approach) and thus might not need to be that central in the paper (esp. proposition 3.2).
Weaknesses
In the experiments, the training based on query generation should not be the last model modification – it should rather the query/document embedding part, since query generation is not the core method proposed by the authors. This would also allow to see what exactly is brought by aligning the embeddings. The cross-encoder based distillation (which is the main one used by dense models) is quite disappointing in term of performance (even the best CE model performs worse than the best dense model, which should not be the case)… and suggests that to properly train a dense model with their method, one needs a very well trained dense models (but still, distilling to a lower-capacity model is interesting in that case, although the difference between the two procedures is not discussed or experimented with. The discussion 196-206 is not really needed (it is quite obvious that aligning with the `[CLS]` does not make sense).
Questions
1. Please clarify what is the $u$-covering number and how important it is in the case of the problem 2. what is “inheriting the document tower” (l. 189) 3. How are the CE models trained (this should be in the main text), especially in case of the modified CE models that allow computing meaningful document/query representations.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The limitation section is quite generic apart from the first sentence that could be more discussed. Also, the fact the CE distillation is not that performant is (at least for me) a major limitation (since cross-encoders are so much powerful), but it is also a challenge for this type of method.
Summary
This work proposed to use embedding matching task for knowledge distillation for IR: unlike in traditional knowledge distillation where the matching score of query and document is used, the embedding matching tasks try to align the embedding representation of query and document. It works both for dual-encoder model (DE) and cross-encoder model (CE). The student model learned achieves 1/10th of size while retraining 95-97% of the teacher performance. This work presented a theoretical analysis of the teacher-student generalization gap for IR settings. The student DE model proposed has an asymmetric configuration where the query encoder is smaller than document encoder, and the later is frozen during knowledge distillation. This could reduces inference latency at query time. To validate the effective of the proposed method, experiments on common benchmark for IR is used. The model achieves competitive results on Natural Questions, MSMARCO and BEIR (zero-shot IR benchmark).
Strengths
* In-depth theoretical analysis of the teacher-student generalization gap in IR models, which inspired the embedding alignment methods. * Extensive experimental results to validate the proposed method: it achieves competitive results on both Natural Questions, MSMARCO datasets and also the BAIR datasets in zero-shot cases. * Overall well written and easy to follow.
Weaknesses
Overall, the technical contribution presented in this work isn’t super strong in my opinion. Knowledge Distillation is known to be effective for IR, both for DE and CE models. The embedding alignment techniques has also be explored in other settings. Although specific technical challenges needed to be addressed in applying embedding alignment for knowledge distillation of IR models, it does not seem to be groundbreaking.
Questions
* Q1: Has aligning other layers of the network being considered/explored?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The authors addressed the limitations of the work where it only studied transformer-based models (but expect it to work for other models such as MLP). It only studied IR in text and did not extend to multi-modal cases.
Summary
The present study introduces a new technique, EmbedDistill, designed to transfer knowledge from large-scale neural networks to smaller versions for information retrieval (IR) purposes. This approach utilizes the relative geometric connections between queries and documents, as learned by the more extensive teacher model, to synchronize the representations of both teacher and student models. Furthermore, it scrutinizes the data manifold to diminish disparities between the student and teacher, particularly in areas where the training data is scant. The findings suggest that this proposed method can effectively distill both dual-encoder and cross-encoder teacher models into 1/10th size asymmetric students, maintaining 95-97% of the original teacher's performance. The study's key contributions consist of a groundbreaking geometric distillation technique, a unique query creation method for enhancing distillation quality, and practical results underlining the success of the proposed method.
Strengths
The paper presents a new approach to distilling knowledge from strong neural models to weaker ones for information retrieval. The paper makes several contributions that are noteworthy. - The paper's approach is original in that it leverages the relative geometry among queries and documents learned by the large teacher model to align the representations of the teacher and student models. - The paper is of good quality, with a well-defined problem statement, clear methodology, and adequate experimental evaluation. The paper's contributions are remarkable in that they provide a new approach to distilling knowledge from strong neural models to weaker ones for information retrieval. The authors provide a detailed analysis of the proposed approach and compare it with existing methods. The experiments are well-designed and the results are presented in a clear and concise manner. - The paper is well-written, well-organized, and easy to follow. The authors provide clear explanations of the proposed approach and the experimental setup.
Weaknesses
Although there are several distillation methods proposed in this paper, the core idea of this paper is to distill representation, via L2 distance, from a strong representation model to smaller ones. And the derivation of the representation is also straightforward and brute-forcing, i.e., CLS representation or mean pooling over corresponding text. The backbone of the retriever is too weak, which cannot verify the generality of the proposed methodology. That is, it’s required a stronger student model to check if the validity of proposed methodology will be mitigated with stronger students. Meantime, the neural retrievers and rankers used in this work is relatively out-of-date. It’s recommended to leverage more state-of-the-art retrievers/rankers to make the experiments more convincing.
Questions
Refer to the Weaknesses
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Yes
Thanks for your answer and further results related to the cross-encoder to dense-encoder distillation which are indeed interesting. Overall, I think the paper makes a contribution to the information retrieval field – although I agree with the other reviewers this is not a strong contribution.
Thank you for the feedback. Certain aspects of my concerns (Weakness 2) have been acknowledged and tackled within the response. However, I remain worried about the novelty of the embedding-matching approach, given that it's a general technique applicable to **ALL tasks**. In this regard, I'm inclined to respectfully disagree with the assertion that, "representation matching + X," wherein X={IR, NER, Entity linking, ...}, could be characterized as a novel idea, even as the first attempt within these specific domains. So, I tend to keep my initial score.
We would like to emphasize that our contribution is not about *blindly* applying embedding-alignment based distillation to an IR setting. In this work, we started by analyzing the performance gap between the teacher and student IR models and found the gap is bounded by the misalignment between the query and document representations (*novel analysis*). This motivated us to develop embedding-matching based solutions as well as **asymmetric DE** architecture since it realizes perfect alignment for the document embedding spaces between the teacher and student models without increasing runtime latency (novel formulation and novel architecture for IR). Moreover, we developed **a novel mechanism** for extraction of query and document representations from the CE model (*dual pooling*) that enables CE to DE distillation via embedding alignment. Hence, we would like the reviewer to reconsider their opinion. Even though the work’s algorithms look similar (e.g. alignment itself), it was motivated differently and is solving a different problem with a different formulation.
Thanks for the detailed comment. > Our contribution is not about blindly applying embedding-alignment based distillation to an IR. I agree with the author's perspective that within this study, alignment isn't simply imposed upon Information Retrieval (IR). But actually for all tasks, we need to do some adaptations to make alignment work accordingly. **My opinion is that adapting a very classic algorithm to a new setting is not novel, which can not reach the bar of NeurIPS.** It's not surprising that alignment works for IR, EL, NER, etc. > Moreover, we developed a novel mechanism for the extraction of query and document representations from the CE model (dual pooling) that enables CE to DE distillation via embedding alignment. Pooling is a very simple and general solution for heterogeneous KD [1,2]. For example, Attention KD is a very early work in this field, which just pools the channel dimensions to facilitate KD [1]. I don't think the novelty of this paper should lie in general KD techniques. Conclusion: This paper should pay more attention to IR, rather than adapting several KD techniques from other fields and show their effectiveness. Besides, the theoretical part is not strongly connected to the IR task. My concern about the novelty is not addressed. [1] Zagoruyko, Sergey, and Nikos Komodakis. "Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer." arXiv preprint arXiv:1612.03928 (2016). [2] Gao, Lei, and Hui Gao. "Feature Decoupled Knowledge Distillation via Spatial Pyramid Pooling." Proceedings of the Asian Conference on Computer Vision. 2022.
Thank you again for taking the time to engage with us during the discussion phase and providing a detailed response. While we feel that some of the opinions expressed by the reviewer are subjective, we would like to respond to other specific points raised by them. > Pooling is a very simple and general solution for heterogeneous KD [1,2]. For example, Attention KD is a very early work in this field, which just pools the channel dimensions to facilitate KD [1]. I don't think the novelty of this paper should lie in general KD techniques. > [1] Zagoruyko, Sergey, and Nikos Komodakis. "Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer." arXiv preprint arXiv:1612.03928 (2016). > [2] Gao, Lei, and Hui Gao. "Feature Decoupled Knowledge Distillation via Spatial Pyramid Pooling." Proceedings of the Asian Conference on Computer Vision. 2022. Thank you for bringing up [1,2]. We would like to point out that both of these works are from the computer vision/image classification domain and differ from our motivation for utilizing **dual pooling** – to extract out meaningful **separate** representations for queries and documents from a cross-encoder (CE) model which are otherwise not available in a vanilla CE model (**see Appendix G for the embedding analysis**). Please note that successful utilization of dual pooling also required modifying the training procedure for the CE model (**see Line 214 -- 217**). > Besides, the theoretical part is not strongly connected to the IR task. Please note that we specifically focus on an overall function class that involves two encoders and the scoring function is defined by the dot-product between the embeddings produced by the two encoders, which is precisely capturing the dual-encoder architecture used in IR (cf. Theorem 3.1). Similarly, our uniform deviation analysis (Proposition 3.2) is again specifically focusing on this dual-encoder architecture. Thus, we do not agree with the conclusion that the theoretical part is not strongly connected to the IR task. If the reviewer feels that our analysis can also be utilized/adapted to study distillation in other domains (which we tend to agree with), that should not be counted as a negative for our contribution. ### Studying distillation for IR even when distillation is a successful technique in other domains As a more general comment, we would also like to bring the reviewers’ attention to the growing literature on even **score-based distillation** in IR setting (see, e.g., [1] -- [5] which have appeared in venues like ICML, ICLR, CVPR and NAACL), despite knowledge distillation being extensively explored in the ML literature. This alludes to the fact that **distillation for IR** presents its unique challenges and addressing those requires novel contributions. Given this, we strongly believe that, if accepted, our work (which, to the best of our knowledge, is the first comprehensive treatment of embedding-alignment based distillation for IR) will be a valuable addition to NeurIPS. [1] Izacard et al. Distilling knowledge from reader to retriever for question answering. In ICLR 2021. [2] Menon et al. In defense of dual-encoders for neural ranking. ICML 2022. [3] Qu et al. RocketQA: An optimized training approach to dense passage retrieval for open-domain question answering. NAACL 2021. [4] Santhanam et al. ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction. NAACL 2022. [5] Miech et al. Thinking Fast and Slow: Efficient Text-to-Visual Retrieval with Transformers. CVPR 2021.
Thank you for your response. The new results of CE to DE distillation seem to be better. However, the gap between the CE teacher and the DE student still remains large - more significant than the gap in DE to DE distillation. As I pointed out before, the reason is that the authors didn't come up with a solution that makes the embeddings from a DE well-distributed. Thus I think the contribution is still limited. I've raised my score from 4 to 5.
Thank you for the response. We have significantly improved the performance of the DE student, surpassing the results of previous papers, and also reducing the CE teacher - DE student gap significantly. For example Table 3 from [1], their student (based on DistilBERT -- the same as ours) shows MRR@10 33.2 distilled from their ensemble teacher (which has MRR@10 39.9), and the gap is 6.7. In contrast, the teacher-student gap for our methods is only 3.1(40.0 vs 36.9). Also we believe that there is a big difference between the expressive power of the CE model and DE model and the gap is primarily coming from that. In particular, the late interaction between query and document representation via only a dot product is very limiting. If we keep the same embeddings but replace dot product with a richer late interaction like ColBERT [2], we can eliminate most of the gap (now the gap is only 0.9): | Method | MRR@10 | nDCG@10 | |:---------------------------------------------|---------:|---------:| | | | | | Teacher | 40.0 | 45.8 | | DE student with EmbedDistill | 36.9 | 42.6 | | + **Replace dot with ColBERT** | 39.1 | 45.0 | | | | | This shows the embedding quality might not be the main culprit for the gap. [1] Hofstätter et al. "Improving efficient neural ranking models with cross-architecture knowledge distillation." arXiv preprint arXiv:2010.02666 (2020). [2] Khattab et al. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT, page 39–48. Association for Computing Machinery, 490 New York, NY, USA, 2020.
Thank the authors for the details responses. I've read the authors comments as well as other reviewer's feedback and discussions. I decide to keep the original score.
Decision
Reject