LightThinker：中間思維的動態壓縮以提高大型語言模型推理的效率

提升大型語言模型的推理效率

像是思考鏈（Chain-of-Thought, CoT）提示的方法，可以幫助我們把複雜的問題分解成一個個小步驟，這樣更容易理解。最近的一些進展，例如o1類似的思考模式，則引入了試錯、回溯、修正和重複等能力，來改善模型在困難問題上的表現。不過，這些改進需要大量的計算資源。因為生成的標記（tokens）增加，會造成記憶體的負擔，這是由於變壓器架構（Transformer architecture）的限制，當上下文長度增加時，注意力機制的複雜度會以平方的方式增長，而KV快取的存儲則是線性增長。例如，當Qwen32B的上下文長度達到10,000個標記時，KV快取所需的記憶體與整個模型相當。

加速大型語言模型推理的方法

目前加速大型語言模型（LLM）推理的方法主要有三個類別：模型量化、生成較少的標記和減少KV快取。模型量化包括參數和KV快取的量化技術。在減少KV快取的類別中，基於修剪的選擇和基於合併的壓縮成為主要策略。基於修剪的策略會實施特定的驅逐政策，以在推理過程中只保留重要的標記。基於合併的策略則會引入錨標記，來壓縮歷史上重要的信息。這兩種方法的不同在於，基於修剪的方法不需要訓練，但需要對每個生成的標記應用驅逐政策，而基於合併的方法則需要模型訓練。

LightThinker的介紹

來自浙江大學、螞蟻集團和浙江大學－螞蟻集團知識圖譜聯合實驗室的研究人員提出了LightThinker，旨在讓大型語言模型在推理過程中動態壓縮中間思考。LightThinker受到人類認知的啟發，將冗長的推理步驟壓縮成簡潔的表示，並丟棄原始的推理鏈，顯著減少了存儲在上下文窗口中的標記數量。研究人員還引入了依賴性（Dependency, Dep）指標，來量化壓縮的有效性，通過測量生成過程中對歷史標記的依賴程度來進行評估。此外，LightThinker在保持競爭性準確度的同時，減少了峰值記憶體使用量和推理時間，為提高大型語言模型在複雜推理任務中的效率提供了一個有前景的方向。

LightThinker的評估

LightThinker方法使用Qwen2.5-7B和Llama3.1-8B模型進行評估。研究人員使用Bespoke-Stratos-17k數據集進行了完整的參數指令調整，生成的模型被稱為Vanilla。實施了五個比較基準：兩種無需訓練的加速方法（H2O和SepLLM）、一種基於訓練的方法（AnLLM），以及對指令和R1-Distill模型應用的CoT提示。評估在四個數據集（GSM8K、MMLU、GPQA和BBH）上進行，測量了效果和效率（通過推理時間、峰值標記數和依賴性指標）。實施了兩種壓縮方法：標記級壓縮（將每6個標記轉換為2個）和思考級壓縮（使用“\n\n”作為分隔符來劃分思考）。

評估結果

在四個數據集上，對這兩個模型的四個指標進行的評估結果顯示了幾個重要的發現。Distill-R1在所有數據集上的表現始終低於CoT，這一表現差距歸因於Greedy Decoding造成的重複問題。H2O有效地保留了模型性能，同時減少了記憶體使用，驗證了其對長文本生成的貪婪驅逐政策。然而，由於其逐標記的驅逐政策為每個生成的標記創造了額外的負擔，H2O顯著增加了推理時間（Qwen增加51%，Llama增加72%）。此外，LightThinker在相似的壓縮率下，與H2O的性能相當，同時減少了推理時間，Qwen減少了52%，Llama減少了41%。

結論

在這篇論文中，研究人員介紹了LightThinker，這是一種通過在生成過程中動態壓縮中間思考來提高大型語言模型在複雜推理任務中效率的新方法。通過訓練模型學習最佳的壓縮時機和方法，LightThinker顯著減少了記憶體負擔和計算成本，同時保持了競爭性的準確度。不過，仍然存在一些限制：尚未探索與像LoRA或QLoRA這樣的參數高效微調方法的兼容性，對於更大訓練數據集的潛在好處尚不清楚，並且在小數據集上進行下一個標記預測訓練時，Llama系列模型的性能下降顯著。

查看 論文。這項研究的所有功勞都歸於這個項目的研究人員。此外，歡迎在Twitter上關注我們，並別忘了加入我們的80k+ ML SubReddit。

🚨 推薦閱讀 – LG AI Research發布NEXUS：一個先進的系統，整合代理AI系統和數據合規標準，以解決AI數據集中的法律問題。