Mamba

Mercuryadore發表於2024-04-16

Graph-Mamba

論文：Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces

對於Graph-Mamba，需要了解一些相關知識。

Mamba結構

Mamba建立在狀態空間模型的概念之上，但引入了一個新的變化。它利用選擇性狀態空間 (SSM: selective state space model)，支援跨長序列更高效和有效地捕獲相關資訊。

Mamba結構最初是在Mamba: Linear-Time Sequence Modeling with Selective State Spaces中提出的，

image-20240310103140494

為了解決 Transformer 在長序列上的計算效率低下問題，人們開發了許多亞二次時間架構，如線性注意力、門控摺積和遞迴模型以及結構化狀態空間模型（SSM）。此類模型的一個關鍵弱點是無法進行基於內容的推理。

作者提出了一類新的選擇性狀態空間模型，該模型在多個維度上改進了先前的工作，從而在序列長度線性縮放的同時，實現了Transformer的建模能力。主要的改進有以下幾點：

  1. 選擇機制 先前的模型，特別是那些用於處理序列資料的模型，可能在有效選擇資料方面存在侷限，它們可能不夠有效地關注重要的輸入資訊或忽略不相關的輸入資訊。選擇性複製和歸納頭等合成任務在理解和改進模型的選擇機制方面提供了重要的直覺。於是作者設計了一種簡單的選擇機制，根據輸入對 SSM 引數進行引數化（讓 SSM 引數成為輸入的函式）。這樣，模型就能過濾掉無關資訊，並無限期地記住相關資訊。
  2. 硬體感知演算法 先前的SSM模型為了計算效率，必須是時間不變和輸入不變的，這意味著它們的計算方式不隨時間或輸入資料的變化而改變。而上一點提到的改進，雖然提高了模型的靈活性和效果，但也帶來了技術上的挑戰。 於是作者開發了一種硬體感知演算法。這種演算法考慮了硬體的特性和限制，特別是在使用GPU進行計算時。演算法使用遞迴的方式透過掃描來計算模型，而不是使用傳統的卷積。 演算法設計考慮到了GPU記憶體層次結構，為了避免不同級別之間的IO訪問，它不會實體化擴充套件的狀態，這一步驟有助於減少記憶體訪問和相關的時間延遲。 在理論上，這種實現方式使得處理速度隨序列長度線性增長，與基於卷積的所有SSM模型的偽線性增長相比有顯著優勢。
  3. 模型架構 作者將先前的 SSM 架構設計（Dao、Fu、Saab 等人，2023 ）與 Transformers 的 MLP 模組合併為一個模組，整合到一個簡化的端到端神經網路架構中，該架構沒有注意力，甚至沒有 MLP 塊，簡化了先前的深度序列模型架構，形成了一種包含選擇性狀態空間的簡單、同質的架構設計（Mamba）。

其演算法的核心內容可以用下圖表示：

image-20240310103528687

公式以及結構的推導可見下連結：

一文讀懂Mamba：具有選擇狀態空間的線性時間序列建模 - 知乎 (zhihu.com)

程式碼實現：

挑戰Transformer的新架構Mamba解析以及Pytorch復現 - 知乎 (zhihu.com)

GNN and GCN

Graph Neural Networks (GNN): Graphical-Based Learning Environments for Pattern Recognition | SpringerLink

Graph Convolutional Networks (GCN):[1609.02907] Semi-Supervised Classification with Graph Convolutional Networks (arxiv.org)

Graph Neural Networks (GNN)

原論文過於枯燥，下面這個專題的文章講解很詳細只粗略瀏覽了一些

GNN綜述——從入門到入門 - 知乎 (zhihu.com)

Graph Convolutional Networks (GCN)

image-20240310105616511

核心原理圖如下：

image-20240310105922025

Graph Transformers

[2012.09699v2] A Generalization of Transformer Networks to Graphs (arxiv.org)

  • Transformer and GNN
image-20240310111520449
  • Graph Transformer結構
image-20240310110954609

GraphGPS

[2205.12454] Recipe for a General, Powerful, Scalable Graph Transformer (arxiv.org)

the replacement of fully-connected Transformer attention with its sparse alternatives

image-20240310112012319 image-20240310112112073

State Space Models(SSM)

image-20240310112917370

Utilizing a recurrent scan, SSM stores context in its hidden states, and updates the output by combining these hidden states with input.

image-20240310113120089

Graph-Mamba

image-20240310113542799

image-20240310113622570

Graph-Mamba採用選擇性SSM來實現依賴於輸入的圖形稀疏化，設計了一個Graph-Mamba塊(GMB)，並將其合併到流行的GraphGPS框架中，以便與其他圖形注意實現進行公平的比較。GMB利用迴圈掃描序列建模和選擇機制來實現兩個級別的圖形稀疏化。第一級涉及Mamba模組中的選擇機制，它有效地過濾遠端上下文中的相關資訊。第二級是透過提出的節點優先順序方法實現的，允許圖中的重要節點訪問更多上下文。因此，這些序列建模特徵為圖稀疏化提供了一種結合資料依賴和啟發式資訊選擇的有前途的途徑。

Structured state space models for sequence

image-20240310113858027 image-20240310113954830

概述Graph-Mamba架構，透過合併GMB來取代GraphGPS框架中的注意力模組。

  • A) GMB層，是對GPS層的改進，結合了基於邊緣的MPNN和以節點為中心的GMB來輸出更新的節點和邊緣嵌入。
  • B) Graph-Mamba將GatedGCN模型作為MPNN的預設模型。
  • C) GMB使用節點優先順序和排列技術的專門訓練配方，執行圖的稀疏化。
  • D)帶有Mamba模組的選擇機制，方便了依賴於輸入的上下文過濾

Graph-dependent selection mechanism

image-20240310115010604

在以節點為輸入序列的圖學習中，選擇機制允許隱藏狀態基於先前序列中的相關節點進行更新，由當前輸入節點進行門控，並隨後影響當前節點的輸出嵌入。Gt的範圍在0到1之間，允許模型在需要時完全過濾掉不相關的上下文。選擇和重置的能力使Mamba能夠在給定長期上下文的情況下提取相關的依賴關係，同時最大限度地減少復發的每個步驟中不重要節點的影響。因此，它提供了一種上下文感知的替代方案，透過僅在長輸入序列中保留相關依賴關係來稀疏圖的注意力。

相關文章