Pense com Esforço: Iterações Latentes Seletivas para Aprimorar Modelos de Linguagem de Raciocínio

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

November 11, 2025

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

cs.AI

Resumo

A melhoria das capacidades de raciocínio de Modelos de Linguagem de Grande Porte (LLMs), especialmente sob restrições de parâmetros, é crucial para aplicações do mundo real. Trabalhos anteriores propõem transformadores recorrentes, que alocam um número fixo de iterações extras por token para melhorar a qualidade da geração. Após a primeira passagem direta padrão, em vez da verbalização, os estados ocultos da última camada são realimentados como entradas para iterações adicionais, refinando as previsões dos tokens. No entanto, identificamos um fenômeno de *overthinking* latente: previsões fáceis de tokens que já estão corretas após a primeira passagem são, por vezes, revisadas para erros em iterações adicionais. Para resolver isto, propomos o Think-at-Hard (TaH), um método de pensamento latente dinâmico que itera mais profundamente apenas em tokens difíceis. Ele emprega um decisor neural leve para acionar iterações latentes apenas em tokens que provavelmente estão incorretos após a passagem direta padrão. Durante as iterações latentes, módulos de Adaptação de Baixa *Rank* (LoRA) deslocam o objetivo do LLM da previsão geral do próximo token para o refinamento focado de tokens difíceis. Introduzimos ainda um mecanismo de atenção duo-causal que estende a atenção da dimensão da sequência de tokens para uma dimensão adicional de profundidade de iteração. Isto permite o fluxo de informação entre iterações, mantendo o paralelismo sequencial total. Experiências mostram que o TaH aumenta o desempenho de raciocínio de LLMs em cinco *benchmarks* desafiadores, mantendo a mesma contagem de parâmetros. Em comparação com linhas de base que iteram duas vezes para todos os tokens de saída, o TaH proporciona ganhos de precisão de 8,1-11,3%, enquanto isenta 94% dos tokens da segunda iteração. Contra fortes modelos Qwen3 de iteração única ajustados com os mesmos dados, também proporciona ganhos de precisão de 4,0-5,0%. Ao permitir menos de 3% de parâmetros adicionais do LoRA e do decisor de iteração, os ganhos aumentam para 8,5-12,6% e 5,3-5,4%, respetivamente. O nosso código está disponível em 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.