在大模型領域，一直穩站 C 位的 Transformer 最近似乎有被超越的趨勢。

這個挑戰者就是一項名為「Mamba」的研究，其在語言、音頻和基因組學等多種模態中都達到了 SOTA 性能。在語言建模方面，無論是預訓練還是下游評估，Mamba-3B 模型都優于同等規模的 Transformer 模型，并能與兩倍于其規模的 Transformer 模型相媲美。

論文一經發表，引起了不小的轟動。驚嘆之余，大家發現論文作者只有兩位，一位是卡內基梅隆大學機器學習系助理教授 Albert Gu，另一位是 Together.AI 首席科學家、普林斯頓大學計算機科學助理教授（即將上任）Tri Dao。

這項研究的一個重要創新是引入了一個名為「選擇性 SSM（ selective state space model）」的架構，相比于 Transformer 中的自注意力機制的計算量會隨著上下文長度的增加呈平方級增長，比如上下文增加 32 倍時，計算量可能會增長 1000 倍，Mamba 可以隨上下文長度的增加實現線性擴展，其性能在實際數據中可提高到百萬 token 長度序列，并實現 5 倍的推理吞吐量提升。而這些，都離不開選擇性 SSM。

看到 Mamba 如此優異的表現后，很多研究者對 SSM（state space model） 相關研究產生了好奇。

在最近的一次采訪中，來自艾倫人工智能研究所（ AI2）的機器學習研究員 Nathan Lambert 與 Mamba 論文作者之一 Tri Dao，以及同樣來自 Together.AI 的科學家 Michael Poli 進行了深入交流。

他們主要討論了 LLM 架構的未來，此外，這三位研究者還探討了狀態空間模型（SSM）在新興的 LLM 市場中的應用前景。對話中涉及的知識點也比較密集，比如為什么 Transformer 中的注意力機制有效、其擴展限制是什么、Mamba 介紹及其硬件優化以及對未來架構預測的討論等。

以下是對話內容。

為什么注意力機制有效

Nathan Lambert：我們首先來討論一下，為什么注意力機制有效，以及注意力機制的局限性在哪里。Transformer 有多少部分是基于注意力機制構建的，還有沒有其他的機制在起作用，在這方面可能會遇到哪些挑戰？

Tri Dao: 沒錯，所謂 Transformer，就是目前推動我們看到的大多數激動人心的應用的架構。如你所說，注意力機制是其中的核心層。其實，注意力機制早在 2014 年到 2015 年就受到關注了，隨后 Transformer 的概念出現，融合了注意力機制，注重多層感知機（MLP）和注意力機制的交織使用。

我認為它的成功很大程度上在于，這些模型似乎能夠很好地進行擴展，你可以通過增加更多的參數和數據來擴大模型規模。這就是成功的秘訣。雖然現在看起來顯而易見，但我認為五年前這還不是一個明確的概念。

Transformer 成功的幾個原因包括：首先，它足夠通用，能夠從大量數據中學到很多東西。其次，它對硬件非常友好。與之前的循環神經網絡（RNN）不同，它沒有順序依賴性。

因此，它可以在 GPU、TPU 上運行得很好，可以進行擴展，同時非常高效地利用硬件。我個人也致力于讓它更加高效地利用硬件。所以，這就是成功的秘訣 —— 作一個既通用又能很好地進行擴展的架構。如果你是研究 NLP 的，也許你會考慮加入一些歸納偏置來增強模型。就我個人而言，我認為 Transformer 是一個非常通用的架構，可擴展性很好，而且對硬件很友好。

Nathan Lambert：是的，是的。現在回想起來，這一切似乎都顯得很明顯。現在，在研究它的替代方案時，一個有趣的維度是上下文長度。Michael，你怎么看？

Michael Poli：是的，我有幾點想說。首先，目前仍然有很多優秀的研究試圖從第一性原理出發解釋 Transformer，為什么它能學習這些有趣的回路？人們會拆解計算過程，比如不同的 transformer 中的 head 組合等。

有一些工作會將 Transformer 作為一種編碼的編程語言來理解。但我認為，正如 Trey 提到的，Transformer 中有一些非常有趣的設計選擇。注意力和 MLP 的交織使用是相當重要的。而且，Transformer 一開始之所以成功，是因為它采用了一些已經為 RNN 和其他傳統 NLP 模型所開發的技術，比如使用門控機制來調節模型吸收哪些信息，以及決定某些內容在這種并行形式中被忘記的速度。這就像是一些可以在 GPU 上進行優化的寶石，雖然不是很容易，但確實可以優化。

注意力計算量呈平方級增長

Nathan Lambert：是的，這些都很棒。我想說的比較具體的一點是，注意力機制最終呈現出一種隨輸入序列長度成二次方增長的計算成本。假設你有一個長度為 L 的輸入序列，并且你想輸出一個長度也為 L 的序列。如果你深入到數學細節中去，觀察大多數庫在推理時發生的情況，你會發現有這樣一個上三角注意力矩陣，其中，你只能考慮文本的過去部分。隨著處理的進行，你會發現它形成了一個 L 平方的關系，其中第一個 token 只考慮一個元素，然后每個后續的 token 需要考慮的過去 token 逐漸增多。剛才我們已經討論了 RNN，以及一些非注意力的方法是如何做到不去查看序列中所有文本歷史的。當你給聊天機器人 GPT 寫一個長長的提示時，你真的希望所有那些信息都被編碼進去嗎？除了這種密集的注意力矩陣，我們還有什么其他的選擇？  

Tri Dao：循環神經網絡可以追溯到 80 年代，也許其中一些更著名的是長短期記憶網絡（LSTM）、門控循環單元（GRU）。大約在 2012 年到 2016 年左右，它們在翻譯、語音識別等方面非常流行，當時它們是 NLP 領域的 SOTA 技術。

它們以一種序列方式處理文本：逐個觀察 token，然后改變隱藏狀態，每次看到一個新 token 時都會更新隱藏狀態。我認為從某種意義上來說，這模仿了人腦處理信息的方式，就像你閱讀一句話或一段話，就像你在大腦中存儲一些信息。當你讀完一個文檔時，你可能能夠回答關于那個文檔的問題，而無需再次參考該文檔。所以，RNN 就是這樣工作的。它們處理文本，然后改變隱藏狀態，隱藏狀態是可以用來生成新 token 或對文檔進行分類的表示。

大約在 2016 年左右，這些方法效果很好。但是，它們逐漸不受青睞，從實驗結果來看，它們的表現并不像 Transformer 那么好。正如你提到的，由于 Transformer 存在這種二次方擴展特性，每個 token 都會與之前的各個 token 進行比較，這為信息傳播提供了一種非常簡單的方式。我認為這是 Transformer 和注意力機制表現出色的部分原因。

但最近，一些新的 RNN 架構似乎表現得相當不錯，RWKV 是其中較早的一個。我非常欽佩那個項目，那是一位名叫 Bo Peng 的研究者主導研發的。它似乎在以一己之力對抗 Transformer 統治地位，證明了 RNN 可以非常強大。

Nathan Lambert：是的。我之前也讀過這篇論文。在技術層面上，他們嘗試通過兩個線性 RNN 來復制類似于注意力機制中的查詢鍵值查找的東西，本質上是為了消除像特定注意力擴展這樣的潛在問題。這兩個 RNN 擁有更好的長上下文行為和不同的實現規則。他們還訓練了參數量高達 140 億的模型。這也引出了我接下來想問的一些問題，包括 Mamba 和 Striped Hyena。我們可以一個一個來聊。

Striped Hyena 是個什么模型？

Nathan Lambert：我進入了 Together API，做了一次 Mistral 和 Striped Hyena 的對比測試。結果顯示，Striped Hyena 是一個不錯的語言模型。它能回答大多數問題，沒有明顯的失敗模式。Michael，你怎么看待這個模型？

Michael Poli：首先我想說，這些新方法之間有一個有趣的聯系。存在一種凸集，它有一個中心點，線性注意力（即沒有 softmax 的注意力）、線性 RNN 以及基于狀態的模型（SSM）之間的關聯都在這個凸集里。在某種程度上，這種基礎模型的數學表述是相同的，我這里不是指基礎架構，而是指基礎模型。

然后你可以朝不同的方向發展，每個方向都有自己的權衡，比如特征映射方向、 kernel 方向。所以當你分解或去除掉 softmax 時，你可以在處理查詢（queries）和鍵（keys）時采取不同的方法。這些查詢和鍵是構成你的注意力矩陣的基本實體。在去除 softmax 之后，你可以構建其他類似 kernel 的函數，或者其他你希望能夠近似實現注意力機制功能的函數。

你可以做一些類似泰勒近似（Taylor approximation）或泰勒展開（Taylor expansion）這樣的事情。你會得到一個略顯不同的視角，但你得到的東西又非常相似。你可以轉向 Time variance。這意味著你修改 RNN，使其計算更多地依賴于輸入序列。也就是說，在線性 RNN 中的計算是由輸入序列所決定的。你可以使用類似門控（gates）的東西，我們已經看到了很多工作，例如，用額外的門控來更新內部 tension，讓你更好地利用你的固定狀態維度。第三個方向 —— 至少在我看來 —— 是使用卷積形式的，更多地使用其他類型的線性運算符，這些運算符仍然是可結合的，仍然允許你并行訓練。

所以這里的內容包括時間不變系統。我可以詳細解釋這些點，但是存在一些可以在卷積和循環之間切換的模型，這些模型還配備了額外的門控機制。我參與的一個項目，它誕生于我剛才提到的這第三類架構。我們真正試圖做到的是，打造出每浮點運算性能最佳的架構。我們反復驗證的一個原則是，似乎將不同層、不同類別的模塊甚至是全注意力層結合起來，得到的東西比單獨的組件要好。

因此，我們正在試圖更深入地理解這些模型的組合方面。這種理解有助于我們創造每浮點運算性能更優的預訓練模型。使用這種模型，我們運行了一整套的規模擴展法則。混合化也給了我們一些優勢，因為我們想要的是一種可以即開即用的東西，它使得過程變得更加簡單。

在為更長的上下文進行微調時，我們可以應用一些為 Transformers 開發的技術，而且有點出人意料的是，這些技術對于混合體也同樣有效。比如，用于旋轉嵌入的線性縮放等等，你可以深入了解細節。所以這主要是一個嘗試性項目，試圖弄清楚，在當前的環境下，我們最好能做到什么程度。

什么是 Mamba

Nathan Lambert：Striped Hyena 使用了一組新的模型移植技術（model grafting）進行了優化，使我們能夠在訓練期間更改模型架構，對我來說，感覺好像有很多事情正在發生，比如一些你可能不能談論太多的內容如數據。

所以，我不認為所有數據都得到了很好的解釋，比如較長的上下文數據是什么，你能從模型的角度解釋一下嗎？我認為即使只是總結一下，對人們來說也是很好的體驗。

該領域有很多很酷的工作，因此，AI 領域有很多新的項目在展開，例如，有些人試圖將 Lama 模型拆解并繼續訓練。事實上，這種做法有點狂野，人們嘗試采用強大的模型并試圖讓它們變得更小，同時還想獲得媲美大模型的性能優勢。

雖然這是個題外話，但我沒有預料到的是當你關注社交媒體的時候，你會看到人們會說，哦，最終還是狀態非注意力模型贏了。在我看來，這種說法掩蓋了很多有趣的細節。

好吧，讓我們回到 Mamba。如果沒記錯的話，我認為 Mamba 套件中最大的模型是 2. 8 億參數，在 NLP 基準測試中，包括 GPT J、Pythia 模型套件給出的基準分數都非常強。

Tri Dao：Mamba 是我與 Albert Gu 合作完成的，當時他是斯坦福大學的博士生，我們就是在那里認識的，他現在是 CMU 的助理教授。所以這是一次美妙的合作，Mamba 的成功要歸功于他。Albert 一直致力于狀態空間模型的研究，從某種意義上說，正如前面提到的，線性張量、線性 RNN、卷積、神經網絡等領域他都有涉及。

在我過去參與的幾個項目中，我還致力于空間、狀態空間的研究，我研究的角度是如何使狀態空間更加硬件高效，并提高其性能。所以和 Albert Gu 一起工作真是太棒了。我認為參與 Mamba 的研究過程更多的是一個概念證明，狀態空間實際上可以像 NLP 領域的 transforme 一樣好嗎？因此，才有了 Mamba，這項研究表明狀態空間在音頻方面可能會更好。但是，對于狀態空間模型來說，語言一直是最難獲得、最難做好的。

而且，語言也是人們現在最關心的事情，所以我做的更多的是一個概念驗證，也就是說，我們想表明狀態空間模型也可以具有競爭力，甚至可以和 Transformer 一較高下。我們實驗中驗證的 token 數從 3B 到 300 B 不等。

所以從絕對意義上講，這些都不是非常強大的模型，這些還不是我們真正想要的模型。我認為我們正在做的更多的是進行學術比較。比如訓練相同數量的 token，狀態空間模型可能比 transformer 稍微好一些。

這件事對我們來說尤其令人興奮，我認為 Albert 已經推動這件事有一段時間了。

結果是，我們的研究在推理上可能會更快，也許我們會有不同的方式來理解上下文學習是如何發生的。我很期待未來的工作。

Mamba 硬件優化

Nathan Lambert：您能否談一談實現這些新 CUDA 內核實際上需要做什么？

Tri Dao：關于狀態空間的研究，它在某種意義上是循環神經網絡。狀態大小是你在遍歷或處理序列時，用來存儲信息的緩沖區。

從某種意義上講，Transformer 也可以這樣理解，它保存的整個歷史記錄通常稱為 KV 緩存，Transformer 保留歷史并不斷引用它。對于 RNN，它們具有固定大小的狀態；對于 transformer，你可以認為狀態大小正在增加。而且，我們的直覺是，狀態規模越大，模型就表現越好。

所以基本上，你需要有更多的空間來存儲你需要記住的東西。因此，之前的模型（例如 S4 等）具有相當大的隱式狀態大小，它們使用卷積視圖來避免將狀態具體化。

我們希望在循環中加入更多的輸入依賴性，然而，這樣做阻止了我們使用可以提高效率的卷積視圖。

因此，我們必須找到一種不同的方法來提高效率，所以我們把重點專注于在 GPU 上提高效率。我們的想法是，我們要有一個大的狀態大小，但我們不需要用到實際的 GPU 內存，比如 HBM，我們可以把大的狀態保存在一個更快的內存中，稱之為 SRAM，你可以把它想象成一個緩存。如果你更熟悉 CPU，這通常是一個緩存和 RAM。

所以，如果你有較大的狀態，你可以將它保存在緩存中，這樣一來你不會遭受太大的損失。

2024 年架構預測

Nathan Lambert：我現在對 GPU 與 TPU 最有力的洞察是，MoE 在 TPU 中效果不佳，因為你必須在基本層上加入一些 MoE。

在分布式訓練中，前饋層可能最終分布在不同的 TPU 節點上，而 TPU 之間是通過鄰近節點進行通信的。因此，與 GPU 相比，TPU 在這方面會受到較大的影響。2024 年這個領域會發生什么？

Tri Dao：我認為 Transform 仍然是一個非常強大的架構，現在可以擴展到一萬億級別的參數，人們往往想要性能最好的模型，這些模型在硬件上運行效率最高，并且在軟件方面擁有最多的支持。

我最近有些新想法，比如狀態空間。我們已經看到，正如 Michael 提到的，混合這些組件似乎可以提高性能，我認為在 7B 大小的模型上得到了驗證，也許狀態空間模型可以在更大規模的模型上起作用。

當前大多數人關注的是基于 Lime 架構的數據和基礎設施建設，盡管現有的 Transformer 架構在生產環境中仍然非常強大且得到廣泛支持，但也存在一些邊緣領域，比如長上下文、音頻、基因組學等，研究這些領域中的替代架構會非常有趣。這些領域提出了一些有意義的科學問題，例如模型是否像人類一樣理解指令和直覺、是否能與量化方法協同工作等。

此外，即使現在人們還在使用 Transformer 架構，未來可能會融入更多的新思想和組件，例如增加更多的層和注意力機制等，盡管它們可能仍然被稱為 Transformer。

簡而言之，盡管目前的人工智能領域傾向于保守，專注于現代架構，但新的架構和思想正在逐漸涌現，這些新穎的觀點和方法可能會為人工智能的發展帶來新的動力和方向。

Michael Poli：是的，我百分之百同意 Tri Dao 的看法，注意力機制作為計算原語依然很重要。注意力機制作為一種高效便捷的方式，可以有效增加序列處理器的狀態容量。

狀態維度與序列長度之間存在權衡關系，當模型尺寸變大，也就是模型變得更寬時，將有效地引入更多的狀態和序列長度。這樣一來，一些邊際效應可能會消失，而某些權衡將會發生改變，尤其是對于那些非常大的模型，如 14B、30B 這類模型。

未來，架構設計將變得更加有趣和復雜，將會有更多的創新發生。無論是混合模型還是引入新的模塊，我們都將看到更多激動人心的創新。

對 AI 更多的預測

Nathan Lambert：混合專家模型（Mixture of Experts, MoE）和狀態空間模型在近期內呈現出一種流行的趨勢。

然而，在開源和學術界，目前還沒有人真正嘗試在混合專家模型上做早期的嘗試和改進。模型嫁接（Model Grafting）現在正在變得更加實用。

跟進這些進展非常有趣，希望這些發展能夠為學術界和科學家們提供更多影響行業對話的方法，特別是在目前行業更多關注于擴大模型規模的情況下。我建議開源公司應該在他們的語言模型中進行特定的改進，以獲得商業優勢。

你們在機器學習方面還關注其他什么嗎？它不一定是關于狀態空間模型的。明年最讓你們興奮的是什么？

Tri Dao：我個人認為數據仍然是最重要的因素。我們正在深入研究數據如何影響模型性能，例如通過一些與模型性能高度相關的合成任務來進行研究。這種方法一直是我們論文和研究工作中的主要動力和示例。未來一段時間我們將專注于數據領域。

雖然所有的架構工作都很有趣，使其在硬件上高效運行也很有趣，但最終還是關于數據的。如果你了解擴展率（scaling law），你就會知道不同的模型架構通常會有相同的斜率，只是偏移量不同，似乎唯一能改變斜率的是數據質量。

Michael Poli：是的，我們添加了數據。數據確實很有趣，比如在架構設計上的微型化，找出并分解語言建模等任務所涉及的各個方面，我們試圖將它們打包成可以用來迭代的東西，這非常令人興奮。

我個人對新應用，尤其是基因組學工作感到非常興奮，但更多地是從工程角度出發，我們正看到一個轉變。目前，語言仍是獲得最多點擊和最大興趣的領域，但我認為這將隨著時間的推移而改變。

Nathan Lambert：是的，每個人都在談論語言，但我覺得圖像、視頻將會是產生巨大價值的東西。我不知道語言的上限在哪里。我很興奮，我已經開始嘗試這些了，比如我會拿博客的文本，讓模型轉換成圖片，然后轉換成帶有音頻的視頻，所有這些都是用一個 Python 腳本完成的，這真的很容易做到。所以我同意你的觀點，超越語言的東西很有趣。

Tri Dao：在你的經驗中，當你將所有這些東西拼湊在一起時，它們確實能夠合理地工作嗎？

Nathan Lambert：它還不是那么完美，DALL?E 生成的圖片比較相似，但我的做法非常簡單，就是直接把文本拿過來，然后用一個系統提示，讓模型生成各種各樣的圖片，我認為我可以做得更好。據我知道，大概在一年內，就會有一個文本到視頻的 API，到時我會切換到 API，那種體驗會很棒。

Tri Dao：是的，我認為這些進步確實產生了大量的經濟價值，我們已經看到了這一點。現在很多公司都開始轉向使用這些技術。我認為這將改變我們的工作方式，正如你提到的，改變我們的工作方式和娛樂方式。所以這是一個非常令人興奮的未來。

原文鏈接：https://www.interconnects.ai/p/interviewing-tri-dao-and-michael?cnotallow=5d10d34c97637bebcfeba6470c0f0d9b