AI 領(lǐng)域的研究者應(yīng)該還記得，在 Transformer 誕生后的三年，谷歌將這一自然語(yǔ)言處理屆的重要研究擴(kuò)展到了視覺(jué)領(lǐng)域，也就是 Vision Transformer。后來(lái)，ViT 被廣泛用作計(jì)算機(jī)視覺(jué)中的通用骨干。

這種跨界，對(duì)于前不久發(fā)布的 xLSTM 來(lái)說(shuō)同樣可以實(shí)現(xiàn)。最近，享譽(yù)數(shù)十年的 LSTM 被擴(kuò)展到一個(gè)可擴(kuò)展且性能良好的架構(gòu) ——xLSTM，通過(guò)指數(shù)門(mén)控和可并行化的矩陣內(nèi)存結(jié)構(gòu)克服了長(zhǎng)期存在的 LSTM 限制?，F(xiàn)在，這一成果已經(jīng)擴(kuò)展到視覺(jué)領(lǐng)域。

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

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

• 論文題目：Vision-LSTM: xLSTM as Generic Vision Backbone
• 論文鏈接：https://arxiv.org/abs/2406.04303
• 項(xiàng)目鏈接: https://nx-ai.github.io/vision-lstm/

正如 xLSTM 誕生之時(shí)，作者希望新架構(gòu)能夠撼動(dòng) Transformer 在語(yǔ)言模型領(lǐng)域的江山。這一次，闖入視覺(jué)領(lǐng)域的 Vision-LSTM 也被寄予厚望。

研究者在論文中表示：「我們的新架構(gòu)優(yōu)于基于 SSM 的視覺(jué)架構(gòu)，也優(yōu)于 ImageNet-1K 分類(lèi)中的優(yōu)化 ViT 模型。值得注意的是，在公平的比較中，ViL 的表現(xiàn)優(yōu)于經(jīng)過(guò)多年超參數(shù)調(diào)整和 Transformer 改進(jìn)的 ViT 訓(xùn)練 pipeline?！?/span>

對(duì)于需要高分辨率圖像以獲得最佳性能的任務(wù)，如語(yǔ)義分割或醫(yī)學(xué)成像， ViL 極具應(yīng)用潛力。在這些情況下，Transformer 因自注意力的二次復(fù)雜性而導(dǎo)致計(jì)算成本較高，而 ViL 的線性復(fù)雜性不存在這種問(wèn)題。研究者還表示，改進(jìn)預(yù)訓(xùn)練方案（如通過(guò)自監(jiān)督學(xué)習(xí)）、探索更好的超參數(shù)設(shè)置或從 Transformer 中遷移技術(shù)（如 LayerScale ）都是 ViL 的可探索方向。

ViT 與 ViL

語(yǔ)言建模架構(gòu) —— 如 Transformer 或最近的狀態(tài)空間模型 Mamba，通常被應(yīng)用到計(jì)算機(jī)視覺(jué)領(lǐng)域，以利用它們強(qiáng)大的建模能力。

然而，在自然語(yǔ)言處理中，通過(guò)離散詞匯表（Discrete vocabulary），輸入的句子通常被編碼成代表詞或常見(jiàn)子詞的 token。

為了將圖像編碼成一組 token，Vision Transformer（ViT）提出將輸入圖像分組成非重疊的補(bǔ)?。ɡ?16x16 像素），將它們線性投影成所謂的補(bǔ)丁 token 序列，并向這些 token 添加位置信息。

然后，這個(gè)序列就可以被語(yǔ)言建模架構(gòu)處理了。

擴(kuò)展長(zhǎng)短期記憶（xLSTM）最近被引入作為一種新的語(yǔ)言建模架構(gòu)，可以說(shuō)是 LSTM 在 LLM 時(shí)代的復(fù)興，與 Transformer 和狀態(tài)空間模型（SSMs）等相媲美。

現(xiàn)有的 Transformer 或狀態(tài)空間模型的視覺(jué)版本，例如 ViT 或 Vision Mamba，已經(jīng)在各種計(jì)算機(jī)視覺(jué)任務(wù)中取得了巨大成果。

使用 xLSTM 作為核心組建的 ViL 使用簡(jiǎn)單的交替設(shè)計(jì)，從而可以有效地處理非序列輸入（如圖像），而無(wú)需引入額外的計(jì)算。

類(lèi)似于 SSMs 的視覺(jué)適應(yīng)，ViL 展示了關(guān)于序列長(zhǎng)度的線性計(jì)算和內(nèi)存復(fù)雜度，這使得它在高分辨率圖像的任務(wù)中展現(xiàn)極佳的作用，如醫(yī)學(xué)成像、分割或物理模擬。

相比之下，ViT 的計(jì)算復(fù)雜度由于自注意力機(jī)制而呈二次方增長(zhǎng)，使得它們?cè)趹?yīng)用于高分辨率任務(wù)時(shí)成本高昂。

交替 mLSTM 塊并行

Vision-LSTM（ViL）是一個(gè)用于計(jì)算機(jī)視覺(jué)任務(wù)的通用骨干，它從 xLSTM 塊殘差構(gòu)建，如圖 1 所示。

與 ViT 類(lèi)似，ViL 首先通過(guò)共享線性投影將圖像分割成非重疊的補(bǔ)丁，然后向每個(gè)補(bǔ)丁 token 添加可學(xué)習(xí)的定位嵌入。ViL 的核心是交替的 mLSTM 塊，它們是可完全并行化的，并配備了矩陣內(nèi)存和協(xié)方差更新規(guī)則。

奇數(shù) mLSTM 塊從左上到右下處理補(bǔ)丁 token，而偶數(shù)塊則從右下到左上。

ImageNet-1K 實(shí)驗(yàn)

研究團(tuán)隊(duì)在 ImageNet-1K 上進(jìn)行了實(shí)驗(yàn)：它包含 130 萬(wàn)張訓(xùn)練圖像和 5 萬(wàn)張驗(yàn)證圖像，每張圖像屬于 1000 個(gè)類(lèi)別之一。

對(duì)比實(shí)驗(yàn)集中在使用序列建模骨干的模型上，而該模型在大致相當(dāng)?shù)膮?shù)數(shù)量上是可比較的。

他們?cè)?224x224 分辨率上訓(xùn)練 ViL 模型，使用余弦衰減調(diào)度，1e-3 的學(xué)習(xí)率訓(xùn)練了 800 個(gè)周期（tiny, tiny+）或 400 個(gè)周期（small, small+, base），具體見(jiàn)下方表 5.

為了對(duì) Vision Mamba（Vim）進(jìn)行公平比較，研究人員向模型內(nèi)添加了額外的塊以匹配 tiny 和小型變體（分別表示為 ViL-T + 和 ViL-S+）的參數(shù)數(shù)量。

需要注意的是，由于 ViL 以交替的方式遍歷序列，而 Vim 則在每個(gè)塊中遍歷序列兩次，因此 ViL 所需的計(jì)算量遠(yuǎn)少于 Vim。

盡管 Vim 使用了優(yōu)化的 CUDA 內(nèi)核（而 mLSTM 目前還沒(méi)有這樣的內(nèi)核），但這仍然成立，并且會(huì)進(jìn)一步加速 ViL 的速度。

如表 4 所示的運(yùn)行時(shí)間對(duì)比，在其中兩項(xiàng)的比較重，ViL 比 Vim 快了 69%。

新秀 ViL 相比于 ViTs 如何？

雖然 ViL 首次出場(chǎng)，但仍是展現(xiàn)了極佳的潛力。

由于 ViTs 在視覺(jué)社區(qū)中已經(jīng)得到了廣泛的認(rèn)可，它們?cè)谶^(guò)去幾年經(jīng)歷了多次優(yōu)化周期。

因?yàn)檫@項(xiàng)工作是首次將 xLSTM 應(yīng)用于計(jì)算機(jī)視覺(jué)，研究人員并不期望在所有情況下都超過(guò) ViTs 多年的超參數(shù)調(diào)整。

即便如此，表 1 中的結(jié)果顯示，ViL 在小規(guī)模上相比于經(jīng)過(guò)大量?jī)?yōu)化的 ViT 協(xié)議（DeiT, DeiT-II, DeiT-III）仍是顯示出較良好的結(jié)果，其中只有訓(xùn)練時(shí)間是 ViL-S 兩倍的 DeiT-III-S 表現(xiàn)略好一點(diǎn)。

在「base」規(guī)模上，ViL 超越了最初的 ViT 模型，并取得了與 DeiT 相當(dāng)?shù)慕Y(jié)果。

需要注意的是：由于在這個(gè)規(guī)模上訓(xùn)練模型的成本很高，ViL-B 的超參數(shù)遠(yuǎn)非最佳。作為參考，訓(xùn)練 ViL-B 大約需要 600 個(gè) A100 GPU 小時(shí)或在 32 個(gè) A100 GPU 上的 19 個(gè)小時(shí)。

通過(guò)在「長(zhǎng)序列微調(diào)」設(shè)置中微調(diào)模型，性能可以進(jìn)一步提高，該設(shè)置通過(guò)使用連續(xù)補(bǔ)丁 token 之間 50% 的重疊，將序列長(zhǎng)度增加到 729，對(duì)模型進(jìn)行 30 個(gè)周期的微調(diào)。

盡管沒(méi)有利用卷積固有的歸納偏置，ViL 還是展現(xiàn)出了與基于 CNN 的模型（如 ConvNeXt）相當(dāng)?shù)男阅堋?/span>

塊設(shè)計(jì)

該團(tuán)隊(duì)研究了設(shè)計(jì) ViL 塊的不同方式，如圖 2 所示。

•  普通且單向的 xLSTM 塊未能達(dá)到期待的性能，因?yàn)?xLSTM 的自回歸性質(zhì)并不適合圖像分類(lèi)。
•  以雙向方式遍歷塊 —— 即在每個(gè)塊中引入第二個(gè) mLSTM 層，該層向后遍歷序列（類(lèi)似于 Vim），提高了性能，但也需要更多的參數(shù)和 FLOPS。
•  共享前向和后向 mLSTM 的參數(shù)使模型在參數(shù)上更高效，但仍然需要更多的計(jì)算并超載這些參數(shù)，而這最終也會(huì)導(dǎo)致性能下降。
•  使用交替塊在保持計(jì)算和參數(shù)效率的同時(shí)提高了性能。

該團(tuán)隊(duì)還探索了四向設(shè)計(jì)，這指的是按行（兩個(gè)方向）和按列（兩個(gè)方向）遍歷序列。雙向僅按行遍歷序列（兩個(gè)方向）。

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

圖片

由于雙向和四向塊的成本增加，這項(xiàng)研究是在設(shè)置大幅減少的條件中進(jìn)行的。

研究人員在 128x128 分辨率下，對(duì)包含僅來(lái)自 100 個(gè)類(lèi)別的樣本的 ImageNet-1K 的一個(gè)子集進(jìn)行 400 個(gè)周期的訓(xùn)練。這是特別必要的，因?yàn)樗南驅(qū)崿F(xiàn)方法與 torch.compile（來(lái)自 PyTorch 的一個(gè)通用速度優(yōu)化方法）不兼容，這會(huì)導(dǎo)致更長(zhǎng)的運(yùn)行時(shí)間，如表 2 最后一列所示。

由于此技術(shù)限制，該團(tuán)隊(duì)最終了選擇交替雙向塊作為核心設(shè)計(jì)。

分類(lèi)設(shè)計(jì)

為了使用 ViT 進(jìn)行分類(lèi)，需要將 token 序列匯集成一個(gè) token，然后將其作為分類(lèi)頭的輸入。

最常見(jiàn)的匯集方法是：(i) 在序列的開(kāi)頭添加一個(gè)可學(xué)習(xí)的 [CLS] token，或 (ii) 平均所有補(bǔ)丁 token，生成一個(gè) [AVG] token。使用 [CLS] 還是 [AVG] token 通常是一個(gè)超參數(shù)，兩種變體的性能相當(dāng)。相反，自回歸模型通常需要專(zhuān)門(mén)的分類(lèi)設(shè)計(jì)。例如，Vim 要求 [CLS] token 位于序列的中間，如果采用其他分類(lèi)設(shè)計(jì)，如 [AVG] token 或在序列的開(kāi)始和結(jié)束處使用兩個(gè) [CLS] token，則會(huì)造成嚴(yán)重的性能損失。

基于 ViL 的自回歸特性，研究者在表 3 中探討了不同的分類(lèi)設(shè)計(jì)。

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

可以發(fā)現(xiàn)的是， ViL 分類(lèi)設(shè)計(jì)相對(duì)穩(wěn)健，所有性能都在 0.6% 以?xún)?nèi)。之所以選擇 「Bilateral [AVG]」而不是 「Middle [CLS]」，因?yàn)?ImageNet-1K 有中心偏差，即物體通常位于圖片的中間。通過(guò)使用 「Bilateral [AVG]」，研究者盡量避免了利用這種偏差，從而使模型更具通用性。

為了與之前使用單個(gè) token 作為分類(lèi)頭輸入的架構(gòu)保持可比性，研究者對(duì)第一個(gè)和最后一個(gè) token 進(jìn)行了平均處理。為了達(dá)到最佳性能，建議將兩個(gè)標(biāo)記合并（「Bilateral Concat」），而不是取平均值。

這類(lèi)似于 DINOv2 等自監(jiān)督 ViT 的常見(jiàn)做法，這些是通過(guò)分別附加在 [CLS] 和 [AVG] token 的兩個(gè)目標(biāo)來(lái)進(jìn)行訓(xùn)練的，因此可以從連接 [CLS] 和 [AVG] token 的表征中獲益。視覺(jué) SSM 模型也探索了這一方向，即在序列中分散多個(gè) [CLS] token，然后將其作為分類(lèi)器的輸入。此外，類(lèi)似的方向也可以提高 ViL 的性能。

更多研究細(xì)節(jié)，請(qǐng)參考原論文。