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

Problema

I Large Language Models (LLM) risolvono compiti complessi generando catene di pensiero (Chain-of-Thought, CoT) multi-step. Tuttavia, un singolo passo di ragionamento errato può compromettere l'intera soluzione. Per migliorare le prestazioni, è stata introdotta la tecnica di Test-Time Scaling (TTS), che prevede la generazione di multiple varianti di passi intermedi, la loro verifica e la selezione dei percorsi migliori (tramite strategie come Best-of-N o Beam Search).
L'approccio attuale per la verifica di questi passi si basa sui Process Reward Models (PRM), modelli critici separati che valutano la qualità di ogni stato parziale. Tuttavia, i PRM presentano tre limiti fondamentali:

1. Costo Computazionale Elevato: Richiedono l'esecuzione di un modello LLM aggiuntivo (spesso grande, 1.5B-8B parametri), aumentando significativamente l'uso di memoria GPU e i costi di inferenza.
2. Dipendenza da Annotazioni Costose: L'addestramento richiede grandi quantità di dati etichettati, spesso ottenuti tramite roll-out Monte-Carlo costosi o annotazioni umane.
3. Scarsa Generalizzazione: I PRM sono spesso addestrati su domini specifici (es. matematica) e faticano a generalizzare su compiti fuori dominio (OOD).

Metodologia: ReProbe

Gli autori propongono ReProbe, un'alternativa leggera e basata sull'addestramento che verifica i passi di ragionamento sondando gli stati interni dell'LLM target invece di analizzare solo il testo generato.

• Concetto Chiave: ReProbe assume che gli stati interni di un LLM (stati nascosti, pesi di attenzione, logit) contengano segnali di "fiducia" o credibilità intrinseci sul processo di ragionamento in corso.
• Architettura: ReProbe è un modulo plug-and-play che si appoggia a un LLM target "congelato" (frozen).
  • Input: Riceve le rappresentazioni interne dei token (es. stati nascosti di tutti i layer o pesi di attenzione + logit) generati dall'LLM target durante la produzione di un passo di ragionamento.
  • Elaborazione: Utilizza un piccolo encoder Transformer (meno di 10 milioni di parametri) per aggregare le caratteristiche a livello di token in un vettore a livello di passo.
  • Output: Un classificatore a due livelli che stima la probabilità che il passo corrente sia corretto (0 o 1).
• Addestramento e Dati:
  • ReProbe può essere addestrato in modo self-supervised (l'LLM stesso annota i propri passi) o utilizzando un LLM più grande come giudice (es. DeepSeek-R1).
  • I dati di addestramento provengono da PRM800K, ma il processo è altamente efficiente: il costo di annotazione è di circa 200$ e richiede solo 4 ore di GPU (GH200).
  • A differenza dei PRM, non richiede roll-out Monte-Carlo costosi né annotazioni umane massicce.

Contributi Chiave

1. Efficienza Estrema: ReProbe utilizza meno di 10M di parametri, rendendolo fino a 810 volte più leggero dei PRM più grandi (7-8B) e fino a 150 volte più leggero dei PRM medi, mantenendo prestazioni competitive o superiori.
2. Generalizzazione Superiore: Essendo basato su segnali interni piuttosto che su conoscenza esterna specifica del dominio, ReProbe mostra una capacità di generalizzazione (OOD) nettamente superiore ai PRM, eccellendo in compiti di pianificazione e domande di conoscenza generale dove i PRM falliscono.
3. Flessibilità di Addestramento: Dimostra che è possibile addestrare verificatori efficaci sia con annotazioni esterne che in modalità self-supervised, riducendo drasticamente la dipendenza da dati etichettati costosi.
4. Integrazione con TTS: ReProbe si integra perfettamente nelle strategie di Test-Time Scaling (Best-of-N e Beam Search), guidando la selezione dei percorsi di ragionamento migliori con un overhead computazionale minimo.

Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi domini: Matematica (In-Domain), Pianificazione e QA (Out-of-Domain), utilizzando modelli come Qwen3 e Phi-4.

• Verifica a Livello di Passo: ReProbe raggiunge un PR-AUC (Area Under the Curve per la Precision-Recall) competitivo o superiore rispetto ai PRM più grandi. In particolare, sui task Out-of-Domain, ReProbe supera significativamente i PRM, che tendono a overfittare sui dati matematici di addestramento.
• Test-Time Scaling:
  • Nel setting Best-of-N, ReProbe permette a Qwen3-8B di superare il suo fratello più grande (Qwen3-14B) su diversi benchmark, dimostrando che una verifica efficiente può compensare la mancanza di parametri nel modello base.
  • Nel setting Beam Search, ReProbe ottiene le migliori prestazioni di accuratezza finale su tutti i dataset, superando i PRM di stato dell'arte.
• Efficienza Temporale: ReProbe offre un speedup di 2.6x - 25x rispetto ai PRM durante l'inferenza, riducendo drasticamente i costi di calcolo.
• Sinergia: Combinare i punteggi di ReProbe (fiducia interna) con quelli di un PRM (conoscenza esterna) in una strategia ibrida porta a ulteriori miglioramenti, suggerendo che i due approcci catturano aspetti complementari della qualità del ragionamento.

Significato e Impatto

Il lavoro di ReProbe segna un cambio di paradigma nella verifica del ragionamento degli LLM. Sposta l'attenzione dall'uso di modelli critici esterni e costosi (PRM) all'analisi introversa degli stati interni del modello generatore.

• Scalabilità: Rende il Test-Time Scaling economicamente ed energeticamente sostenibile, permettendo di scalare le capacità di ragionamento senza dover addestrare o eseguire modelli di verifica enormi.
• Accessibilità: La bassa barriera all'ingresso (pochi parametri, costi di addestramento minimi) rende questa tecnologia accessibile anche per modelli più piccoli o in contesti con risorse limitate.
• Futuro: Apre la strada a sistemi di verifica ibridi e a LLM più "introversi" (introspective), capaci di valutare la propria affidabilità in tempo reale, un passo cruciale verso sistemi di IA più robusti e affidabili.