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

Q1: Analysis of the learned structures.

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.

Q2: Consistency of tree structures across layers.

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.

Tabel 11. Ablation study about tree scanning algorithm.

Q3: Discussion of previous hierarchical structure approaches.

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.

[1] Ordered Neurons: Integrating Tree Structures into Recurrent Neural Networks.

Q1: Missing efficiency performance.

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.

Table 10. Comparison of the throughputs during training.

Q1: The traverse time of all vertices using dynamic programming algorithm.

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.

Q2: There are no speed comparisons.

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.

Q3: Some details are missing.

Ans: Thanks for your valuable suggestions!

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

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

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

4. The cross scan in Table 4 is the 4-directional raster-scanning strategy shown on the left side of Figure 1 in the manuscript.

We will revise our paper to provide more details.

Q4: Explanation for why some vertices will be visited many times.

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.

Q5: Is there any specific reason the algorithm is not used during pre-training?

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.

Q6: How is Eq 5 derived from Eq 4?

Ans: The original Mamba state transition formula (Equation 4) can be easily derived into the form of Equation 5 as follows:

$$\begin{aligned} h _ {i}&=\bar{\mathbf{A}} _ {i}h _ {i-1}+\bar{\mathbf{B}} _ {i}x _ {i}\\\\ &=\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}\\\\ &\cdots\\\\ &={\textstyle\sum _ {j=1}^{i}}{\textstyle\prod _ {k=j+1}^{i}}\bar{\mathbf{A}} _ {k}\bar{\mathbf{B}} _ {j}x _ {j} \end{aligned}$$

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

$$h _ {i}={\textstyle\sum _ {j\in \\{ k|\text{par(k)=}i \\}}}\bar{\mathbf{A}} _ {j}h _ {j}+\bar{\mathbf{B}} _ {i}x _ {i}$$

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

Q7: The semantic segmentation results seem marginal compared to other tasks. Can this be due to aggregation instead of processing all pixels by Mamba?

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.

Q1: More efficiency analysis.

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.

Q2: There is no significance test and unclear illustrations on the performance deviations.

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.

Q3: How to relate and/or differentiate between the Graph-based mamba models?

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.

Q4: How many times of evaluations were performed?

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.

[1] Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces

[2] A framework for few-shot language model evaluation