COSTA Contrastive Span Prediction Pre-Training for Dense Retrieval

A decoder-free self-supervised pre-training algorithm for dense text retrieval that samples T spans at each of four granularities (word, phrase, sentence, and paragraph) per input document d_i using Beta-distributed span lengths p_span ~ Beta(alpha, beta), l_span = p_span (l_max - l_min) + l_min, encodes the whole document d_i = ([CLS], x_1, ..., x_n) in a single forward pass of a bidirectional Transformer encoder into token hidden states h_0, h_1, ..., h_n, transforms the whole-text [CLS] representation through a nonlinear projector z_i = Tanh(FFN(h_0)) while pooling each sampled span p in S(i) via token average pooling z_p = AvgPooling([h_start, ..., h_{end-1}]), and jointly optimizes Masked Language Modeling L_MLM and a Group-Wise Contrastive loss L_GWC = sum_{i=1}^N -(1/(4T)) sum_{p in S(i)} log(exp(sim(z_i, z_p)/tau) / sum_{j=1}^{N(4T+1)} 1_{[i != j]} exp(sim(z_i, z_j)/tau)) that pulls each whole-text representation toward all 4T of its own multi-granularity spans while pushing it away from all other whole-text and span representations in the mini-batch.