如何評價最新的Mamba?真的能超越Transformer嗎?最新綜述!
本文經自動駕駛之心公眾號授權轉載,轉載請聯系出處。
寫在前面&筆者的個人理解
深度學習作為一項重要技術,引發了人工智能(AI)的一場非凡革命,導致了人類生活方式的巨大變化。作為最具代表性的架構,Transformers已經為眾多高級模型提供了支持,特別是包含數十億個參數的大型語言模型(LLM),成為深度學習的基石。盡管取得了令人矚目的成就,但Transformer仍然面臨著固有的局限性,特別是注意力計算的二次計算復雜性導致的耗時推理。最近,一種名為Mamba的新型架構從經典狀態空間模型中汲取靈感,成為構建基礎模型的有前景的替代方案,在保持序列長度接近線性可擴展性的同時,為Transformers提供了可比的建模能力。這引發了越來越多的研究,積極探索Mamba在各個領域取得令人印象深刻的性能的潛力。鑒于這種快速發展,迫切需要進行系統審查,鞏固現有的Mamba授權模型,全面了解這種新興的模型架構。因此,在本次調查中,我們對最近的曼巴相關研究進行了深入調查,主要從三個方面進行:基于曼巴的模型的進展、使曼巴適應各種數據的技術以及曼巴可以擅長的應用。具體來說,我們首先回顧各種代表性深度學習模型的基礎知識以及Mamba-1和2的細節作為估計。然后,為了展示Mamba對人工智能的意義,我們全面回顧了有關Mamba模型架構設計、數據適應性和應用的相關研究。最后,我們討論了當前的局限性,并探索了各種有前景的研究方向,為未來的研究提供了更深入的見解。
本文簡介
在過去的二十年里,深度學習作為最突出的人工智能(AI)技術,在醫療保健、自主系統、推薦系統和金融服務等各個領域帶來了革命。這一時期見證了眾多深度神經網絡(DNN)的出現,這些網絡極大地改變了人類的生活方式,為個人提供了巨大的便利。一個值得注意的例子是U-Net,這是一種視野內的強大深度學習模型,廣泛應用于醫學成像,用于檢查MRI和CT等放射學掃描。其應用有助于疾病的識別和診斷,展示了其在這一關鍵醫療保健領域的有效性。此外,圖神經網絡(GNN)用于處理圖結構數據以支持智能服務,例如向用戶推薦個性化內容、產品或服務的推薦系統。此外,循環神經網絡(RNN)在機器翻譯中被廣泛采用,因為它們能夠捕獲準確翻譯所必需的順序和上下文信息,使來自不同語言背景的個人能夠有效地交流和包含彼此的想法、觀點和信息。
在各種深度學習架構中,Transformers最近脫穎而出,并在廣泛的應用領域確立了其主導地位。例如,作為最具代表性的大型基礎模型,像ChatGPT這樣的大型語言模型(LLM)從根本上建立在Transformer架構上。通過將模型規模擴展到數十億,并在各種數據源的混合上進行訓練,這些基于Transformer的模型展示了人類水平的智能,在語言理解、常識推理和內容學習方面具有令人印象深刻的能力。這一顯著的成功得益于注意力機制,它使基于Transformer的模型能夠專注于輸入序列的相關部分,并促進更好的上下文理解。然而,注意力機制也引入了顯著的計算開銷,該開銷隨著輸入大小的平方而增加,這給處理縱向輸入帶來了挑戰。例如,計算成本的快速增長使Transformers處理大量序列變得不切實際或不可行,從而限制了它們在文檔級機器翻譯或長文檔摘要等任務中的適用性。
最近,一種有前景的架構,結構化狀態空間序列模型(SSM),已經出現,可以有效地捕獲序列數據中的復雜依賴關系,成為Transformer的強大競爭對手。這些模型受經典狀態空間模型的啟發,可以被認為是遞歸神經網絡和卷積神經網絡的融合。它們可以通過遞歸或卷積運算高效計算,實現序列長度的線性或近線性縮放,從而顯著降低計算成本。更具體地說,作為最成功的SSM變體之一,Mamba實現了與Transformer相當的建模能力,同時保持了序列長度的線性可擴展性,將其推向了焦點領域。Mamba首先引入了一種簡單而有效的選擇機制,通過基于輸入對SSM參數進行參數化,使模型能夠過濾掉不相關的信息,同時無限期地保留必要和相關的數據。然后,Mamba提出了一種硬件感知算法,通過掃描而不是卷積來定期計算模型,在A100 GPU上實現了高達3倍的計算速度。如圖1所示,Mamba對復雜和長度順序數據的強大建模能力,以及近乎線性的可擴展性,使其成為一種新興的基礎模型,有望徹底改變各種研究和應用領域,如計算機視覺、自然語言處理、醫療保健等。例如,Zhu等人提出了Vim,它比DeiT快2.8倍,在提取高分辨率圖像的特征時節省了86.8%的GPU內存。Dao和Gu展示了SSM和注意力變體之間的聯系,并提出了一種新的架構,該架構改進了選擇性SSM,使語言建模速度提高了2-8倍。
在Mamba強大的長序列建模能力及其高效性的推動下,出現了大量文獻,重點是在各種下游任務中使用和改進Mamba。鑒于與曼巴相關的研究大幅增加,對現有文獻進行全面回顧并考慮未來研究的潛在方向至關重要。因此,在這項調查中,我們從幾個角度對曼巴進行了全面的回顧,為新來者提供了對曼巴內部運作的基本了解,同時幫助經驗豐富的從業者了解其最新發展。具體來說,剩下的調查組織如下:在第2節中,我們回顧了各種代表性深度神經網絡的背景知識,包括RNN、Transformers和狀態空間模型,而Mamba的詳細信息在第3節中介紹。接下來,我們在第4節中從塊設計、掃描模式和內存管理的角度總結了基于Mamba的研究的最新進展。然后,第5節介紹了使Mamba適應各種數據的技術,包括順序和非順序數據。此外,第6節介紹了Mamba模型的代表性應用,第7節介紹了挑戰和未來方向。最后,我們在第8節總結了整個調查。
與我們的調查同時,發布了幾項相關調查,純粹關注狀態空間模型和視覺曼巴。與這些調查不同,本文主要關注曼巴的相關研究。它從一個新的角度系統地分析了現有文獻,以探索Mamba架構的演變以及基于Mamba的模型中使用的數據自適應方法。
Mamba
為了解決傳統SSM在上下文感知能力方面的上述缺點,[55]提出了Mamba作為一種潛在的替代方案,有望成為通用的序列基礎模型骨干。最近,Mamba-2[提出了結構化空間狀態對偶性(SSD),建立了一個將結構化SSM和各種形式的注意力聯系起來的穩健理論框架,使我們能夠將最初為變壓器開發的算法和系統優化轉移到SSM。在本節中,我們將對曼巴和曼巴-2進行簡潔明了的介紹。
傳統的SSM在建模文本和其他信息密集型數據方面的效果有限,阻礙了它們在深度學習方面的進展。在追求賦予SSM Transformer建模能力的過程中,Gu和Dao引入了三種基于結構化狀態空間模型的創新技術,即基于高階多項式投影算子(HiPPO)的內存初始化、選擇機制和硬件感知計算,如圖3所示。這些技術旨在增強SSM在長距離線性時間序列建模中的能力。特別是,初始化策略建立了一個連貫的隱藏狀態矩陣,有效地促進了長距離記憶。然后,選擇機制使SSM能夠獲取內容感知表示。最后,Mamba設計了兩種硬件感知計算算法,并行聯想掃描和內存重構,以提高訓練效率。
HiPPO-based Memory Initialization:從序列數據中建模和學習是當代機器學習的基本挑戰,是各種任務的基礎,包括語言建模、語音識別和視頻處理。對復雜和長期的時間依賴性進行建模的一個基本組成部分在于內存,包括存儲和整合來自先前時間步驟的信息的能力。與RNN類似,保留和忘記歷史隱藏狀態(即矩陣A)在SSM中起著至關重要的作用,以實現令人滿意的性能。在之前的結構化狀態空間序列模型(SSM)中,有人建議進行特殊的初始化,特別是在復值模型的情況下。事實證明,這些特殊的初始化在各種情況下都是有益的,包括數據可用性有限的情況。同樣,Mamba主要關注隱藏狀態矩陣A的初始化,以捕獲復雜的時間依賴關系。這是通過利用HiPPO理論和創新的勒讓德尺度(LegS)來實現的,確保全面考慮完整的歷史背景,而不是有限的滑動窗口。具體來說,HiPPO LegS為所有歷史數據點分配了統一的權重,可以表示為:
基于HiPPO理論,Mamba為復雜和真實的情況引入了兩種簡單的初始化方法,即S4D-Lin和S4D-real,如
Selection Mechanism:由于時間不變性的特性,傳統的狀態空間模型對于基于特定模型輸入(即內容感知建模能力)產生個性化輸出是必不可少的。為了給SSM提供類似于注意力機制的能力,Mamba設計了一種時變選擇機制,根據模型輸入對權重矩陣進行參數化。這種創新使SSM能夠過濾掉無關的信息,同時無限期地保留相關細節。形式上,選擇機制涉及將區間D和矩陣B、C設置為輸入x的函數,可以表示為:
Hardware-aware Computation:該選擇機制旨在克服線性時不變模型的局限性。盡管如此,它還是對高效訓練提出了挑戰:SSM的卷積核變得依賴于輸入,導致無法執行并行計算。為了解決這個問題,Mamba使用了兩種計算技術,即并行關聯掃描(也稱為并行前綴和)和內存重新計算。首先,并行關聯掃描利用線性關聯計算的特性和現代加速器(GPU和TPU)的并行性,以內存高效的方式執行選擇性SSM的計算。更具體地說,并行關聯掃描將模型訓練的計算復雜度從降低到。掃描的核心是在給定的輸入上構建一個平衡的二叉樹,并從根開始來回掃描。換句話說,并行關聯掃描首先從葉子遍歷到根(即,向上掃描),在樹的內部節點處創建部分和。然后,它反轉遍歷,從根向上移動到樹上,使用部分和構建整個掃描(即向下掃描)。另一方面,Mamba利用傳統的重新計算方法來減少訓練選擇性SSM層的總體內存需求。特別是,在并行關聯掃描的前向傳遞過程中,Mamba與大小為BLDN的干擾中間狀態保持距離,以防止內存擴展。相反,它在反向過程中重新計算這些中間狀態以進行梯度計算。通過這樣做,重新計算避開了在GPU內存單元之間讀取元素的必要性。除了優化掃描操作的內存需求外,Mamba-1還擴展了重新計算的使用,以提高整個SSM層的效率。這種優化包括投影、卷積和激活,這些通常需要大量的內存資源,但可以快速重新計算。
Mamba-2: State Space Duality
Transformers在各個領域的深度學習成功中發揮了至關重要的作用,激發了各種技術的發展,如參數高效微調[95]、災難性遺忘緩解[96]和模型量化[195],旨在從各個角度提高模型性能。為了使狀態空間模型能夠訪問并受益于最初為Transformer開發的有價值的技術,Mamba-2[29]引入了一個名為結構化狀態空間對偶性(SSD)的綜合框架,該框架在SSM和不同形式的注意力之間建立了理論聯系。公式化:
基于SSD,Mamba-2通過塊分解矩陣乘法算法開發了一種更具硬件效率的計算方法。具體來說,通過矩陣變換將狀態空間模型視為半可分矩陣,Mamba-2將計算分解為矩陣塊,其中對角塊表示塊內計算。相比之下,非對角塊表示通過SSM的隱藏狀態分解的塊間計算。這種方法使Mamba-2的訓練過程比Mamba-1的并行聯想掃描快2-8倍,同時保持與Transformer的競爭力。
Mamba Block
在本小節中,我們總結了Mamba-1和Mamba-2的區塊設計。圖4顯示了這兩種架構的比較。Mamba-1的動機是以SSM為中心的觀點,其中選擇性SSM層的任務是從輸入序列X到Y進行映射。在這種設計中,(a,B,C)的線性投影是在創建X的初始線性投影之后應用的。然后,輸入令牌和狀態矩陣通過選擇性SSM單元,利用并行關聯掃描產生輸出Y。之后,Mamba-1使用跳過連接來鼓勵特征重用和模型訓練過程中經常出現的所有退化問題。最后,通過將該塊與標準歸一化和殘差連接交織堆疊,構建了Mamba模型。
至于Mamba-2,它引入了SSD層,旨在創建一個從[X,A,B,C]到Y的映射。這是通過在塊的開頭用單個投影同時處理[X,A,B,C]來實現的,類似于標準注意力架構并行生成Q,K,V投影的方式。換句話說,Mamba-2塊通過去除連續的線性投影簡化了Mamba-1塊。與Mamba-1中的并行選擇性掃描相比,這使得SSD結構的計算速度更快。此外,在跳躍連接后添加了一個歸一化層,旨在提高訓練穩定性。
Advancements in Mamba Models
狀態空間模型和Mamba最近得到了探索,并成為一種有前景的替代方案,作為基礎模型骨干。雖然Mamba在自然語言處理方面表現出了熟練程度,但它仍然遇到了挑戰,例如記憶喪失、對各種任務的泛化,以及基于Transformer的語言模型捕獲復雜模式的能力較差。為了克服這些挑戰,已經做出了大量努力來改進曼巴結構。現有的研究主要集中在修改塊設計、掃描模式和內存管理方面。本節將從這三個方面介紹幾種重要技術,相關研究總結見表1。
Block Design
曼巴block的設計和結構對曼巴模型的整體性能有重大影響,使其成為新興的研究熱點。如圖5所示,基于構建新曼巴區塊的不同方法,現有研究可分為三類:a)整合方法旨在將曼巴區塊與其他知名模型整合,以在有效性和效率之間取得平衡;b) 替代方法試圖利用Mamba塊來替代高級模型框架中的主要層;c)修改方法側重于修改經典曼巴塊內的組件。
Scanning Mode
并行關聯掃描操作是Mamba模型中的一個關鍵組成部分,旨在解決選擇機制引起的計算問題,加速訓練過程,減少內存需求。它通過利用時變SSM的線性特性在硬件級別設計核融合和重新計算來實現這一點。然而,Mamba的單向序列建模范式阻礙了對各種數據(如圖像和視頻)的全面學習過程。為了緩解這一問題,幾項研究側重于設計高效的掃描方法,以提高模型性能并促進Mamba模型的訓練過程。如圖6所示,專注于開發掃描模式技術的現有研究可分為兩類:1)平面掃描方法從令牌序列的平面角度處理模型輸入;以及2)立體掃描方法跨維度、通道或比例掃描模型輸入。
Memory Management
與RNN一樣,狀態空間模型中隱藏狀態的記憶有效地存儲了之前步驟的信息,因此在SSM的整體功能中起著至關重要的作用。雖然Mamba引入了基于HiPPO的內存初始化方法,但SSM單元的內存管理仍然存在挑戰,包括在層之間傳輸隱藏信息和實現無損內存壓縮。為此,一些開創性的研究提出了不同的解決方案,包括內存初始化、壓縮和連接。例如,Ezoe和Sato試圖通過在模型再訓練期間使用平衡截斷方法來改進選擇性SSM的初始化過程。此外,DGMamba引入了一種隱藏狀態抑制方法,以增強狀態空間模型中隱藏狀態的領域泛化能力。這種方法可以消除這些隱藏狀態的負面影響,從而縮小不同領域隱藏狀態之間的差距。同樣,DenseMamba提出了一種密集連接方法,以增強SSM中各層之間隱藏信息的傳播。該策略旨在通過有選擇地將淺層的隱藏狀態集成到深層來減輕內存退化,并為輸出生成保留詳細信息。
Adapting Mamba to Diverse Data
Mamba架構代表了選擇性狀態空間模型的擴展,它具有循環模型的基本特性,使其非常適合作為操作文本、時間序列、語音等序列的通用基礎模型。與此同時,最近的開創性研究將Mamba架構的利用擴展到了序列數據之外,包括圖像和圖形等領域,如圖7所示。這些研究旨在利用Mamba在捕獲長期依賴關系方面的卓越能力,同時利用其在學習和推理過程中的效率。因此,在本節中,我們的目標是研究使Mamba適應各種類型數據的新興技術。
Sequential Data
序列數據是指以特定順序收集或組織的數據,其中數據點的順序具有重要意義。為了探索利用Mamba作為序列數據任務基礎模型的潛力,我們在以下部分進行了全面的回顧,涵蓋了各種序列數據,包括自然語言、視頻、時間序列、語音和人體運動。
Non-Sequential Data
非序列數據與順序數據的不同之處在于不遵循特定的順序。其數據點可以按任何順序組織或訪問,而不會顯著影響數據的含義或解釋。這種固有順序的缺失給循環模型帶來了困難,例如專門用于捕獲數據中時間依賴性的RNN和SSM。令人驚訝的是,代表SSM的Mamba在最近的發展中在有效處理非連續數據方面取得了顯著的成功。在本節中,我們將回顧有關Mamba如何有效處理非連續數據(包括圖像、圖形和點云)的相關研究。
Multimodal Data
整合多種模態,如語言(順序數據)和圖像(非順序數據),為人工智能感知和場景理解提供了有價值和互補的信息。最近,多模態大型語言模型(MLLM)受到了廣泛的研究關注,MLLM固有了大型語言模型的高級功能,包括強大的語言表達和邏輯推理。雖然Transformers一直是該領域的主導方法,但Mamba已經成為一個強大的競爭對手,在對齊混合源數據和實現序列長度的線性復雜度縮放方面表現出了令人印象深刻的性能,這使得Mamba成為Transformers多模態學習的有前景的替代品。例如,Qiao等人提出VL Mamba,探索利用Mamba的高效架構來解決視覺語言任務,利用預先訓練的Mamba模型進行語言理解,并結合連接器模塊將視覺補丁與語言標記對齊。Wang等人提出了文本控制運動Mamba,它利用Mamba基于文本查詢動態捕獲全局時間信息,以增強對人類運動的理解。此外,Fusion-Mamba和Sigma試圖融合來自不同模態(如熱、深度和RGB)的互補信息。Fusion-Mamba專注于改進對象檢測,而Sigma旨在增強語義分割。
應用
Natural Language Processing
在自然語言處理領域,最近出現了一些基于Mamba的模型作為基于Transformer的語言建模模型的替代品,特別是在涉及廣泛上下文的應用中,如問答系統和文本摘要。
Computer Vision
除了NLP應用外,基于Mamba的模型在計算機視覺領域也顯示出潛力,如疾病診斷和運動識別與生成等代表性應用。
Speech Analysis
語音信號本質上由數千個樣本組成。雖然這種寬泛的時間背景提供了豐富的聲學特征,但它也帶來了巨大的計算需求。為了有效地處理語音信號,幾種基于Mamba的模型已成功應用于各種語音應用,特別是語音分離和標記以及語音增強。
Drug Discovery
蛋白質設計、分子設計和基因組分析在推進藥物發現和生物技術方面至關重要。利用基于MAMBA的模型顯著降低了這些域中長序列建模的復雜性。具體來說,PTM-Mamba]和ProtMamba是基于Mamba架構的蛋白質語言模型。PTM-Mamba使用雙向門控Mamba塊和結構化狀態空間模型,有效地處理長序列,同時減少計算需求。ProtMamba被設計為具有同源性意識但無比對,足以處理數百個蛋白質序列的廣泛背景。這兩種模型即使在大數據集的情況下也能保持高效率和準確性,為蛋白質設計提供了關鍵工具。與此同時,生成式分子設計旨在模擬具有特定分布的定制特性的分子。然而,目前的模型缺乏優化高保真神諭所需的效率,直接導致成功率低。Saturn應用Mamba架構,利用其線性復雜性和計算效率在藥物發現方面超越了22個競爭模型。此外,了解基因組對于深入了解細胞生物學至關重要。基因組建模的挑戰包括捕捉遠距離標記之間的相互作用,考慮上游和下游區域的影響,并確保DNA序列的互補性。Caduceus和MSAMamba都利用了Mamba模型,在應對這些挑戰方面表現出色。Caduceus是一種DNA基礎模型,它通過BiMamba和MambaDNA組件增強了Mamba架構,用于雙向建模并確保反向補體等效性,在長期基因組任務中顯著優于現有模型。同樣,MSAMamba通過沿序列維度實施選擇性掃描操作,解決了基于變換器的DNA多序列比對模型的局限性。這種設計將以前方法的訓練上下文長度延長了八倍,從而可以對廣泛的DNA序列進行更全面的分析。
Recommender Systems
推薦系統廣泛應用于電子商務和社交網絡,旨在捕捉用戶不斷變化的偏好及其過去行為之間的相互依賴關系。盡管基于變換器的模型在推薦系統中表現出了有效性,但由于注意力機制的平方復雜性,它們面臨著計算效率的挑戰,尤其是在處理較長的行為序列時。最近,幾個基于Mamba的模型已被應用于分析個性化推薦的長期用戶行為。例如,Mamba4Rec率先使用選擇性狀態空間模型進行高效的順序推薦,在保持推理效率的同時提高了模型性能。同樣,RecMamba探討了Mamba在終身順序推薦場景(即序列長度≥2k)中的有效性,實現了與基準模型相當的性能,同時將訓練時間縮短了70%,內存成本降低了80%。此外,EchoMamba4Rec將雙向Mamba模塊與頻域濾波集成在一起,以準確捕捉用戶交互數據中的復雜模式和相互依賴關系。它展示了優于現有模型的性能,提供了更精確和個性化的建議。此外,Mamba4KT專為智能教育中的知識追蹤而設計,利用Mamba模型來捕捉練習和學生知識水平之間的持久相關性。隨著教育數據集的擴展,這種方法為提高知識追蹤研究中的預測準確性、模型效率和資源利用率提供了一條有前景的途徑。
Robotics and Autonomous Systems
機器人和自主系統的主要目標是開發能夠適應全面視覺環境并執行復雜動作的模型。目前用于機器人技術的多模態大語言模型(MLLM)在兩個主要方面面臨著重大挑戰:1)處理需要高級推理的復雜任務的能力有限,2)微調和推理任務的計算費用巨大。由于其在推理速度、內存利用率和整體效率方面的優勢,基于Mamba的模型正在成為自主和智能系統的有前景的基礎,有望實現卓越的性能和巨大的可擴展性潛力。例如,RoboMamba將視覺編碼器與Mamba集成在一起,創建了一個端到端的機器人MLLM。該方法通過協同訓練將視覺數據與語言嵌入對齊,用視覺常識和機器人特定的推理增強模型,同時確保高效的微調和推理能力。同樣,Jia等人介紹了MaIL,這是一種使用Mamba作為骨干的模仿學習(IL)策略架構。MaIL彌合了處理觀測序列的效率和性能之間的差距。對真實機器人實驗的廣泛評估表明,MaIL為傳統、大型和復雜的基于Transformer的IL策略提供了一種有競爭力的替代方案。
未來挑戰
- Mamba-based Foundation Models
- Hardware-Awareness Computation
- Trustworthy Mamba Models
- Applying Emerging Techniques from Transformer to Mamba
結論
Mamba是一種新興的深度學習架構,由于其強大的建模能力和計算效率,在語言生成、圖像分類、推薦和藥物發現等各個領域都取得了顯著的成功。最近,人們越來越努力地開發基于Mamba的深度學習模型,這些模型具有更強大的表示學習能力和更低的計算復雜度。鑒于曼巴的快速進展,迫切需要一個系統的概述。為了彌合這一差距,在本文中,我們對Mamba進行了全面的回顧,重點介紹了它的架構進步、數據適應性和應用領域,為研究人員提供了對Mamba最新發展的深入了解和概述。此外,鑒于曼巴研究仍處于起步階段,我們還討論了當前的局限性,并為未來的研究提出了有前景的方向。