Rescaling MLM-Head for Neural Sparse Retrieval

기술 요약: 신경망 희소 검색(Neural Sparse Retrieval)을 위한 MLM-Head 리스케일링

문제 정의
학습된 희소 검색(Learned Sparse Retrieval, LSR) 모델, 특히 SPLADE는 사전 학습된 인코더의 마스크드 언어 모델링(Masked Language Modeling, MLM) 헤드에 의존하여 희소 어휘 표현을 생성합니다. 기존의 BERT 백본을 더 강력한 현대적 인코더(예: ModernBERT, Ettin, RoBERTa)로 교체하는 것이 검색 성능을 향상시킬 것이라는 점은 직관적이지만, 저자들은 역설적인 실패 현상을 관찰했습니다. 즉, 표준 SPLADE 학습 레시피를 사용할 때, MLM-head의 L2 노름(norm)이 큰 백본들은 심각한 성능 저하나 완전한 학습 붕괴를 겪습니다.

본 논문은 그 근본 원인을 **스케일 불일치(scale mismatch)**로 지목합니다. SPLADE에서 희소 벡터는 MLM-head의 출력($H W_{head}^T$)으로부터 직접 구축되며 비정규화된 내적(unnormalized dot product)을 통해 비교됩니다. 이때 MLM-head 행렬 $W_{head}$의 노름이 크면 희소 활성화(sparse activations)와 검색 점수가 팽창하게 됩니다. 이러한 팽창은 대조 학습 랭킹 손실(contrastive ranking loss)의 최적화 역학을 왜곡하고, 희소성 정규화(FLOPS)와 충돌을 일으켜 학습 불안정 및 비효과적인 검색을 초래합니다.

방법론
모델 아키텍처나 학습 목적 함수를 변경하지 않고 이 스케일 불일치 문제를 해결하기 위해, 저자들은 간단하고 비용이 들지 않는 초기화 시점의 교정 방법인 **MLM-Head 리스케일링(MLM-Head Rescaling)**을 제안합니다.

학습이 시작되기 전, MLM-head 투영 행렬을 상수 $k$에 의해 리스케일링합니다:
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

이 방법의 주요 특징은 다음과 같습니다:

• 비용 제로 조정 (Zero-Cost Adjustment): 추가적인 파라미터를 도입하지 않으며, 아키텍처 변경이 필요 없고, 학습 또는 추론 시 계산 오버헤드를 발생시키지 않습니다.
• 정보 보존 (Preservation of Information): 사전 학습된 어휘 투영의 방향성과 어휘 정보를 그대로 유지합니다.
• 적용 가능성 (Applicability): 결합된(tied) MLM-head를 사용하는 모델의 경우, 가중치 결합(weight tying)을 유지하기 위해 공유된 가중치 행렬에 이 연산을 적용합니다.

실험 설정
저자들은 고정된 SPLADE 레시피를 사용하여 MS MARCO 트리플렛 데이터셋에서 7개의 인코더 백본(BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT, Ettin)을 평가했습니다. 다양한 리스케일링 계수($k$)를 테스트하였으며, 다음 지표로 성능을 평가했습니다:

• 도메인 내 (In-domain): MS MARCO Dev 및 TREC 2019.
• 도메인 외 (Out-of-domain): BEIR-13 (평균 nDCG@10).
• 희소성-효율성 트레이드오프 (Sparsity-Efficiency Trade-off): 활성화된 FLOPs 대비 NanoBEIR 성능을 통해 분석.

주요 결과

1. 대형 노름 백본의 성능 회복: 표준 초기화($k=1$)는 노름이 큰 현대적 인코더들의 성능을 크게 떨어뜨렸습니다. 예를 들어, ModernBERT는 BEIR-13 nDCG@10이 0.127로 급락했습니다(BERT의 0.370과 비교 시). 리스케일링($k=16$)을 적용하자 ModernBERT의 성능은 0.405로 회복되어 BERT 베이스라인을 넘어섰습니다. 마찬가지로 Ettin은 0.221에서 0.391로 개선되었습니다.
2. 비단조적 민감도 (Non-Monotonic Sensitivity): 최적의 리스케일링 계수는 백본의 초기 노름에 따라 달라집니다. ModernBERT와 Ettin은 강한 리스케일링($k=16$)이 필요했던 반면, RoBERTa는 $k=2$에서 정점을 찍었습니다. 과도한 리스케일링(예: RoBERTa에 $k=16$ 적용)이나 이미 노름이 작은 모델(예: ALBERT)에 대한 리스케일링은 성능 붕괴를 초래했습니다.
3. 학습 안정성: 리스케일링은 학습 역학을 크게 안정화했습니다. 리스케일링되지 않은 ModernBERT와 Ettin 모델은 극도로 높은 손실 값으로 시작하여 좋지 않은 로컬 미니마(local minima)로 수렴했습니다. 반면 리스케일링된 변체들은 더 매끄러운 손실 곡선과 더 낮은 최종 손실 값을 보였습니다.
4. 개선된 트레이드오프: 리스케일링된 모델은 성공적인 BERT 기반 SPLADE 모델들이 점유하는 영역인 '높은 효과성, 낮은 FLOPs' 영역에서 작동하며 더 나은 효과성-희소성 트레이드오프를 달 achievement 했습니다.

의의 및 주장
본 논문은 사전 학습된 인코더를 학습된 희소 검색(LSR)에 적응시키는 데 있어 병목 현상이 단순히 인코더의 용량이나 사전 학습 품질 때문이 아니라, MLM-head 스케일의 캘리브레이션(calibration) 때문이라고 주장합니다.

저자들의 주장은 다음과 같습니다:

• MLM head의 스케일은 LSR에서 중요하지만 이전에 간과되었던 요소입니다.
• 더 강력한 인코더가 반드시 SPLADE를 위한 "드롭인(drop-in)" 교체 대상이 되는 것은 아닙니다. 이들의 투영 스케일은 희소 검색의 특정 역학에 맞춰 캘리브레이션되어야 합니다.
• 단순한 리스케일링 교정은 불안정한 ModernBERT/Ettin 학습 과정을 경쟁력 있는 희소 검색기로 변모시켜, 종종 클래식한 BERT-SPLADE 베이스라인과 일치하거나 이를 능가할 수 있습니다.

이 연구는 향후 사전 학습된 인코더를 LSR에 적응시킬 때, 인코더의 표현력뿐만 아니라 어휘 투영 레이어의 기하학적 구조와 스케일을 반드시 고려해야 함을 시사합니다.