【源头活水】让Mamba涨点！树形SSM来了！清华和腾讯提出GrootVL：多功能多模态框架

“问渠那得清如许，为有源头活水来”，通过前沿领域知识的学习，从其他研究领域得到启发，对研究问题的本质有更清晰的认识和理解，是自我提高的不竭源泉。为此，我们特别精选论文阅读笔记，开辟“源头活水”专栏，帮助你广泛而深入的阅读科研文献，敬请关注！

论文：https://arxiv.org/pdf/2406.02395

代码（已开源）：

https://github.com/EasonXiao-888/GrootVL

【Highlight】：

使用递归范式传播特征的状态空间模型展示了与 Transformer 模型相当的强大表征能力和卓越的推理效率。然而，受序列固有几何约束的限制，它在建模远程依赖方面稍显不足。为了解决这个问题，我们提出了 GrootVL 网络，它首先基于空间关系和输入特征动态生成树形拓扑结构。然后，基于该无环图执行特征传播，从而打破原始序列约束以实现更强的表征能力。此外，我们引入了一种线性复杂的度动态规划算法，在不增加计算成本的情况下增强远程交互。GrootVL 是一个多功能的多模态框架，可以应用于视觉和文本任务。大量实验表明，我们的方法在图像分类、目标检测和分割方面明显优于现有的结构化状态空间模型。此外，通过微调大语言模型，我们的方法在较小的训练成本下在多个文本任务中取得了一致的性能提升。

1 回顾状态空间模型

主流基础模型主要基于 CNN和 Transformer 架构，它们在视觉和语言任务中占主导地位。然而，CNN 的小感受野和 Transformer 的高度复杂性给平衡性能和效率带来挑战。状态空间模型(SSM)试图打破这种僵局，它以循环范式对序列进行建模。与之前的循环神经网络不同，这些方法利用结构参数初始化实现稳定的优化和优越的计算性能。然而，它仍然容易受到循环神经网络共有的内在缺陷的影响，即捕获远程依赖的缺陷。

最近，Mamba[1]提出了一种改进的选择机制来减轻SSM遇到的挑战。这种方法在传播过程中引入了权重调制，扩大了有效感受野，并在 NLP 任务中取得了令人印象深刻的效果。此外，许多研究[2][3][4]旨在通过使用各种人工预定义的策略将 2D 图像特征映射到 1D 序列，从而将 Mamba 扩展到计算机视觉中。

尽管这些方法成功地将 Mamba 嵌入视觉信号输入，但如图1(a)所示，光栅扫描和局部扫描策略都引入了相邻像素之间的空间不连续，但Mamba中的特征变换依赖于特征关系，因此这种扫描范式会阻碍序列中有效的信息流。此外，连续扫描策略试图通过简单地调整不连续位置的传播方向来缓解这个问题。然而，所有这些方法都依赖于固定的传播轨迹，忽略了固有的空间结构，不能根据输入动态调整拓扑。

2 基于自适应树形拓扑结构的状态空间模型（Tree SSM）

本文试图探索一个新的视角:为状态空间模型中的特征传播引入输入感知的拓扑网络。为了实现这一点，我们开发了一个树状态空间模型如图2所示，并提出了一个称为 GrootVL 的新框架，该框架根据输入特征自适应地生成树拓扑结构，然后在其上执行特征传播。具体来说，我们分别为视觉和语言任务设计了两个子网络 GrootV 和 GrootL，如图 1(b) 和图 1(d) 所示。对于视觉任务，我们首先基于相邻特征之间的差异在四连通平面图上构造一个最小生成树。这个过程可以自适应地将空间和语义信息编码到树结构中。然后，我们迭代地遍历每个像素，将其视为根顶点，并使用 Mamba 的状态转移方式聚合其他像素的特征。直观地说，此操作需要对整个像素集进行两级遍历，导致相对于视觉像素级输入信号不可接受的二次复杂度。然而，鉴于树图是非循环的，我们提出了一种动态规划算法来实现线性复杂度的传播。通过这种输入感知树拓扑，我们的方法能够实现更有效的远程交互，并保持与 Mamba 一致的线性复杂度。此外，我们的方法还可以通过基于标记特征之间的差异构建树结构来应用于语言任务，这克服了文本序列的几何约束。使用与 GrootV 类似的聚合过程，GrootL 可以显着提高预训练大语言模型的语言表征能力。

2.1 树拓扑扫描算法

、

、

2.2 适应于视觉和文本信号的高效部署

2.3 子网络架构设计

基于上述基础算子，我们可以将其有效应用于各类视觉和文本任务。

GrootL. 我们基于树扫描算法提出了一种有效的应用于预训练Mamba的微调范式。具体来说，基于树的拓扑分支与单向滚动通过调制因子来合并。值得注意的是，这种范式不会引入任何额外的训练参数。相反，它利用预训练的状态转换参数通过结合拓扑结构以进行语义聚合。实验结果表明了我们方法的有效性。

3 实验结果与分析

3.1 视觉任务

我们在图像分类，目标检测，实例分割与语义分割下游视觉任务上进行验证。如图6、7、8中结果所示，我们的方法显著优于各类基于状态空间模型的方法，并能与最先进的基于CNN或者Transformer架构的方法相媲美。

3.2 文本任务

我们将130M参数量的 Mamba作为基础模型。为了验证我们的 GrootL 在自然语言理解中的有效性，我们首先用相同的训练设置在Alpaca 数据上分别通过LoRA 和 GrootL对预训练Mamba进行微调。接着在公开lm-evaluation-harness项目中对多个NLP benchmark进行zero-shot验证，结果如图9所示。

3.3 消融实验

我们对树拓扑算法中多个配置进行消融实验。我们首先验证了树拓扑扫描相对于人工预设的光栅扫描和交叉扫描方式的优越性。接着，我们探索了相邻节点间边权重的刻画方式，尝试了曼哈顿距离，欧式距离和余弦距离三种方式。最后，我们仅针对单个节点进行状态聚合进行了实验从而验证了每个节点均需获取来自其余顶点信息的必要性。

3.4 定性分析

为了更好地说明我们的扫描策略的优越性，我们可视化了每个输入图像中红色十字标记的不同位置的亲和图。例如，当我们设置天空左上角的锚点，如图11(a)第二行所示。我们的方法可以很容易地识别白色房屋、旗杆和天空，光栅扫描则无法实现。这证明了我们的算法具有保留详细结构信息的能力。更多的对比结果可见于论文的附录D中。