MAIS: Memory-Attention for Interactive Segmentation

TL;DR

The paper tackles the inefficiency of current ViT-based interactive segmentation methods in 3D medical images, where interactions are treated independently, leading to diminishing gains. It proposes MAIS, a memory-attention mechanism that stores past prompts and segmentation masks in a FIFO memory bank and conditions a 3D SAM-Med3D backbone on this history to enable coherent, incremental refinements. The approach uses a lightweight memory-attention block and only fine-tunes the prompt encoder and mask decoder, achieving strong improvements across CT and MR datasets in low-data regimes and sometimes matching or surpassing task-specific baselines like nn-UNet. The findings demonstrate data-efficient adaptation with near real-time inference and suggest broader applicability to foundation-model-assisted segmentation in clinical workflows.

Abstract

Interactive medical segmentation reduces annotation effort by refining predictions through user feedback. Vision Transformer (ViT)-based models, such as the Segment Anything Model (SAM), achieve state-of-the-art performance using user clicks and prior masks as prompts. However, existing methods treat interactions as independent events, leading to redundant corrections and limited refinement gains. We address this by introducing MAIS, a Memory-Attention mechanism for Interactive Segmentation that stores past user inputs and segmentation states, enabling temporal context integration. Our approach enhances ViT-based segmentation across diverse imaging modalities, achieving more efficient and accurate refinements.

Figures (5)

• Figure 1: Segmentation accuracy (Dice $\%$) vs. user interactions on Han-SegHan-seg. SAM-Med3D (3D medical-specific) shows poor zero-shot performance and plateaus after few clicks (fine-tuned), revealing limited task specificity. Our method (MAIS:Ft ), leveraging memory of past interactions, sustains improvement with increasing clicks, approaching Oracle performance even with 10$\%$ training data.
• Figure 2: Overview of the proposed segmentation model incorporating memory attention: A 3D image encoder processes the input image generating a 3D embedding. Prompt encoders transform user interactions—positive or negative 3D coordinates—and previous mask predictions into prompt embeddings, which are stored in a memory bank for future interactions. The memory attention mechanism conditions the image embedding on the stored memory before passing it to the mask decoder, which produces the final segmentation output.
• Figure 3: Memory bank and prompt type analysis : Subplot (a) Impact of memory bank size on segmentation performance. Dice score ($\%$) is shown as a function of the number of clicks for different memory sizes. Subplots (b) Sparse vs dense memory embeddings comparison for the different number of clicks and images used for training.
• Figure 4: Qualitative result on AMOS-MR dataset For Liver with varying numbers of user clicks. The first column shows the original MRI images, and the second column presents the ground-truth segmentations. The remaining columns compare segmentations produced by SAM3D, fine-tuned SAM3D (FT-SAM3D), and MAIS.
• Figure 5: Qualitative result on the HaN-SEG dataset (CT images) for Esophagus segmentation with different numbers of user clicks. The first column shows the original CT images, while the second column presents the ground-truth segmentations. The remaining columns compare segmentations produced by SAM3D, fine-tuned SAM3D (FT-SAM3D), and MAIS. It can be observed that SAM3D tends to oversegment the esophagus, while FT-SAM3D undersegments it. MAIS provides more balanced results as the number of clicks increases.