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

Although SSM serves as a powerful tool in systems and control engineering, its time-continuous nature poses challenges for integration into deep learning architectures. To alleviate this issue, most methods utilize the zero-order hold rule [18] to discretize the continuous system described by Eq. 1 and convert continuous variables ($\mathbf{A}$, $\mathbf{B}$, $\mathbf{C}$, $\mathbf{D}$) into corresponding discrete parameters ($\bar{\mathbf{A}}$, $\bar{\mathbf{B}}$, $\bar{\mathbf{C}}$, $\bar{\mathbf{D}}$) over the specified sampling time-scale $\Delta\in\mathbb{R}^{D}$:

In addition, many improved methods [38, 18] use an approximation of $\bar{\mathbf{B}}$ based on the first-order Taylor Series:

Previous SSMs store information through finite states and inherent time-invariance, which limits their effectiveness. Therefore, Mamba [18] introduces a dynamic mechanism to selectively filter out input into a sequential state. Specifically, it utilizes Linear Projection to calculate the parameters $\{\mathbf{B}_{i}\}_{i=1}^{L}$, $\{\mathbf{C}_{i}\}_{i=1}^{L}$ and $\{\mathbf{\Delta}_{i}\}_{i=1}^{L}$ from the input sequence $\{x_{i}\}_{i=1}^{L}$ with $x_{i}\in\mathbb{R}^{1\times D}$ directly to improve the context-aware ability. Then the output sequence $\{y_{i}\}_{i=1}^{L}$ can be computed with those input-adaptive discretized parameters as follows:

Given an input feature $X=\{x_{i}\}_{i=1}^{L}$ where $L$ is the sequence length (or the number of input pixels), we construct an undirected $m$-connected graph $G=(V,E)$ for the feature. $m$ is a hyper-parameter that indicates the number of adjacent tokens. Following [64, 50], we set $m=4$ for visual tasks, meaning each pixel is connected to its four neighboring pixels. For language tasks, we set $m=3$ by default, meaning each token is connected to the previous three tokens. In addition, the vertices $V$ represent the pixel (or token) embeddings, and the $E$ indicates the edges of the graph. The edge weight is calculated by the feature dissimilarity between adjacent vertices. Besides, the metric of dissimilarity uses cosine distance by default, and the comparison with other metrics refers to Table 6.

We use the Contractive Boruvka algorithm [2] to prune the edges with significant dissimilarity, which generates a minimum spanning tree (MST) $\mathcal{G}_{T}$ whose sum of dissimilarity weights is minimum out of all spanning trees. In the propagation process, we iteratively traverse each vertex, treating it as the root, and aggregate the features of the remaining vertices. Intuitively, applying state propagation within such a geometric configuration makes its preferential interactions among vertices with small spatial and feature distances. Following the Mamba, we employ the data-dependent transition matrix for state propagation. For a vertex $k$, we denote the transition matrix with its parent as $\bar{\mathbf{A}}_{k}$. Furthermore, following the Eq. 4, the state aggregation process for the $i$-th vertex can be formulated as:

where $\Omega$ denotes the index set of all vertices in the tree. $S(E_{ij})$ represents the path weight of hyperedge $E_{ij}$ traced from $j$-th vertex to $i$-th vertex in the tree $\mathcal{G}_{T}$, and $N_{ij}$ indicates the index set of all vertices on this hyperedge. For visual tasks, we iterate over each vertex, treating it as the root of the spanning tree $\mathcal{G}_{T}$, and aggregate the states from the other vertices, thereby obtaining the transformed states $\{h_{i}\}_{i=1}^{L}$. For textual tasks, because of the causal prediction manner in large language models, we only take the last token as root and aggregate from other tokens. To achieve end-to-end training, we derive the derivative of the output hidden state $h_{i}$ to the input variables $\bar{\mathbf{A}}_{k}$, $\bar{\mathbf{B}}_{j}$ and $x_{j}$ as follows:

where $C_{k}^{i}$ indicates the children of vertex $k$ in tree $\mathcal{G}_{T}$ whose root is the vertex $i$, and $n$ denotes the parent of vertex $k$ in Eq. 7. Finally, the output feature $Y$ can be formulated as:

where $Y$, $H$ and $X$ indicate the stack of $\{y_{i}\}_{i=1}^{L}$, $\{h_{i}\}_{i=1}^{L}$ and $\{x\}_{i=1}^{L}$ respectively. $\odot$ denotes the element-wise multiplication.

For visual tasks, the tree scanning algorithm requires two levels of traversal across the entire pixel set, resulting in an unacceptable quadratic complexity relative to the number of pixels $\mathcal{O}(L^{2})$. To alleviate this issue, we utilize a dynamic programming procedure to accelerate the inference and training processes as elaborated in Algorithm 1, which results in linear complexity $\mathcal{O}(L)$. For textual tasks, we perform a unidirectional aggregation approach (shown in Algorithm 2 of Appendix B) in adherence to the causal nature of language. Moreover, we provide the back-propagation process for both Vision Tree Scanning and Language Tree Scanning processes, whose detailed proofs refer to Appendix C.

Given an image with a shape of $H\times W\times 3$, our goal is to obtain high-quality visual features for downstream tasks. To this end, we propose an effective vision architecture GrootV which consists of a stem module, several basic blocks and downsampling layers to generate hierarchical representations illustrated in Fig. 3. Overall, our GrootV comprises four stages similar to previous general vision backbones [41, 40, 57, 38]. We integrate the stem module before the first stage to decrease the resolution of the input image signal by a factor of $4$, resulting in a feature map with a shape of $\frac{H}{4}\times\frac{W}{4}\times C$. It includes two convolutions, two Layer Normalization (LN) layers and one GELU activation function. The kernel size for both convolutions is $3$ with a stride of $2$ and padding of $1$. Similarly, a downsampling layer consists of a $3\times 3$ convolution with a stride of $2$ and padding of $1$ and an LN layer. Positioned between two stages, it serves to downsample the input feature map by a factor of $2$. Motivated by [57, 38], we devise a residual block with skip connections to integrate our fundamental Tree State Space Model in Sec. 3.2. In detail, we first normalize the input features with LN layer. Spatial priors and long-range dependencies are then obtained through our tree scanning algorithm with residual connections established alongside the input features. Finally, a feedforward neural network is utilized to project the normalized features to output signals as shown in Fig. 3. Based on the above origin components, we develop our GrootV in three scales, $i.e.$, GrootV-Tiny, GrootV-Small and GrootV-Base.

Recurrent neural networks rely on fixed memory to preserve past information, which poses limitations when handling long contexts where relevant words are distant from the current moment. While Mamba [18] employs a selection mechanism to enhance context awareness, its fixed memory size cannot expand over time, resulting in restricted state space. Therefore, the ability to extrapolate decreases during scrolling as the prompt extends. To mitigate this issue, we propose an effective fine-tuning paradigm. Specifically, the tree-based topology branch is built upon one-way scrolling with a scaling factor, enabling state transitions within such a structure. This arrangement facilitates the preferential interaction of semantically related tokens. It is noteworthy that this paradigm does not introduce any additional training parameters. Instead, it utilizes pretrained state transformation parameters to conduct semantic aggregation by incorporating topological structures. Experimental results demonstrate the effectiveness of our approach.

We assess the classification performance of GrootV on the ImageNet-1k dataset [12]. Following previous practices [40, 41, 57, 38], all GrootV models are trained for $300$ epochs from scratch using AdamW optimizer with a warm-up strategy of $20$ epochs. During training, we utilize a Cosine Scheduler with an initial learning rate of $1\times 10^{-3}$ and weight decay of $0.05$. In addition, the exponential moving average (EMA) is also applied.

The comparison results summarized in Table 1 show GrootV leading all SSM-based models and competitive with advanced CNNs and Transformers across tiny, small, and base scales. Specifically, GrootV-T achieves $83.4\%$ Top-1 Acc. boosting ViM-S by $2.9\%$, LocalVim-S by $2.2\%$, PlainMamba-L2 by $1.8\%$ and VMamba-T by $0.9\%$ with similar FLOPs. Additionally, it surpasses ConvNeXt-T by $1.3\%$ and Swin-T by $2.2\%$, demonstrating the effectiveness of our method.

We verify the detection performance of GrootV on the MSCOCO 2017 dataset [36] with MMDetection library [3]. We follow previous works [38, 57, 40, 31, 49, 51, 67, 63, 6] to validate object detection and instance segmentation tasks with Mask-RCNN [26]. Specifically, We adopt the AdamW optimizer with a learning rate of $1\times 10^{-4}$ and batch size of $16$ to optimize the model built upon our pre-trained classification backbones on ImageNet-1K. The training schedules include $1\times$ ($12$ epochs) and $3\times$ ($36$ epochs) with multi-scale data augmentation.

As depicted in Table 7 (in Appendix A.), our method outperforms existing methods on most evaluation metrics, especially for instance segmentation. Under $1\times$ schedule, GrootV-T achieves $47.0$ in box mAP (AP^b), which is $1.1$ points higher than ViM-S and $0.5$ points higher than VMamba-T. It is worth noting that GrootV-T outperforms ViM-S by $1.7$ points with $1\times$ schedule and LocalVMamba-T by $0.4$ points with $3\times$ schedule in mask mAP (AP^m). Moreover, the best AP^b $50.1$ and AP^m $44.6$ are obtained by GrootV-S in $3\times$ schedule with multi-scale training.

To evaluate the semantic segmentation performance of our GrootV series, we train our models with UperNet [60] initialized by pre-trained classification weights on ADE20K[68] for 160k iterations, following common practices without additional augmentations for fair comparison.

Our method performs exceptionally well on segmentation tasks shown in Fig. 4. GrootV-T yields a clear improvement of $+3.6$ in single-scale mIoU compared to ViM-S and $+1.9$ in multi-scale mIoU compared to LocalViM-S. Furthermore, GrootV-S boosts InterImage-S by $0.6$ and $0.8$ in single-scale and multi-scale respectively. We consider the preservation of intricate structural details through tree topology scanning to be particularly advantageous for segmentation tasks that require pixel-level perception.

We conduct a head-to-head comparison of different scanning strategies, as shown in Table 6. The tree topology scanning outperforms previous strategies by $0.8\%$ and $0.3\%$, highlighting the superiority of our algorithm in vision recognition.

Before generating a minimum spanning tree from a connected graph, it is important to measure the edge weights between vertices. Therefore, we validate several distance metrics as illustrated in Table 6. The results indicate that $Cosine$ distance most effectively represents the relationship between vertices, performing $0.5\%$ better than $Manhattan$ and $0.2\%$ better than $Euclidean$.

We traverse all vertices, treating each as a root, and perform state transitions along the topological path from the other vertices toward the root. This traversal ensures that each vertex captures long-range dependencies. To verify the effectiveness of this operation, we consider only the first and last vertices as the root in Table 6. The results show reductions of $0.5\%$ and $0.4\%$, respectively.

To better illustrate the superiority of our scanning strategy, we visualize the affinity maps of different positions marked by the red cross in each input image. For example, we set the anchor point in the upper left corner of the sky as shown in the second row of in Fig. 4(a). Our method can easily identify white houses, flagpoles, and the sky, which raster scanning fails to achieve. This demonstrates the capability of our algorithm to preserve detailed structural information. More comparisons can be seen in Fig. 6 (in Appendix D.)