https://research.google/blog/introducing-nested-learning-a-new-ml-paradigm-for-continual-learning/

嵌套學習：深度學習架構的幻象

作者：Ali Behrouz、Meisam Razaviyayn、Peilin Zhong、Vahab Mirrokni
機構：Google Research

📖 全文翻譯與章節重點總結

📌 摘要

核心觀點：

  • 提出了 嵌套學習 這一新的學習典範，將機器學習模型表示為一組嵌套的、多層次的優化問題，每一層都有自己的「上下文流」。
  • 從 NL 的角度看，現有深度學習方法通過壓縮自身的上下文流來從數據中學習，而大型模型中的上下文學習自然湧現。
  • NL 提出了一種哲學思想，即設計具有更多「層次」、更富表達力的學習演算法，從而實現更高階的上下文學習，並可能解鎖有效的持續學習能力。

主要貢獻：

  1. 可表達的最佳化器：展示了 Adam、帶動量的 SGD 等基於梯度的最佳化器實際上是試圖壓縮梯度資訊的關聯記憶體模組，並在此基礎上提出了具有更深記憶或更強大學習規則的「更富表達力」的最佳化器。
  2. 自修改學習模組：利用 NL 對學習演算法的洞察，提出了一種通過學習自己的更新演算法來學習如何修改自身的序列模型。
  3. 連續體記憶系統：提出了一種新的記憶系統表述，推廣了傳統的「長/短期記憶」觀點。結合自修改序列模型和連續體記憶系統，提出了名為 Hope 的持續學習模組，在語言建模、知識整合、少樣本泛化、持續學習和長上下文推理任務中顯示出良好效果。

1️⃣ 引言

核心問題：

  • 儘管深度學習（尤其是大型語言模型）取得了巨大成功，但它們在初始部署後基本上是靜態的，無法在上下文之外持續獲取新能力，只能通過上下文學習進行快速適應。
  • 當前模型（尤其是 LLMs）的知識僅限於其上下文視窗內的即時上下文或預訓練結束時儲存在 MLP 中的「長期過去」知識，類似於「順行性遺忘症」，無法形成新的長期記憶。

神經科學動機：

  • 人腦的持續學習能力 歸功於神經可塑性和神經振盪（腦波）。
  • 多時間尺度處理系統：腦波（如 Gamma、Beta、Theta）以不同頻率協調神經計算和資訊處理，對應不同的認知功能和記憶鞏固時間尺度。
  • 大腦的均勻可重用結構：大腦結構是均勻且可重用的（如半球切除術後仍能正常運作），記憶是跨多個區域分佈的神經迴路，而非孤立於特定區域。

貢獻與路線圖：

  1. 介紹 嵌套學習典範。
  2. 論證 最佳化器和架構都是學習模組。
  3. 重新審視 持續學習、上下文學習、預訓練和學習 等常見概念。
  4. 提出 連續體記憶系統、自引用 Titans 和 Hope 架構。
  5. 提供 實驗評估，支持 NL 觀點和提出的元件的有效性。

2️⃣ 預備知識

回顧概念：

  • 梯度下降 及其變體的公式化表述。
  • 元學習：將模型設計自動化建模為一個兩級優化過程。
  • 快速權重程式設計器 (FWPs)：一種循環神經網路，其記憶是矩陣值的，作為短期記憶。
  • 上下文學習 (ICL)：模型根據給定上下文自適應和學習的能力。本文使用最廣義的定義。

3️⃣ 嵌套學習 (Nested Learning)

核心思想：

  • 將機器學習模型（包括其訓練過程）表示為一個相互關聯的、多層次的、嵌套的優化問題系統。
  • 每個元件（架構、最佳化器）都有自己的上下文流和優化目標，它們在不同層次以不同頻率更新。

關鍵定義：

  1. 關聯記憶體：一個將一組鍵映射到一組值的算子。學習是獲取有效記憶體（映射）的過程。
  2. 更新頻率：定義每個元件（參數或非參數）的更新頻率。
  3. 嵌套系統：一個有序的多層次系統，每層包含一組用梯度下降優化的（關聯記憶體）問題。
  4. 神經學習模組：架構和優化過程共同確定的模型的聯合表示。

層次間的知識轉移方式：

  • 參數化直接連接（如 FWPs）
  • 非參數化直接連接（如注意力）
  • 通過反向傳播的知識轉移
  • 通過初始化（如 MAML）
  • 通過生成（權重或上下文生成，如超網路、優化過程）

重要洞見：

  • 訓練一個線性層：通過反向傳播訓練一個線性層可以被視為一個基於「驚奇」的記憶體，它試圖將數據樣本映射到其預測誤差。
  • 帶動量的梯度下降：可以視為一個兩級關聯記憶體（優化過程），其中內層學習將梯度壓縮到其參數中。
  • 架構分解：像線性注意力這樣的架構可以分解為嵌套的優化問題（例如，投影層在低頻優化，記憶體矩陣在高頻優化）。
  • 堆疊層次：嵌套學習允許計算模型由多個（多層）層次組成，以不同的抽象層次和更新頻率學習和處理數據。

4️⃣ 最佳化器作為學習模組

核心觀點：

  • 反向傳播作為關聯記憶體：使用反向傳播訓練神經網路可以看作一個壓縮過程，其中每一層都在儲存其輸入與對應的局部誤差信號（驚奇）之間的映射。這是一個 自我指涉過程，因為用於記憶體的值是由記憶體本身生成的。
  • 基於動量的最佳化器作為關聯記憶體：動量項可以看作一個旨在將過去梯度壓縮到其參數中的關聯記憶體。它是一個兩級優化過程。
  • 流行的最佳化器如 Adam、AdaGrad 都可以被重新表述為關聯記憶體，它們試圖將梯度映射到其某種全域屬性（如方差）。
  • 為持續學習設計更具表達力的最佳化器：現有動量設計（如指數移動平均）容量有限，對長遠過去梯度記憶不佳。NL 視角啟發我們設計：
    • 更具表達力的關聯（如添加預處理）。
    • 更具表達力的目標（如用 L2 回歸損失代替點積相似度，得到 Delta Momentum）。
    • 更具表達力的記憶體結構（如用 MLP 代替線性層作為動量記憶體，得到 Deep Momentum GD）。
    • 更高階的特徵映射。
    • 非線性輸出（如 Muon 最佳化器中的牛頓-舒爾茨正交化）。
  • Delta 梯度下降 (DGD)：用 L2 回歸損失代替點積相似度作為內部目標，得到的更新規則不僅依賴於當前數據，還依賴於權重的先前狀態，從而能捕獲數據樣本間的依賴關係。
  • 廣義梯度下降 (GGD)：將梯度下降定義為任何旨在壓縮數據樣本並將其映射到一組自生成值的自我指涉關聯記憶體。

5️⃣ 現有架構作為神經學習模組

核心觀點：

  • 現代序列模型（Transformer, RNNs）都可以被視為關聯記憶體，它們旨在根據某種內部目標（注意力偏差）和最佳化演算法學習從鍵到值的映射。
  • 從 NL 視角看，所有現代架構都是均勻的：它們都由一組前饋網路（線性或非線性 MLP 塊）組成，每個網路根據自己的上下文流和優化問題在不同層次和時間尺度上進行訓練。我們觀察到的異質性源於只看到了這些優化問題的最終解，造成了「深度學習架構的幻象」。
  • 具體示例：
    • Softmax 注意力：是某個 L2 回歸目標的非參數解（Nadaraya-Watson 估計器）。
    • 基於 Hebbian 規則的 RNNs（如線性注意力）：是使用點積相似度作為內部目標，用梯度下降優化的關聯記憶體。
    • 基於 Delta 規則的 RNNs（如 DeltaNet）：是使用 MSE（L2 回歸）作為內部目標，用梯度下降優化的關聯記憶體。
    • 其他學習規則：如 Oja‘s rule, Omega rule 等，也都是特定內部目標和最佳化演算法的實例。
  • 門控機制的作用：在現代序列模型中，門控線性層的輸出可以被視為 MLP 塊的持久記憶體，或者當記憶體初始狀態未元學習時，作為記憶體模組的初始化。

6️⃣ 要點總結與概念重審

對常見術語的 NL 視角解讀：

  • 記憶體與學習：記憶體是由輸入引起的神經更新，學習是獲取有用記憶體的過程。記憶體不是孤立系統，而是分佈在所有參數中。任何由梯度下降（或其他演算法）引起的更新都是記憶體的一種形式。
  • 模型參數：模型的參數不僅包括在預訓練級別優化的那些。NL 表示中出現的所有參數（如動量項、RNN 的隱藏狀態）都對其性能和表達能力有貢獻。
  • 每個神經元的更多計算：堆疊層次可以增加每個低頻參數的計算深度（如 Muon 最佳化器中的正交化步驟）。
  • 上下文學習 (ICL)：從 NL 視角看，每個塊或層次都有自己的上下文流，因此任何對該上下文的神經更新或適應都被視為一種 ICL。Transformer 的 ICL 是非參數 ICL 的一個特例，而 RNN 的記憶體執行的是參數化 ICL。
  • 測試時訓練/記憶體：是參數化 ICL 的實例，當上下文移除時，獲取的上下文知識也隨之消失。
  • 預訓練與測試時間：預訓練只是 ICL 的一個可能實例，其上下文是整個預訓練數據。訓練和測試時間的區分是由於缺乏從最高頻率層（如 Transformer 的上下文）到低頻層（如預訓練）的知識轉移過程。
  • 持續學習：從 NL 視角看，任何機器學習模型都在進行持續學習。知識可能無法持久或轉移到更持久的層次，主要是由於層次間缺乏知識轉移。
  • 現有架構骨幹與混合模型：從 NL 視角看，所有現代架構都是均勻的。混合架構可以看作是傳統的 Transformer 模型，只是為某些 MLP 塊增加了新的計算層次。循環模型正在取代 MLP 塊。
  • 神經學習模組作為互聯繫統：架構為最佳化器生成上下文（即梯度）。因此，梯度的適當記憶體管理依賴於架構的選擇。未來需要設計架構特定的最佳化器，以使這個互聯繫統完美和諧地工作。
  • 最佳化器 vs. 學習到的最佳化器：兩者是同一概念的不同實例，只是頻率和上下文流不同。

7️⃣ 連續體多時間尺度記憶系統

核心觀點：

  • 推廣傳統的「長/短期記憶體」觀點，提出 連續體記憶系統 (CMS)：將記憶體視為一個具有更新頻率譜的分佈式互聯繫統。
  • 高頻神經元負責快速適應但儲存記憶體時間短，低頻神經元負責更持久的知識。
  • 與 LSM 相比，這種多頻率設計形成了一個記憶體循環過程，知識在被遺忘後可以部分恢復。

CMS 設計：

  • 由一系列 MLP 塊鏈組成，每個塊與一個更新頻率（塊大小）相關聯。塊參數根據其頻率定期更新（例如，每 C 步更新一次）。
  • 變體：
    1. 嵌套 CMS：高層級 MLP 塊的初始狀態由低層級元學習，允許高階 ICL。
    2. 順序 CMS：所有 MLP 塊的初始狀態都在最低頻率層通過反向傳播元學習。
    3. 獨立（頭-wise）CMS：獨立的塊具有不同的上下文長度，然後通過聚合函數組合。
  • CMS 有助於持續學習：當一個塊更新並遺忘舊知識時，這些知識仍可能儲存在其他低頻塊中，並且通過反向傳播的知識轉移可以循環回來，形成一個時間維度的循環，從而很難遺忘重要知識。
  • CMS 的效率：每次只更新一小部分參數（根據其頻率），並且其設計解鎖了高層級的序列並行化。
  • CMS 在最佳化器中的應用：作為概念驗證，設計了 多尺度動量 Muon (M3) 最佳化器，結合了 Adam、Muon 和 CMS 的思想，具有多個動量項以更好地壓縮長期梯度資訊。
  • 臨時層次堆疊：用預訓練模型初始化 CMS：可以利用預訓練模型的權重來初始化 CMS 塊的參數，從而幫助模型更快地適應其層次的上下文流。

8️⃣ Hope：具有連續體記憶體的自指涉學習模組

設計動機：

  • Transformer 是兩層級元件（投影和 MLP 在第一層，注意力是第二層的非參數解），計算深度有限，且參數在上下文中是靜態的，缺乏自我修改能力。
  • 需要一種允許所有元件在上下文中自適應，並能學習自我修改的模型。

Hope 架構組成：

  1. 自修改深度 Titans：
    • 一個完全自適應的記憶體系統，其中所有元件（生成鍵、值、查詢、學習率、衰減率的映射）都能在上下文中適應。
    • 關鍵創新：模型學習根據上下文生成自己的潛在值（目標值），從而修改自身。
    • 使用 Delta 梯度下降 (DGD) 作為內部優化演算法，因為令牌空間中的元素高度相關。
    • 所有記憶體（映射函數）的架構都是任意的，本文使用 2 層 MLP。
    • 採用分塊訓練演算法實現快速並行化訓練。
  2. 連續體記憶系統 (CMS)：
    • 為模型提供更持久的記憶體儲存。

Hope 的前向傳播：

  • 輸入首先通過自修改 Titans 模組進行處理，其輸出再通過 CMS 鏈。
  • 輸出是這兩個系統協同工作的結果。

Hope-Attention 變體：將自修改 Titans 替換為 softmax 全域注意力。

9️⃣ 實驗評估

實驗評估了 NL 觀點及其元件（如 CMS、Hope）在多個任務上的有效性：

  1. Hope：持續學習與長上下文理解
    • 類增量學習（CLINC, Banking, DBpedia）：Hope 在持續學習任務上優於 ICL、EWC、InCA 等基線，表明 CMS 設計對增強持續學習能力有效。
    • 層級對上下文學習的影響：在 LongHealth、QASPER、MK-NIAH 任務上，更多記憶層級和適當的更新頻率能提升模型的長上下文理解和 ICL 能力。
    • 上下文學習新語言 (CTNL)：Hope 在持續翻譯新語言任務上顯著優於 ICL，能有效緩解災難性遺忘。
  2. Hope：長上下文理解
    • 大海撈針 (NIAH) 任務：Hope 在各種設置（單針、多鍵、多查詢、多值）和長度下，性能優於其他注意力無關模型（RWKV-7, Comba, Titans）及 Transformer。Hope-Attention 也優於普通 Transformer。
    • BABILong 基準測試：Hope 在長達 10M 的上下文長度上保持良好性能，優於其他小型微調模型和大型零樣本模型。
  3. Hope：語言建模與常識推理
    • 在 Wikitext、LAMBADA、PIQA 等多個基準測試上，Hope 在平均性能上優於所有基線模型（包括 Transformer 和各種 RNNs），並且在參數規模擴大時表現出更高的性能增益。
  4. Hope：上下文召回任務與 MAD 合成基準測試
    • 在多個上下文召回數據集上，Hope 優於所有注意力無關模型，縮小了與 Transformer 的差距。
    • 在 MAD 合成基準測試上，Hope 在壓縮、記憶、複製等任務上優於所有基線，包括 Transformer。
  5. 語言識別任務
    • 在形式語言識別任務上，Hope 取得了與 LSTM、SRWM 等非線性循環模型相當的完美分數，同時保持了可並行訓練的優勢。
  6. Hope：消融研究與擴展
    • 消融研究表明，Hope 的所有元件（DGD、動量、權重衰減、CMS、內部投影）都對其性能有積極貢獻。
    • 上下文使用越多，模型困惑度越低，表明其具有強大的記憶體管理能力。
  7. 可表達的最佳化器
    • ImageNet：在 ImageNet-21K 上訓練 ViT，M3 最佳化器比 AdamW 和 Muon 獲得了更低的訓練/測試損失。
    • 大模型效率：M3 由於使用多個動量（記憶體），訓練速度相對 Muon 較慢，但與 AdaMuon 效率相當。

🔚 結論

  • 本文提出了 嵌套學習 (NL)，這是一個新的學習典範，將現代機器學習系統建模為相互關聯的、多層次的優化問題。
  • 從這個角度看，架構和最佳化器都是嵌套關聯記憶體系統的實例，它們壓縮自己的上下文（令牌、梯度或高層信號）。
  • NL 將預訓練、上下文學習和持續學習重新建構為同一底層機制的表現：在不同層次和時間尺度上學習和重用上下文。
  • 基於 NL 的觀點，推導了 廣義的基於梯度的更新（如 DGD、Delta Momentum、M3），並將現代序列架構重新解釋為嵌套關聯記憶體。
  • 為了增強記憶體處理，引入了 連續體記憶系統 (CMS)。基於自修改 Titans 和 CMS 的 Hope 架構 提高了持續學習和長上下文推理能力，同時作為通用骨幹具有競爭力。
  • 災難性遺忘解決了嗎？ 從 NL 視角看，災難性遺忘是壓縮的自然結果，因為網路的有限容量迫使模型遺忘以保留新資訊的容量。NL 是一個路線圖，它表明在持續學習、長上下文推理、現代最佳化器和自修改模型方面的進展將來自於更好地利用「層次」這一額外設計軸，而不是來自更深的靜態網路。