Transformer繼任者「Titans」來了，上下文記憶瓶頸被打破

長江后浪推前浪。

原標題：Transformer繼任者「Titans」來了，上下文記憶瓶頸被打破
文章來源：Founder Park
內容字數：10921字

谷歌推出Titans架構：Transformer的繼任者？

谷歌在Transformer架構發布八年后，推出了全新的Titans架構，旨在解決LLM長上下文處理的難題。該架構的核心在于其創新的長期神經記憶模塊，允許模型在測試時學習和記憶信息，從而將上下文窗口擴展到200萬tokens。

Titans架構的核心：長期神經記憶模塊

1. **挑戰與思路:** 傳統的注意力機制雖然能精確建模依賴關系，但計算成本高，限制了上下文窗口長度。Titans通過引入長期神經記憶模塊來解決這個問題，該模塊能夠學習如何在測試時記憶或忘記信息，避免過擬合并提升泛化能力。

2. **學習過程與意外指標:** 模型通過在線學習的方式，學習記憶函數。它利用“意外指標”來衡量輸入數據與過去數據的偏差，從而決定哪些信息需要記憶。該指標包含過去意外和瞬時意外兩個方面，基于一個損失函數進行學習，目標是實現聯想記憶（鍵值對存儲）。

3. **遺忘機制:** 為了管理有限的記憶容量，Titans采用了自適應遺忘機制，允許模型忘記不再需要的信息。

4. **記憶架構與檢索:** 長期記憶模塊采用多層MLP架構，檢索信息則通過簡單的線性層投影和前向傳遞實現。

5. **并行化訓練:** 為了高效訓練，Titans利用了塊式梯度下降、張量化以及參數作為塊函數的簡化方法，充分利用硬件加速器。

Titans架構的三個變體

1. **Memory as a Context (MAC):** 將記憶作為當前信息的上下文，結合注意力機制，選擇性地利用歷史信息。

2. **Memory as a Gate (MAG):** 使用滑動窗口注意力作為短期記憶，神經記憶模塊作為衰減記憶，兩者協同工作。

3. **Memory as a Layer (MAL):** 將循環模型和具有注意力機制的深度神經網絡進行堆疊。

實驗結果與結論

在語言建模、常識推理、基因組學和時序預測等任務中，Titans架構均表現優于Transformer和其他的現代線性循環模型，尤其在“大海撈針”任務中，其在超過200萬tokens的上下文窗口下取得了更高的準確率。實驗還表明，更深的記憶模塊和更大的參數量可以提升模型性能。盡管訓練速度略慢于部分基線模型，但Titans (MAL) 的訓練速度仍然優于其他一些基線模型和神經記憶模塊。

總而言之，Titans 架構提出了一種新穎的將記憶機制融入深度學習架構的方法，為處理超長上下文和提升LLM性能提供了新的思路。

聯系作者

文章來源：Founder Park
作者微信：
作者簡介：來自極客公園，專注與科技創業者聊「真問題」。