不分割成token，直接從字節中高效學習，Mamba原來還能這樣用
在定義語言模型時，通常會使用一種基本分詞方法，把句子分為詞（word）、子詞（subword）或字符（character）。其中，子詞分詞法一直是最受歡迎的選擇，因為它在訓練效率和處理詞匯表外單詞的能力之間實現了自然的折中。然而，一些研究指出了子詞分詞法的問題，如對錯別字、拼寫和大小寫變化以及形態變化缺乏穩健性。
因此，有些研究人員另辟蹊徑，采用了一種使用字節序列的方法，即從原始數據到預測的端到端映射，中間不進行任何分詞。與子詞模型相比，基于字節級的語言模型能夠更容易地在不同的書寫形式和形態變化之間進行泛化。當然，將文本建模為字節意味著生成的序列要比對應的子詞長得多。如此一來，效率的提升就要依靠架構的改進來實現了。
自回歸 Transformer 在語言建模中占主導地位，但效率問題尤為突出：計算成本隨序列長度呈二次方增長，因此對長（字節）序列的擴展能力很差。研究人員壓縮了 Transformer 的內部表示，以便處理長序列，例如開發了長度感知建模方法，在這種方法中，token 組在中間層內合并。最近，Yu 等人 [2023] 提出了 MegaByte Transformer，它使用固定大小的字節片段作為子詞的模擬壓縮形式。因此，MegaByte 可以降低計算成本。不過，這可能還不是最好的方法。
在一份新論文中，來自康奈爾大學的研究者介紹了一種高效、簡單的字節級語言模型 MambaByte。該模型對最近推出的 Mamba 架構進行了直接改造。Mamba 建立在狀態空間模型（SSM）開創的方法基礎上，引入了對文本等離散數據更有效的選擇機制，并提供了高效的 GPU 實現。作者的簡單觀察結果是，使用 Mamba（不做修改）可以緩解語言建模中的主要計算瓶頸，從而消除 patching 并有效利用可用的計算資源。
- 論文標題：MambaByte: Token-free Selective State Space Model
- 論文鏈接：https://arxiv.org/pdf/2401.13660.pdf
他們在實驗中將 MambaByte 與 Transformers、SSM 和 MegaByte（patching）架構進行了比較，這些架構都是在固定參數和固定計算設置下，并在多個長篇文本數據集上進行比較的。圖 1 總結了他們的主要發現。
與字節級 Transformers 相比，MambaByte 能更快地實現更好的性能，計算效率也明顯更高。作者還考慮了無 token 語言模型與現有最先進的子詞模型相比的可行性。在這方面，他們發現 MambaByte 與各種子詞基線模型相比具有競爭力，但它能處理更長的序列。研究結果表明，MambaByte 是現有依賴分詞器（ tokenizer）的模型的有力替代品，有望用來促進端到端學習。
背景：選擇性狀態空間序列模型
SSM 通過一階微分方程對隱藏狀態的跨時間演變進行建模。線性時不變（time-invariant） SSM 在幾種模態的深度學習中顯示出了良好的效果。然而，Mamba 作者 Gu 和 Dao 最近認為，這些方法的恒定動態缺乏隱藏狀態中依賴輸入的上下文選擇，而這可能是語言建模等任務所必需的。為此，他們提出了 Mamba，該方法將給定輸入 x (t) ∈ R、隱藏狀態 h (t) ∈ R^n 和輸出 y (t) ∈ R 在時間 t 的時變連續狀態動態定義為：
其參數為對角時不變系統矩陣 A∈R^(n×n)，以及隨時間變化的輸入和輸出矩陣 B (t)∈R^(n×1) 和 C (t)∈R^(1×n)。
要對字節等離散時間序列建模，必須通過離散化來逼近 (1) 中的連續時間動態。這就產生了離散時間隱態 recurrence，每個時間步都有新矩陣 A、B 和 C，即
請注意，(2) 類似于循環神經網絡的線性版本，可以在語言模型生成過程中以這種循環形式應用。離散化要求每個輸入位置都有一個時間步，即 ?[k]，對應于 
的 x [k] = x (t_k)。然后就可以根據 ?[k] 計算出離散時間矩陣 A、B 和 C。圖 2 展示了 Mamba 如何為離散序列建模。
在 Mamba 中，SSM 項是輸入選擇性的，即 B、C 和 ? 被定義為輸入 x [k]∈R^d 的函數：
其中 W_B ∈ R^(n×d)（C 的定義類似），W_? ∈ R^(d×r) 和 W_R ∈ R^(r×d)（對于某個 r ?d）是可學習的權重，而 softplus 則確保正向性。請注意，對于每個輸入維度 d，SSM 參數 A、B 和 C 都是相同的，但時間步數 ? 是不同的；這導致每個時間步數 k 的隱藏狀態大小為 n × d。
Mamba 將這個 SSM 層嵌入到一個完整的神經網絡語言模型中。具體來說，該模型采用了一系列門控層，其靈感來源于之前的門控 SSM。圖 3 顯示了將 SSM 層與門控神經網絡相結合的 Mamba 架構。
線性 recurrence 的并行掃描。在訓練時，作者可以訪問整個序列 x，從而更高效地計算線性 recurrence。Smith et al. [2023] 的研究證明，使用工作效率高的并行掃描可以高效計算線性 SSM 中的順序 recurrence。對于 Mamba，作者首先將 recurrence 映射到 L 個元組序列，其中 e_k =
，然后定義一個關聯算子 
使得 
。最后，他們應用并行掃描計算序列 
。一般來說，這需要 
時間，使用 L/2 個處理器，其中 
是矩陣乘法的成本。注意，A 是一個對角矩陣，線性 recurrence 可在 
時間和 O (nL) 空間內并行計算。使用對角矩陣進行并行掃描的運行效率也很高，只需 O (nL) FLOPs。
實驗結果
表 2 顯示了每個數據集的每字節比特數（BPB）。在本實驗中，MegaByte758M+262M 和 MambaByte 模型使用相同的每字節 FLOP 數（見表 1）。作者發現，在所有數據集上，MambaByte 的性能始終優于 MegaByte。此外，作者注意到，由于資金限制，他們無法對 MambaByte 進行完整的 80B 字節訓練，但 MambaByte 在計算量和訓練數據減少 63% 的情況下仍優于 MegaByte。此外，MambaByte-353M 還優于字節級 Transformer 和 PerceiverAR。
在如此少的訓練步驟中，MambaByte 為什么比一個大得多的模型表現得更好？圖 1 通過觀察參數數量相同的模型進一步探討了這種關系。圖中顯示，對于參數大小相同的 MegaByte 模型，輸入 patching 較少的模型表現更好，但在計算歸一化后，它們的表現類似。事實上，全長的 Transformer 雖然在絕對意義上速度較慢，但在計算歸一化后，其性能也與 MegaByte 相似。相比之下，改用 Mamba 架構可以顯著提高計算使用率和模型性能。
根據這些發現，表 3 比較了這些模型在 PG19 數據集上的較大版本。在這個實驗中，作者將 MambaByte-972M 與 MegaByte-1.3B+350M 和其他字節級模型以及幾個 SOTA 子詞模型進行了比較。他們發現，MambaByte-972M 即使只訓練了 150B 字節，其性能也優于所有字節級模型，并與子詞模型相比具有競爭力。
文本生成。Transformer 模型中的自回歸推理需要緩存整個上下文，這會大大影響生成速度。MambaByte 不存在這一瓶頸，因為它每層只保留一個隨時間變化的隱藏狀態，因此每生成一步的時間是恒定的。表 4 比較了 MambaByte-972M 和 MambaByte-1.6B 與 MegaByte-1.3B+350M 在 A100 80GB PCIe GPU 上的文本生成速度。雖然 MegaByte 通過 patching 大大降低了生成成本，但他們觀察到 MambaByte 由于使用了循環生成，在參數相似設置下速度達到了前者的 2.6 倍。
