Tree-Mamba: Tree-Based Neural SSMs

Updated 3 July 2026

Tree-Mamba is a family of neural architectures that combine linear-time state space models with tree-aware serialization to preserve topological context in spatial data.
They integrate methods like sparse tokenization, graph-regulated feature disentanglement, and geometry-guided sequencing to enhance performance in remote sensing, 3D segmentation, and underwater vision.
Empirical results demonstrate state-of-the-art accuracies, faster inference, and reduced memory usage compared to traditional attention models, underscoring their practical impact.

Tree-Mamba is a collective term for advanced neural architectures that integrate Mamba-style Structured State Space Models (SSMs) and tree-based serialization or tree-aware mechanisms to achieve efficient, topologically consistent, and context-sensitive modeling in remote sensing, 3D point cloud, and computer vision applications. These models leverage linear-time SSM recurrence, data-driven or domain-guided hierarchical structuring of spatial features, and token sparsification to overcome inefficiencies and context limitations typical of standard attention- or SSM-based models. Tree-Mamba variants have demonstrated state-of-the-art performance in tree species classification from MODIS time series, 3D forest point cloud segmentation, underwater instance segmentation, and underwater monocular depth estimation (Alkayid et al., 7 May 2026, Nguyen et al., 1 Jun 2026, Zhuang et al., 10 Jul 2025, Cong et al., 1 Aug 2025, Tu et al., 10 Oct 2025).

1. Foundational Principles and Model Variants

Tree-Mamba encompasses a family of architectures characterized by the joint use of state-space models and tree-structured, geometry- or similarity-driven serializations. The unifying principles are:

Linear-time state modeling: SSMs (as in Mamba) process input sequences with global receptive fields and memory of all past inputs, but with $O(L)$ complexity, where $L$ is the sequence length. The continuous-to-discrete state space mapping is given by

$h_i = A h_{i-1} + B_i x_i, \quad y_i = C_i h_i + D x_i,$

with input-dependent, selectively parameterized projections. This contrasts with the $O(L^2)$ cost of self-attention (Cong et al., 1 Aug 2025).

Tree-aware or geometry-guided sequencing: Rather than arbitrary raster or scanline traversal, these models employ data-driven minimum spanning tree (MST) serialization or ecological geometry cues to aggregate spatial features, preserving the underlying object or ecological topology (Zhuang et al., 10 Jul 2025, Nguyen et al., 1 Jun 2026).
Sparse tokenization: Adaptive selection of the most informative features (spectral bands, time steps, spatial tokens, or regions) sharply reduces computational redundancy and addresses the correlation decay issue in dense SSMs (Alkayid et al., 7 May 2026).
Multi-modal or multi-branch disentanglement: Many variants explicitly separate spectral, spatial, and temporal streams, processing them independently before fusion to exploit the independence of modes (Alkayid et al., 7 May 2026).

Model instantiations include:

Variant	Domain	Key Mechanisms
GDS-Mamba	MODIS tree species classification	Graph-regulated, sparse disentangling branches
ForestMamba	3D forest LiDAR segmentation	Vertical slab serialization, CHM-guided queries, SSM
UIS-Mamba	Underwater instance segmentation	Dynamic tree scan, Ncut-based state attenuation
Tree-Mamba-UMDE	Underwater monocular depth estimation	MST scan by feature similarity, dual-pass SSM
Minkowski-MambaNet	LiDAR biomass quantification	SSM within Minkowski backbone, channel recalibration

2. Graph- and Tree-Regulated Encoding for Remote Sensing

Graph-regulated Disentangled Sparse Mamba (GDS-Mamba) addresses the challenge of classifying tree species from long-term MODIS time series, where inter-class signature differences are subtle and high-dimensional spatial-spectral-temporal correlations are strong. The core innovations are:

Mini-batch graph-regulated stream: Each MODIS patch is treated as a node in a batch-graph; an RBF-kernel adjacency $A_{ij} = \exp(-\|x_i-x_j\|^2/\sigma^2)$ is computed, and a Graph Convolutional Network (GCN) propagates features via a normalized Laplacian $L$ . This encodes cross-patch topology beyond single-patch context (Alkayid et al., 7 May 2026).
Disentangling sparse Mamba stream: Input cubes are decomposed into spectral, temporal, and spatial tokens via group convolution. For each stream:
- Importance-based pruning (self-attention for spectral/temporal, cosine similarity for spatial) retains only the top- $k$ tokens.
- Dedicated Mamba blocks model long-range dependencies within each mode.
- Outputs are scattered back (via masking) and fused for classification.

This architectural design explicitly breaks spatial-spectral-temporal entanglement and enables the model to learn class-distinguishing features from subtle MODIS signatures.

Practical impact: GDS-Mamba achieves 93.94% overall accuracy on Alberta MOD13Q1, outperforming 12 baselines and providing sharper, less contaminated class boundaries. Robustness is validated through cross-provincial transfer (80.19% OA on Saskatchewan zero-shot) (Alkayid et al., 7 May 2026).

3. Structure-Aware Decoding in 3D Forest Segmentation

ForestMamba extends the Tree-Mamba paradigm to 3D forest point cloud segmentation, exploiting forest geometry:

Vertical-priority slab serialization: The point cloud is voxelized and grouped into vertical slabs (e.g., ground to canopy), then serialized for Mamba-based sequence processing. This respects ecological hierarchy and enables long-range context propagation within the SSM (Nguyen et al., 1 Jun 2026).
Geometry-guided query initialization: Candidate tree queries are generated from peaks in on-the-fly multi-scale Canopy Height Models (CHMs), which reflect the physical structure of forest canopies. Understory coverage is ensured by supplementing with Farthest Point Sampling (FPS) on discriminative embeddings.
Mamba-based query decoder: Each query aggregates features from local kNN voxels and exchanges information with other queries through a dual-path spatial Mamba module (two orderings), enabling efficient inter-tree disambiguation at $O(K)$ cost (with $K$ queries).

Performance: ForestMamba yields a 1.1pp instance F1 improvement and a 3 $\times$ speedup over Transformer baselines, with 2.3 $L$ 0 lower GPU memory on million-voxel scenes (Nguyen et al., 1 Jun 2026).

4. Tree-Structured Sequence Modeling: MST and Dynamic Tree Scan

Tree-aware and dynamic tree scan strategies replace rigid spatial traversals with hierarchies determined by feature similarity or geometric adaptation:

Tree-Mamba for underwater monocular depth estimation (Zhuang et al., 10 Jul 2025):
- Constructs a minimum spanning tree (MST) over a feature graph (cosine distance) on the 2D feature map.
- Aggregates features via bottom-up (leaf→root) and top-down (root→leaf) SSM traversals.
- Embeds this operator in a CNN-based multi-scale fusion framework for depth regression.
- Demonstrates superior quantitative (e.g., RMSE, $L$ 1 accuracy) and qualitative performance on the BlueDepth benchmark, with real-time throughput.
UIS-Mamba for instance segmentation (Cong et al., 1 Aug 2025):
- "Dynamic Tree Scan" (DTS): grid-aligned patches are adaptively offset and scaled, then linked into a deformable MST, aligning the scan with object boundaries distorted by underwater imaging.
- "Hidden State Weaken" (HSW): normalized cut (Ncut) partitions foreground/background; a multiplicative suppression attenuates background patches in the Mamba SSM recurrence.

These approaches exploit the MST to enforce semantic continuity and facilitate both local aggregation and global context via SSMs, offering performance advantages on tasks where spatial structure is non-Euclidean or highly variable.

5. Sparse and Selective Token Mechanisms

Tree-Mamba models leverage token sparsification to address the "correlation decay" bottleneck that arises when long, dense sequences are fed into conventional SSMs. The process comprises:

Importance scoring: Spectral and temporal importance scores are derived from multi-head self-attention maps; spatial tokens are scored by similarity to a central reference.
Token pruning and masking: Only the $L$ 2 most informative tokens are selected per mode/branch, resulting in a greatly reduced sequence length (and thus computation) within the SSM block. Output is scattered back to the original tensor size through binary masks (Alkayid et al., 7 May 2026).
Selection in 3D point clouds: In ForestMamba, geometry-informed and embedding-informed sampling guide which regions of large volumetric data are emphasized by the query decoder (Nguyen et al., 1 Jun 2026).

Ablation studies confirm substantial drops in accuracy when any sparsification or selection mechanism is omitted.

6. Cross-Domain Applications and Empirical Impact

Tree-Mamba architectures set state-of-the-art benchmarks across multiple domains:

Tree species mapping: GDS-Mamba achieves OA=93.94% (Alberta), OA=80.19% (Saskatchewan zero-shot), with sharply improved minority class accuracy (Alkayid et al., 7 May 2026).
Forest LiDAR segmentation: ForestMamba attains instance F1=83.4% and accelerates inference by $L$ 3 over Transformer-based models (Nguyen et al., 1 Jun 2026).
Underwater instance segmentation: UIS-Mamba surpasses WaterMask and USIS-SAM, while using only 16% of the latter's parameters (Cong et al., 1 Aug 2025).
Underwater monocular depth estimation: Tree-Mamba achieves major reductions in RMSE and S.Rel over previous methods and shows real-time speed (Zhuang et al., 10 Jul 2025).
Biomass regression: Minkowski-MambaNet improves $L$ 4 by 2.5pp and reduces RMSE by 1.41 Mg/ha over MSENet50 (Tu et al., 10 Oct 2025).

Across all areas, the models enable efficient, large-scale feature interaction, robust context propagation, and domain-informed topological consistency, with computational advantages over quadratic-complexity attention baselines.

7. Limitations and Future Directions

While Tree-Mamba architectures address key computational and topological modeling deficits of prior SSM and attention models, several areas remain for further investigation:

Generalization beyond forest and underwater imagery: Current evidence of efficacy is strong within ecological and underwater vision modalities. Extension to urban, biomedical, or planetary-scale datasets remains an open area.
Adaptation to variable data densities: For highly irregular or sparse domains (e.g., LiDAR with voids), serialization and query initialization strategies may require further refinement.
Unified framework: Tree-Mamba is a family rather than a single model; formalizing reusable interfaces for tree-aware serialization, token pruning, and graph-context integration will be important for broader adoption.

A plausible implication is that the tree-based serialization and SSM combination paradigm will exert domain-wide influence wherever efficient, global, and topology-aware context modeling is required at scale, especially in settings where quadratic transformers are infeasible.