Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling

Resumo Técnico: Amostragem Adaptativa Guiada por RL para Escalonamento em Tempo de Execução

Definição do Problema
O escalonamento em tempo de execução (test-time scaling) emergiu como uma técnica crítica para aprimorar as capacidades de raciocínio de Grandes Modelos de Linguagem (LLMs) sem treinamento adicional. Métodos como Autoconsistência (SC), Árvore de Pensamentos (Tree-of-Thoughts) e amostragem Best-of-N melhoram a qualidade das respostas ao alocar mais computação durante a inferência. No entanto, esses ganhos acarretam custos substanciais tanto em computação total quanto em latência. Métodos de amostragem adaptativa existentes, como a Autoconsistência Adaptativa (ASC) e a Autoconsistência com Parada Antecipada (ESC), tentam mitigar esses custos decidindo dinamicamente quando interromper a amostragem. Contudo, essas abordagens sofrem de limitações significativas: frequentemente dependem de heurísticas desenhadas por humanos ou suposições de distribuição, exigem sinais auxiliares (ex: confiança interna do modelo, estados ocultos ou dificuldade da questão) que podem estar indisponíveis, e podem ser invasivas para pipelines de inferência padrão. Há uma necessidade de um método leve e principiado que equilibre correção da resposta, latência e custo computacional sem exigir recursos auxiliares ou intervenção no modelo.

Metodologia
Os autores propõem a Amostragem Guiada por RL (RL-Guided Sampling), um framework que formula a amostragem adaptativa como um Processo de Decisão de Markov (MDP) de horizonte finito.

• Formulação do MDP:

  • Estado ($s_t$): O estado é construído exclusivamente a partir das estatísticas do pool atual de respostas amostradas ($D_t$). Ele inclui as contagens das $K$ classes de respostas mais frequentes, o número total de amostras geradas e a entropia dessas contagens. Crucialmente, o estado não depende de informações específicas da consulta ou sinais internos do modelo.
  • Ação ($a_t$): Em cada rodada, o controlador seleciona uma ação de um conjunto discreto: ou interrompe a amostragem (ação 0) ou gera um número específico de respostas candidatas adicionais em paralelo (ações $k_1, k_2, \dots$).
  • Função de Recompensa: A recompensa equilibra a qualidade da resposta final contra os custos de inferência.
    • Penalidade por Passo: Uma penalidade é aplicada para cada rodada de amostragem adicional ($\lambda_{lat}$) e para cada nova amostra gerada ($\lambda_{comp}$).
    • Recompensa Terminal: Ao término (seja pela interrupção pelo controlador ou ao atingir o orçamento máximo), o modelo recebe uma recompensa baseada na correção da resposta de voto majoritário. A correção alvo ($y^\star$) é definida como o voto majoritário que seria obtido se a amostragem continuasse até o orçamento máximo $N$, desacoplando a recompensa de rótulos de verdade fundamental (ground-truth).
  • Otimização: O controlador é treinado usando Otimização de Política Próxima (PPO) para maximizar a recompensa cumulativa esperada. Os autores observam que este objetivo pode ser interpretado como a relaxação de Lagrange de um problema de otimização restrita, onde o objetivo é maximizar a precisão sujeita a orçamentos de latência e computação.

• Arquitetura e Implementação:
  O controlador é um Perceptron de Camadas Múltiplas (MLP) leve de quatro camadas. Ele requer apenas as respostas finais geradas para funcionar, tornando-o não invasivo e implantável em CPUs. Não requer o ajuste fino (fine-tuning) do LLM subjacente.

Principais Contribuições

1. Formulação de MDP para Amostragem Adaptativa: O artigo introduz uma formulação inovadora de amostragem adaptativa como um MDP, afastando-se de regras heurísticas em direção a uma política aprendida que navega explicitamente pelo equilíbrio entre desempenho e custo.
2. Design Leve e Não Invasivo: O método baseia-se puramente nas estatísticas do pool de respostas amostradas, não exigindo sinais auxiliares (como pontuações de confiança) e nenhuma modificação no processo de raciocínio do LLM. O controlador é pequeno o suficiente para ser treinado e implantado em CPUs.
3. Interpretação Teórica: Os autores fornecem uma conexão teórica entre seu objetivo de RL e a otimização restrita, mostrando que a função de recompensa ponderada atua como uma relaxação de Lagrange para maximizar a precisão sob rigorosas restrições de recursos.
4. Generalização: O framework demonstra forte transferibilidade, onde políticas treinadas em um conjunto de dados ou modelo (ex: Qwen-3-0.6B) podem ser aplicadas a diferentes benchmarks e até a modelos diferentes e mais capazes (ex: GPT-4.1-nano) com degradação mínima de desempenho.

Resultados Experimentais
O método foi avaliado em três benchmarks de raciocínio matemático (AIME24, AIME25, HMT 2025) utilizando vários amostradores de LLM (família Qwen-3 e GPT-4.1-nano).

• Desempenho vs. Baselines: Comparado a baselines fortes como ASC e ESC, a Amostragem Guiada por RL alcança trade-offs superiores:
  • vs. ASC: Reduz as rodadas de amostragem em aproximadamente 3–4$\times$ e o total de amostras em ~30% mantendo uma precisão comparável. Isso aborda o gargalo de alta latência da amostragem sequencial da ASC.
  • vs. ESC: Reduz os requisitos totais de amostras em ~33% e as rodadas de amostragem em ~10%, alcançando maior precisão ao evitar paradas antecipadas excessivamente agressivas.
  • vs. SC: Embora todos os métodos adaptativos reduzam o total de contagens de amostras em relação à SC de orçamento fixo, a Amostragem Guiada por RL alcança isso com significativamente menos rodadas que a ASC e melhor eficiência que a ESC.
• Comportamento de Escalonamento: As curvas de escalonamento indicam que a Amostragem Guiada por RL supera consistentemente os baselines em várias configurações de parâmetros, oferecendo melhor precisão para o mesmo número ou menos amostras e rodadas.
• Alocação de Recursos: A análise mostra que a política aprendida intuitivamente aloca mais amostras para consultas com maior entropia de resposta (diversidade), alinhando-se com a necessidade de mais amostras para convergir em um voto majoritário confiável.
• Estudos de Ablação: O alvo de "Maioria Corrente" (voto majoritário no orçamento máximo $N$) mostrou-se superior ao uso de rótulos de verdade fundamental (Real Label), que degradou o desempenho ao introduzir ruído e impedir o aprendizado de critérios de parada universais.

Significância e Alegações
O artigo afirma que a Amostragem Guiada por RL oferece uma alternativa principiada e leve aos métodos de amostragem adaptativa existentes. Ao formular o problema como um MDP e treinar um controlador pequeno via RL, o método equilibra efetivamente a correção da resposta, a latência e o custo de computação sem exigir modificações invasivas no LLM ou acesso aos estados internos do modelo. Os autores enfatizam que a abordagem é particularmente valiosa para implantação prática, podendo ser treinada em modelos de código aberto econômicos e transferida para reduzir os custos de inferência de sistemas mais fortes, proprietários ou de código fechado. O trabalho representa um passo inicial em direção à amostragem adaptativa baseada em RL, com direções futuras sugeridas para incorporar sinais de estado mais ricos e métricas de custo do mundo real.