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_cc_2024_Conference
Published: 2024-05-11 10:42:20

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.

Main paper: 
 introduction: 
 Mainstream fundamental models are primarily based on CNN [30, 62, 44, 32, 13] and Transformer architectures [15, 43, 42, 59, 14] , 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 a good balance between effectiveness and efficiency. The state space models (SSMs) [22, 24, 52] attempt to 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 to attain stable optimization and superior computing performance. Nevertheless, it remains susceptible to the intrinsic flaw shared by recurrent neural networks, i.e., a 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 in NLP tasks. Besides, numerous studies aim to extend Mamba into computer vision, by employing various pre-defined strategies to map 2D image features into 1D sequences. ViM [77] and VMamba [41] utilize a multi-directional raster-scanning strategy, while LocalMamba [34] further confines its Figure 1 : Comparison of different 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 a diagonal matrix, prompting a subsequent series of research on structured SSMs [24, 21, 26, 19] . Recently, the Selective State Space Model [19] , known as Mamba, strikes a 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 to showcasing superior outcomes in sequence modeling, Mamba has been seamlessly incorporated into the visual domain [77, 41, 34, 69, 68] . These studies often rely on handcrafted fixed scanning mechanisms to 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 MambaTree in this work to enhance long-range modeling for both vision and language tasks by introducing an input-aware tree-based topological structure. 
method: 
 In this section, we first revisit the selective state space model [19] and then elaborate on our inputaware topology scanning algorithm for state space modeling. With this superior algorithm, we develop a tree SSM and propose a novel framework called MambaTree, which consists of two sub-networks: MambaTreeV for visual tasks and MambaTreeL for fine-tuning a pre-trained language model [19] .State Space Models (SSMs) are commonly regarded as continuous linear time-invariant systems [64] that map input stimulation x(t) ∈ R 1×D to output signal y(t) ∈ R 1×D through a state vector h(t) ∈ R 1×N , where t, 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:EQUATIONwhereA ∈ R N ×N , B ∈ R N ×D , C ∈ R N ×D and feedthrough coefficient D ∈ R D .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 EQUATIONIn addition, many improved methods [41, 19] use an approximation of B based on the first-order Taylor Series:EQUATIONSelective Mechanism . Previous SSMs store information through finite states and inherent timeinvariance, which limits their effectiveness. Therefore, Mamba [19] introduces a dynamic mechanism to selectively filter out input into a sequential state. Specifically, it utilizes Linear Projection to calculate the parameters {B i } L i=1 , {C i } L i=1 and {∆ i } L i=1 from the input sequence {x i } L i=1 with x i ∈ R 1×D directly to improve the context-aware ability. Then the output sequence {y i } L i=1 can be computed with those input-adaptive discretized parameters as follows:h i = Āi h i-1 + Bi x i , y i = C i h i + Dx i (4)Mamba [19] has showcased remarkable performance in modeling the dependencies of consecutive words in a sequence. However, its applicability in long-context 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] , we construct a transform block as a 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. In addition, our algorithm can effectively enhance language representations by incorporating such a tree topology during text processing, which overcomes the geometrical constraints of text sequences. In the following, we elaborate on the proposed tree scanning algorithm and its applications for multi-modal tasks.Tree Scanning Algorithm. Given an input feature X = {x i } L i=1 where L is the sequence length (or the number of input pixels), we construct an undirected m-connected graph G = (V, E) for the feature. m is a hyper-parameter that indicates the number of adjacent tokens. Following [71, 54] , we set m = 4 for visual tasks, meaning each pixel is connected to its four neighboring pixels. For language tasks, we set m = 3 by default, meaning each token is connected to the previous three tokens. In addition, the vertices V represent the pixel (or token) embeddings, and the E indicates the edges of the graph. The edge weight is calculated by the feature dissimilarity between adjacent vertices. Besides, the metric of dissimilarity uses cosine distance by default, and the comparison with other metrics refers to Table 5 .We use the Contractive Boruvka algorithm [2] to prune the edges with significant dissimilarity, which generates a minimum spanning tree (MST) G T whose sum of dissimilarity weights is minimum out of all spanning trees. In the propagation process, we iteratively traverse each vertex, treating it as the root, and aggregate the features of the remaining vertices. Intuitively, applying state propagation within such a geometric configuration makes its preferential interactions among vertices with small spatial and feature distances. Following the Mamba, we employ the data-dependent transition matrix for state propagation. For a vertex k, we denote the transition matrix with its parent as Āk . Furthermore, following the Eq. ( 4), the state aggregation process for the i-th vertex can be formulated as:EQUATIONwhere Ω denotes the index set of all vertices in the tree. S(E ij ) represents the path weight of hyperedge E ij traced from j-th vertex to i-th vertex in the tree G T , and N ij indicates the index set of all vertices on this hyperedge. For visual tasks, we iterate over each vertex, treating it as the root of the spanning tree G T , and aggregate the states from the other vertices, thereby obtaining the transformed states {h i } L i=1 . For textual tasks, because of the causal prediction manner in large language models, we only take the last token as root and aggregate from other tokens. To achieve end-to-end training, we derive the derivative of the output hidden state h i to the input variables Āk , Bj and x j as follows:EQUATION)EQUATIONwhere C i k indicates the children of vertex k in tree G T whose root is the vertex i, and n denotes the parent of vertex k in Eq. (7) . Finally, the output feature Y can be formulated as:EQUATIONwhere Y , H and X indicate the stack of {y i } L i=1 , {h i } L i=1 and {x} L i=1 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(L 2 ). To alleviate this issue, we utilize a dynamic Algorithm 1 Vision Tree ScanningInput: Input feature {x i } L i=1 ; Input matrix { Bi } L i=1 ; State matrix { Āi } L i=1 ; Gradient of loss to hidden states { ∂Loss ∂hi } L i=1 ; Minimum Spanning Tree G T . Traverse Path: Root, . . . , Leaf ← BF S(G T )▷ Breadth-first topological order of G T Forward:Initialization: {ξ i } L i=1 ← {x i } L i=1 2: for i ← Leaf to Root do ξ i = Bi x i + ∀j∈{t|Par(t)=i} ξ j Āj 4: end for for i ← Root to Leaf do 6: if i is Root then h i = ξ i 8:elseh i = Āi (h Par(i) -Āi ξ i ) + ξ i = (1 -Ā2 i )ξ i + Āi h Par(i) 10:end if end for Backward: 12: Initialization:{η i } L i=1 ← { ∂Loss ∂hi } L i=1for i ← Leaf to Root do 14:η i = Bi ∂Loss ∂hi + ∀j∈{t|Par(t)=i} η j Āj end for 16: for i ← Root to Leaf do if i is Root then 18: ∂Loss ∂xi = η i Bi , ∂Loss ∂ Bi = η i x i , ∂Loss ∂ Āi = 0 else 20: ∂Loss ∂xi = (1 -Ā2 i )η i Bi + Āi ∂Loss ∂x Par(i) Bi , ∂Loss ∂ Bi = (1 -Ā2 i )η i x i + Āi ∂Loss ∂ BPar(i) x i ∂Loss ∂ Āi = η i * (h i -Āi ξ i ) + ξ i * ( ∂Loss ∂xi -Āi η i ) = η i h i + ξ i ∂Loss ∂xi -2η i ξ i Āi 22:end if end for Output: Hidden states {h i } L i=1 ; Grad. of loss to input feature { ∂Loss ∂xi } L i=1 ; Grad. of loss to input matrix { ∂Loss ∂ Bi } L i=1 ; Grad. of loss to state matrix{ ∂Loss ∂ Āi } L i=1 .programming procedure to accelerate the inference and training processes as elaborated in Algorithm 1, which results in linear complexity O(L). For textual tasks, we perform a unidirectional aggregation approach (shown in Algorithm 2 of Appendix B) in adherence to the causal nature of language. Moreover, we provide the back-propagation process for both Vision Tree Scanning and Language Tree Scanning processes, whose detailed proofs refer to Appendix C.MambaTreeV Given an image with a shape of H × W × 3, our goal is to obtain high-quality visual features for downstream tasks. To this end, we propose an effective vision architecture MambaTreeV which consists of a stem module, several basic blocks and downsampling layers to generate hierarchical representations illustrated in Fig. 3 . Overall, our MambaTreeV 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 a feature map with a shape of H 4 × W 4 × C. It includes two convolutions, two Layer Normalization (LN) layers and one GELU activation function. The kernel size for both convolutions is 3 with a stride of 2 and padding of 1. Similarly, a downsampling layer consists of a 3 × 3 convolution with a stride of 2 and padding of 1 and an LN layer. Positioned between two stages, it serves to downsample the input feature map by a factor of 2. Motivated by [62, 41] , we devise a residual block with skip connections to integrate our fundamental Tree State Space Model in Sec. 3.2. In detail, we first normalize the input features with LN layer. Spatial priors and long-range dependencies are then obtained through our tree scanning algorithm with residual connections established alongside the input features. Finally, a feedforward neural network is utilized to project the normalized features to output signals as shown in Fig. 3 . Based on the above origin components, we develop our MambaTreeV in three scales, i.e., MambaTreeV-Tiny, MambaTreeV-Small and MambaTreeV-Base.MambaTreeL Recurrent neural networks rely on fixed memory to preserve past information, which poses limitations when handling long contexts where relevant words are distant from the current moment. While Mamba [19] employs a selection mechanism to enhance context awareness, its fixed memory size cannot expand over time, resulting in restricted state space. Therefore, the ability to extrapolate decreases during scrolling as the prompt extends. To mitigate this issue, we propose an effective fine-tuning paradigm. Specifically, the tree-based topology branch is built upon one-way scrolling with a scaling factor, enabling state transitions within such a structure. This arrangement facilitates the preferential interaction of semantically related tokens. It is noteworthy that this paradigm does not introduce any additional training parameters. Instead, it utilizes pretrained state transformation parameters to conduct semantic aggregation by incorporating topological structures.Experimental results demonstrate the effectiveness of our approach. 
experiments: 
 We conduct extensive experiments to evaluate the effectiveness of MambaTreeV and compare it with advanced CNN-based, Transformer-based, 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.Settings. We assess the classification performance of MambaTreeV on the ImageNet-1k dataset [12] . Following previous practices [43, 44, 62, 41] , all MambaTreeV models are trained for 300 epochs from scratch using AdamW optimizer with a warm-up strategy of 20 epochs. During training, we utilize a Cosine Scheduler with an initial learning rate of 1 × 10 -3 and weight decay of 0.05. In addition, the exponential moving average (EMA) is also applied. Table 2 : Semantic segmentation performance on ADE20K val set. The crop size is all set to 512 2 . SS and MS denote singlescale and multi-scale testing, respectively.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] to validate object detection and instance segmentation tasks with Mask-RCNN [29] . Specifically, We adopt the AdamW optimizer with a learning rate of 1 × 10 -4 and batch size of 16 to optimize the model built upon our pre-trained classification backbones on ImageNet-1K. The training schedules include 1× (12 epochs) and 3× (36 epochs) with multi-scale data augmentation.Results. As depicted in Table 8 (in Appendix A.), our method outperforms existing methods on most evaluation metrics, especially for instance segmentation. Under 1× schedule, MambaTreeV-T achieves 47.0 in box mAP (AP b ), which is 1.1 points higher than ViM-S and 0.5 points higher than VMamba-T. It is worth noting that MambaTreeV-T outperforms ViM-S by 1.7 points with 1× schedule and LocalVMamba-T by 0.4 points with 3× schedule in mask mAP (AP m ). Moreover, the best AP b 50.1 and AP m 44.6 are obtained by MambaTreeV-S in 3× schedule with multi-scale training.Settings. To evaluate the semantic segmentation performance of our MambaTreeV series, we train our models with UperNet [65] initialized by pre-trained classification weights on 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 2 . MambaTreeV-T yields a clear improvement of +3.6 in single-scale mIoU compared to ViM-S and +1.9 in multi-scale mIoU compared to LocalViM-S. Furthermore, MambaTreeV-S boosts InternImage-S by 0.6 and 0.8 in single-scale and multi-scale respectively. We consider the preservation of intricate structural details through tree topology scanning to be particularly advantageous for segmentation tasks that require pixel-level perception.Method PIQA ↑ Arc-E ↑ SST ↑ WG ↑ L-ppl ↓ Race ↑ BQA ↑ Average ↑ Acc.Mamba [19] 64 We regard Mamba [19] with 130M parameters as 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 MambaTreeL 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 in the open-sourced lm-evaluation-harness project [18] for evaluation, including PIQA [1] , AI2-ARC [8] , SST [60] , WinoGrande, LAMBADA [49] , Race [36] and Openbookqa [47] . The results in Table 3 demonstrate that our 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.In this section, we conduct analysis experiments on ImageNet-1K 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 in vision recognition.Distance Metric. Before generating a minimum spanning tree from a connected graph, it is important to measure the edge weights between vertices. Therefore, we validate several distance metrics as illustrated in Table 5 . 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 in Table 6 . 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 on an Nvidia V100 GPU. MambaTreeV-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.Qualitative Results. To better illustrate the superiority of our scanning strategy, we visualize the affinity maps of different positions marked by the red cross in each input image. For example, we Table 7 : Runtime comparison on an Nvidia V100 GPU during inference.set the anchor point in the upper left corner of the sky as shown in the second row of in Fig. 4 (a).Our method can easily identify white houses, flagpoles, and the sky, which raster scanning fails to achieve. This demonstrates the capability of our algorithm to preserve detailed structural information. More comparisons can be seen in Fig. 6 (in Appendix D.) 
conclusion & limitations: 
 In this paper, we propose a tree state space model to perform feature propagation on an 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, we establish the general multi-modal networks to break the original sequence constraints and achieve stronger representation capabilities. Extensive experiments demonstrate the effectiveness of our method in both visual and language tasks. The limitation of our method is that the tree structure is not a common paradigm, and it needs to be specifically optimized according to the hardware device.  
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?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 a 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, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to 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. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?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 or 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). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a 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 no such risks.• Released models that have a 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 a 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]Justification: The creators or original owners of assets (e.g., code, data, models), used in the paper, are properly credited as well as the license and terms of use explicitly are mentioned and properly respected.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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution 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. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?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. • Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: In this section, we present detailed proofs for our tree scanning algorithm. The definitions of symbols are consistent with those in the main paper.C.1 Proof for Algorithm 1.We randomly take a vertex in the MST G T as the root. According to the definition of the tree scanning algorithm introduced in Sec. 3.2, we can derive h root as follows:EQUATIONwhich 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 m as an example, the Aggr m can be derived as:EQUATION)We assume that one of the child of m is n and h n can be derived as following:EQUATIONwhere Aggr m (x) indicates the aggregation value from the vertices ∈ Ω\C root m to vertex m. Therefore, we can obtain the propagation relationship between the hidden state of parent m and child n:EQUATIONThrough the above derivation, we can calculate {h i } L i=1 with only two traversals (i.e., the aggregation from leaf 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(L 2 ) to O(L).Next, we analyze the backpropagation process in Algorithm 1. According to the chain rule, we can easily calculate the derivative of loss with respect to x i :EQUATIONSimilarly, the derivative of loss with respect to Bi is:EQUATIONThe above formulas are equivalent to replacing the input x with ∂loss ∂h during the forward process.Subsequently, we assume that the vertex k is the child of vertex l and define C k l indicates the children of vertex l with the root of vertex k. ∂loss ∂ Āk is formulated as follows:EQUATIONSo far we have completed the proof of forward and back-propagation of Algorithm 1.We only take the last token as root and replace the transition source from Ω to C i in sequence modeling tasks like nature language understanding to ensure causality. Therefore, only one traversal (from leaf 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 1. Fig. 6 displays additional qualitative comparisons between our algorithm and previous scanning strategies (e.g., cross-scanning and raster-scanning), which shows our advanced capability to perceive detailed structural information and capture long-range dependencies. 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 lm-evaluation-harness project as shown in Table 9 .Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?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/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Justification: Our work conforms with the NeurIPS 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: 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)

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
2. 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.
questions: 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].
2. How many times of evaluations were performed to achieve the results?

[1] Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces
limitations: limitations: Please see the weakness for details.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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? 

3. 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.

1. Some details are missing in the paper. 
    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? 
    2. In Figure 4, how is the specific position defined? Is it based on the root vertex setting? 
    3. The cross scan used in Table 4 is not described.
questions: 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? 

2. 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?

3. How is Eq 5 (state aggregation process) derived from Eq 4 (SSM state computation)? Could authors provide more explanation about it? 

4. 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: limitations: The limitation is discussed and isreasonable.
paper_summary: paper_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.
strength_weakness: strength_weakness: Weakness:

1. 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.
questions: questions: Please see the comments above.
limitations: limitations: Please see the comments above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. Analysis on the learned structures

While 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.

---
2. Consistency of tree structures across layers

If 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?

---
3. Discussion of previous hierarchical structure approaches

The 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.

[1] Socher et al. Parsing Natural Scenes and Natural Language with Recursive Neural Networks. ICML 2011.\
[2] Wang et al. Tree Transformer: Integrating Tree Structures into Self-Attention. EMNLP 2019.\
[3] Bolya et al. Token Merging: Your ViT But Faster. ICLR 2023.\
[4] Ke et al. Learning Hierarchical Image Segmentation For Recognition and By Recognition. ICLR 2024.\
[5] Shen et al. Ordered Neurons: Integrating Tree Structures into Recurrent Neural Networks. ICLR 2019.
questions: questions: N/A
limitations: 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.

Comments: title: Response to the rebuttal
Response to the Rebuttal
comment: Thank you for the answers. 
My major concern was the efficiency of the algorithm. The answer to the speed comparison and the algorithm's details helped me understand the algorithm and resolved my concerns. Please make sure to add these details to the updated version. 
The paper includes clear strengths and improves SSMs for vision and language models. I increase my rating.
Thank you for the rebuttal. I find this paper interesting, and since all other reviewers have anonymously recommended acceptance, I stand by my original evaluation.

I'm looking forward to seeing the analysis of the learned structures (Q1) in the revised paper. It's also exciting to see that the consistency of tree structures across layers significantly improves throughput while maintaining accuracy (Q2). This makes sense, as the structure of data should generally remain consistent across layers unless there's a clear reason for it to change.