Resumo Técnico: PS-PPO: Prefix-Sampling PPO para RLHF sem Crítico

1. Definição do Problema

O Aprendizado por Reforço com Feedback Humano (RLHF) é o método padrão para alinhar Grandes Modelos de Linguagem (LLMs) com as preferências humanas. No entanto, o treinamento de LLMs com RLHF enfrenta gargalos computacionais significativos, particularmente em abordagens sem crítico (ex: GRPO, RLOO) que evitam o treinamento de uma rede de valor separada.

Nesses métodos sem crítico, uma única recompensa escalar é atribuída a uma trajetória completa (conclusão) após a geração. Essa recompensa é transmitida uniformemente para todos os tokens na trajetória para calcular as vantagens. Consequentemente, as atualizações da política exigem a retropropagação de trajetória completa para cada rollout. Os autores argumentam que isso é ineficiente porque, em muitas tarefas de raciocínio (como o passo a passo matemático), prefixos intermediários frequentemente contêm informações suficientes para determinar o resultado final. A análise empírica dos autores (Figura 1) mostra que a taxa de sucesso de uma trajetória frequentemente se estabiliza bem antes do fim da conclusão, implicando que os tokens no sufixo carregam sinais de aprendizado redundantes. Atualizações de trajetória completa desperdiçam, assim, um poder computacional e memória de GPU substanciais nesses sufixos redundantes.

2. Metodologia: PS-PPO

Os autores propõem o Prefix-Sampling Proximal Policy Optimization (PS-PPO), um método de RLHF sem crítico e eficiente em termos computacionais que explora a redundância temporal ao truncar a passagem para trás (backward pass) mantendo um estimador de gradiente não enviesado.

Mecanismo Central

Em vez de retropropagar através de toda a conclusão $o = [o_1, \dots, o_T]$, o PS-PPO amostra um passo de corte (cutoff timestep) $H$ para cada trajetória. A atualização da política é computada apenas no prefixo $o_{\le H} = [o_1, \dots, o_H]$.

Para garantir que este truncamento não introduza viés, o PS-PPO utiliza um esquema de reponderação:

1. Distribuição de Corte: Uma distribuição condicionada ao prompt $\xi_{1:T}$ é definida, onde $\xi_t = \Pr(H \ge t | x)$ representa a probabilidade de o $t$-ésimo token ser retido. Esta distribuição é não crescente ($\xi_1 \ge \xi_2 \ge \dots \ge \xi_T$).
2. Estimador Não Enviesado: O gradiente para uma trajetória truncada é reponderado pelo inverso da probabilidade de inclusão $1/\xi_t$. O estimador resultante $\hat{G}(\theta)$ é:
   $$\hat{G}(\theta) := \frac{1}{K} \sum_{k=1}^K \sum_{t=1}^{H^{(k)}} \frac{1}{\xi_t} g_t^{(k)}(\theta)$$
   onde $g_t^{(k)}$ é o gradiente por passo de tempo. A esperança deste estimador sobre a amostragem de $H$ recupera a atualização de trajetória completa.

Otimizando a Distribuição de Corte

Os autores formulam a seleção de $\xi_{1:T}$ como um problema de otimização convexa. O objetivo é minimizar a variância do estimador de gradiente truncado sujeito a um orçamento de computação $B$ (o número esperado de tokens retropropagados).

• Surrogate de Variância: Os autores derivam um substituto tratável para a variância induzida pela amostragem de importância, aproximando-a como $\sum w_t(x) (1/\xi_t - 1)$, onde $w_t(x)$ representa a importância da atualização no passo de tempo $t$.
• Proxy de Importância: Como computar a norma exata do score $\|\nabla_\theta \log \pi_\theta\|^2$ requer uma passagem para trás completa (invalidando o propósito), os autores utilizam um proxy de apenas avanço (forward-only) baseado nas ativações do cabeçalho de saída e na incerteza da recompensa.
  • Incerteza da Recompensa ($u_t(x)$): Estimada usando a variância das recompensas dado o estado do prefixo, aproximada pela distância entre as distribuições de próximo token de rollouts bem-sucedidos e malsucedidos.
  • Proxy de Norma de Gradiente ($\bar{\gamma}_t$): Derivado da norma do gradiente do cabeçalho de saída.
• Otimização: O problema minimiza $\sum \frac{\bar{\gamma}_t(x,t) u_t(x)}{\xi_t}$ sujeito à restrição de orçamento e monotonicidade. A solução é encontrada usando o algoritmo Pool Adjacent Violators (PAV) para impor a restrição de monotonicidade.

Fluxo do Algoritmo

1. Gerar $K$ conclusões para um prompt.
2. Computar recompensas terminais e transmitir vantagens.
3. Estimar a incerteza por passo de tempo e proxies de gradiente a partir dos rollouts.
4. Resolver a otimização convexa para obter as probabilidades de corte monotônicas ótimas $\xi^*_{1:T}$.
5. Amostrar um corte $H^{(k)}$ para cada trajetória a partir da distribuição definida por $\xi^*$.
6. Realizar a retropropagação apenas nos tokens $t \le H^{(k)}$, reponderando a perda por $1/\xi^*_t$.

3. Principais Contribuições

1. Estrutura PS-PPO: Um novo método de RLHF sem crítico que introduz o truncamento estocástico condicionado ao prompt da passagem para trás, reduzindo significativamente os custos de computação e memória.
2. Estimador de Truncamento Não Enviesado: Uma derivação teórica de um estimador reponderado pela probabilidade de inclusão que recupera a atualização de transmissão total em esperança, sem exigir rollouts adicionais ou modelos de valor auxiliares.
3. Estratégia de Corte Otimizada: Uma abordagem fundamentada para determinar onde truncar, resolvendo um problema de otimização convexa que equilibra a redução de variância contra um orçamento de computação, utilizando proxies de apenas avanço para a importância do gradiente e incerteza da recompensa.
4. Validação Empírica: Demonstração de que o PS-PPO alcança precisão comparável às fortes bases sem crítico (GRPO, DAPO, RLOO), enquanto reduz substancialmente o tempo de treinamento e o pico de memória da GPU.

4. Resultados Experimentais

Os autores avaliaram o PS-PPO em benchmarks de raciocínio matemático (MATH500, AMC 2023, Minerva Math, AIME 2024/2025) usando Llama-3.1-8B-Instruct e Qwen2.5-Math-7B.

• Desempenho: O PS-PPO (Otimizado) alcançou precisão Pass@1 comparável ou ligeiramente superior às bases fortes como DAPO e S-GRPO em todos os benchmarks. Por exemplo, no MATH500 com Llama-3.1-8B, o PS-PPO alcançou 47,6% contra 46,8% do DAPO.
• Eficiência:
  • Tempo de Treinamento: O PS-PPO reduziu o tempo de treinamento por etapa em 33%–45% em comparação com as bases. Isso se deveu principalmente aos custos reduzidos de avanço/retrocesso, superando o overhead de computar a distribuição de corte.
  • Memória: O uso de pico de memória da GPU diminuiu em 15%–17% devido à redução no armazenamento de ativações para sequências truncadas.
  • Escalabilidade: Os ganhos de eficiência tornaram-se mais pronunciados conforme o comprimento máximo da conclusão ($T_{max}$) aumentava. Em $T_{max}=4096$, o PS-PPO foi 2,8x mais rápido que o S-GRPO e 3,3x mais rápido que o DAPO na etapa de atualização.
• Ablações:
  • Orçamento ($B$): Um orçamento de $B=128$ tokens ofereceu o melhor equilíbrio entre precisão e velocidade. Orçamentos menores aumentaram a variância; orçamentos maiores renderam retornos de precisão decrescentes com custos computacionais mais altos.
  • Estratégia de Corte: A estratégia "Otimizada" (resolvendo o problema convexo) superou significativamente as estratégias de truncamento uniforme, de prioridade temporal e heurística, confirmando que a seleção adaptativa de prefixos informativos é crucial.
  • Rollouts ($K$): Aumentar $K$ melhorou a estabilidade e o desempenho, mas aumentou o custo computacional por etapa. $K=8$ foi identificado como um compromisso prático.

5. Significância e Alegações

O artigo afirma que o PS-PPO oferece um caminho escalável para alinhar grandes modelos de linguagem sob requisitos substancialmente reduzidos de recursos. A principal significância reside na observação de que os prefixos intermediários em traços de raciocínio frequentemente contêm informação suficiente para determinar o resultado final, tornando a atualização de trajetória completa computacionalmente redundante.

Ao desacoplar o sinal de aprendizado do comprimento total da sequência através de amostragem de prefixo estocástica e ponderação de importância, o PS-PPO permite:

• Redução de Barreiras de Hardware: Menores custos de pico de memória e computação tornam o RLHF mais acessível a grupos acadêmicos e com recursos limitados.
• Treinamento Eficiente de Contexto Longo: O método é particularmente eficaz para traços de raciocínio longos, onde o custo da retropropagação completa é proibitivo.

Os autores mantêm um tom modesto, observando que o PS-PPO é um método de otimização que melhora a eficiência, mas não determina inerentemente o comportamento ou a segurança do modelo; estes permanecem dependentes das especificações de recompensa, dados de treinamento e salvaguardas de implantação. O método preserva as propriedades teóricas do objetivo de sequência total em esperança, garantindo que os ganhos de eficiência não ocorram à custa de um aprendizado enviesado.