原作者带队,LSTM卷土重来之Vision-LSTM出世
与 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 的性能。
更多研究细节,请参考原论文。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道:[email protected]
微信扫码关注该文公众号作者