SegMAN takes a distinct approach to semantic segmentation by integrating State Space Models (SSMs) with Neighborhood Attention mechanisms in one architecture. This combination targets the challenge of capturing both omni-scale long-range context and fine-grained spatial details efficiently — a balance often tricky in dense prediction tasks where accuracy and computational cost tug in opposite directions.
hybrid encoder-decoder architecture leveraging ssm and natten
At its core, SegMAN builds on MMSegmentation v0.30.0, a popular open-source framework for dense prediction, to deliver a semantic segmentation model that fuses two complementary modules in the encoder:
The State Space Model (SSM) branch implemented via Selective Scan 2D (from the Mamba lineage) aims to model long-range dependencies across the entire image with linear computational complexity relative to image size. This is significant since pure transformer-based attention mechanisms typically have quadratic complexity, which is expensive for high-resolution inputs common in segmentation.
Neighborhood Attention (Natten) captures local spatial interactions, crucial for preserving fine spatial structures and accurate boundary delineation.
This hybrid encoder is available in four scale variants: Tiny (T), Small (S), Base (B), and Large (L), spanning parameter counts from 6.4 million to 92.6 million. This range allows practitioners to trade off between accuracy and resource constraints. The decoder leverages MMSegmentation’s existing dense prediction framework, integrating multi-scale features from the encoder to produce segmentation maps.
Besides semantic segmentation, the SegMAN encoder can be used standalone as an ImageNet-1k pretrained backbone, with reported top-1 accuracy of 85.5% for the Large variant. This versatility adds practical value if you want to use the encoder features for other vision tasks.
balancing omni-scale context and spatial precision with ssm + natten
What sets SegMAN apart is the architectural bet on combining SSM for omni-scale, efficient long-range context modeling with Natten for local attention. The SSM branch, realized through Selective Scan 2D, provides linear complexity by processing state transitions over 2D spatial data, allowing it to capture dependencies at multiple scales without the quadratic cost of transformers.
Natten complements this by focusing on neighborhood-level attention windows, preserving spatial detail that SSM alone might blur or miss. This dual-branch encoder design reflects a tradeoff: while pure transformers offer powerful context modeling, their compute cost grows quadratically with input size, making them challenging to scale for dense prediction. Pure SSMs, on the other hand, may lose some fine spatial precision. SegMAN aims to get the best of both worlds.
The codebase is implemented in Python, leveraging PyTorch and building on MMSegmentation’s modular design. The SSM component is integrated via the Selective Scan 2D kernel (part of the Mamba lineage), and Natten is installed as a dependency with versioning tightly coupled to PyTorch and CUDA compatibility.
Benchmarks reported in the README reflect competitive performance on the ADE20K semantic segmentation dataset:
- SegMAN-L (Large): 53.2 mIoU with 92.6M parameters and 97.1G FLOPs
- SegMAN-T (Tiny): 43.0 mIoU with 6.4M parameters and 6.2G FLOPs
The encoder’s ImageNet-1k classification accuracy further supports the quality of learned representations:
- Large encoder achieves 85.5% top-1 accuracy
- Base and Small variants reach 85.1% and 84.0% respectively
- Tiny variant scores 76.2%
These figures suggest that the model scales well with size, allowing users to pick the best model size for their compute budget.
installation and data preparation steps
The installation process involves setting up a Python 3.10 environment with specific PyTorch and torchvision versions, followed by installing MMSegmentation v0.30.0 and its dependencies. Because the repo uses torch version >= 2.1.0, a small patch to the MMCV parallel functions is required to ensure compatibility.
Here are the exact installation commands from the README:
# Step 1: create a new environment and install torch
conda create -n segman python=3.10
conda activate segman
pip install torch==2.1.2 torchvision==0.16.2
# Step 2: install MMSegmentation and prepare datasets
pip install -U openmim
mim install mmcv-full
cd segmentation
pip install -v -e .
The patch mentioned in the README requires adding this import:
from packaging import version
Then, in the file /miniconda3/envs/segman/lib/python3.10/site-packages/mmcv/parallel/_functions.py at line 75, replace the code with:
if version.parse(torch.__version__) >= version.parse('2.1.0'):
streams = [_get_stream(torch.device("cuda", device)) for device in target_gpus]
else:
streams = [_get_stream(device) for device in target_gpus]
Next, install the Natten dependency compatible with your PyTorch and CUDA version:
pip install natten==0.17.3+torch210cu121 -f https://shi-labs.com/natten/wheels/
Then install the Selective Scan 2D kernel:
cd kernels/selective_scan && pip install .
Finally, install other required dependencies:
pip install -r requirements.txt
This setup reflects the repo’s close coupling with specific PyTorch and CUDA versions, so environment management is essential for smooth usage.
verdict: a solid hybrid architecture for semantic segmentation research and experimentation
SegMAN offers a clear architectural approach to semantic segmentation by combining SSM and Natten within a single encoder. This hybrid design tackles the tension between modeling omni-scale context efficiently and retaining local spatial precision.
The model scales across four sizes, making it adaptable from resource-constrained setups to high-end GPUs. The accompanying benchmarks on ADE20K and ImageNet-1k classification accuracy suggest the approach is competitive.
The codebase builds on MMSegmentation, a well-known modular framework, which helps with integration into existing workflows. The installation requires careful environment setup, including a patch to MMCV for torch 2.1 compatibility and specific versions of dependencies like Natten and Selective Scan 2D.
If you are researching semantic segmentation architectures or need a versatile encoder pretrained on ImageNet, SegMAN is worth exploring. The code is Pythonic and leverages standard deep learning libraries, but it’s not a plug-and-play solution for beginners — expect some environment management and familiarity with MMSegmentation.
In practice, the tradeoff between the hybrid encoder’s complexity and the performance gains is worth understanding if you are balancing accuracy and compute for dense vision tasks. SegMAN’s design points to a broader trend of mixing architectural paradigms to address the scaling challenges in vision models.
→ GitHub Repo: yunxiangfu2001/SegMAN ⭐ 260 · Python