Resumo Técnico: Best-of-Better-N (BoBN)

Declaração do Problema

Métodos de alinhamento em tempo de inferência (ITA), como o Best-of-N (BoN), visam alinhar Grandes Modelos de Linguagem (LLMs) com preferências humanas ou requisitos de tarefas específicas sem modificar os pesos do modelo. Esses métodos operam em um framework de "amostragem e avaliação": eles amostram $N$ respostas de uma política de referência $\pi_{ref}$ e selecionam aquela com a maior recompensa de acordo com um modelo de recompensa $r$.

A eficácia do ITA é fundamentalmente limitada por dois pressupostos:

1. Cobertura (A1): A política de referência $\pi_{ref}$ deve amostrar pelo menos uma resposta de alta qualidade (alta recompensa).
2. Seleção (A2): O critério de seleção deve identificar com sucesso a resposta de alta qualidade entre as amostras.

A pesquisa atual em ITA foca amplamente em melhorar o (A2) através do refinamento dos critérios de seleção para mitigar o reward hacking. No entanto, um gargalo crítico permanece no (A1). Se $\pi_{ref}$ não estiver alinhada com a tarefa de destino (por exemplo, um modelo que não foi treinado para segurança ao tentar recusar prompts tóxicos), ela atribui probabilidade negligenciável a respostas de alta recompensa. Nesses casos, nenhuma estratégia de seleção pode ter sucesso porque a região de alta recompensa do espaço de saída é efetivamente inacessível para o modelo de referência. As soluções existentes geralmente exigem ajuste fino (fine-tuning, que é custoso e exige acesso aos pesos) ou utilizam exemplos estáticos de contexto que podem carecer de relevância específica para a tarefa.

Metodologia: Best-of-Better-N (BoBN)

Os autores propõem o Best-of-Better-N (BoBN), um framework de tempo de inferência projetado para deslocar a distribuição de amostragem de $\pi_{ref}$ em direção à região de alta recompensa sem necessidade de ajuste fino ou acesso aos pesos. O BoBN utiliza Aprendizado em Contexto (ICL) combinado com um pipeline de recuperação e reestilização (restyling).

O fluxo de trabalho do BoBN consiste em três estágios:

1. Recuperação (Retrieval): Dado um prompt de entrada $x$, o sistema recupera os top-$K$ pares prompt-resposta $(x_i, y_i)$ mais similares de um conjunto de dados $D$ usando um modelo de embedding. Esses exemplos servem como demonstrações de comportamentos de alta recompensa relevantes para a consulta específica.
2. Reestilização (Restyling): Para garantir que os exemplos recuperados se alinhem com as restrições, o formato e a persona da tarefa alvo, o modelo de referência $\pi_{ref}$ é usado para "reestilizar" as respostas recuperadas $y_i$. Especificamente, $\pi_{ref}$ reescreve $y_i$ em $\hat{y}_i$ condicionado a um prompt como "Reescreva $y_i$ para corresponder ao estilo $T$". Este passo garante a consistência distributiva ao nível de token entre o contexto e o modelo de referência, enquanto preserva a lógica de alinhamento dos dados recuperados.
3. Geração e Seleção: Os exemplos reestilizados $C_K = \{(x_i, \hat{y}_i)\}$ são prefixados ao prompt de entrada $x$ como exemplos de contexto. O modelo de referência então amostra $N$ respostas da distribuição condicionada $\pi_{ref}(\cdot | C_K, x)$. Finalmente, um critério de seleção (ex: BoN) escolhe a melhor resposta com base no modelo de recompensa.

Ao condicionar em $C_K$, o BoBN efetivamente desloca a massa de probabilidade de $\pi_{ref}$ para a região de alta recompensa, aumentando a probabilidade de que candidatos de alta qualidade sejam gerados primordialmente.

Principais Contribuições

1. Framework Inovador: A proposta do BoBN, uma abordagem baseada em ICL que melhora a cobertura de $\pi_{ref}$ sobre respostas de alta qualidade utilizando exemplos recuperados e reestilizados.
2. Análise Teórica: Uma caracterização teórica de como a qualidade da recuperação controla a cobertura e o desempenho downstream. Os autores demonstram que a recuperação relevante para a tarefa reduz comprovadamente a cobertura (definida como a divergência qui-quadrado entre a política de ouro e a política de referência) em comparação com ICL estático ou sem ICL. Eles mostram que essa redução na cobertura se combina com a seleção BoN para melhorar o desempenho da tarefa.
3. Validação Empírica: Experimentos extensos nos benchmarks de alinhamento de segurança (StrongReject, SORRY-Bench) e raciocínio matemático (GSM8k, MATH500). Os resultados mostram que o BoBN melhora o desempenho em múltiplos LLMs de referência (zephyr-7b, tulu3-8b), tanto quando $N=1$ (qualidade da resposta individual) quanto quando $N>1$ (combinado com seleção).

Resultados Experimentais

Os autores avaliaram o BoBN contra o BoN padrão (usando $\pi_{ref}$ diretamente) e o URIAL (ICL vanilla com exemplos estáticos).

• Alinhamento de Segurança: No benchmark StrongReject, o BoBN superou significativamente os baselines. Para o modelo zephyr-7b, o BoBN com $N=16$ alcançou uma pontuação de 0.033 (próxima do ótimo), comparável a modelos de fronteira como o Claude Sonnet-3.7, enquanto o $\pi_{ref}$ sozinho obteve 0.425. O BoBN também demonstrou cobertura superior (menor $\log C$) em comparação ao URIAL e $\pi_{ref}$.
• Raciocínio Matemático: No GSM8k e MATH500, o BoBN gerou melhorias consistentes de acurácia. Para o zephyr-7b no GSM8k, o BoBN alcançou 58.7% de acurácia com $N=16$, um ganho significativo sobre os 43.9% alcançados pelo $\pi_{ref}$ com BoN. Notavelmente, o BoBN com $N=4$ frequentemente igualou ou excedeu o desempenho dos baselines com $N=16$, indicando maior eficiência de amostragem.
• Custo e Latência: Embora o BoBN incorra em latência adicional devido ao passo de reestilização (exigindo $K$ chamadas ao $\pi_{ref}$ antes da geração), ele oferece uma eficiência de tokens significativa. O BoBN com $N=4$ economizou até 50% do uso de tokens em comparação ao URIAL com $N=16$, alcançando desempenho similar ou superior.
• Estudos de Ablação:
  • Reestilização: A remoção do passo de reestilização resultou em um alinhamento de segurança mais fraco e maior cobertura (pior desempenho), confirmando a necessidade de alinhar os exemplos recuperados ao estilo da tarefa alvo.
  • Contagem de Recuperação ($K$): O desempenho melhorou conforme $K$ aumentou de 1 para 8, mas degradou quando $K=16$, sugerindo que recuperar muitos exemplos introduz dados menos relevantes que diluem o sinal de contexto.

Significância e Alegações

O artigo afirma que o BoBN aborda uma limitação fundamental do alinhamento em tempo de inferência: a incapacidade de modelos de referência não alinhados de amostrar respostas de alta recompensa. Ao deslocar a distribuição de amostragem via exemplos de contexto recuperados e reestilizados, o BoBN permite:

• Maior Qualidade com Orçamento Fixo: Melhor desempenho de tarefa para um número fixo de amostras $N$.
• Eficiência: Alcançar o desempenho alvo com um $N$ menor, reduzindo custos de inferência.
• Alinhamento Sem Treinamento: Melhorar o alinhamento para modelos que não foram explicitamente ajustados para a tarefa de destino, sem exigir acesso aos pesos ou verossimilhanças do modelo.

Os autores enfatizam que sua análise teórica, baseada em transformadores lineares, fornece uma ligação comprovável entre a qualidade da recuperação, a redução da cobertura e o desempenho da tarefa downstream, ofereando uma abordagem principista para melhorar o ITA além de meras heurísticas de seleção.