Nd-BiMamba2: A Unified Bidirectional Architecture for Multi-Dimensional Data Processing
This addresses the need for simplified development and maintenance in multi-dimensional data processing, though it appears incremental as it builds on the Mamba2 module.
The paper tackles the problem of deep learning models requiring specialized architectures for different data dimensions by proposing Nd-BiMamba2, a unified bidirectional architecture that efficiently handles 1D, 2D, and 3D data, with experimental results showing it runs efficiently on multiple hardware platforms like CPU, GPU, and mobile devices.
Deep learning models often require specially designed architectures to process data of different dimensions, such as 1D time series, 2D images, and 3D volumetric data. Existing bidirectional models mainly focus on sequential data, making it difficult to scale effectively to higher dimensions. To address this issue, we propose a novel multi-dimensional bidirectional neural network architecture, named Nd-BiMamba2, which efficiently handles 1D, 2D, and 3D data. Nd-BiMamba2 is based on the Mamba2 module and introduces innovative bidirectional processing mechanisms and adaptive padding strategies to capture bidirectional information in multi-dimensional data while maintaining computational efficiency. Unlike existing methods that require designing specific architectures for different dimensional data, Nd-BiMamba2 adopts a unified architecture with a modular design, simplifying development and maintenance costs. To verify the portability and flexibility of Nd-BiMamba2, we successfully exported it to ONNX and TorchScript and tested it on different hardware platforms (e.g., CPU, GPU, and mobile devices). Experimental results show that Nd-BiMamba2 runs efficiently on multiple platforms, demonstrating its potential in practical applications. The code is open-source: https://github.com/Human9000/nd-Mamba2-torch