與 DeiT 等使用 ViT 和 Vision-Mamba (Vim) 方法的模型相比，ViL 的效能更勝一籌。
AI 領域的研究者應該還記得，在 Transformer 誕生後的三年，谷歌將這一自然語言處理屆的重要研究擴充套件到了視覺領域，也就是 Vision Transformer。後來，ViT 被廣泛用作計算機視覺中的通用骨幹。
這種跨界，對於前不久釋出的 xLSTM 來說同樣可以實現。最近，享譽數十年的 LSTM 被擴充套件到一個可擴充套件且效能良好的架構 ——xLSTM，透過指數門控和可並行化的矩陣記憶體結構克服了長期存在的 LSTM 限制。現在，這一成果已經擴充套件到視覺領域。
在最近的這篇論文中，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/
對於需要高解析度影像以獲得最佳效能的任務，如語義分割或醫學成像， 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 的效能。
更多研究細節，請參考原論文。