使用這些基本 REST API 最佳實踐構建出色的 API
Mamba: Linear-Time Sequence Modeling with Selective State Spaces一文中提出了Mamba，我們在之前的文章中也有詳細的介紹。
在本篇文章中，通過將繪制RNN，transformer，和Mamba的架構圖，并進行詳細的對比，這樣我們可以更詳細的了解它們之間的區別。
為了說明為什么Mamba是這樣一個有趣的架構，讓我們先介紹Transformer。
Transformer將任何文本輸入視為由令牌組成的序列。
transformer的一個主要優點是，無論它接收到多長的輸入，它都使用序列中的任何令牌信息（無論序列有多長）來對輸入數據進行處理。
這就是我們在論文中看到的注意力機制的作用，但是為了獲得全局信息，注意力機制在長序列上非常耗費顯存，這個我們后面說。
Transformer由兩個結構組成，一組用于表示文本的編碼器塊和一組用于生成文本的解碼器塊。這些結構可以用于多種任務，包括翻譯。
我們可以采用這種結構來創建僅使用解碼器的生成模型。比如基于Transformer的GPT，使用解碼器塊來完成一些輸入文本。
單個解碼器塊由兩個主要部分組成，一個是自注意力模塊，另一個是前饋神經網絡。
注意力創建一個矩陣，將每個令牌與之前的每個令牌進行比較。矩陣中的權重由令牌對之間的相關性決定。
它支持并行化，所以可以極大地加快訓練速度!
但是當生成下一個令牌時，我們需要重新計算整個序列的注意力，即使我們已經生成了一些新的令牌。
為長度為L的序列生成令牌大約需要L2的計算量，如果序列長度增加，計算量可能會很大。并且在這里需要計算所有令牌的注意力，所以如果序列很長，那么內存占用也會很大。所以需要重新計算整個序列是Transformer體系結構的主要瓶頸。當然也有很多技巧來提升注意力機制的效率，這里我們暫時不提，只看最經典的原始論文。
下面我們介紹更早的序列模型RNN。循環神經網絡(RNN)是一種基于序列的網絡。它在序列的每個時間步長取兩個輸入，即時間步長t的輸入和前一個時間步長t-1的隱藏狀態，以生成下一個隱藏狀態并預測輸出。
RNN有一個循環機制，允許它們將信息從上一步傳遞到下一步。我們可以“展開”這個可視化，使它更明確。
在生成輸出時，RNN只需要考慮之前的隱藏狀態和當前的輸入。這樣不會重新計算以前的隱藏狀態，這正Transformer沒有的。
這種流程可以讓RNN進行快速推理，因為的時間與序列長度線性擴展!并且可以有無限的上下文長度（理論上），因為每次推理他只取一個隱藏狀態和當前輸入，內存的占用是非常穩定的。
我們將RNN應用于之前使用過的輸入文本。
每個隱藏狀態都是以前所有隱藏狀態的聚合。但是這里就出現了問題，在生成名稱“Maarten”時，最后一個隱藏狀態不再包含關于單詞“Hello”的信息（或者說最早的信息會被坐進的信息覆蓋）。這會導致隨著時間的推移，rnn會忘記信息，因為它們只考慮前一個狀態。
并且rnn的這種順序性產生了另一個問題。訓練不能并行進行，因為它需要按順序完成每一步。
與Transformer相比，rnn的問題完全相反!它的推理速度非常快，但不能并行化導致訓練很慢。
人們一直在尋找一種既能像Transformer那樣并行化訓練，能夠記住先前的信息，并且在推理時間還是隨序列長度線性增長的模型，Mamba就是這樣宣傳的。
在介紹Mamba之前，讓我們還需要介紹以下狀態空間模型
狀態空間模型(SSM)，像Transformer和RNN一樣，可以處理序列信息，比如文本，也包括信號。
狀態空間是包含能夠完全描述一個系統的最少數量變量的概念。它是一種通過定義系統可能的狀態來數學表示問題的方式。
比如說我們正在通過一個迷宮。”狀態空間” 就是所有可能位置（狀態）的地圖。每個點代表迷宮中的一個獨特位置，具有特定的細節，比如你離出口有多遠。
“狀態空間表示” 是對這個地圖的簡化描述。它展示了你當前所處的位置（當前狀態），以及下一步可以去哪里（可能的未來）。
雖然狀態空間模型使用方程和矩陣來跟蹤這種行為，描述狀態的變量，在我們的例子中是X和Y坐標以及到出口的距離，可以表示為“狀態向量”。
聽起來熟悉嗎?這不就是強化學習中的狀態嗎，我個人認為是可以這么理解的，那么怎么和序列有關呢？
因為語言模型中的嵌入或向量也經常用于描述輸入序列的“狀態”。例如，你當前位置的向量(狀態向量)可能看起來像這樣:
在神經網絡中，“狀態”通常是指其隱藏狀態，在大型語言模型的背景下，這是生成新標記的一個最重要的方面之一。
狀態空間模型（SSMs）是用于描述這些狀態表示并根據某些輸入進行下一個狀態預測的模型。
在時間t，狀態空間模型(SSMs)：
這里就與強化學習中使用離散序列（如僅向左移動一次）不同，它將連續序列作為輸入并預測輸出序列。
ssm假設動態系統，例如在三維空間中移動的物體，可以通過兩個方程從時間t的狀態預測。
通過求解這些方程，假設可以揭示基于觀測數據(輸入序列和先前狀態)預測系統狀態的統計原理。
它的目標是找到這個狀態表示h(t)這樣我們就可以從一個輸入序列到一個輸出序列。
這兩個方程就是是狀態空間模型的核心。狀態方程描述了基于輸入如何影響狀態(通過矩陣B)的狀態變化(通過矩陣A)。
h(t)表示任意時刻t的潛在狀態表示，而x(t)表示某個輸入。
輸出方程描述了狀態如何轉化為輸出(通過矩陣C)，以及輸入如何影響輸出(通過矩陣D)。
矩陣A、B、C和D通常被稱為參數，因為它們是可學習的。將這兩個方程可視化，我們可以得到如下架構:
下面我們看看這些矩陣如何影響學習過程。
假設我們有一個輸入信號x(t)這個信號首先乘以矩陣B它描述了輸入如何影響系統。
更新狀態(h)是包含環境核心“知識”的潛在空間。我們將狀態與矩陣A相乘，矩陣A描述了所有內部狀態是如何連接的，因為它們代表了系統的潛在表示。
這里可以看到，在創建狀態表示之前應用矩陣A，并在狀態表示更新之后更新矩陣A。
然后使用矩陣C來描述如何將狀態轉換為輸出。
最后利用矩陣D提供從輸入到輸出的直接信號。這通常也被稱為跳過（殘差）連接。
由于矩陣D類似于跳過連接，所以SSM通常被視為為不進行跳過連接的部分
回到我們的簡化視圖，現在可以將重點放在矩陣A、B和C上，它們是SSM的核心。
更新原始方程并添加一些顏色來表示每個矩陣的目的
這兩個方程根據觀測數據預測系統的狀態。由于期望輸入是連續的，SSM是連續時間表示。
但是因為文字都是離散的輸入，我們還需要將模型離散化。這里就要使用* Zero-order hold * 技術
每次我們接收到一個離散信號，都會保證他的值不變，直到接收到一個新的離散信號再改變。這個過程創建了一個SSM可以使用的連續信號:
我們保持該值的時間由一個新的可學習參數表示，稱為步長?。這樣就得到了一個連續的信號并且可以只根據輸入的時間步長對值進行采樣。
這些采樣值就是我們的離散輸出!在數學上，我們可以應用Zero-order hold如下:
因為我們SSM處理的是離散信號，所以這里不是一個函數到函數，x(t)→y(t)，而是一個序列到序列，x?→y?，我們用公式表示如下:
矩陣A和B現在表示模型的離散參數，用k代替t來表示離散的時間步長。
離散化的SSM允許在特定的時間步中處理信息。就像我們之前在循環神經網絡（RNNs）中看到的那樣，循環方法在這里也非常有用，可以將問題重新表述為時間步驟：
在每個時間步長，我們計算當前輸入(Bx?)如何影響前一個狀態(Ah??)，然后計算預測輸出(Ch?)。
這種表示看起來是不是有點熟悉？其實他的處理方法和RNN一樣
也可以這樣展開:
這種技術與RNN類似，快速推理和慢速訓練。
另一種ssm的表示是卷積的表示。我們應用過濾器(核)來獲得聚合特征:
因為我們處理的是文本而不是圖像，所以我只要一維的視角:
我們用來表示這個“過濾器”的核是由SSM公式推導出來的:
可以使用SSM核遍歷每一組令牌并計算輸出:
上圖也說明了padding 可能對輸出產生的影響，所以我們一般都會在末尾padding而不是在前面。第二步核被移動一次來執行下一步計算:
在最后一步，我們可以看到核的完整效果:
卷積的一個主要好處是它可以并行訓練。但是由于核大小是固定，它們的推理不如rnn快速并且對序列長度有限制。
上面的三種SMM都有各自的優缺點
這里可以使用一個簡單的技巧，即根據任務選擇表示。在訓練過程中使用可以并行化的卷積表示，在推理過程中，我們使用高效的循環表示:
聽起來有點奇幻，但是有人就是實現出來了，這個模型叫做Linear State-Space Layer (LSSL)
它結合了線性動態系統理論和神經網絡的概念，可以有效地捕獲數據中的時序信息和動態特征。LSSL 基于線性動態系統理論，這種系統可以用狀態空間模型表示。在這個模型中，系統的行為由狀態變量的演化和外部控制信號的影響決定。狀態變量是系統的內部表示，可以捕獲系統的動態特性。
這些表示都有一個重要的特性，即線性時不變性(LTI)。LTI表示ssm參數A、B和C對于所有時間步長都是固定的。這意味著對于SSM生成的每個令牌，矩陣A、B和C都是相同的。
也就是說無論給SSM的序列是什么，A、B和C的值都保持不變。這樣就得到了一個不感知內容的靜態表示。但是靜態表示沒有任何意義對吧，所以Mamba解決的就是這個問題，但是在介紹Mamba之前，我們還有一個知識點需要強調，那就是矩陣A
因為SSM公式中最重要的就是矩陣a。正如我們之前在循環表示中看到的那樣，它捕獲了關于前一個狀態的信息來構建新狀態，如果矩陣a如果跟RNN一樣會遺忘掉非常靠前的信息那么SMM將沒有任何的意義，對吧。
如何保留大上下文大小的方式創建矩陣A呢?
HiPPO 的模型結合了遞歸記憶（Recurrent Memory）和最優多項式投影（Optimal Polynomial Projections）的概念，這種投影技術可以顯著改善遞歸記憶的性能，特別是在處理長序列和長期依賴關系時。
https://proceedings.neurips.cc/paper/2020/hash/102f0bb6efb3a6128a3c750dd16729be-Abstract.html
使用矩陣A來構建一個狀態表示，該狀態表示可以很好地捕獲最近的令牌并衰減較舊的令牌。其公式可表示為:
具體的詳細內容我們就不介紹了，有興趣的查看原論文。
這樣我們就基本上解決了所有的問題：1、狀態空間模型；2、處理遠程依賴關系；3、離散化和并行計算
如果想深入了解有關如何計算HiPPO矩陣和自己構建S4模型建議您閱讀注釋的S4。
https://srush.github.io/annotated-s4/
上面介紹完所有必要的基礎知識，最后就是我們的重點了
Mamba 有兩個主要貢獻:
1、選擇性掃描算法，模型可以過濾有關和無關的信息
2、硬件感知算法，通過并行掃描、核融合和重計算有效地存儲(中間)結果。
在探討這兩個主要貢獻之前，我們先看看一下為什么它們是必要的。
狀態空間模型，S4(Structured State Space Model)，在語言建模和生成中的某些任務上表現不佳
比如在選擇性復制任務中，SSM的目標是按順序復制輸入和輸出的部分:
(循環/卷積)SSM在這個任務中表現不佳，因為它是線性時不變的。對于SSM生成的每個令牌，矩陣A、B和C都是相同的。
因為它將每個令牌平等地視為固定的a、B和C矩陣的結果，所以SSM不能執行內容感知推理
SSM表現不佳的第二個任務是重現輸入中發現的模式:
我們的提示在“教”模型在每個“Q:”之后提供“A:”響應。但是由于ssm是時間不變的，它不能選擇從其歷史中獲取先前的令牌。
以矩陣B為例不管輸入x是什么，矩陣B保持完全相同，并且與x無關:
同理無論輸入是什么，A和C也不變，這就是我們上面說的靜態。
而Transformers 可以根據輸入序列動態地改變注意力。可以選擇性地“看”或“注意”序列的不同部分，再加上位置編碼，這使得Transformers對于這種任務非常的簡單。
ssm在這些任務上的糟糕性能說明了定常ssm的潛在問題，矩陣A、B和C的靜態特性導致了內容感知問題。
SSM的循環表示創建了一個非常有效的小狀態，因為它壓縮了整個歷史信息，所以與不壓縮歷史(注意力矩陣)的Transformer模型相比，它的功能要弱得多。
Mamba 的目標是獲得Transformer一樣強大的“小”狀態
通過有選擇地將數據壓縮到狀態，當輸入一個句子時，通常會有一些信息，比如停頓詞，這些信息沒有太多的意義。
我們先看看SSM在訓練期時的輸入和輸出維度:
在結構化狀態空間模型(S4)中，矩陣a、B和C獨立于輸入，因為它們的維度N和D是靜態的，不會改變。
而Mamba通過結合輸入的序列長度和批量大小，使矩陣B和C，甚至步長?依賴于輸入:
這意味著對于每個輸入標記，有不同的B和C矩陣，這解決了內容感知問題!這里矩陣A保持不變，因為希望狀態本身保持靜態，但影響它的方式(通過B和C)是動態的。
也就是說它們一起選擇性地選擇將什么保留在隱藏狀態中，什么需要忽略，這都是由輸入確定的。
較小的步長?導致忽略特定的單詞，而是更多地使用之前的上下文，而較大的步長?則更多地關注輸入單詞而不是上下文:
這些矩陣現在是動態的了所以它們不能使用卷積表示來計算，只能使用循環進行處理，這就使得無法進行并行化。
為了實現并行化，我們先看看循環的輸出:
每個狀態都是前一個狀態(乘以A)加上當前輸入(乘以B)的和。這被稱為掃描操作，可以很容易地通過for循環計算出來。但是并行化似乎是不可能的，因為每個狀態只有在我們有前一個狀態時才能計算出來。
但是Mamba使用并行掃描算，通過關聯屬性假定執行操作的順序無關緊要。這樣就可以計算部分序列并迭代組合它們:
這樣還有一個好處是因為順序不重要，也可以省略掉Transformer的位置編碼。
最近gpu的一個缺點是它們在小但高效的SRAM和大但效率稍低的DRAM之間的傳輸(IO)速度有限。在SRAM和DRAM之間頻繁地復制信息成為瓶頸。
Mamba的DRAM和SRAM分配的具體實例如下:
中間狀態不被保存，但對于反向傳播計算梯度是必要的。作者重新計算了反向傳遞過程中的中間狀態。盡管這看起來效率很低，但它比從相對較慢的DRAM讀取所有這些中間狀態的成本要低得多。
這里我們就不詳細說明了，因為這部分我也沒太研究過
Mamba 塊
選擇性SSM可以作為一個塊，就像在Transformer中的的注意力模塊一樣。我們可以堆疊多個塊，并使用它們的輸出作為下一個曼巴塊的輸入:
最后一個端到端（輸入到輸出）的例子包含了歸一化層和選擇輸出標記softmax。
這樣就得到了快速的推理和訓練，而且是“無限”長度上下文的模型
看完這篇文章，我希望你能對Mamba 和狀態空間模型有一定的了解，最后我們以作者的發現為結尾：
作者發現模型與相同尺寸的Transformer模型的性能相當，有時甚至超過了它們!
本文章轉載微信公眾號@算法進階