Topo-VM-UNetV2: Encoding Topology into Vision Mamba UNet for Polyp Segmentation

TL;DR

The paper tackles polyp segmentation by addressing the deficiency of capturing topological structure in existing CNN/Transformer/SSM-based models. It proposes a two-stage method, where Stage 1 derives topology attention maps from VM-UNetV2 probability maps using persistence diagrams and a sigmoid transformation of persistence, yielding $T$ maps aligned with image space. Stage 2 injects these topology cues into a topology-guided SDI module (Topo-SDI) inside VM-UNetV2 to improve boundary fidelity and reduce mis-segmentation. This approach achieves state-of-the-art results across five public datasets, with notable gains on challenging cases (e.g., ETIS) and practical computational costs, and the authors plan to release the code for broader use.

Abstract

Convolutional neural network (CNN) and Transformer-based architectures are two dominant deep learning models for polyp segmentation. However, CNNs have limited capability for modeling long-range dependencies, while Transformers incur quadratic computational complexity. Recently, State Space Models such as Mamba have been recognized as a promising approach for polyp segmentation because they not only model long-range interactions effectively but also maintain linear computational complexity. However, Mamba-based architectures still struggle to capture topological features (e.g., connected components, loops, voids), leading to inaccurate boundary delineation and polyp segmentation. To address these limitations, we propose a new approach called Topo-VM-UNetV2, which encodes topological features into the Mamba-based state-of-the-art polyp segmentation model, VM-UNetV2. Our method consists of two stages: Stage 1: VM-UNetV2 is used to generate probability maps (PMs) for the training and test images, which are then used to compute topology attention maps. Specifically, we first compute persistence diagrams of the PMs, then we generate persistence score maps by assigning persistence values (i.e., the difference between death and birth times) of each topological feature to its birth location, finally we transform persistence scores into attention weights using the sigmoid function. Stage 2: These topology attention maps are integrated into the semantics and detail infusion (SDI) module of VM-UNetV2 to form a topology-guided semantics and detail infusion (Topo-SDI) module for enhancing the segmentation results. Extensive experiments on five public polyp segmentation datasets demonstrate the effectiveness of our proposed method. The code will be made publicly available.

Figures (5)

• Figure 1: Examples of polyp segmentations generated by VM-UNetV2 zhang2024vm, showing that topological features (e.g., connected components) are not well captured.
• Figure 2: The pipeline of our proposed framework. (a) Stage 1: VM-UNetV2 zhang2024vm is used to generate probability maps for the training and test images, which are then used to produce corresponding topology attention maps. (b) Stage 2: The training images and their topology attention maps are used to train our Topo-VM-UNetV2 model, which consists of an Encoder, the Topo-SDI (topology-guided semantics and detail infusion) module, and a Decoder. (c) The structure of the Topo-SDI module. For simplicity, only the refinement of the third scale features ($l=3$) is shown.
• Figure 3: Illustrating the sub-level filtration process. (a) An example image and its probability map generated by VM-UNetV2. (b) Four thresholded binary images of the probability map. As the threshold value increases or decreases, some connected components or loops are born or die. (c) The computed persistence diagram of the probability map, where red points denote 0-D persistent homology ($\text{PD}_0$) and blue points denote 1-D persistent homology ($\text{PD}_1$).
• Figure 4: Examples of segmentations from five polyp segmentation datasets.
• Figure 5: Examples of different attention maps computed from probability maps.