MambaTree: Tree Topology is All You Need in State Space Model
Authors: Yicheng Xiao
Lin Song
Shaoli Huang
Jiangshan Wang
Siyu Song
Yixiao Ge
Xiu Li
Ying Shan
Venue: NeurIPS 2024 spotlight (2024)
Keywords: State Space Model
Vision Recognition
Language Understanding
Primary area: deep_learning_architectures

Abstract: The state space models, employing recursively propagated features, demonstrate strong representation capabilities comparable to Transformer models and superior efficiency.
However, constrained by the inherent geometric constraints of sequences, it still falls short in modeling long-range dependencies.
To address this issue, we propose the MambaTree network, which first dynamically generates a tree topology based on spatial relationships and input features.
Then, feature propagation is performed based on this graph, thereby breaking the original sequence constraints to achieve stronger representation capabilities.
Additionally, we introduce a linear complexity dynamic programming algorithm to enhance long-range interactions without increasing computational cost.
MambaTree is a versatile multimodal framework that can be applied to both visual and textual tasks.
Extensive experiments demonstrate that our method significantly outperforms existing structured state space models on image classification, object detection and segmentation.
Besides, by fine-tuning large language models, our approach achieves consistent improvements in multiple textual tasks at minor training cost.

MambaTree: Tree Topology is All You Need in State Space Model
Yicheng Xiao t. Lin Song 2,es
Shaoli Huang', Jiangshan Wang', , Siyu Song", Yixiao Ge Xiu Li Ying Shan
Tsinghua Shenzhen International Graduate School, Tsinghua University
ARC Lab, Tencent PCG Tencent Al Lab
"South China Normal University
xiaoyc23@mailss tsinghua. enu.cn
romnysong@tencent. com
Abstract
The state space models, employing recursively propagated features, demonstrate strong representation capabilities comparable to Transformer models and supe- rior efficiency. However, constrained by the inherent geometric constraints sequences, it still falls short in modeling long-range dependencies. To address this issue, we propose the MambaTrree network, which first dynamically generates tree topology based on spatial relationships and input features. Then, feature propagation is performed based on this graph, thereby breaking the original se- quence constraints achieve stronger representation capabilities. Additionally, we introduce a linear complexity dynamic programming algorithm to enhance long-range interactions without increasing computational cost. MambaTree is versatile multimodal framework that can be applied to both visual and textual tasks. Extensive experiments demonstrate that our method significantly outperforms existing structured state space models on image classification, object detection and segmentation. Besides, by fine-tuning large language models, our approach achieves consistent improvements in multiple textual tasks at minor training cost. Code is available at https://github.. comm/asonXiao 1aaonniiioo--oooo 888/GrootVL.
Introduction
Mainstream fundamental models are primarily based CNN [30, 62, 44, 32, 13] and Transformer architectures [15, 43, 42, 59, 14J, which dominate in visual and language tasks. However, the small receptive field of CNNs and the high complexity of Transformers make it challenging to strike good balance between effectiveness and efficiency. The state space models (SSMs) [22, 24, 52) attempt disrupt this impasse, which model sequences in a recurrent form. Different from the previous recurrent neural networks [31, 7]. these approaches draw inspiration from control systems, leveraging structural parameter initialization attain stable optimization and superior computing performance. Nevertheless, it remains susceptible to the intrinsic flaw shared by recurrent neural networks, i.c., deficiency in capturing long-range dependencies.
Recently, an improved selection mechanism known as Mamba [19] is proposed to mitigate the challenges of SSMs. This approach introduces weight modulation during the propagation process, which substantially enlarges the effective receptive field and achieves impressive performance NLP tasks. Besides, numerous studies aim extend Mamba into computer vision, by employing various pre-defined strategies to map 2D image features into ID sequences. ViM [77] and VMamba [41] utilize multi-directional raster-scanning strategy, while LocalMamba [34] further confines its
"Equal contribution. Work done during an internship at Tencent. se Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

propagation range within a local window. They have successfully adapted Mamba to image inputs. Nevertheless, as shown in Fig. ((a), both raster-scanning and -ecaninc strategies introduce spatial discontinuities between adjacent pixels, and feature transformations in Mamba rely on the feature relationships, thereby impeding the effective information flow in a sequence. Additionally, PlainMamba [69] introduces continuous scanning strategy, aiming to alleviate this issue by simply adjusting the propagation direction at discontinuous positions. However, all these methods rely on fixed propagation trajectories, which ignore the inherent spatial structure and cannot dynamically adjust the topology based on input. Therefore, this paper endeavors to explore a new perspective: introducing an input-aware topological network feature propagation in state space models.
To achieve it, we develop a tree state space model and propose new framework, termed MambaTree, which adaptively generates a tree topology based on the input feature and then performs feature propagation on it. Specifically, two sub-networks, Vaeea and MambaTrreel., are designed for visual and language tasks respectively, which are illustrated in Fig. 1(b) and Fig. ((d). For visual tasks, motivated by [71, 54], we first utilize the dissimilarity between adjacent features to construct minimum spanning tree
four-connected planner graph. This process can adaptively encode the spatial and semantic information into a tree graph [4]. Then, we iteratively traverse each pixel, considering it as the root vertex, and aggregate the features of other pixels using the state transition function of Mamba. Intuitively, this operation requires two levels of traversal across the entire pixel set, resulting in unacceptable quadratic complexity relative to the number of pixels. However, given that the tree graph is acyclic, we propose a dynamic programming algorithm to achieve linear complexity propagation. With such input-awart tree topology, our approach enables more effective long-range interactions while maintaining consistent linear complexity with Mamba. Furthermore, our method can also be applied language tasks by constructing a tree topology based on the dissimilarity between token features, which overcomes the geometrical constraints of the text sequence. Using a similar aggregation process as MambaTreek, Mamba TreeL can significantly enhance the language representation of a pre-trained Large Language Model [19].
We conduct extensive experiments to validate the effectiveness of Mamba on multiple visual benchmarks, i.e. image classification on ImageNet [12]. object detection and instance segmentation on MSCOCO [39] as well semantic segmentation on ADE20K [75]. Results show that our method notably outperforms existing SSM-based methods for all benchmarks and achieves competitive performance with CNN and Transformer-base based approaches. Moreover, with LoRA finetuning [33], MambaTreel. demonstrates consistent improvements for pre-trained large language model at minor training cost.
2
Related Work
2.1 Conventional Vision Foundation Models
The evolution of deep neural networks has been a significant catalyst in machine vision perception. CNN-based models [30, 51, 35, 25, 61, 72, 38, 55, 73] firstly emerge as pivotal landmarks, with ResNet [30] notably standing out for its inventive residual connection module, gareering widespread adoption across diverse domains of visual recognition. Furthermore, more efficient convolution operations are formulated, such as depth-wise convolutions introduced by MobileNet [32], paving the way for lightweight models. Additionally, deformable convolution [10] has been proposed to enhance the receptive field. Subsequently, VIT [15] has significantly improved the vision recognition paradigm. reformulates the architecture design and training mechanism by combining transformer architecture in natural language processing, aiming to improve computational efficiency and broaden the scope of applications. After research discourse is centred on hierarchical ViTs [43, 42, 11, 63, 14, 56,5) 5] which design networks by decreasing feature resolution across the backbone gradually. Furthermore, recent research built on CNN serves to re-emphasize the capabilities of convolutional networks. For example, Internlmaee [62] presents a large model based on deformable CNN, while UniRepLKNet [13] exhibits significant performance through large kernel convolution.
2.2 Explorations about State Space Models
State space models (SSMs) have emerged as a novel class of models within the deep learning paradigm, showing significant potential for sequence transforming [23, 22, 52]. These methods have attracted significant attention due to their linear scalability with sequence length. The carly method,

---

Raster Scan
Continuous Scan
Local Scan
Tree Topology Scan
Previous Visual SSMs
(b) Treba
What
ingredients
do
need
what
ingredients
shortoncst
n-ddll
pies
pies
puiin
shortonust
(c) Previous Textual SSMs
(d) Mamba
Figure 1: Comparison of differrnt propagation strategies for multi-modal tasks. For visual tasks, the previous strategies (a) are based on fixed patterns, while our method can adaptively generate the propagation topology according to input features, For textual tasks, compared to previous methods (c), our approach (d) can break the inherent constraints of text sequences, facilitating the effective transmission of long-range information.
LSSL [23], draws inspiration from continuous state space models in control systems and attempts to address the long-range dependency problem through a combination with HIPPO [20] initialization. S4 [22] proposes to normalize the parameters into diagonal matrix, prompting subsequent series of research on structured SSMs [24, 21, 26, 19]. Recently, the Selective State Space Model [19], known Mamba, strikes balance between effectiveness and efficiency through the design of an input-dependent parameter initialization strategy, which has emerged as a formidable competitor to both transformer and CNN structures. In addition showcasing superior outcomes in sequence modeling, Mamba has been seamlessly incorporated into the visual domain [77, 41, 34, 69, 68). These studies often rely handerafted fixed scanning mechanisms mitigate the execution bias of the selective state space model on 2D non-causal images. However, such simplistic approaches cannot effectively capture spatial relationships in an input-dependent paradigm. To address this limitation, we propose an effective framework Mamba Tree in this work to enhance long-range modeling for both vision and language tasks by introducing an input-aware tree-based topological structure.
3 Method
In this section, we first revisit the selective state space model [19] and then elaborate on our input- aware topology scanning algorithm for state space modeling. With this superior algorithm,, we develop a tree SSM and propose novel framework called MambaTrree which consists of two sub-networks: MambaTreee for visual tasks and MambaTreed for fine-tuning pre-trained language model [ [19].
3.1 Revisiting Selective State Space Model
State Space Models (SSMs) commonly regarded continuous linear time-invariant systems [64] that map input stimulation xxtt) E R RIXD to output signal y(t) E RIID through a state vector M(t) E RIIN where i, D and N indicate the time step, channel number of the signal and state size, respectively. These models can be formulated as the following linear ordinary differential equations:
(t( = Ah(t) + Bz((), y(t) Ch(() + Dz(t).
where E NXN BE NXD, C NDD and feedthrough coefficient E P.
Discretization. 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 [19] to discretize the continuous system described by Eq. (1) and convert continuous variables (A, B, C, D) into corresponding discrete

---

Lineur
Linear
DWConv
Duilding a Femmeecie Gruph
Lineur
Lincar
Embedding
Parameters Generator
Prusing
SiLU
TSA
SoftPlus
A.B, C,D
Spanning
SiL.U
y
= exp(AA)
IN
for i
B (AB)
h Eveo Sl,r,,, A C, Norma(d)) Dr.
Linear
A,B.C,D Parameters Generator
Tree Stute Space Model
Tree Scanning Algorithm
Figure Illustration of Tree State Space Model. With an image feature map z, we perform Tree Scanning Algorithm (TSA) to construct 4-connected graph with edge weights measured by dissimilarity between pixels. Then, we obtain an MST with vertices Ω through pruning algorithm and perform the state transition for each vertex in this topology (detailed in Sec. 3.2). Red arrows describe the propagation source vertex i.
parameters (A, B. C, D) over the specified sampling time-scale A I::;
A = A B (E 1) A 'B, C, D
(2)
addition, many improved methods [41, 19] use an approximation of B based on first-order Taylor Series:
B (E A (AA) AB AB
Selective Mechanism. Previous SSMs store information through finite states and inherent time- invariance, which limits their effectiveness. Therefore, Mamba [19] introduces dynamic mechanism to selectively filter out input into a sequential state. Specifically, it utilizes Linear Projection to calculate the parameters (B.) (C,) and (A,) from the input sequence {aa) with I Ix directly to improve the context-aware ability. Then the output sequence {u can be computed with those anputaddativv discretized parameters as follows:
h A + B y1 = C,h,, Daa
3.2 Tree State Space Model
Mamba [19] has showcased remarkable performance in modeling the dependencies of consecutive words in sequence. However, its applicability in long-contex. tasks, especially visual modeling, still poses certain challenges. For visual tasks, many methods attempt to address this problem by employing fixed scanning strategies, such as multi-directional raster scan [41, 77], local scan [34], and continuous scan [69]. However, these handcrafted scanning methods fail to effectively preserve the 2D structural information of images.
Following the design in Mamba [19], construct transform block tree state space model, which is presented in Fig. 2. The only difference between our block and Mamba lies in the replacement of the structured state space block with the proposed tree scanning algorithm. In the tree scanning algorithm, we generate a tree topology and then propagate the state of each vertex along the topological path to obtain strong feature representations. representations by incorporating such tree topology during text processing, which overcomes the geometrical constraints of text sequences. the following, we elaborate on the proposed tree scanning algorithm and its applications for multi-modal tasks.
addition, our algorithm can effectively enhance language
Tree Scanning Algorithm. Given an input feature X (2x) where I. is the sequence length (or the number of input pixels), construct undirected m-connected graph G (V, E) for the

---

#hdW3
H
W
# 2
Down Sampling
Down Sampling ling
Downn Sampling
STAGE Basic Block X4a
STAGE 2 Basic Block Xlg
STAGE 3 Basic xlg
STAGE 4 Basic Block Xh4
Hedd
Stem
Head Cl.
Basic Block
Conv: 3x3, S2. LN + GELU
Del.
Sog.
Tree
Down Sampling S2,
LN
Space Model
LN
FTN
Conv: 3x3, 57,
LN
LN
Figure 3: Overview of MambaTreeV. LN means LayerNorm and FFN is a forward network in the basic block. S2 and P1 denote stride of 2 and padding size of in convolution, respectively.
feature. m is we set m = 4 for visual tasks, meaning each pixel is connected to its four neighboring pixels. For language tasks, we set m by default, meaning each token is connected 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 Table
that indicates the number of adjacent tokens. Following [71, 54],
We use the Contractive Boruvka algorithm [2] to prune the edges with significant dissimilarity, which generates minimum spanning tree (MST) Gr 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 Aş. Furthermore, following the Eq. (4), the state aggregation process for the i-th vertex can be formulated as:
h; Σ S(E,)BBBj,, VjER
S(E = II A kEN
(5)
where S denotes the index set of all vertices in the tree. S(Fr,) represents the path weight of hyperedge Eij traced from j-th vertex to i-th vertex in the tree Gr. and Ni, indicates the index set of all vertices on this hyperedge. For visual tasks, we iterate over each vertex, treating as the root of the spanning tree Gr. and aggregate the states from the other vertices, thereby obtaining the transformed states [ For textual tasks, because of the causal prediction manner in large language models, we only take the last token root and aggregate from other tokens. To achieve end-to-end training, we derive the derivative of the output hidden state h2 to the input variables Ag. B, and zy as follows:
ah; = S (Fag)) Bj, Oz,
ah, S (Fas) B
0h:
Σ YJEGE
Bya,
0A
where Cj indicates the children of vertex tree Gr whose root is the vertex and n denotes the parent of vertex l in Eq. (7). Finally, the output feature Y can be formulated as:
C Norm H) + D X,
(8)
where Y. Hl and X indicate the stack of (uiten (hijen and (2) respectively. @ denotes the element-wise multiplication.
Efficient Implementation for Multi-Modality. 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 O To alleviate this issue, utilize dynamic
5

---

Algorithm Vision Tree Scanning
Input: Input feature { Input matrix (B.) State matrix {A.) Gradient of loss to hidden states (HLose
0h, ji= Minimum Spanning Tree Gr. Breadth first topological order of Gr
Traverse Path: Root,.. .oot,..... Leaf t BFS((T)
Forward:
Initialization: {EiHe1 & {ai)i=1 2: for i +- Leaff to Root do
E -: Exje((Pan()):)) A
4: end for
for i Root to Leaf do
6:
if is Root then
h4 =3
8:
else
h4 A, (hpaar(()
A,E:) + & A?) A,hpar(t)
10: end if
end for
Backward:
12:
Initialization: {nE t {1 (Loss }
-m Bh; ji=1
for Leaf to Root do
14: = B, aLoss The
(()e(()a(t)) Aj
end for
for Root Leaf do if is Rool then
0Loss B: Sz,
aLass BB.
BLoss 8A:
Mi
else
20:
0Loss = (1 A + As OLoss B,
ÔLoss (1 A7) A; aLoss 22
dz,
0B.
(BP
aLoss Ti (h+ A + + BA.
OLoss Or;;
Am)) ni, + + di, 27,E,
ÔLoss
end if end for
Output: Hidden states {ha Grad. of loss to input feature { (BLoss 1 matrix { @ooss
Dzz Ji=1: Grad. of loss to input Ji= Grad. of loss to state matrix @oos 14 DA, Ji=1
programming procedure accelerate the inference and training processes as elaborated in Algo- rithm 1, which results in linear complexity (LL). For textual tasks, we perform unidirectional aggregation approach (shown in Algorithm of Appendix B) in adherence the causal nature of language. Moreover, we provide the back prooppppppggtiinn process for both Vision Tree Scanning and Language Tree Scanning processes, whose detailed proofs refer Appendix C.
3.3 Application for Vision and Language
Given an image with a shape of H x W x 3, our goal is to obtain high-quality visual features for downstream tasks. To this end, propose effective vision architecture MambaTrreeV which consists of a stem module, several basic blocks and downsampling layers to generate hierarchical representations illustrated in Fig. 3. Overall, our MambaTTree comprises four stages similar to previous general vision backbones [44, 43, 62, 41]. 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 feature map with a shape of H. W. layers and one GELU activation function. The kernel size for both convolutions is with a stride of and padding of 1. Similarly, a downsampling layer consists of 3 x 3 convolution with a stride of 2 and padding of and an LN layer. Positioned between two stages, serves to downsample the input feature map by a factor of 2. Motivated by [62, 41], we devise residual block with skip connections to integrate our fundamental Tree State Space Model Sec. 3.2. 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, feedforward neural network is utilized to project the normalized features to output signals as shown
MambaTreeV
4 4 x C.. It includes two convolutions, two Layer Normalization (LN)

---

Top-1
Top-1
Method
Type #Param. #FLOPs
Method
Type #Param. #FLOPs
Acc.
Acc.
Deit-S [59]
22M
4.6G
79.9
ConvNeXt-S [44]
50M
8.7G
83.1
Swin-T [43]
28M
4.6G
81.3
SLaK-S [40]
55M
9.8G
83.8
CoAINet-O [1]]
T
25M
4.0G
81.6
UniRepL.KNe-S [13]
56M
9.1G
83.9
SG-Former-S [50]
T
23M
4.8G
83.2
InternImage-S [62]
50M
8.0G
84.2
ConvNeXt-T [44]
29M
4.5G
82.1
HyenaViT-B [17]
88M
78.5
(LaK-T[40)
30M
5.0G
82.5
S4ND- [48]
89M
80.4
UniRepLKNee-T
31M
4.9G
83.2
PlainMamba- L3 [69]
50M
14,4G
82.3
Internlmage-T [62]
C
30M
5.0G
83.5
VMamba- S [41]
50M
8.7G
83.6
(IM-S(77)
26M
5.1G
80.5
LocalVMamba-SS [34]
50M
11,4G
83.7
LocalViM-S [34]
28M
4.8G
81.2
MambaTree (Ours)
51M
8.5G
84.2
PlainMamba--12 [69]
25M
8.1G
81.6
Deit-B [59]
86M
55,4G
83.1
Mamba-2D-S [37]
24M
81.7
Swin-B [43]
88M
15.4G
83.5
LND-Conv.NeXt-T [48]
S
30M
82.2
CoAtNet-2 [11]
75M
16.0G
84.1
VMamba-T [41]
31M
4.9G
82.5
ConvNeXt-B [44]
C
89M
15.4G
83.8
-o [34]
26M
5.7G
82.7
SLaK-B - [40]
C
95M
17.0G
84.0
MambaTroeVas (Ours)
30M
4.8G
83.4
Mamba-2D-B [37]
92M
83.0
Swin-S [43]
50M
8.7G
83.0
VMamba-B [41]
89M
15.4G
83.9
CoAINet-1 [1]]
T
42M
8.0G
83.3
MambaTreeeVB (Ours)
91M
15.1G
84.8
Table Image classification performance on the ImageNet-1K 1K validation set. T, C and indicate the model type of Transformer, CNN and SSM, respectively. All models take a scale of 224 as input.
in Fig. 3. Based on the above origin components, we develop our Mamba TreeV in three scales, i.e., VambaTrreV-Tings MambaTreeM V-Small and V-
MambaTreel. Recurrent neural networks rely fixed memory preserve past information, which poses limitations when handling long contexts where relevant words are distant from the current moment. While Mamba [19] employs selection mechanism enhance context awareness, its fixed memory size cannot expand over time, resulting in restricted state space. Therefore, the ability to extrapolate decreases during scrolling 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 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, utilizes pretrained state transformation parameters conduct semantic aggregation by incorporating topological structures. Experimental results demonstrate the effectiveness of our approach.
4 Experiments
We conduct extensive experiments to evaluate the effectiveness MambaTrreV and compare it with advanced CNN-basede Transformer-baseer and SSM-based models covering various downstream tasks, including image classification, object detection and semantic segmentation. Furthermore, we validate the capability of MambaTreel in the field of natural language understanding.
4.1 Image Classification
Settings. We assess the classification performance of MambaTreeV the ImageNet-Ik dataset [12]. Following previous practices [43, 44, 62, 41J, all MambaTreeV models are trained for 300 epochs from scratch using AdamW optimizer with warm-up uum strategy of 20 epochs. During training, utilize a Cosine Scheduler with an initial learning rate of x 10 and weight decay 0.05. In addition, the exponential moving average (EMA) also applied.
The comparison results summarized in Table show MambaTred leading all SSMebased models and competitive with advanced CNNs and Transformers across tiny, small, and base scales. Specifically, MambaTreat Tree-- achieves 83. 33.4% Top-1 Acc. boosting ViM-S by 2.9%, Local Vim-S by 2.2%, PlainMamba-LL2 by 1.8% and VMamba-T by 0.9% with similar FLOPs. Additionally, surpasses Con-NeXi-T by 1.3% and Swin-T by 2.2%, demonstrating the effectiveness of our method.
Results.
7

---

(a) Input Signal
(b) TP Scan
(c) Raster Scan
mloU
mloU
Method
Type #FLOPs
SS
MS
Swin-T [43]
945G
44.5
45.8
ConvNeXt-T [44]
C
939G
46.0
46.7
SLaK-T [40]
C
936G
47.6
InternImage-T [62]
C
944G
47.9
48.1
UniRepLKKNet-T [13]
C
946G
48.6
49.1
ViM-S [77]
S
44,9
LocalViM-S [34]
S
297G
46.4 47.5
PlainMamba-Le L2 [69]
S
285G
46.8
VMamba-T [41]
964G
47.3 48.3
TomalVMamba-T [41]
970G
47.9 49.1
MambaTreeVas (Ours)
S
941G
48.5
49.4
Swin-S [43]
1038G
47.6
49.5
ConvNeXt-S [44]
C
1027G
48.7 49.6
SLaK-S [40]
C
1028G
49.4
Internlmage-S [62]
C
1017G
50.1 50.9
UniRepL.KNet-S [13]
C
1036G
50.5 51.0
PlainMamba-L33 [69]
S
419G
49.1
VMamba-S [41]
S
1081G
49.5 50.5
LocalVMamba-S [34]
S
1095G
50.0 51.0
Figure 4: Visualization of aflinity maps in the specific position. The Location is marked by the red cross each input (a). TP is our tree topol- ogy scanning algorithm (b), which captures more detailed structural information and has larger re- ceptive field compared to raster scanning (c).
Vamb (Ours)
S
1019G
50.7 51.7
Table 2: Semantic segmentation perfor- mance on ADE20K val set. The crop size is all 512 SS and MS denote single- scale and multi-scale testing, respectively.
4.2 Object Detection
Settings. We verify the detection performance of MambaTreeV on the MSCOCO 2017 dataset [39] with MMDetection library [3]. We follow previous works [41, 62, 43, 34, 53, 55, 74, 70, 6] validate object detection and instance segmentation tasks with Mask-RCNN [29]. Specifically, We adopt the AdamW optimizer with learning rate of 1 x 10-4 and batch size of 16 to optimize the model built upon our pre-trained classification backbones on ImageNet-IK. The training schedules include 1 (12 epochs) and 3x (36 epochs) with multi-scale data augmentation.
Results. As depicted in Table (in Appendix A.), method outperforms existing methods most evaluation metrics, especially for instance segmentation. Under 1 x schedule, MambaTTreeV-T achieves 47.0 in box mAP (AP"), which is 1.1 points higher than ViM-S and 0.5 points higher than VMamba-T. It is worth noting that MambaTreeVa. outperforms ViM-S by 1.7 points with 1x schedule and AocallMamba+ by 0.4 points with 3x schedule mask mAP (AP"). Moreover, the best APt 50.1 and APm 44.6 are obtained by MambaTree V-S in 3x schedule with multi-scale training.
4.3 Semantic Segmentation
Settings. To evaluate the semantic segmentation performance our MambaTreeV series, we train models with UperNet [65] initialized by pre-trained classification weights ADE20K[75] for 160k iterations, following common practices without additional augmentations for fair comparison.
Results. Our method performs exceptionally well on segmentation tasks shown in Table MambaTreeeV-T yields clear improvement of +3.6 in single-scale mloU compared to ViM-S and +1.9 multi-scale mloU compared to LocaIViM-S. Furthermore, MambaTreeV-S boosts InternImage-S by 0.6 and 0.8 in single-scale and multi-scale respectively. We consider the preserva- tion of intricate structural details through tree topology scanning to be particularly advantageous for segmentation tasks that require pixel-level perception.

---

Average
Method
PIQA f Arc-E1
SST
WG 1
L-ppl 4
Race
BQA
Acc.
Mamba [19]
64.5
48.0
65.6
51.8
16.1
27.4
16.8
45.7
LoRA [33]
64.7
48.3
52.2
17.7
28.6
17.8
46.1
Mamba TreeL. (Ours)
65.0
49.8
69.5
51.1
15.9
28.9
19.2
47.2
Table 3:
Evaluation language model benchmarks. Arc-E, WG, L-ppl and BQA indicate Arc- easy [8], WinoGrande, LAMBADA [49] and Openbookqa [47] benchmark, respectively.
Scanning Strategy
Acc
Distance Metric
Acc.
Root Setting
Acc.
Raster Scan
82.6
Manhallan 82.9
First vertex
82.9
Cross Scan
83.1
Euclidean
83.2
Last vertex
83.0
Tree Topology Scan
83.4
Cosine
83.4
All vertices
83.4
Table 4: Effectiveness of algorithm.
Table 5: Impact of different distance Metries.
Table 6: Superiority of traversing all vertices.
4.4 Language Understanding
We regard Mamba [19] with 130M parameters the base model. To verify the effectiveness of our MambaTreeL in nature language understanding, we first fine-tune pre-trained Mamba via LoRA [33] and MambaTreck under the same setting with the Alpaca data [58]. which contains 52000 instruction tuning data for supervised fine-tuning. Then we utilize popular language benchmarks provided the open-sourced Im-evaluation harnnsss project [18] for evaluation, including PIQA [1], A12-ARC [8], SST [60], WinoGrande, LAMBADA [49], Race [36] and Openbookqa [47]. The results Table demonstrate that MambaTreeL provides a benefit of +1.1% in average Acc. compared to LoRA. Since the short prompt length of WinoGrande dataset, the performance degrades with a marginal gap.
4.5 Ablation Study & Qualitative Results
In this section, we conduct analysis experiments on ImageNet- dataset and present some visual results to illustrate the effectiveness of our algorithm.
Scanning Strategy. We conduct a head-to-head comparison of different scanning strategies, as shown in Table 4. The tree topology scanning outperforms previous strategies by 0.8% and 0.3%, highlighting the superiority of our algorithm vision recognition.
Distance Metric. Before generating minimum spanning tree from connected graph, it is important to measure the edge weights between vertices. Therefore, we validate several distance metrics as illustrated Table The results indicate that Cosine distance most effectively represents the relationship between vertices, performing 0.5% better than M anhattan and 0.2% better than Euclidean.
Root Setting. 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 Table The results show reductions of 0.5% and 0.4%, respectively.
Inference speed comparison. As shown in Table 7, we report the inference throughputs of our method Nvidia V100 GPU. MambaTreeVa. refers each stage sharing the same topology structure, which enhances efficiency without compromising accuracy. To achieve better practical inference speed, also introduce cuda implementation optimized for GPUs. Compared with other counterparts, our approach exhibits superior effectiveness and faster inference speed.
Qualitative Results. 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,

---

Method
Throughput (img/s)
GPU Memory
FLOPs
#Params.
Acc. (Top1.)
(224x224)
PlainMamba--12 [69]
363
4204M
8.1G
25M
81.6
VMamba-T [41]
374
8646M
4.9G
31M
82.5
LocalVMamba-T [34]
311
11298M
5.7G
26M
82.7
MambaTree V-T root)
283
6012M
4.8G
30M
83.0
V-T
281
6471M
4.8G
30M
83.4
Man
392
4800M
4.8G
30M
83.4
Table 7: Runtime comparison on an Nvidia V100 GPU during inference.
set the anchor point in the upper left corner of the sky 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.)
Conclusion & Limitations
In this paper, we propose tree state space model to perform feature propagation input-aware topology. Besides, we introduce a linear complexity dynamic programming algorithm to enhance long-range interactions without increasing computational cost. With the proposed techniques, establish the general multi-modal networks to break the original sequence constraints and achieve stronger representation capabilities. Extensive experiments demonstrate the effectiveness of method in both visual and language tasks. The limitation of our method is that the tree structure common paradigm, and it needs to be specifically optimized according to the hardware device.
Acknowledgments and Disclosure of Funding
This work was supported by the STI 2030-Major Projects under Grant 20212D0201104.
References
[1] Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al.: Piqa: Reasoning about physical commonsense in natural language. In: AAAL pp. 7432-7439 (2020)
[2] Borûvka, O.: jistém problému minimálním (1926) 5
[3] Chen, K., Wang, J., Pang, J., Cao, Y., Xiong. Y., Li, X., Sun, S., Feng, W., Liu, Z., Xu, J., et al.: Mmdetection: Open mmlab detection toolbox and benchmark. arXiv preprint arXiv:1906.07155
(2019)* 8
[4] Chen, Z., Duan, Y., Wang, W., He, J., Lu, T., Dai, J., Qiao, Y.: Vision transformer adapter for dense predictions. arXiv preprint arXiv:2205.088 200022205.08534 (2022) 16
[5] Cheng, C., Song, L., Xue, R., Wang, H., Sun, H., Ge, Y., Shan, Y.: Meta-adapter: online few-shot learner for vision-language model. arXiv preprint (XXX:23311.03777 (2023) 2
[6] Cheng. T., Song, L., Ge, Y., Liu, W. Wang. X., Shan, Y.: Yolo-world: Real-time open- vocabulary object detection. arXiv preprint arXiv:2401.17270 (2024) 8
[7] Chung. J., Gulcehre, C., Cho. K., Bengio, Y.: Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.338. (2014)
[8] Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., Tafjord, O.: Think you have solved question answering? arc, the ai2 reasoning challenge. arXiv preprint SrXiv:1803.0545 (2018) 9
[9] Cubuk, F.D., Zoph, B., Mane, D., Vasudevan, V., Le, Q.V.: Autoaugment: Learning augmenta- tion strategies from data. In: CVPR, pp. 113-123 (2019) 15
10

---

[10] Dai. J.. Qi, H., Xiong, Y., Li. Y., Zhang, G., Hu, H., Wei, Y.: Deformable convolutional networks.
In: ICCV. pp. 764 164-773 (2017) 2
[11] Dai, Z., Liu, H., Le, Q.V., Tan, M.: Coatnet: Marrying convolution and attention for all data sizes. NeurlPS 34, 3965-3977 (2021) 2,7
[12] Deng, J., Dong, W., Socher, R., Li, LJJ, Li, K., Fei-Fei, L.: Imagenet: A large-scale hierarchical image database. In: CVPR. pp. 248 248-255. leee (2009) 2,7
[13] Ding, X., Zhang. Y., Ge. Y., Zhao, S., Song, L., Yue, X., Shan, Y.: Unireplknet: universal perception large-kernel convnet for audio, video, point cloud, time-series and image recognition. CVPR (2023) 2, 7,8
[14] Dong, X., Bao, Chen, D., Zhang, W., Yu, N., Yuan, L.,, Chen, D., Guo, B.: Cswin transformer: A general vision transformer backbone with cross-shaped windows, In: CVPR. pp.
(2022) 1, 2, 16
[15] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani. M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., Houlsby, N.: An image is worth 16x16 words: Transformers for image recognition at scale. In: ICLR (2021) 1,
[16] Fang, C., He, C., Xiao, F., Zhang, Y., Tang, L.., Zhang, Y., Li, K., Li, X.: Real-world image dehazing with coherence-based label generator and cooperative unfolding network. arXiv preprint arXivv2406.07966 (2024) 15
[17] Fu, D., Arora, S., Grogan, J., Johnson, I., Eyuboglu, E.S., Thomas, A., Spector, B., Poli, M., Rudra, A., Ré, C.: Monarch mixer: A simple sub-quadratic gemm-based architecture. NeurIPS 36 (2023) 7
[18] Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac"h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., Zou, A.:
A framework for few-shot language model evaluation (12 2023) 9
[19] Gu, A., Dao, T.: Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 (2023) 1, 2, 3, 4, 7,9
[20] Gu, A., Dao, T., Ermon, S., Rudra, A., Ré, C.: Hippo: Recurrent memory with optimal polynomial projections. NeurlPS 33, 1474-1487 (2020) 3
[21] Gu, A., Goel, K., Gupta, A., Ré, C.: On the parameterization and initialization of diagonal state space models. NeurIPS 35, 35971-35983 (2022)* 3
[22] Gu, A., Goel, K., Ré, C.: Efficiently modeling long sequences with structured state spaces. In: ICLR (2022) 1, 2,3
[23] Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A., Ré, C.: Combining recurrent, convolutional, and continuous-time models with linear state space layers. NeurlPS 34, 572-585 (2021) 2,
[24] Gupta, A., Gu, A., Berant, J.: Diagonal state spaces are effective as structured state spaces. NeurIPS 35, 22982-22994 (2022) 1.3
[25] Han, K., Wang, Y., Xu, C., Guo, J., Xu, C., Wu, E., Tian, Q.: Ghostnets on heterogeneous devices via cheap operations. IJCV 130(4), 1050-1069 (2022) 2
[26] Hasani, R., Lechner, M., Wang, T.H., Chahine, M., Amini, A., Rus, D.: Liquid structural state-space models. arXiv preprint arXiv:2209.1295/ (2022) 3
[27] He, C., Li, K., Zhang, Y., Xu, G., Tang, I.,, Zhang, Y., Guo, 7.,, Li, X.: Weakly-supervisee concealed object segmentation with sam-based pseudo labeling and multi-scale feature grouping. Advances in Neural Information Processing Systems 36 (2024)
[28] He, C., Shen, Y., Fang, C., Xiao, F., Tang, I.,, Zhang, Y., Zuo, W., Guo, 7., Li, X.: Diffusion models in low-level vision: A survey, arXiv preprint arXiv:2406. 11138 (2024) 15
11

---

[29] He, K., Gkioxari, G., Dollár, P., Girshick, R.: Mask r-cmn. In: ICCV. pp. 2961-2969 (2017) 8
[30] He, k., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR. pp. 770-778 (2016) 1,2
[31] Hochreiter, S., Schmidhuber, J.: Long short-term memory, Neural computation 98), 1735-1780
(1997) 1
[32] Howard, A.G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andrectto, M., Adam, H.: Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv:1704.0483 (2017) 1,2
[33] Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang. S., Wang, L., Chen, W.: Lora: Low-rank adaptation of large language models. In: ICLR (2022) 2, 9
[34] Huang, T., Pei, X., You, S., Wang. F., Qian. C., Xu, C.: Localmamba: Visual state space model with windowed selective scan. arXiv preprint arXiv:2403.09336 (2024) 1, 3, 4, 7, 8, 10, 15, 16
[35] Krizhevsky, A., Sutskever, I., Hinton, G.E.; Imagenet classification with deep convolutional
neural networks. NeurlPS 25 (2012) 2
[36] Lai, G., Xie, Q. Liu, H., Yang, Y., Hovy, E.H.: RACE: large-scale reading comprehension dataset from examinations. In: EMNLP. pp. 785-794. Association for Computational I Linguistics (2017)* 9
[37] Li, S., Singh, H., Grover, A.: Mamba-nd: Selective state space modeling for multi nulti-dimensional data, arXiv preprint arXiv:2402.05892 (2024) 7
[38] Li, Y., Song, L.., Chen, Y., Li, Z.,, Zhang, X., Wang, X., Sun, J.: Learning dynamic routing for semantic segmentation, In: CVPR (2020) 2
[39] Lin, T.Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., Zitnick, C.I.: Microsoft coco: Common objects in context. In: ECCV, pp. 740-755. Springer (2014) 2,
[40] Liu, S., Chen, T., Chen, X., Chen, X., Xiao, Q. Wu, B., Kärkkäinen, T., Pechenizkiy, M., Mocanu, D., Wang, Z.: More convnets in the 2020s: Scaling up kernels beyond 51x51 using sparsity, arXiv preprint arXiv:2207.0m 03620 (2022) 7.8
[41] Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang, Y., Ye, Q., Liu, Y.: Vmamba: Visual state space model. arXiv preprint arXiv:2401.10166 (2024) 1, 3, 4, 6, 7, 8, 10, 15, 16
[42] Liu, Z., Hu, H., Lin, Y.. Yao, Z., Xie, Z., Wei, Y.. Ning, J., Cao, Y., Zhang. Z., Dong, L., et al.: Swin transformer v2: Scaling up capacity and resolution. In: CVPR. pp. 12009-12019 (2022) 1.2
[43] Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y. Zhang, Z., Lin, S., Guo, B.: Swin transformer: Hierarchical vision transformer using shifted windows. In: ICCV. pp. 10012-10022 (2021) 1, 2,
6, 7, 8, 15, 16
[44] Liu, Z., Mao, H., Wu, C.Y., Feichtenhofer, C., Darrell, T., Xie, S.: A convnet for the 2020s. In: CVPR. pp. 11976-11988 (2022) 6, 7, 16
[45] Loshchilov, I., Hutter, F: Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 (2017) 15
[46] Luo, Z., Xiao, Y., Liu, Y., Li, S., Wang. Y., Tang. Y., Li, X., Yang. Y.: Soc: Semantic- assisted object cluster for referring video object segmentation. Advances in Neural Information Processing Systems 36 (2024) 15
[47] Mihaylov, T., Clark, P., Khot, T., Sabharwal, A.: Can suit of armor conduct electricity? A new dataset for open book question answering. In: EMNLP, pp. 2391- Association for Computational Linguistics (2018) 9
12

---

[48] Nguyen, E., Goel, K., Gu, A., Downs, G.W., Shah, P., Dao, T., Baccus, S.A., Ré, C.: S4nd: Modeling images and videos as multidhmennional signals using state spaces, arXiv preprint arXiv:2210.06583 (2022) 7
[49] Radford, A., Wu, J.. Child, R., Luan, D., Amodei, D., Sutskever, I., et al.: Language models are unsupervised multitask learners. OpenAl blog 1(8), 9 (2019) 9, 19
[50] Ren, S., Yang, X., Liu, S., Wang. X.: Sg-former: Self-guided transformer with evolving token reallocation. In: ICCV. pp. 6003-6014 (2023) 7
[51] Simonyan, K., Zisserman, A.: Very deep convolutional networks for large-scale image recogni- tion. In: Bengio, Y., LeCun, Y. (eds.) ICLR (2015)
[52] Smith, J.T., Warrington, A., Linderman, S. S.W.: Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.0499: (2022) 1,2
[53] Song, L., Li, Y., Jiang, Z., Li, Z.., Sun, H., Sun, J., Zheng, N.: Fine-grained dynamic head for object detection. NIPS (2020) 8
[54] Song, L.,, Li, Y., Li, Z., Yu, G., Sun, H., Sun, J., Zheng, N.: Learnable tree filter for structure- preserving feature transform. NeurlPS 32 (2019) 2,
[55] Song, L.,, Zhang, S., Yu, G., Sun, H.: Tacnet: Transition-aware context network for spatio- temporal action detection. In: CVPR (2019) 2, 8
[56] Song, L., Zhang, S., Liu, S., Li, Z., He, X., Sun, H., Sun, J., Zheng, N.: Dynamic grained encoder for vision transformers. NIPS (2021) 2
[57] Tang, L., Li, K., He, C., Zhang, Y., Li, X.: Source-free domain adaptive fundus image seg- mentation with class-balanced mean teacher. In: International Conference on Medical Image Computing and Computer- -asiuter-Assisted Intervention. pp. 684 694. Springer (2023) 15
[58] Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., Hashimoto, T.B.: Stanford alpaca: An inslluction-follovinn Ilama model (2023) 9
[59] Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., Jégou, H.: Training data-efficient image transformers & distillation through attention. In: ICMI.. pp. 10347-10357. PMLR (2021) 1.7
[60] Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., Bowman, S.R.: GLUE: A multi-task benchmark and analysis platform for natural language understanding. In: ICLR (2019) 9
[61] Wang. J., Song, L.,, Li, Z., Sun, H., Sun, J., Zheng, N.: End-to-end object detection with fully convolutional network. In: CVPR (2021) 2
[62] Wang, W., Dai, J., Chen, Z., Huang. Z., Li, 7., Zhu, X., Hu, X., Lu, T., Lu, I.., Li, H., et al.: Internimage: Exploring large-scale vision foundation models with deformable convolutions. In: CVPR. pp. 14408 14419 (2023) 1, 6, 7, 8, 15, 16
[63] Wang, W., Xie, E., Li, X., Fan, D.P., Song, K., Liang, D., La, T., Luo, P., Shao, L.: Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In: ICCV pp. 568-578 (2021)2 2
[64] Williams, R.L., Lawrence, D.A., et al.: Linear spaee control systems. John Wiley & Sons (2007) 3
[65] Xiao, T., Liu, Y., Zhou, B., Jiang Y. Sun, J.: Unified perceptual parsing for scene understanding. In: ECCV. pp. 418 434 (2018) 8
[66] Xiao, Y., Luo, Z., Liu, Y., Ma, Y., Bian, H., Ji, Y., Yang, Y., Li, X.: Bridging the gap: unified video comprehension framework for moment retrieval and highlight detection. CVPR (2024) 15
[67] Xiao, Y., Ma, Y., Li, S., Zhou, H., Liao, R., Li, X.: Semanticac: Semantics-assisted framework for audio classification. In: ICASSP 2023-2023 IEEE International Conference on Acousties, Speech and Signal Processing (ICASSP). pp. 1-5. IEEE (2023) 15
13

---

[68] Xu, Z., Lin, Y., Han, H., Yang, S., Li. R., Zhang, Y., Li, X.: Mambatalk: Efficient holistic gesture synthesis with selective state space models. arXiv preprint arXiv:2403.09471 (2024)
[69] Yang, C., Chen, Z., Espinosa, M., Eriesson, L., Wang, Z., Liu, J., Crowley, E.J.: Plainmamba: Improving non-hierarchical mamba in visual recognition. arXiv preprint arXiv:2403.17693 (2024) 2, 3, 4, 7, 8, 10
[70] Yang. J., Song. L.., Liu, S., Li, Z., Li, X., Sun, H., Sun, J., Zheng, N.: Dbq-ssd: Dynamic ball query for efficient 3d object detection. arXiv preprint arXiv:2207.10909 (2022) 8
[71] Yang, Q.: Stereo matching using tree filtering. IEEE TPAMI 37(4), 834-846 (2014) 2,
[72] Yang. R., Song. L., Ge. Y., Li, X.: Boxsnake: Polygonal instance segmentation with box supervision. In: Proceedings of the IEEE/CVF International Conference on Computer Vision
(2023) 2
[73] Zhang, S., Song, L., Gao, C., Sang, N.: Glnet: Global local network for weakly supervised action localization. IEEE Transactions on Multimedia 22(10), 2610-2622 (2019)
[74] Zhang, S., Song. L., Liu, S., Ge, Z., Li, Z., He, X., Sun, J.: Workshop autonomous driving at cvpr 2021: Technical report for streaming perception challenge. arXiv preprint arXiv:2108.0422. 04.200000000. (2021)* 8
[75] Zhou, B., Zhao, H., Puig, X., Fidler, S.. Barriuso, A., Torralba, A.: Scene parsing through ade20k dataset. In: CVPR. pp. 633-641 (2017) 2, 8
[76] Zhou, H., Yang, R., Zhang, Y., Duan, H., Huang, Y., Hu, R., Li, X., Zheng, Y.: Unihead: unifying multi multi-perception for detection heads. TNNLS (2023) 15
[77] Zhu, L., Liao, B., Zhang, Q. Wang, X., Liu, W., Wang, X.: Vision mamba: Efficient visual
representation learning with bidirectional state space model. arXiv preprint arXiv:2401.04417 (2024) 1, 3, 4, 7, 16
14

---

Appendix
A Detailed Training Settings and Results
A.1 Image Classification.
We follow the previous works [62, 41, 43] to conduct the experiments. The models are trained with thirty-two 32GB V100 GPUs by default. We set betas and momentum of the AdamW [45, 76, 66] optimizer with (0.9, 0.999) and 0.9, respectively. During training. we utilize Cosine Scheduler with an initial learning rate of 1 10-3 and weight decay of 0.05. We adopt the common training data aug- mentation strategies following [34, 62], including Auto/Aummet [9] with 1onddm9-mstddl incc... MixUp strategy with a ratio of 0.8 also adopted each batch. Horizontal flip and Random resized crop strategy are both used in the process of training.
MommhTTneV-5
ManbaT
VMuambuSS
VMamba-R
LoceMMmbasSS
Top-1 Accuracy
Mamba TTeeV
LacalMarnba-TT
MainMomio-L33
VManis-T
PlainMemiss
LoulDaashs
FLOPP
Figure 5: Classification performance comparison among SSM-based vision foundation models.
Performance Comparison. We compare various SSM-based visual foundation models as shown in Fig. 5, with different colors representing different models and different shapes indicating different model scales. The size of each shape indicates the number of model parameters. The horizontal axis denotes FLOPs and the vertical axis represents the Top-1 accuracy of the corresponding method on ImageNet-11 K val dataset. Fig. demonstrates that Mamba Vree is the best choice in terms of efficiency and effectiveness.
A.2 Object Detection.
For fair comparison, conduct the evaluation following common practice [62, 41, 433. The models are trained with eight 32GB V100 GPUs by default. The input image is resized so that the shorter side is 800 pixels, while the longer side does not exceed 1333 pixels during the 1 x schedule. The number of warmup steps is set to 500 in the 1 x schedule. For 3 schedule, the shorter side resized to 480-800 pixels and the longer side does not exceed 1333 pixels. The number of warmup steps is set to 1000 x schedule. Results shown in Table demonstrate the effectiveness of Mamba'Tres TreeV in object detection and instance segmentation on COCO val2017.
A.3 Semantic Segmentation.
We optimize our Mamba V-T/S using AdamW optimizer with an initial learning rate of 6 x 10
which is decayed by a rate 1.0 with the polynomial decay schedule following [62, 27, 57, 16, 28]. The number of warmup iters is set to 1600 with an initial learning rate of 1 x 10-6 [41, 34, 46, 67]. The default input resolution 512 x 512 as well as FLOPSs are calculated with an input size of 512 x 2048. The models are trained with eight 32GB V V100 GPUs by default.
15

---

Mask R-CNN 1x Schedule
Mask R-CNN 3x MS Schedule
Method
#FLOPS.
AP"
A0
AP
AP"
AP
AP
AP"
AP
AP7s
APm
AP
AP76
Swin-T T [43]
267G
42.7
65.2
46.8
39.3
62.2
42.2
46.0
68.1
50.3
41.6
63.1
44.9
ConvNcXt-T [44]
262G
44.2
66.6
48.3
40.1
63.3
42.8
46.2
67.9
50.8
41.7
65.0
44.9
CSWin-T [14]
279G
46.7
68.6
51.3
42.2
65.6
45.4
49.0
70.7
53.7
43.6
67.9
46.6
ViM-S [77]
218G
44.9
67.1
49.3
41.0
64.2
44.1
VMamba-T [41]
286G
68.5
50.7
42.1
65.5
45.3
48.5
69.9
43.2
66.8
46.3
L- L-Vmamba-T [34]
291G
46.7
68.7
50.8
42.2
65.7
45.5
48.7
70.1
53.0
43.4
67.0
46.4
MambaTreeV-T (Ours)
265G
47.0
69.4
51.5
42.7
66.4
46.0
49.0
70.8
54.0
43.8
67.6
47.1
Vit-Adapter-S [4]
403G
44.7
65.8
48.3
39,9
62.5
42.8
48.2
69.7
52.5
42.8
66.4
45,9
Swin-S [43]
354G
66.6 44.8
48.9
40.9
63.4
44.2
48.2
69.8
52.8
43.2
67.0
46.1
ConvNeXt-T [44]
348G
67.9 45.4
50.0
41.8
65.2
45.1
47.9
70.0
52.7
42.9
66.9
46.2
InternImage-S [62]
340G
47.8
69.8
52.8
43.3
67.1
46.7
49.7
71.1
54.5
44.5
68.5
47.8
VMamha-S [41]
400G
48.2
69.7
52.5
43.0
66.6
46.4
49.7
70.4
54.2
44.0
67.6
47.3
L-Vmamba-S [34]
414G
48.4
69.9
52.7
43.2
66.7
46.5
49,9
70.5
54.4
44.1
67.8
47.4
Va (Ours)
341G
48.6 70.3
53.5
43.6
67.5
47.1
50.1
71.2
54.9
44.6
68.7
47.8
Table 8: Object detection and instance segmentation performance on COCO val2017. AP and APM indicate the mAP of detection and segmentation, respectively. MS indicates the multi-scale training strategy.
B Language Tree Topology Scanning Operator
Algorithm 2 Language Tree Scanning
Input: Input feature { Input matrix (B.) State matrix {A,) Gradient of loss to hidden states (LLoss @hh Ji=1 Minimum Spanning Tree Gr.
Traverse Path: Root, ...., Leaf + BFS(Gr)
Breadth -frst topological order of Gr
Forward:
Initialization: (Eiter & {ai)te
2: for +- Leaf to Root do
E Bi ExsE((Pa(()=)) Aj
4: end for
Backward:
Root Leaf do 6: is Root then
for i
01.oss Sz.
B,B.
BLass
77
0l.oss DA;
UB.
8: else 2Loss
BLose
B; + A,
&Loas
B,
BLoss ØB.
BLoss + A aLoss BBPu(i)
Oz.
Hh.
8Loss
8Loss hi
DA,
end if
12: end for Output: Hidden states ( Grad. of loss input feature
{ (LLss dx: Ji=1i Grad, of loss input
matrix { (Booss) L
Timis Grad. of loss to state matrix { (osss L
0A. liml.
C Algorithm Proof
In this section, we present detailed proofs our tree scanning algorithm. The definitions of symbols consistent with those in the main paper.
16

---

C.1 Proof for Algorithm 1.
We randomly take vertex in the MST Gr the root. According the definition of the tree scanning algorithm introduced Sec. 3.2, we can derive froot as follows:
Rroot
Σ
S(Eroot.)) Byz ((rrott,))
II
Ag.
(9)
AjECrool
kENrooot.
which shows a process of aggregation from all leaf vertices to the root. Therefore, each vertex is only related to its child in this period. Taking vertex as example, the Aggr, can be derived as:
Aggr, (x)
Bm +
Σ
Aggr, (x)
(10)
(kE())Par()=s}
We assume that one of the child of m is n and hn can be derived following:
hn Ager, (x) + An (x),
(11)
where Aggr,m (z) indicates the aggregation value from the vertices E Q Croot to vertex m. Therefore, we can obtain the propagation relationship between the hidden state of parent and child n:
hn Aggr,,((0 (x) + AnAggr, (x) Aggr, (x) + An AnAggr, (x))
(12)
Am + +(1 Aggr,, (x)
Through the above derivation, we can calculate { with only two traversals (i.e., the aggre- gation from f to root and the propagation from root to leaf) in the forward process as shown in Algorithm 1, thereby reducing the computational complexity from O to O(L).
Next, analyze the backpropagation process in Algorithm 1. According to the chain rule, easily calculate the derivative loss with respect to xz
Ōloss
i loss Oh, Ohg
ou:
Dz.
jER
(13)
Mloss
B Σ S (E
Th,
jen
Similarly, the derivative of loss with respect B, is:
Ōloss
loss Th
0B;
Oh 0B,
jen
(14)
Ôloss
Σ S (E)))
Th
jen
The above formulas equivalent to replacing the input with Gloss during the forward process.
17

---

Subsequently, we assume that the vertex & is the child of vertex and define Cf indicates the children of vertex with the root of vertex k. Ōloss is formulated as follows: 8A
Ōloss aA
Ōloss dhg
Ōhg
dA
JEn
Ōloss
aS (S(Ejp) B
Ōhj
0A
jen
Pen2
Σ
aloss
Σ
S(Exp)S(EE)) Bp + Σ Σ S(Ek))S Bpp))
Aloss
ohj
Ohg
jECt
pEc
jEC
PECt
Õloss
Ōloss
Σ
S(Eji)
Σ S(Ekp) Bpa +
S(E)))
Σ S(pp))Bpp
Thy
Jhy
jECt
pEci
jEC
PEOf
ØLoss AgggT Oh
òloss
bloss
}) Aggr (x) + Aggr
* (h Aggr
Ozk
Oh
ôLoss
bloss hh AAg
Ōloss
Aggr (x) Aggr
Aggr, (x)
Oak
Oh
0Lo8s
EK + mchk 2A
(definition in Algorithm 1)
Ozk
(15)
So far we have completed the proof forward and back-propagation of Algorithm
C.2 Proof for Algorithm 2.
We only take the last token as root and replace the transition source from Ω to C₁ in sequence modeling tasks like nature language understanding to ensure causality. Therefore, only one traversal (from leaf f to root) is required for the forward process, and another traversal (from root to leaf) is needed for the backpropagation process. The proof is similar to the Algorithm
Input
TP
Raster
Cross
TP
Raster
(a)
(b)
(d)
Figure 6: Visualization of affinity maps the specific position. The Location marked by the red cross in each affinity map. TP represents our Tree Scanning Algorithm.
18

---

D More Qualitative Results
Fig. displays additional qualitative comparisons between our algorithm and previous scanning strategies (e.g., cross- cross-scanning and raster-scanning), which shows our advanced capability to perceive detailed structural information and capture long-range dependencies.
E
Statistical Significance
Method
PIQA
Arc-Easy
SST
WinoGrande
L.M--pll
Race
Openhookqu
MambaTreel. (Ours)
0.011
0.010
0.016
0.014
0.553
0.014
0.018
Table 9: Standard error on language model benchmarks. LAM-ppl indicates LAMBADA [49].
We calculate the standard deviation of our MambaTreeL. on language model benchmarks in the open-sourced ievtlluaaiio-hhaneess project as shown Table
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims in the abstract and introduction accurately reflect motivation and contribution.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations of our work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
20

---

Justification:
We provide detailed proofs for our tree topology scanning algorithm in both the main paper and the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We describe the experimental details both the main paper and the appendix. We will make our code publicly available, along with detailed instructions for reproduction.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
21

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide the core part of our code in the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
//niss httppsccccpbccccccceececccccclcc /tt.dde/ooooeeummmiieeee /Cooeeeooiissseeeeeee wwssssnoolicgo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: The complete experiment details are introduced in the main paper and the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We provide the statistical analysis results in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have described the resources required to perform our experiments in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: : //te////neeee http//nnriippppe http//bbll//cchccctsiieieeeeeeee
Answer: [Yes]
Justification: Our work conforms with the NeurlPS Code of Ethics in every respect
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: Our work mainly focuses on the research of basic neural network architecture. We have identified no potential negative social impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
The creators or original owners of assets (e.g., code, data, models), used the paper, properly credited well as the license and terms of use explicitly are mentioned and properly respected.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our work does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
25

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "vS6lBTp1KL", "reviewer": "Reviewer_NQEG", "summary": "This paper studies the optimization of selective stat space modeling by particularly proposing the GrootVL model. Specifically, it firstly constructs the tree topology based on spatial information and then aggregates the features to enhance the representation informativeness. The proposed methods are versatile for both visual and textual tasks. The experimental results generally demonstrate the effectiveness of the proposed model.", "questions": "1. Since the Tree structure is a graph as well, how to relate and/or differentiate the connections between graph-based mamba models, e.g., [1].\n2. How many times of evaluations were performed to achieve the results?\n\n[1] Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces", "limitations": "Please see the weakness for details.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper studies the problem of mamba framework optimization with a particular effort in the design tree structure learning module, which is an interesting investigation for the related domain.\n2. The paper is written with good clarity and thus it is easy to follow.\n3. The experimental results generally demonstrate the effectiveness of the proposed methods to support the claims made by this paper.", "weaknesses": "1. One of the major concerns is the insufficient efficiency analysis of the proposed Tree Scanning Algorithm in terms of complexity analysis and/or runtime costs for this related module. \n2. Another concern is there is no significance test for the proposed model evaluation, especially when the metrics achieved by Grootvl are numerically close to the baselines. Furthermore, it is unclear the performance deviations in evaluation, which is also important to indicate the stability of the proposed model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: More efficiency analysis.**\n\n**Ans:** Thanks for your valuable suggestions! We have introduced the complexity and optimized version of our method in Section 3.1 and Section B of our manuscript. For the comparison of inference time, please refer to \"All Reviewer\" section at the top of this rebuttal page. We will further clarify it in the revision.\n\n**Q2: There is no significance test and unclear illustrations on the performance deviations.**\n\n**Ans:** We adhere strictly to the same benchmark evaluation protocols to ensure a fair and standardized comparison with previous models. For language tasks, we have included comprehensive significance results in the Appendix of the manuscript (Section E). For vision tasks, as the reviewer mentioned, we have retrained our models three times for the segmentation task, whose standard variation is about 0.11%. The results demonstrate that our approach consistently improves performance compared to other counterparts. We will add more significance tests in the revision.\n\n**Q3: How to relate and/or differentiate between the Graph-based mamba models?**\n\n**Ans:** This is an insightful question. The primary difference between our GrootVL and Graph-Mamba[1] lies in the input type and topology construction manner. Graph-Mamba directly utilizes a graph structure as input, which includes both node and edge embeddings, and keeps the topology through the whole process. In contrast, GrootVL takes images or text as input and dynamically constructs the topology structure based on input feature.\n\n**Q4: How many times of evaluations were performed?**\n\n**Ans:** For language tasks, the evaluation is calculated three times by using the most popular benchmark, lm-evaluation-harness[2]. For vision tasks, we have additionally retrained Groot V-T three times on semantic segmentation and got 0.11% standard variation. We will add more clarification in the revision.\n\n[1] Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces\n\n[2] A framework for few-shot language model evaluation", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cJbjQhuGRd", "reviewer": "Reviewer_xmLw", "summary": "The paper introduces a tree scanning algorithm for state space models specifically for Mamba. The naive and fixed scan patterns like raster or local scans commonly used for vision tasks do not consider the topological structure of 2D image input. The proposed algorithm generates a minimum spanning tree which can help Mamba to model the semantic information of the input. The paper further introduces a dynamic programming procedure to avoid the quadratic complexity of the tree scanning algorithm. The experiments on various vision and language modeling tasks show that the tree topology scan helps improve accuracy.", "questions": "1. In eq 5, Ω is the index set of all vertices. As $S(E_{ij})$ aggregates the connected vertices, some vertices will be visited many times during the recurrent scan in Mamba. It seems like redundant computations. Could authors either empirically or theoretically justify why this is needed? \n\n2. In language understanding, the tree scanning algorithm is applied during finetuning. Intuitively, this will change the causality of tokens trained during pretraining. Is there any specific reason the algorithm is not used during pre-training?\n\n3. How is Eq 5 (state aggregation process) derived from Eq 4 (SSM state computation)? Could authors provide more explanation about it? \n\n4. The semantic segmentation results seem marginal compared to other tasks, and this is the only dense prediction task. Can this be due to aggregation instead of processing all pixels by Mamba?", "limitations": "The limitation is discussed and isreasonable.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The motivation and the proposed scan algorithm make sense and are easy to understand. \n2. The proposed method can be useful not only for SSMs but also for other models that require sequential scans. \n3. As the proposed scan already generates a structure of the input based on the relevance of tokens, the sequence learning process can potentially be minimal. \n4. Using the proposed method, GrootV outperforms many recent baselines for multiple tasks (image classification/segmentation, object detection, and language understanding).", "weaknesses": "1. Regarding the root setting, the authors show that the root setting to all vertices outperforms the ones with only the first or the last. However, it increases the traverse time (1 vertex vs all vertices), and the accuracy improvement using all vertex is marginal (only 0.4% on Imagenet-1K). I understand that the dynamic programming procedure improves the speed but still increases by the sequence length $L$. I am still not fully convinced how effective this is. Can authors comment on it? \n\n3. Related to point 2, there are no speed comparisons, especially with different scanning strategies, root settings, and the use of the dynamic programming procedure in practice.\n\n1. Some details are missing in the paper. \n    1. Figure 1 is not explained in the paper. What exactly is the parameter generator?  Is it the initialization stage of the state, input, and output matrices (A, B, C, and D), or is the projection of these parameters the same as in Mamba? In the tree state space model part, is the main difference using TSA instead of eq 4 compared to the original Mamba? \n    2. In Figure 4, how is the specific position defined? Is it based on the root vertex setting? \n    3. The cross scan used in Table 4 is not described.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: The traverse time of all vertices using dynamic programming algorithm.**\n\n**Ans:** Thanks for the comments, but we believe there exists some misunderstanding. Given a sequence with the length of $L$ with an established corresponding minimum spanning tree, for the case of single-vertex setting, we treat it as the root of a tree and aggregate features from other vertices, which operate in $O(L)$ complexity. While for the all-vertices setting, a naive approach treats each vertex as a root separately, resulting in $O(L^2)$ complexity. In contrast, this paper proposes a dynamic programming algorithm where a random vertex is chosen as the root, features are aggregated from leaf vertices to the root, followed by propagation from the root to the leaves, achieving the same effect. Therefore, the complexity of GrootV for all nodes remains linear at $O(L)$. We will further clarify it in the revision.\n\n**Q2: There are no speed comparisons.**\n\n**Ans:** For the comparison of inference speed, please refer to \"All Reviewer\" section at the top of this rebuttal page. We will further clarify it in the revision.\n\n**Q3: Some details are missing.**\n\n**Ans:** Thanks for your valuable suggestions!\n\n1) Figure 1 illustrates a comparison of different propagation strategies for vision tasks and language tasks as discussed in Line 33 and Line 44 of the main manuscript.\n\n2) The parameters generator in Figure 2 utilizes the same projection network as Mamba to unleash the context-aware capability of state space modeling. The only difference between Tree SSM and Mamba lies in the replacement of the structured state space block with the proposed tree scanning algorithm (referring to Line 130 of the main manuscript).\n\n3) The anchor points shown in Figure 4 are randomly selected. We visualize the affinity maps of different positions to illustrate the capability of TSA to preserve detailed structural information. More qualitative results are shown in Sec. D of Appendix in the manuscript. Benefiting from the superiority of TSA, all pixels have equal access to long-range context.\n\n4) The cross scan in Table 4 is the 4-directional raster-scanning strategy shown on the left side of Figure 1 in the manuscript. \n\nWe will revise our paper to provide more details.\n\n**Q4: Explanation for why some vertices will be visited many times.**\n\n**Ans:** There could be some misunderstandings. For causal reasoning in text tasks, each node is visited only once. In image tasks, due to the use of the proposed dynamic programming algorithm, each node is only visited twice. The detailed mechanism refers to Section 3.2 of our manuscript and the answer in Q1.\n\n**Q5: Is there any specific reason the algorithm is not used during pre-training?**\n\n**Ans:** The core objective of this paper is to introduce a new structure for state space models. Through strict ablation studies, we validate the effectiveness of the method using LoRA fine-tuning for language tasks. Scaling up the models and exploring pre-training settings are left for future work.\n\n**Q6: How is Eq 5 derived from Eq 4?**\n\n**Ans:** The original Mamba state transition formula (Equation 4) can be easily derived into the form of Equation 5 as follows:\n$$\n\\begin{aligned}\nh _ {i}&=\\bar{\\mathbf{A}} _ {i}h _ {i-1}+\\bar{\\mathbf{B}} _ {i}x _ {i}\\\\\\\\\n     &=\\bar{\\mathbf{A}} _ {i}\\bar{\\mathbf{A}} _ {i-1}h _ {i-2}+\\bar{\\mathbf{A}} _ {i}\\bar{\\mathbf{B}} _ {i-1} x _ {i-1}+\\bar{\\mathbf{B}} _ {i}x _ {i}\\\\\\\\\n     &\\cdots\\\\\\\\\n     &={\\textstyle\\sum _ {j=1}^{i}}{\\textstyle\\prod _ {k=j+1}^{i}}\\bar{\\mathbf{A}} _ {k}\\bar{\\mathbf{B}} _ {j}x _ {j}\n\\end{aligned}\n$$\nThe feature aggregation described by the above formula is based on a linear topology (from the first vertex to $i$-$th$ vertex). If we propagate the state of each vertex along the built tree-topological path, Equation 4 is firstly transformed to the following formula (from children vertices to their parent vertex):\n$$\nh _ {i}={\\textstyle\\sum _ {j\\in \\\\{ k|\\text{par(k)=}i \\\\}}}\\bar{\\mathbf{A}} _ {j}h _ {j}+\\bar{\\mathbf{B}} _ {i}x _ {i}\n$$\nThen it can be easily derived into Equation 5 in a similar way as above, the detailed definition can be seen in Sec 3.2 of the manuscript.\n\n**Q7: The semantic segmentation results seem marginal compared to other tasks. Can this be due to aggregation instead of processing all pixels by Mamba?**\n\n**Ans:** As the same as Mamba-based methods, we aggregate features for all vertices, which is elaborated in Q1. Besides, as shown in Table 2 and Table 7 in the manuscript, our method shows consistent improvements over other SSM-based methods in terms of accuracy and efficiency.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bnzXO1sEBS", "reviewer": "Reviewer_g9ag", "summary": "This paper proposes a tree state space model (SSM) to perform feature propagation on an input-aware topology. The author explores tree topology in SSM from both vision and language sides, leading to GrootV and GrootL respectively. The proposed method exhibits strong empirical performances on mainstream tasks. Extensive ablations are conducted to verify the approach's effectiveness. Missing training efficiency and inference throughput.", "questions": "Please see the comments above.", "limitations": "Please see the comments above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strength：\n\n1. The motivation is clear and tree topology in building the scanning mechanism is novel to me.\n2. Extensive experiments demonstrate the strong empirical performances of the proposed method.\n3. Abundant ablation studies are conducted.", "weaknesses": "Weakness:\n\n1. Missing efficiency performance. Mamba is known for its linear complexity and efficiency. However, there is no efficiency experiments to report this property of Groot-V/L. It's important for me to know the training efficiency and inference throughput especially for Groot-V.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: Missing efficiency performance.**\n\n**Ans:** Thanks for the suggestion. For inference throughputs, please refer to the \"All Reviewer\" section at the top of this rebuttal page. Besides, we provide the training throughputs of our method in Table 10, which are measured on a Nvidia V100 GPU with an input image scale of $224\\times224$. We believe the efficiency can be further improved with more sophisticated optimization.\n\n**Table 10. Comparison of the throughputs during training.**\n|| GrootV-T | GrootV-S | GrootV-B |\n|:-:|:-:|:-:|:-:|\n| Throughput (img/s) | 171      | 101      | 92       |\nWe will update it in the revision!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VTs2zoZuel", "reviewer": "Reviewer_LPAH", "summary": "This paper integrates tree structures into SSMs, enabling hierarchical data processing instead of the traditional sequential approach. It identifies a minimum spanning tree that considers locality, utilizing dynamic programming to prune edges from the grid graph. By adopting this tree structure, the length of the original sequence is reduced, allowing SSMs to handle long-range dependencies more effectively. As a result, the proposed GrootVL outperforms previous SSMs in image classification, detection, segmentation, and language understanding.", "questions": "N/A", "limitations": "The paper clearly states its limitation: the current architecture is not implemented in a hardware-efficient manner. However, I believe the paper provides enough academic insights, and practical extensions could be left for future work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Integrating hierarchical structures into sequential models is an important yet underexplored direction.\n- This paper proposes an intuitive yet novel algorithm that implements tree scanning into SSM architectures.\n- Enhancing long-range dependencies is a significant benefit of the tree structure. In theory, it can reduce the context length from $L$ to $\\log L$.\n- The experiments are extensive, demonstrating results in various image and language tasks.\n- The paper is well-structured, featuring a clear logical flow and ablation studies.", "weaknesses": "1. Analysis on the learned structures\n\nWhile the paper primarily highlights performance benefits, tree structures offer additional advantages, such as interpretability. For example, as illustrated in Figure 1, one can identify syntax structures in sentences or scene structures in images. Examining the learned tree structures for individual data in depth would be highly insightful.\n\n---\n2. Consistency of tree structures across layers\n\nIf I understand correctly, the current architecture constructs a tree for each block independently. Do the layers generate consistent trees? If not, how can one justify the validity of the discovered structure? Would it be helpful to regularize the tree structures to ensure consistency across layers?\n\n---\n3. Discussion of previous hierarchical structure approaches\n\nThe related work section lacks a comparison with previous hierarchical structure approaches. Classic recursive neural networks [1] pioneered learning hierarchical structures from images and texts, and this concept has been extended to modern Transformer architectures [2-4]. However, those approaches build trees in a bottom-up manner by gradually grouping the original tokens, while this paper proposes a top-down approach that prunes edges from grid graphs. This methodological difference is noteworthy and merits discussion. Additionally, other approaches to encoding trees into sequential architectures, such as [5], are also worth discussing.\n\n[1] Socher et al. Parsing Natural Scenes and Natural Language with Recursive Neural Networks. ICML 2011.\\\n[2] Wang et al. Tree Transformer: Integrating Tree Structures into Self-Attention. EMNLP 2019.\\\n[3] Bolya et al. Token Merging: Your ViT But Faster. ICLR 2023.\\\n[4] Ke et al. Learning Hierarchical Image Segmentation For Recognition and By Recognition. ICLR 2024.\\\n[5] Shen et al. Ordered Neurons: Integrating Tree Structures into Recurrent Neural Networks. ICLR 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: Analysis of the learned structures.**\n\n**Ans:** Thanks for the valuable suggestions! We acknowledge that input-aware tree structures provide significant interpretability benefits, including preserving the intricate structural details in the vision content and enhancing the long-range modeling in a sequence. In addition, tree scanning algorithm captures the higher-order relationships between each unit, significantly expanding the feature space of the model compared to the second-order relationships handled by attention mechanisms. We will revise our paper by providing more analysis.\n\n**Q2: Consistency of tree structures across layers.**\n\n**Ans:** In our manuscript, we dynamically generate the tree topology based on the specific input feature for each block separately, which achieves higher performance. We have experimented to explore the effect of a regularized tree structure as noted by the reviewer. The results are shown in Table 11. GrootV-T* refers to each stage sharing the same tree structure. This approach enhances efficiency with only a minimal compromise in accuracy. \n\n**Tabel 11. Ablation study about tree scanning algorithm.**\n| **Method (224x224)**   | **Throughput (img/s)** | **Acc. (Top1)** |\n|--------------------|:-:|:-:|\n| Baseline (w/o TSA) | 373                | 82.60       |\n| GrootV-T           | 281                | 83.44       |\n| GrootV-T*          | 392                | 83.41       |\nWe will update this ablation study in the revision.\n\n**Q3: Discussion of previous hierarchical structure approaches.**\n\n**Ans:** Thanks very much for the insightful perspective and suggestions. It is the main difference from the previous bottom-up hierarchical structure that our tree-based topology retains original vertices to propagate features in a top-down manner. Method[1] as mentioned by the reviewer, captures the hierarchical representation by encoding an ordered tree structure to the sequence. Compared to it, our tree topology is essentially an undirected and acyclic graph, which can be dynamically constructed based on the input signal. In a word, it’s an interesting topic. We will supplement the related work section in our revised manuscript.\n\n[1] Ordered Neurons: Integrating Tree Structures into Recurrent Neural Networks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
**To All Reviewers (Reply about efficiency performance):**

We sincerely appreciate all reviewers and ACs for their precious time and valuable feedback. Given that reviewers NQEG, xmLw, and g9ag have raised concerns regarding efficiency comparison, we will respond to this issue in this section. All the source code will be made public.

**Q1: Inference speed comparison.**

**Ans:** As shown in Table 9, we report the inference throughputs of our method on a Nvidia V100 GPU. The GrootV-T* refers to each stage sharing the same tree topology structure, which enhances efficiency without compromising accuracy. To achieve better practical inference speed, we also introduce a cuda implementation optimized for GPUs. Compared with other counterparts, our approach exhibits superior effectiveness and faster inference speed. We will add this table to the revision and release the optimized cuda code.

**Table 9: Runtime comparison on a Nvidia V100 GPU during inference.**
|**Method(224x224)**|**Throughput (img/s)**|**GPU Memory**|**FLOPs**|**#Params.**|**Acc. (Top1)**|
|-|:-:|:-:|:-:|:-:|:-:|
| PlainMamba-L2      | 363             | 4204M               | 8.1G      | 25M          | 81.6     |
| VMamba-T           | 374             | 8646M               | 4.9G      | 31M          | 82.5     |
| LocalVMamba-T      | 311             | 11298M              | 5.7G      | 26M          | 82.7     |
| GrootV-T(one root) | 283             | 6012M               | 4.8G      | 30M          | 83.0     |
| GrootV-T           | 281             | 6471M               | 4.8G      | 30M          | 83.4     |
| GrootV-T*          | **392**         | 4800M               | 4.8G      | 30M          | **83.4** |

Meta-review: The authors proposed GrootVL -- a tree-structured state-space model (SSM) which efficiently models long-range interactions across a range of tasks including image classification, semantic segmentation and language understanding.

Reviewers appreciated that the method can be applied to different SSMs, and achieves strong empirical results. As a result, they all recommended acceptance.

A common concern among reviewers was about the efficiency analysis. The table in the rebuttal answered this, so please include it in the camera-ready version of the paper. Different architectures have different efficiency characteristics, and hence it is important to report multiple efficiency indicators to give a fair representation (as also mentioned in more detail in https://arxiv.org/abs/2110.12894)
Decision: Accept (spotlight)