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