@inproceedings{
xiao2024mambatree,
title={MambaTree: Tree Topology is All You Need in State Space Model},
author={Yicheng Xiao and Lin Song and Shaoli Huang and Jiangshan Wang and Siyu Song and Yixiao Ge and Xiu Li and Ying Shan},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=W8rFsaKr4m}
}
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 Lin Song Shaoli Huang Jiangshan Wang Siyu Song Yixiao Ge Xiu Li Ying Shan