Rescaling MLM-Head for Neural Sparse Retrieval

Resumo Técnico: Redimensionamento da Cabeça de MLM para Recuperação Esparsa Neural

Definição do Problema
Modelos de Recuperação Esparsa Aprendida (LSR), particularmente o SPLADE, dependem da cabeça de Modelagem de Linguagem Mascarada (MLM) de um codificador pré-treinado para gerar representações lexicais esparsas. Embora seja intuitivo que substituir o backbone BERT padrão por codificadores modernos mais fortes (ex: ModernBERT, Ettin, RoBERTa) deva melhorar o desempenho de recuperação, os autores observam uma falha contraintuitiva: sob as receitas padrão de treinamento do SPLADE, backbones com normas L2 da cabeça de MLM elevadas sofrem uma degradação severa de desempenho ou colapso total do treinamento.

O artigo identifica a causa raiz como um desajuste de escala (scale mismatch). No SPLADE, os vetores esparsos são construídos diretamente das saídas da cabeça de MLM ($H W_{head}^T$) e comparados via produto escalar não normalizado. Quando a matriz da cabeça de MLM $W_{head}$ possui uma norma grande, ela infla a magnitude das ativações esparsas e dos escores de recuperação. Essa inflação distorce a dinâmica de otimização da perda de classificação contrastiva e cria tensão com a regularização de esparsidade (FLOPS), levando a um treinamento instável e uma recuperação ineficaz.

Metodologia
Para abordar esse desajuste de escala sem alterar a arquitetura do modelo ou o objetivo de treinamento, os autores propõem uma correção simples de tempo de inicialização e custo zero: Redimensionamento da Cabeça de MLM (MLM-Head Rescaling).

Antes do início do treinamento, a matriz de projeção da cabeça de MLM é redimensionada por um fator constante $k$:
$$W_{head} \leftarrow \frac{W_{head}}{k}$$

As principais características deste método incluem:

• Ajuste de Custo Zero: Não introduz parâmetros adicionais, não requer mudanças arquiteturais e não adiciona sobrecarga computacional durante o treinamento ou inferência.
• Preservação de Informação: A operação preserva a direção e a informação lexical da projeção do vocabulário pré-treinado.
• Aplicabilidade: Para modelos com cabeças de MLM compartilhadas (tied), a operação é aplicada à matriz de peso compartilhada para manter o vínculo de pesos (weight tying).

Configuração Experimental
Os autores avaliaram sete backbones de codificadores (BERT, DistilBERT, ALBERT, GTE-MLM, RoBERTa, ModernBERT e Ettin) treinados em tripletos MS MARCO usando uma receita fixa de SPLADE. Eles testaram vários fatores de redimensionamento ($k$) e avaliaram o desempenho em:

• Domínio interno: MS MARCO Dev e TREC 2019.
• Fora de domínio: BEIR-13 (média de nDCG@10).
• Compromisso entre Esparsidade-Eficiência: Analisado via desempenho NanoBEIR contra FLOPs ativados.

Resultos Principais

1. Recuperação de Backbones de Grande Norma: A inicialização padrão ($k=1$) fez com que codificadores modernos com normas elevadas tivessem um desempenho significativamente inferior. Por exemplo, o ModernBERT caiu para um nDCG@10 de 0,127 no BEIR-13 (comparado aos 0,370 do BERT). Aplicar o redimensionamento ($k=16$) recuperou o desempenho do ModernBERT para 0,405, superando a linha de base do BERT. Da mesma forma, o Ettin melhorou de 0,221 para 0,391.
2. Sensibilidade Não Monotônica: O fator de redimensionamento ideal depende da norma inicial do backbone. Enquanto o ModernBERT e o Ettin exigiram um redimensionamento forte ($k=16$), o RoBERTA atingiu o pico em $k=2$. O redimensionamento excessivo (ex: $k=16$ para o RoBERTa) ou o redimensionamento de modelos com norma já pequena (como o ALBERT) levou ao colapso do desempenho.
3. Estabilidade de Treinamento: O redimensionamento estabilizou significativamente a dinâmica de treinamento. Os modelos não redimensionados de ModernBERT e Ettin começaram com valores de perda extremamente altos e convergiram para mínimos locais ruins. As variantes redimensionadas exibiram curvas de perda mais suaves e perdas finais menores.
4. Melhoria nos Compromissos: Os modelos redimensionados alcançaram um melhor compromisso entre eficácia e esparsidade, operando na região de alta eficácia e baixos FLOPs tipicamente ocupada por modelos bem-sucedidos baseados em BERT-SPLADE.

Significância e Alegações
O artigo argumenta que o gargalo na adaptação de codificadores pré-treinados para a Recuperação Esparsa Aprendida não é meramente a capacidade do codificador ou a qualidade do pré-treinamento, mas sim a calibração da escala da cabeça de MLM.

Os autores afirmam que:

• A escala da cabeça de MLM é um fator crítico e anteriormente negligenciado em LSR.
• Codificadores mais fortes não são automaticamente substitutos diretos (drop-in replacements) para o BERT no SPLADE; suas escalas de projeção devem ser calibradas para a dinâmica específica da recuperação esparsa.
• Uma simples correção de redimensionamento pode transformar execuções de treinamento instáveis de codificadores modernos em recuperadores esparsos competitivos, muitas vezes igualando ou superando a linha de base clássica do BERT-SPLADE.

Este trabalho sugere que as futuras adaptações de codificadores pré-treinados para LSR devem levar em conta a geometria e a escala da camada de projeção do vocabulário, e não apenas o poder de representação do codificador.