Resumo Técnico: Quando a Distância Distrai: Viés de Distância de Representação no BT-Loss para Modelos de Recompensa

1. Problema

Os modelos de recompensa (RMs) são centrais para o Aprendizado por Reforço com Feedback Humano (RLHF), servindo como o substituto para as preferências humanas que guiam o alinhamento de Grandes Modelos de Linguagem (LLMs). O objetivo padrão para o treinamento desses modelos é a perda de Bradley-Terry (BT), que aprende a partir de dados pareados consistindo em respostas "escolhidas" ($y_w$) e "rejeitadas" ($y_l$).

Os autores identificam uma falha crítica na dinâmica de atualização da perda BT padrão: o Viés de Distância de Representação. Embora o gradiente da BT seja destinado a ser impulsionado pelo erro de predição (a magnitude do erro na classificação), a análise revela que a norma do gradiente escala com dois componentes distintos:

1. Erro de Predição: A diferença nas recompensas previstas entre as respostas escolhida e rejeitada.
2. Distância de Representação: A distância entre as duas respostas no espaço de saída da camada final ($\|h_w - h_l\|$).

Este acoplamento cria um sinal de aprendizado espúrio onde a magnitude da atualização é ditada por quão distantes as representações do par são, em vez de ser determinada unicamente pelo quão errada está a predição.

• Pares de Pequena Distância: Pares que são semanticamente ou estruturalmente semelhantes (por exemplo, um trecho de código correto versus um com um único erro lógico) frequentemente possuem pequenas distâncias de representação. Sob a perda BT padrão, esses pares recebem atualizações assustadoramente fracas, mesmo quando o modelo os classifica incorretamente. Isso é problemático porque distinções finas são frequentemente as mais críticas para tarefas como raciocínio e codificação.
• Pares de Grande Distância: Pares que são vastamente diferentes (por exemplo, uma recusa versus uma resposta a uma consulta nociva) naturalmente possuem grandes distâncias de representação. Esses pares recebem atualizações desproporcionalmente fortes, potencialmente ofuscando os sinais de aprendizado de pares mais sutis e finos.

Consequentemente, o modelo subtreina nas distinções finas e difíceis onde ele mais falha, enquanto sobrecarrega diferenças superficiais.

2. Metodologia: NormBT

Para abordar este viés, os autores propõem o NormBT, um esquema de normalização pareada adaptativa que reescala a contribuição do gradiente de cada par de preferência.

Base Teórica

Os autores decompõem a norma do gradiente BT $\|\nabla_\theta L_{BT}\|$ e mostram que ela é proporcional a:
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
onde $d$ é a diferença de recompensa e $\|h_w - h_l\|$ é a distância de representação. Para fazer com que o tamanho da atualização dependa primariamente do termo de erro de predição $|\sigma(d) - 1|$, os autores introduzem um peso $w_i$ que escala inversamente com a distância de representação.

Detalhes de Implementação

Computar diretamente a norma exata do gradiente para ponderação é computacionalmente caro. Em vez disso, o NormBT usa a distância de representação da camada final ($\|h_w - h_l\|$) como um substituto (proxy), que se mostra um substituto exato para a cabeça de pontuação linear e uma forte aproximação para o modelo completo.

A função de perda proposta é:
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

Onde o peso $\tilde{w}$ é definido como:
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

Escolhas de design fundamentais incluem:

• Stop-Gradient: O termo de distância é computado com um operador stop-gradient para evitar que o próprio peso de normalização seja otimizado.
• Estabilização por EMA: Para evitar o desvio de escala durante o treinamento (já que as magnitudes dos embeddings podem mudar), os autores utilizam uma Média Móvel Exponencial (EMA) da média de representação do lote ($\mu_t$) para normalizar os pesos. Isso garante que os pesos permaneçam estáveis e adaptáveis às mudanças distributivas.
• Integração Leve: O NormBT não requer mudanças arquiteturais, preserva a base probabilística da BT e acarreta um overhead computacional insignificante.

3. Principais Contribuições

1. Análise de Gradiente: O artigo fornece uma análise teórica e empírica detalhada mostrando que as normas do gradiente BT são determinadas conjuntamente pelo erro de predição e pela distância de representação, levando a sinais de aprendizado enviesados.
2. Identificação de Viés: Os autores demonstram que este acoplamento faz com que pares de pequena distância (frequentemente os mais críticos para tarefas finas) recebam atualizações fracas, enquanto pares de grande distância dominam o sinal de aprendizado.
3. Proposta do NormBT: Eles introduzem o NormBT, uma modificação direta (drop-in) que normaliza as atualizações pela distância de representação, garantindo que os sinais de aprendizado foquem nos erros de predição.
4. Validação Empírica: Experimentos extensos em diferentes backbones de LLM (Gemma-2b-it, Llama-3.2-3B) e conjuntos de dados (Unified-Feedback, Skywork-Reward) mostram melhorias de desempenho consistentes.

4. Resultados Experimentais

Os autores avaliaram o NormBT contra a linha de base BT padrão e várias variantes (BT + Margin, BT + Label Smoothing) no benchmark RewardBench e em tarefas de seleção Best-of-N (BoN) subsequentes.

• Desempenho no RewardBench: O NormBT superou consistentemente a linha de base BT em todas as configurações. Os ganhos mais significativos foram observados na categoria de Raciocínio (Reasoning), onde a precisão melhorou em mais de 5% em média. Isso corrobora a hipótese de que as tarefas de Raciocínio dependem de distinções finas (pequenas distâncias de representação) que são suprimidas pela BT padrão.
• RLHF Subsequente (BoN): Em experimentos de seleção Best-of-N, os modelos treinados com NormBT alcançaram "scores de ouro" (qualidade das respostas selecionadas) mais altos em comparação com as linhas de base BT.
• Estudos de Ablação:
  • O NormBT superou outras técnicas de treinamento como LayerNorm e Gradient Clipping, que falharam em abordar o problema do peso relativo entre os pares.
  • O método provou ser complementar, melhorando o desempenho mesmo quando aplicado sobre outras variantes de BT (ex: NormBT + Margin).
  • Remover a estabilização por EMA ou usar métricas de similaridade alternativas (similaridade de cosseno, embeddings de pooling médio) resultou em degradação de desempenho, validando as escolhas de design específicas.

5. Significância e Alegações

O artigo afirma que o NormBT oferece uma solução fundamentada e leve para uma limitação estrutural na perda BT padrão. Ao desacoplar a magnitude da atualização da distância de representação, o NormBT garante que o modelo de recompensa aprenda de forma mais fiel as distinções finas, que são frequentemente as mais desafiadoras, porém críticas para um alinhamento de alta qualidade.

Os autores enfatizam que seu trabalho não aborda diretamente a qualidade dos dados ou o viés do anotador, mas melhora a eficiência e a fidelidade dos sinais de aprendizado extraídos dos dados pareados existentes. Eles posicionam o NormBT como uma melhoria fundamental para o paradigma de modelagem de recompensa discriminativa, potencialmente levando a sinais de recompensa mais confiáveis em aplicações de RLHF subsequentes, particularmente em domínios que exigem raciocínio preciso e detecção de erros sutis.