신경망 희소 검색을 위한 MLM 헤드 재조정

Rescaling MLM-Head for Neural Sparse Retrieval

SPLADE와 같은 학습 기반 희소 검색 모델은 일반적으로 BERT 스타일의 마스크 언어 모델(MLM)을 기본 인코더로 사용합니다. 자연스러운 기대는 더 강력한 사전 훈련된 인코더로 BERT를 대체하면 검색 성능이 향상될 것이라는 점입니다. 그러나, 표준 SPLADE 훈련 방법을 사용할 때, 큰 MLM 헤드 L2 노름 값을 가진 인코더는 성능 저하 및 심지어 훈련 실패를 초래할 수 있다는 것을 발견했습니다. 우리는 이러한 실패의 원인이 MLM 헤드의 스케일 불일치에 있다고 판단했습니다. SPLADE는 MLM 헤드 출력을 직접 사용하여 희소 어휘 표현을 구성하고, 질의-문서 간의 관련성은 이러한 표현에 대한 정규화되지 않은 내적 곱으로 계산됩니다. 결과적으로, 과도하게 확장된 MLM 헤드의 스케일은 희소 활성화를 증폭시키고, 매칭 점수를 왜곡하며, 일반적인 훈련 환경에서 대조 학습을 불안정하게 만들 수 있습니다. 이 문제를 해결하기 위해, 우리는 SPLADE 훈련 전에 MLM 헤드 투영을 일정한 계수로 재조정하는 간단한 초기화 시간 수정 방법을 제안합니다. 이 무작위 비용 조정은 모델 아키텍처나 훈련 목표를 변경하지 않고도 훈련 안정성을 향상시킵니다. 인-도메인 및 아웃-오브-도메인 검색 벤치마크에서, 이 간단한 수정 사항은 ModernBERT 및 Ettin과 같은 큰 노름 값을 가진 인코더의 성능을 크게 향상시켜 불안정한 훈련 과정을 경쟁력 있는 희소 검색기로 전환합니다. 몇 가지 설정에서는, 수정된 모델이 기존 BERT-SPLADE 기준에 부합하거나 능가하는 결과를 보여줍니다. 이러한 결과는 사전 훈련된 인코더를 희소 검색에 적용하는 데 있어서 병목 현상이 단순히 인코더의 용량만이 아니라, 희소 어휘 표현을 구성하는 데 사용되는 MLM 헤드 스케일의 조정에 있다는 것을 시사합니다.

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.