Resumo Técnico: Além da Fronteira: Backtracking Estocástico para Escalonamento Eficiente em Tempo de Teste

Declaração do Problema

O escalonamento em tempo de teste (TTS) visa melhorar o raciocínio de Modelos de Linguagem de Grande Escala (LLMs) alocando computação adicional de inferência para explorar múltiplas trajetórias de solução. Embora os Modelos de Recompensa de Processo (PRMs) tenham surgido como um mecanismo padrão para guiar essa busca ao pontuar prefixos intermediários, os métodos atuais de última geração sofrem de uma limitação crítica: busca apenas na fronteira.

Em abordagens de apenas fronteira (por exemplo, Busca em Feixe padrão, SMC), o estado da busca é restrito aos filhos recém-gerados em cada etapa. Prefixos previamente gerados são podados irreversivelmente com base em pontuações ruidosas do PRM. Isso cria um modo de falha de "comprometimento prematuro":

1. Decisões Irreversíveis: Um prefixo que é corretamente pontuado baixo devido ao ruído do PRM (ou uma incompatibilidade entre o rollout de treinamento do PRM e o LLM base implantado) é descartado permanentemente, mesmo que admita uma continuação correta.
2. Colapso da Diversidade: Prefixos incorretos superpontuados dominam a busca, fazendo com que o sistema desperdice computação em caminhos improdutivos enquanto perde o acesso a estados históricos potencialmente corretos.
3. Ineficiência: A incapacidade de revisitar estados históricos força o modelo a gastar tokens excessivos para se recuperar de erros de poda precoce.

Metodologia

O artigo propõe Backtracking Estocástico sobre um Pool Persistente, um framework que mantém um pool de memória linear de prefixos históricos, permitindo que a busca revisite e expanda estados previamente gerados, em vez de estar confinada à fronteira atual. Os autores instanciam esse framework com dois algoritmos complementares:

1. Seleção de Subpool (SPS)

A SPS aborda o "problema do bloqueador" na seleção gananciosa, onde prefixos incorretos superpontuados bloqueiam permanentemente o acesso a prefixos históricos corretos, mas com pontuação mais baixa.

• Mecanismo: Em vez de selecionar os $M$ melhores candidatos de todo o pool persistente (o que favorece bloqueadores de alta pontuação), a SPS constrói um subpool aleatório $D_t$ de tamanho $K_t$ a partir do pool histórico $P_{t-1}$.
• Seleção: Em seguida, seleciona os $M$ melhores candidatos estritamente dentro desse subpool aleatório.
• Razão Adaptativa: O tamanho do subpool é determinado por uma razão adaptativa $\rho_t$, calculada com base na pontuação média do PRM do pool atual. Se o pool tiver baixa confiança (pontuação média baixa), o subpool é menor para evitar prefixos incorretos de alta pontuação; caso contrário, é maior.
• Efeito: Essa amostragem estocástica dá aos prefixos históricos "bloqueados" uma chance de serem selecionados e expandidos, contornando a dominância dos candidatos da fronteira superpontuados.

2. Power Backtrack Sequential Monte Carlo (PB-SMC)

O PB-SMC estende os métodos de Monte Carlo Sequencial (SMC) para o pool persistente, permitindo recuperação estocástica de decisões de poda anteriores.

• Alvo do Pool Persistente: Diferentemente do SMC padrão, que restringe partículas à fronteira atual, o PB-SMC mantém um pool contendo tanto filhos recém-gerados quanto um subconjunto retido de prefixos históricos.
• Pontuações de PRM Potencializadas: O método aplica uma potência $\beta_t$ às pontuações do PRM ($r_{PRM}(z)^{\beta_t}$) para afiar a distribuição de recompensa, acentuando prefixos de alta pontuação enquanto reduz o peso dos de baixa pontuação.
• Pesos Corrigidos por Mistura: O algoritmo trata a busca como um problema de Amostragem por Importância Múltipla (MIS). Ele define uma distribuição proposta de mistura $q_{mix}$ combinando:
  • $q_{new}$: Novos filhos gerados a partir da etapa anterior.
  • $q_{hist}$: Prefixos históricos retidos de etapas anteriores.
• Ponderação: Deriva pesos específicos $W_t(z)$ que corrigem a proposta de mistura, garantindo que o estimador convirja para a distribuição proxy alvo potencializada. Isso permite que o sistema reamostre estocasticamente partículas históricas que haviam sido previamente podadas.

Eficiência de Memória e Computação

• Memória Linear: O pool persistente armazena apenas estados de texto e pontuações do PRM, não caches completos de Chave-Valor (KV). Os caches KV são recomputados apenas quando um prefixo histórico é selecionado para expansão.
• Sobrecarga: O artigo argumenta que a sobrecarga computacional de recomputar caches KV para estados de backtracking é negligenciável em comparação com os ganhos em precisão e eficiência de tokens.

Principais Contribuições

1. Identificação da Falha Apenas na Fronteira: O artigo identifica a seleção irreversível apenas na fronteira como um modo de falha primário no TTS guiado por PRM, onde pontuações ruidosas levam à perda permanente de caminhos de solução corretos.
2. Framework de Pool Persistente: Introdução de uma alternativa de memória linear que mantém prefixos históricos elegíveis para expansão futura, desacoplando a elegibilidade da busca da fronteira imediata.
3. Dois Algoritmos Novos:
   • Seleção de Subpool (SPS): Um aprimoramento de busca gananciosa que usa subpools aleatórios para contornar bloqueadores superpontuados.
   • Power Backtrack SMC (PB-SMC): Uma extensão do SMC usando pontuações de PRM potencializadas e pesos corrigidos por mistura para permitir backtracking estocástico.
4. Garantias Teóricas: Prova de que o PB-SMC converge em probabilidade para o alvo proxy definido pelo PRM potencializado, garantindo correção assintótica em relação ao objetivo proxy.

Resultados Experimentais

Os métodos foram avaliados em quatro benchmarks de raciocínio matemático (MATH500, AMC23, AIME24, Minerva Math) usando três LLMs (Qwen2.5-7B, Qwen2.5-3B, Phi-4-mini) e o Qwen2.5-Math-PRM-7B.

• Eficiência de Tokens: Tanto a SPS quanto o PB-SMC alcançaram precisão competitiva ou superior usando apenas $N=8$ candidatos por etapa, enquanto bases fortes (Busca em Feixe, SMC, DVTS) exigiam $N=32$.
  • SPS: Reduziu o consumo de tokens em 49,4% em média, ao mesmo tempo que melhorou a precisão.
  • PB-SMC: Reduziu o consumo de tokens em 36,7% em média, ao mesmo tempo que melhorou a precisão.
• Ganhos de Precisão:
  • No Qwen2.5-7B-Instruct, a SPS alcançou 54,66% de precisão média, superando o DVTS (51,40%) e a Busca em Feixe (52,12%).
  • No Phi-4-mini-Instruct, o PB-SMC alcançou a maior precisão geral (48,69%).
• Estudos de Ablação:
  • Backtracking: Adicionar backtracking ao SMC (Backtrack SMC) melhorou a precisão média de 49,11% para 52,27%.
  • Estocasticidade: Substituir a seleção global gananciosa pela Seleção de Subpool (SPS) melhorou a precisão média de 51,61% para 54,66%.
• Tempo de Execução: Apesar da sobrecarga potencial de recomputar caches KV, os métodos foram aproximadamente 2,5x mais rápidos que a Busca em Feixe devido ao tamanho reduzido do conjunto de candidatos ($N=8$ vs $N=32$), com uma penalidade de tempo de execução negligenciável (~10%) em comparação com as bases competitivas mais próximas.

Significado e Alegações

O artigo alega que o backtracking estocástico de pool persistente fornece um mecanismo simples e eficaz para melhorar o trade-off precisão–token no escalonamento em tempo de teste. O significado central reside na mudança do foco de otimização de meramente pontuar prefixos para manter o acesso a eles. Ao permitir que a busca revisite estados históricos, os métodos mitigam os riscos de pontuações ruidosas do PRM e de comprometimento prematuro.

Os autores posicionam seu trabalho como uma solução prática para implantações com restrições de recursos, demonstrando que o escalonamento eficaz em tempo de teste não requer aumentos massivos de computação, mas sim uma alocação mais inteligente da computação existente por meio de uma busca consciente do histórico. Os resultados sugerem que o acesso histórico é um fator crítico, muitas vezes negligenciado, na maximização da eficiência do raciocínio de LLMs.