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