LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

CV技术指南(公众号)發表於2024-05-14
前言 LSTM:這次重生,我要奪回 Transformer 拿走的一切。

本文轉載自新智元

僅用於學術分享,若侵權請聯絡刪除

歡迎關注公眾號CV技術指南,專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀、CV招聘資訊。

CV方向的準研究生們,未來三年如何度過?

招聘高光譜影像、語義分割、diffusion等方向論文指導老師

20 世紀 90 年代,長短時記憶(LSTM)方法引入了恆定誤差選擇輪盤和門控的核心思想。三十多年來,LSTM 經受住了時間的考驗,併為眾多深度學習的成功案例做出了貢獻。然而,以可並行自注意力為核心 Transformer 橫空出世之後,LSTM 自身所存在的侷限性使其風光不再。

當人們都以為 Transformer 在語言模型領域穩坐江山的時候,LSTM 又殺回來了 —— 這次,是以 xLSTM 的身份。

5 月 8 日,LSTM 提出者和奠基者Sepp Hochreiter在 arXiv 上傳了 xLSTM 的預印本論文。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

論文的所屬機構中還出現了一家叫做「NXAI」的公司,Sepp Hochreiter 表示:「藉助 xLSTM,我們縮小了與現有最先進 LLM 的差距。藉助 NXAI,我們已開始構建自己的歐洲 LLM。」

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!
  • 論文標題:xLSTM: Extended Long Short-Term Memory
  • 論文連結:

具體來說,xLSTM 從三個層面解決了 LSTM 以往所存在的侷限性:

(i) 無法修改儲存決策。

可以透過「最近鄰搜尋」(Nearest Neighbor Search)問題來舉例說明這一侷限性:在給定參考向量的情況下,必須按順序掃描序列,尋找最相似的向量,以便在序列末端提供其附加值。圖 2 左側顯示了這項任務的均方誤差。當發現更相似的向量時,LSTM 難以修改儲存的值,而新的 xLSTM 透過指數門控彌補了這一限制。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

(ii) 儲存容量有限,即資訊必須壓縮成標量單元狀態。

圖 2 右側給出了 Wikitext103 上不同 token 頻率的 token 預測困惑度。由於 LSTM 的儲存容量有限,它在不常見 token 上的表現較差。xLSTM 透過矩陣記憶體解決了這一問題。

(iii) 由於記憶體混合而缺乏可並行性,需要進行順序處理。例如,從一個時間步到下一個時間步的隱藏狀態之間的隱藏 - 隱藏連線。

與此同時,Sepp Hochreiter 和團隊在這篇新論文中回答了一個關鍵問題:如果克服這些侷限性並將 LSTM 擴充套件到當前大語言模型的規模時,能實現怎樣的效能?

將 LSTM 擴充套件到數十億引數

為了克服 LSTM 的侷限性,xLSTM 對等式(1)中的 LSTM 理念進行了兩項主要修改。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

在原來的 LSTM 中,恆定誤差選擇輪盤是由單元輸入 z_t 對單元狀態 c_(t-1)(綠色)進行的加法更新,並由 sigmoid 門(藍色)進行調節。輸入門 i_t 和遺忘門 f_t 控制這一更新,而輸出門 o_t 控制儲存單元的輸出,即隱藏狀態 h_t。儲存單元的狀態被 ψ 歸一化或壓縮,然後輸出門控得到隱藏狀態。

xLSTM 的修改包括指數門控和新穎的記憶體結構,因此豐富了 LSTM 家族的兩個成員:

(i) sLSTM(第 2.2 節),具有標量記憶體、標量更新和記憶體混合功能;

(ii) mLSTM(第 2.3 節),具有矩陣記憶體和協方差(外積)更新規則,完全可並行處理。

sLSTM 和 mLSTM 都透過指數門控增強了 LSTM。為了實現並行化,mLSTM 放棄了記憶體混合,即隱藏 - 隱藏遞迴連線。mLSTM 和 sLSTM 都可以擴充套件到多個儲存單元,其中 sLSTM 具有跨單元記憶體混合的特點。此外,sLSTM 可以有多個頭,但不存在跨頭的記憶體混合,而只存在每個頭內單元間的記憶體混合。透過引入 sLSTM 頭和指數門控,研究者建立了一種新的記憶體混合方式。對於 mLSTM 而言,多頭和多單元是等價的。

將這些新的 LSTM 變體整合到殘差塊模組中,就得到了 xLSTM 塊。將這些 xLSTM 塊剩餘堆疊到架構中,就形成了 xLSTM 架構。xLSTM 架構及其元件見圖 1。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

xLSTM 塊應在高維空間中對過去進行非線性總結,以便更好地分離不同的歷史或上下文。分離歷史是正確預測下一個序列元素(如下一個 token)的先決條件。研究者在此採用了 Cover 定理,該定理指出,在高維空間中,非線性嵌入模式比在原始空間中更有可能被線性分離。

他們考慮了兩種殘差塊結構:(i) post up-projection 的殘差塊(如 Transformer),它非線性地概括了原始空間中的歷史,然後線性地對映到高維空間,應用非線性啟用函式,再線性地對映回原始空間(圖 3 左側和圖 1 第三欄,更詳細的版本見圖 9)。(ii) pre up-projection 的殘差塊(如狀態空間模型),它線性地對映到高維空間,在高維空間中非線性地總結歷史,然後線性地對映回原始空間。對於包含 sLSTM 的 xLSTM 塊,研究者主要使用了 post up-projection 塊。對於包含 mLSTM 的 xLSTM 塊,使用 pre up-projection 塊,因為在高維空間中記憶體容量會變大。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

實驗

隨後,研究者對 xLSTM 進行了實驗評估,並將其與現有的語言建模方法進行了比較。

第 4.1 節討論了 xLSTM 在合成任務中的具體能力。首先,研究者測試了 xLSTM 的新指數門控與記憶體混合在形式化語言上的有效性。然後,他們評估了 xLSTM 的新矩陣記憶體在多次查詢聯想記憶任務(Multi-Query Associative Recall,MQAR)中的有效性。最後,研究者評估了 xLSTM 在 Long Range Arena(LRA)中處理長序列的效能。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

第 4.2 節比較了當前各種語言建模方法的驗證集複雜度,包括在同一資料集上對 xLSTM 進行消融研究,然後對不同方法的縮放行為進行評估。

研究者在自迴歸語言建模設定中使用 SlimPajama 的 15B token 訓練了 xLSTM、Transformers、狀態空間模型(SSM)等模型。表 1 中的結果顯示,xLSTM 在驗證複雜度方面優於所有現有方法。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

圖 6 顯示了該實驗的擴充套件結果,表明 xLSTM 對於更大規模的模型也有良好的表現。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

消融研究則表明,效能改進源於指數門控和矩陣記憶體。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

第 4.3 節進行了更深入的語言建模實驗。

研究者增加了訓練資料量,對來自 SlimPajama 的 300B 個 token 進行了訓練,並比較了 xLSTM、RWKV-4、Llama 和 Mamba。他們訓練了不同大小的模型(125M、350M、760M 和 1.3B),進行了深入的評估。首先,評估這些方法在推斷較長語境時的表現;其次,透過驗證易混度和下游任務的表現來測試這些方法;此外,在 PALOMA 語言基準資料集的 571 個文字域上評估了這些方法;最後,評估了不同方法的擴充套件行為,但使用的訓練資料多了 20 倍。

可以看出,xLSTM 在效能和擴充套件性上都更勝一籌。

LSTM捲土重來!xLSTM:一舉超越Mamba、Transformer!

更多研究細節,可參考原論文。

歡迎關注公眾號CV技術指南,專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀、CV招聘資訊。

計算機視覺入門1v3輔導班

【技術文件】《從零搭建pytorch模型教程》122頁PDF下載

QQ交流群:470899183。群內有大佬負責解答大家的日常學習、科研、程式碼問題。

其它文章

分享一個CV知識庫,上千篇文章、專欄,CV所有資料都在這了

明年畢業,還不知道怎麼做畢設的請抓緊機會了

LSKA注意力 | 重新思考和設計大卷積核注意力,效能優於ConvNeXt、SWin、RepLKNet以及VAN

CVPR 2023 | TinyMIM:微軟亞洲研究院用知識蒸餾改進小型ViT

ICCV2023|漲點神器!目標檢測蒸餾學習新方法,浙大、海康威視等提出

ICCV 2023 Oral | 突破性影像融合與分割研究:全時多模態基準與多互動特徵學習

聽我說,Transformer它就是個支援向量機

HDRUNet | 深圳先進院董超團隊提出帶降噪與反量化功能的單幀HDR重建演算法

南科大提出ORCTrack | 解決DeepSORT等跟蹤方法的遮擋問題,即插即用真的很香

1800億引數,世界頂級開源大模型Falcon官宣!碾壓LLaMA 2,效能直逼GPT-4

SAM-Med2D:打破自然影像與醫學影像的領域鴻溝,醫療版 SAM 開源了!

GhostSR|針對影像超分的特徵冗餘,華為諾亞&北大聯合提出GhostSR

Meta推出畫素級動作追蹤模型,簡易版線上可玩 | GitHub 1.4K星

CSUNet | 完美縫合Transformer和CNN,效能達到UNet家族的巔峰!

AI最全資料彙總 | 基礎入門、技術前沿、工業應用、部署框架、實戰教程學習

計算機視覺入門1v3輔導班

計算機視覺交流群

相關文章