Nd-BiMamba2：面向多维数据处理的统一双向架构
机器学习
2024-11-26 v1 人工智能
摘要
深度学习模型通常需要为处理不同维度的数据（如 1D 时间序列、2D 图像和 3D 体数据）设计专门的架构。现有的双向模型主要关注序列数据，难以有效扩展到更高维度。在解决此问题方面，我们提出了一种新型的多维双向神经网络架构，命名为 Nd-BiMamba2，有效处理 1D、2D 和 3D 数据。Nd-BiMamba2 基于 Mamba2 模块，引入创新的双向处理机制和自适应填充策略，以在保持计算效率的同时捕获多维数据的双向信息。不同于需要为不同维度数据设计特定架构的方法，Nd-BiMamba2 采用统一的架构，具有模块化设计，简化了开发和维护成本。为验证 Nd-BiMamba2 的可移植性和灵活性，我们成功地将其导出到 ONNX 和 TorchScript，并在不同的硬件平台（如 CPU、GPU 和移动设备）上进行测试。实验结果表明，Nd-BiMamba2 在多个平台上运行效率较高，展示了其在实际应用中的潜力。该代码已开源：https://github.com/Human9000/nd-Mamba2-torch
关键词
引用
@article{arxiv.2411.15380,
title = {Nd-BiMamba2: A Unified Bidirectional Architecture for Multi-Dimensional Data Processing},
author = {Hao Liu},
journal= {arXiv preprint arXiv:2411.15380},
year = {2024}
}