Strategic Scaling of Test-Time Compute: A Bandit Learning Approach

Título: Escalonamento Estratégico de Computação em Tempo de Teste: Uma Abordagem de Aprendizado de Bandit

1. Problema e Motivação

O avanço recente em Grandes Modelos de Linguagem (LLMs) deslocou o foco do compute (capacidade computacional) de treinamento para o compute em tempo de teste (Test-Time Compute - TTC). Técnicas como Best-of-N (amostragem repetida) e verificação de consistência melhoram a qualidade da saída gerando múltiplas respostas e selecionando a melhor.

No entanto, a maioria dos métodos atuais utiliza uma alocação uniforme de computação: atribuem o mesmo orçamento de inferência (número de gerações) para todas as consultas, independentemente da sua dificuldade.

• Ineficiência: Isso desperdiça recursos em consultas fáceis (que poderiam ser resolvidas com menos esforço) e subestima consultas difíceis (que precisam de mais tentativas).
• Limitação de Trabalhos Anteriores: Métodos adaptativos existentes focam ou na alocação dentro de uma única consulta ou dependem de procedimentos de duas etapas que exigem o treinamento de um modelo auxiliar, incorrendo em custos adicionais.

O objetivo deste trabalho é formular a alocação de computação como um problema de aprendizado de bandit (bandit learning), permitindo a alocação adaptativa de recursos entre um conjunto de consultas para maximizar a precisão global dentro de um orçamento fixo.

2. Metodologia

2.1 Formulação como Problema de Bandit

Os autores reformulam o problema de alocação de computação em tempo de teste como um problema de exploração pura (pure exploration) no contexto de aprendizado de bandit:

• Ação: Cada consulta $x$ no conjunto de dados é tratada como um "braço" (arm) do bandit.
• Feedback: Ao "puxar" um braço (alocar uma unidade de computação), o modelo gera uma resposta e recebe uma recompensa de um oráculo (verificador de verdade ou modelo de recompensa).
• Objetivo: Maximizar a fração de consultas respondidas corretamente dentro de um orçamento total de computação $B$, sem conhecer a priori a dificuldade de cada consulta.

Diferente dos bandits de limiar (thresholding bandits) que buscam identificar braços com recompensa acima de um limite, o objetivo aqui é garantir que pelo menos uma resposta de alta qualidade seja gerada para todas as consultas, independentemente da sua dificuldade esperada.

2.2 Algoritmo Proposto (Eliminação Adaptativa)

O artigo propõe um framework algorítmico (Algoritmo 1) baseado em regras de exploração e eliminação:

1. Conjunto Ativo: Mantém um conjunto de consultas que ainda não foram "confiavelmente" resolvidas.
2. Regra de Exploração: Em cada rodada, o algoritmo gera novas respostas para as consultas no conjunto ativo.
3. Regra de Eliminação: Se uma consulta atingir uma pontuação de recompensa acima de um limiar $\gamma$ (indicando que foi resolvida com confiança), ela é removida do conjunto ativo, liberando seu orçamento restante para consultas mais difíceis.
4. Estratégias de Exploração: Além da exploração uniforme básica, o framework suporta:
   • UCB (Upper Confidence Bound): Prioriza consultas com alta recompensa empírica e alta incerteza.
   • Gap-based: Foca em consultas próximas ao limiar de eliminação.
   • Entropia (Novo): Propõe selecionar consultas com padrões de resposta mais diversos (alta entropia). Isso é crucial para identificar consultas "solúveis" em conjuntos difíceis, onde consultas insolúveis tendem a gerar respostas inválidas ou pouco diversas.

2.3 Análise Teórica

Os autores provam teoremas que demonstram a superioridade teórica da alocação adaptativa sobre a uniforme:

• Para garantir uma probabilidade de sucesso de $1-\delta$ em todas as consultas, o algoritmo adaptativo requer um orçamento total da ordem de $\Theta(\sum \frac{1}{\Delta_x})$, onde $\Delta_x$ é a probabilidade de sucesso da consulta $x$.
• A alocação uniforme requer um orçamento da ordem de $\Omega(\frac{|S|}{\min \Delta_x})$.
• Conclusão: Em cenários heterogêneos (mistura de consultas fáceis e difíceis), a abordagem adaptativa pode ser exponencialmente mais eficiente (fator de $n$ ou mais) do que a alocação uniforme.

3. Contribuições Principais

1. Nova Formulação: Conecta a comunidade de scaling em tempo de teste com a teoria de aprendizado de bandit, oferecendo um framework de decisão preciso.
2. Framework Algorítmico Flexível: Um algoritmo geral que suporta diferentes estratégias de exploração (incluindo a nova regra baseada em entropia), métodos de agregação (como Self-Consistency) e cenários de streaming (consultas chegando sequencialmente).
3. Garantias Teóricas e Empíricas: Prova a vantagem de eficiência teórica e valida empiricamente que o método supera consistentemente as bases de comparação (uniforme e duas etapas) em benchmarks de matemática e código.

4. Resultados Experimentais

Os experimentos foram realizados em benchmarks padrão de matemática (MATH-500, AIME25) e código (LiveCodeBench), utilizando diversos modelos (Llama-3, Qwen, Gemini, DeepSeek-R1).

• Desempenho Geral: O algoritmo proposto (variante Elim) superou consistentemente as bases Uniform e Two-Stage.
  • MATH-500: Melhoria de até 11,10% (absoluta) na precisão, com ganho de eficiência de até 3,93x (ou seja, o método uniforme precisaria de quase 4x mais computação para atingir o mesmo desempenho).
  • AIME25: Melhoria de até 10,82% (absoluta) na precisão, com ganho de eficiência de 4,00x.
  • LiveCodeBench: Melhoria de até 11,23% (absoluta) na cobertura, com ganho de eficiência de 4,00x.
• Análise de Alocação:
  • Em conjuntos com consultas fáceis e difíceis, o algoritmo aloca menos recursos para as fáceis e mais para as difíceis.
  • Em conjuntos com consultas solúveis e insolúveis (MATH-500-Hard), a estratégia baseada em Entropia aprendeu a priorizar consultas solúveis, evitando desperdício em consultas que geram apenas respostas inválidas (baixa entropia).
• Robustez: O método manteve a superioridade mesmo em cenários de streaming (sem acesso ao conjunto total de consultas) e sob restrições de orçamento de tokens (não apenas número de gerações).

5. Significado e Conclusão

Este trabalho demonstra que a alocação inteligente e adaptativa de recursos computacionais em tempo de teste é fundamental para escalar o desempenho de LLMs de forma eficiente.

• Eficiência de Recursos: Permite obter ganhos de desempenho significativos sem aumentar o custo total de inferência, realocando o orçamento de consultas fáceis para as difíceis.
• Generalização: A abordagem é agnóstica ao modelo e ao tipo de tarefa (matemática, código, raciocínio), funcionando bem com diferentes oráculos de recompensa.
• Futuro: Abre caminho para sistemas de inferência que "pensam" de forma estratégica, decidindo dinamicamente quanto esforço dedicar a cada problema, aproximando-se de uma forma mais humana e eficiente de resolução de problemas.

Em resumo, o artigo estabelece que tratar a computação em tempo de teste como um problema de decisão sequencial (bandit) é uma estratégia superior à abordagem "tamanho único" (uniforme), oferecendo ganhos substanciais de precisão e eficiência computacional.