Rescaling MLM-Head for Neural Sparse Retrieval

Technische Samenvatting: Rescaling MLM-Head voor Neural Sparse Retrieval

Probleemstelling
Learned Sparse Retrieval (LSR) modellen, in het bijzonder SPLADE, vertrouwen op de Masked Language Modeling (MLM) head van een vooraf getrainde encoder om ijle (sparse) lexicale representaties te genereren. Hoewel het intuïtief is dat het vervangen van de standaard BERT-backbone door sterkere moderne encoders (bijv. ModernBERT, Ettin, RoBERTa) de retrieval-prestaties zou moeten verbeteren, observeren de auteurs een contra-intuïtief falen: onder standaard SPLADE-trainingsrecepten lijden backbones met grote MLM-head L2-normen vaak aan ernstige prestatiedegradatie of volledige trainingsinstabiliteit.

Het onderzoek identificeert de wortel van het probleem als een schaalverschil (scale mismatch). In SPLADE worden ijle vectoren direct geconstrueerd uit de MLM-head outputs ($H W_{head}^T$) en vergeleken via een ongenormaliseerd inproduct. Wanneer de MLM-head matrix $W_{head}$ een grote norm heeft, blaast dit de magnitude van de ijle activaties en retrieval-scores op. Deze inflatie verstoort de optimalisatiedynamiek van de contrastieve ranking loss en creëert spanning met de sparsity-regularisatie (FLOPS), wat leidt tot instabiele training en ineffectieve retrieval.

Methodologie
Om dit schaalverschil aan te pakken zonder de modelarchitectuur of het trainingsdoel te wijzigen, stellen de auteurs een eenvoudige, kosteloze correctie tijdens de initialisatie voor: MLM-Head Rescaling.

Voordat de training begint, wordt de MLM-head projectiematrix geschaald door een constante factor $k$:
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

Belangrijke kenmerken van deze methode zijn:

• Kosteloze Aanpassing: Het introduceert geen extra parameters, vereist geen architecturale wijzigingen en voegt geen computationele overhead toe tijdens training of inferentie.
• Behoud van Informatie: De operatie behoudt de richting en de lexicale informatie van de vooraf getrainde vocabulaire-projectie.
• Toepasbaarheid: Voor modellen met gedeelde (tied) MLM-heads wordt de operatie toegepast op de gedeelde gewichtsmatrix om weight tying te behouden.

Experimentele Opzet
De auteurs evalueerden zeven encoder-backbones (BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT en Ettin) getraind op MS MARCO triplets met een vaste SPLADE-recept. Ze testten verschillende rescaling-factoren ($k$) en evalueerden de prestaties op:

• In-domain: MS MARCO Dev en TREC 2019.
• Out-of-domain: BEIR-13 (gemiddelde nDCG@10).
• Sparsity-Efficiency Trade-off: Geanalyseerd via NanoBEIR prestaties tegenover geactiveerde FLOPs.

Belangrijkste Resultaten

1. Herstel van Large-Norm Backbones: Standaard initialisatie ($k=1$) zorgde ervoor dat moderne encoders met grote normen aanzienlijk onderpresteerden. Zo daalde ModernBERT naar een BEIR-13 nDCG@10 van 0,127 (vergeleken met 0,370 voor BERT). Het toepassen van rescaling ($k=16$) herstelde de prestaties van ModernBERT naar 0,405, waarmee het de BERT-baseline overtrof. Ook Ettin verbeterde van 0,221 naar 0,391.
2. Niet-monotone Gevoeligheid: De optimale rescaling-factor hangt af van de initiële norm van de backbone. Terwijl ModernBERT en Ettin sterke rescaling vereisten ($k=16$), bereikte RoBERTa een piek bij $k=2$. Overmatige rescaling (bijv. $k=16$ voor RoBERTa) of rescaling van reeds kleine-norm modellen (zoals ALBERT) leidde tot prestatie-instorting.
3. Trainingsstabiliteit: Rescaling stabiliseerde de trainingsdynamiek aanzienlijk. Ongerescaleerde ModernBERT en Ettin begonnen met extreem hoge loss-waarden en convergeerden naar slechte lokale minima. Gerescaleerde varianten vertoonden gladdere loss-curves en lagere uiteindelijke losses.
4. Verbeterde Trade-offs: Gerescaleerde modellen bereikten een betere effectiviteit-sparsity trade-off, waarbij ze opereerden in de regio van hoge effectiviteit en lage FLOPs die typisch wordt ingenomen door succesvolle BERT-gebaseerde SPLADE-modellen.

Betekenis en Claims
Het artikel betoogt dat de bottleneck bij het aanpassen van vooraf getrainde encoders aan Learned Sparse Retrieval niet louter de encoder-capaciteit of pretraining-kwaliteit is, maar de kalibratie van de schaal van de MLM-head.

De auteurs claimen dat:

• De schaal van de MLM-head een kritische, eerder over het hoofd geziene factor is in LSR.
• Sterkere encoders niet automatisch "drop-in" vervangers zijn voor BERT in SPLADE; hun projectieschalen moeten worden gekalibreerd voor de specificiteit van de dynamiek van sparse retrieval.
• Een eenvoudige rescaling-correctie instabiele trainingsruns van moderne encoders kan transformeren naar competitieve sparse retrievers die vaak de klassieke BERT-SPLADE baseline evenaren of zelfs overtreffen.

Dit werk suggereert dat toekomstige adaptaties van vooraf getrainde encoders naar LSR rekening moeten houden met de geometrie en schaal van de vocabulaire-projectielaag, en niet alleen met de representatieve kracht van de encoder.