Rescaling MLM-Head for Neural Sparse Retrieval

Technical Summary: Rescaling MLM-Head for Neural Sparse Retrieval

Problem Statement
Learned Sparse Retrieval (LSR) models, particularly SPLADE, rely on the Masked Language Modeling (MLM) head of a pretrained encoder to generate sparse lexical representations. While it is intuitive that replacing the standard BERT backbone with stronger modern encoders (e.g., ModernBERT, Ettin, RoBERTa) should improve retrieval performance, the authors observe a counterintuitive failure: under standard SPLADE training recipes, backbones with large MLM-head L2 norms often suffer from severe performance degradation or complete training collapse.

The paper identifies the root cause as a scale mismatch. In SPLADE, sparse vectors are constructed directly from the MLM-head outputs ($H W_{head}^T$) and compared via an unnormalized dot product. When the MLM-head matrix $W_{head}$ has a large norm, it inflates the magnitude of sparse activations and retrieval scores. This inflation distorts the optimization dynamics of the contrastive ranking loss and creates tension with sparsity regularization (FLOPS), leading to unstable training and ineffective retrieval.

Methodology
To address this scale mismatch without altering the model architecture or training objective, the authors propose a simple, zero-cost initialization-time correction: MLM-Head Rescaling.

Before training begins, the MLM-head projection matrix is rescaled by a constant factor $k$:
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

Key characteristics of this method include:

• Zero-Cost Adjustment: It introduces no additional parameters, requires no architectural changes, and adds no computational overhead during training or inference.
• Preservation of Information: The operation preserves the direction and lexical information of the pretrained vocabulary projection.
• Applicability: For models with tied MLM heads, the operation is applied to the shared weight matrix to maintain weight tying.

Experimental Setup
The authors evaluated seven encoder backbones (BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT, and Ettin) trained on MS MARCO triplets using a fixed SPLADE recipe. They tested various rescaling factors ($k$) and evaluated performance on:

• In-domain: MS MARCO Dev and TREC 2019.
• Out-of-domain: BEIR-13 (mean nDCG@10).
• Sparsity-Efficiency Trade-off: Analyzed via NanoBEIR performance against activated FLOPs.

Key Results

1. Recovery of Large-Norm Backbones: Standard initialization ($k=1$) caused modern encoders with large norms to underperform significantly. For instance, ModernBERT dropped to a BEIR-13 nDCG@10 of 0.127 (compared to BERT's 0.370). Applying rescaling ($k=16$) recovered ModernBERT's performance to 0.405, surpassing the BERT baseline. Similarly, Ettin improved from 0.221 to 0.391.
2. Non-Monotonic Sensitivity: The optimal rescaling factor depends on the backbone's initial norm. While ModernBERT and Ettin required strong rescaling ($k=16$), RoBERTa peaked at $k=2$. Excessive rescaling (e.g., $k=16$ for RoBERTa) or rescaling already small-norm models (like ALBERT) led to performance collapse.
3. Training Stability: Rescaling significantly stabilized training dynamics. Unrescaled ModernBERT and Ettin models started with extremely high loss values and converged to poor local minima. Rescaled variants exhibited smoother loss curves and lower final losses.
4. Improved Trade-offs: Rescaled models achieved a better effectiveness-sparsity trade-off, operating in the high-effectiveness, low-FLOPs region typically occupied by successful BERT-based SPLADE models.

Significance and Claims
The paper argues that the bottleneck in adapting pretrained encoders to Learned Sparse Retrieval is not merely encoder capacity or pretraining quality, but the calibration of the MLM-head scale.

The authors claim that:

• The scale of the MLM head is a critical, previously overlooked factor in LSR.
• Stronger encoders are not automatically "drop-in" replacements for BERT in SPLADE; their projection scales must be calibrated to the specific dynamics of sparse retrieval.
• A simple rescaling correction can transform unstable training runs of modern encoders into competitive sparse retrievers, often matching or surpassing the classic BERT-SPLADE baseline.

This work suggests that future adaptations of pretrained encoders to LSR must account for the geometry and scale of the vocabulary projection layer, not just the encoder's representational power.