谷歌Titans模型論文解析，Transformer迎來變革拐點—

一、引入 Titans 模型

我們將深入探討谷歌研究院的一篇新論文《Titans: Learning to Memorize at Test Time》，該論文介紹了一種名為 Titans 的新模型架構。
在這里插入圖片描述

Titans 在緩解 Transformer 二次方成本問題的同時，展現出了令人期待的成果。Titans 模型的設計靈感來源于人類大腦的記憶運作方式。論文中有一句有趣的話提到，記憶是一種基本的心理過程，是人類學習不可或缺的組成部分。如果沒有正常運作的記憶系統，人類和動物將局限于基本的反射和刻板行為。

在這里插入圖片描述

二、深度神經長期記憶模塊

Titans 論文的一個關鍵貢獻是深度神經長期記憶模塊（deep neural long-term memory module）。我們先來了解一下它是什么，然后再探討它是如何融入 Titans 模型的。

與循環神經網絡中將記憶編碼為固定向量不同，神經長期記憶模塊本身就是一個模型，是一個多層神經網絡，它將過去的歷史抽象編碼到自身的參數中。要訓練這樣一個模型，一種思路是訓練它記憶訓練數據。然而，眾所周知，記憶會限制模型的泛化能力，可能導致性能不佳。

三、避免過擬合的記憶機制

研究人員設計了一種巧妙的方法，來創建一個能夠記憶，但又不會在訓練數據上過擬合的模型。這種方法的靈感來自人類記憶的一個類比。當我們遇到令人驚訝的事件時，我們更有可能記住它。神經長期記憶模塊的學習過程正是基于這一理念設計的。

3.1 對 “驚訝” 進行建模

Updating the neural memory using a surprise element

通過論文中的上述定義，我們可以了解研究人員是如何對 “驚訝” 進行建模的。 Mt代表時刻t的神經長期記憶模塊。它通過上一時刻的參數以及一個被建模為梯度的 “驚訝” 元素進行更新。如果梯度較大，說明模型對輸入更 “驚訝”，這會導致模型權重得到更顯著的更新。

然而，這個定義仍不理想，因為模型可能會錯過 “驚訝” 時刻之后緊接著發生的重要信息。

3.2 對 “過去的驚訝” 進行建模

從人類的角度來看，一個令人驚訝的事件不會在很長一段時間內持續讓我們感到驚訝，盡管它仍然令人難忘。我們通常會適應這個令人驚訝的事件。不過，該事件可能足夠令人驚訝，從而在很長一段時間內吸引我們的注意力，使我們記住整個時間段。

Updating the neural memory using both past and momentary surprise

通過論文中的上述定義，我們可以了解到改進后的建模方式，其中包括對 “過去的驚訝” 的建模。現在，我們使用上一時刻的權重狀態以及一個 “驚訝” 分量St來更新神經長期記憶的權重。“驚訝” 分量現在是隨時間測量的，它由帶有衰減因子的上一時刻 “驚訝” 和我們在上一節討論過的當下 “驚訝” 組成。

還有一個尚未建模的因素是遺忘。

3.3 對遺忘進行建模

Allowing the model to forget past information

在處理非常長的序列（例如數百萬個標記 tokens）時，管理哪些過去的信息應該被遺忘至關重要。通過論文中的上述定義，我們可以看到最終的建模方式。這些定義與上一節的定義相同，只是我們添加了一個自適應遺忘機制，用 alpha 表示，也稱為門控機制（gating mechanism）。這使得記憶能夠遺忘不再需要的信息。

3.4 損失函數

Neural memory loss function

損失函數由上述公式定義。該損失旨在通過將過去的數據存儲為鍵值對來模擬聯想記憶，并訓練模型在鍵和值之間進行映射。與 Transformer 類似，線性層將輸入投影為鍵和值。然后，損失值loss衡量記憶模塊學習鍵值對之間關聯的程度。

需要說明的是，模型不是一次性處理整個序列，而是逐步處理，在其權重中積累記憶信息。

四、Titans 架構

4.1 Titans 架構 1—— 記憶作為上下文（MAC）

論文提出了幾種不同的架構。第一種稱為記憶作為上下文（MAC）。我們可以通過論文中的下圖來了解它。

Titans Memory as a Context Architecture

該模型由 3 個組件組成：持久記憶、核心和上下文記憶。給定一個長輸入序列，我們將其劃分為較小的序列。核心組件中顯示的序列實際上是一個分塊后的序列。

持久記憶

我們首先引入持久記憶，在序列開頭添加與數據無關的可學習權重。實際上，這些是可學習的標記，用于存儲不受輸入影響的全局信息。序列還會經過一個注意力塊，在序列開頭添加可學習標記有助于避免注意力分散，因為由于偏差，注意力權重往往會過度關注初始標記。英偉達的 Hymba 模型添加可學習元標記的做法也取得了成功，我后續會介紹。

上下文記憶

上下文記憶（Contextual memory）由深度神經長期記憶模塊構成，因為它依賴于上下文。我們使用神經長期記憶模塊的前向傳遞，以分塊后的輸入序列作為輸入，從長期記憶中檢索記憶標記（memory tokens）。由于即使在測試階段神經記憶也在不斷更新，當我們處理長序列的一個分塊時，神經記憶已經更新了同一序列先前分塊的信息。檢索到的記憶標記會在持久記憶標記之后添加到輸入序列中。

核心組件

核心組件將所有信息整合在一起，使用一個注意力塊，該注意力塊的輸入是擴展后的序列，其中包含來自持久記憶、上下文記憶和輸入本身的信息。通過這種方式，模型可以利用多種類型的記憶。注意力塊的輸出用于更新神經記憶。注意力機制有助于根據給定的輸入確定是否使用長期記憶。此外，注意力機制有助于長期記憶僅存儲當前上下文中有用的信息。最終輸出基于注意力塊的輸出和神經記憶的輸出確定。

4.2 Titans 架構 2—— 記憶作為門控（MAG）

Titans 的下一個架構版本稱為記憶作為門控（MAG: Memory as a Gate），我們可以通過論文中的下圖來了解它。

Titans Memory as a Gate Architecture

在這個版本中，我們同樣有代表持久記憶、核心和上下文記憶的 3 個分支。與前一個版本不同的是，這里的序列沒有分塊，輸入序列是完整的輸入序列。這是通過在注意力塊中使用滑動窗口注意力實現的。持久記憶的可學習權重同樣添加到序列開頭。但與之前不同的是，神經記憶不會為注意力塊的上下文提供數據。相反，神經記憶根據輸入序列進行更新，其輸出通過門控機制與核心分支相結合。

4.3 Titans 架構 3—— 記憶作為層（MAL）

Titans 架構的第三個變體稱為記憶作為層（MAL）。我們可以通過論文中的下圖來了解它。

Titans Memory as a Layer Architecture

與前一個版本類似，這里的序列也沒有分塊，并且我們使用滑動窗口注意力。

在這個版本中，我們將神經記憶用作模型的一層，輸入序列與可學習權重首先經過神經記憶，然后再經過注意力塊。這種設計允許堆疊多個神經記憶模塊和注意力塊，就像通常堆疊 Transformer 層一樣。然而，這種順序設計使得模型的能力受到每一層能力的限制。這與 Hymba 論文中的觀察結果類似，出于同樣的原因，Hymba 論文中的不同組件采用并行而非順序放置。