Hyena成下一代Transformer?StripedHyena-7B開源:最高128k輸入,訓(xùn)練速度提升50%
最近幾年發(fā)布的AI模型,如語言、視覺、音頻、生物等各種領(lǐng)域的大模型都離不開Transformer架構(gòu),但其核心模塊「注意力機(jī)制」的計(jì)算復(fù)雜度與「輸入序列長(zhǎng)度」呈二次方增長(zhǎng)趨勢(shì),這一特性嚴(yán)重限制了Transformer在長(zhǎng)序列下的應(yīng)用,例如無法一次性處理一整本書,或是處理千兆像素級(jí)別的圖像。
即便強(qiáng)如GPT-4也難以擺脫這種缺陷。
最近,Together Research開源了一個(gè)全新的語言模型StripedHyena,采用了針對(duì)「長(zhǎng)上下文」的新架構(gòu),可以處理高達(dá)128k個(gè)token的長(zhǎng)上下文,并且改進(jìn)了Transformer架構(gòu)在訓(xùn)練和推理上的性能,為目前的主流架構(gòu)提供了一種可選方案。
開源鏈接:https://github.com/togethercomputer/stripedhyena
StripedHyena也是「首個(gè)」在短上下文和長(zhǎng)上下文評(píng)估中,以相同模型尺寸,實(shí)現(xiàn)了與最佳開源Transformer模型性能相匹敵的模型:在OpenLLM基準(zhǔn)任務(wù)上與Llama-2, Yi和Mistral 7B實(shí)現(xiàn)了相當(dāng)?shù)男阅?,并且在長(zhǎng)上下文摘要上表現(xiàn)更出色。
StripedHyena是一種混合架構(gòu),由多頭、分組查詢注意力和排列在Hyena塊中的門控卷積組成,不同于傳統(tǒng)的decoder-only的Transformer:通過將卷積表示為狀態(tài)空間模型(SSM,state-space model)或截?cái)酁V波器,在Hyena塊中進(jìn)行常數(shù)內(nèi)存解碼。
實(shí)驗(yàn)結(jié)果顯示,StripedHyena在32k個(gè)token、64k個(gè)token和128k個(gè)token序列的端到端訓(xùn)練中比傳統(tǒng)transformer快30%、50%和100%以上。
SH 7B的另一個(gè)優(yōu)點(diǎn)是,與Transformer相比,自回歸生成期間的內(nèi)存占用減少了50%以上;在Transformers中,每個(gè)層的鍵和值在預(yù)填充階段被緩存,以避免重新計(jì)算并加快增量解碼。
Hyena塊
現(xiàn)有的基于低秩和稀疏近似的次二次(subquadratic)方法需要與稠密的注意層相結(jié)合才能匹配Transformer,也就是說二者在表達(dá)能力上存在差距。
也就是說,注意力機(jī)制在語言處理中只利用了其二次方能力的一小部分,所以研究問題在于,是否存在一個(gè)次二次算子,在大規(guī)模訓(xùn)練時(shí)的性能可以與注意力機(jī)制相匹敵?
今年2月,來自斯坦福大學(xué)和蒙特利爾大學(xué)(Mila and Université de Montréal)的研究人員提出了一個(gè)次二次下降的注意力替代品Hyena:在對(duì)數(shù)千到數(shù)十萬個(gè)token的序列進(jìn)行召回和推理任務(wù)時(shí),Hyena比依賴于狀態(tài)空間和其他隱式和顯式方法的運(yùn)算符提高了50多個(gè)點(diǎn)的準(zhǔn)確性,匹配基于注意力的模型。
論文鏈接:https://arxiv.org/abs/2302.10866
研究人員在標(biāo)準(zhǔn)數(shù)據(jù)集(WikiText 103和The Pile)的語言建模上設(shè)置了一個(gè)新的免密集注意力架構(gòu),達(dá)到了Transformer的質(zhì)量,在序列長(zhǎng)度為2k時(shí)所需的訓(xùn)練計(jì)算減少了20%;在序列長(zhǎng)度為8k時(shí),Hyena算子的速度是高度優(yōu)化注意力的兩倍,在序列長(zhǎng)度為64k時(shí),速度快100倍。
研究人員對(duì)高效次二次primitive進(jìn)行組合,如元素乘法(門控)和長(zhǎng)卷積(即濾波器大小與輸入一樣長(zhǎng)的卷積),最終從實(shí)驗(yàn)結(jié)果中得到了肯定的答案。
根據(jù)最近在機(jī)制可解釋性方面所做的工作(如召回recall和歸納induction),研究人員制定了一系列有針對(duì)性的推理任務(wù),以提煉出注意力與其性能相關(guān)的三個(gè)特性,以及與現(xiàn)有次二次元方法之間的質(zhì)量差距:
1. 數(shù)據(jù)控制(Data control)
注意力機(jī)制實(shí)現(xiàn)了一種富有表現(xiàn)力的數(shù)據(jù)控制線性算子,在單個(gè)塊中對(duì)整個(gè)線性函數(shù)族進(jìn)行編碼。
2. 次線性參數(shù)縮放(sublinear parameter scaling)
將注意力層的參數(shù)數(shù)量與序列長(zhǎng)度脫鉤,允許Transformers在注意力層之間的其他地方分配更多參數(shù),例如前饋神經(jīng)網(wǎng)絡(luò)(FFN)。
3. 無限制的上下文(unrestricted context)
對(duì)于給定的輸入,注意力具有不受限制的上下文,可以近似任意兩個(gè)輸入之間的依賴關(guān)系,而不受任意限制,如局部性(locality);但使用掩碼的情況除外,如自回歸模型。
Hyena層次結(jié)構(gòu)
基于上述三個(gè)發(fā)現(xiàn),研究人員提出Hyena層次結(jié)構(gòu),由兩個(gè)高效的二次基元遞歸定義的算子(長(zhǎng)卷積和元素乘法門控)組成。
遞歸的指定深度(即步數(shù))可控制算子的大小;對(duì)于短遞歸,現(xiàn)有模型可作為特例。
通過將 Hyena 遞歸中的每一步映射到相應(yīng)的矩陣形式,研究人員發(fā)現(xiàn)Hyena算子可以等價(jià)地定義為數(shù)據(jù)控制矩陣的分解,即entries為輸入函數(shù)的矩陣。
此外,研究人員還展示了如何利用快速卷積算法,在不具體化全矩陣的情況下高效地評(píng)估 Hyena 算子。
從經(jīng)驗(yàn)上看,Hyena 算子能夠顯著縮小與大規(guī)模注意力的質(zhì)量差距,以更少的計(jì)算成本實(shí)現(xiàn)了相似的困惑度和下游性能,而且無需混合注意力。
縮小能力差距
設(shè)計(jì)Hyena的初衷是「標(biāo)準(zhǔn)稠密注意力」和「次二次運(yùn)算符」之間的存在質(zhì)量差距,并且可以通過與大規(guī)模語言建模性能相關(guān)的推理任務(wù)來確定這一差距。
研究人員擴(kuò)展了一套基本的機(jī)械可解釋性基準(zhǔn)(歸納和召回),并增加了額外的任務(wù)以探究當(dāng)任務(wù)復(fù)雜度增加(如詞匯量增加)時(shí),模型性能會(huì)如何快速下降。
此外,文中還研究了 Hyena 中長(zhǎng)卷積的最佳參數(shù)化。
在具有數(shù)十萬詞條的最具挑戰(zhàn)性的設(shè)置中,隱式參數(shù)化方案比其他利用狀態(tài)空間、頻域參數(shù)化或標(biāo)準(zhǔn)卷積的算子提高了50%以上的準(zhǔn)確率。
語言和視覺中的擴(kuò)展
研究人員還驗(yàn)證了推理基準(zhǔn)套件中的排名是否能預(yù)測(cè)大規(guī)模質(zhì)量,對(duì)Hyena 自回歸語言建模進(jìn)行了十億以下參數(shù)規(guī)模的測(cè)試,在標(biāo)準(zhǔn)數(shù)據(jù)集(WikiText103 和 The Pile)中的無稠密注意力架構(gòu)中實(shí)現(xiàn)了新sota,并與Transformer的質(zhì)量相當(dāng)。
在3.35億參數(shù)規(guī)模的The Pile數(shù)據(jù)集上,該系統(tǒng)以減少20%浮點(diǎn)運(yùn)算(FLOPs)總數(shù)達(dá)到了與Transformer相當(dāng)?shù)睦Щ蠖戎笜?biāo)。
作為擴(kuò)展,研究人員在大規(guī)模圖像識(shí)別中測(cè)試了Hyena運(yùn)算符的通用性,并在視覺Transformer(ViT)中取代了注意力。
在圖像分類中,當(dāng)在ImageNet-1k數(shù)據(jù)集上從頭開始訓(xùn)練時(shí),Hyena 的準(zhǔn)確率可以與注意力相媲美。
面向更長(zhǎng)的上下文
研究人員還對(duì) Hyena 在長(zhǎng)序列上的效率進(jìn)行了基準(zhǔn)測(cè)試。在長(zhǎng)度為 8192 的序列上,測(cè)得的速度是稠密自注意力的 5 倍,是高度優(yōu)化的 FlashAttention的 2 倍,而在長(zhǎng)度為 64k 的序列上,測(cè)得的速度是 FlashAttention 的 100 倍。