Resumo Técnico: Recompensa Consciente da Distribuição para Regressão em LLM

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) demonstraram a capacidade de prever quantidades de valor real a partir de entradas heterogêneas (texto, código, strings moleculares). No entanto, os objetivos de treinamento padrão para regressão em LLM — como Ajuste Fino Supervisionado (SFT) ou Aprendizado por Reforço (RL) com recompensas pontuais (por exemplo, erro quadrático médio negativo) — tratam cada número de ponto flutuante decodificado como uma previsão independente.

Essa abordagem otimiza estimativas pontuais individuais sem garantir que o conjunto de previsões amostradas forme uma distribuição preditiva calibrada. Consequentemente, os modelos frequentemente produzem previsões que estão individualmente próximas do alvo, mas coletivamente mal calibradas, excessivamente estreitas ou unilaterais. Esse comportamento de "busca pela média" colapsa a distribuição preditiva, falhando em capturar a incerteza necessária ou a ordenação relativa. Essa limitação impede aplicações que exigem classificação de candidatos, estimativa de incerteza ou tomada de decisão robusta sob incerteza, onde a distribuição de valores possíveis importa tanto quanto a única melhor estimativa.

Metodologia: Recompensa Consciente da Distribuição (DAR)

Os autores propõem a Recompensa Consciente da Distribuição (DAR), um objetivo de aprendizado por reforço on-policy projetado para otimizar a qualidade de toda a distribuição preditiva formada por múltiplas amostras decodificadas, em vez de otimizar execuções individuais isoladamente.

Mecanismo Central

1. Distribuição Preditiva Empírica: Para uma dada entrada $x$, o modelo amostra $K$ execuções, produzindo um conjunto de previsões $\{p_{ik}\}_{k=1}^K$. Esse conjunto é tratado como uma distribuição preditiva empírica $F_i$.
2. Pontuação via CRPS: A qualidade dessa distribuição é avaliada contra o alvo de verdade fundamental $y_i$ usando a Pontuação de Probabilidade Ranqueada Contínua (CRPS). O CRPS é uma regra de pontuação própria que recompensa distribuições que são precisas (concentradas perto do alvo) e apropriadamente dispersas (evitando degenerescência).
   • A recompensa para o conjunto completo é definida como $S_i = -\text{CRPS}(F_i, y_i)$.
3. Atribuição de Crédito Leave-One-Out: Como o RL requer uma recompensa escalar para cada execução individual para atualizar a política, o DAR converte a pontuação no nível da distribuição em feedback no nível da execução. Ele calcula a contribuição marginal de cada execução $p_{ik}$ comparando o CRPS do conjunto completo ($S_i$) com o CRPS do conjunto com aquela execução removida ($S_i^{(-k)}$).
   • A recompensa da execução é: $R_{ik}^{\text{DAR}} = S_i - S_i^{(-k)}$.
   • Se remover uma execução degradar a distribuição (reduzir o CRPS negado), essa execução recebe uma recompensa positiva, mesmo que seu erro pontual individual seja alto. Por outro lado, execuções redundantes que não melhoram a dispersão ou o centramento da distribuição recebem recompensas menores ou negativas.

Distinção Teórica

O artigo contrasta isso com a perda quadrática pontual (MSE), que é teoricamente "buscadora de média" e incentiva a distribuição preditiva a colapsar para um único ponto (a média condicional). Em contraste, o CRPS é "buscador de distribuição", otimizando para a distribuição condicional completa. O DAR aproveita isso usando atribuição de crédito leave-one-out para incentivar um conjunto de execuções que sejam bem centradas e utilmente dispersas.

Contribuições Principais

1. Novo Objetivo de RL: Introdução do DAR, que desloca o objetivo de otimização da minimização de erros pontuais individuais para a maximização da qualidade da distribuição preditiva empírica.
2. Atribuição de Crédito Leave-One-Out: Um mecanismo específico para atribuir recompensas a execuções individuais com base em sua contribuição marginal para a qualidade da distribuição coletiva, permitindo que o modelo recompense "dispersão útil" mesmo quando amostras individuais são menos precisas.
3. Mitigação do Colapso de Diversidade: O método aborda a tendência do ajuste fino por RL de reduzir a diversidade de saída (entropia da política) e colapsar previsões em direção à média, garantindo que as execuções amostradas permaneçam diversas o suficiente para cobrir a verdade fundamental.

Resultados Experimentais

Os autores avaliaram o DAR em três tarefas distintas: uma tarefa de regressão de mistura gaussiana sintética, previsão de desempenho de código e previsão de propriedades moleculares.

1. Regressão Distribucional Sintética

• Tarefa: Prever alvos a partir de uma mistura gaussiana de dois componentes heterocedástica.
• Descobertas: O DAR alcançou a maior correlação de Spearman entre os métodos baseados em LLM, superando significativamente o SFT e recompensas MSE pontuais. Ele preservou melhor a ordenação global da função alvo e adaptou-se bem a regiões de extrapolação onde outras bases neurais lutaram.

2. Previsão de Desempenho de Código

• Tarefa: Prever latência de kernel (KBSS) e uso de pico de memória (APPS) a partir de código-fonte.
• Descobertas: O DAR alcançou a correlação de rank mais forte (Spearman $\rho$) e erro absoluto médio (MAE) competitivo ou inferior em comparação com o SFT e bases de RL pontuais.
  • No KBSS, o DAR melhorou a correlação de Spearman em 6 pontos sobre a base mais forte.
  • No APPS, melhorou a correlação de Spearman em 4 pontos.
  • Isso indica sensibilidade superior à variação dependente de entrada e classificação de candidatos mais confiável.

3. Previsão de Propriedades Moleculares (MoleculeNet)

• Tarefa: Prever propriedades físico-químicas (FreeSolv, ESOL, Lipofilicidade) a partir de strings SMILES.
• Descobertas: Apesar de usar apenas strings SMILES (faltando vieses indutivos explícitos de grafos ou 3D usados por fortes bases GNN), o DAR permaneceu competitivo.
  • Alcançou melhorias estatisticamente significativas no RMSE no FreeSolv e Lipofilicidade em comparação com outras estratégias de treinamento de LLM (SFT, MSE, Value Head).
  • Demonstrou que a otimização no nível da distribuição pode produzir comportamento de regressão robusto mesmo sem representações moleculares especializadas.

Diagnósticos de Incerteza

• Calibração: A análise mostrou uma forte correlação log-Pearson (0,91) entre o desvio padrão da execução e o erro de previsão para o DAR, indicando que as estimativas de incerteza do modelo estavam bem calibradas aos erros reais.
• Diversidade: Durante o treinamento, o DAR manteve maior entropia de execução e uma maior "taxa de parêntese" (a fração de exemplos onde a verdade fundamental está entre as previsões mínima e máxima da execução) em comparação com o RL baseado em MSE, confirmando a mitigação do colapso de diversidade.

Significado e Alegações

O artigo afirma que otimizar diretamente distribuições preditivas torna a regressão baseada em LLM mais robusta e melhor calibrada. Ao tratar múltiplas amostras decodificadas como uma distribuição e recompensar sua qualidade coletiva, o DAR melhora o desempenho baseado em rank (crucial para seleção de candidatos) e a estimativa de incerteza sem sacrificar a precisão pontual.

Os autores posicionam este trabalho como um passo rumo ao uso de LLMs não apenas como geradores de texto ou estimadores pontuais, mas como regressores probabilísticos flexíveis para domínios científicos. Eles enfatizam que, embora a regressão em LLM ofereça flexibilidade no manuseio de entradas heterogêneas (texto, código, strings), a abordagem proposta consciente da distribuição é necessária para desbloquear o potencial total desses modelos para tarefas que exigem quantificação de incerteza e ordenação confiável, áreas onde o treinamento pontual tradicional falha. O trabalho não afirma substituir modelos de regressão especializados (como GPs ou GNNs) em todos os domínios, mas demonstra que o RL consciente da distribuição pode tornar LLMs de propósito geral competitivos em configurações de regressão científica.