Rescaling MLM-Head for Neural Sparse Retrieval

Résumé Technique : Redimensionnement de la tête MLM pour la recherche parcimonieuse neuronale

Énoncé du problème
Les modèles de recherche parcimonieuse apprise (LSR - Learned Sparse Retrieval), particulièrement SPLADE, s'appuient sur la tête de modélisation de masquage de langue (MLM) d'un encodeur pré-entraîné pour générer des représentations lexicales parcimonieuses. Bien qu'il soit intuitif que le remplacement du backbone BERT standard par des encodeurs modernes plus puissants (ex. : ModernBERT, Ettin, RoBERTa) devrait améliorer les performances de recherche, les auteurs observent un échec contre-intuitif : sous les recettes d'entraînement standard de SPLADE, les backbones présentant de grandes normes L2 de la tête MLM souffrent d'une dégradation sévère des performances ou d'un effondrement complet de l'entraînement.

Le papier identifie la cause profonde comme étant un déséquilibre d'échelle (scale mismatch). Dans SPLADE, les vecteurs parcimonieux sont construits directement à partir des sorties de la tête MLM ($H W_{head}^T$) et comparés via un produit scalaire non normalisé. Lorsque la matrice de la tête MLM $W_{head}$ possède une norme élevée, cela gonfle la magnitude des activations parcimonieuses et des scores de recherche. Cette inflation fausse la dynamique d'optimisation de la perte de classement contrastive et crée une tension avec la régularisation de la parcimonie (FLOPS), entraînant une instabilité de l'entraînement et une recherche inefficace.

Méthodologie
Pour remédier à ce déséquilibre d'échelle sans modifier l'architecture du modèle ni l'objectif d'entraînement, les auteurs proposent une correction simple au moment de l'initialisation, sans coût supplémentaire : le Redimensionnement de la tête MLM (MLM-Head Rescaling).

Avant le début de l'entraînement, la matrice de projection de la tête MLM est redimensionnée par un facteur constant $k$ :
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

Les caractéristiques clés de cette méthode incluent :

• Ajustement sans coût : Elle n'introduit aucun paramètre supplémentaire, ne nécessite aucun changement architectural et n'ajoute aucune surcharge de calcul pendant l'entraînement ou l'inférence.
• Préservation de l'information : L'opération préserve la direction et l'information lexicale de la projection du vocabulaire pré-entraîné.
• Applicabilité : Pour les modèles avec des têtes MLM liées (tied MLM heads), l'opération est appliquée à la matrice de poids partagée pour maintenir le lien des poids (weight tying).

Configuration expérimentale
Les auteurs ont évalué sept backbones d'encodeurs (BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT, et Ettin) entraînés sur des triplets MS MARCO en utilisant une recette SPLADE fixe. Ils ont testé divers facteurs de redimensionnement ($k$) et ont évalué les performances sur :

• Domaine interne (In-domain) : MS MARCO Dev et TREC 2019.
• Hors domaine (Out-of-domain) : BEIR-13 (moyenne nDCG@10).
• Compromis Parcimonie-Efficacité : Analysé via la performance NanoBEIR par rapport aux FLOPs activés.

Résultats clés

1. Récupération des backbones à grande norme : L'initialisation standard ($k=1$) a provoqué une sous-performance significative des encodeurs modernes possédant de grandes normes. Par exemple, ModernBERT est tombé à un nDCG@10 de 0,127 sur BEIR-13 (comparé à 0,370 pour BERT). L'application du redimensionnement ($k=16$) a permis de récupérer la performance de ModernBERT à 0,405, surpassant la ligne de base BERT. De même, Ettin est passé de 0,221 à 0,391.
2. Sensibilité non monotone : Le facteur de redimensionnement optimal dépend de la norme initiale du backbone. Alors que ModernBERT et Ettin nécessitaient un fort redimensionnement ($k=16$), RoBERTA atteignait un pic à $k=2$. Un redimensionnement excessif (ex. : $k=16$ pour RoBERTA) ou un redimensionnement de modèles ayant déjà une petite norme (comme ALBERT) entraînait un effondrement des performances.
3. Stabilité de l'entraînement : Le redimensionnement a considérablement stabilisé la dynamique d'entraînement. Les modèles ModernBERT et Ettin non redimensionnés commençaient avec des valeurs de perte extrêmement élevées et convergeaient vers de mauvais minima locaux. Les variantes redimensionnées ont présenté des courbes de perte plus lisses et des pertes finales plus faibles.
4. Amélioration des compromis : Les modèles redimensionnés ont atteint un meilleur compromis efficacité-parcimonie, opérant dans la région haute efficacité/bas FLOPs typiquement occupée par les modèles BERT-SPLADE réussis.

Signification et Revendications
Le papier soutient que le goulot d'étranglement pour l'adaptation des encodeurs pré-entraînés à la recherche parcimonieuse apprise n'est pas seulement la capacité de l'encodeur ou la qualité du pré-entraînement, mais le calibrage de l'échelle de la tête MLM.

Les auteurs affirment que :

• L'échelle de la tête MLM est un facteur critique et précédemment négligé dans la LSR.
• Les encodeurs plus puissants ne sont pas automatiquement des remplacements directs ("drop-in") de BERT pour SPLADE ; leurs échelles de projection doivent être calibrées pour la dynamique spécifique de la recherche parcimonieuse.
• Une simple correction de redimensionnement peut transformer les sessions d'entraînement instables des encodeurs modernes en des récupérateurs parcimonieux compétitifs, égalant ou dépassant souvent la ligne de base classique BERT-SPLADE.

Ce travail suggère que les futures adaptations d'encodeurs pré-entraînés à la LSR doivent tenir compte de la géométrie et de l'échelle de la couche de projection du vocabulaire, et non seulement de la puissance de représentation de l'encodeur.