日本高清色午夜com,色综合国产精品视频,午夜亚洲在在线观看,国产午夜在线网站

      <td id="p7kjh"></td>
      <td id="p7kjh"></td>

      Mamba-2新架構出世一統(tǒng)江湖!普林斯頓CMU華人再出神作,性能狂飆8倍

      發(fā)布時間:2024-06-05 09:00:33 編輯: 來源:
      導讀 相信很多大家對Mamba-2新架構出世一統(tǒng)江湖!普林斯頓CMU華人再出神作,性能狂飆8倍還不知道吧,今天菲菲就帶你們一起去了解一下~.~! 【新...

      相信很多大家對Mamba-2新架構出世一統(tǒng)江湖!普林斯頓CMU華人再出神作,性能狂飆8倍還不知道吧,今天菲菲就帶你們一起去了解一下~.~!

      【新智元導讀】在開源社區(qū)引起「海嘯」的Mamba架構,再次卷土重來!這次,Mamba-2順利拿下ICML。通過統(tǒng)一SSM和注意力機制,Transformer和SSM直接成了「一家親」,Mamba-2這是要一統(tǒng)江湖了?

      年前,Mamba被頂會ICLR拒稿的消息曾引起軒然大波。

      甚至有研究人員表示:如果這種工作都被拒了,那我們這些「小丑」要怎么辦?

      這次,新一代的Mamba-2卷土重來、再戰(zhàn)頂會,順利拿下了ICML2024!

      仍是前作的兩位大佬(換了個順序),仍是熟悉的配方:

      論文地址:https://arxiv.org/pdf/2405.21060

      開源代碼和模型權重:https://github.com/state-spaces/mamba

      不同的是,作者在更高的視角上,統(tǒng)一了狀態(tài)空間模型(SSM)和注意力機制(Attention),也就是文章標題所說的「Transformers are SSMs」。

      ——這下咱們都是一家人了,不用動不動就「打生打死」了。

      性能方面,Mamba-2采用了新的算法(SSD),比前代提速2-8倍,對比FlashAttention-2也不遑多讓,在序列長度為2K時持平,之后便一路遙遙領先。

      在Pile上使用300B token訓練出的Mamba-2-2.7B,性能優(yōu)于在同一數(shù)據集上訓練的Mamba-2.8B、Pythia-2.8B,甚至是更大的Pythia-6.9B。

      從理論上整合了SSM和Transformer,同等性能下,模型更小,消耗更低,速度更快。

      更重要的是,能夠利用GPU的硬件資源(矩陣乘法單元),以及針對Transformer的一系列優(yōu)化。

      ——Mamba-2大有一統(tǒng)江湖之勢。

      1代Mamba,爆發(fā)式占領AI社區(qū)

      事實上,關于1代Mamba的各種研究一直在爆發(fā)性地增長,arxiv已經被各種Mamba所占領,谷歌學術的引用量也達到了350多。

      后續(xù)工作如雨后春筍一般冒出,包括視覺、基因組學、圖表等的直接應用,以及回憶能力、上下文學習能力、形式語言表達能力等方面的研究。

      作者興奮地表示:「我們多年來一直在追求的高效序列模型研究路線,真正引起了機器學習社區(qū)的共鳴?!?/p>

      唯一遺憾的是,Mamba遭到ICLR拒稿,所以關于Mamba到底有沒有前途這個事也就被打上了問號。

      現(xiàn)在,問題解決了,不但論文被接收了,而且還證明了Transformer和Mamba其實是一家人——

      「你說我不行?那Transformer到底行不行?」

      值得注意的是,之前很火的Vision Mamba以及另一篇關于Mamba的研究也殺入了ICML2024。

      對于改進Mamba的初衷,作者表示,當前AI社區(qū)的大家都在努力解決Transformer的問題,盡管SSM的特性和效果都相當好,但卻跟社區(qū)的努力方向不一致。

      這次的Mamba-2可以把針對Transformer的優(yōu)化都用上,不浪費大家的努力。

      新架構一統(tǒng)江湖

      在介紹新架構之前,小編先幫大家簡單理一下背景。

      狀態(tài)空間模型SSM之所以如此令人著迷,是因為它們顯得如此之「基礎」。

      比如,它們與序列模型的許多主要范式,都有著豐富的聯(lián)系。

      它們似乎抓住了連續(xù)、卷積和循環(huán)序列模型的本質,把所有這些元素都包含在了一個簡單優(yōu)雅的模型里。

      不過,另一個主要的序列模型范式——注意力機制的變體,卻更加無所不在。

      然而SSM卻總感覺和Attention是脫節(jié)的。

      在這里,研究者們發(fā)出了「靈魂拷問」——SSM和注意力之間的概念聯(lián)系是什么?有無可能將二者結合起來?

      那就要從公式說起了。

      狀態(tài)空間模型SSM可以這么定義:

      這是個微分方程,利用導數(shù)定義進行代換:

      可以得到SSM的解:

      這個東西就跟RNN一毛一樣了:

      所以可以認為SSM等價于RNN。

      如果將RNN的遞歸結構展開,那么它又可以等價于卷積:

      此時,便可以利用卷積的特性進行并行訓練,而進行推理時又可以享受RNN帶來的O(1)復雜度。

      當然,好事不能讓你全占了,這種結構仍然逃不過固有的梯度爆炸(或消失),以及難以勝任選擇性復制和上下文學習等任務。

      為此,Mamba在SSM的基礎上加入了能夠隨輸入變化的參數(shù)。

      不過這樣做的代價是失去了固定kernel帶來的并行性,所以作者另辟蹊徑,使用前綴和的方式來加速RNN的訓練。

      不過,從計算角度來看,Mamba在硬件效率上仍然遠不如注意力機制。

      原因在于,目前常用的GPU、TPU等加速器,是為矩陣乘法進行過專門優(yōu)化的。

      1代Mamba吃不到硬件矩陣運算單元的紅利,盡管推理時有速度優(yōu)勢,但訓練時問題就大了。

      所以作者就想,我能不能把Mamba的計算重構成矩陣乘法呢?

      于是,新一代的Mamba誕生了。

      結構化狀態(tài)空間對偶性:SSD

      Mamba-2的核心,是結構化狀態(tài)空間對偶性(State Space Duality,SSD)的概念:

      1.SSD模型指的是一個特定的獨立層,比如注意力層或狀態(tài)空間模型(SSM),可以被整合到深度神經網絡中;

      2.SSD框架是一個用于推理該模型(以及更多理論連接)的通用框架;

      3.SSD算法是一種比以前的SSM更高效地計算SSD層的算法。

      SSD框架(紅色,藍色):狀態(tài)空間模型(即半分離矩陣)和結構化掩碼注意力涵蓋了大量高效的序列模型。它們的交集就是SSD模型(紫色)

      原始的Mamba(或更準確地說,其核心「S6」層)實際上是一個具有對角結構的選擇性狀態(tài)空間模型(SSM)。

      Mamba-2的SSD層只做了一個小改動:它進一步限制了對角矩陣??,使其成為標量乘以單位矩陣的結構。換句話說,??的對角元素必須都是相同的值。

      在這種情況下,??可以表示為形狀(??),并且還可以將????識別為一個標量(有時會表示為????)。

      所謂「對偶性」是指,方程(1)(標量-恒等結構????的情況)和(3)中定義的兩個模型實際上是完全相同的模型。

      因此,我們可以將其視為一個特定函數(shù):

      SSD vs. SSM

      與之前的狀態(tài)空間模型(SSM)相比,SSD在遞歸矩陣??上增加了更多結構:

      1. Mamba-1(S6)在矩陣??上使用對角結構,而Mamba-2(SSD)在矩陣??上使用標量乘以單位矩陣的結構;

      2. Mamba-1的頭維度是??=1(即所有通道完全由獨立的SSM控制),而Mamba-2使用的頭維度是??>1(默認情況下類似于??=64)。

      特別是,這可以通過兩種方式視為權重共享:

      1. 通過將矩陣??的對角結構限制為標量乘以單位矩陣,遞歸動態(tài)在狀態(tài)空間的所有??元素之間共享;

      2. 這些動態(tài)也在給定頭的所有??通道之間共享。

      換句話說,一個單一的SSM頭的總狀態(tài)大小為??×??,在Mamba-1中由獨立的標量遞歸控制,而在Mamba-2中由單一的共享遞歸控制。

      而這些變化,主要就是為了提高效率——讓模型能夠以「雙重注意形式」查看,從而允許使用矩陣乘法。

      因此,與Mamba-1相比,Mamba-2支持更大的狀態(tài)維度(從N=16提升到了N=64、N=256甚至更高),同時在訓練期間速度更快。

      SSD vs. Attention

      與標準(自)注意力機制相比,SSD只有兩點不同:

      1. 取消了softmax歸一化;

      2. 以乘法方式應用單獨的元素級掩碼矩陣。

      第一個不同之處在于,它將模型的有效狀態(tài)大小從線性減少到常數(shù),并將效率從二次方提升到了線性。

      第二個不同之處是SSD與標準線性注意力的區(qū)別。一種理解掩碼的方法是將其視為依賴于輸入的相對位置編碼,由于掩碼??的存在,標準的注意力得分????????會被一個權重????:??×=?????????+1所衰減,這可以理解為基于位置??和??之間距離的「折現(xiàn)細數(shù)」(discount factor)。

      在注意力形式中,這種依賴輸入的位置掩碼可以解釋為Mamba「選擇性」的關鍵因素!

      SSD算法

      由于Mamba-1的算法和實現(xiàn)沒有使用張量核心,因此只能進行小規(guī)模的狀態(tài)擴展(通常為??=16)。

      相比之下,矩陣乘法的FLOPs要比非矩陣乘法快得多(最多快16倍):

      - A100GPU有312TFLOPS的BF16矩陣乘法性能,但只有19TFLOPS的FP32算術性能;

      - H100有989TFLOPS的BF16矩陣乘法性能,但只有67TFLOPS的FP32算術性能。

      這次,Mamba-2的一個主要目標,便是利用張量核心來加速SSM。

      由于SSD連接了SSM和結構化矩陣,計算SSM或線性注意力的高效算法,可以直接對應于「token混合」或「序列混合」矩陣??的不同分解。

      如今,這個算法不僅速度更快,而且比原始的Mamba選擇性掃描更容易實現(xiàn),僅需大約25行代碼!

      defsegsum(x):"""Naivesegmentsumcalculation.exp(segsum(A))producesa1-SSmatrix,whichisequivalenttoascalarSSM."""T=x.size(-1)x_cumsum=torch.cumsum(x,dim=-1)x_segsum=x_cumsum[...,:,None]-x_cumsum[...,None,:]mask=torch.tril(torch.ones(T,T,device=x.device,dtype=bool),diagonal=0)x_segsum=x_segsum.masked_fill(~mask,-torch.inf)returnx_segsumdefssd(X,A,B,C,block_len=64,initial_states=None):"""Arguments:X:(batch,length,n_heads,d_head)A:(batch,length,n_heads)B:(batch,length,n_heads,d_state)C:(batch,length,n_heads,d_state)Return:Y:(batch,length,n_heads,d_head)"""assertX.dtype==A.dtype==B.dtype==C.dtypeassertX.shape[1]%block_len==0#Rearrangeintoblocks/chunksX,A,B,C=[rearrange(x,"b(cl)...->bcl...",l=block_len)forxin(X,A,B,C)]A=rearrange(A,"bclh->bhcl")A_cumsum=torch.cumsum(A,dim=-1)#1.Computetheoutputforeachintra-chunk(diagonalblocks)L=torch.exp(segsum(A))Y_diag=torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp",C,B,L,X)#2.Computethestateforeachintra-chunk#(righttermoflow-rankfactorizationofoff-diagonalblocks;Bterms)decay_states=torch.exp((A_cumsum[:,:,:,-1:]-A_cumsum))states=torch.einsum("bclhn,bhcl,bclhp->bchpn",B,decay_states,X)#3.Computetheinter-chunkSSMrecurrence;producescorrectSSMstatesatchunkboundaries#(middletermoffactorizationofoff-diagblocks;Aterms)ifinitial_statesisNone:initial_states=torch.zeros_like(states[:,:1])states=torch.cat([initial_states,states],dim=1)decay_chunk=torch.exp(segsum(F.pad(A_cumsum[:,:,:,-1],(1,0))))new_states=torch.einsum("bhzc,bchpn->bzhpn",decay_chunk,states)states,final_state=new_states[:,:-1],new_states[:,-1]#4.Computestate->outputconversionperchunk#(lefttermoflow-rankfactorizationofoff-diagonalblocks;Cterms)state_decay_out=torch.exp(A_cumsum)Y_off=torch.einsum('bclhn,bchpn,bhcl->bclhp',C,states,state_decay_out)#Addoutputofintra-chunkandinter-chunkterms(diagonalandoff-diagonalblocks)Y=rearrange(Y_diag+Y_off,"bclhp->b(cl)hp")returnY,final_state

      Mamba-2架構

      Mamba-2架構的核心貢獻是提出了新的SSD層,及其理論,與此同時,研究者也對Mamba的神經網絡架構做了一些小改變。

      Mamba-2塊通過刪除連續(xù)線性映射來簡化Mamba塊:SSM參數(shù)??, ??, ??是在塊的開頭生成的,而不是作為SSM輸入??的函數(shù)。如NormFormer中一樣,添加了一個額外的歸一化層,以提高穩(wěn)定性。B和C映射只有一個在??頭之間共享的頭,類似于多值注意力(MVA)

      主要的變化是,??輸入并行生成(??, ??, ??)SSM參數(shù),而非按順序生成。

      之所以這樣的做,與注意力有一定關系。但實際上,它更加簡潔,易于使用張量并行等擴展技術。

      此外,模型架構還有一些其他不同之處。不過,研究作者想要強調的是,這些架構更改并不是模型的真正要點。

      論文中,研究人員主要討論了兩種設計選擇,最終形成Mamba-2架構。

      首先是,塊設計。

      1.并行參數(shù)映射

      在Mamba-2中,SSD層被視為從??, ??, ??, → ?? 的映射。與標準注意力架構的類比,其中??, ??, ??對應于并行創(chuàng)建的Q, K, V投影。

      2.額外歸一化

      在初步實驗中,研究者發(fā)現(xiàn)在較大的模型中容易出現(xiàn)不穩(wěn)定性。

      他們通過在最終輸出映射前的數(shù)據塊中添加一個額外的歸一化層(如LayerNorm、GroupNorm或RMSNorm)來緩解這一問題。

      這種歸一化的用法與NormFormer架構有最直接的關系,后者也在MLP和MHA塊的末尾添加了歸一化層。

      此外,研究者還發(fā)現(xiàn),這種變化與最近從線性注意力角度衍生而來,并與Mamba-2相關的模型類似。

      最初的線性注意力公式通過一個分母項,進行了歸一化,它模仿了標準注意力中softmax函數(shù)的歸一化。

      在TransNormerLLM和RetNet研究中,卻發(fā)現(xiàn)了這種歸一化不穩(wěn)定性。因此在線性注意力層之后增加了一個額外的LayerNorm或GroupNorm。

      研究者提出的「額外歸一化層」與此前研究策略有不同,這是在「乘法門」分支之后加入的歸一化層。

      其次是,序列轉換多頭模式。

      此前,曾提到了SSM被定義為序列轉換,其中:

      - ??, ??, ??參數(shù)的狀態(tài)維度是N;

      - 它們定義了序列轉換

      ,比如可以表示為矩陣

      ;

      - 這一轉換只針對輸入序列

      進行操作,并獨立于P軸。

      總而言之,大家可以將其視為,定義序列轉換的一個頭。

      多頭序列變換由H個獨立的頭組成,模型總維度為D=d_model。這些參數(shù)在多頭之間是共享的,從而形成一個頭模式(head pattern)。

      狀態(tài)大小N和頭維度P,分別類似于注意力的??頭維度和??頭維度。

      正如現(xiàn)代Transformer架構,在Mamba-2中,研究牛人員通常將這些維度,選擇為64或128左右的常數(shù)。

      當模型維度D增加時,便會增加頭數(shù)量,同時保持頭維度N和P不變。

      為了說明如何做到這一點,研究人員可以將多頭注意力的思想進行移植和擴展,從而為SSM或任何一般序列變換定義類似的模式。

      1. 多頭SSM(MHS) / 多頭注意力(MHA)模式

      經典多頭注意力(MHA)模式假定頭維度P,能夠被模型維度D整除。頭的數(shù)量H被定義為H=D/P。

      然后,通過為每個參數(shù)創(chuàng)建H個獨立的副本,就構建出了H個獨立序列變換的「頭」副本。

      值得注意的是,雖然MHA模式最初只是針對注意力序列變換而描述的,但其可以應用到任何符合定義的序列變換上。

      比如,上圖中,多頭SSD層將接受形狀如公式 (17) 所示的輸入,其中SSD算法被復制到了H=n_heads維度上。

      2. 多合約SSM(MCS)/ 多查詢注意(MQA)模式

      多查詢注意力,是對注意力的一種巧妙優(yōu)化,可以顯著提高自回歸推理的速度,它依賴于緩存??張量。

      這一技術只需避免給??和??額外的頭維度,或者換句話說,將(??, ??)的單個頭廣播到??的所有查詢頭。

      利用狀態(tài)空間對偶性(SSD),便可以多頭注意力(MQA)定義為與方程 (18) 等效的的狀態(tài)空間模型(SSM)。

      在這里,??和??(分別對應注意力機制中的V和K)在H個頭之間共享。

      由于控制SSM狀態(tài)收縮的??參數(shù),在每個頭中都有獨立的副本,因此研究人員也將其稱之為多合約SSM(MCS)頭模式。

      類似地,研究人員還可以定義一種多鍵注意力(MKA),或多擴展SSM(MES)頭模式。其中??(控制SSM擴展)在每個頭中是獨立的,而??和??則在所有的頭中共享。

      3.多輸入SSM(MIS)/多值注意力(MVA)模式

      雖然MQA因其KV緩存而受到關注,但它并不是SSM的自然選擇。

      相反,在Mamba中,??被視為SSM的主要輸入,因此??和??是輸入通道共享的參數(shù)。

      研究人員在公式(20)中定義了新的多輸入SSM(MIS)模式的多值注意(MVA),它同樣可以應用于任何序列變換,如SSD。

      有了以上詞匯,現(xiàn)在便可以更準確地描述最初的Mamba架構。

      Mamba架構的選擇性SSM(S6)層具有的特征是:

      - 頭維度??=1:每個通道都有獨立的SSM動態(tài)??;

      - 多輸入SSM(MIS)或多值注意力(MVA)頭結構:矩陣??、??(對應于注意力對偶性中的K、Q應)在輸入??(對應于注意力中的V)的所有通道中共享。

      當然,作者表示,也可以在應用SSD時,去掉這些頭模式的辯題。

      有趣的是,盡管在參數(shù)數(shù)量和總狀態(tài)維度上都有所控制,但在下游性能上卻存在明顯差異。他們根據經驗發(fā)現(xiàn)Mamba最初使用的MVA模式性能最佳。

      第三是,分組頭模式。

      多查詢注意力的理念可以擴展到分組查詢注意力:與使用1個K和V頭不同,它可以創(chuàng)建G個獨立的K和V頭,其中1<G,而且G可以整除H。

      這樣做有兩個動機:一是彌合多查詢注意力和多頭注意力性能差距,二是通過將G設置為分片數(shù)(shards)的倍數(shù),以實現(xiàn)更高效的張量并行。

      最后,研究人員還提到了線性注意力的其他SSD擴展項。

      比如,核注意力近似于Softmax注意力,指數(shù)核特征圖。

      語言建模

      論文中,雖沒有像Mamba-1那樣廣泛地測試Mamba-2,但作者認為新架構總體上可與第一代性能相當,或者更好。

      另外,研究稱,全語言模型結果使用與Mamba相同的協(xié)議,并且在Chinchilla Law上的擴展性略好于Mamba。

      Pile數(shù)據集上的充分訓練的模型,以及標準的零樣本下游任務評估中,也看到了類似的趨勢。

      即使在性能相當?shù)那闆r下,Mamba-2的訓練速度也比初代Mamba快得多!

      合成語言建模:MQAR

      更有趣的是,研究者針對Mamba-2再次嘗試了一項合成任務。

      初代Mamba論文中,曾研究了「合成復制」和「誘導頭」等合成任務后,后續(xù)研究中開始研究更難的聯(lián)想回憶任務。

      目前,由Zoology和Based團隊引入的多查詢聯(lián)想回憶(MQAR),已經成為行業(yè)里的事實標準。

      這次,研究人員測試了一個更難的版本,結果發(fā)現(xiàn),Mamba-2的性能顯著優(yōu)于Mamba-1。

      其中一個原因是,新架構的「狀態(tài)」要大得多——最多是Mamba-1的16倍。這也是Mamba-2的設計初衷之一

      另外,即便是在控制狀態(tài)大小的情況下,Mamba-2在這一特定任務上的表現(xiàn)也明顯優(yōu)于Mamba-1。

      系統(tǒng)和擴展優(yōu)化

      好在,Transformer誕生后,整個研究界和大公司已經對它進行了長達7年的系統(tǒng)優(yōu)化。

      SSD框架在SSM和注意力之間建立聯(lián)系后,也可以讓我們?yōu)镸amba-2等模型實現(xiàn)很多類似的優(yōu)化。

      為此,研究者的重點,就是用于大規(guī)模訓練的張量并行和序列并行,以及用于高效微調和推理的變長序列。

      張量并行

      使用張量并行(TP)進行Mamba-1的大規(guī)模訓練時,一個難點在于,它每層需要進行2次全歸約(all-reduce),而Transformer中的注意力或MLP層,每層只需1次全歸約。

      這是因為,一些SSM參數(shù)是內部激活的函數(shù),而不是層輸入的函數(shù)。

      在Mamba-2中,采用了「并行投影」結構,所有SSM參數(shù)都是層輸入的函數(shù),因此,就可以輕松地將TP應用于輸入投影。

      將輸入投影和輸出投影矩陣,根據TP的程度分成2、4、8個分片。

      使用分組歸一化,分組數(shù)量可被TP程度整除,這樣每個GPU都能單獨進行歸一化。就是這些改變,使得每層只需1次全歸約,而不是2次。

      序列并行

      在訓練非常長的序列時,可能需要沿序列長度進行拆分,并將不同部分分配給不同的設備。

      有兩種主要的序列并行(SP)形式:對于殘差和歸一化操作,這種形式將TP中的全歸約替換為規(guī)約-散布、殘差+歸一化,然后是all-gather。

      由于Mamba-2使用與Transformer相同的殘差和歸一化結構,這種SP形式可以直接應用,無需修改。

      對于注意力或SSM操作,也稱為上下文并行(CP)。

      對于注意力機制,可以使用Ring注意力沿序列維度進行拆分。

      對于Mamba-2,SSD框架又再次幫了大忙:使用相同的塊分解,就可以讓每個GPU計算其本地輸出和最終狀態(tài),然后在更新每個GPU的最終輸出之前,在GPU之間傳遞狀態(tài)。

      可變長度

      在微調和推理過程中,同一批次中經常會出現(xiàn)不同長度的序列。

      對于Transformer,通常會采用填充方式使所有序列長度相同(雖然會浪費計算資源),或者專門為可變長度序列實現(xiàn)注意力機制,并進行負載平衡。

      而對于SSM,就可以將整個批次視為一個長「序列」,并通過將每個序列末尾token的狀態(tài)轉移????設置為0,避免在批次中的不同序列之間傳遞狀態(tài)。

      結果

      結果顯示,更快的SSD算法,直接能讓我們將狀態(tài)維度增加到64或128!而在Mamba-1中,維度僅為16。

      盡管從技術角度看,對于相同的??,Mamba-2比Mamba-1受到的限制會更多,然而更大的狀態(tài)維度,帶來的結果通常就是模型質量的提升。

      更受限制,但更大的狀態(tài)維度通常會提升模型質量。

      比如我們開頭所見的,在Pile上訓練3000億tokens,Mamba-2的表現(xiàn)就明顯優(yōu)于Mamba-1和Pythia。

      而混合模型的表現(xiàn),也很令人滿意。

      從最近的Jamba和Zamba的工作中,研究者發(fā)現(xiàn),將Mamba層與注意力層結合,可以超過純Transformer或Mamba模型的性能。

      在2.7B參數(shù)和3000億tokens規(guī)模上驗證一個僅包含6個注意力塊(和58個SSD塊)的混合模型后可以發(fā)現(xiàn),其表現(xiàn)優(yōu)于64個SSD塊以及標準的Transformer++基線模型(32個門控MLP和32個注意力塊)。

      混合Mamba/注意力模型的下游評估

      而且,對于相同的狀態(tài)維度,SSD算法比Mamba-1的選擇性掃描算法快得多,并且在計算上更能擴展到更大的狀態(tài)維度。

      其中的關鍵就在于,要充分利用張量核心的強大計算能力!

      序列長度2K的效率基準

      未來方向

      如今,線性注意力和SSM連接起來后,前途一片大好,更快的算法、更好的系統(tǒng)優(yōu)化,就在眼前了。

      作者提出,接下來AI社區(qū)需要探索的,有以下三個方向——

      理解:含有少量(4-6)注意力層的混合模型表現(xiàn)非常出色,甚至超過了純Mamba(-2)或Transformer++。

      這些注意力層的作用是什么?它們能被其他機制替代嗎?

      訓練優(yōu)化:盡管SSD可能比注意力機制更快,但由于Transformer中的MLP層非常適合硬件,整體上Mamba-2在短序列長度(例如2K)上,可能仍然比Transformer慢。

      未來,是不是可以讓SSD利用H100的新特性,讓SSM在2-4K序列長度的大規(guī)模預訓練中,比Transformer還快?

      推理優(yōu)化:有許多針對Transformers的優(yōu)化方法,特別是處理KV緩存(量化、推測性解碼)。

      如果,模型狀態(tài)(如SSM狀態(tài))不再隨著上下文長度擴展,KV緩存不再是瓶頸,那時的推理環(huán)境,會如何變化?

      以上就是關于【Mamba-2新架構出世一統(tǒng)江湖!普林斯頓CMU華人再出神作,性能狂飆8倍】的相關內容,希望對大家有幫助!

      免責聲明:本文由用戶上傳,如有侵權請聯(lián)系刪除!

      熱點推薦

      精選文章