Rescaling MLM-Head for Neural Sparse Retrieval

Resumen Técnico: Reescalado del MLM-Head para la Recuperación Dispersa Neuronal

Planteamiento del Problema
Los modelos de Recuperación Dispersa Aprendida (LSR, por sus siglas en inglés), particularmente SPLADE, dependen de la cabeza de Modelado de Lenguaje con Máscara (MLM) de un codificador preentrenado para generar representaciones léxicas dispersas. Aunque es intuitivo que reemplazar el backbone estándar de BERT por codificadores modernos más potentes (p. ej., ModernBERT, Ettin, RoBERTa) debería mejorar el rendimiento de la recuperación, los autores observan un fallo contraintuitivo: bajo las recetas de entrenamiento estándar de SPLADE, los backbones con normas L2 de la cabeza MLM elevadas sufren una degradación severa de rendimiento o un colapso total del entrenamiento.

El artículo identifica la causa raíz como un desajuste de escala (scale mismatch). En SPLADE, los vectores dispersos se construyen directamente a partir de las salidas de la cabeza MLM ($H W_{head}^T$) y se comparan mediante un producto punto no normalizado. Cuando la norma de la matriz de la cabeza MLM $W_{head}$ es grande, infla la magnitud de las activaciones dispersas y de las puntuaciones de recuperación. Esta inflación distorsiona la dinámica de optimización de la pérdida de clasificación contrastiva y crea una tensión con la regularización de dispersión (FLOPS), lo que conduce a un entrenamiento inestable y a una recuperación ineficaz.

Metodología
Para abordar este desajuste de escala sin alterar la arquitectura del modelo ni el objetivo de entrenamiento, los autores proponen una corrección simple de costo cero en el momento de la inicialización: Reescalado de la Cabeza MLM (MLM-Head Rescaling).

Antes de que comience el entrenamiento, la matriz de proyección de la cabeza MLM se reescala por un factor constante $k$:
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

Las características clave de este método incluyen:

• Ajuste de Costo Cero: No introduce parámetros adicionales, no requiere cambios arquitectónicos y no añade sobrecarga computacional durante el entrenamiento o la inferencia.
• Preservación de la Información: La operación preserva la dirección y la información léxica de la proyección del vocabulario preentrenado.
• Aplicabilidad: Para modelos con cabezas MLM compartidas (tied MLM heads), la operación se aplica a la matriz de pesos compartida para mantener el vínculo de pesos (weight tying).

Configuración Experimental
Los autores evaluaron siete backbones de codificadores (BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT y Ettin) entrenados en tripletas de MS MARCO utilizando una receta fija de SPLADE. Probaron varios factores de reescalado ($k$) y evaluaron el rendimiento en:

• Dominio interno: MS MARCO Dev y TREC 2019.
• Fuera de dominio: BEIR-13 (media de nDCG@10).
• Compromiso entre Dispersión y Eficiencia: Analizado mediante el rendimiento de NanoBEIR frente a los FLOPs activados.

Resultos Clave

1. Recuperación de Backbones de Norma Grande: La inicialización estándar ($k=1$) causó que los codificadores modernos con normas grandes rindieran significativamente por debajo de lo esperado. Por ejemplo, ModernBERT cayó a un nDCG@10 de 0.127 en BEIR-13 (comparado con el 0.370 de BERT). Aplicar el reescalado ($k=16$) recuperó el rendimiento de ModernBERT a 0.405, superando la línea base de BERT. Del mismo modo, Ettin mejoró de 0.221 a 0.391.
2. Sensibilidad No Monotónica: El factor de reescalado óptimo depende de la norma inicial del backbone. Mientras que ModernBERT y Ettin requirieron un reescalado fuerte ($k=16$), RoBERTA alcanzó su punto máximo en $k=2$. Un reescalado excesivo (p. ej., $k=16$ para RoBERTa) o el reescalado de modelos con normas ya pequeñas (como ALBERT) provocó un colapso en el rendimiento.
3. Estabilidad del Entrenamiento: El reescalado estabilizó significativamente la dinámica de entrenamiento. Los modelos de ModernBERT y Ettin sin reescalar comenzaron con valores de pérdida extremadamente altos y convergieron a mínimos locales pobres. Las variantes reescaladas exhibieron curvas de pérdida más suaves y pérdidas finales más bajas.
4. Mejora de los Compromisos (Trade-offs): Los modelos reescalados lograron un mejor compromiso entre efectividad y dispersión, operando en la región de alta efectividad y bajos FLOPs que típicamente ocupan los exitosos modelos SPLADE basados en BERT.

Significado y Reivindicaciones
El artículo argumenta que el cuello de botella para adaptar codificadores preentrenados a la Recuperación Dispersa Aprendida no es simplemente la capacidad del codificador o la calidad del preentrenamiento, sino la calibración de la escala de la cabeza MLM.

Los autores afirman que:

• La escala de la cabeza MLM es un factor crítico y previamente pasado por alto en LSR.
• Los codificadores más fuertes no son automáticamente reemplazos directos (drop-in replacements) de BERT en SPLADE; sus escalas de proyección deben calibrarse para la dinámica específica de la recuperación dispersa.
• Una simple corrección de reescalado puede transformar ejecuciones de entrenamiento inestables de codificadores modernos en recuperadores dispersos competitivos, igualando o incluso superando la línea base clásica de BERT-SPLADE.

Este trabajo sugiere que las futuras adaptaciones de codificadores preentrenados a LSR deben tener en cuenta la geometría y la escala de la capa de proyección del vocabulario, no solo el poder de representación del codificador.