Resumo Técnico: Alinhando Políticas de Busca em Árvore com Orçamentos Fixos de Tokens no Escalonamento de Tempo de Teste de LLMs

Declaração do Problema

Grandes Modelos de Linguagem (LLMs) podem melhorar a qualidade da resposta através do Escalonamento de Tempo de Teste (Test-Time Scaling), que aloca computação adicional durante a inferência sem atualizar os parâmetros do modelo. Embora a decodificação por busca em árvore (ex: Monte Carlo Tree Search, MCTS) seja uma forma eficaz de escalonamento de tempo de teste, implementações do mundo real frequentemente impõem um orçamento fixo de tokens por consulta.

As políticas de busca em árvore existentes são amplamente agnósticas ao orçamento. Elas tipicamente tratam o orçamento de tokens meramente como uma condição de terminação, em vez de uma restrição dinâmica que guia o comportamento da busca. Esse descompasso leva a dois modos primários de falha:

1. Ramificação excessiva no estágio tardio (Late-stage over-branching): A busca continua a expandir nós rasos no final do processo, esgotando o orçamento antes que possa refinar ou completar candidatos promissores.
2. Terminação prematura: A busca para precocemente, deixando uma parte do orçamento alocado sem uso.

Variantes atuais de parada antecipada focam em economizar tokens, mas carecem de uma política para mudar dinamicamente a estratégia de busca de uma exploração ampla para o refinamento conforme o orçamento restante diminui.

Metodologia: MCTS Guiado por Orçamento (BG-MCTS)

Os autores propõem o BG-MCTS (Budget-Guided MCTS), um algoritmo de decodificação de busca em árvore que condiciona explicitamente sua política de busca ao orçamento de tokens restante. O mecanismo central baseia-se em uma razão de suficiência de orçamento $\rho \in [0, 1]$, definida como:
$$\rho = 1 - \frac{C_{used}}{B}$$
onde $C_{used}$ é o acumulado de tokens gerados e $B$ é o orçamento total. À medida que $\rho$ diminui (aproximando-se de 0), a política muda da exploração para a explotação e conclusão.

O BG-MCTS introduz dois mecanismos de controle condicionados ao orçamento:

1. Seleção Guiada por Orçamento (BG-PUCT)

A seleção PUCT padrão é modificada para amortecer o bônus de exploração e introduzir um viés de conclusão:
$$\text{BG-PUCT}(p, s, \rho) = \frac{\tilde{W}(s, \rho)}{m_s} + \rho \cdot c \cdot P(s|p) \sqrt{\frac{\ln(m_p)}{m_s}}$$

• Amortecimento de Exploração: O termo de exploração é multiplicado por $\rho$. Conforme o orçamento se esgota ($\rho \to 0$), o incentivo para explorar novos ramos diminui.
• Viés de Conclusão: O valor acumulado $\tilde{W}(s, \rho)$ é corrigido por um termo proporcional à profundidade do nó $d(x)$ e $(1-\rho)$. Isso direciona a seleção para nós mais profundos que estão mais próximos de gerar uma resposta final, sem aumentar artificialmente nós que já contêm uma resposta completada.

2. Alargamento de Árvore Guiado por Orçamento (Budget-Guided Tree Widening)

Para evitar ramificações dispendiosas no estágio tardio, o BG-MCTS trata o "alargamento" (gerar um novo filho a partir de um pai) como uma ação selecionável ao lado dos filhos existentes. Um filho generativo virtual $s_{gen}(p)$ é adicionado ao conjunto de seleção. Seu escore é:
$$E_{gen}(p, \rho) = \mu(p) + \lambda \cdot \rho \cdot \sigma^2(p)$$

• $\mu(p)$ e $\sigma^2(p)$ representam a média e a variância dos escores dos filhos existentes.
• O termo é escalonado por $\rho$. Quando o orçamento é amplo, o alargamento é incentivado (especialmente se os filhos discordarem). Quando o orçamento fica escasso, o incentivo para o alargamento é suprimido, forçando a busca a aprofundar trajetórias promissoras existentes.

Algoritmo Unificado

O algoritmo segue um loop padrão de MCTS, mas torna tanto a decisão de seleção quanto a de alargamento dependentes de $\rho$. Isso induz um cronograma de amplo-para-profundo (wide-to-deep schedule): exploração ampla no início da busca, seguida de refinamento focado e conclusão da resposta conforme o orçamento é exaurido.

Principais Contribuições

1. Design de Política Consciente de Orçamento: O artigo introduz o primeiro framework de decodificação de busca em árvore que alinha dinamicamente o comportamento de busca (exploração vs. explotação vs. alargamento) com o orçamento de tokens restante, em vez de tratar o orçamento apenas como um sinal estático de parada.
2. Mecanismos Inovadores: A proposta do BG-PUCT (com exploração condicionada ao orçamento e correção de valor enviesada pela profundidade) e de Filhos Generativos Virtuais para alargamento consciente do orçamento.
3. Validação Empírica: Avaliação abrangente em benchmarks de raciocínio matemático (MATH500, AIME24/25) e um benchmark de raciocínio de física (UGPhysics) usando LLMs de pesos abertos (Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B).

Resultados Experimentais

Os autores avaliaram o BG-MCTS contra baselines agnósticos ao orçamento (MCTS Padrão, AB-MCTS-M, LiteSearch) e baselines de escalonamento (Amostragem Repetida, Refinamento Sequencial) sob orçamentos fixos de 10k, 20k e 30k tokens.

• Desempenho: O BG-MCTS superou consistentemente todos os baselines em todos os modelos e orçamentos. Por exemplo, no MATH500 com Qwen2.5-7B, o BG-MCTS alcançou 71,1% de precisão com 30k tokens, comparado aos 65,9% do segundo melhor baseline de busca em árvore (MCTS).
• Utilização de Orçamento: Ao contrário do LiteSearch, que frequentemente para prematuramente, o BG-MCTS continua a melhorar a precisão à medida que se aproxima da exaustão do orçamento.
• Comportamento Qualitativo: A análise das árvores de busca revelou que o BG-M_CTS mantém uma largura ampla no início, mas muda para o aprofundamento de ramos específicos no final da busca. Em contraste, o MCTS padrão frequentemente continua alargando nós rasos até que o orçamento seja esgotado, falhando em completar caminhos promissores.
• Raciocínio de Física: A eficácia do método estendeu-se ao UGPhysics, demonstrando que o comportamento de amplo-para-profundo alinhado ao orçamento é benéfico além do puro raciocínio matemático.

Significância e Alegações

O artigo afirma que condicionar as decisões de busca em árvore ao orçamento restante é mais eficaz do que tratar o orçamento meramente como uma condição de parada.

• Escalonamento Previsível: O BG-MCTS oferece uma interface simples para um escalonamento de tempo de teste previsível, onde os usuários especificam um orçamento de tokens e a política adapta automaticamente seu cronograma de exploração-para-conclusão dentro desse envelope de custo.
• Implicações para Síntese de Dados: Os autores observam que o BG-MCTS tende a produzir menos nós respondidos por instância de problema, porém com precisão significativamente maior (corretude). Este regime de "menos instâncias, melhores respostas" é destacado como particularmente valioso para a síntese de dados, onde traços de raciocínio de alta qualidade são mais críticos do que a cobertura bruta.
• Robustez: O método permanece eficaz mesmo ao usar modelos de recompensa de processo (PRMs) imperfeitos que fornecem sinais fracos ou saturados, sugerindo que mudanças estruturais conscientes do orçamento podem compensar sinais de avaliação ruidosos.

O trabalho conclui que alinhar as políticas de busca com restrições fixas de tokens é um passo crucial para implantar LLMs em ambientes do mundo real com recursos limitados.