Pequeno Controlador RL, Grande Modelo de Linguagem: Amostragem Adaptativa Guiada por RL para Escalonamento em Tempo de Teste

Resumo

A escala em tempo de teste melhora o desempenho de raciocínio de grandes modelos de linguagem, mas incorre em custos substanciais tanto em computação total quanto em latência. Métodos existentes de amostragem adaptativa mitigam parcialmente esse problema ao decidir dinamicamente quando parar a amostragem, mas geralmente dependem de regras heurísticas ou de suposições sobre a distribuição. Neste trabalho, formulamos a amostragem adaptativa como um processo de decisão de Markov (MDP). Treinamos um controlador de amostragem leve com aprendizado por reforço (RL) para equilibrar conjuntamente a correção da resposta, a latência e o custo computacional. Em cada rodada, o controlador decide parar a amostragem ou adquirir amostras adicionais. Nosso método é leve, baseando-se apenas em estatísticas das respostas finais, e pode ser treinado e implantado em CPU. Mostramos ainda que o arcabouço resultante admite uma interpretação como a relaxação Lagrangiana de um problema de otimização com restrições orçamentárias explícitas. Experimentos contra baselines fortes, como ASC e ESC, mostram que nosso método alcança trade-offs melhorados entre correção da resposta, rodadas de amostragem e número total de amostras necessárias.

English

Test-time scaling improves the reasoning performance of large language models but incurs substantial cost in both total computation and latency. Existing adaptive sampling methods partially mitigate this issue by dynamically deciding when to stop sampling, yet they typically rely on heuristic rules or rely on distribution assumptions. In this work, we formulate adaptive sampling as a Markov decision process (MDP). We train a lightweight sampling controller with reinforcement learning (RL) to jointly balance answer correctness, latency, and computation cost. At each round, the controller decides to stop sampling or to acquire additional samples. Our method is lightweight which only relies on statistics of final answers, and can be trained and deployed on CPU. We further show that the resulting framework admits an interpretation as the Lagrangian relaxation of a constrained optimization problem with explicit budget constraints. Experiments against strong baselines such as ASC and ESC show that our method achieves improved trade-offs among answer correctness, sampling rounds, and total samples required.