Rescaling MLM-Head for Neural Sparse Retrieval

技術要約：ニューラル疎な検索のためのMLMヘッドの再スケーリング

問題提起
学習された疎な検索（Learned Sparse Retrieval: LSR）モデル、特にSPLADEは、事前学習済みエンコーダのMasked Language Modeling (MLM) ヘッドを利用して、疎な語彙表現を生成する。標準的なBERTバックボーンをより強力な現代的なエンコーダ（ModernBERT、Ettin、RoBERTaなど）に置き換えることが、検索性能の向上につながることは直感的であるが、著者らはある直感に反する失敗を観察している。すなわち、標準的なSPLADEの訓練レシピの下では、MLMヘッドのL2ノルムが大きいバックボーンは、深刻な性能低下や訓練の完全な崩壊を招くということである。

本論文は、その根本原因が**スケールの不一致（scale mismatch）**にあることを特定している。SPLADEにおいて、疎なベクトルはMLMヘッドの出力（$H W_{head}^T$）から直接構築され、正規化されていないドット積によって比較される。MLMヘッドの行列 $W_{head}$ のノルムが大きい場合、疎な活性化および検索スコアが膨張してしまう。この膨張は、コントラスティブ・ランキング損失の最適化ダイナミクスを歪め、スパース性への正則化（FLOPS）との間に緊張関係を生み出し、訓練の不安定化や効果的な検索の阻害を引き起こす。

手法
モデルのアーキテクチャや訓練目的を変更することなく、このスケールの不一致に対処するために、著者らはシンプルかつコストのかからない初期化時の補正策であるMLM-Head Rescalingを提案する。

訓練を開始する前に、MLMヘッドの投影行列を定数因子 $k$ によって再スケーリングする：
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

本手法の主な特徴は以下の通りである：

• ゼロコストの調整: 追加のパラメータを導入せず、アーキテクチャの変更を必要とせず、訓練中や推論中の計算オーバーヘッドも発生させない。
• 情報の保持: プリトレーニングされた語彙投影の方向性と語彙情報を保持する。
• 適用可能性: 重みがタイイング（共有）されているMLMヘッドを持つモデルの場合、重みのタイイングを維持するために、共有された重み行列に対してこの操作を適用する。

実験設定
著者らは、固定された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)。
• スパース性・効率性のトレードオフ: 活性化された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. 非単調な感度: 最適な再スケーリング因子は、バックボンドの初期ノルムに依存する。ModernBERTやEttinは強力な再スケーリング（$k=16$）を必要としたが、RoBERTaは$k=2$でピークに達した。過度な再スケーリング（例：RoBERTaに対して$k=16$）や、すでにノルムが小さいモデル（ALBERTなど）への再スケーリングは、性能の崩壊を招いた。
3. 訓練の安定性: 再スケーリングは訓練のダイナミクスを大幅に安定させた。再スケーリングされていないModernBERTおよびEttinのモデルは、極めて高い損失値から始まり、不十分な局所解に収束した。再スケーリングされたバリアントは、より滑らかな損失曲線とより低い最終損失を示した。
4. 改善されたトレードオフ: 再スケーリングされたモデルは、成功したBERTベースのSPLADEモデルが通常位置する「高効果・低FLOPs」の領域において、より優れた効果・スパース性のトレードオフを実現した。

意義と主張
本論文は、事前学習済みエンコーダを学習された疎な検索（LSR）に適応させる際のボトルネックは、単なるエンコーダの容量や事前学習の質ではなく、**MLMヘッドのスケールの較正（calibration）**にあると主張している。

著者らは以下のことを主張している：

• MLMヘッドのスケールは、LSRにおける重要かつ、これまで見落とされていた要因である。
• 強力なエンコーダは、必ずしもSPLADEにおけるBERTの単純な「ドロップイン（そのまま置き換え可能）」な代替品ではない。その投影スケールは、疎な検索の特定のダイナミクスに合わせて較正される必要がある。
• シンプルな再スケーリングによる補正は、現代的なエンコーダの不安定な訓練プロセスを、競争力のある疎な検索器へと変貌させ、しばしば古典的なBERT-SPLADEのベースラインに匹賛、あるいはそれを凌駕させることができる。

本研究は、事前学習済みエンコーダをLSRへ適応させる将来の試みにおいては、エンコーダの表現能力だけでなく、語彙投影層の幾何学的構造とスケールを考慮しなければならないことを示唆している。