ICML2025|大模型深度思考新正規化

2025-07-30 02:39 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：機器之心
作者介紹：本文第一作者是豐田工業大學芝加哥 PhD 學生楊晨曉，研究興趣是機器學習理論和大模型推理，在 ICML，NeurIPS，ICLR 等頂級會議上發表過論文。
本文提出一個交替「推理 – 擦除」的深度思考新正規化 PENCIL，比傳統 CoT 更高效地解決更復雜的推理任務。理論上，我們證明 PENCIL 可用最優空間與最優時間下解決所有可計算問題，而這對於傳統的 CoT 是不可能的！該工作已被機器學習頂會 ICML 2025 收錄。
  • 題目： PENCIL: Long Thoughts with Short Memory
  • 連結： https://arxiv.org/pdf/2503.14337
  • 程式碼： https://github.com/chr26195/PENCIL 
最近的大模型（如 OpenAI 的 o1/o3、DeepSeek 的 R1）發現能透過在測試階段深度思考（Test-Time Scaling）來大幅提高模型的推理能力。目前實現深度思考的關鍵在於使用長鏈思維鏈（Long Chain-of-Thought，CoT），即讓模型生成更長中間結果得到最終答案。然而，傳統「只寫不擦」的方法在處理高難度、大規模任務時面臨以下瓶頸： 
  1. 超出上下文視窗：一旦鏈條過長，就會觸及模型的最大上下文長度限制；
  2. 資訊檢索困難：隨著上下文不斷累積，模型難以從冗長曆史中 Retrieve 關鍵線索； 
  3. 生成效率下降：上下文越長，每步生成新 token 的計算量越大。
不過實際上，並非所有中間思路都後續推理有用：例如定理證明裡，引理一旦驗證透過，其具體推導可被丟棄；解數學題時，已知某條思路走不通就無需保留那段「嘗試」的細節。縱觀計算機科學的發展歷史，這一「隨時清理」的理念早已滲透到幾乎所有計算模型之中：從最早的圖靈機模型中，已讀寫的磁帶符號可以被覆蓋或重寫，直到現在高階程式語言中，垃圾回收機制會自動清理不再可達的記憶體單元。
基於這樣的動機，我們提出一個新的深度思考正規化 PENCIL，迭代地執行生成（Generation）和擦除（Reduction），即在生成的過程中動態地擦除不再需要的中間結果，直到得到最後的答案。
一、交替「生成 – 擦除」的深度思考正規化
下圖以一個簡單的算術題為例展示了 PENCIL 的工作機制： 
  • CoT 將每步推理串聯到上下文中直到給出答案並返回整個序列。
  • PENCIL 交替執行生成（圖中加粗部分）和 擦除（圖中綠色高亮部分）：模型先寫出新的思考過程，再刪掉對之後的推理無用片段，只保留對後續的推理過程有用的部分，內部形成一系列隱式思維，最後僅返回最終答案。
PENCIL 擦除機制的設計借鑑了邏輯學與經典自動定理證明中的重寫規則（Rewriting Rule 和函數語言程式設計語言中的棧幀記憶體管理（Stack Frame）。 具體地，我們引入三個特殊字元（Special Token），叫做 [CALL], [SEP], [RETURN]，並用以下的規則（Reduction Rule）來實現擦除： 
其中 C（Context）表示上下文，T（Thoughts）表示中間思考，A（Answer）表示回答。每當生成的序列與左側模式完全匹配時，PENCIL 即觸發一次擦除，丟棄 T。重要的是，C、T、A 本身均可包含其他特殊標記，從而支援類似多層函式呼叫的遞迴結構。
PENCIL 的擦除機制能夠靈活支撐多種推理模式，例如：
  • 任務分解（Decomposition）：透過 [CALL] 啟動子任務，完成後用 [RETURN] 合併輸出並擦除子任務推理細節；
  • 搜尋與回溯（Search and Backtrack）：在搜尋樹中，用特殊字元管理探索分支，衝突或失敗時擦除無效路徑； 
  • 摘要與總結（Summarization）：將冗長的思考片段歸納為簡潔摘要，類似程式設計中的尾遞迴（Tail Recursion）：
其中 T 表示原始的複雜思考過程（或更難的問題），T' 歸納或簡化後的摘要（或等價的、更易處理的問題）。
示例： 布林可滿足性（SAT）是經典的 NP-Complete 問題：給定一個 n 個變數布林公式，判斷是否存在一組變數賦值使其為真。這個問題（廣泛認為）需要指數時間但僅需多項式空間來解決，其中最簡單的做法是構造一個深度為 n 的二叉搜尋樹遍歷所有可能。傳統 CoT 將每步計算附加到上下文，長度與搜尋樹節點數成正比 (O (exp (n)))，導致指數爆炸；PENCIL 在遞迴分支嘗試時，遇到衝突立即回溯並擦除該分支所有思考，僅保留關鍵結果，使上下文長度僅與搜尋深度成正比 (O (n))。
如圖所示，對比 CoT 無擦除（藍）與 PENCIL 擦除（紅）兩種思考模式下的最大上下文長度，隨著問題規模增大，PENCIL 能將所需序列長度控制在千級或百級，而傳統 CoT 則迅速攀升至數萬甚至數十萬。即使在複雜的 Einstein's Puzzle 中，PENCIL 也能將需要幾十萬 token 的上下文壓縮到幾千 token。
二、訓練和實驗結果
訓練和測試：在訓練時，CoT 每個新 token 的損失計算都基於完整的歷史上下文；PENCIL 在每輪「寫 — 擦」迴圈結束後只在被擦除後的短序列上計算損失。即使兩者生成 token 數量相同，PENCIL 每一個 token 對應的上下文長度卻大幅縮短；另一方面，在每次 Reduction 後，C 部分的 KV cache 可以直接複用，只需為更短的 A 部分重新計算快取。這樣， PENCIL 在訓練和測試時能顯著減少自注意力計算開銷。 
實驗設定：我們針對三種具有代表性的高難度推理任務構建資料集：3-SAT（NP-Complete）、QBF（PSPACE-Complete）和 Einstein’s Puzzle（自然語言推理）。所有實驗均在相同配置下從隨機初始化開始進行預訓練和評估，採用小型 Transformer（10.6M 引數和 25.2M 引數），訓練超引數保持一致。 
1. 準確率
相比 CoT，PENCIL 能解決更大規模的推理問題。如下圖所示，在 SAT（左圖）和 QBF（右圖）任務中，當問題規模較小時，CoT 與 PENCIL 均能完美解決問題；但隨著規模增大，傳統 CoT 的準確率顯著下降（例如 SAT 在 n=10 時僅約 50%），而 PENCIL 始終保持 ≥ 99% 的高準確率。 
2. 計算效率
PENCIL 還能顯著節省計算資源。如圖所示，我們在相同 FLOPs 預算下對比了 CoT（藍色）與 PENCIL（紅色）的訓練收斂表現。PENCIL 訓練早期迅速達到 100% 準確率，訓練損失更快穩定；CoT 因上下文膨脹需投入更多資源才能接近最優。隨著問題規模增加，兩者之間的差距愈發明顯。 
3. 自然語言推理任務：Einstein’s Puzzle
我們測試了 PENCIL 在極具挑戰性的 Einstein's Puzzle 上的表現。該問題要求從一系列線索（如「綠房子在養鳥者右側」、「養狗者住在紅房子」等）推斷出五個房屋中人們的全部屬性（顏色、國籍、飲品、香菸和寵物）。即使是 GPT-4 也難以解決此類邏輯推理問題 [1]。下圖展示了 n=3 時的問題簡化：
如圖所示，對於該大模型也難以解決的問題，而 PENCIL 僅用一個 25.2M 引數的小模型將準確率提升至 97%；相比較之下，傳統 CoT 準確率僅 25%，接近隨機猜測的準確率。
三、理論：PENCIL 用最優的空間 / 時間實現圖靈完備
我們進一步從理論表達能力的角度展示 PENCIL 相較於傳統 CoT 的根本性優勢。具體地，我們證明：使用一個固定的、有限大小的 Transformer，PENCIL 可以用最優的時間和空間複雜度模擬任意圖靈機的運算過程（即實現圖靈完備），從而高效地解決所有可計算問題： 
具體而言，若任意圖靈機在某輸入上需 T 步計算和 S 空間，PENCIL 僅需生成 O (T) 個 token 並保持上下文長度至多為 O (S) 即可輸出相同結果。值得注意的是，大多數演算法的空間複雜度都遠小於其時間複雜度，即 S << T。
相比之下，傳統 CoT 雖能實現圖靈完備 [2] —— 思維鏈的每一步表示圖靈機的一步中間計算過程，因此思維鏈足夠長就可以解決所以可計算問題。但這意味著其生成序列的上下文長度必須與執行步數 T 成正比，代價十分昂貴：對於中等難度任務也許尚可承受，一旦面對真正複雜需要深度思考的問題，這種指數級的上下文爆炸就變得不切實際。
例如，一系列（公認）無法在多項式時間內解決卻可在多項式空間內解決的 NP-Complete（如旅行商等等），對於使用有限精度 Transformer 的 CoT 而言至少需要超越多項式（例如 exp (n)）規模的上下文長度，在真實應用中由於記憶體的限制完全不可行；而 PENCIL 只需 poly (n) 規模的上下文就能高效求解，讓「深度思考」變得切實可行。 
證明思路：證明關鍵在用一系列「思考 — 總結」迴圈來替代持續累積的思維鏈。
具體地，如上圖左圖所示，我們先將圖靈機狀態轉移編碼為三元組 token（新狀態、寫入符號、移動方向）。模型透過自注意力計算讀寫頭位置，並從上下文回溯讀取符號。未經最佳化時，需保留 T 步完整歷史，上下文長度為 O (T)。
PENCIL 能夠實現空間 / 時間最優的核心是利用交替「思考 – 總結」的生成方式：
  1. 思考 （Simulation）：生成連續狀態轉移 token，模擬圖靈機計算； 
  2. 總結 （Summarization）：當新 token 數超過實際所需空間兩倍時，用不超過 S 個的 token 總結當前狀態，觸發擦除規則丟棄中間過程。
透過這種策略，PENCIL 生成總 token 數仍為 O (T)，卻把最大上下文長度嚴格限制在 O (S)，達到了空間與時間的雙重最優。
最後，我們需要證明這種「思考 – 總結」的生成方式可以被現實中的 Transformer 實現。為此，我們設計了 Full-Access Sequence Processing (FASP) 程式語言，並證明所有用 FASP 寫的程式都可被 Transformer 表達。透過構造能執行「思考 – 總結」操作的 FASP 程式，我們證明了等價存在固定大小 Transformer 完成相同功能，從而理論上證明 PENCIL 可用最優複雜度模擬任意計算過程。
參考文獻
[1] Dziri, Nouha, et al. "Faith and fate: Limits of transformers on compositionality." in NeurIPS 2023.
[2] Merrill, William, and Ashish Sabharwal. "The expressive power of transformers with chain of thought." in ICLR 2024.

何愷明在MIT授課的課件PPT下載
在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）人數破萬！如果你想要了解最新最快最好的CV/DL/AI論文、實戰專案、行業前沿、從入門到精通學習教程等資料，一定要掃描下方二維碼，加入CVer知識星球！最強助力你的科研和工作！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看


相關文章