Rescaling MLM-Head for Neural Sparse Retrieval

2026-06-17

AI Summary

This work addresses the instability and occasional collapse observed when employing strong pretrained encoders—such as ModernBERT and Ettin—for neural sparse retrieval, which stems from scale mismatch in the masked language modeling (MLM) head outputs. To resolve this issue without altering model architecture or training objectives, the authors propose a zero-overhead MLM head initialization rescaling strategy that calibrates the output scale via a constant factor, substantially improving training stability. Integrated within the SPLADE framework and combined with contrastive learning and sparse activation mechanisms, the approach yields significant retrieval performance gains on both in-domain and out-of-domain benchmarks. Notably, under certain configurations, it even surpasses the established BERT-SPLADE baseline, highlighting MLM head scale calibration as a critical bottleneck in adapting powerful encoders to sparse retrieval tasks.

Abstract

Learned sparse retrieval (LSR) models such as SPLADE have traditionally used BERT-style masked language models 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.

Problem

• sparse retrieval
• MLM-head scale
• SPLADE
• pretrained encoders
• training stability

Innovation

• sparse retrieval
• MLM-head rescaling
• SPLADE
• pretrained encoder adaptation
• training stability