Title: EmbedDistill: A Geometric Knowledge Distillation for Information Retrieval
Venue: NeurIPS 2023
Accepted: no
Keywords: Knowledge distillation

teacher-student gap

dual encoder

cross encoder

information retrieval

query generation

embedding matching

retrieval

re-ranking
Abstract: 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.

Abstract

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 dualencoder (DE) and cross-encoder (CE) teacher models to 1/10th size asymmetric students that can retain 95-97% of the teacher performance.

1 Introduction

Neural models for information retrieval (IR) are increasingly used to model the true ranking function in various applications, including web search [38], recommendation [65], and question-answering (QA) [6]. Notably, the recent success of Transformers [59]-based pre-trained language models [11, 30, 49] on a wide range of natural language understanding tasks has also prompted their utilization in IR to capture query-document relevance [see, e.g., 10, 34, 43, 26, 20].

A typical IR system comprises two stages: (1) A retriever first selects a small subset of potentially relevant candidate documents (out of a large collection) for a given query; and (2) A re-ranker then identifies a precise ranking among the candidates provided by the retriever. Dual-encoder (DE) models are the de-facto architecture for retrievers [26, 20]. Such models independently embed queries and documents into a common space, and capture their relevance by simple operations on these embeddings such as the inner product. This enables offline creation of a document index and supports fast retrieval during inference via efficient maximum inner product search implementations [12, 19], with online query embedding generation primarily dictating the inference latency. Cross-encoder (CE) models, on the other hand, are preferred as re-rankers, owing to their excellent performance [43, 9, 62]. A CE model jointly encodes a query-document pair while enabling early interaction among query and document features. Employing a CE model for retrieval is often infeasible, as it would require processing a given query with every document in the collection at inference time. In fact, even in the re-ranking stage, the inference cost of CE models is high enough [22] to warrant exploration of efficient alternatives [14, 22, 37]. Across both architectures, scaling to larger models brings improved performance at increased computational cost [41, 39].

Knowledge distillation [5, 13] provides a general strategy to address the prohibitive inference cost 38 associated with high-quality large neural models. In the IR literature, most existing distillation 39 methods only rely on the teacher's query-document relevance scores [see, e.g., 31, 14, 8, 51, 56] or 40 their proxies [16]. However, given that neural IR models are inherently embedding-based, it is natural 41 to ask: Is it useful to go beyond matching of the teacher and student models' scores , and directly aim 42 to align their embedding spaces ? 43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

With this in mind, we propose a novel distillation method for IR models that utilizes an embedding matching task to train student models. The proposed method is inspired by our rigorous treatment of the generalization gap between the teacher and student models in IR settings. Our theoretical analysis of the teacher-student generalization gap further suggests novel design choices involving asymmetric configurations for student DE models, intending to further reduce the gap by better aligning teacher and student embedding spaces. Notably, our proposed distillation method supports cross-architecture distillation and improves upon existing (score-based) distillation methods for both retriever and re-ranker models. When distilling a large teacher DE model into a smaller student DE model, for a given query (document), one can minimize the distance between the query (document) embeddings of the teacher and student (after compatible projection layers to account for dimension mismatch, if any). In contrast, a teacher CE model doesn't directly provide document and query embeddings, and so to effectively employ embedding matching-based distillation requires modifying the scoring layer with dual-pooling [61] and adding various regularizers. Both of these changes improve geometry of teacher embeddings and facilitate effective knowledge transfer to the student DE model via embedding matching-based distillation.

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

Our key contributions toward improving IR models via distillation are:

We provide the first rigorous analysis of the teacher-student generalization gap for IR settings which captures the role of alignment of embedding spaces of the teacher and student towards reducing the gap (Sec. 3).

Inspired by our analysis, we propose a novel distillation approach for neural IR models, namely EmbedDistill , that goes beyond score matching and aligns the embedding spaces of the teacher and student models (Sec. 4). We also show that EmbedDistill can leverage synthetic data to improve a student by further aligning the embedding spaces of the teacher and student (Sec. 4.3).

Our analysis motivates novel distillation setups. Specifically, we consider a student DE model with an asymmetric configuration, consisting of a small query encoder and a frozen document encoder inherited from the teacher. This significantly reduces inference latency of query embedding generation, while leveraging the teachers' high-quality document index (Sec. 4.1).

We provide a comprehensive empirical evaluation of EmbedDistill (Sec. 5) on two standard IR benchmarks - Natural Questions [23] and MSMARCO [40]. We also evaluate EmbedDistill on BEIR benchmark [57] which is used to measure the zero-shot performance of an IR model.

Note that prior works have utilized embedding alignment during distillation for non-IR setting [see, e.g., 52, 55, 18, 1, 64, 7]. However, to the best of our knowledge, our work is the first to study embedding matching-based distillation method for IR settings which requires addressing multiple IR-specific challenges such as cross-architecture distillation, partial representation alignment, and enabling novel asymmetric student configurations. Furthermore, unlike these prior works, our proposed method is theoretically justified to reduce the teacher-student performance gap.

2 Background

Let Q and D denote the query and document spaces, respectively. An IR model is equivalent to 81 a scorer s : Q × D → R , i.e., it assigns a (relevance) score s ( q, d ) for a query-document pair 82 ( q, d ) ∈ Q × D . Ideally, we want to learn a scorer such that s ( q, d ) > s ( q, d ′ ) iff the document d is 83 more relevant to the query q than document d ′ . We assume access to n labeled training examples 84 S n = { ( q i , d i , y i ) } i ∈ [ n ] . Here, d i = ( d i, 1 , . . . , d i,L ) ∈ D L , ∀ i ∈ [ n ] , denotes a list of L documents 85 and y i = ( y i, 1 , . . . , y i,L ) ∈ { 0 , 1 } L denotes the corresponding labels such that y i,j = 1 iff the 86 document d i,j is relevant to the query q i . Given S n , we learn an IR model by minimizing 87

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

While this learning framework is general enough to work with any IR models, next, we formally introduce two families of Transformer-based IR models that are prevalent in the recent literature.

2.1 Transformer-based IR models: Cross-encoders and Dual-encoders

Let query q = ( q 1 , . . . , q m 1 ) and document d = ( d 1 , . . . , d m 2 ) consist of m 1 and m 2 tokens, respectively. We now discuss how Transformers-based CE and DE models process the ( q, d ) pair.

Cross-encoder model. Let p = [ q ; d ] be the sequence obtained by concatenating q and d . Further, let ˜ p be the sequence obtained by adding special tokens such [CLS] and [SEP] to p . Given an encoder-only Transformer model Enc , the relevance score for the ( q, d ) pair is where w is a d -dimensional classification vector, and pool( · ) denotes a pooling operation that transforms the contextualized token embeddings Enc(˜ p ) to a joint embedding vector emb q,d . [CLS] -pooling is a common operation that simply outputs the embedding of the [CLS] token as emb q,d .

Dual-encoder model. Let ˜ q and ˜ d be the sequences obtained by adding appropriate special tokens to q and d , respectively. A DE model comprises two (encoder-only) Transformers Enc Q and Enc D , which we call query and document encoders, respectively. 1 Let emb q = pool ( Enc Q (˜ q ) ) and emb d = pool ( Enc D ( ˜ d ) ) denote the query and document embeddings, respectively. Now, one can define s ( q, d ) = 〈 emb q , emb d 〉 to be the relevance score assigned to the ( q, d ) pair by the DE model.

2.2 Score-based distillation for IR models

Most distillation schemes for IR [e.g., 31, 14, 8] rely on teacher relevance scores. Given a training set S n and a teacher with scorer s t , one learns a student with scorer s s by minimizing

where /lscript d captures the discrepancy between s s and s t . See Appendix A for common choices for /lscript d .

3 Teacher-student generalization gap: Inspiration for embedding alignment

Our main objective is to devise novel distillation methods to realize high-performing student DE models. As a first step in this direction, we rigorously study the teacher-student generalization gap as realized by standard (score-based) distillation in IR settings. Informed by our analysis, we subsequently identify novel ways to improve the student model's performance. In particular, our analysis suggests two natural directions to reduce the teacher-student generalization gap: 1) enforcing tighter alignment between embedding spaces of teacher and student models; and 2) exploring novel asymmetric configuration for student DE model.

Let R ( s ) = E [ /lscript ( s q, d , y )] be the population version of the empirical risk in Eq. 1, which measures the test time performance of the IR model defined by the scorer s . Thus, R ( s s ) -R ( s t ) denotes the teacher-student generalization gap . In the following result, we bound this quantity (see Appendix C.1 for a formal statement and proof). We focus on distilling a teacher DE model to a student DE model and L = 1 (cf. Sec. 2) as it leads to easier exposition without changing the main takeaways. Our analysis can be extended to L > 1 or CE to DE distillation with more complex notation.

Theorem 3.1 (Teacher-student generalization gap (informal)) . Let F and G denote the function classes for the query and document encoders for the student model, respectively. Suppose that the score-based distillation loss /lscript d in Eq. 3 is based on binary cross entropy loss (Eq. 12 in Appendix A). Let one-hot (label-dependent) loss /lscript in Eq. 1 be the binary cross entropy loss (Eq. 10 in Appendix A). Further, assume that all encoders have the same output dimension and embeddings have their /lscript 2 -norm bounded by K . Then, we have

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

Special tokens

Query & doc tokens

Student Query

Encoder

Teacher Doc

Encoder

…

…

Doc tokens

Query tokens

Score

Teacher Doc

Encoder

…

…

Doc tokens

Query tokens

Score

Encoder

≈

Embedding matching

Pooling

Pooling

Pooling

Pooling

Score-based distillation

(a) DE to DE distillation

Teacher

Cross Encoder

Student Query

Encoder

Student Doc

Encoder

…

…

Doc tokens

Query tokens

Score

Score

Special tokens

Query & doc tokens

…

…

Doc tokens

Query tokens

Pooling

Pooling

Score-based

distillation

≈

Embedding matching

(b) CE to DE distillation

where E n ( F , G ) := sup s s ∈ F × G ∣ ∣ R ( s s , s t ; S n ) -E /lscript d ( s s q,d , s t q,d )∣ ∣ ; σ denotes the sigmoid function; and ∆ ( s t ; S n ) denotes the deviation between the empirical risk (on S n ) and population risk of the teacher s t . Here, R Emb ,Q (t , s; S n ) and R Emb ,D (t , s; S n ) measure misalignment between teacher and student embeddings by focusing on queries and documents, respectively (cf. Eq. 7 & 8 in Sec. 4.1).

The last three quantities in the bound in Thm. 3.1, namely ∆ ( s t ; S n ) , E [ | σ ( s t q,d ) -y | ] , and 1 n ∑ i ∈ [ n ] | σ ( s t q i ,d i ) -y i | , are independent of the underlying student model. These terms solely depend on the quality of the underlying teacher model s t . That said, the teacher-student gap can be made small by reducing the following three terms: 1) uniform deviation of the student's empirical distillation risk from its population version E n ( F , G ) ; 2) misalignment between teacher student query embeddings R Emb ,Q (t , s; S n ) ; and 3) misalignment between teacher student document embeddings R Emb ,D (t , s; S n ) .

The last two terms motivate us to propose an embedding matching -based distillation that explicitly aims to minimize these terms during student training. Even more interestingly, these terms also inspire an asymmetric DE configuration for the student which strikes a balance between the goals of reducing the misalignment between the embeddings of teacher and student (by inheriting teacher's document encoder) and ensuring serving efficiency (small inference latency) by employing a small query encoder. Before discussing these proposals in detail in Sec. 4 and Fig. 1, we explore the first term E n ( F , G ) and highlight how our proposals also have implications for reducing this term. Towards this, the following result bounds E n ( F , G ) . Due to space constraints, we present an informal statement of the result (see Appendix C.2 for a more precise statement and proof).

Proposition 3.2. Let /lscript d be a distillation loss which is L /lscript d -Lipschitz in its first argument. Let F and G denote the function classes for the query and document encoders, respectively. Further assume that, for each query and document encoder in our function class, the query and document embeddings have their /lscript 2 -norm bounded by K . Then,

Furthermore, with a fixed document encoder, i.e., G = { g ∗ } ,

Here, N ( u, · ) is the u -covering number of a function class.

Note that Eq. 5 and Eq. 6 correspond to uniform deviation when we train without and with a frozen 156 document encoder, respectively. It is clear that the bound in Eq. 6 is less than or equal to that in 157 Eq. 5 (because N ( u, G ) ≥ 1 for any u ), which alludes to desirable impact of employing a frozen 158 document encoder as one of our proposal seeks to do via inheriting teacher's document encoder (for 159 instance in an asymmetric DE configuration). Furthermore, our proposal of employing an embedding160 matching task will regularize the function class of query encoders; effectively reducing it to F ′ with 161 | F ′ | ≤ | F | . The same holds true for document encoder function class when document encoder is 162 trainable (as in Eq. 5), leading to an effective function class G ′ with | G ′ | ≤ | G | . Since we would have 163 N ( u, F ′ ) ≤ N ( u, F ) and N ( u, G ′ ) ≤ N ( u, G ) , this suggests desirable implications of embedding 164 matching for reducing the uniform deviation bound. 165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

4 Embedding-matching based distillation

Informed by our analysis of teacher-student generalization gap in Sec. 3, we propose EmbedDistill - a novel distillation method that explicitly focuses on aligning the embedding spaces of the teacher and student. Our proposal goes beyond existing distillation methods in the IR literature that only use the teacher scores. Next, we introduce EmbedDistill for two prevalent settings: (1) distilling a large DE model to a smaller DE model; 2 and (2) distilling a CE model to a DE model.

4.1 DE to DE distillation

Given a ( q, d ) pair, let emb t q and emb t d be the query and document embeddings produced by the query encoder Enc t Q and document encoder Enc t D of the teacher DE model, respectively. Similarly, let emb s q and emb s d denote the query and document embeddings produced by a student DE model with (Enc s Q , Enc s D ) as its query and document encoders. Now, EmbedDistill optimizes the following embedding alignment losses in addition to the score-matching loss from Sec. 2.2 to align query and document embeddings of the teacher and student:

Asymmetric DE. We also propose a novel student DE configuration where the student employs the teacher's document encoder (i.e., Enc s D = Enc t D ) and only train its query encoder, which is much smaller compared to the teacher's query encoder. For such a setting, it is natural to only employ the embedding matching loss in Eq. 7 as the document embeddings are aligned by design (cf. Fig. 1a).

Note that this asymmetric student DE does not incur an increase in latency despite the use of a large teacher document encoder. This is because the large document encoder is only needed to create a good quality document index offline, and only the query encoder is evaluated at inference time. Also, the similarity search cost is not increased as the projection layer ensures the same small embedding dimension as in the symmetric DE student. Thus, for DE to DE distillation, we prescribe the asymmetric DE configuration universally. Our theoretical analysis (cf. Sec. 3) and experimental results (cf. Sec. 5) suggest that the ability to inherit the document tower from the teacher DE model can drastically improve the final performance, especially when combined with query embedding matching task (cf. Eq. 7).

4.2 CE to DE distillation

Given that CE models jointly encode query-document pairs, individual query and document embeddings are not readily available to implement embedding matching losses as per Eq. 7 and 8. This makes it challenging to employ EmbedDistill for CE to DE distillation.

As a naïve solution, for a ( q, d ) pair, one can simply match a joint transformation of the student's query embedding emb s q and document embedding emb s d to the teacher's joint embedding emb t q,d , produced by (single) teacher encoder Enc t . However, we observed that including such an embedding matching task often leads to severe over-fitting, and results in a poor student. Since s t ( q, d ) = 〈 w, emb t q,d 〉 , during CE model training, the joint embeddings emb t q,d for relevant and irrelevant ( q, d ) pairs are encouraged to be aligned with w and -w , respectively. This produces degenerate embeddings that do not capture semantic query-to-document relationships. We notice that even the final query and document token embeddings lose such semantic structure (cf. Appendix G.2). Thus, a teacher CE model with s t ( q, d ) = 〈 w, emb t q,d 〉 does not add value for distillation beyond score-matching; in fact, it hurts to include naïve embedding matching. Next, we propose a modified CE model training strategy that facilitates EmbedDistill .

CE models with dual pooling. A dual pooling scheme is employed in the scoring layer to produce two embeddings emb t q ← ( q,d ) and emb t d ← ( q,d ) from a CE model that serve as the proxy query and document embeddings, respectively. Accordingly, we define the relevance score as s t ( q, d ) = 〈 emb t q ← ( q,d ) , emb t d ← ( q,d ) 〉 . We explore two variants of dual pooling: (1) special token-based pooling that pools from [CLS] and [SEP] ; and (2) segment-based weighted mean pooling that separately

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

performs weighted averaging on the query and document segments of the final token embeddings. See Appendix B for details.

In addition to dual pooling, we also utilize a reconstruction loss during the CE training, which measures the likelihood of predicting each token of the original input from the final token embeddings. This loss encourages reconstruction of query and document tokens based on the final token embeddings and prevents the degeneration of the token embeddings during training. Given proxy embeddings from the teacher CE, we can perform EmbedDistill with the embedding matching loss defined in Eq. 7 and Eq. 8 (cf. Fig. 1b).

4.3 Task-specific online data generation

Data augmentation as a general technique has been previously considered in the IR literature [see, e.g., 45, 47, 17], especially in data-limited, out-of-domain, or zero-shot settings. As EmbedDistill aims to align the embeddings spaces of the teacher and student, the ability to generate similar queries or documents can naturally help enforce such an alignment globally on the task-specific manifold. Given a set of unlabeled task-specific query and document pairs U m , we can further add the embedding matching losses R Emb , Q (t , s; U m ) or R Emb , D (t , s; U m ) to our training objective. Interestingly, for DE to DE distillation setting, our approach can even benefit from a large collection of task-specific queries Q ′ or documents D ′ . Here, we can independently employ embedding matching losses R Emb , Q (t , s; Q ′ ) or R Emb , D (t , s; D ′ ) that focus on queries and documents, respectively. Please refer to Appendix E describing how the task-specific data were generated.

5 Experiments

We now conduct a comprehensive evaluation of the proposed distillation approach. Specifically, we highlight the utility of the approach for both DE to DE and CE to DE distillation. We also showcase the benefits of combining our distillation approach with query generation methods.

5.1 Setup

Benchmarks and evaluation metrics. Weconsider two popular IR benchmarks - Natural Questions (NQ) [24] and MSMARCO [40], which focus on finding the most relevant passage/document given a question and a search query, respectively. NQ provides both standard test and dev sets, whereas MSMARCO provides only the dev set that are widely used for common benchmarks. In what follows, we use the terms query (document) and question (passages) interchangeably. For NQ, we use the standard full recall ( strict ) as well as the relaxed recall metric [20] to evaluate the retrieval performance. For MSMARCO, we focus on the standard metrics Mean Reciprocal Rank (MRR)@10, and normalized Discounted Cumulative Gain (nDCG)@10 to evaluate both re-ranking and retrieval performance. For the re-ranking, we restrict to re-ranking only the top 1000 candidate document provided as part of the dataset to be fair, while some works use stronger methods to find better top 1000 candidates for re-ranking (resulting in higher evaluation numbers) See Appendix D for a detailed discussion on these evaluation metrics. Finally, we also evaluate EmbedDistill on the BEIR benchmark [57] in terms of nDCG@10 and recall@100 metrics.

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

Model architectures. We follow the standard Transformers-based IR model architectures similar to Karpukhin et al. [20], Qu et al. [48], O˘ guz et al. [47]. We utilized various sizes of DE models based on BERT-base [11] (12-layer, 768 dim, 110M parameters), DistilBERT [55] (6-layer, 768 dim, 67.5M parameters -∼ 2/3 of base), or BERT-mini [58] (4-layer, 256 dim, 11.3M parameters -∼ 1/10 of base). For query generation (cf. Sec. 4.3), we employ BART-base [27], an encoder-decoder model, to generate similar questions from each training example's input question (query). We randomly mask 10% of tokens and inject zero mean Gaussian noise with σ = { 0 . 1 , 0 . 2 } between the encoder and decoder. See Appendix E for more details on query generation and Appendix F.1 for hyperparameters.

5.2 DE to DE distillation

We employ AR2 [63] 3 and SentenceBERTv5 [50] 4 as teacher DE models for NQ and MSMARCO. Note that both models are based on BERT-base. For DE to DE distillation, we consider two kinds of configurations for the student DE model: (1) Symmetric : We use identical question and document encoders. We evaluate DistilBERT and BERT-mini on both datasets. (2) Asymmetric : The student inherits document embeddings from the teacher DE model and are not trained during the distillation. For query encoder, we use DistilBERT or BERT-mini which are smaller than document encoder.

Student DE model training. We train student DE models using a combination of (i) one-hot loss (cf. Eq. 9 in Appendix A) on training data; (ii) distillation loss in (cf. Eq. 11 in Appendix A); and (iii) em-

bedding matching loss in Eq. 7. We used [CLS] -pooling for all student encoders. Unlike DPR [20] or AR2, we do not use hard negatives from BM25 or other models, which greatly simplifies our distillation procedure.

Results and discussion. To understand the impact of various proposed configurations and losses, we train models by sequentially adding components and evaluate their retrieval performance on NQ and MSMARCO dev set as shown in Table 1 and Table 3 respectively. (See Table 6 in Appendix F.2 for performance on NQ in terms of the relaxed recall and Table 7 in Appendix F.3 for MSMARCO in terms of nDCG@10.)

We begin by training a symmetric DE without distillation. As expected, moving to distillation brings in considerable gains. Next, we swap the student document encoder with document embeddings from the teacher (non-trainable), which leads to a good jump in the performance. Now we can introduce EmbedDistill with Eq. 7 for aligning query representations between student and teacher. The two losses are combined with weight of 1 . 0 (except for BERT-mini models in the presence of query generation with 5 . 0 ). This improves performance significantly, e.g.,it provides ∼ 3 and ∼ 5 points increase in recall@5 on NQ with students based on DistilBERT and BERT-mini, respectively (Table 1). We further explore the utility of EmbedDistill in aligning the teacher and student embedding spaces in Appendix G.1.

On top of the two losses (standard distillation and embedding matching), we also use R Emb , Q (t , s; Q ′ ) from Sec. 4.3 on 2 additional questions (per input question) generated from BART. We also try a variant where we eliminate the standard distillation loss and only employ the embedding matching loss in Eq. 7 along with inheriting teacher's document embeddings. This configuration without the standard distillation loss leads to excellent performance (with query generation again providing additional gains in most cases.)

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

It is worth highlighting that DE models trained with the proposed methods (e.g., asymmetric DE with embedding matching and generation) achieve 99% of the performance in both NQ/MSMARCO tasks with a query encoder that is 2/3rd the size of that of the teacher. Furthermore, even with 1/10th size of the query encoder, our proposal can achieve 95-97% of the performance. This is particularly useful for latency critical applications with minimal impact on the final performance.

Finally, we take our best student models, i.e., one trained using with additional embedding matching loss and using data augmentation from query generation, and evaluate on test sets. We compare with various prior work and note that most prior work used considerably bigger models in terms of parameters,

depth (12 or 24 layers), or width (upto 1024 dims). For NQ test set results are reported in Table 2, but as MSMARCO does not have any public test set, we instead present results for the BEIR benchmark in Table 4. Note we also provide evaluation of our SentenceBERT teacher achieving very high performance on the benchmark which can be of independent interest (please refer to Appendix F.4 for details). For both NQ and BEIR, our approach obtains competitive student model with fewer than 50% of the parameters: even with 6 layers, our student model is very close (98-99%) to its teacher.

5.3 CE to DE distillation

We consider two CE teachers for MSMARCO reranking task 5 : a standard [CLS] -pooled CE teacher, and the Dual-pooled CE teacher (cf. Sec. 4.2). Both teachers are based on RoBERTa-base and trained on triples in the training set for 300K steps with crossentropy loss.

Student DE model training. We considered the following distillation variants: standard score-based distillation from the [CLS] -pooled teacher, and our novel Dual-pooled CE teacher (with and without embedding matching loss). For each variant, we initialize encoders of the student DE model with two RoBERTabase models and train for 500K steps on the training triples. We performed the naïve joint embedding

Table 5: Performance of DE models distilled from [CLS] -pooled and Dual-pooled CE models on MSMARCO re-ranking task (original top1000 dev). While both teacher models perform similarly, embedding matching-based distillation only works with the Dual-pooled teacher. See Appendix F for nDCG@10 metric.

matching for the [CLS] -pooled teacher (cf. Sec. 4.2) and employed the query embedding matching (cf. Eq.7) for the Dual-pooled CE teacher. In either case, embedding-matching loss is added on top of the standard cross entropy loss with the weight of 1 . 0 (when used).

Results and discussion. Table 5 evaluates the effectiveness of the dual pooling and the embedding matching for CE to DE distillation. As described in Sec. 4.2, the traditional [CLS] -pooled teacher did not provide any useful embedding for the embedding matching (see Appendix G.2 for the further analysis of the resulting embedding space). However, with the Dual-pooled teacher, embedding matching does boost student's performance.

6 Related work

Here, we position our EmbedDistill work with respect to prior work on distillation and data augmentation for Transformers-based IR models. We also cover prior efforts on aligning representations during distillation for non-IR settings. Unlike our problem setting where the DE student is factorized, these works mainly consider distilling a single large Transformer into a smaller one.

Distillation for IR. Traditional distillation techniques have been widely applied in the IR literature, often to distill a teacher CE model to a student DE model [28, 8]. Recently, distillation from a DE

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

model (with complex late interaction) to another DE model (with inner-product scoring) has also been considered [29, 15]. As for distilling across different model architectures, Lu et al. [31], Izacard and Grave [16] consider distillation from a teacher CE model to a student DE model. Hofstätter et al. [14] conduct an extensive study of knowledge distillation across a wide-range of model architectures. Most existing distillation schemes for IR rely on only teacher scores; by contrast, we propose a geometric approach that also utilizes the teacher embeddings . Many recent efforts [48, 51, 56] show that iterative multi-stage (self-)distillation improves upon single-stage distillation [48, 51, 56]. These approaches use a model from the previous stage to obtain labels [56] as well as mine harder-negatives [60]. We only focus on the single-stage distillation in this paper. Multi-stage procedures are complementary to our work, as one can employ our proposed embedding-matching approach in various stages of such a procedure. Interestingly, we demonstrate in Sec. 5 that our proposed EmbedDistill can successfully benefit from high quality models trained with such complex procedures [50, 63]. In particular, our single-stage distillation method can transfer almost all of their performance gains to even smaller models. Also to showcase that our method brings gain orthogonal to how teacher was trained, we conduct experiments with single-stage trained teacher in Appendix F.5.

Distillation with representation alignments. Outside of the IR context, a few prior works proposed to utilize alignment between hidden layers during distillation [52, 55, 18, 1, 64]. Chen et al. [7] utilize the representation alignment to re-use teacher's classification layer for image classification. Unlike these works, our work is grounded in a rigorous theoretical understanding of the teacher-student (generalization) gap for IR models. Further, our work differs from these as it needs to address multiple challenges presented by an IR setting: 1) cross-architecture distillation such as CE to DE distillation; 2) partial representation alignment of query or document representations as opposed to aligning for the entire input, i.e., a query-documents pair; and 3) catering representation alignment approach to novel IR setups such as asymmetric DE configuration. To the best of our knowledge, our work is first in the IR literature that goes beyond simply matching scores (or its proxies) for distillation.

Semi-supervised learning for IR. Data augmentation or semi-supervised learning has been previously used to ensure data efficiency in IR [see, e.g., 35, 66]. More interestingly, data augmentation have enabled performance improvements as well. Doc2query [45, 44] performs document expansion by generating queries that are relevant to the document and appending those queries to the document. Query expansion has also been considered, e.g., for document re-ranking [67]. Notably, generating synthetic (query, passage, answer) triples from a text corpus to augment existing training data for QA systems also leads to significant gains [2, 47]. Furthermore, even zero-shot approaches, where no labeled query-document pairs are used, can also perform competitively to supervised methods [26, 17, 33, 54]. Unlike these works, we utilize query-generation capability to ensure tighter alignment between the embedding spaces of the teacher and student.

Richer transformers-based architectures for IR. Besides DE and CE models (cf. Sec. 2), intermediate configurations [36, 22, 42, 32] have been proposed. Such models independently encode query and document before applying a more complex late interaction between the two. Nogueira et al. [46] explore generative encoder-decoder style model for re-ranking. In this paper, we focus on basic DE/CE models to showcase the benefits of our proposed geometric distillation approach. Exploring embedding matching for aforementioned architectures is an interesting avenue for future work.

7 Conclusion

Wepropose EmbedDistill -anovel distillation method for IR that goes beyond simple score matching. En route, we provide a theoretical understanding of the teacher-student generalization gap in an IR setting which not only motivated EmbedDistill but also inspired new design choices for the student DE models: (a) reusing the teacher's document encoder in the student and (b) aligning query embeddings of the teacher and student. This simple approach delivers consistent quality and computational gains in practical deployments and we demonstrate them on MSMARCO, NQ, and BEIR benchmarks. Finally, we found EmbedDistill retains 95-97% of the teacher performance to with 1/10th size students.

Limitations. As discussed in Sec. 4.2 and 5.3, EmbedDistill requires modifications in the CE scoring function to be effective. In terms of underlying IR model architectures, we only explore Transformerbased models in our experiments; primarily due to their widespread utilization. That said, we expect our results to extend to non-Transformer architectures such as MLPs. Finally, we note that our experiments only consider NLP domains, and exploring other modalities (e.g., vision) or multi-modal settings (e.g., image-to-text search) is left as an interesting avenue for future work.

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

References

Gustavo Aguilar, Yuan Ling, Yu Zhang, Benjamin Yao, Xing Fan, and Chenlei Guo. Knowledge distillation from internal representations. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 7350-7357, 2020.

Chris Alberti, Daniel Andor, Emily Pitler, Jacob Devlin, and Michael Collins. Synthetic QA corpora generation with roundtrip consistency. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 6168-6173, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1620. URL https://aclanthology.org/P 19 -1620 .

Yoshua Bengio and Jean-SÉbastien Senecal. Adaptive importance sampling to accelerate training of a neural probabilistic language model. IEEE Transactions on Neural Networks , 19 (4):713-722, 2008. doi: 10.1109/TNN.2007.912312.

Olivier Bousquet, Stéphane Boucheron, and Gábor Lugosi. Introduction to Statistical Learning Theory , pages 169-207. Springer Berlin Heidelberg, Berlin, Heidelberg, 2004. ISBN 978-3-540-28650-9. doi: 10.1007/978-3-540-28650-9_8. URL https://doi.org/ 10 . 1007 / 978 -3 -540 -28650 -9 _ 8 .

Cristian Bucilˇ a, Rich Caruana, and Alexandru Niculescu-Mizil. Model compression. In Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , KDD '06, pages 535-541, New York, NY, USA, 2006. ACM.

Danqi Chen, Adam Fisch, Jason Weston, and Antoine Bordes. Reading Wikipedia to answer open-domain questions. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1870-1879, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10.18653/v1/P17-1171. URL https://aclanthology.org/P 17 -1171 .

Defang Chen, Jian-Ping Mei, Hailin Zhang, Can Wang, Yan Feng, and Chun Chen. Knowledge distillation with the reused teacher classifier. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11933-11942, 2022.

Xuanang Chen, Ben He, Kai Hui, Le Sun, and Yingfei Sun. Simplified tinybert: Knowledge distillation for document retrieval. In Djoerd Hiemstra, Marie-Francine Moens, Josiane Mothe, Raffaele Perego, Martin Potthast, and Fabrizio Sebastiani, editors, Advances in Information Retrieval , pages 241-248, Cham, 2021. Springer International Publishing. ISBN 978-3-03072240-1.

Zhuyun Dai and Jamie Callan. Deeper text understanding for IR with contextual neural language modeling. In Benjamin Piwowarski, Max Chevalier, Éric Gaussier, Yoelle Maarek, Jian-Yun Nie, and Falk Scholer, editors, Proceedings of the 42nd International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR 2019, Paris, France, July 21-25, 2019 , pages 985-988. ACM, 2019.

Zhuyun Dai and Jamie Callan. Context-aware sentence/passage term importance estimation for first stage retrieval. arXiv preprint arXiv:1910.10687 , 2019.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACLHLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers) , pages 4171-4186. Association for Computational Linguistics, 2019.

Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning , 2020. URL https://arxiv.org/abs/ 1908 . 10396 .

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015.

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

Sebastian Hofstätter, Sophia Althammer, Michael Schröder, Mete Sertkan, and Allan Hanbury. Improving efficient neural ranking models with cross-architecture knowledge distillation. CoRR , abs/2010.02666, 2020. URL https://arxiv.org/abs/ 2010 . 02666 .

Sebastian Hofstätter, Sheng-Chieh Lin, Jheng-Hong Yang, Jimmy Lin, and Allan Hanbury. Efficiently teaching an effective dense retriever with balanced topic aware sampling. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval , SIGIR '21, page 113-122, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450380379. doi: 10.1145/3404835.3462891. URL https://doi.org/ 10 . 1145 / 3404835 . 3462891 .

Gautier Izacard and Edouard Grave. Distilling knowledge from reader to retriever for question answering. In International Conference on Learning Representations , 2021. URL https: //openreview.net/forum?id = NTEz6 wysdb .

Gautier Izacard, Mathild Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. arXiv preprint arXiv:2112.09118 , 2021.

Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. TinyBERT: Distilling BERT for natural language understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pages 4163-4174, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.findings-emnlp.372. URL https://aclanthology.org/ 2020 .findings-emnlp. 372 .

Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data , 7(3):535-547, 2021. doi: 10.1109/TBDATA.2019.2921572.

Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 6769-6781, Online, November 2020. Association for Computational Linguistics.

Vladimir Karpukhin, Barlas O˘ guz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04906 , 2020.

Omar Khattab and Matei Zaharia. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT , page 39-48. Association for Computing Machinery, New York, NY, USA, 2020. ISBN 9781450380164.

Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics , 7:452-466, 2019. doi: 10.1162/tacl_a_00276. URL https://aclanthology.org/Q 19 -1026 .

Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics , 7:453-466, 2019.

Michel Ledoux and Michel Talagrand. Probability in Banach spaces . Springer-Verlag, 1991.

Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Anna Korhonen, David R. Traum, and Lluís Màrquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers , pages 6086-6096. Association for Computational Linguistics, 2019.

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequence-to-sequence pretraining for natural language generation, translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics , pages 7871-7880, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.acl-main. 703. URL https://aclanthology.org/ 2020 .acl-main. 703 .

Canjia Li, Andrew Yates, Sean MacAvaney, Ben He, and Yingfei Sun. Parade: Passage representation aggregation for document reranking. arXiv preprint arXiv:2008.09093 , 2020.

Sheng-Chieh Lin, Jheng-Hong Yang, and Jimmy Lin. In-batch negatives for knowledge distillation with tightly-coupled teachers for dense retrieval. In Proceedings of the 6th Workshop on Representation Learning for NLP (RepL4NLP-2021) , pages 163-173, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.repl4nlp-1.17. URL https://aclanthology.org/ 2021 .repl 4 nlp1 . 17 .

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

Wenhao Lu, Jian Jiao, and Ruofei Zhang. Twinbert: Distilling knowledge to twin-structured compressed bert models for large-scale retrieval. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management , CIKM '20, page 2645-2652, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 9781450368599. doi: 10.1145/3340531.3412747. URL https://doi.org/ 10 . 1145 / 3340531 . 3412747 .

Yi Luan, Jacob Eisenstein, Kristina Toutanova, and Michael Collins. Sparse, dense, and attentional representations for text retrieval. Transactions of the Association for Computational Linguistics , 9:329-345, 2021. doi: 10.1162/tacl_a_00369. URL https://aclanthology.org/ 2021 .tacl1 . 20 .

Ji Ma, Ivan Korotkov, Yinfei Yang, Keith Hall, and Ryan McDonald. Zero-shot neural passage retrieval via domain-targeted synthetic question generation. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume , pages 1075-1088, Online, April 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.eacl-main.92. URL https://aclanthology.org/ 2021 .eacl-main. 92 .

Sean MacAvaney, Andrew Yates, Arman Cohan, and Nazli Goharian. CEDR: Contextualized embeddings for document ranking. In Proceedings of the 42nd International ACM SIGIR Conference on Research and Development in Information Retrieval , SIGIR'19, page 1101-1104, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450361729. doi: 10.1145/3331184.3331317. URL https://doi.org/ 10 . 1145 / 3331184 . 3331317 .

Sean MacAvaney, Andrew Yates, Kai Hui, and Ophir Frieder. Content-based weak supervision for ad-hoc re-ranking. In Proceedings of the 42nd International ACM SIGIR Conference on Research and Development in Information Retrieval , SIGIR'19, page 993-996, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450361729. doi: 10.1145/ 3331184.3331316. URL https://doi.org/ 10 . 1145 / 3331184 . 3331316 .

Sean MacAvaney, Franco Maria Nardini, Raffaele Perego, Nicola Tonellotto, Nazli Goharian, and Ophir Frieder. Efficient Document Re-Ranking for Transformers by Precomputing Term Representations , page 49-58. Association for Computing Machinery, New York, NY, USA, 2020. ISBN 9781450380164.

Aditya Menon, Sadeep Jayasumana, Ankit Singh Rawat, Seungyeon Kim, Sashank Reddi, and Sanjiv Kumar. In defense of dual-encoders for neural ranking. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 15376-15400. PMLR, 17-23 Jul 2022. URL https: //proceedings.mlr.press/v 162 /menon 22 a.html .

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

Bhaskar Mitra and Nick Craswell. An introduction to neural information retrieval. Foundations and Trends® in Information Retrieval , 13(1):1-126, 2018. ISSN 1554-0669. doi: 10.1561/ 1500000061. URL http://dx.doi.org/ 10 . 1561 / 1500000061 .

Arvind Neelakantan, Tao Xu, Raul Puri, Alec Radford, Jesse Michael Han, Jerry Tworek, Qiming Yuan, Nikolas Tezak, Jong Wook Kim, Chris Hallacy, et al. Text and code embeddings by contrastive pre-training. arXiv preprint arXiv:2201.10005 , 2022.

Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. MS MARCO: A human generated machine reading comprehension dataset. In Tarek Richard Besold, Antoine Bordes, Artur S. d'Avila Garcez, and Greg Wayne, editors, Proceedings of the Workshop on Cognitive Computation: Integrating neural and symbolic approaches 2016 , volume 1773 of CEUR Workshop Proceedings . CEUR-WS.org, 2016.

Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernandez Abrego, Ji Ma, Vincent Zhao, Yi Luan, Keith Hall, Ming-Wei Chang, and Yinfei Yang. Large dual encoders are generalizable retrievers. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 9844-9855, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https://aclanthology.org/ 2022 .emnlp-main. 669 .

Ping Nie, Yuyu Zhang, Xiubo Geng, Arun Ramamurthy, Le Song, and Daxin Jiang. DC-BERT: decoupling question and document for efficient contextual encoding. In Jimmy Huang, Yi Chang, Xueqi Cheng, Jaap Kamps, Vanessa Murdock, Ji-Rong Wen, and Yiqun Liu, editors, Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval, SIGIR 2020, Virtual Event, China, July 25-30, 2020 , pages 1829-1832. ACM, 2020. doi: 10.1145/3397271.3401271. URL https://doi.org/ 10 . 1145 / 3397271 . 3401271 .

Rodrigo Nogueira and Kyunghyun Cho. Passage re-ranking with BERT. CoRR , abs/1901.04085, 2019. URL http://arxiv.org/abs/ 1901 . 04085 .

Rodrigo Nogueira, Jimmy Lin, and AI Epistemic. From doc2query to doctttttquery. Online preprint , 6, 2019.

Rodrigo Nogueira, Wei Yang, Jimmy Lin, and Kyunghyun Cho. Document expansion by query prediction. arXiv preprint arXiv:1904.08375 , 2019.

Rodrigo Nogueira, Zhiying Jiang, Ronak Pradeep, and Jimmy Lin. Document ranking with a pretrained sequence-to-sequence model. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pages 708-718, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.findings-emnlp.63. URL https: //aclanthology.org/ 2020 .findings-emnlp. 63 .

Barlas O˘ guz, Kushal Lakhotia, Anchit Gupta, Patrick Lewis, Vladimir Karpukhin, Aleksandra Piktus, Xilun Chen, Sebastian Riedel, Wen-tau Yih, Sonal Gupta, et al. Domain-matched pre-training tasks for dense retrieval. arXiv preprint arXiv:2107.13602 , 2021.

Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. RocketQA: An optimized training approach to dense passage retrieval for open-domain question answering. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tür, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021 , pages 5835-5847. Association for Computational Linguistics, 2021.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 21(140):1-67, 2020. URL http://jmlr.org/papers/v 21 / 20 -074 .html .

Nils Reimers, Iryna Gurevych, and Iryna Gurevych. Sentence-BERT: Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing . Association for Computational Linguistics, 11 2019. URL http://arxiv.org/abs/ 1908 . 10084 .

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

Ruiyang Ren, Yingqi Qu, Jing Liu, Wayne Xin Zhao, Qiaoqiao She, Hua Wu, Haifeng Wang, and Ji-Rong Wen. Rocketqav2: A joint training method for dense passage retrieval and passage re-ranking. In Proceedings of EMNLP , 2021.

Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint arXiv:1412.6550 , 2014.

Devendra Sachan, Mostofa Patwary, Mohammad Shoeybi, Neel Kant, Wei Ping, William L. Hamilton, and Bryan Catanzaro. End-to-end training of neural retrievers for open-domain question answering. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) , pages 6648-6662, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.519. URL https: //aclanthology.org/ 2021 .acl-long. 519 .

Devendra Singh Sachan, Mike Lewis, Mandar Joshi, Armen Aghajanyan, Wen-tau Yih, Joelle Pineau, and Luke Zettlemoyer. Improving passage retrieval with zero-shot question generation. arXiv preprint arXiv:2204.07496 , 2022.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 , 2019.

Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction. CoRR , abs/2112.01488, 2021.

Nandan Thakur, Nils Reimers, Andreas Rücklé, Abhishek Srivastava, and Iryna Gurevych. BEIR: A heterogeneous benchmark for zero-shot evaluation of information retrieval models. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2) , 2021. URL https://openreview.net/forum?id = wCu 6 T 5 xFjeJ .

Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv:1908.08962 , 2019.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems , NIPS'17, page 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.

Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul N. Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id = zeFrfgyZln .

Nishant Yadav, Nicholas Monath, Rico Angell, Manzil Zaheer, and Andrew McCallum. Efficient nearest neighbor search for cross-encoder models using matrix factorization. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 2171-2194, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https://aclanthology.org/ 2022 .emnlp-main. 140 .

Zeynep Akkalyoncu Yilmaz, Wei Yang, Haotian Zhang, and Jimmy Lin. Cross-domain modeling of sentence-level evidence for document retrieval. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 3490-3496, Hong Kong, China, November 2019. Association for Computational Linguistics.

Hang Zhang, Yeyun Gong, Yelong Shen, Jiancheng Lv, Nan Duan, and Weizhu Chen. Adversarial retriever-ranker for dense text retrieval. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id = MR 7 XubKUFB .

Linfeng Zhang and Kaisheng Ma. Improve object detection with feature-based knowledge 657 distillation: Towards accurate and efficient detectors. In International Conference on Learning 658 Representations , 2020. 659

660

661

662

663

664

665

666

667

668

669

670

671

672

Shuai Zhang, Lina Yao, Aixin Sun, and Yi Tay. Deep learning based recommender system: A survey and new perspectives. ACM Comput. Surv. , 52(1), feb 2019. ISSN 0360-0300. doi: 10.1145/3285029. URL https://doi.org/ 10 . 1145 / 3285029 .

Chen Zhao, Chenyan Xiong, Jordan Boyd-Graber, and Hal Daumé III. Distantly-supervised dense retrieval enables open-domain question answering without evidence annotation. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 9612-9622, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.emnlp-main.756. URL https://aclanthology.org/ 2021 .emnlp-main. 756 .

Zhi Zheng, Kai Hui, Ben He, Xianpei Han, Le Sun, and Andrew Yates. BERT-QE: Contextualized Query Expansion for Document Re-ranking. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pages 4718-4728, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.findings-emnlp.424. URL

https://aclanthology.org/ 2020 .findings-emnlp. 424 . 673

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n2. The experimental findings presented in this study reveal certain positive outcomes when applying EmbedDistill to IR models.\n\nWEAKNESSES:\n1. 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.\n2. 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.\n3. 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. \n\n[1] Jiao, Xiaoqi, et al. \"Tinybert: Distilling bert for natural language understanding.\" arXiv preprint arXiv:1909.10351 (2019).  \n[2] Romero, Adriana, et al. \"Fitnets: Hints for thin deep nets.\" arXiv preprint arXiv:1412.6550 (2014).  \n[3] Hofstätter, Sebastian, et al. \"Improving efficient neural ranking models with cross-architecture knowledge distillation.\" arXiv preprint arXiv:2010.02666 (2020).\n\nQUESTIONS:\nMy main concern lies in the novelty of this work. Please refer to the weaknesses.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The authors show theoretically that embedding matching can close the teacher-student generalization gap in IR setting, and support the use of asymmetric conﬁgurations.\n\nWEAKNESSES:\n- 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.\n    - 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.\n- 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.\n- The document needs improvement in terms of writing quality.\n    - Section 5, which covers the result analysis, is difficult to follow.\n        - The authors use too many parentheses to add additional explanations, which interrupts the reading flow.\n        - The captions of the tables contain too much information.\n        - 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.\n    - The author uses too much italics in the introduction, which obscures the key points of the paper.\n\nQUESTIONS:\nPlease explain Figure 4 in Appendix G.\n\nLIMITATIONS:\nAs pointed out in the Weaknesses, EmbedDistill does not seem to work in CE to DE distillation, which limits the contribution.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to learn dense neural IR models by distilling not only the teacher scores but also by aligning the learned representations.\n\nThe 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.\n\nExperiments 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.\n\nThe experiments are conducted on the Natural Questions and MS-Marco datasets.\n\nSTRENGTHS:\nThis 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.\n\nThere is also a nice result showing that just using embedding alignment (without the other part of the distillation loss) words pretty well.\n\nThe 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).\n\nWEAKNESSES:\nIn 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.\n\nThe 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.\n\nThe discussion 196-206 is not really needed (it is quite obvious that aligning with the `[CLS]` does not make sense).\n\nQUESTIONS:\n1. Please clarify what is the $u$-covering number and how important it is in the case of the problem\n\n2. what is “inheriting the document tower” (l. 189)\n\n3. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nThis 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.\n\nTo 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).\n\nSTRENGTHS:\n* In-depth theoretical analysis of the teacher-student generalization gap in IR models, which inspired the embedding alignment methods.\n* 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. \n* Overall well written and easy to follow.\n\nWEAKNESSES:\nOverall, 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.\n\nQUESTIONS:\n* Q1: Has aligning other layers of the network being considered/explored?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\n- 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. \n- 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.\n- The paper is well-written, well-organized, and easy to follow. The authors provide clear explanations of the proposed approach and the experimental setup.\n\nWEAKNESSES:\nAlthough 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. \n\nThe 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.\n\nQUESTIONS:\nRefer to the Weaknesses\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}