文章介绍了一种称为VMRNN（Vision Mamba RNN）的新模型，该模型通过整合Vision Mamba模块与长短期记忆网络（LSTM），以应对时空预测挑战，特别是在处理视频序列预测等任务时，能够有效地建模长依赖关系并保持计算效率。文章强调了传统的卷积神经网络（CNNs）和视觉变换器（ViTs）在处理此类任务时的局限性，如感受野受限及计算需求高，并展示了VMRNN在网络规模较小的情况下，在多种时空预测任务中取得了有竞争力的结果。
1 VMRNN的架构
一个基于VMRNN Cell的基本模型（VMRNN-B）和一个更深层的模型（VMRNN-D）。在每个时间步骤中，图像被分割为非重叠的补丁，并通过展平和初步线性转换进入后续处理阶段。
1.VMRNN-B模型：
- 处理流程：VMRNN层接收嵌入后的图像patch以及前一时刻的状态信息（隐藏状态Ht-1和细胞状态Ct-1），进而生成当前的隐藏状态Ht和细胞状态Ct。
- 多用途隐藏状态：Ht被复制用于两个目的，一是送入重构层，二是与Ct一起为下一时间步的VMRNN层提供输入。
2.VMRNN-D模型：
- 深度扩展：相比VMRNN-B，VMRNN-D包含更多的VMRNN单元，并引入了Patch Merging和Patch Expanding层。
- 下采样：Patch Merging层用于减少空间维度，降低计算复杂度，同时提取更加抽象的全局特征。
- 上采样：Patch Expanding层则用于增加空间维度，恢复细节，并在重构阶段准确定位特征。
VSS Block（Visual Spatial Selective Block）
- 输入处理：输入首先通过一个线性嵌入层进行处理，这一步骤将输入的形状从[B, L, C]转换为[B, H, W, C]，即将一维的补丁标记转换回二维图像的形式。
- 2D-selective-scan (SS2D)：之后，VSS块利用2D-selective-scan来处理2D图像数据，该方法通过四个不同的方向展开图像补丁：从左上角到右下角、从右下角到左上角、从右上角到左下角、以及从左下角到右上角。这样创建出四组不同的序列。
- S6 Block处理：每一组特征序列会通过S6块处理，S6块是一种SSM（Selective Scan Space State Sequential Model）运算符，允许一维数组中的每一个元素通过压缩的隐藏状态与之前扫描的样本互动。
- 重组：最后，这些序列会重新组合回单个图像中，这一过程称为Scan Merge Stage。
- 输出：给定输入特征z，SS2D的输出特征z¯可以通过一系列公式计算得出，其中expand(·)对应于scan expand操作，merge(·)对应于scan merge操作。
2 结语
本文提出了一种名为VMRNN的新架构，通过融合Vision Mamba模块与LSTM，有效解决了视频级别未来帧预测任务中的时空动态捕捉问题，并在多个数据集上展示了其优越的性能和效率。
论文题目：VMRNN: Integrating Vision Mamba and LSTM for Efficient and Accurate Spatiotemporal Forecasting
论文链接：https://arxiv.org/abs/2403.16536
PS: 欢迎大家扫码关注公众号_，我们一起在AI的世界中探索前行，期待共同进步！
精彩回顾
1. 从图灵测试到“通”测试
