超越 Transformer 與 Mamba 之爭:一個統一模型架構的新起點

長期以來,AI 模型架構的發展彷彿一場路線之爭:究竟是選擇 Transformer 強大的表達能力,還是擁抱 Mamba 等狀態空間模型(SSM)的線性效率?一篇突破性的研究論文揭示,這兩者並非對立,而是一個更深層結構的兩種表現形式。這項發現不僅終結了長期的辯論,更為下一代 AI 系統的設計開闢了全新的可能性。

超越 Transformer 與 Mamba 之爭:一個統一模型架構的新起點

過去幾年,大型語言模型的發展路線,似乎分裂成兩條截然不同的道路。一邊是我們所熟知的 Transformer 架構,以其強大的注意力機制(Attention)稱霸武林,幾乎是所有頂尖大型語言模型的基石;另一邊則是 Mamba 所代表的狀態空間模型(State Space Models, SSM),憑藉其處理長序列時的線性時間複雜度與卓越效率,成為一股不可忽視的新興力量,挑戰著既有格局。

這場「路線之爭」的核心,在於一個難以兩全的權衡取捨。Transformer 的成功源於其二次方複雜度的注意力機制,它允許模型中的每個詞元(token)都能直接與其他所有詞元互動,從而捕捉複雜的語義關聯,但這也使其在處理超長文本時,運算與記憶體成本變得極為高昂。SSM 則像是一種對傳統方法的「回歸」,其結構類似於循環神經網路(RNN),以線性複雜度處理序列,雖然在速度與記憶體效率上大獲全勝,卻也讓許多人懷疑它是否能真正企及 Transformer 那樣強大的表達能力。

然而,來自史丹佛大學與 Together AI 的研究者 Tri Dao 等人發表的一篇論文,徹底顛覆了這個「二選一」的傳統思維。他們的研究指出,Transformer 與 SSM 並非對立的架構,而是在數學上具有「對偶性」(duality)的同一類結構。這項發現不僅為長久以來的爭論劃下句點,更重要的是,它為我們指引了一個統一兩者的理論基礎,也催生了新一代更高效的模型架構。

回顧兩種架構的權衡取捨

在深入探討這項新發現之前,讓我們先清晰地回顧一下兩種架構各自的優劣。這場看似路線之爭的背後,其實是模型設計者在「能力」與「效率」之間不斷尋求平衡的結果。

Transformer:全局視野的代價

Transformer 架構的核心是自注意力機制(Self-Attention)。它的優點在於能建立全局依賴關係,模型在處理任何一個詞元時,都能「看見」並權衡輸入序列中的所有其他詞元。這種強大的上下文建模能力,是它在各種自然語言處理任務上取得巨大成功的原因。但缺點也同樣顯著,其運算量與記憶體需求會隨著序列長度(N)呈二次方(O(N²))增長,這成為了它擴展到更長上下文時的根本瓶頸。

SSM:線性效率的潛力

以 Mamba 為代表的狀態空間模型(SSM),採用了一種更接近傳統訊號處理與 RNN 的方式。它維護一個隱藏狀態(hidden state),並依序處理輸入序列。優點是其運算與記憶體複雜度都是線性的(O(N)),在處理長達數百萬詞元的序列時,依然能保持極高的效率。過去大家對它的疑慮在於,這種循序漸進的處理方式,是否會限制其捕捉長距離依賴的能力,從而弱於 Transformer。

過去,我們傾向於認為必須在這兩者之間做出選擇。要麼為了極致的表達能力而忍受二次方的成本,要麼為了效率而接受一種看似不同的、能力或許有所折衷的架構設計。但新的研究告訴我們,這個前提假設本身可能就是錯的。

SSD 框架:從對偶性看見統一的數學基礎

這篇論文的核心貢獻,是提出了一個名為「結構化狀態空間對偶性」(Structured State Space Duality, SSD)的理論框架。這個框架巧妙地揭示了 Transformer 與 SSM 之間隱藏的數學聯繫。

簡單來說,研究者發現這兩種模型的核心計算,都可以被抽象為對某種「結構化矩陣」(structured matrix)的操作。而 Transformer 的注意力機制和 SSM 的循環計算,只是計算這個矩陣與輸入序列交互作用的兩種「不同方法」而已。

這意味著,注意力機制並非唯一解。它只是實現某種深層運算的一種「演算法」,而 SSM 提供了另一種更高效的演算法來達成相同的目標。

我們可以這樣理解:想像一下,你想從 A 點移動到 B 點。Transformer 的方法像是搭乘一架能超光速飛行的飛機,它可以瞬間讓你看到整張地圖上所有點與 B 點的關係,但這架飛機的燃料消耗極其驚人。SSM 的方法則像是駕駛一輛高速列車,它沿著一條精心設計的軌道前進,每到一站都會更新自己的狀態資訊,最終高效地抵達 B 點。過去我們以為這是兩種完全不同的交通工具,但 SSD 框架告訴我們,它們其實都是在解決同一個「位移問題」,只是遵循了不同的「物理定律」或「計算路徑」。

這個發現的意義是深遠的。它將模型架構的討論,從「選擇哪種架構」提升到了「如何更有效地計算這個底層的結構化矩陣」。這不再是一個「站隊」問題,而是一個純粹的演算法優化問題。

Mamba-2 的誕生:理論統一的實踐價值

一個優雅的理論若不能轉化為實際的效益,終究只是紙上談兵。幸運的是,基於 SSD 框架,研究團隊直接推出了一個名為 Mamba-2 的新模型,完美展示了理論統一的實踐價值。

Mamba-2 並不是對 Mamba 的簡單改良,而是根據 SSD 理論重新設計的產物。它在設計上融合了對偶性的理解,採用了更高效的演算法來實現 SSM 的核心計算。結果令人振奮,證明了理論的強大指導力:

Mamba-2 的突破性表現

首先,Mamba-2 實現了速度大幅提升。相較於已經非常高效的 Mamba,Mamba-2 在硬體上的運行速度又提升了 2 到 8 倍。更重要的是,它在性能上毫不妥協,在保持甚至超越頂級 Transformer 模型性能的同時,實現了這樣的速度飛躍。

Mamba-2 的成功證明,SSD 框架不僅僅是一個數學上的漂亮理論,它直接指導了如何建構出速度更快、效率更高且能力同樣強大的新一代模型。它讓我們看到,未來的模型架構,或許不再需要被貼上「Transformer」或「SSM」的標籤,而是可以被視為在同一個統一框架下,針對不同硬體和任務需求,選擇最優計算策略的靈活體系,實現真正的「架構即演算法」。

對我而言,這項研究的重要性在於它為整個領域的發展提供了一個更具普遍性的抽象層次。我們終於可以停止爭論哪條路是「唯一正確」的,轉而專注於一個更根本的問題:如何設計出能統一速度、記憶體與表達能力的模型架構。這不僅為解決 AI 的效率瓶頸開闢了新途徑,也預示著一個模型設計更加靈活、更加第一性原理的時代即將到來。

延伸閱讀

我是江中喬,一位具有 TPM 與產品管理背景的 AI 系統建構者,目前專注於 AI 認知增強系統與多 Agent 協作架構的設計與實踐。