Resumo Técnico: REAR: Realinhamento de Preferência em Tempo de Teste através da Decomposição de Recompensa

Declaração do Problema

Alinhar Grandes Modelos de Linguagem (LLMs) com preferências diversas e personalizadas do usuário é um desafio crítico. Embora métodos de pós-treinamento como o Aprendizado por Reforço com Feedback Humano (RLHF) sejam eficazes, eles exigem curadoria de dados dispendiosa e treinamento adicional, tornando-os inflexíveis para necessidades dinâmicas ou específicas do usuário. O Escalonamento em Tempo de Teste (TTS) oferece uma alternativa sem treinamento, aproveitando computação adicional durante a inferência para melhorar a qualidade da saída. No entanto, os métodos atuais de TTS têm sido amplamente limitados a domínios verificáveis (ex: matemática, programação) onde a correção é facilmente julgada. Aplicar o TTS ao alinhamento de preferências é difícil porque a qualidade da resposta é holística e não redutível a uma resposta verificável simples. As abordagens atuais de alinhamento em tempo de teste frequentemente impõem sobrecargas computacionais fixas ou requerem o treinamento de cabeças de valor, limitando sua escalabilidade e eficiência.

Metodologia

O artigo introduz o REAR (REAlignment Reward), um novo framework que enquadra o alinhamento de preferência em tempo de teste como um problema de realinhamento. A ideia central é que o objetivo de treinamento de um modelo base pode não satisfazer de forma ótima a preferência de um usuário específico no momento da inferência. O REAR aborda isso decompondo a função de recompensa subjacente em dois componentes e reescalonando dinamicamente sua importância.

1. Decomposição de Recompensa

Os autores formalizam a estrutura de recompensa de um LLM pré-treinado dentro do framework de RL de Máxima Entropia. Eles postulam que a recompensa $r(s \oplus x_p, a)$ para um estado $s$ (pergunta) e preferência $x_p$ pode ser decomposta em:

• Recompensa relacionada à pergunta ($r_0$): A recompensa implícita que o modelo base maximiza para responder à pergunta corretamente.
• Recompensa relacionada à preferência ($r_p$): O componente de recompensa que captura o quão bem a resposta satisfaz a preferência $x_p$.

A recompensa composta é modelada como $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$, onde $\alpha$ é um coeficiente de equilíbrio inerente.

2. Recompensa de Realinhamento (REAR)

Para realinhar o modelo para um usuário específico, o REAR introduz um coeficiente controlável $\hat{\alpha}$ (ou $\lambda = \hat{\alpha}/\alpha$) para reescalonar o componente de preferência durante a inferência:
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

Crucialmente, os autores derivam uma formulação computacionalmente eficiente para esta recompensa. Ao aproveitar a relação entre políticas ótimas e funções de valor em RL de Máxima Entropia, eles mostram que $r_{REAR}$ pode ser expressa como uma combinação linear de log-probabilidades de política ao nível de token:
$$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$$
Esta função de pontuação depende apenas das log-probabilidades do modelo base sob dois contextos: o prompt contendo apenas a pergunta e o prompt completo incluindo as preferências. Isso não requer o treinamento de um modelo de recompensa separado ou a geração de tokens adicionais para pontuação.

3. Integração com Escalonamento em Tempo de Teste

A pontuação REAR é integrada em dois algoritmos de TTS:

• Amostragem Best-of-N (BoN): Gera $N$ respostas e seleciona aquela com a maior pontuação $S_{REAR}$.
• Busca em Árvore de Verificador Diverso (DVTS): Um algoritmo de busca em árvore que utiliza pontuações REAR para guiar a expansão e seleção de ramos, promovendo diversidade enquanto otimiza o alinhamento de preferência.

Principais Contribuições

1. Formalização do Realinhamento: O artigo formaliza o alinhamento de preferência em tempo de teste como um problema de realinhamento, propondo o REAR como uma recompensa computacionalmente eficiente derivada de um objetivo decomposto.
2. Pontuação Sem Treinamento: O REAR é derivado exclusivamente das probabilidades internas do modelo base, não exigindo treinamento adicional, treinamento de modelo de recompensa ou etapas de geração especulativa.
3. Algoritmos Escaláveis: Os autores desenvolvem dois algoritmos de tempo de teste escaláveis (BoN e DVTS) guiados pelo REAR que podem aumentar o desempenho da geração com orçamentos de amostragem variáveis.
4. Generalizabilidade: O método mostra-se capaz de generalizar além do alinhamento de preferência explícito para raciocínio matemático e tarefas visuais ao injetar texto de preferência específico da tarefa.

Resultados Experimentais

Os autores avaliaram o REAR em vários benchmarks usando modelos Qwen e Llama:

• Benchmarks de Alinhamento de Preferência: No PrefEval (cobrindo preferências explícitas/implícitas e escolhas), Multifaceted Bench e Ping-Pong (role-playing), os métodos guiados pelo REAR (tanto BoN quanto DVTS) superaram consistentemente os baselines, incluindo greedy decoding, Amulet, Linear Alignment (LA) e Best-of-N com Modelos de Recompensa Generativos (GenRM). Por exemplo, o DVTS com REAR melhorou em média 11,6% em relação ao greedy decoding em termos de métricas.
• Robustez de Contexto Longo: O REAR demonstrou superior robustez em cenários de contexto longo (até 16k tokens) comparado ao GenRM, que degradou significativamente conforme o comprimento do contexto aumentava.
• Raciocínio Matemático: Quando aplicado a tarefas matemáticas (MATH500, AIME, AMC) com o texto de preferência apropriado (ex: "Atue como um matemático rigoroso"), a seleção guiada pelo REAR superou a votação majoritária, mostrando ganhos de escalonamento consistentes.
• Tarefas Visuais: No MMHal-Bench para detecção de alucinação visual, o REAR melhorou as pontuações gerais e reduziu as taxas de alucinação em comparação ao greedy decoding e GenRM.
• Transferência Entre Modelos: O método generalizou-se para o Llama-3.1-8B-Instruct sem ajuste por modelo, confirmando que a decomposição de recompensa não é específica da família Qwen.
• Eficiência: Os métodos guiados pelo REAR foram significativamente mais rápidos que os baselines que exigem modelos de recompensa externos ou otimização iterativa (como o Amulet), pois evitam o carregamento de modelos separados e dependem de uma única passagem direta para pontuação.

Significância e Alegações

O artigo afirma que o REAR fornece uma solução escalável e geral para alinhar LLMs com as preferências do usuário em tempo de teste sem o overhead de treinamento ou modelos de recompensa externos. Ao decompor a recompensa em componentes de pergunta e preferência, o REAR permite o controle dinâmico sobre o equilíbrio entre responder à pergunta e aderir às preferências.

Os autores enfatem que o REAR permite que o escalonamento em tempo de teste se estenda a domínios subjetivos e abertos (como alinhamento de preferência e role-playing) onde verificadores tradicionais falham. Eles observam que, embora um $\lambda = 20$ padrão funcione robustamente em muitas tarefas, o método permite o ajuste da força do realinhamento. O trabalho posiciona-se como um passo para tornar o alinhamento em tempo de teste mais acessível e eficiente, preenchendo a lacuna entre o sucesso do TTS em domínios verificáveis e as complexidades do alinhamento de preferência subjetiva.

Limitações Reconhecidas: Os autores observam que o REAR expõe um hiperparâmetro ajustável ($\lambda$) ao usuário, o que pode exigir pequenas varreduras de validação para tarefas com distribuições de preferência significativamente diferentes. Além disso, embora o TTS seja mais leve que o ajuste fino (fine-tuning), identificar o orçamento computacional ideal continua sendo uma direção para trabalhos futuros.