隨著 AI 技術(shù)的飛速發(fā)展，從「快思考」到 「慢思考」，大語言模型（LLMs）在處理復(fù)雜推理任務(wù)上展現(xiàn)出驚人的能力。無論是我們熟知的思維鏈（CoT），還是更復(fù)雜的深度思考模式（Thinking），都讓 AI 的回答日益精準(zhǔn)、可靠。

然而，這種性能的提升并非沒有代價(jià)。模型在推理過程中會(huì)產(chǎn)生大量的中間步驟和文本（tokens），這不僅極大地拖慢了計(jì)算速度，還對(duì)內(nèi)存和計(jì)算資源造成了巨大的壓力。簡(jiǎn)單來說，就是「想得越多，算得越慢，耗得越多」。

為了解決這一難題，研究者們從人類的認(rèn)知過程中汲取靈感。想象一下人類在解決一個(gè)復(fù)雜數(shù)學(xué)題時(shí)的情景：我們通常會(huì)在草稿紙上寫下關(guān)鍵的計(jì)算步驟（如下圖 a 中的黃色高亮部分），而將一些輔助性的思考過程（非高亮部分）放在腦中。

圖 1：(a) 展示了一個(gè)典型的思維鏈推理過程，黃色部分為關(guān)鍵步驟。(b) 對(duì)比了傳統(tǒng)方案 Vanilla 與 LightThinker 的推理流程。

本文中，來自浙江大學(xué)、螞蟻集團(tuán)等機(jī)構(gòu)的研究者提出了 LightThinker，它模仿了這一高效的思考模式。它訓(xùn)練 LLM 在推理過程中動(dòng)態(tài)地將冗長(zhǎng)的中間思考步驟壓縮成緊湊的表示（gist tokens /cache tokens），然后「扔掉」原始的、繁瑣的推理鏈，僅保留核心摘要以繼續(xù)下一步的思考。 這樣一來，存放在上下文窗口中的 tokens 數(shù)量被大幅削減，從而顯著降低了內(nèi)存占用和計(jì)算成本。

• 論文標(biāo)題：LightThinker: Thinking Step-by-Step Compression
• 論文鏈接: https://arxiv.org/abs/2502.15589 
• 代碼鏈接: https://github.com/zjunlp/LightThinker

LightThinker 概覽

LightThinker 通過訓(xùn)練的方式讓模型具備這種能力。這涉及到兩個(gè)關(guān)鍵問題：「何時(shí)壓縮？」和「如何壓縮？」。整個(gè)過程可以概括為以下三個(gè)關(guān)鍵步驟：

第一步：數(shù)據(jù)重構(gòu) —— 在思考流程中植入「壓縮指令」

LightThinker 的第一步就是改造訓(xùn)練數(shù)據(jù)，讓 LLM 明白「壓縮」這個(gè)動(dòng)作的存在和時(shí)機(jī) 。具體操作是：

1. 步驟劃分：首先，將模型原本冗長(zhǎng)的完整回答 Y，按照語義或段落（即一個(gè)完整的「想法」）切分成若干個(gè)思維步驟 S1, S2, S3, ...。
2. 插入特殊指令符：在這些思維步驟之間，插入一組特殊的「指令令牌」4。這組指令符主要包含兩個(gè)部分：

• 緩存令牌 (Cache Tokens, C)：這是一組特殊的、用于存儲(chǔ)壓縮后信息的「摘要令牌」。它的作用就像是為即將產(chǎn)生的「思想摘要」預(yù)留的空白便簽。
• 輸出令牌 (Output Token, [o])：這是一個(gè)強(qiáng)制性的輸出信號(hào)，它的作用是告訴模型：「好了，摘要寫完了，現(xiàn)在請(qǐng)基于這份摘要繼續(xù)你下一步的思考」 。

經(jīng)過這樣的改造，原本一條完整的思考鏈，就變成了一個(gè)「思考步驟 1 S1 → 進(jìn)行壓縮 → 繼續(xù)思考步驟 S2 → 再次壓縮 → ...」的全新格式。這等于是在模型的學(xué)習(xí)材料中明確地標(biāo)注出了「何時(shí)」需要進(jìn)行壓縮。

注意，研究者在具體實(shí)現(xiàn)中，采用換行符作為思維步驟的劃分，此處不存在任何數(shù)據(jù)精心構(gòu)造的過程。

第二步：注意力改造 —— 學(xué)會(huì)壓縮與理解壓縮的內(nèi)容

教會(huì)了模型「何時(shí)」壓縮，下一步就是最關(guān)鍵的如何壓縮。這主要通過一種名為 「Thought-based Attention Mask」的技術(shù)來實(shí)現(xiàn)，如圖 2 (b) 所示。精確地控制著模型在思考時(shí) “能看什么” 和 “不能看什么” 。

這個(gè)過程分為兩個(gè)階段：

• 壓縮階段（生成思維摘要）：當(dāng)模型需要將思維步驟 Si 壓縮進(jìn)緩存令牌 C 時(shí)，注意力掩碼會(huì)強(qiáng)制這些 C 令牌只能「看到」三個(gè)東西：
• 最初的問題 X；
• 先前已經(jīng)壓縮好的歷史摘要；
• 當(dāng)前正在處理的思維步驟 Si。

其他所有原始的、未壓縮的思維步驟都會(huì)被「遮蔽」。這迫使模型必須將 Si 中的所有關(guān)鍵信息高度濃縮并存儲(chǔ)到 C 中 。

• 生成階段（基于摘要生成思維）：當(dāng)思維步驟 Si 被成功壓縮進(jìn) C 之后，更關(guān)鍵的一步來了。在生成下一個(gè)思緒片段 S (i+1) 時(shí)，注意力掩碼會(huì)徹底「遮蔽」掉原始的思維步驟 Si。此時(shí)，模型只能「看到」最初的問題 X 和包括剛剛生成的摘要在內(nèi)的所有歷史摘要 。

通過這種方式，模型被迫學(xué)會(huì)僅依賴緊湊的「思想摘要」來進(jìn)行連貫的、層層遞進(jìn)的推理，而不是依賴越來越長(zhǎng)的原始思考全文。

第三步：動(dòng)態(tài)推理 ——「即用即棄」的高效循環(huán)

經(jīng)過以上兩個(gè)步驟的訓(xùn)練，LightThinker 模型在實(shí)際推理時(shí)，就會(huì)形成一種高效的動(dòng)態(tài)循環(huán)，如圖 1 (b) 和圖 2 (c) 所示，清晰地展示了「生成→壓縮→拋棄」的動(dòng)態(tài)循環(huán)過程。下面以圖 1 (b) 為例進(jìn)行分析：

1. 模型接收問題，生成第一段思考（Thought 1）。
2. 觸發(fā)壓縮，將 Thought 1 中的核心信息壓縮成緊湊的摘要（CT1）。
3. 拋棄原文，將冗長(zhǎng)的 Thought 1 從上下文中丟棄。
4. 模型基于問題和摘要（CT1），生成第二段思考（Thought 2）。
5. 再次壓縮，將 Thought 2 壓縮為摘要（CT2），并丟棄 Thought 2 原文。
6. 如此循環(huán)，直到問題解決。

通過這種「即用即棄」的機(jī)制，LightThinker 確保了模型的上下文窗口始終保持在一個(gè)非常小的尺寸，從而解決了因上下文過長(zhǎng)導(dǎo)致的內(nèi)存爆炸和計(jì)算緩慢問題，實(shí)現(xiàn)了效率與性能的完美平衡。

圖 3 展示了不同方法在推理過程中上下文長(zhǎng)度的變化，其中曲線和坐標(biāo)軸圍城的面積為我們定義的新指標(biāo) Dependency，其意義生成 token 時(shí)需要關(guān)注 token 的數(shù)量總和。

實(shí)驗(yàn)結(jié)果

研究者在四個(gè)數(shù)據(jù)集和兩個(gè)不同的模型上對(duì) LightThinker 進(jìn)行了廣泛的測(cè)試，結(jié)果如表 1 所示。

表 1 主要實(shí)驗(yàn)結(jié)果。Acc 為準(zhǔn)確率，Time 為平均推理耗時(shí)，Peak 為平均峰值 token 占用數(shù)量，Dep 為生成 token 時(shí)需要關(guān)注 token 的數(shù)量總和（如圖 3）所示。

結(jié)果表明，在 Qwen 系列模型上，與傳統(tǒng)模型（Vanilla）相比：

• 峰值內(nèi)存使用減少 70%：LightThinker 極大地節(jié)約了寶貴的內(nèi)存資源。
• 推理時(shí)間縮短 26%：在保證結(jié)果準(zhǔn)確性的前提下，思考速度得到了顯著提升。 
• 取得了準(zhǔn)確度和效率的平衡。

此外，在 Llama 上，也取得了準(zhǔn)確度和效率的平衡。

相關(guān)工作

當(dāng)前關(guān)于加速大語言模型（LLMs）推理過程的研究主要集中在四類方法：模型量化、輔助解碼、生成更少的 Token 和減少 KV 緩存。模型量化包括參數(shù)量化 [1-2] 和 KV 緩存量化 [3-4]，輔助解碼主要包括投機(jī)采樣，本節(jié)將重點(diǎn)關(guān)注后兩類方法。

需要注意的是，生成長(zhǎng)文本和理解長(zhǎng)文本代表著不同的應(yīng)用場(chǎng)景，因此，專門針對(duì)長(zhǎng)文本生成階段的加速方法（例如，預(yù)填充階段加速技術(shù)如 AutoCompressor [5]、ICAE [6]、LLMLingua [7]、Activation Beacon [8]、SnapKV [9] 和 PyramidKV [10]）不在此處討論。以下是后兩類方法的詳細(xì)概述。

生成更少的 Token

這一類別可以根據(jù)推理過程中使用的 token 數(shù)量和類型進(jìn)一步分為三種策略：

• 離散 Token 減少通過提示工程 Prompt [11-13]、指令微調(diào) [14-15] 或強(qiáng)化學(xué)習(xí) [16-17] 等技術(shù)來引導(dǎo) LLM 在推理過程中使用更少的離散 token。例如，TALE [11] 提示 LLM 在預(yù)定義的 token 預(yù)算內(nèi)完成任務(wù)。Arora 和 Zanette [16] 構(gòu)建特定數(shù)據(jù)集并采用強(qiáng)化學(xué)習(xí)獎(jiǎng)勵(lì)機(jī)制來鼓勵(lì)模型生成簡(jiǎn)潔準(zhǔn)確的輸出，從而減少 token 使用量。
• 連續(xù) Token 替換這些方法 [18-19] 探索使用連續(xù)空間 token 代替?zhèn)鹘y(tǒng)的離散詞匯 token。一個(gè)代表性例子是 CoConut [18]，它利用課程學(xué)習(xí)來訓(xùn)練 LLM 使用連續(xù) token 進(jìn)行推理。
• 無 Token 使用通過在模型層之間內(nèi)化推理過程，在推理過程中直接生成最終答案而不需要中間 token [20-21]。

這三種策略都是在模型訓(xùn)練后實(shí)施的，推理過程中不需要額外干預(yù)。從技術(shù)上講，這些方法的加速效果依次遞增，但代價(jià)是 LLM 的泛化性能逐漸下降。此外，第一種策略并不能顯著減少 GPU 內(nèi)存使用。

減少 KV 緩存

這一類別可以分為兩種策略類型：基于剪枝的離散空間 KV 緩存選擇和基于合并的連續(xù)空間 KV 緩存壓縮。

• 基于剪枝的策略設(shè)計(jì)特定的淘汰策略 [22-25] 在推理過程中保留重要的 token。例如，StreamingLLM [23] 認(rèn)為初始的 sink token 和最近的 token 是重要的；H2O [22] 關(guān)注具有高歷史注意力分?jǐn)?shù)的 token；SepLLM [24] 強(qiáng)調(diào)對(duì)應(yīng)于標(biāo)點(diǎn)符號(hào)的 token 是重要的。
• 基于合并的策略引入錨點(diǎn) token，訓(xùn)練 LLM 將歷史重要信息壓縮到這些 token 中，從而實(shí)現(xiàn) KV 緩存合并 [26]。

這兩種策略都需要在推理過程中進(jìn)行干預(yù)。關(guān)鍵區(qū)別在于：第一種策略是無需訓(xùn)練的，但對(duì)每個(gè)生成的 token 都要應(yīng)用淘汰策略；而第二種策略是基于訓(xùn)練的方法，允許 LLM 自主決定何時(shí)應(yīng)用淘汰策略。

局限性

受限于自身的數(shù)據(jù)重構(gòu)方案（目前分割思維步驟是依賴規(guī)則，而不是基于語義）和訓(xùn)練數(shù)據(jù)（約 16K 訓(xùn)練數(shù)據(jù)），本文方法在數(shù)學(xué)相關(guān)的任務(wù)上表現(xiàn)并不出色。

如下圖所示，展示了 LightThinker 在 GSM8K 上的一個(gè) Bad Case。研究者觀察到，盡管 LLM 在思考過程中得出了正確答案（見上圖中的 Model's Thoughts 字段），但在最終輸出中卻出現(xiàn)了錯(cuò)誤（見圖中的 Model's Solution 字段）。

具體來說，在 Model's Solution 字段的第三句話中，第一次出現(xiàn)的「4000」是錯(cuò)誤的。這表明在第二次壓縮步驟中發(fā)生了信息丟失（理論上，「8000」、「4000」和「24000」都應(yīng)該被壓縮，但 LLM 只壓縮了「4000」和「24000」），導(dǎo)致后續(xù)的推理錯(cuò)誤。這類錯(cuò)誤在 GSM8K 數(shù)據(jù)集中頻繁出現(xiàn)，表明當(dāng)前的壓縮方法對(duì)數(shù)值的敏感度還不夠。