LLM未來架構：誰有可能動搖Transformer的統治地位？

在大模型領域，一直穩站 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具有二次方擴展特性，使得每個標記與先前的所有標記進行比較，這為資訊傳播提供了一種非常簡單的方式。我相信這正是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 更多的預測

Mixture of Experts (MoE) and State Space Models have recently emerged as a popular trend, according to Nathan Lambert.

然而，在開源和學術界，目前還沒有人真正嘗試在混合專家模型上做早期的嘗試和改進。模型嫁接（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

以上是LLM未來架構：誰有可能動搖Transformer的統治地位？的詳細內容。更多資訊請關注PHP中文網其他相關文章！