Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models

Título

ReProbe: Escalabilidade Eficiente em Tempo de Teste de Raciocínio Multi-etapa Sondando Estados Internos de Grandes Modelos de Linguagem (LLMs)

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) modernos, especialmente os Modelos de Grande Raciocínio (LRMs), resolvem tarefas complexas gerando cadeias de pensamento (Chain-of-Thought - CoT). No entanto, um único passo de raciocínio incorreto pode invalidar toda a solução.
Para mitigar isso, utiliza-se a Escalabilidade em Tempo de Teste (Test-Time Scaling - TTS), onde o modelo gera múltiplas variações de raciocínio e um verificador seleciona os melhores passos. A abordagem atual dominante para verificação são os Modelos de Recompensa de Processo (PRMs).
Limitações dos PRMs:

• Custo Computacional: São modelos grandes (1.5B a 8B parâmetros), exigindo recursos significativos de GPU e memória durante a inferência.
• Dependência de Anotação: Requerem grandes quantidades de dados anotados (humanos ou via Monte-Carlo), o que é caro e difícil de escalar, especialmente em tarefas onde a correção de um passo não pode ser inferida apenas pelo resultado final (ex: provas matemáticas, planejamento).
• Generalização: Frequentemente treinados em domínios específicos (como matemática), falham ao generalizar para outros domínios (OOD - Out-of-Domain).

2. Metodologia: ReProbe

Os autores propõem o ReProbe, uma alternativa leve e baseada em treinamento que verifica a credibilidade dos passos de raciocínio sondando os estados internos do LLM gerador, em vez de analisar apenas o texto gerado.

• Mecanismo de Funcionamento:
  • O ReProbe é um módulo "plug-and-play" (um classificador leve) treinado sobre um LLM alvo congelado.
  • Ele utiliza sinais internos do modelo gerador, como pesos de atenção (attention weights), logits das candidatas e estados ocultos (hidden states) de todas as camadas.
  • A arquitetura consiste em uma projeção linear dos recursos, seguida por camadas de Transformer (para capturar dependências contextuais) e uma rede neural de classificação que prevê a probabilidade de um passo estar correto.
• Treinamento e Dados:
  • Anotação: Pode ser feita de duas formas:
    1. Supervisionada Externa: Um LLM maior (ex: DeepSeek-R1) avalia os passos.
    2. Auto-supervisionada: O próprio modelo gera e avalia seus próprios passos (self-supervised), eliminando a necessidade de anotação externa cara.
  • Eficiência: O modelo ReProbe possui menos de 10 milhões de parâmetros (comparado aos 1.5B-8B dos PRMs). O custo de treinamento é baixo (apenas ~4 horas de GPU GH200 e ~$200 em custos de anotação).

3. Contribuições Principais

1. Verificação Baseada em Estados Internos: Demonstra que os estados internos de um LLM contêm sinais de confiança suficientes para verificar a correção de passos de raciocínio, sem depender de um modelo externo de recompensa.
2. Eficiência Extrema: Redução drástica no tamanho do modelo (até 810x menor que os PRMs) e no custo de inferência, mantendo ou superando a performance.
3. Generalização Robusta: Ao contrário dos PRMs que sofrem de overfitting em domínios específicos (como matemática), o ReProbe generaliza bem para tarefas fora do domínio de treinamento (planejamento, QA geral).
4. Flexibilidade de Treinamento: Funciona tanto com anotação externa quanto auto-supervisionada, tornando o processo de criação de verificadores acessível e escalável.

4. Resultados Experimentais

Os autores avaliaram o ReProbe em três domínios: Matemática (In-Domain), Planejamento (Out-of-Domain) e Perguntas e Respostas Gerais (Out-of-Domain), utilizando modelos alvo como Qwen3 e Phi-4.

• Avaliação de Correção de Passos (PR-AUC):
  • O ReProbe igualou ou superou os PRMs mais avançados (7B-8B) em tarefas fora do domínio (OOD).
  • Em tarefas de matemática (In-Domain), o ReProbe foi competitivo com os melhores PRMs, embora estes tenham uma leve vantagem devido ao overfitting no domínio específico.
  • O ReProbe auto-supervisionado teve desempenho comparável ao supervisionado externamente, validando a viabilidade do treinamento sem anotação humana/experta.
• Escalabilidade em Tempo de Teste (TTS):
  • Best-of-N e Busca em Feixe (Beam Search): O uso do ReProbe para guiar a seleção de trajetórias resultou em melhores taxas de acerto final do que os PRMs em vários benchmarks.
  • Em alguns casos, o uso de ReProbe permitiu que um modelo menor (Qwen3-8B) superasse um modelo irmão maior (Qwen3-14B) em benchmarks de raciocínio.
• Eficiência Computacional:
  • O ReProbe é 2.6x a 25x mais rápido na inferência do que os PRMs de última geração.
  • Consome significativamente menos memória GPU.

5. Significado e Conclusão

O trabalho do ReProbe representa um avanço significativo na direção de sistemas de raciocínio mais eficientes e introspectivos.

• Viabilidade Prática: Torna a escalabilidade em tempo de teste acessível para aplicações com recursos limitados, substituindo a necessidade de modelos de verificação pesados.
• Insight Teórico: Confirma que a "confiança" do modelo em seu próprio raciocínio está codificada em seus estados internos e pode ser extraída de forma eficiente.
• Futuro: Sugere que a combinação de sinais de confiança internos (ReProbe) com conhecimento externo (PRMs) em uma estratégia híbrida pode oferecer o melhor equilíbrio entre custo e precisão.

Em resumo, o ReProbe oferece uma solução leve, escalável e generalizável para a verificação de raciocínio em LLMs, superando as limitações de custo e generalização dos modelos de recompensa de processo tradicionais.