Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models

1. O Problema

Os Modelos de Linguagem (LMs) modernos, especialmente após o treinamento com Aprendizado por Reforço (RL), tendem a colapsar sua distribuição interna de respostas em um único modo dominante (a resposta mais provável). Embora isso seja eficaz para tarefas de benchmark com uma única resposta correta, ele falha em cenários do mundo real onde:

• Existem múltiplas respostas válidas: Como em diagnósticos médicos (diferenciais) ou soluções de código com diferentes abordagens algorítmicas.
• Há incerteza irreduzível: Devido a informações incompletas ou ambíguas (ex: um paciente com sintomas sobrepostos).
• A diversidade é crucial: Métodos atuais de inferência (como sampling repetido ou best-of-k) são computacionalmente caros e frequentemente redundantes, pois o modelo foi treinado para "escolher" uma resposta, não para explorar um espaço de hipóteses.

O objetivo do trabalho é treinar LMs para gerar explicitamente distribuições de respostas plausíveis (conjuntos de hipóteses) com estimativas de confiança calibradas, em uma única passagem de forward, em vez de depender de amostragem externa.

2. Metodologia: Multi-Answer Reinforcement Learning (Multi-Answer RL)

Os autores propõem uma nova abordagem de RL que modifica o objetivo de treinamento para otimizar a geração de conjuntos de respostas e distribuições de confiança.

A. Multi-Answer RLVR (Reinforcement Learning with Verifiable Rewards)

• Objetivo: Treinar o modelo para gerar um conjunto de $K$ respostas candidatas distintas ($A = \{a_1, ..., a_K\}$) em uma única cadeia de pensamento.
• Função de Recompensa: Utiliza uma recompensa no nível do conjunto que conta quantas das respostas geradas pertencem ao conjunto de respostas verdadeiras ($Y^*$).
  $$R_{RLVR}^{multi}(A, Y^*) = \sum_{i=1}^{K} \mathbb{1}[a_i \in Y^*]$$
• Generalização: Este objetivo engloba casos clássicos como Pass@K (quando há uma única resposta correta, mas o modelo gera múltiplas) e recuperação de conjuntos parciais ou completos (quando há múltiplas respostas corretas).
• Restrição de Formato: Adiciona uma recompensa de formato para garantir que as $K$ respostas geradas sejam únicas (evitando duplicatas).

B. Multi-Answer RLCR (Reinforcement Learning with Calibration Rewards)

• Objetivo: Estender o RLVR para incluir estimativas de confiança calibradas para cada resposta no conjunto.
• Saída: O modelo gera o conjunto de respostas $A$ e um conjunto correspondente de valores de confiança $Q = \{q_1, ..., q_K\}$, onde $q_i \in [0, 1]$.
• Função de Recompensa: Combina a recompensa de cobertura do conjunto com uma penalidade baseada em regras de pontuação adequadas (proper scoring rules), especificamente o Multi-Brier Score.
  $$R_{RLCR}^{multi} = R_{RLVR}^{multi} - \text{Multi-Brier}(A, Q, Y^*)$$
• Interpretação: Isso incentiva o modelo a não apenas encontrar as respostas corretas, mas também a atribuir probabilidades que reflitam a veracidade real de cada hipótese, permitindo uma representação explícita da distribuição $P(y|x)$.

3. Contribuições Principais

1. Novo Paradigma de Treinamento: Introduz o Multi-Answer RL, que internaliza a busca por múltiplas hipóteses no processo generativo do modelo, eliminando a necessidade de sampling repetido na inferência para obter diversidade.
2. Eficiência Computacional: Demonstra que gerar um conjunto de respostas em uma única passagem é significativamente mais eficiente em termos de tokens do que amostrar múltiplas vezes independentemente (redução de redundância no raciocínio).
3. Calibração de Incerteza: Mostra como integrar regras de pontuação adequadas (como Brier Score) no treinamento por RL para obter estimativas de confiança calibradas em nível de conjunto e individual.
4. Validação Empírica: Aplica a metodologia em três domínios distintos:
   • Diagnóstico Médico (DDXPlus): Onde múltiplos diagnósticos são simultaneamente corretos.
   • Perguntas Ambíguas (HotPotQA-Modified): Onde a informação incompleta exige múltiplas hipóteses.
   • Codificação (MBPP): Onde existem múltiplas implementações corretas para a mesma tarefa.

4. Resultados Chave

Os experimentos compararam modelos treinados com Single-Answer RL (baselines) contra os novos modelos Multi-Answer RL:

• Cobertura e Diversidade:

  • Os modelos Multi-Answer recuperaram significativamente mais respostas corretas (maior Coverage) em tarefas médicas e de codificação.
  • Em tarefas de codificação, a precisão Top-1 aumentou em mais de 50% em comparação com baselines de amostragem simples, enquanto a diversidade de soluções únicas foi drasticamente maior.
  • O modelo Single-Answer tendia a colapsar para o mesmo modo dominante repetidamente, mesmo quando amostrado várias vezes.

• Eficiência de Tokens:

  • Para gerar o mesmo número de respostas (ex: 3), o modelo Multi-Answer utilizou menos da metade dos tokens em comparação com o modelo Single-Answer amostrado 3 vezes. Isso ocorre porque o modelo Multi-Answer evita a repetição de cadeias de raciocínio redundantes.

• Calibração:

  • O uso de RLCR Multi resultou em melhores métricas de calibração (Brier Score e ECE) em nível de conjunto.
  • O modelo RLVR Multi (sem calibração explícita) tendia a ser sistematicamente superconfiante, enquanto o RLCR Multi aprendeu a atribuir confidências mais realistas, especialmente em cenários de alta incerteza.

• Estabilidade:

  • O treinamento permaneceu estável à medida que o número de respostas alvo ($K$) aumentava, permitindo que o modelo explorasse mais hipóteses válidas sem degradação de desempenho.

5. Significado e Conclusão

O trabalho posiciona o Multi-Answer RL como uma alternativa principial e eficiente ao escalonamento de inferência (como best-of-k ou sampling massivo).

• Mudança de Paradigma: Em vez de tratar a incerteza e a ambiguidade como falhas a serem contornadas na inferência, o modelo aprende a representá-las estruturalmente durante o treinamento.
• Aplicabilidade em Alta Risco: Para áreas críticas como medicina e segurança, a capacidade de um modelo de listar múltiplos diagnósticos plausíveis com suas respectivas probabilidades é superior a fornecer uma única resposta "confiante" que pode estar errada.
• Eficiência: A redução no custo computacional (tokens) torna viável a implementação de sistemas de raciocínio distribucional em produção, onde a geração de múltiplas hipóteses é necessária para tomada de decisão robusta.

Em suma, o artigo demonstra que otimizar explicitamente para a diversidade e a calibração de conjuntos de respostas durante o treinamento por RL supera as limitações dos métodos atuais, produzindo modelos mais robustos, diversos e eficientes.