Expected Reward Prediction, with Applications to Model Routing

Resumo Técnico: Previsão de Recompensa Esperada com Aplicações em Roteamento de Modelos

1. O Problema

Os Modelos de Recompensa (Reward Models - RMs) são ferramentas padrão para avaliar e classificar respostas de Grandes Modelos de Linguagem (LLMs). Tradicionalmente, esses modelos operam no nível da resposta: recebem um prompt $x$ e uma resposta $y$, retornando uma pontuação $r(x, y)$.

O desafio central abordado neste trabalho é a elevação (lifting) dessa função de recompensa do nível da resposta para o nível do modelo. Especificamente, os autores investigam se é possível prever a priori (antes de gerar qualquer resposta) qual será a Recompensa Esperada que um determinado LLM $\pi$ obterá para um prompt específico $x$.

Matematicamente, o objetivo é prever:
$$ER_\pi(x) := \mathbb{E}_{Y \sim \pi(Y | x)}[r(x, Y)]$$

Existem duas incertezas principais que tornam essa previsão teoricamente difícil:

1. Comportamento do Modelo: A sensibilidade do modelo a detalhes do prompt e a distribuição de suas respostas não determinísticas podem ser imprevisíveis.
2. Arbitrariedade do Modelo de Recompensa: A maioria dos RMs é treinada com o objetivo de Bradley-Terry para classificar pares de respostas ($r(x, y_1) - r(x, y_2)$). Isso permite adicionar um valor arbitrário dependente do prompt a cada recompensa, o que, teoricamente, tornaria a recompensa absoluta (e não apenas a diferença) não identificável.

2. Metodologia

2.1. Configuração Experimental

• Dataset: Utilizaram o open-perfectblend, amostrando 4.000 prompts (1.000 de cada uma das 4 categorias: Chat Geral, Seguimento de Instruções, Raciocínio Matemático e Raciocínio em Código).
• Modelos Geradores (LLMs): Uma variedade de modelos instruídos de peso aberto: Llama 3.1 (8B, 70B), Gemma 2 (9B, 27B) e Gemma 1 (7B).
• Modelos de Recompensa (RMs): Três modelos distintos (OpenAssistant-RM, GRM-2B-RM, InternLM-RM) treinados com o objetivo de Bradley-Terry.
• Cálculo do "Ground Truth": Para cada prompt, geraram $K=32$ respostas de cada modelo, calcularam a recompensa de cada uma e computaram a média empírica como o valor alvo de recompensa esperada.

2.2. Abordagem de Previsão (ERP - Expected Reward Prediction)
Os autores propõem uma abordagem surpreendentemente simples:

1. Embedding: Usam um modelo de embedding pré-treinado (gte-large-en-v1.5) para converter o prompt $x$ em um vetor fixo $v(x)$.
2. Modelo Linear: Treinam um probe linear (regressão linear com regularização Ridge) para mapear diretamente o embedding do prompt para a recompensa esperada empírica.
   • Função de perda: Minimização do erro quadrático entre a previsão linear e a média empírica das recompensas.
3. Roteamento: Utilizam essas previsões para rotear prompts para o modelo mais adequado, equilibrando recompensa esperada e custo computacional.

3. Contribuições Principais

1. Previsibilidade de Alta Fidelidade: Demonstram empiricamente que a recompensa esperada de um LLM para um prompt é altamente previsível, mesmo usando apenas um modelo linear simples sobre embeddings genéricos. Isso contradiz a intuição de que o comportamento do modelo seria muito complexo ou estocástico para ser previsto sem gerar a resposta.
2. Elevação de Nível de Recompensa: Validam a ideia de transformar funções de recompensa de nível de resposta em funções de recompensa de nível de modelo, permitindo a comparação direta de modelos para um prompt específico sem necessidade de geração de amostras.
3. Protocolo de Roteamento Escalável (ERP): Propõem um protocolo de roteamento que supera métodos baseados em comparação de preferências (pairwise).
   • Vantagem de Escala: Enquanto métodos de preferência exigem dados de comparação entre todos os pares de modelos (escala $O(N^2)$), o ERP requer apenas um modelo de previsão por modelo individual (escala $O(N)$). Adicionar um novo modelo ao pool exige apenas treinar um novo preditor para ele, sem retreinar o sistema inteiro.
4. Explicação Teórica para Métodos Existentes: Mostram que o sucesso de métodos complexos de roteamento (como o Zooter, que prevê a distribuição de vencedores) pode ser explicado pela simples previsibilidade da recompensa esperada.

4. Resultados

• Desempenho Preditivo: Modelos lineares alcançaram um $R^2$ significativo (variando de ~0.19 a ~0.59 dependendo do modelo e categoria), indicando que uma fração substancial da variação na qualidade da resposta é capturada apenas pelo embedding do prompt.
• Roteamento vs. Baselines:
  • O roteamento baseado em ERP superou a estratégia de "melhor modelo por categoria" (que usa conhecimento órfão das categorias) e roteamento aleatório.
  • Em curvas de Pareto (compromisso entre custo computacional e arrependimento/regret), o ERP dominou a maioria das baselines, oferecendo um trade-off superior.
  • O ERP teve desempenho comparável ao método Zooter (que é mais complexo), sugerindo que a complexidade adicional do Zooter não é necessária se a recompensa esperada for previsível.
• Correlação com Taxa de Vitória: A diferença nas recompensas esperadas previstas correlaciona-se fortemente com a taxa de vitória real de um modelo contra outro em um prompt específico (validado via AUC em tarefas de classificação binária).

5. Significado e Implicações

• Eficiência em Inferência: A capacidade de prever qual modelo performará melhor para um prompt específico antes de gerar a resposta permite otimizar custos computacionais em tempo de inferência, direcionando prompts simples para modelos menores e prompts complexos para modelos maiores.
• Simplicidade e Extensibilidade: A abordagem elimina a necessidade de dados de preferência pareada complexos para novos modelos, tornando a gestão de pools de modelos heterogêneos muito mais prática.
• Novas Aplicações: Além do roteamento, essa propriedade abre portas para outras aplicações em tempo de inferência, como modificação dinâmica de prompts de sistema ou seleção de hiperparâmetros baseada na estimativa de recompensa.

Conclusão: O artigo estabelece que, na prática, modelos de recompensa modernos (fine-tuned em LLMs) codificam informações suficientes nos prompts para permitir a previsão precisa da qualidade esperada de diferentes modelos. Isso transforma o problema de seleção de modelos de um problema de comparação de pares complexo para um problema de regressão simples e escalável.