Alexia Jolicoeur-Martineau
Less is More: Recursive Reasoning with Tiny Networks
少即是多:微模型的遞迴推理
https://arxiv.org/pdf/2510.04871
摘要
層次推理模型(HRM)是一種創新方法,它使用兩個以不同頻率遞迴的小型神經網路。這種受生物學啟發的方法,在諸如數獨、迷宮和ARC-AGI等難題任務上,表現優於大語言模型(LLMs),同時僅使用小型模型(2700萬參數)和少量資料(約1000個示例)進行訓練。HRM在利用小型網路解決難題方面前景廣闊,但其原理尚未被充分理解,且可能並非最優方案。我們提出了微型遞迴模型(TRM),這是一種更為簡潔的遞迴推理方法。在僅使用一個僅含2層的小型網路的情況下,其泛化能力顯著超越HRM。TRM僅擁有700萬參數,就在ARC-AGI-1上達到了45%的測試準確率,在ARC-AGI-2上達到8%,超過了大多數大語言模型(例如Deepseek R1、o3-mini、Gemini 2.5 Pro),而其參數量卻不足這些模型的0.01%。
1 引言
儘管大型語言模型(LLMs)能力強大,但在解決困難的問答問題時仍可能面臨挑戰。由於它們是自回歸地生成答案,所以存在較高的出錯風險,因為單個錯誤的標記就可能導致答案無效。為了提高其可靠性,LLMs依賴於思維鏈(CoT)(Wei et al., 2022)和測試時計算(TTC)(Snell et al., 2024)。思維鏈旨在模仿人類推理,讓LLM在給出答案前採樣出逐步推理的軌跡。這樣做可以提高準確性,但思維鏈成本高昂,需要高品質的推理資料(可能無法獲得),並且可能很脆弱,因為生成的推理可能是錯誤的。為了進一步提高可靠性,可以使用測試時計算,即從K個答案中報告最常見的答案或最高獎勵的答案(Snell et al., 2024)。
在這項工作中,我們表明,遞迴推理帶來的益處可以得到極大提升,其改進遠不止是漸進式的。我們提出了微型遞迴模型(TRM),這是一種改進且簡化的方法,它使用一個僅包含2層、規模小得多的微型網路,在多種問題上實現了比HRM顯著更高的泛化能力。透過這一方法,我們將數獨極限測試集的準確率從55%提升至87%,迷宮困難測試集從75%提升至85%,ARC-AGI-1從40%提升至45%,ARC-AGI-2從5%提升至8%。
2. 背景
HRM的演算法描述詳見演算法2。我們將在下文進一步討論該演算法的細節。
2.1. 結構與目標
HRM的研究重點是監督學習。給定一個輸入,生成一個輸出。假定輸入和輸出都具有形狀 [B, L](當形狀不同時,可以添加填充標記),其中 B 是批量大小,L 是上下文長度。
在這項工作中,我們表明,遞迴推理帶來的益處可以得到極大提升,其改進遠不止是漸進式的。我們提出了微型遞迴模型(TRM),這是一種改進且簡化的方法,它使用一個僅包含2層、規模小得多的微型網路,在多種問題上實現了比HRM顯著更高的泛化能力。透過這一方法,我們將數獨極限測試集的準確率從55%提升至87%,迷宮困難測試集從75%提升至85%,ARC-AGI-1從40%提升至45%,ARC-AGI-2從5%提升至8%。
2.2. 兩種不同頻率的遞迴
2.3. 使用一步梯度近似的定點遞迴
2.7 HRM總結
HRM利用兩個不同頻率(高頻與低頻)網路的遞迴和深度監督,學習在多個監督步驟中改進其答案(並使用ACT減少每個資料樣本的處理時間)。這使得模型能夠模仿極深的網路,而無需對所有層進行反向傳播。該方法在常規監督模型難以應對的困難問答任務上取得了顯著更高的性能。然而,該方法相當複雜,過於依賴不確定的生物學論據和無法保證適用的不動點定理。在下一節中,我們將討論這些問題以及HRM潛在的改進目標。
3. 層次推理模型的改進目標
在本節中,我們確定了HRM的關鍵改進目標,這些目標將由我們提出的方法——微型遞迴模型(TRM)來解決。
3.1. 隱函數定理與一步梯度近似
HRM僅對6次遞迴中的最後2次進行反向傳播。作者透過應用隱函數定理和一步近似來證明這種做法的合理性,該定理指出:當遞迴函數收斂到不動點時,可以在該平衡點處透過單步進行反向傳播。
因此,儘管將隱函數定理和一步梯度近似應用於HRM有一定依據,因為殘差通常會隨時間推移而減小,但在實際應用該定理時,模型很可能並未達到不動點。
在下一節中,我們將展示可以繞過對隱函數定理和一步梯度近似的需求,從而徹底避免這個問題。
3.2. 自適應計算時間(ACT)導致前向傳播次數加倍
HRM在訓練期間使用自適應計算時間(ACT)來優化每個資料樣本所花費的時間。如果不使用ACT,每個資料樣本將需要花費
然而,ACT是有代價的。這一代價沒有直接在HRM的論文中顯示,但在其官方程式碼中有所體現。Q學習目標依賴於一個停止損失和一個繼續損失。繼續損失需要對HRM進行額外的一次前向傳播(包含全部6次函數評估)。這意味著,雖然ACT能更有效地優化每個樣本的時間,但每個優化步驟需要2次前向傳播。具體公式如演算法2所示。
在下一節中,我們將展示如何避免ACT中兩次前向傳播的需求。
3.3. 基於複雜生物學論據的層次化解釋
HRM的作者基於生物學論據來證明兩個潛在變數和兩個在不同層次運行的網路的設計是合理的,但這些論據與人工神經網路相距甚遠。他們甚至嘗試將HRM與對小鼠的實際大腦實驗聯繫起來。儘管這很有趣,但這種解釋使得理解HRM為何如此設計變得極其困難。考慮到其論文中缺乏消融實驗表,以及對生物學論據和不動點定理(其並非完全適用)的過度依賴,很難確定HRM的哪些部分起到了什麼作用以及原因。此外,不清楚他們為何使用兩個潛在特徵,而不是其他特徵組合。
在下一節中,我們將展示遞迴過程可以得到極大的簡化,並且可以用一種簡單得多的方式來理解,這種方式不需要任何生物學論據、不動點定理、層次化解釋,也不需要兩個網路。這也解釋了為什麼2是最佳特徵數(即)。
4. 微型遞迴模型
在本節中,我們介紹微型遞迴模型(TRM)。與HRM相反,TRM不需要複雜的數學定理、層級結構或生物學論據。它在泛化能力上表現更優,同時僅需一個微型網路(而非兩個中等規模網路),並且ACT(自適應計算時間)僅需單次前向傳播(而非2次)。我們的方法在演算法3中描述,並在圖1中展示。我們還在數獨極限資料集(一個僅有1K訓練樣本但包含423K測試樣本的困難數獨資料集)上進行了消融實驗,結果如表1所示。下文將闡述TRM的關鍵組成部分。
4.1. 無需不動點定理
雖然這種解釋很直觀,但我們仍想驗證使用更多或更少的特徵是否有所幫助。結果如表2所示。
單一特徵
:類似地,我們測試了僅使用一個特徵的想法,即在監督步驟之間只傳遞。該方法在演算法4中描述。透過這種方式,我們發現性能有所下降。這也是預期的結果,因為正如前面所討論的,這迫使模型將解
因此,我們在數獨極限任務上探索了使用更多或更少潛在變數的情況,但發現僅使用
4.4. 少即是多
我們曾嘗試透過增加層數來提高容量,以擴展模型。令人驚訝的是,我們發現增加層數會因過擬合而降低泛化能力。相反,在按比例增加遞迴次數(
較小的網路表現更好,這相當令人驚訝,但 2 層似乎是最優選擇。Bai & Melas-Kyriazi (2024)在深度平衡擴散模型的背景下也觀察到了 2 層網路的最佳性能;然而,他們的性能與更大網路相似,而我們則觀察到 2 層網路性能更好。這可能看起來有些不同尋常,因為對於現代神經網路,泛化能力往往直接與模型規模相關。然而,當資料過於稀缺且模型規模過大時,可能會出現過擬合的懲罰(Kaplan 等人, 2020)。這可能表明資料量過少。因此,使用具有深度遞迴和深度監督的微小網路似乎能讓我們避免很多過擬合問題。
4.5. 針對固定小上下文長度任務的無注意力架構
4.7. 指數移動平均
在小資料集(如 Sudoku-Extreme 和 Maze-Hard)上,HRM 往往會迅速過擬合並隨後發散。為了減少這個問題並提高穩定性,我們採用了權重指數移動平均(EMA),這是 GAN 和擴散模型中用於提高穩定性的常見技術(Brock 等人, 2018; Song & Ermon, 2020)。我們發現它能防止急劇崩潰並帶來更高的泛化能力(從 79.9% 到 87.4%;見表 1)。
4.8. 優化遞迴次數
在下一節中,我們將展示 HRM、TRM 和 LLMs 在多個資料集上的主要結果對比。
5. 結果
遵循 Wang 等人 (2025) 的方法,我們在以下資料集上測試了我們的方法:Sudoku-Extreme (Wang 等人,2025)、Maze-Hard (Wang 等人,2025)、ARC-AGI1 (Chollet,2019) 和 ARC-AGI-2 (Chollet 等人,2025)。結果展示在表 4 和表 5 中。超參數詳見第 6 節。資料集討論如下。
Sudoku-Extreme
包含極其困難的數獨謎題(Dillion, 2025; Palm 等人, 2018; Park, 2018)(9x9 網格),僅使用 1K 訓練樣本來測試小樣本學習能力。在 423K 個樣本上進行測試。
Maze-Hard
包含由 Lehnert 等人 (2024) 提出的程式生成的 30x30 迷宮,其最短路徑長度超過 110;訓練集和測試集各包含 1000 個迷宮。
ARC-AGI-1
和 ARC-AGI-2 是包含獎金的幾何謎題。每個謎題都設計得對人類來說容易,但對當前的人工智慧模型來說卻很困難。每個謎題任務包含 2-3 個輸入-輸出演示對和 1-2 個待求解的測試輸入。最終分數計算為在所有測試輸入上經過兩次嘗試生成正確輸出網格的準確率。最大網格尺寸為 30x30。ARC-AGI-1 包含 800 個任務,而 ARC-AGI-2 包含 1120 個任務。我們還使用來自緊密相關的 ConceptARC 資料集(Moskvichev 等人,2023)的 160 個任務來增強我們的資料。我們提供了 ARC-AGI-1 和 ARC-AGI-2 在公共評估集上的結果。
雖然這些資料集規模較小,但為了改善泛化能力,使用了大量的資料增強。Sudoku-Extreme 對每個資料樣本使用了 1000 次洗牌增強(在不違反數獨規則的情況下進行)。Maze-Hard 對每個資料樣本使用了 8 種二面體變換。ARC-AGI 對每個資料樣本使用了 1000 次資料增強(顏色置換、二面體群變換和平移變換)。二面體群變換包括隨機的 90 度旋轉、水平/垂直翻轉和反射。
從結果中可以看出,不使用自注意力的 TRM 在 Sudoku-Extreme 上獲得了最佳的泛化能力(87.4% 測試準確率)。同時,使用自注意力的 TRM 在其他任務上泛化得更好(可能是由於歸納偏差以及 MLP 在大的 30x30 網格上的過擬合傾向)。使用自注意力的 TRM 在 Maze-Hard 上獲得 85.3% 的準確率,在 ARC-AGI-1 上獲得 44.6% 的準確率,在 ARC-AGI-2 上獲得 7.8% 的準確率,參數量為 7M。這顯著高於使用 4 倍參數量(27M)的 HRM 所獲得的 74.5%、40.3% 和 5.0%。
6. 結論
我們提出了微型遞迴模型(TRM),這是一種簡單的遞迴推理方法,透過在潛在的推理特徵上進行遞迴並使用單個微型網路逐步改進最終答案,從而在困難任務上實現了強大的泛化能力。與分層推理模型(HRM)相比,TRM 不需要不動點定理、複雜的生物學論證,也不需要層次結構。它透過將層數減半並用單個微型網路替代兩個網路,顯著減少了參數量。它還簡化了停止過程,無需額外的前向傳播。總的來說,TRM 比 HRM 簡單得多,同時實現了更好的泛化能力。
儘管我們的方法在 4 個基準測試中帶來了更好的泛化性能,但我們所做的每一個選擇都不能保證在所有資料集上都是最優的。例如,我們發現用 MLP 替換自注意力在 Sudoku-Extreme 上效果極好(測試準確率提升了 10%),但在其他資料集上表現不佳。不同的問題設置可能需要不同的架構或參數量。需要擴展定律來優化這些網路的參數化。雖然我們簡化和改進了深度遞迴,但為什麼遞迴比使用更大更深的網路幫助更大,這個問題仍有待解釋;我們懷疑這與過擬合有關,但我們沒有理論來支持這一解釋。並非我們所有的想法都取得了成功;我們將在第 6 節簡要討論一些我們嘗試過但未成功的想法。
目前,遞迴推理模型(如 HRM 和 TRM)是監督學習方法,而不是生成模型。這意味著給定一個輸入問題,它們只能提供一個確定的答案。然而,在許多情況下,一個問題可能存在多個答案。因此,將 TRM 擴展到生成任務將是很有意義的研究方向。