The Pre-Training Study of Expanded-SPLADE Models on Web Document Titles

Technical Summary: The Pre-Training Study of Expanded-SPLADE Models on Web Document Titles

Problem Statement
Masked Language Modeling (MLM) pre-training is a standard initialization method for Neural Information Retrieval (IR) models. However, prior research indicates that MLM pre-trained models often exhibit limited readiness and transfer learning issues when fine-tuned into Neural Bi-Encoder models for retrieval tasks. Specifically, the objectives of MLM (intra-instance prediction) and IR ranking (inter-instance discrimination) are not always aligned. While Learned Sparse Retrieval (LSR) models like SPLADE utilize the MLM layer during fine-tuning to generate sparse representations, the optimal pre-training strategy for these models—particularly Expanded-SPLADE (ESPLADE) models which employ custom, high-dimensional output vocabularies—remains under-explored. This paper investigates the impact of different pre-training datasets and hyperparameters on the downstream retrieval effectiveness and efficiency of ESPLADE models.

Methodology
The authors conducted a series of experiments involving the pre-training of Expanded Masked Language Models (EMLM) and their subsequent fine-tuning into ESPLADE models. The study utilized in-house web document titles as the primary data source.

• Model Architecture: The study employs the ESPLADE architecture (Dudek et al., 2023), which reconstructs the output MLM fully connected layer to support a custom, expanded vocabulary (100K unigrams). This allows for a larger dimensionality than standard BERT vocabularies.
• Pre-Training Configurations: The authors varied two primary factors:
  1. Datasets:
     • Overlap-Repeat: Pre-trained on 614M web titles derived from the fine-tuning training set (trainset-small), with repeated masking patterns.
     • Indep-Repeat: Pre-trained on 614M records from a distinct in-house collection (79M unique titles), using a vocabulary aligned with the fine-tuning set.
     • Indep-Uniq: Pre-trained on 614M unique web titles from the in-house collection.
  2. Learning Rates and Steps:
     • Low Learning Rate (-lr-l): Trained for 600K steps.
     • High Learning Rate (-lr-h): Trained for 2.9M steps but checkpointed at 600K steps, resulting in a higher "effective" learning rate compared to the low-rate models.
• Fine-Tuning and Evaluation: Models were fine-tuned on a search log containing query-document clicks using in-batch negative loss and FLOPS regularization. Evaluation was conducted on a held-out test set with optional static pruning of sparse vectors (query and document terms) based on MLM logit scores to simulate strict efficiency constraints.
• Metrics: Retrieval effectiveness was measured via MRR@10 and R@10. Efficiency was approximated using a FLOPS metric based on the traversal cost of postings lists. The study also analyzed the statistical properties of EMLM logit vectors (e.g., logit-score-std) and postings list length distributions.

Key Results
The experimentation yielded three primary observations:

1. Pre-Training Data and Learning Rate Impact: Fine-tuned models achieving the highest retrieval effectiveness in both unpruned and strictly pruned settings were predominantly those pre-trained on general, independent corpora (indep-) rather than the fine-tuning corpus (overlap-). Furthermore, models pre-trained with a higher effective learning rate (-lr-h) generally outperformed those with lower rates. Notably, these high-performing models exhibited lower MLM accuracy and higher logit-score-std (variance in logit scores) during pre-training, suggesting that high MLM accuracy does not necessarily correlate with retrieval readiness.
2. Trade-off in Strict Pruning: In the most strict pruning settings (e.g., keeping only the top 5 query terms and 10 document terms), the high-performing models (indep- with -lr-h) maintained better retrieval effectiveness but incurred higher retrieval costs (FLOPS) and exhibited higher variance in the length of individual postings lists. This indicates a trade-off where maintaining effectiveness under strict sparsity constraints requires a less uniform distribution of term frequencies, leading to higher computational traversal costs.
3. Dataset Repetition: Repeating the general pre-training dataset (comparing indep-repeat vs. indep-uniq) had minimal impact on retrieval effectiveness. This suggests that the diversity of the corpus for the MLM classification task is not a critical factor for ranking fine-tuning when FLOPS regularization is applied.

Significance and Claims
The paper empirically identifies potential limitations in aligning standard MLM pre-training with ESPLADE fine-tuning. The authors claim that:

• Misalignment of Objectives: The findings reinforce the hypothesis that MLM pre-training and IR ranking fine-tuning are not always aligned. Models that perform well on MLM tasks (high accuracy, low variance) may not be optimal for retrieval, whereas models with higher logit-score variance and lower MLM accuracy (often resulting from general corpus pre-training with higher learning rates) demonstrate better adaptability for ESPLADE fine-tuning.
• Efficiency-Effectiveness Trade-off: The study provides empirical evidence that in strict pruning scenarios, maintaining retrieval effectiveness often necessitates higher retrieval costs and a higher variance in postings list lengths. This trade-off is inherent to the pre-training configuration and the resulting sparse representation characteristics.
• Limited Transfer of Contextual Knowledge: The lack of significant performance difference between repeated and unique general datasets implies a limitation (or perhaps a necessity) in transferring contextual knowledge from diverse MLM pre-training corpora to the specific ranking fine-tuning of ESPLADE models.

The authors conclude that while MLM pre-training is useful, the specific configuration (corpus type, learning rate) significantly influences the model's readiness for sparse retrieval fine-tuning, and that standard MLM metrics may not be reliable indicators of downstream retrieval performance.