Think-at-Hard: Iterazioni Latenti Selettive per Migliorare i Modelli Linguistici di Ragionamento

Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models

November 11, 2025

Autori: Tianyu Fu, Yichen You, Zekai Chen, Guohao Dai, Huazhong Yang, Yu Wang

cs.AI

Abstract

Migliorare le capacità di ragionamento dei Large Language Model (LLM), specialmente in contesti con vincoli parametrici, è cruciale per le applicazioni nel mondo reale. I lavori precedenti propongono transformer ricorrenti, che allocano un numero fisso di iterazioni aggiuntive per token per migliorare la qualità della generazione. Dopo il primo passaggio in avanti standard, invece di procedere alla verbalizzazione, gli stati nascosti dell'ultimo layer vengono reimmessi come input per iterazioni aggiuntive per affinare le previsioni dei token. Tuttavia, identifichiamo un fenomeno di *overthinking* latente: previsioni di token semplici, già corrette dopo il primo passaggio, vengono talvolta modificate in errori durante le iterazioni aggiuntive. Per affrontare questo problema, proponiamo Think-at-Hard (TaH), un metodo di pensiero latente dinamico che itera più a fondo solo sui token difficili. Esso impiega un decisore neurale leggero per attivare iterazioni latenti solo sui token che sono probabilmente errati dopo il passaggio in avanti standard. Durante le iterazioni latenti, i moduli Low-Rank Adaptation (LoRA) spostano l'obiettivo dell'LLM dalla generica previsione del token successivo al raffinamento mirato dei token difficili. Introduciamo inoltre un meccanismo di attenzione *duo-causale* che estende l'attenzione dalla dimensione della sequenza di token a una dimensione aggiuntiva di profondità iterativa. Ciò consente un flusso di informazioni cross-iterazione mantenendo il pieno parallelismo sequenziale. Gli esperimenti mostrano che TaH aumenta le prestazioni di ragionamento degli LLM su cinque benchmark complessi mantenendo lo stesso numero di parametri. Rispetto ai baseline che iterano due volte per tutti i token di output, TaH fornisce un miglioramento in accuratezza dall'8.1% all'11.3%, esentando contemporaneamente il 94% dei token dalla seconda iterazione. Rispetto a forti modelli Qwen3 a singola iterazione addestrati con gli stessi dati, fornisce anche miglioramenti in accuratezza dal 4.0% al 5.0%. Quando si consentono meno del 3% di parametri aggiuntivi da LoRA e dal decisore di iterazione, i miglioramenti aumentano rispettivamente all'8.5-12.6% e al 5.3-5.4%. Il nostro codice è disponibile all'indirizzo https://github.com/thu-nics/TaH.

English

Improving reasoning capabilities of Large Language Models (LLMs), especially under parameter constraints, is crucial for real-world applications. Prior work proposes recurrent transformers, which allocate a fixed number of extra iterations per token to improve generation quality. After the first, standard forward pass, instead of verbalization, last-layer hidden states are fed back as inputs for additional iterations to refine token predictions. Yet we identify a latent overthinking phenomenon: easy token predictions that are already correct after the first pass are sometimes revised into errors in additional iterations. To address this, we propose Think-at-Hard (TaH), a dynamic latent thinking method that iterates deeper only at hard tokens. It employs a lightweight neural decider to trigger latent iterations only at tokens that are likely incorrect after the standard forward pass. During latent iterations, Low-Rank Adaptation (LoRA) modules shift the LLM objective from general next-token prediction to focused hard-token refinement. We further introduce a duo-causal attention mechanism that extends attention from the token sequence dimension to an additional iteration depth dimension. This enables cross-iteration information flow while maintaining full sequential parallelism. Experiments show that TaH boosts LLM reasoning performance across five challenging benchmarks while maintaining the same parameter count. Compared with baselines that iterate twice for all output tokens, TaH delivers 8.1-11.3% accuracy gains while exempting 94% of tokens from the second iteration. Against strong single-iteration Qwen3 models finetuned with the same data, it also delivers 4.0-5.0% accuracy gains. When allowing less than 3% additional parameters from LoRA and the iteration decider, the gains increase to 8.5-12.6% and 5.3-5.4%, respectively. Our code is available at https://github.com/thu-nics/TaH.