Think-at-Hard: Селективные латентные итерации для улучшения языковых моделей рассуждений

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

November 11, 2025

Авторы: Tianyu Fu, Yichen You, Zekai Chen, Guohao Dai, Huazhong Yang, Yu Wang

cs.AI

Аннотация

Повышение способностей к рассуждению у больших языковых моделей (LLM), особенно при ограничениях на количество параметров, крайне важно для практического применения. Предыдущие работы предлагали рекуррентные трансформеры, которые выделяют фиксированное количество дополнительных итераций на токен для улучшения качества генерации. После первого стандартного прямого прохода вместо вербализации скрытые состояния последнего слоя подаются обратно на вход для дополнительных итераций с целью уточнения предсказаний токенов. Однако мы выявили феномен латентного «перемудривания»: предсказания простых токенов, уже верные после первого прохода, иногда исправляются на ошибочные в ходе дополнительных итераций. Для решения этой проблемы мы предлагаем метод Think-at-Hard (TaH) — динамическое латентное мышление, которое углубленно итерирует только на сложных токенах. Он использует легковесный нейросетевой модуль принятия решений для запуска латентных итераций только на тех токенах, которые, вероятно, неверны после стандартного прямого прохода. Во время латентных итераций модули Low-Rank Adaptation (LoRA) смещают цель LLM с общего предсказания следующего токена на целенаправленное уточнение сложных токенов. Мы также вводим механизм дуо-каузального внимания, который расширяет внимание от последовательности токенов до дополнительного измерения глубины итерации. Это обеспечивает поток информации между итерациями при сохранении полной последовательной параллельности. Эксперименты показывают, что TaH повышает производительность LLM в рассуждениях на пяти сложных бенчмарках при сохранении того же количества параметров. По сравнению с базовыми методами, которые выполняют две итерации для всех выходных токенов, TaH обеспечивает прирост точности на 8,1–11,3%, освобождая 94% токенов от второй итерации. По сравнению с сильными одноитерационными моделями Qwen3, дообученными на тех же данных, он также дает прирост точности на 4,0–5,0%. При использовании менее 3% дополнительных параметров от LoRA и модуля принятия решений прирост увеличивается до 8,5–12,6% и 5,3–5,4% соответственно. Наш код доступен по адресу 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.