Riepilogo Tecnico: Scaling del Test-Time per Modelli di Linguaggio Efficienti e Addestrabili tramite Local Branch Routing

Definizione del Problema

I recenti progressi nel ragionamento dei modelli di linguaggio (LLM) dimostrano che allocare ulteriore computazione durante l'inferenza (scaling del test-time) può migliorare significativamente le prestazioni. Tuttavia, gli approcci esistenti affrontano un compromesso critico:

• Campionamento Chain-of-Thought (CoT): Sebbene efficace, la CoT standard rimane single-threaded, impegnandosi su un singolo token prima di osservare le conseguenze delle alternative.
• Metodi Basati sulla Ricerca: Approcci come Tree of Thoughts o Monte Carlo Tree Search (MCTS) introducono una larghezza attraverso l'esplorazione di alberi di pensieri intermedi. Tuttavia, questi sono spesso a grana grossa, computazionalmente costosi e difficili da ottimizzare come una singola politica di decoding addestrabile end-to-end.
• Metodi Soft-Token: Tecniche come Soft Thinking o Multiplex Thinking offrono una larghezza a livello di token fondendo gli embedding dei candidati in miscele continue. Sebbene addestrabili, questi metodi sfumano l'identità discreta dei rami candidati, potenzialmente oscurando le informazioni specifiche del concetto.

Il paper sostiene che i modelli di linguaggio potrebbero beneficiare di un meccanismo di scaling del test-time a livello di token, leggero, che consenta alle decisioni di essere informate dagli stati nascosti dei candidati locali futuri senza incorrere nel costo di una ricerca completa a livello di soluzione o nell'ambiguità delle miscele soft.

Metodologia: Local Branch Routing (LBR)

Gli autori propongono il Local Branch Routing (LBR), un framework di decoding branch-and-route a livello di token e addestrabile. LBR aumenta il normale decoding autoregressivo con un albero di lookahead locale rotante.

Pipeline di Decoding

Ad ogni step di decoding $t$, dato un prefisso impegnato $x_{<t}$, LBR esegue un processo "prune–shift–grow" (potatura–spostamento–crescita) in quattro fasi:

1. Grow (Crescita): Invece di campionare un singolo token successivo, LBR espande un albero locale $\mathcal{A}_t$ di larghezza $K$ e profondità $L$ radicato in $x_{<t}$. Ogni nodo campionato viene inoltrato attraverso il modello di linguaggio per calcolare il suo stato nascosto $h_t(v)$. Questo crea un albero di stati nascosti locali piuttosto che un semplice insieme di stringhe candidate.
2. Route (Instradamento): Un router leggero $\rho_\phi$ osserva l'albero inoltrato e seleziona un sotto albero di profondità 1 da impegnare. Il router utilizza un meccanismo di set-attention per codificare ogni sotto albero candidato (ad esempio, un singolo token per $L=1$, o un token più le sue continuazioni per $L>1$) in un vettore. Questi vettori vengono confrontati tramite cross-subtree attention per produrre un punteggio relativo, permettendo al router di prendere una decisione contrastiva tra i fratelli futuri.
3. Prune (Potatura): Il router seleziona il miglior sotto albero (indice $k^*_t$). Tutti gli altri sotto alberi vengono scartati. Il token radice del sotto albero selezionato diventa il prossimo token impegnato $x_t$.
4. Shift and Regrow (Spostamento e Ricrescita): Il sotto albero selezionato viene spostato in avanti (la sua radice viene rimossa e i suoi discendenti diventano il nuovo prefisso per lo step successivo). Un nuovo livello viene fatto crescere per ripristinare la profondità dell'albero $L$, e il ciclo si ripete.

Obiettivo di Addestramento

LBR definisce una verosimiglianza di traiettoria d'albero trattabile che consente l'apprendimento per rinforzo end-to-end con ricompense verificabili (RLVR).

• Stocasticità: Solo due operazioni sono stocastiche: il campionamento dei nuovi nodi cresciuti $G_t$ dal modello base e il campionamento della decisione del router $k^*_t$. Potatura, spostamento e riutilizzo dei nodi sono deterministici.
• Verosimiglianza: Il log-probabilità di una traiettoria $F$ è la somma dei log-probabilità dei nuovi token cresciuti e delle probabilità di selezione del router:
  $$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$$
• Ottimizzazione: Il framework è addestrato utilizzando un obiettivo di rapporto di verosimiglianza (ad esempio, GRPO) dove l'advantage $A(F)$ da un verificatore premia sia la generazione di token del modello base che le decisioni di selezione del router.

Contributi Chiave

1. Branching a Livello di Token con Identità Discreta: A differenza dei metodi soft-token che fondono i candidati in miscele continue, LBR preserva i candidati come rami inoltrati discreti. Ciò consente al router di sfruttare gli stati nascosti post-candidato—informazioni rivelate solo dopo che un token candidato è stato generato ed elaborato dal modello.
2. Addestramento End-to-End Trattabile: Definendo una specifica verosimiglianza per la traiettoria dell'albero, LBR consente al modello base e al router di essere ottimizzati congiuntamente sotto lo stesso principio RLVR utilizzato nei sistemi di ragionamento a token discreti.
3. Scaling del Test-Time Efficiente: Il framework introduce ampiezza a livello di token con un lookahead superficiale ($L$ è piccolo, ad esempio 1 o 2), evitando l'overhead computazionale di una ricerca completa ad alta profondità, pur sfruttando comunque le evidenze oltre la distribuzione del next-token radice.
4. Router Set-Attention: Gli autori introducono un'architettura di router che utilizza la cross-subtree attention per confrontare i candidati fratelli, dimostrando che l'instradamento è una decisione contrastiva piuttosto che un reranking indipendente.

Risultati Sperimentali

Pianificazione Gerarchica Sintetica

Su un compito sintetico riguardante la raggiungibilità di grafi con traslazione radice:

• Prestazioni: LBR ha raggiunto la massima accuratezza target (0.967), superando Discrete CoT (0.914) e Soft Thinking (0.845).
• Analisi del Meccanismo: Esperimenti di probing hanno rivelato che gli stati nascosti post-candidato in LBR sono significativamente più predittivi del target finale rispetto agli stati pre-branching. Ciò conferma che l'inoltro dei candidati permette al modello di accedere a informazioni a valle che non sono disponibili per il decoding autoregressivo standard.
• Rappresentazione: A differenza di Soft Thinking, che sfuma le identità dei concetti tramite miscele continue, LBR preserva le identità discrete dei rami, rendendo il percorso corretto più recuperabile.

Benchmark di Ragionamento Matematico

Valutato su modelli DeepSeek-R1-Distill-Qwen (1.5B e 7B) su sei benchmark (Minerva, AIME, MATH500, ecc.):

• Guadagni di Prestazioni: LBR ha migliorato sia Pass@1 che Pass@32 rispetto a:
  • Discrete Chain-of-Thought.
  • Vanilla discrete-token RLVR.
  • Baselines di branching soft-token compatibili con RL (ad esempio, Multiplex Thinking).
• Ablazione della Profondità: Aumentare la profondità di lookahead da $L=1$ a $L=2$ ha ulteriormente migliorato le prestazioni, in particolare sul backbone 7B, suggerendo che continuazioni locali più profonde forniscono maggiori evidenze di routing.
• Ablazione del Router: Rimuovere la cross-subtree attention (usando un router indipendente) ha portato a prestazioni inferiori, validando il fatto che il confronto contrastivo tra i rami fratelli è cruciale.
• Dinamiche di Addestramento: Durante l'addestramento RLVR, l'entropia del router è diminuita man mano che il router diventava più deciso, mentre anche l'entropia del modello base nelle posizioni di routing è diminuita, indicando un co-adattamento in cui il modello impara a generare rami locali più distinguibili.

Significato e Rivendicazioni

Il paper sostiene che il Local Branch Routing offre una forma efficiente, addestrabile e discreta di scaling del test-time per i modelli di linguaggio. Il suo significato primario risiede nel dimostrare che:

1. Le evidenze post-candidato sono preziose: Le decisioni di routing beneficiano degli stati nascosti indotti dai token candidati, non solo della distribuzione pre-branching.
2. Il branching discreto è superiore alla fusione soft: Preservare le identità discrete dei rami permette un processo decisionale più chiaro rispetto alle miscele di token soft, che possono oscurare l'identità del concetto.
3. Scalabilità: Il framework fornisce una via di mezzo tra la CoT single-threaded e la costosa ricerca ad albero completa, consentendo uno scaling efficace attraverso il branching locale leggero.

Gli autori riconoscono il limite della località: LBR è vincolato dalla finestra di lookahead superficiale e non esegue una ricerca globale su soluzioni complete. Si suggerisce come lavoro futuro la combinazione di LBR con profondità adattiva o ricerca globale selettiva.