原作者帶隊，LSTM捲土重來之Vision-LSTM出世

机器之心發表於2024-06-09
與 DeiT 等使用 ViT 和 Vision-Mamba (Vim) 方法的模型相比，ViL 的效能更勝一籌。

AI 領域的研究者應該還記得，在 Transformer 誕生後的三年，谷歌將這一自然語言處理屆的重要研究擴充套件到了視覺領域，也就是 Vision Transformer。後來，ViT 被廣泛用作計算機視覺中的通用骨幹。

這種跨界，對於前不久釋出的 xLSTM 來說同樣可以實現。最近，享譽數十年的 LSTM 被擴充套件到一個可擴充套件且效能良好的架構 ——xLSTM，透過指數門控和可並行化的矩陣記憶體結構克服了長期存在的 LSTM 限制。現在，這一成果已經擴充套件到視覺領域。

圖片

xLSTM和 Vision-LSTM 兩項研究均由 LSTM 原作者帶隊，也就是 LSTM 的提出者和奠基者 Sepp Hochreiter。

在最近的這篇論文中，Sepp Hochreiter 等人推出了 Vision-LSTM（ViL）。ViL 包含一堆 xLSTM 塊，其中奇數塊從上到下、偶數塊則從下到上處理補丁 token 序列。

圖片

  • 論文題目：Vision-LSTM: xLSTM as Generic Vision Backbone

  • 論文連結：https://arxiv.org/abs/2406.04303

  • 專案連結: https://nx-ai.github.io/vision-lstm/

正如 xLSTM 誕生之時，作者希望新架構能夠撼動 Transformer 在語言模型領域的江山。這一次，闖入視覺領域的 Vision-LSTM 也被寄予厚望。

研究者在論文中表示：「我們的新架構優於基於 SSM 的視覺架構，也優於 ImageNet-1K 分類中的最佳化 ViT 模型。值得注意的是，在公平的比較中，ViL 的表現優於經過多年超引數調整和 Transformer 改進的 ViT 訓練 pipeline。」

對於需要高解析度影像以獲得最佳效能的任務，如語義分割或醫學成像， ViL 極具應用潛力。在這些情況下，Transformer 因自注意力的二次複雜性而導致計算成本較高，而 ViL 的線性複雜性不存在這種問題。研究者還表示，改進預訓練方案（如透過自監督學習）、探索更好的超引數設定或從 Transformer 中遷移技術（如 LayerScale ）都是 ViL 的可探索方向。

ViT 與 ViL

語言建模架構 —— 如 Transformer 或最近的狀態空間模型 Mamba，通常被應用到計算機視覺領域，以利用它們強大的建模能力。

然而，在自然語言處理中，透過離散詞彙表（Discrete vocabulary），輸入的句子通常被編碼成代表詞或常見子詞的 token。

為了將影像編碼成一組 token，Vision Transformer（ViT）提出將輸入影像分組成非重疊的補丁（例如 16x16 畫素），將它們線性投影成所謂的補丁 token 序列，並向這些 token 新增位置資訊。

然後，這個序列就可以被語言建模架構處理了。

擴充套件長短期記憶（xLSTM）最近被引入作為一種新的語言建模架構，可以說是 LSTM 在 LLM 時代的復興，與 Transformer 和狀態空間模型（SSMs）等相媲美。

現有的 Transformer 或狀態空間模型的視覺版本，例如 ViT 或 Vision Mamba，已經在各種計算機視覺任務中取得了巨大成果。

使用 xLSTM 作為核心組建的 ViL 使用簡單的交替設計，從而可以有效地處理非序列輸入（如影像），而無需引入額外的計算。

類似於 SSMs 的視覺適應，ViL 展示了關於序列長度的線性計算和記憶體複雜度，這使得它在高解析度影像的任務中展現極佳的作用，如醫學成像、分割或物理模擬。

相比之下，ViT 的計算複雜度由於自注意力機制而呈二次方增長，使得它們在應用於高解析度任務時成本高昂。

交替 mLSTM 塊並行

Vision-LSTM（ViL）是一個用於計算機視覺任務的通用骨幹，它從 xLSTM 塊殘差構建，如圖 1 所示。

圖片

與 ViT 類似，ViL 首先透過共享線性投影將影像分割成非重疊的補丁，然後向每個補丁 token 新增可學習的定位嵌入。ViL 的核心是交替的 mLSTM 塊，它們是可完全並行化的，並配備了矩陣記憶體和協方差更新規則。

奇數 mLSTM 塊從左上到右下處理補丁 token，而偶數塊則從右下到左上。

ImageNet-1K 實驗

研究團隊在 ImageNet-1K 上進行了實驗：它包含 130 萬張訓練影像和 5 萬張驗證影像，每張影像屬於 1000 個類別之一。

對比實驗集中在使用序列建模骨幹的模型上，而該模型在大致相當的引數數量上是可比較的。

他們在 224x224 解析度上訓練 ViL 模型，使用餘弦衰減排程，1e-3 的學習率訓練了 800 個週期（tiny, tiny+）或 400 個週期（small, small+, base），具體見下方表 5.

圖片

為了對 Vision Mamba（Vim）進行公平比較，研究人員向模型內新增了額外的塊以匹配 tiny 和小型變體（分別表示為 ViL-T + 和 ViL-S+）的引數數量。

需要注意的是，由於 ViL 以交替的方式遍歷序列，而 Vim 則在每個塊中遍歷序列兩次，因此 ViL 所需的計算量遠少於 Vim。

儘管 Vim 使用了最佳化的 CUDA 核心（而 mLSTM 目前還沒有這樣的核心），但這仍然成立，並且會進一步加速 ViL 的速度。

如表 4 所示的執行時間對比，在其中兩項的比較重，ViL 比 Vim 快了 69%。

圖片

新秀 ViL 相比於 ViTs 如何？

雖然 ViL 首次出場，但仍是展現了極佳的潛力。

由於 ViTs 在視覺社群中已經得到了廣泛的認可，它們在過去幾年經歷了多次最佳化週期。

因為這項工作是首次將 xLSTM 應用於計算機視覺，研究人員並不期望在所有情況下都超過 ViTs 多年的超引數調整。

即便如此，表 1 中的結果顯示，ViL 在小規模上相比於經過大量最佳化的 ViT 協議（DeiT, DeiT-II, DeiT-III）仍是顯示出較良好的結果，其中只有訓練時間是 ViL-S 兩倍的 DeiT-III-S 表現略好一點。

在「base」規模上，ViL 超越了最初的 ViT 模型，並取得了與 DeiT 相當的結果。

需要注意的是：由於在這個規模上訓練模型的成本很高，ViL-B 的超引數遠非最佳。作為參考，訓練 ViL-B 大約需要 600 個 A100 GPU 小時或在 32 個 A100 GPU 上的 19 個小時。

圖片

透過在「長序列微調」設定中微調模型，效能可以進一步提高，該設定透過使用連續補丁 token 之間 50% 的重疊，將序列長度增加到 729，對模型進行 30 個週期的微調。

儘管沒有利用卷積固有的歸納偏置，ViL 還是展現出了與基於 CNN 的模型（如 ConvNeXt）相當的效能。

塊設計

該團隊研究了設計 ViL 塊的不同方式，如圖 2 所示。

  • 普通且單向的 xLSTM 塊未能達到期待的效能，因為 xLSTM 的自迴歸性質並不適合影像分類。

  • 以雙向方式遍歷塊 —— 即在每個塊中引入第二個 mLSTM 層，該層向後遍歷序列（類似於 Vim），提高了效能，但也需要更多的引數和 FLOPS。

  • 共享前向和後向 mLSTM 的引數使模型在引數上更高效，但仍然需要更多的計算並超載這些引數，而這最終也會導致效能下降。

  • 使用交替塊在保持計算和引數效率的同時提高了效能。

該團隊還探索了四向設計，這指的是按行（兩個方向）和按列（兩個方向）遍歷序列。雙向僅按行遍歷序列（兩個方向）。

圖 2 視覺化了不同的遍歷路徑。

圖片

由於雙向和四向塊的成本增加，這項研究是在設定大幅減少的條件中進行的。

研究人員在 128x128 解析度下，對包含僅來自 100 個類別的樣本的 ImageNet-1K 的一個子集進行 400 個週期的訓練。這是特別必要的，因為四向實現方法與 torch.compile（來自 PyTorch 的一個通用速度最佳化方法）不相容，這會導致更長的執行時間，如表 2 最後一列所示。

由於此技術限制，該團隊最終了選擇交替雙向塊作為核心設計。

圖片

分類設計

為了使用 ViT 進行分類，需要將 token 序列彙整合一個 token，然後將其作為分類頭的輸入。

最常見的彙集方法是：(i) 在序列的開頭新增一個可學習的 [CLS] token，或 (ii) 平均所有補丁 token，生成一個 [AVG] token。使用 [CLS] 還是 [AVG] token 通常是一個超引數，兩種變體的效能相當。相反，自迴歸模型通常需要專門的分類設計。例如，Vim 要求 [CLS] token 位於序列的中間，如果採用其他分類設計，如 [AVG] token 或在序列的開始和結束處使用兩個 [CLS] token，則會造成嚴重的效能損失。

基於 ViL 的自迴歸特性，研究者在表 3 中探討了不同的分類設計。

圖片

[AVG] 是所有補丁 token 的平均值，「Middle Patch 」使用中間的補丁 token，「Middle [CLS]」使用序列中間的一個 [CLS] token，「Bilateral [AVG]」使用第一個和最後一個補丁 token 的平均值。

可以發現的是， ViL 分類設計相對穩健，所有效能都在 0.6% 以內。之所以選擇 「Bilateral [AVG]」而不是 「Middle [CLS]」，因為 ImageNet-1K 有中心偏差，即物體通常位於圖片的中間。透過使用 「Bilateral [AVG]」，研究者儘量避免了利用這種偏差，從而使模型更具通用性。

為了與之前使用單個 token 作為分類頭輸入的架構保持可比性，研究者對第一個和最後一個 token 進行了平均處理。為了達到最佳效能，建議將兩個標記合併（「Bilateral Concat」），而不是取平均值。

這類似於 DINOv2 等自監督 ViT 的常見做法，這些是透過分別附加在 [CLS] 和 [AVG] token 的兩個目標來進行訓練的，因此可以從連線 [CLS] 和 [AVG] token 的表徵中獲益。視覺 SSM 模型也探索了這一方向，即在序列中分散多個 [CLS] token，然後將其作為分類器的輸入。此外，類似的方向也可以提高 ViL 的效能。

更多研究細節，請參考原論文。

相關文章