Resumo Técnico: ETS: Escalonamento Guiado por Energia no Tempo de Teste para Alinhamento de RL sem Treinamento

Declaração do Problema

O pós-treinamento por Aprendizado por Reforço (RL) tornou-se um paradigma central para alinhar Modelos de Linguagem de Grande Escala (LLMs) com a intenção humana, particularmente para aprimorar o raciocínio e o seguimento de instruções. No entanto, essa abordagem sofre de limitações práticas significativas: exige modelagem de recompensa custosa, depende de dados de preferência humana em grande escala, exibe dinâmicas de treinamento instáveis e é sensível a hiperparâmetros. Além disso, qualquer mudança no design da recompensa exige retreinamento, levando a custos computacionais e humanos substanciais.

Embora a distribuição ótima de RL admita uma solução de forma fechada teoricamente, os métodos existentes dependem de otimização iterativa baseada em gradiente para aproximar essa solução, em vez de explorar sua estrutura diretamente. Isso levanta uma questão fundamental: Podemos amostrar diretamente da distribuição ótima de RL no momento da inferência sem qualquer treinamento adicional?

Metodologia

Os autores propõem o Escalonamento Guiado por Energia no Tempo de Teste (ETS), um método de inferência sem treinamento projetado para amostrar diretamente da política ótima de RL. O método é construído sobre um framework unificado de Modelagem de Linguagem Mascaramento (MLM) que abrange tanto Modelos Autoregressivos (ARMs) quanto Modelos de Linguagem de Difusão (DLMs).

1. Fundamento Teórico: Transição Reponderada por Energia

A principal insight teórica é que o kernel de transição reversa ótimo para o framework MLM decompõe-se em dois fatores:

1. Transição de Referência: O kernel de transição fornecido pelo modelo de referência ($p_{ref}$).
2. Termo de Energia: Uma expectativa condicional de recompensas exponenciadas, atuando como um termo de energia.

Especificamente, a transição reversa alvo $p(s | t, y)$ é proporcional a:
$$p_{ref}(s | t, y) \cdot \mathbb{E}_{p_{ref}(x_0|y,s)} \left[ \exp\left(\frac{r(y, x_0)}{\lambda}\right) \right]$$
onde $r(y, x_0)$ é a função de recompensa e $\lambda$ é uma constante de regularização. Essa formulação permite amostragem sem treinamento estimando o termo de energia via métodos de Monte Carlo.

2. Algoritmo: Escalonamento Guiado por Energia no Tempo de Teste (ETS)

O algoritmo opera como um processo iterativo de orientação:

• Amostragem: Em cada etapa de orientação, o algoritmo amostra $M$ continuação candidatas do modelo de referência.
• Estimativa de Energia: Para cada candidato, o termo de energia é estimado usando amostragem de Monte Carlo. Especificamente, $K$ conclusões independentes são geradas para cada candidato para computar um estimador de amostragem por importância auto-normalizado da recompensa esperada.
• Seleção: Os candidatos são reponderados com base em suas energias estimadas, e um é selecionado para prosseguir para a próxima etapa.
• Design de Recompensa: Para operar sem recompensas de verdade durante a inferência, o ETS emprega uma recompensa de autoconsistência. Ele extrai respostas finais das $K$ conclusões, identifica a votação majoritária como consenso e atribui uma recompensa de 1 para conclusões que correspondem ao consenso e 0 caso contrário.

3. Aceleração via Amostragem por Importância (ETS-IS)

Para abordar o gargalo de latência causado por consultas repetidas ao modelo para estimativa de energia, os autores introduzem o ETS-IS.

• Modelo de Proposta: Em vez de amostrar do grande modelo de referência ($p_{ref}$) para estimativa de energia, um modelo de proposta mais leve e computacionalmente mais barato ($p_{small}$) é utilizado.
• Correção: Pesos de amostragem por importância são aplicados para corrigir o viés introduzido pelo modelo de proposta, garantindo que o estimador permaneça não tendencioso e teoricamente consistente com a distribuição alvo.
• Implementação:
  • Para ARMs: Um modelo menor e alinhado em escala (por exemplo, Qwen3-1.7B para Qwen3-8B) serve como proposta.
  • Para DLMs: O Fast-dLLM é utilizado, aproveitando o cache KV por blocos e a decodificação paralela para acelerar a inferência.

Contribuições Principais

1. Alinhamento sem Treinamento: O artigo apresenta um método para amostrar diretamente da política ótima de RL sem atualizações de gradiente iterativas ou ajuste fino, resolvendo os problemas de instabilidade e custo do pós-treinamento de RL tradicional.
2. Framework Unificado: A abordagem é aplicável tanto a ARMs quanto a DLMs sob uma formulação MLM unificada, fornecendo uma solução geral para escalonamento sem treinamento.
3. Garantias Teóricas: Os autores fornecem provas rigorosas sobre a taxa de convergência da distribuição de amostragem para a distribuição ótima alvo em termos de Distância de Variação Total. O limite de erro é mostrado como escalando linearmente com o número de etapas de orientação e inversamente com a raiz quadrada do tamanho da amostra.
4. Aceleração Eficiente: Ao integrar amostragem por importância com modelos de proposta leves e frameworks de aceleração modernos (como Fast-dLLM), o método reduz significativamente a latência de inferência enquanto preserva a qualidade da amostragem.

Resultados Experimentais

Os autores avaliaram o ETS em benchmarks de raciocínio (MATH500, GSM8K), codificação (HumanEval) e ciência (GPQA) usando tanto ARMs (Qwen3-1.7B, Qwen3-8B) quanto DLMs (LLaDA-8B-Instruct).

• Desempenho: O ETS melhora consistentemente a qualidade de geração sobre a inferência padrão e outras bases de comparação de escalonamento no tempo de teste (por exemplo, Busca em Feixe, Best-of-N, Amostragem de Potência).
• Comparação com RL: Notavelmente, o ETS iguala ou supera o desempenho de políticas de RL pós-treinadas (como GRPO) sem qualquer atualização de parâmetros.
• Eficiência: O método de aceleração (ETS-IS) alcança um compromisso favorável entre latência e precisão. Por exemplo, no Qwen3-8B, o ETS-IS alcança precisão comparável ao ETS padrão, mas com latência significativamente reduzida.
• Escalonamento: Os experimentos mostram que aumentar o número de candidatos ($M$) é mais eficaz para melhorar a precisão do que aumentar o número de conclusões ($K$), e que etapas intermediárias de orientação ($I$) fornecem melhorias proporcionais na precisão.

Significado e Afirmações

O artigo afirma que o ETS estabelece um caminho prático e sem treinamento em direção à geração consciente de recompensas. Ao demonstrar que alinhamento de alta qualidade pode ser alcançado através de escalonamento no tempo de inferência em vez de pós-treinamento custoso, o trabalho oferece uma alternativa econômica e estável para configurações com recursos limitados. Os autores afirmam que seu método democratiza técnicas de alinhamento e reduz a pegada ambiental da implantação de LLMs, eliminando a necessidade de ciclos extensos de retreinamento. O framework é apresentado como uma fundação para pesquisas futuras sobre inferência eficiente e escalonamento no tempo de teste em modelos generativos.