GrootVL: Tree Topology is All You Need in State Space Model

论文： 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 树拓扑扫描算法

具体来说，对于输入特征 ，其中 是序列长度（或输入像素的数量）。我们为该特征构建一个无向 连通图 。其中， 是一个超参数，表示邻节点的数量。在视觉任务中 预设为 4 , 表示每个像素都与其四个相邻像素相连。对于语言任务, 我们默认设置 , 表示每个标记都与前三个标记相连。此外，顶点 表示像素（或标记）嵌入集， 表示图的边集。边权重由相邻顶点之间的特征相异性计算得出。此外，顶点之间的距离度量默认使用余弦距离。我们使用收缩Boruvka算法来修剪具有显著差异的边，从而获得一个最小生成树 (MST)，其差异权重之和在所有生成树中最小。在传播过程中，我们迭代遍历每个顶点，将其视为根，并聚合剩余顶点的特征。直观地讲，在这样的几何配置中应用状态传播会优先在空间和特征距离较小的顶点之间进行交互。参考Mamba，我们使用与数据相关的转换矩阵进行状态传播。对于顶点 ，我们对其向父节点的转换矩阵表示为 。因此，第 个顶点的状态聚合过程可以表示为：

其中 表示树中所有顶点的索引集。 表示超边 从MST中第 个顶点到第 个顶点的路径权重, 表示此超边上所有顶点的索引集。对于视觉任务, 我们迭代每个顶点, 将其视为生成树的根, 并聚合来自其他顶点的状态, 从而获得变换后的状态 。对于文本任务, 由于大型语言模型中的因果预测范式, 我们仅将最后一个标记作为根, 并从其他标记中进行聚合。为了实现端到端训练, 我们推导出输出隐藏状态 对输入变量 和 的导数, 如下所示:

其中 表示以顶点 为根且为顶点 的子孙节点集合, 顶点为 的父节点。最后, 输出特征 可以表示为:

其中 和 分别表示 、 和 代表元素乘法（哈达玛积）。

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

对于视觉任务, 树扫描算法需要对整个像素集进行两级遍历, 导致相对于像素数量不可接受的二次复杂度 。为了缓解这个问题，我们利用动态规划过程来加速推理和训练过程，如算法1（图 3）所示，其实现了线性复杂度 的传播，从而在保持计算效率的同时，实现更有效的远距离建模以及空间感知。对于文本任务，我们按照语言的因果性质执行单向聚合方法如算法2（图4）所示。此外, 我们为视觉树扫描和语言树扫描过程提供了反向传播过程以进行端到端的训练, 其详细证明见论文附录C。

2.3 子网络架构设计

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

GrootV. 对于一个形状为