MiniMax豪賭線性注意力:專訪架構負責人鍾毅然

非主流技術路徑上的先驅者

您能簡單介紹一下自己嗎?

我是鍾毅然,MiniMax 的資深研究總監,主要負責網路架構的設計和多模態理解大型模型的開發。在 MiniMax,我的主要職責是領導 MiniMax-01 網路結構的設計。

之前,我曾在上海人工智慧實驗室擔任新架構探索組的 PI,專注於非 Transformer 架構的高效訓練建模方法以及視覺-音訊-語言多模態融合的研究。

您是什麼時候開始研究線性注意力的?為什麼選擇這條技術路線?

我大約在 2021 年 7 月開始研究線性注意力。這源於我在 2020 年攻讀博士學位時撰寫的一篇論文《可逆注意力》(Invertible Attention)。當時,可逆神經網路和注意力機制都非常流行,因此我們將它們結合起來進行研究。

後來,我們團隊的一些成員對數學產生了濃厚的興趣。像線性注意力這樣的高效序列建模方法需要紮實的數學基礎,並涉及大量的公式推導,這與團隊的興趣完美契合,因此我們選擇了這個方向。

當時線性注意力在業界的狀況如何?

非常非主流,很少有人研究。大多數研究人員都專注於 Transformer,它基本上已經成為 NLP 領域的主導力量。

我們認為,與其做一個 Transformer 研究的普通人,不如做一些不同的事情。

您如何評估線性注意力路線的技術潛力?

我們最初的動機很簡單:解決 Transformer 的二次方計算複雜度。我們測試了各種方法,包括稀疏 Transformer 和線性注意力。

我們發現稀疏 Transformer 確實有效,與 Transformer 相比,它提供了更快的速度和更低的記憶體使用率。然而,線性注意力的表現很差,而且速度也很慢。儘管如此,我們還是選擇了追求線性注意力。

原因之一是它在數學上的吸引力——我們相信它的性能應該更好。另一個原因是,我們認為稀疏注意力的上限是完全注意力,因此很難超越。另一方面,線性注意力有潛力超越它。

您能解釋一下什麼是線性注意力嗎?

線性注意力本質上是一種核技巧。在 Transformer 中,將 Q、K 和 V 矩陣相乘涉及不同的計算複雜度,具體取決於您是先乘 QK 還是先乘 KV,這是由於維度不同造成的。

先乘 KV 可以將計算複雜度降低到線性。然而,問題是 QK 相乘之後會進行 softmax 運算,softmax 運算不滿足交換律,因此無法輕易地拆分為先乘 KV。因此,線性注意力的第一步是移除 softmax。

但移除 softmax 會影響結果。後續的任務是在沒有 softmax 的情況下保持結果的一致性,這就是線性注意力的目標。

線性注意力、稀疏注意力和線性 RNN 架構之間的基本區別是什麼?

稀疏注意力本質上仍然是一種 softmax 注意力。它只是計算比密集注意力矩陣更少的點。例如,滑動視窗注意力僅計算視窗內的注意力分數,通過減少計算量來實現加速。

線性 RNN 和線性注意力本質上是同一件事,只是有些人稱之為 RNN,有些人稱之為注意力。

一切都可以寫成 RNN 形式。例如,lightning attention 對應於 RWKV-4,而 RWKV-7 是 gating delta net 的改進版本。儘管它們本質上相似,但它們的實現細節有所不同。

線性注意力機制研究的關鍵里程碑是什麼?

大約在 2018-19 年,研究表明可以使用核技巧來降低 Transformer softmax 注意力的理論計算複雜度,但結果很差,效率也很低。

在 2019-20 年,稀疏注意力佔據主導地位,像 Google 這樣的公司提出了許多稀疏注意力變體。後來,線性注意力開始出現,但它面臨著性能差和速度慢的挑戰。

研究人員主要採用兩種方法來改進:一種是近似 softmax 函數,使分佈符合 softmax;另一種是我們選擇的方法,使用完全不同的方法進行建模,而不關心近似 softmax。

我們於 2021 年 10 月發表了第一篇論文《COSFORMER: RETHINKING SOFTMAX IN ATTENTION》,該論文用餘弦函數替換了 softmax 運算,從而可以拆分計算。

在 2022 年上半年,我們發表了第二篇論文《The Devil in Linear Transformer》,該論文分析了線性注意力性能下降的原因並提供了解決方案。這是 lightning attention 的前身。

後來,我們還專門針對線性注意力和長卷積研究了位置編碼,發表了 TNN,《TOEPLITZ NEURAL NETWORK FOR SEQUENCE MODELING》,這是一種與 S4(Mamba 的前身)類似的方法。

最後,我們推出了 lightning attention,通過改進的衰減方法和網路結構,使其性能與 Transformer 相匹配。我們還使用了一種平鋪技術使其速度更快。

您對當前非 Transformer 架構的技術路線有何看法?

線性注意力實際上是一種非 Transformer 方法。目前,除了類似 RNN 的方法外,其他非 Transformer 架構正在衰落。

例如,像長卷積和大核卷積這樣的 CNN,感覺它們已經由於性能不佳而逐漸被淘汰,但它們在某些方面實際上非常強大,在序列建模中仍然具有一定的作用,例如異常檢測任務。

實際上只有三種非 Transformer 架構:線性注意力、長卷積和線性 RNN。

但實際上,這三者可以統一為一個,我們稱之為線性複雜度模型。我們寫了一篇文章,涵蓋了所有這三種。

lightning attention 與 Mamba 和 RWKV 之間的核心區別是什麼?

最核心的區別是 lightning attention 是最簡單的線性注意力。Mamba 和 RWKV 都使用數據依賴衰減,而 lightning attention 使用手工製作的衰減來提高速度。

儘管可學習衰減可以實現更好的結果,但它會犧牲速度。例如,RWKV-7 比 gating delta net 慢 10-15%,而 gating delta net 的速度約為 lightning attention 的一半。

RWKV 的建模效果確實比 lightning attention 更好,但它速度較慢,並且尚未解決檢索問題。

現在業界是否普遍認為線性注意力具有高且可行的上限?

不,如果這是共識,那麼每個人都會擴大規模線性注意力模型。而且現在也不是共識。如果是的話,每個人都會做線性,但正如你所看到的,情況並非如此。

但對我們來說,我們在 2023 年下半年就已經看到了這一點。當時,我問過很多人,也和很多人談過,他們提出的最常見的觀點是,他們知道線性注意力在小規模上確實有效,但他們覺得一旦擴大規模就會失敗。

當時,我想我會擴大規模讓大家看看。現在 MiniMax-01 已經推出,沒有人懷疑線性注意力在大規模上的能力。

從小型實驗到大規模實現

您認為線性注意力的上限可以超越完全注意力嗎?

我們現在可以看到混合架構比純 Transformer 更好。但純線性注意力最大的問題是檢索能力,這對學術界來說是一個難以解決的問題。

現有的方法雖然複雜且速度慢,但仍然無法完全解決它,這就是為什麼有必要轉向混合架構的原因。

您觀察到哪個節點讓您決定走出實驗室?

在 2023 年 5 月至 6 月,我們內部已經有了 lightning attention 2,這是世界上第一個比 Flash attention 更快的線性注意力實現。

我們認為它已經跨越了工業紅線,其技術成熟度非常高,可以擴大規模。

您如何定義這條工業紅線?

首先,效果優於 Transformer,其次,速度快於 Transformer。這使其具有取代 Transformer 的能力。當時我們在一個 15B 規模的密集模型上驗證了這一點。

在您走出實驗室的節點上,為什麼最終選擇與 MiniMax 合作?

實際上,當時我與一些大型公司進行了洽談。但最終,我還是與 MiniMax 實現了這一目標。

首先,cosformer 是我與 Junjie 合作撰寫的文章。我們有合作的基礎。Junjie 在 SenseTime 工作時是我的老闆。23 年底,Junjie 邀請我共進晚餐。他對這些前沿技術的可能性更有信心。我的理解是,他當時也在尋找技術突破。

當時,MiniMax 已經完成了 Moe 的研究,實際上下一步的技術突破點非常少。當時,lightning attention 已經發布,Mamba 也很流行,所以在他的眼中,這是一個可行的方向。

這與 MiniMax 的互動伴侶產品有關嗎?

沒有聯繫。顏俊傑更關心模型的上限以及如何進一步突破這個天花板。

在大眾眼中,線性注意力可能更多的是一種突破效率的方向,而不是突破天花板。

這裡的重點是,首先,每個製造商的計算能力是恆定的。模型加速得越快,它可以吃的數據就越多,產生的模型就越好。當計算能力恆定時,模型越快越好。

您是否觀察到數據已經達到峰值的情況?

還沒有,對吧?數據仍然處於持續擴展的階段,但可能不像 23 年那樣激進。

因為數據總是在增加,而且每天都有新數據產生。對於模型來說,它每天都有新數據需要處理。網際網路每天產生的數據量非常大。通過清理,我們仍然可以獲得新數據。

與人類發展這麼多年來存在的數據相比,數據增長率是否有所放緩?

實際上,不一定。看看中國五千年的歷史,只積累了那幾本書。但隨著網際網路的發展,數據量的增加是一條非常陡峭的曲線。網際網路之前產生的總體數據可能不如之後一年產生的數據多。

在擴大規模的過程中,lightning attention 面臨著哪些挑戰?

為了驗證其可擴展性,我們首先進行了縮放定律實驗,從小型模型逐漸擴展到 7B、9B,最終擴展到具有超過 400B 參數的模型。

而且我們從理論上證明了線性的容量大於 Transformer。

我們將容量定義為 RNN 當前狀態的大小。對於 Transformer,容量大小為 O(d),其中 d 是大小;對於線性注意力,容量大小為 d²/h。由於 d 遠大於 h,因此容量更大。

最後,我們還驗證了混合模型優於純 Transformer。

4M 長度的序列視窗是如何實現的?

對於 lightning,訓練長度可以是任意的。只要充分利用計算能力,訓練 8K、32K 或 128K 的速度是相同的,並且 TGS(每 GPU 每秒處理的 token 數)是相同的。

因為 Transformer 是 n² 計算複雜度,序列越長,計算複雜度增長得越快,並且延遲以二次曲線增長。在 1M 長度時,softmax 注意力的延遲是 lightning attention 的 2,700 倍。

未來要實現無限上下文視窗,還需要解決哪些技術挑戰?

在我們目前的混合架構中,仍然有 1/8 的 softmax 注意力。這在 1M 長度時是一個瓶頸。這 1/8 帶來的延遲遠高於其餘 7/8 的線性注意力。

如果我們想優化長文本,我們必須考慮優化 softmax 注意力部分。我們可以從稀疏注意力方法中學習,使其更快更輕。

此外,我們還在考慮使 softmax 和線性注意力的混合比例更加極端,不再是 1/8,而是可能是 1/16 或 1/32。最激進的解決方案是在整個模型中只放置一層 softmax,但為了保險起見,我們沒有採用它,主要考慮到對檢索能力的影響。

為什麼檢索能力對模型如此重要?

檢索是上下文學習的基礎,是必要條件。

你必須記住上下文中的資訊才能進行上下文學習,而上下文學習是當前大型模型所有高級能力的基礎,例如 CoT(Chain of Thought),尤其是長 CoT,它們都依賴於檢索能力。

決定性的新架構

您是否關注了業界在 FFN 和注意力方面的最新架構改進?

FFN 的改進是 Moe。我也關注了 Byte 的 Ultra Mem,但我認為這是一種有損的東西,一種有損壓縮。如果將來擴大規模可能會出現問題,但我們沒有擴大規模,所以我只能說可能會出現問題。

因為 FFN 基本上就是這些。我們在 Moe 領域的改進無非是從以前的大專家模式改為現在的小專家模式,使其更加稀疏,然後進行一些加速,這需要進一步研究。

如果你想進一步優化它,因為 FFN 是矩陣乘法,優化只能在 Nvidia 的 CUDA 層面上進行,進行一些矩陣乘法的底層優化。

您是否關注了業界在注意力架構方面的改進?

注意力方面的改進基本上是線性的。我們也在考慮未來是否要做更強大的 Linear,並在目前的基礎上進一步加速線性注意力。

有很多方法可以改進,一種是改變衰減,另一種是改變內部的一些小技巧。你可以期待我們的新論文。

我們目前的上下文長度和推理成本的比例是否相對先進?

一旦涉及到延長序列長度,我們就具有非常明顯的計算能力成本優勢。越長,成本優勢就越明顯,無論是推理還是訓練。

例如,在 1M 上,線性注意力消耗的計算能力是完全注意力的 1/2700。相比之下,因為我們仍然有 1/8 的完全注意力,所以基本上是 Transformer 架構的 1/8,因為線性注意力基本上不計為費用。

如果計算成本如此之低,它可以達到計算瓶頸嗎?

現在它確實是記憶體訪問瓶頸。解碼是一個記憶體訪問瓶頸,而不是計算瓶頸。因為 lightning 非常快,它太快了,無法讓記憶體訪問佔用與計算一樣少的資源。這主要是因為實際應用中的序列長度不夠長。

未來如何使其成為計算瓶頸取決於如何優化記憶體訪問。這些將是工程部門需要負責的事情。

如果線性架構成為下一代的主流架構,那麼哪些硬體適應性改進更適合它?

這裡有一個非常棘手的問題,我們需要考慮序列長度。如果你的序列長度集中在 8K 或 32K,那麼注意力只佔百分之十多一點,其餘的百分之八十是 FFN 部分。

即使你將注意力優化到極致,優化到 0,你也只優化了百分之十多一點的延遲。但如果你延長序列長度,注意力的比例會越來越大。這是與完全注意力相比,但對於線性注意力,其比例是不變的。

因為 FFN 也是線性的,並且線性注意力也是線性的,所以它的比例約為 10%,即使在 1M 的情況下也幾乎不變。

但如果是完全注意力,注意力計算可能佔 99%,而下面的 FFN 只佔 1%。因此,線性注意力僅在長文本中具有優勢。

如果線性架構成為主流,那麼追求的可能是低能耗硬體,僅降低能耗。包括 Spiking Neural Network (SNN) 晶片可能更合適,並且實際上有些人正在做。

展望通往 AGI 的道路

您對模型開源效果有什麼期望?

首先是宣傳效果。我個人認為,除了展示一些肌肉外,開源最重要的是看大家未來如何使用它。我認為小型模型開源可能是我們未來更需要考慮做的事情。

以及如何為大家微調做一些基礎設施也可能需要考慮。開源對我們來說是未來的一件長期的事情,旗艦模型應該繼續開源。

未來是否有可能出現非混合的純血架構?

目前,沒有任何方法可以做得比混合更好,尤其是在速度方面。添加一小部分 softmax 注意力,當序列長度不是特別長時,速度優勢非常明顯,尤其是在 flash attention 出現之後。

對純血架構的研究仍在進行中,但非常困難,並且沒有更多的唾手可得的果實。我們有一些技術解決方案,但實現起來並不簡單,最終取決於我們需要實現多長的序列長度。

另一個問題是,對超長文本是否有強烈的需求?儘管像 Claude 這樣的模型已經達到了 200K 上下文,但用戶似乎對目前的長度非常滿意。Agent 應用程序未來可能會帶來對超長序列的需求,但目前還沒有成熟的基準。

但我認為這個問題就像 Nvidia 為未來的遊戲開發先進性能的圖形卡一樣,即使現在不需要它們,它們也是面向未來的技術。

例如,深入的研究需要模型閱讀數十個網站的內容,並且處理時間以數十分鐘為單位,這可能是長文本的應用方向。

你認為 CoT 之後的下一個重大事件可能是什麼?

我們考慮過這個問題。首先,目前的推理模型比較流行,今年主流仍然是推理部分。之後,我們很難想到純語言模型在未來會有任何特別大的變化。

我也與其他老師進行了交談,他們的感覺是,每個人都會重新降低模型的成本,從而使推理速度越來越快,其價格越來越低,並且在降低成本的同時保持效果。

因為天花板正在迅速逼近,絕大多數情況都是檢查和填補大型模型的能力缺口。但如果出現更大的技術突破,它們在短期內可能相對罕見,我們還沒有看到它們。

在 MiniMax 探索線性注意力之後,下一個探索方向可能是什麼?

接下來可能是探索多模態的架構,特別是我們是否想要做這種本機生成和理解統一的大型模型架構。

以 AGI 為終點,哪種計算複雜度為 O(n²) 或 O(n) 的模型會是更好的答案?

當然是 O(n)。從擬人化的角度來看,人一定是 O(n) 複雜度的。例如,如果一個人的複雜度是 O(n²),那麼我對你說話的速度就會越來越慢。

因為對於 Transformer 來說,它的推理複雜度是 O(n²) 計算複雜度,也就是說,吐出第一個 token 和吐出第 100 個 token 的延遲是不同的。

我們人類無法想像這樣的事情,因為人從出生以來就從未重新啟動過,並且一直在吐東西,所以人的計算複雜度是恆定的。

人一定是智慧的最佳解決方案嗎?

我們目前只能這樣認為。也有一些人在走仿生智慧的路線,但我們沒有過多關注那些方向。

以 AGI 為最終遊戲,哪些模型改進領域是最重要的事情?

除了語言建模之外,還有學習方法的問題。你如何學習,以及從環境中學習,從與環境的互動中學習非常重要。畢竟,目前的多模態理解仍然非常缺乏數據。

即使是機器的小樣本學習目前也是有標籤的,但人類的學習是無標籤的。因此,如何將一切統一在一個自我構建的框架下也是一個問題。