![]()
來自清華大學、無問芯穹、上海交通大學等機構的研究團隊提出Think-at-Hard(TaH):一種面向小模型的選擇性潛空間迭代方法,讓 Looped Transformer 只在真正困難的 token 上多想一步,在跳過 93% 的額外迭代的同時,于 9 個數學、問答、代碼基準上取得 3.0%–6.8% 的穩定提升。
本文的共同第一作者為清華大學電子系 NICS-EFC 實驗室四年級直博生傅天予和大四本科生尤憶晨,并與無問芯穹合作完成。
隨著 o1/R1 等推理模型的發展[1][2],「讓模型多想一會兒」幾乎成了提升復雜推理能力的標準方案。更長的 Chain-of-Thought、更大的測試時計算、更深的內部推理,都在用更多計算換取更可靠的答案。
但一個問題很少被認真討論:模型真的有必要在每個 token 上都多想嗎?
對于參數受限的小模型來說[4][5],這個問題尤其關鍵。小模型成本低、速度快、適合邊緣部署,但在數學、代碼和問答任務中,往往會因為少數關鍵 token 預測錯誤,讓整條推理路徑偏離正確方向。已有的 Looped Transformer[6][7][8]試圖緩解這一問題:他們在生成每個 token 前,把最后一層的隱藏狀態送回模型做額外的潛空間迭代,相當于在不增加參數的情況下為每個 token 增加計算深度。
來自清華大學、無問芯穹、上海交通大學等機構的研究團隊在論文中發現,事情沒這么簡單:相當一部分 token 在第一次前向時已經預測正確,后續的潛空間迭代反而可能把正確的預測改錯。論文將這一現象命名為latent overthinking,也就是「潛空間過度思考」。
基于這一觀察,作者提出Think-at-Hard(TaH):一種選擇性潛空間迭代的 Looped Transformer。通過后訓練,讓標準模型變為 Looped Transformer,且只在真正困難的 token 上增加算力。本工作入選 ICLR LIT Workshop Best Paper Shortlist,并被 ICML 2026 接收。
![]()
潛空間迭代可以把錯誤預測改對,也會把正確預測改錯
論文的核心貢獻在于:
- 揭示并量化了 Looped Transformer 中的潛空間過度思考 (latent overthinking) 現象,指出統一深度的潛空間迭代會同時帶來「改對」和「改錯」。
- 提出 TaH,通過輕量級 iteration decider、duo-causal attention 與 depth-aware LoRA,實現 token 級動態思考。
- 在 9 個數學、問答、代碼基準測試上穩定取得提升;TaH 平均只讓約 7% 的 token 進入第二輪迭代,相比所有 token 均二次迭代,其基準測試精度反而提升 3.8-4.4%。
本工作現已開源,歡迎交流討論。
![]()
- 論文標題:Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models
- 代碼鏈接:https://github.com/thu-nics/TaH
- 主頁鏈接:https://fuvty.github.io/TaH_project_page
- 論文鏈接:https://arxiv.org/pdf/2511.08577
核心洞見
迭代計算簡單 token 反而損害模型性能
R2R 等前序工作指出,在語言模型的推理過程中,并不是所有 token 都同等重要[10][11]。真正決定推理路徑的,往往是少數關鍵位置:轉折、因果連接、中間結論等。
為了量化這種選擇性迭代的潛力,作者構造了一個oracle 策略:僅當模型第一次預測某個 token 出錯時,才讓它繼續在潛空間迭代;如果第一次已經預測正確,就直接輸出。實驗顯示,僅靠這個 oracle,模型在下游任務上最多就能換來 7.3% 的性能提升,而且只需要讓大約11–19%的 token 二次迭代;如果再換上為選擇性迭代優化的 TaH 架構,oracle 帶來的提升甚至超過25%。
這意味著推理時的計算動態分配應該細化到 token 級別:難題里也有簡單 token,簡單題里同樣可能出現關鍵 token。更關鍵的是,對簡單 token 強行多算不僅浪費算力,還會讓一部分本來預測對的 token 在第二次迭代里被改錯,即latent overthinking。
TaH
在困難處停下來多想
為了解決這一問題,TaH 的思路非常直接:簡單 token 快速輸出,困難 token 繼續迭代思考。
具體來說,TaH 在模型中加入一個輕量級 iteration decider(小型 MLP)。每完成一輪潛空間迭代后,decider 會基于backbone(模型骨干)的狀態 *,預測一個繼續概率。如果低于閾值,模型直接輸出下一個 token;如果高于閾值,則進入下一輪潛空間迭代。
在實際推理中,TaH 平均每個 token 只執行 1.07 次迭代,相當于跳過了約 93% token 的二次迭代計算。相比「所有 token 都想兩遍」的策略,TaH 把算力集中到了更可能出錯、更可能影響推理方向的位置。
![]()
TaH 的 duo-causal 注意力機制以及模型架構;*decider 判定繼續迭代的輸入是 backbone 的淺層、中層和最終層的隱藏狀態向量
為了讓這種動態深度策略真正提升精度和效率,TaH 在模型架構和訓練策略上都做了專門設計。
Duo-causal attention 架構。選擇性迭代會把模型處理的序列結構從一維 token 序列,變成「token 位置 × 迭代深度」的二維網格。TaH 將大模型序列維度的因果注意力(causal attention)擴展到 token 維度和迭代深度的二維平面。如圖所示,對于 token i 的第 d 次迭代,它的 query 可以注意到前序位置中深度不超過 d 的 key 和 value。
![]()
這樣既允許跨迭代深度的信息流動,也保留了訓練時序列維度計算的全并行性。
Depth-aware LoRA 架構。我們觀察到,模型的第一次迭代主要負責常規下一個 token 預測(next-token prediction,NTP),更深層迭代則是在修正當前的困難 token。TaH 因此只在 d>1 的更深迭代中啟用 LoRA adapter[12],讓 LoRA 專注于學習困難 token 的修正方向。配合跨迭代的殘差連接,深層迭代被自然地學習為「在前一輪基礎上做修正」,而不是從零再預測一次。
兩階段訓練。注意到,Decider 的判斷目標依賴 backbone 的預測質量,而 backbone 的訓練目標又依賴 decider 決定的迭代深度。因為這兩者緊耦合,所以端到端訓練并不穩定。TaH 采用解耦的兩階段方案:先用靜態 oracle 策略訓練 backbone 模型,再凍結 backbone,單獨訓練 decider 去模仿 oracle 的繼續 / 停止決策。這種訓練方法大大提升了訓練的穩定性和收斂速度。
![]()
在 Qwen3-0.6B-Base 基礎上訓練,TaH 的收斂速度更快
結果
更少迭代,更強推理
論文在 Qwen3-0.6B、1.7B 和 4B 三個規模[3]的 backbone 上驗證 TaH,訓練數據來自 Open-R1[13]中數學、問答和代碼的均衡混合,并在 GSM8K、MATH500、AMC23、AIME25、OlympiadBench、GPQA-Diamond、MMLU-STEM、HumanEval++ 和 MBPP++ 共 9 個基準上評測。所有方法 * 均使用完全相同的訓練數據,在相同的預訓練 backbone 模型上后訓練。
準確性提升:在不增加參數預算的前提下顯著超越基線。在 0.6B 和 1.7B 的嚴格參數限制下,TaH 相比標準 Qwen3 模型提升 3.0%-3.8%;TaH+ 在增加不超過 3% 額外參數(來自 decider 等)的情況下,將提升進一步擴大到 5.3%-6.2%。相比同類 Looped Transformer 方法 Ouro,TaH 取得 3.8–4.4% 的優勢,TaH+ 的優勢達到 6.1–6.8%。
![]()
TaH 在幾乎所有評測基準和模型尺寸上都實現了性能提升;*SoftThink 是推理時優化方法,無需訓練
計算效率更高。TaH 平均每個 token 只執行 1.07 次迭代,完成問答的平均 FLOPs 和顯存訪問相比標準模型只多 4–5%。在真實解碼測試中,TaH 相比始終迭代的 AlwaysThink 顯存低 1.48 倍、解碼快 2.48 倍,同時準確率反而更高。
迭代選擇的語義解釋。一個有意思的現象是,TaH 自動學到帶有明顯語義偏好的迭代行為。在驗證集上,But 和 So 是最常被觸發額外迭代的 token,概率分別達到 34% 和 18%。這些詞大多對應轉折、因果和推理方向的切換,正是復雜推理中最容易決定后續路徑的位置。
![]()
模型預測在兩次迭代之間的變化
消融實驗
為了驗證 TaH 中每一項設計的必要性,我們做了系統的消融實驗。
模型架構。把動態迭代深度的 decider 換成靜態深度的 Always-1 或 Always-2,基準測試性能平均分別下降 6.1% 和 16.4%,說明選擇性迭代本身就比統一深度更優;把 duo-causal 注意力機制替換為不同種的傳統因果注意力,測試性能下降 5.4%-8.5%,說明跨迭代深度信息流動的重要性;移除 depth-aware LoRA 與跨迭代殘差,效果下降了 4.9%,確認了在輔助不同迭代進行優化目標轉變時,架構的重要性。
訓練策略。相比 TaH 的兩階段訓練,簡單監督所有深度的預測會使基準測試性能下降 4.3%,說明不同迭代應承擔不同優化目標;把訓練時的迭代策略從靜態 oracle 換成 decider-based 或動態 oracle 會因 backbone 與 decider 的耦合而不穩定甚至崩潰,證明了 TaH 兩階段訓練的必要性
![]()
TaH 在模型架構和訓練策略上的消融實驗
總結與展望
TaH 的意義不止是提出了一個新的 Looped Transformer 變體和后訓練方法,更重要的是,它探索了如何將 test-time scaling 推向更細的 token 粒度。TaH 表明,更智能的動態算力分配甚至可以比單純使用更高算力的模型帶來更好的效果,為后續研究帶來啟發。
參考文獻
[1] Jaech, A., Kalai, A., Lerer, A., et al. OpenAI o1 system card. arXiv preprint arXiv:2412.16720, 2024.
[2] Guo, D., Yang, D., Zhang, H., et al. DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
[3] Yang, A., Li, A., Yang, B., et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.
[4] Abdin, M., Aneja, J., Awadalla, H., et al. Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404.14219, 2024.
[5] Team, M., Xiao, C., Li, Y., et al. MiniCPM4: Ultra-efficient LLMs on end devices. arXiv preprint arXiv:2506.07900, 2025.
[6] Hutchins, D., Schlag, I., Wu, Y., Dyer, E., and Neyshabur, B. Block-recurrent transformers. Advances in Neural Information Processing Systems, 35:33248–33261, 2022.
[7] Saunshi, N., Dikkala, N., Li, Z., Kumar, S., and Reddi, S. J. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint arXiv:2502.17416, 2025.
[8] Zhu, R.-J., Wang, Z., Hua, K., et al. Scaling latent reasoning via looped language models. arXiv preprint arXiv:2510.25741, 2025.
[9] Wu, Y., Wang, Y., Ye, Z., Du, T., Jegelka, S., and Wang, Y. When more is less: Understanding chain-of-thought length in LLMs. arXiv preprint arXiv:2502.07266, 2025.
[10] Wang, S., Yu, L., Gao, C., et al. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for LLM reasoning. arXiv preprint arXiv:2506.01939, 2025.
[11] Fu, T., Ge, Y., You, Y., et al. R2R: Efficiently navigating divergent reasoning paths with small-large model token routing. arXiv preprint arXiv:2505.21600, 2025.
[12] Hu, E. J., Shen, Y., Wallis, P., et al. LoRA: Low-rank adaptation of large language models. ICLR, 2022.
[13] Hugging Face. Open R1: A fully open reproduction of DeepSeek-R1, January 2025. URL https://github.com/huggingface/open-r1.
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.