Rescaling MLM-Head for Neural Sparse Retrieval

Abstract: Learned sparse retrieval (LSR) models such as SPLADE have traditionally used BERT-style masked LLMs as backbone encoders. A natural expectation is that replacing BERT with stronger pretrained encoders should improve retrieval effectiveness. However, we find that under standard SPLADE training recipes, backbones with large MLM-head L2 norms can suffer performance degradation and even training collapse under standard SPLADE training recipes. We identify this failure as a scale mismatch in the MLM head: SPLADE directly uses MLM-head outputs to construct sparse lexical representations, and query-document relevance is computed by an unnormalized dot product over these representations. As a result, an inflated MLM-head scale can amplify sparse activations, distort matching scores, and destabilize contrastive training under common training settings. To address this issue, we introduce a simple initialization-time correction that rescales the MLM-head projection by a constant factor before SPLADE training. This zero-cost adjustment improves training stability without modifying the model architecture or training objective. Across both in-domain and out-of-domain retrieval benchmarks, this simple correction substantially improves large-norm backbones such as ModernBERT and Ettin, turning unstable training runs into competitive sparse retrievers. In several settings, the corrected models further match or surpass the classic BERT-SPLADE baseline. These findings suggest that the bottleneck in adapting pretrained encoders to LSR is not encoder capacity alone, but the calibration of the MLM-head scale used to construct sparse lexical representations.

• The paper demonstrates that a simple initialization-time rescaling of the MLM-head effectively addresses training collapse and boosts retrieval performance.
• Empirical analysis on benchmarks like BEIR-13 shows that large-norm backbones achieve up to 215% relative performance gain after proper rescaling.
• The findings underscore that calibration of MLM-head scale is crucial for leveraging modern Transformer encoders in sparse retrieval.

Rescaling MLM-Head for Neural Sparse Retrieval: Technical Summary and Analysis

Introduction

This paper rigorously examines the interplay between masked language modeling (MLM) head initialization and the effectiveness of learned sparse retrieval (LSR) models, specifically SPLADE, when built on modern Transformer encoder architectures. While recent pretrained encoders such as ModernBERT and Ettin offer advances in architecture and pretraining, direct transfer to LSR under standard SPLADE recipes often results in subpar retrieval effectiveness, training collapse, and unstable optimization. Through systematic experimentation, the authors demonstrate that the culprit is the scale of the MLM-head projection. Large L2 norms in the MLM-head matrix destabilize sparse representation and contrastive training; consequently, SPLADE fails to leverage the purported benefits of modern encoders without an explicit calibration step.

Problem Analysis: The Role of MLM-Head Norm in SPLADE Failures

Unlike dense retrievers, SPLADE directly utilizes the vocabulary projection head (MLM head) of the backbone encoder to construct sparse lexical representations. The final retrieval score in SPLADE is computed as an unnormalized dot product in the sparse vocabulary space. As the paper shows, encoders with large MLM-head norms (e.g., ModernBERT, Ettin, RoBERTa) produce magnified activations, leading to uncontrolled, inflated retrieval scores and an unstable optimization landscape, especially in the presence of the FLOPS regularizer.

Figure 1: Large-norm MLM-heads in modern backbones empirically degrade zero-shot and fine-tuned retrieval performance for SPLADE compared to small-norm backbones.

Empirical results concretely demonstrate this trend: BERT, DistilBERT, and GTE-MLM (small-norm) yield effective SPLADE models, while ModernBERT and Ettin (large-norm) underperform despite stronger pretraining and architecture, contradicting the prevalent drop-in-backbone assumption in the community.

Proposed Solution: Initialization-Time MLM-Head Rescaling

The authors propose a zero-cost, initialization-time correction—scaling the MLM-head matrix by a constant factor before SPLADE training begins. This approach preserves the geometry of the pretrained mapping and lexical coverage, while controlling the absolute scale of projected representations. Importantly, no architecture changes, parameter increases, or extra computation during training or inference are required.

Figure 2: Effectiveness versus sparsity trade-off on NanoBEIR as a function of different FLOPS regularization strengths; rescaled models recover both efficiency and effectiveness.

This simple rescaling corrects the scale mismatch in the initialization regime, allowing large-norm backbones to converge to competitive retrieval effectiveness, in line with or exceeding the classical BERT-SPLADE baseline.

Quantitative Results

Experiments are conducted across MSMARCO, TREC DL-2019, and BEIR-13 (zero-shot) benchmarks. Key results include:

• For ModernBERT, rescaling (with $k=16$) yields a 215% relative gain on BEIR-13 nDCG@10 (from 0.127 to 0.405).
• Ettin backbone achieves a 77% increase, and RoBERTa regains competitive performance under moderate rescaling.
• Small-norm backbones (e.g., ALBERT) are degraded by further rescaling, highlighting the requirement for an optimal scale range.

  Figure 3: Retrieval effectiveness (BEIR) as a function of MLM-head rescaling factor $k$; optimal gains for large-norm models, harmful for already small-norm models.

  

  Figure 4: Loss curves across training for ModernBERT and Ettin-based SPLADE backbones show rescaling stabilizes training and accelerates convergence.

The work also analyzes the FLOPS-effectiveness trade-off, showing that rescaled models both reduce sparsity overhead and improve retrieval quality. Notably, aggressive downscaling is not universally beneficial—the calibration must bring the head norm into a task-appropriate regime (neither too large nor too small).

Theoretical and Practical Implications

The findings fundamentally challenge the notion that encoder capacity is the limiting factor in LSR adaptation, exposing the importance of normalization and scale calibration when transferring pretrained encoders to sparse lexical spaces. This has several implications:

• Model selection for LSR should explicitly examine MLM-head scale, not merely encoder test set accuracy or perplexity.
• Simple rescaling at initialization time can unlock the benefits of new encoders for sparse retrieval, enabling the rapid evaluation of novel architectures and pretraining recipes in LSR.
• The approach generalizes to any architecture utilizing direct projection to a predefined lexical space for scoring.
• For future LSR systems and toolkit design, including automatic MLM-head norm normalization is necessary to avoid silent degradation and training collapse.

Future Directions

Several future avenues stem from this result:

• Systematic analytic characterization of the optimal head scale as a function of dataset, regularization strength, and backbone architecture.
• Extension to multilingual, cross-modal, and non-lexical sparse retrieval settings, particularly where the mapping is not directly inherited from MLM pretraining.
• Exploration of dynamic scale adaptation or trainable norm constraints as opposed to global static rescaling.
• Integration with recent developments in vocabulary selection, expansion, and hybrid sparse-dense modeling.

Conclusion

This paper demonstrates that the standard SPLADE recipe, when applied naïvely to strong modern encoders, often fails due to an easily remedied initialization mismatch: the scale of the MLM-head matrix. A simple, zero-cost rescaling correction—applied only at initialization—restores training stability and enables substantial improvements in LSR effectiveness for large-norm backbones. The results reframe the core challenge in adapting new encoders for LSR: beyond capacity or pretraining, calibration of projection scale is critical for unlocking the potential of modern lexical encoders.