Rescaling MLM-Head for Neural Sparse Retrieval

Sintesi Tecnica: Riscalamento della Testa MLM per la Recupero Sparso Neurale

Definizione del Problema
I modelli di Recupero Sparso Appreso (LSR), in particolare SPLADE, si affidano alla testa di Masked Language Modeling (MLM) di un encoder pre-addestrato per generare rappresentazioni lessicali sparse. Sebbene sia intuitivo che sostituire il classico backbone BERT con encoder moderni più forti (es. ModernBERT, Ettin, RoBERTa) debba migliorare le prestazioni di recupero, gli autori osservano un fallimento controintuitivo: sotto le ricette standard di addestramento SPLADE, i backbone con norme L2 della testa MLM elevate soffrono di un grave degrado delle prestazioni o di un completo collasso dell'addestramento.

Il documento identifica la causa principale in un disallineamento di scala (scale mismatch). In SPLADE, i vettori sparsi sono costruiti direttamente dagli output della testa MLM ($H W_{head}^T$) e confrontati tramite un prodotto scalare non normalizzato. Quando la norma della matrice della testa MLM $W_{head}$ è elevata, essa gonfia l'entità delle attivazioni sparse e dei punteggi di recupero. Questo gonfiamento distorce la dinamica di ottimizzazione della perdita di ranking contrastiva e crea tensione con la regolarizzazione della sparsità (FLOPS), portando a un addestramento instabile e a un recupero inefficace.

Metodologia
Per affrontare questo disallineamento di scala senza alterare l'architettura del modello o l'obiettivo di addestramento, gli autori propongono una semplice correzione al tempo di inizializzazione a costo zero: il Riscalamento della Testa MLM (MLM-Head Rescaling).

Prima dell'inizio dell'addestramento, la matrice di proiezione della testa MLM viene riscalata per un fattore costante $k$:
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

Le caratteristiche chiave di questo metodo includono:

• Regolazione a Costo Zero: Non introduce parametri aggiuntivi, non richiede modifiche architettoniche e non aggiunge sovraccarico computazionale durante l'addestramento o l'inferenza.
• Preservazione dell'Informazione: L'operazione preserva la direzione e l'informazione lessicale della proiezione del vocabolario pre-addestrato.
• Applicabilità: Per i modelli con teste MLM condivise (tied), l'operazione viene applicata alla matrice dei pesi condivisa per mantenere il legame dei pesi (weight tying).

Configurazione Sperimentale
Gli autori hanno valutato sette backbone encoder (BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT ed Ettin) addestrati su triplette MS MARCO utilizzando una ricetta SPLADE fissa. Hanno testato vari fattori di riscalamento ($k$) e valutato le prestazioni su:

• In-domain: MS MARCO Dev e TREC 2019.
• Out-of-domain: BEIR-13 (media nDCG@10).
• Trade-off Sparsità-Efficienza: Analizzato tramite le prestazioni NanoBEIR rispetto ai FLOPs attivati.

Risultati Chiave

1. Recupero di Backbone con Norma Elevata: L'inizializzazione standard ($k=1$) ha causato un sotto-performance significativo per gli encoder moderni con norme elevate. Ad esempio, ModernBERT è sceso a un nDCG@10 di 0.127 su BEIR-13 (rispetto allo 0.370 di BERT). Applicando il riscalamento ($k=16$), è stato possibile recuperare le prestazioni di ModernBERT portandolo a 0.405, superando la baseline di BERT. Allo stesso modo, Ettin è migliorato da 0.221 a 0.391.
2. Sensibilità Non Monotona: Il fattore di riscalamento ottimale dipende dalla norma iniziale del backbone. Mentre ModernBERT ed Ettin richiedevano un forte riscalamento ($k=16$), RoBERTA ha raggiunto il picco con $k=2$. Un riscalamento eccessivo (es. $k=16$ per RoBERTa) o il riscalamento di modelli con norma già piccola (come ALBERT) ha portato al collasso delle prestazioni.
3. Stabilità dell'Addestramento: Il riscalamento ha stabilizzato significativamente le dinamiche di addestramento. I modelli ModernBERT ed Ettin non riscalati partivano con valori di perdita estremamente alti e convergevano verso minimi locali poveri. Le varianti riscalate hanno mostrato curve di perdita più fluide e perdite finali inferiori.
4. Miglioramento dei Trade-off: I modelli riscalati hanno ottenuto un migliore trade-off tra efficacia e sparsità, operando nella regione ad alta efficacia e bassi FLOPs tipicamente occupata dai modelli successo-base BERT-SPLADE.

Significatività e Rivendicazioni
Il documento sostiene che il collo di bottiglia nell'adattare gli encoder pre-addestrati al Recupero Sparso Appreso non è semplicemente la capacità dell'encoder o la qualità del pre-addestramento, ma il calibrazione della scala della testa MLM.

Gli autori affermano che:

• La scala della testa MLM è un fattore critico, precedentemente trascurato, nell'LSR.
• Gli encoder più forti non sono automaticamente sostituti "plug-and-play" di BERT in SPLADE; le loro scale di proiezione devono essere calibrate sulla specifica dinamica del recupero sparso.
• Una semplice correzione di riscalamento può trasformare sessioni di addestramento instabili di encoder moderni in recuperatori sparsi competitivi, spesso eguagliando o superando la classica baseline BERT-SPLADE.

Questo lavoro suggerisce che i futuri adattamenti di encoder pre-addestrati all'LSR devono tenere conto della geometria e della scala dello strato di proiezione del vocabolario, non solo della potenza rappresentativa dell'encoder.