蘋果「魔改」快閃記憶體的論文,暴露了它想徹底掌控大模型的野心
2017年,來自 Google 的幾個研究員寫下 Attention is All you need 的題目,給世界帶來了 Transformer 模型架構,它成為了今天「大模型」繁榮背後無可爭議的根基,OpenAI 的 GPT,Meta 的 Llama 以及一眾最主流大語言模型都是在 Transformer 的基礎上生長起來,今天還有一批又一批聰明的研究者在不停嘗試提出比Transformer更強的模型架構。
某種程度上,今天所有AI模型層面的研究都在圍繞對Transformer的掌控與超越展開。但這樣一個劃時代的研究在當時並未立刻引起所有人的重視,而這種「嗅覺」的差異也很大程度決定了今天的AI格局——OpenAI在這篇論文出現第二天就立刻徹底轉向了Transformer架構,然後2020年5月OpenAI基於Transformer架構的GPT-3論文發表,2年後ChatGPT出現,一切都不再相同。
「變壓器」這個欄目名來自對Transformer的直譯,我們會拆解和介紹關於AI及相關技術的最新論文和最前沿研究,希望像OpenAI當年看到Transformer一樣,幫助更多人遇到自己的「變壓器」時刻,比一部分人更早進入生成式AI的世代。
大模型領域最新的一個熱門趨勢是把模型塞到手機裡。而最應該做這個研究的公司終於帶著它的論文現身,那就是蘋果。
這家公司的研究團隊最近發佈了一篇論文《LLM in a flash: Efficient Large Language Model Inference with Limited Memory》,簡單說,它嘗試利用快閃記憶體來解決大模型在塞進手機時遇到的記憶體不足的問題。
這是一個對於端側部署模型十分關鍵的問題。
計算機的記憶體(Memory)簡單分為記憶體(Ram)和快閃記憶體(Flash)兩種。記憶體用於臨時儲存那些需要隨時訪問的資料和指令,它提供高速的讀寫,有較高的儲存密度。而快閃記憶體正相反,它讀寫較慢,適用於長期資料的儲存。
因此從特性上看,記憶體更適合需要頻繁讀寫的大模型。然而這帶來一個問題,它成了一個限制死了的搭配,比如一個70億參數的模型就必須需要超過14GB的記憶體才能以半精度浮點格式載入參數,但這超出了大多數邊緣設備的能力。
如這篇論文的標題所示,蘋果想要通過快閃記憶體來解決這個問題。
蘋果版的曹衝稱象
論文為了將大模型搬到快閃記憶體上,一共做了三步。
第一步:先讓快閃記憶體能參與進模型運行中來。論文提到一個概念,大語言模型在前饋網路(FFN)層展現出高度的稀疏性(超過90%)。FFN是一種基本的神經網路架構,其中資訊單向流動,從輸入層流向輸出層,中間可能經過多個隱藏層。在這種網路中,每一層的輸出僅作為下一層的輸入,而沒有任何反饋或循環連接。於是論文把FFN當成是一個篩子,僅迭代傳輸快閃記憶體中必要的、非稀疏資料到DRAM進行處理。
使用修正線性單元(ReLU)前後輸出幅度對比,ReLU用來實現稀疏性
接下來是第二步:論文提出了一種名為「滑動窗口技術」的神經元資料管理方法。把那些在預測模型中產生正輸出的神經元定義為活躍神經元,在記憶體中保留最近一部分輸入標記的神經元資料,僅載入當前輸入標記與其直接前驅不同的神經元資料。這樣做能有效利用記憶體,釋放掉之前分配給已不在滑動窗口內的舊標記神經元資料的記憶體。
滑動窗口技術
第三步:論文還提出了一種增加資料塊大小的策略。論文用OPT和Falcon模型做實驗,把向上投影的第i列和向下投影的第i行捆綁儲存。當激活第i個中間神經元時,這兩部分資料會同時被使用。通過在快閃記憶體中將這些對應的列和行一起儲存,可以將資料整合成更大的塊進行讀取。
記憶體管理策略,首先將最後的元素複製到要刪除的神經元中,以保持連續的記憶體塊,然後將需要的元素堆疊到最後
這些術語看起來依然晦澀?沒關係我們可以做個類比,事實上它的思路與曹衝稱象非常像。
首先論文要解決的問題就是,大模型是大象,沒辦法直接上秤測量(設備記憶體有限,放不了大模型)。
於是用了三個步驟來在特定環節減少對大模型的訪問延遲。
首先找到一個等價方法,讓大象上船,測量水位線,再用石頭壘在船上,船達到同樣的水位線,最後稱這些石頭的重量(也就是上面說的第一步,可以理解為減少資料載入)。
然後,其中體積一樣大的石頭不需要稱第二次(也就是第二步,最佳化資料塊大小以提高快閃記憶體吞吐量)。
此外,搬運石頭的時候使用更大的框,一次可以裝很多塊石頭(就是最後一步的,高效管理載入到記憶體中的資料)。
而這個過程的重點,是最佳化快閃記憶體互動和記憶體管理,以實現記憶體受限設備上的高效推理。使用這個方法來預測FFN的稀疏性並避免載入零化的參數,最佳化成本模型和按需選擇性載入參數,實現了可以運行比設備DRAM容量大兩倍的模型,並在CPU和GPU上分別比傳統方法提速4-5倍和20-25倍。
當然,論文提供的方法只針對60到70億左右參數的模型,如果是幾百億參數的模型,這樣的辦法會造成死鎖或者記憶體溢出。不過他仍然給了便攜使用大模型這事一種可能性,這是非常難得的。
為了證明論文提出方法的實際價值,論文引用了Facebook的OPT 6.7B模型和TII的Falcon 7B模型。下圖在模型的一半記憶體可用時,1個token的推理延遲。在M1 Max上,每個token從快閃記憶體載入需要125毫秒的延遲,記憶體管理需要65毫秒。因此,每個token的總的與記憶體相關的延遲小於190毫秒(兩者總和)。相比之下,傳統方法需要以6.1GB/s的速度載入13.4GB的資料,導致每個token的延遲大約為2330毫秒。因此,這個方法相對於基準方法表示了重大改進。Falcon 7B也是類似,使用論文的方法延遲僅為250毫秒,而傳統方法的延遲為2330毫秒。延遲肯定是越低越好,越低代表從快閃記憶體中載入大模型的速度越快。
各模型1個token的推理延遲
在AI上落後了?蘋果已經悄悄地做了一堆工作
在今年AI的瘋狂裡,蘋果曾被詬病動作很慢,但這篇論文、此前蘋果提出的MLX框架、自動語音識別(ASR)以及它自己的模型Ferret等研究其實說明,蘋果已經目標明確的在做很具體的研究了。看看這幾個重要的但並沒有引起很多重視的研究,也可以感受到蘋果AI上的方向。
MLX框架是蘋果在2023年推出的一個專門運行在蘋果晶片上的機器學習陣列框架。MLX支持可組合的函數變換,用於自動微分、自動向量化和計算圖最佳化,但重點是MLX中的計算只有在需要時,陣列才會被實際計算出來。同時MLX中的計算圖是動態構建的,改變函數參數的形狀不會觸發緩慢的編譯過程。而且MLX中的陣列存在於共享記憶體中,可以在任何支持的設備類型上執行MLX陣列的操作,而不需要資料傳輸。
也就是說,MLX突出一個節省資源且「海陸空」三棲作戰(可以同時調用記憶體、視訊記憶體,可以在手機和電腦運行)。這說明蘋果非常注重模型的可實現性,即便是手機這樣記憶體有限的設備也能跑大模型。當有了這樣的框架後,蘋果就可以將Ferret模型塞進便攜設備裡了。
Ferret模型是蘋果在2023年10月推出的新型多模態大型語言模型(MLLM),它能夠理解圖像內任意形狀或粒度的空間指代,並準確地對開放詞彙的描述進行定位。Ferret採用了一種新穎而強大的混合區域表示方法,將離散座標和連續特徵結合起來表示圖像中的一個區域。為了提取多樣化區域的連續特徵,論文提出了一種空間感知的視覺取樣器,能夠處理不同形狀之間的稀疏性差異。模型這種理解能力,意味著Ferret可以接受各種輸入,比如點、邊界框和自由形狀,像是DALL·E也好,Midjourney也好,都不能完全理解這種提示詞的輸入。
蘋果將要推出的AR設備Vision Pro,對外宣稱是首款採用空間計算的產品。空間計算本質是傳感器的一門學問,通過傳感器來獲取關於物理空間的資料,並通過計算和分析這些資料來理解和處理環境資訊。傳統電子設備螢幕都只是平面二維,但是空間計算作用的是現實中三維空間的物理概念,在Ferret的加持下,空間的邊界感、長寬高三種向量的立體感就會更加明顯。
Ferret模型不一定能按要求畫出最美的畫面,但它一定能符合擁有藝術設計能力創作者的需求。尤其是在視覺識別、配色方案、排版、網格等設計專業領域,Ferret模型的效果將會最為明顯。設計從業者是蘋果最為廣泛的受眾之一,蘋果就像是個狙擊手,專門瞄準使用者最需要它的地方。
此外蘋果也一直在對與Siri相關的AI技術做研究,比如大語言模型在SLU任務上的準確性受限於ASR系統對給定語音輸入的準確性。那為了解決這一問題,蘋果找到了一種方法:使用ASR的n-best假設列表來提示大語言模型,而非僅依賴錯誤率較高的1-best假設。意味著Siri在接入大語言模型後,性能會得到提高。
至此,從硬體的晶片層,到調用系統側,到與空間計算概念相聯繫的自研多模態模型,再到目前看起來最被期待的蘋果的AI能力的入口Siri,蘋果已經有體系有目的有節奏的完成了諸多技術積累。2024年,在討論AI時沒人能忽視蘋果了。