機器之心編輯部
很長一段時間，「長上下文」一直是各大模型廠商軍備競賽的焦點，從 128K 到 1M，再到更長的上下文視窗，業界已然形成一個固有認知，只要視窗足夠大，模型就能記住更多內容，也就能處理更長、更復雜的任務論文。
但問題也隨之而來：上下文越長，KV Cache 越臃腫，不僅導致視訊記憶體瞬間被「吃光」，推理速度愈發緩慢，成本也迅速上升論文。
更關鍵的是論文，把更多 token 放進視窗，並不等於模型真的把這些資訊轉化成了可推理的長期記憶，結果是，榜單分數越刷越高，可在一些需要「深度腦暴」的複雜推理任務中，模型常常因為「記不住細節」，頻頻翻車……
面對這一兩難問題論文，近日，卡內基梅隆大學（CMU）聯合馬里蘭大學等在一篇新論文中提出了有意思的視角：既然人類連續工作久了會變笨，大模型也一樣，既然如此為什麼不讓 LLM 睡一覺呢？
這篇論文的題目言簡意賅，《Language Models Need Sleep》，也就是《語言模型需要睡眠》論文。
展開全文
當然，這裡的「睡眠」不是真的睡眠，更準確地說，是一種類似睡眠的「記憶鞏固機制」論文。
作者認為，基於 Transformer 的大語言模型正越來越多地被用於長程任務，然而，其注意力機制在面對更長上下文時擴充套件性較差論文。為此，他們研究出了這一「記憶鞏固機制」：
在睡眠過程中，模型會對累積的上下文執行 N 次離線遞迴前向傳播，並透過一種學習得到的區域性規則，更新其狀態空間模型（SSM）模組中的快速權重（fast weights）論文。在推理階段，這種方法把額外計算轉移到「睡眠」階段，同時保持模型在「醒著」進行預測時的延遲不變。
換句話說，它不是讓模型一直把所有內容攤在眼前，而是讓模型學會在某些節點「停下來想一想」，把剛剛讀過的內容消化成之後還能呼叫的內部狀態論文。
論文地址：
作者在一系列受控的合成任務上測試了該方法，包括細胞自動機、多跳圖檢索，以及一個更貼近真實場景的數學推理任務論文。在這些任務上，普通 Transformer 和 SSM-attention 混合模型都會失敗，而增加模型的「睡眠」時長 N ，可以提升效能，其中在需要更深層推理的樣本上，提升最為明顯。
接下來，我們來詳細瞭解一下論文。
從動物睡眠中獲得啟發
這篇論文的靈感，來自動物睡眠中的記憶鞏固過程論文。
神經科學的研究認為，動物從短期記憶到長期記憶的轉移，是受到海馬體 replay 機制的支援，尤其是在睡眠期間論文。在這一階段，短期的海馬體記憶會被重新啟用，並鞏固到皮層突觸權重中。睡眠會讓動物無法對外部刺激作出反應，這也說明睡眠必須帶來足夠大的認知收益，才值得付出這一代價。
基於這一認知，作者提出了這種把上下文視窗記憶轉移到持久權重中的方法，即當模型的上下文視窗在推理過程中被填滿時，模型就會進入「睡眠」狀態：對累積的上下文執行多次前向傳播，並透過學習得到的區域性規則遞迴地更新 fast weights，在這個階段，模型不會接收外部輸入 token論文。
鞏固完成後，上下文視窗會被清空，模型則帶著更新後的 fast weights 繼續執行論文。在訓練過程中，模型透過整個過程的反向傳播進行端到端最佳化，以最大化睡眠之後的任務表現。
也就是說論文，大模型的訓練過程被劃分為兩個階段：
「醒著」階段：只負責快速響應，模型就像普通的 Transformer 一樣正常工作，它接收長文字輸入，快速給出預測和回覆，這時候它不需要對資訊進行深度內化，只管「讀」和「答」論文。
「睡眠」階段：每隔一段時間，模型就會進入「離線睡眠狀態」，期間模型會利用專門的後臺時間，對積累的上下文進行 N 次迴圈往復的離線處理（Recurrent passes），將近期上下文中的關鍵細節，轉化為持久的 fast weights，並寫入其狀態空間模型（SSM）模組中論文。
具體如下論文。
當上下文視窗被填滿、模型即將從注意力層中淘汰 token 之前，模型會先進入一個「鞏固階段」，在這一階段執行遞迴計算，透過這種方式擴充套件計算量來處理深度推理任務，對於較大的 時間步 t，仍然滿足預測階段的延遲約束論文。
例如論文，如果在全部 D 個模組上進行迴圈，其形式如下：
其中，N 表示在整個架構上迴圈執行 N 次傳遞論文。
下圖對架構進行了詳細描述，從一個 SSM-Attention 混合模型初始化，該模型具有固定的上下文視窗大小 L，其中注意力快取每 L 個 token 就會被完全淘汰論文。在每 L 個 token 淘汰 KV Cache 之前，模型會執行 N 次遞迴傳遞，根據下面的公式 3 迭代更新 SSM 模組內部的快速權重；當 N = 1 時，它就退化為一個普通的 SSM-Attention 混合模型。模型在迭代更新快速權重的這一階段就是「睡眠階段」。
在遞迴式地細化快速權重之後，KV Cache 會被淘汰，模型隨後處理接下來的 L 個 token論文。
在完整上下文處理完畢後，模型會基於已經細化後的記憶和當前上下文，透過一次前向傳播來預測答案論文。訓練時，模型透過對公式 6 所示的整個計算圖進行反向傳播，最小化預測誤差，這一點與其他深度遞迴模型類似。
不同的是，以往的深度遞迴模型中，梯度會流經遞迴細化後的特徵向量；而在這裡，由於睡眠階段結束後，細化後的特徵會被丟棄，梯度實際上流經的是被細化後的快速權重論文。
完整的訓練流程如下所示論文：
實驗：睡得越久論文，推理越強？
為了驗證：增加睡眠時 N，到底能不能提升模型對「舊」上下文的推理能力？作者進行了系列實驗論文。下面我們來看一個更接近自然語言的數學推理任務 GSM-Infinite。
GSM-Infinite 可以理解為一個長上下文數學推理基準，它會透過新增干擾 token 拉長題目，同時用所需算術運算元控制難度論文。題目越複雜，需要的推理步驟越多。
作者在 Jet-Nemotron 2B 和 Ouro 1.4B 兩個預訓練模型上測試了模型的「睡眠」機制論文。
結果呈現出一個清晰趨勢論文，題目越難，「睡眠」帶來的提升越明顯：
對於 Jet-Nemotron 2B論文，6 次 sleep loop 將 6 步運算題準確率從 0.742 提升到 0.812，將 8 步運算題從 0.351 提升到 0.388；
對於 Ouro 1.4B，4 次 sleep loop 將 6 步運算題準確率從 0.419 提升到 0.615，將 8 步運算題從 0.210 提升到 0.272論文。
也就是說論文，「睡眠」機制對簡單題的幫助相對沒有那麼明顯限，因為模型本來就能做得不錯；但當任務變得複雜，需要更多步推理、更強的上下文組織能力時，「睡眠」階段的額外計算就開始發揮作用了……
侷限性：效果明顯論文，代價同樣明顯
當然，這篇論文並沒有把問題說得過於樂觀論文。
作者坦言，這種方法是透過把額外遞迴計算轉移到鞏固階段，保持了預測階段的單次前向傳播延遲論文。但可這種收益不是免費的：在訓練過程中，需要執行 N 次更深的前向和反向傳播，這會讓訓練變慢，也可能變得不穩定。
而執行 N 次論文，帶來效果明顯提升是真，訓練成本隨其線性增長也是真……
因此，這項工作目前仍主要是方法論探索論文。
作者表示，這一方法主要貢獻是方法論層面的，並且評估主要基於受控合成任務和中等規模預訓練模型論文。目前，它還不是一個已經在超大規模商用模型、真實長程 Agent 系統中充分驗證的成熟方案。
更多詳情論文，可檢視論文了解！
參考連結論文：
