Dual-stage 3D medical image segmentation integrating learnable prompt generation and memory attention

Abstract

Although MedSAM2 achieves 3D medical image segmentation through a memory attention mechanism, its performance declines significantly when manually designed prompts are replaced by automatically generated ones—particularly in the context of rare cases or multi-object segmentation scenarios. Current automatic prompt generation methods often extract prompt cues directly from image features, which typically lack rich spatiotemporal context and semantic information, resulting in suboptimal performance. To overcome these limitations, we propose DSSAM2-LAPG, a dual-stage 3D medical image segmentation network that integrates a Learnable Automatic Prompt-space Generator (LAPG) with memory attention. The LAPG acts as a trainable mapper that transforms raw 3D image features into a semantically-rich and spatially-aligned prompt embedding space. In the preliminary stage, this mapper utilizes learnable object tokens (concept embeddings) to dynamically interact with image features, generating coarse but reliable spatial priors and initial semantic prompts. In the refinement stage, a memory attention mechanism integrates these priors with historical context from a support memory to precisely delineate boundaries and ensure 3D consistency. This integrated approach specifically addresses the failure of existing methods in capturing 3D context, encoding rare-object semantics, and providing instance-aware guidance. Experimental results demonstrate that DSSAM2-LAPG achieves Dice score improvements of 7.2%, 6.0% and 4.1% on the private XYCH-cervical dataset, public dataset CCTH-Cervical and Multi-Organ BTCV datasets, respectively, compared to the strong baseline MedSAM2, and all without requiring any manual prompts. Our code is available https://github.com/liufangcoca-515/APG/tree/main.

Introduction

Three-dimensional (3D) medical images segmentation refers to an automated technique that precisely identifies specific anatomical structures or lesion regions through pixel-wise classification of volumetric data from modalities such as CT and MRI¹. It serves as a critical component in assisted diagnosis and treatment planning. While deep learning algorithms, notably Convolutional Neural Networks (CNNs)^2,3 and Vision Transformers (ViTs)⁴, have significantly advanced performance in this domain, they remain challenged by strong data dependency. Taking cervical cancer (a prevalent gynecological malignancy) as an example⁵, pathological data scarcity and high annotation costs–requiring 3–5 h for expert mask annotation per 3D scan–severely constrain model development⁶. Efficient automated cervical cancer MRI segmentation can quantify tumor volume and 3D spatial distribution, providing objective evidence for FIGO staging and radiotherapy target delineation, which is crucial for clinical decision-making⁷. In multi-organ segmentation scenarios, automation further enables precise spatial relationship analysis between organs, thereby optimizing therapeutic strategies.

Meta’s Segment Anything Model (SAM)⁸ innovatively established a prompt-driven segmentation paradigm, demonstrating exceptional zero-shot generalization capabilities through point, box, or text prompts. This has propelled its exploration in medical imaging: specialized variants like MedSAM have been derived through training on large-scale medical datasets, while fine-tuned adapters have been developed to accommodate diverse imaging modalities. However, these approaches still rely on manually generated high-quality prompts, which demand specialized expertise and are time-intensive. Emerging research explores automatic prompt generation, such as leveraging high-level image encoding features, mask prediction encoder guidance, or employing detection models like YOLO⁹ to generate bounding box prompts. Nevertheless, existing methods exhibit significant limitations: the segmentation performance of models like MedSAM/SAM2 is highly dependent on prompt quality and constrained by training data distribution (domain shift effects), leading to inadequate generalization on unseen modalities, multi-organ segmentation objects, and rare pathologies (e.g., specific cervical cancer subtypes). Consequently, developing learnable automatic prompting methods suitable for specific medical modalities and capable of handling multi-organ segmentation remains an urgent research problem.

SAM2¹⁰, an enhanced version of SAM, introduces a memory mechanism and memory attention to process 3D medical images as video frame sequences (assigning higher weights to adjacent frames), establishing spatiotemporal context for volumetric data segmentation. Its support for mask prompts overcomes the limitations of point/box prompts in precisely defining boundaries within medical images. However, directly substituting manually generated prompts with automatically generated ones leads to significant performance degradation on rare samples or multi-target data. Current automatic methods (e.g., those based solely on image encodings) struggle to accurately predict the coordinate positions required for bounding boxes–essential for such prompts–due to the non-coordinate nature of embeddings and random initialization. Simultaneously, bounding box prompts lack the semantic richness necessary for effective multi-class segmentation, while the binary masks generated by SAM2 inherently lack semantic class information. Current research must address two key challenges: how to fully leverage SAM2’s mask prompting capability and spatiotemporal modeling potential for efficient interaction, and how to effectively fuse automatically generated prompts–such as sparse prompts (points/boxes) and dense prompts (masks)¹¹–via improved memory attention mechanisms to optimize segmentation. Therefore, developing deep learning methods capable of automatically generating high-quality, semantically rich prompts while synergistically integrating cross-frame contextual fusion represents a critical pathway to overcome generalization bottlenecks in 3D medical images segmentation.

To address the limitations of suboptimal prompt quality and inadequate spatiotemporal context in automated segmentation, we introduce the Learnable Automatic Prompt-space Generator (LAPG). Rather than a simple generator, LAPG acts as a trainable mapper that transforms raw image features into a semantically-rich and spatially-aligned prompt embedding space using learnable object tokens. Building upon this, we propose the DSSAM2-LAPG framework, which employs a two-phase mechanism to effectively utilize these embeddings. During Preliminary Stage (Coarse Localization), the LAPG-generated embeddings are processed by a lightweight decoder to produce initial semantic prompts and coarse spatial priors. This stage establishes fundamental spatial awareness and eliminates the dependency on manual prompting. During Refinement Stage (Precise Delineation) to achieve accurate boundary delineation, a memory attention mechanism integrates the preliminary priors with historical context stored in a support memory. Through dynamic progressive fusion, the refinement decoder aggregates these spatiotemporal cues with enhanced image encodings, ensuring 3D structural continuity and robust target segmentation. The key contributions of this work are summarized as follows:

1. Dual-stage mask prediction architecture: The proposed DSSAM2-LAPG network enhances spatial semantic representation through LAPG-generated mask predictions in the preliminary stage, while strengthening inter-frame contextual correlations in the refinement stage. This significantly improves boundary precision and 3D spatial structural continuity in image segmentation.

2. Semantic prompt interaction mechanism: By designing object-specific tokens as learnable parameters within LAPG, the interaction between generated dense/sparse prompts, enhanced image encodings, and object tokens substantially boosts target recognition capability. This approach demonstrates superior performance in rare tumor cases and multi-organ segmentation objects while significantly reducing dependency on manual annotations.

3. Cross-institution performance validation: Comprehensive experiments on public and private datasets demonstrate significant improvements: For medical images segmentation, our method achieves Dice score improvements of 7.2% ,6.0% and 4.1% on the private XYCH- cervical dataset,public dataset CCTH-Cervical and Multi-Organ BTCV datasetsrespectively, confirming superior performance.

Related work

3D medical image segmentation networks

The field of medical images segmentation has evolved significantly from traditional encoder-decoder architectures, UNet^12,13 represents a milestone in this domain, with enhanced variants like nnUnet¹⁴, TransUNet¹⁵, SwinUNETR¹⁶ and nnFormer¹⁷ integrating Transformer modules to improve global context modeling. These architectures achieve remarkable precision in segmenting complex anatomical structures but require extensive annotated training data and exhibit limited generalizability across imaging modalities or anatomical regions^18,19. The Segment Anything Model (SAM)²⁰ has introduced a paradigm shift through its universal architecture, interactive prompting mechanisms (points, boxes, masks), and real-time inference capabilities, this advancement has catalyzed progress in segmentation networks, SAM’s success extended to various medical images applications²¹, exemplified by MedSAM⁸, a foundational model designed for universal medical images segmentation^22,23 that addresses cross-object generalizability challenges. For other 3D medical imaging applications, Medical SAM Adapter²⁴ employs fine-tuning strategies, while 3DSAM-adapter²⁵ proposes hierarchical 2D-to-3D feature mapping with spatial-channel attention mechanisms to align cross-dimensional features. MA-SAM²⁶ use fine-tuning strategy enables the pre-trained 2D backbone to extract thir dimensional information from input data. Current approaches for 3D medical images processing face significant limitations: 2.5D methods suffer from incomplete spatial integration, while pure 3D CNN architectures (e.g., Sam-med3d²⁷, 3d ux-net²⁸) incur substantial computational complexity. Recent research has addressed these challenges through customized adapters, 3DSAM-adapter²⁵, and low-rank adaptation (LoRA) techniques that develop specialized 3D modules for volumetric spatial information capture. SAM2²⁹ incorporates memory mechanisms to capture inter-frame relationships, implemented as both a 3D slicer plugin and Gradio API for efficient 3D segmentation³⁰. Medical variants include Medical SAM2³¹ supporting diverse 3D prompts (points, boxes, masks) with cross-scenario generalization, and SAM2Point³² which interprets 3D data as multi-directional video sequences for spatial segmentation. These developments^33,34 accelerated SAM2-based medical segmentation networks, Memorysam³⁵ let SAM2 better understand semantics and inspired by cross-frame correlation in videos. However, in 3D MRI prompt-based segmentation, current methods struggle with modeling long-range dependencies and maintaining spatial continuity across adjacent slices, lacking to capture spatiotemporal context and semantic coherence. To address these limitations, effectively integrating Segment Anything Model 2’s (SAM2) memory-enhanced mechanisms for propagating semantic and spatial information in 3D volumes remains a persistent challenge.

Segmentation based on prompt embedding

Segmentation based on prompt embedding bridges semantic-visual alignment with interactive guidance, where prompt embeddings dynamically modulate visual features for mask generation. The foundational Segment Anything Model (SAM)⁸, proposed by Meta, represents a paradigm shift through its prompt-driven mechanism, and some improvements³⁶ to prompt-based SAM2 networks for medical image segmentation. However, early methodologies predominantly relied on auxiliary segmentation networks to generate coarse spatial prompts. For instance, AdapterShadow (EfficientNet)³⁷ and UV-SAM (SegFormer)³⁸ utilized pretrained models to produce preliminary masks as input prompts for SAM. While effective, these approaches introduced parameter scales comparable to object-specific models, resulting in computational inefficiency and limited clinical applicability, SurgicalSAM³⁹, proposing a lightweight prototype-based class prompt encoder that generates embeddings directly from class prototypes, eliminating explicit prompts to improve robustness and pipeline simplicity, Auto-prompting SAM⁴⁰, which integrates contextual fully convolutional layers within the SAM framework to directly extract image features for prompt generation, significantly reducing computational overhead. Traditional manual prompting relies on human effort, is time-consuming, and has poor generalization; automated prompt segmentation represents an advanced approach^41,42, several advanced frameworks have been developed, including SAM on medical images⁴³ and automatic prompt generation⁴⁴. AutoSAM Adapter^25,40 automates prompt learning through trainable semantic vectors that dynamically generate object-specific priors for 3D segmentation, eliminating manual prompt engineering. Bbox prompts generated by object detection frameworks such as YOLO-SAM⁹, which provide sparse coordinate inputs by localizing lesion regions. Notable implementations include AutoProSAM⁴⁵, which further optimizes this paradigm for tumor MRI segmentation while balancing parameter efficiency with accuracy preservation; and AM-SAM⁴⁶, introducing automated mask calibration using lightweight architectures and low-rank adapters. Meanwhile, Beyond Adapting SAM⁴⁷ enables joint generation of dense (mask) and sparse (point/box) prompts through prototype-based encoders. Recent advances in adaptive prompt generation include RFMedSAM⁴⁸, which harnesses SAM’s intrinsic multi-stage refinement capability through cascaded decoders to progressively calibrate prompts without external dependencies; the Attention-based Prompt Generation (APG) module in Beyond Adapting SAM⁴⁷, which decouples prompt generation from SAM’s backbone to dynamically predict lesion locations; and SAM2Rad⁴⁹, incorporating a Prompt Prediction Network (PPN) that learns nonlinear mappings from image features to optimal prompt parameters, RFMedSAM 2 to get automatic prompt refinement,enabling automatic prompt generation and dual-stage refinement, Proxy Prompt⁵⁰ enhance the capacity of automated prompting and high-level humanmodel interaction. However, in prompt-based 3D MRI segmentation, current methods heavily rely on prompt quality. Automated prompt generation techniques yield suboptimal results when segmenting untrained rare or multi-object datasets, highlighting the critical need to develop effective sparse-dense prompt interaction mechanisms with imaging data.

Methodology

Working principle

The SAM2³¹ model achieves exceptional performance through a memory mechanism that captures inter-frame relationships. This architecture efficiently processes video data and extends to 3D medical images flow by treating sequential slices as video frames to facilitate memory-enhanced medical images. The core workflow comprises two phases: prompt frame processing and prompt-free frame processing. During memory encoding, predicted masks of the current frame are encoded into memory features containing maskmem features (visual representations) and maskmem positional encodings (spatial context embeddings). Concurrently, object pointers derived from projected SAM decoder tokens encapsulate high-level semantic information from historical frames. These pointers integrate into the memory sequence for subsequent attention computations. The memory attention module facilitates cross-frame interaction via cross-attention, enabling intelligent retrieval of relevant historical information. This process fuses spatiotemporal context to output augmented pixel features.

The baseline model comprises five core components: an image encoder, a prompt encoder , a mask decoder equipped with a feature adapter, a memory encoder, and a memory attention module. This model achieves mask segmentation through these components, the network structure as shown in Fig. 1a. SAM2 further employs temporal context modeling by assigning relative positional encodings to memory frames. Temporal distinction is implemented through time-step encodings, use higher position values indicating closer temporal proximity to the current frame, this prioritization of recent information enables efficient temporal connectivity through attention mechanisms that combine stacked historical memory features with current frame features.

Fig. 1.

(a) Shows the network architecture of MedSAM2, which is built upon the SAM2 framework; (b) shows an automated prompt generation structure, where manual prompts and the prompt encoder are replaced with a learnable automatic prompt generation (LAPG), referred to as APMedSAM2.

To achieve automatic segmentation, we replace manual prompts with learnable automatic prompt generation (LAPG), the network structure is shown in Fig. 1b, named APMedSAM2. LAPG employs learnable object tokens, output tokens, and bidirectional attention operations with image embeddings to automatically generate prompt embeddings, including both dense prompts and sparse prompts. Thereby completely eliminating the manual prompt inputs required by the original model.

Architectural design

Analysis of experimental results from Fig. 1a, b, which showcase two fundamental segmentation network architectures, reveals that on the private cervical cancer XYCH-cervical dataset and the public BTCV dataset, the model’s Dice Similarity Coefficient decreased by 10.1% and 7.8% respectively compared to the baseline model. This outcome primarily stems from SAM2’s dependence on prompt quality and efficient interaction, coupled with the generated prompts lacking contextual spatial and semantic information when directly fed into the mask decoder, leading to performance degradation.

To strengthen the contextual spatial and semantic information, we position the learnable automatic prompt generation (LAPG) before the memory attention component, named SAM2-LAPG, the object tokens are now conceptualized as learnable prototype vectors or concept embeddings. Each token represents the compressed “essence” of an anatomical structure, to build a segmentation network that combines LAPG and attention mechanisms memory encoding mechanism, as shown in Fig. 2a. This configuration leverages SAM2’s pre-trained memory attention to intelligently select relevant historical frames through memory retrieval, facilitating interaction between current frame features and historical memory features. The network features automated prompt generation/propagation, memory-based feature fusion, and temporal consistency enhancement. The new backbone features generated by LAPG, along with complete dense prompt embeddings, are directly fed into the memory attention module. This ensures more accurate support memory updates and significantly enhances target recognition in query images. Simultaneously, the sparse prompt embeddings produced by LAPG are input into the SAM2 mask decoder to predict the final mask through joint processing of image embeddings and prompt features. Experimental results demonstrate that SAM2-LAPG achieves Dice score improvements of only 2.2% ,2.0%and 1.1% on the private XYCH-cervical dataset, public dataset CCTH-Cervical and Multi-Organ BTCV datasets, respectively, compared to a strong baseline. However, in 3D medical images analysis, the model exhibits limitations in modeling long-range dependencies, maintaining inter-slice spatial continuity, and balancing single-frame accuracy with spatiotemporal consistency.

Fig. 2.

(a) Shows the network structure of SAM2-LAPG, in which LAPG is placed before the memory attention mechanism; (b) shows the DSSAM2-LAPG network structure, which consists of two parts: Preliminary Mask Prediction and Refinement Mask Prediction.

To address the aforementioned challenges, we propose a two-stage segmentation architecture, comprising preliminary mask prediction and refinement mask prediction, named DSSAM2-LAPG, as shown in Fig. 2b. In the initialization stage, the output of the learnable automatic prompt generation (LAPG) is fed into a lightweight decoder to generate a coarse mask . During the refinement stage, high-precision segmentation is achieved through an integrated optimization process where the dense prompt interacts with the current frame features and SAM2’s memory features via a memory attention mechanism to reinforce 3D spatial feature representation. Concurrently, point-box prompts derived from the preliminary mask are concatenated with LAPG’s object token to form joint sparse prompts, with positional encoding maintaining inter-slice continuity. Finally, a hierarchical mask decoder generates refined outputs through an adaptive decoding process where the joint prompting mechanism autonomously replaces manual annotation, significantly reducing interaction costs while improving segmentation accuracy. This process can be formulated as follows.

1

where denotes slice positional encoding , represents spatial support memory representations for refinement, and indicates adjacent slice indices. This strategy accommodates auxiliary prompts (e.g., clicks/bboxes) that are merged as sparse embeddings during refinement. Memory-augmented feature propagation strengthens temporal dependency modeling, dynamically adjusting current frame segmentation while ensuring inter-frame temporal consistency. Experiments demonstrate Dice improvements of 5.2% and 3.1% on cervical cancer XYCH-cervical dataset and BTCV datasets than SAM2-LAPG, respectively.

The dual-stage system is trained end-to-end using a unified loss function that jointly optimizes both stages, enabling simultaneous learning of coarse prompt generation and fine-grained segmentation refinement through gradient flow across the entire pipeline. The training objective integrates four complementary loss functions with carefully calibrated weights. Through grid search on the validation set, multiple weighting schemes were compared in ablation experiments. Among the various coefficient configurations of the loss function, the current weights (0.3, 1.0, 0.2, 0.1) were determined to achieve the highest overall segmentation accuracy across the evaluated datasets. The Stage-1 mask loss, denoted as , combines Dice and binary cross-entropy (BCE) to supervise the coarse mask prediction against the ground truth G, and is assigned a weight of 0.3. The Stage-2 mask loss , which serves as the primary optimization objective, applies the same Dice+BCE formulation to the refined mask prediction and carries a weight of 1.0. An auxiliary IoU prediction loss employs mean squared error (MSE) to align the predicted IoU score with the true IoU value , weighted at 0.2. Additionally, a coordinate regression loss uses Smooth L1 loss to regulate the spatial accuracy of sparse prompts generated by the Learnable Automatic Prompt Generation (LAPG) module, minimizing the discrepancy between predicted prompt coordinates and ground truth coordinates , with a weight of 0.1. The total loss is formulated as a weighted sum:

2

where is total loss,This weighting scheme prioritizes the refinement stage while maintaining balanced supervision across auxiliary tasks.

The design of learnable automatic prompt generation

The Learnable Automatic Prompt-space Generator (LAPG) is conceptualized not merely as an automatic prompt generator but as a trainable mapping function–a Learned Prompt-Space Mapper. The mapping function transforms image features into prompt embeddings . Given the learnable object token , output token , and 3D image embeddings (where denote spatial dimensions and the embedding dimension) as inputs, the architecture of the LAPG is shown in Fig. 3. The object token first engages in symmetric cross-interaction with the output token. This bidirectional interaction is implemented through nested combination operations , enhanced via residual connection and layer normalization, producing the refined object token representation . The process is formally defined as:

3

where denotes the combination operation, with softmax function and MLP consisting of two linear layers, is a learnable projection matrix, LayerNorm denotes standard layer normalization. The object token functions as a learnable parameter in the prompt generation, encoding shape, texture invariant target characteristics. In multi-target segmentation, these tokens automatically optimize through backpropagation to learn shared anatomical features. The output token provides category-specific discriminative signals, maintaining semantic consistency across frames to prevent classification discrepancies.

Fig. 3.

The network structure of the learnable automatic prompt generation (LAPG), which comprises two core components: the learnable token block and the learnable automatic prompt generation.

The refined object token and 3D image features undergo feature-object alignment through bidirectional cross-attention, generating enhanced image features and a composite token . This bidirectional mechanism achieves deep semantic-spatial coupling via complementary attention pathways. Image-to-object attention updates semantic representations, this process can be formulated as follows:

4

where denote learnable projection matrices, with as the key dimension normalization factor. And Object-to-image attention extracts spatially enhanced image features adapt the same method.

The composite token is decoupled into two components through slicing operations, one is (comprising the first n vectors ) maintains object-level semantics, the other is (the subsequent m vectors ) enables spatial coordinate mapping to the image domain.

Sparse prompts are generated by processing output tokens through an affine coordinate decoder , this process can be formulated as follows:

5

where denotes the sigmoid activation function mapping coordinates to [0, 1], is a multi-layer perceptron, and represents the MLP output dimension.

Dense prompts are produced by a mask adapter implemented as a lightweight CNN decoder, this process can be formulated as follows:

6

where Mask Adapter is a module to capture spatial information more effectively sequentially, is enhanced image features.

Mechanisms and functions from preliminary mask prediction

During the preliminary mask prediction stage, the coarse mask prediction network, is shown in Fig. 4. Comprises three core components: the learnable automatic prompt generation (LAPG) with prompt encoder, mask decoder, and memory encoder. The LAPG generates a coarse mask via a lightweight decoder, These masks serve dual purposes: On the one hand, as dense prompts for the refinement stage, enhancing tissue boundary delineation through spatial-semantic information on the other hand, as sources for generating point-box prompts via bounding box transformation; on the other hand providing precise contextual localization for interactive segmentation. During prompt encoding, sparse prompts are encoded using point positional encodings combined with learnable embedding vectors; dense prompts are processed through a 3-layer convolutional downsampling network with layer normalization and GELU activation. The lightweight decoder in Stage 1 is a compact MLP-based module with an input token count of and an output prompt token count of . The mask decoding employs a TwoWayTransformer that mediates interactions between sparse prompts, dense prompts, and image features through: Self-attention on sparse inputs, Cross-attention from sparse to dense inputs, MLP processing of sparse inputs, Cross-attention from dense to sparse inputs. Memory encoder producing memory features through fusion of current frame features with the output mask to get support_memory and object pointer.

Fig. 4.

The network structure of preliminary mask prediction, which comprises three core components: the learnable automatic prompt generation (LAPG) with prompt encoder, mask decoder, and memory encoder.

The dense prompts (i.e., the coarse mask) generated in the preliminary stage are fed into a memory encoder, which constructs a support memory repository . This repository is built by fusing resampled backbone features with LAPG-generated dense embeddings through concatenation and encoding operations:

7

where denotes resampled backbone features and represents the dense embeddings from LAPG. Temporal consistency is maintained via an incremental memory update mechanism, where the refined mask at time step is encoded as query_memory and dynamically integrated into the support memory :

8

where is a controlling the update momentum. Simultaneously, denotes the current memory features. An object pointer for cross-frame tracking is extracted by applying a multilayer perceptron to the output token corresponding to the best mask, selected based on the highest IoU score. This process can be formulated as follows:

9

where is the object pointer, and i is the index of the output token corresponding to the mask prediction with the highest IoU score. The object pointer serves as a high-level semantic summary of the target, encoding both the visual features and positional information of the object. Then propagates this information across the temporal dimension via a memory attention mechanism, enabling cross-frame object tracking and facilitating object recognition and segmentation in subsequent frames.

Refinement mask prediction via prompt-image interaction

During the refinement stage, the model progressively optimizes mask predictions through the memory attention mechanism and the prompt-image interaction within the mask decoder. The output of the memory attention module is directly fed into the mask decoder, enabling SAM2 to make segmentation decisions based on rich spatiotemporal context. The Memory Attention Module takes as input the visual features of the current frame, historical memory features (maskmem_features), and object pointers (obj_ptr) from previous frames. It performs cross-frame feature fusion via an attention mechanism and outputs an enhanced representation of the current frame, denoted as . This process is formally expressed as:

10

where denotes the current frame features, represents historical memory features, and refers to historical object pointers, and the index iterates over the historical timesteps from to , representing a sliding window of the most recent frames.

Subsequently, the enhanced feature is fed into the mask decoder for deep integration with sparse prompts. These sparse prompts consist of new object tokens generated by a learnable automatic prompt generation (LAPG) and bounding box prompts derived from coarse masks. They are concatenated and serve as queries, while the enhanced image features act as keys and values. The decoder computes interactions between queries and key-value pairs via a multi-layer transformer, ultimately producing refined mask logits and IoU scores. This decoding process is formulated as:

11

where and the bounding box prompts from coarse mask guide the decoding as queries, enabling the model to effectively capture the structure, contours, and semantic context of the target. The final output is a segmentation result with precise localization and well-delineated boundaries.

The mask decoder produces three primary outputs: low-resolution mask logits , IoU prediction scores , and output tokens . The low-resolution feature layers incorporate memory-conditioned information that carries rich historical context, thereby enhancing temporal coherence. The final mask prediction is selected by choosing the mask with the highest IoU score, This selected mask is then upsampled to high resolution via interpolation,this process can be formulated as follows:

12

where is the best mask prediction. the final mask prediction. This integrated mechanism ensures consistent and accurate 3D segmentation by leveraging both structural cues from the prompt-generated masks and contextual continuity across adjacent slices.

Experiments and results

Dataset and implementation

This study utilizes two types of datasets for 3D medical image segmentation analysis. In the field of medical imaging, compared to the vast number of examinations such as lung and head CT/MR, high-quality, expert-annotated 3D multimodal cervical cancer MRI datasets suitable for research are scarce resources in both public and private domains. One is public datasets: the CCTH (Cervical Cancer Tumor Heterogeneity) data set⁵¹, obtained from the public cervical cancer imaging archive, which provides multiparametric MRI sequences (T1-weighted, T2-weighted, and DWI); and the BCTV data set⁵², a public multi-organ segmentation data set focused on the abdomen. The other is a private dataset, named XYCH-cervical, derived from the cervical cancer MRI sequences at XiangYang Central Hospital. Within such datasets, cases with typical features are “common,” while specific pathological subtypes, extreme sizes, special locations, or cases with complex invasion constitute the “rare proportion.”It comprises 6586 expert-annotated MRI slices prospectively collected from 100 patients between June 2018 and August 2023. This dataset includes four MRI modalities: diffusion-weighted imaging (DWI), contrast-enhanced T1-weighted (T1CE), T2-weighted (T2), and apparent diffusion coefficient (ADC). All samples underwent densification processing and were stored as NumPy compressed archives (.npz) containing 3D multimodal image stacks, binary ground-truth masks, and anonymized metadata. The dataset was partitioned into training , validation , and test sets. Informed consent was obtained from all participants. The study protocol, including the informed consent process, was reviewed and approved by the Medical Ethics Committee of Xiangyang Central Hospital, Affiliated Hospital of Hubei University of Arts and Science. The rarity of cases in the XYCH-cervical dataset is quantitatively defined according to two clinical-pathological criteria: lesion size (tumors with a maximum diameter of < 2 cm or > 7 cm) and tumor location (lesions primarily involving the cervical stroma with full-thickness invasion or extending to the lower uterine segment). The dataset distribution is as follows: 100 cases with standard presentations, 8 cases with rare lesion sizes, and 5 cases with rare tumor locations.

In my methods section, we confirm that all methods were performed in accordance with the relevant guidelines and regulations. Specifically, this study was conducted in full compliance with the ethical principles of the Declaration of Helsinki and the national regulations governing human research in China. Data annotation followed rigorous clinical standards, with tumor contours manually delineated by doctor using RECIST 1.1 criteria, achieving inter-rater consistency. The implementation jointly processes MRI volumes across the sagittal (SAG), coronal (COR) and axial (AX) planes. For MRI data preprocessing, intensity normalization was applied using 0.5-99.5 percentile cropping. Connected component analysis via the cc3d library filtered small objects (3D voxel threshold: 100; 2D pixel threshold: 100). All slices were uniformly resampled to a resolution of 10241024 pixels using third-order interpolation for images and zero-order interpolation for segmentation labels. Optimization employed AdamW(, ) with a weight decay of . To assess the performance of our model, we use the Intersection over Union (IoU) and Dice Similarity Coefficient(Dice).

Results

To ensure a rigorous evaluation of model performance, all imaging data was subjected to standardized spatial preprocessing, including intensity normalization and consistent data augmentation. Using identical pre-trained initialization parameters . We compared the four aforementioned SAM2-based models MedSAM2³¹, APMedSAM2, SAM2-LAPG, and DSSAM2-LAPG on the public CCTH-cervical dataset, the private XYCH-cervical dataset, and the BTCV dataset. Qualitative analysis demonstrates that DSSAM2-LAPG significantly outperforms the other models in both cervical tumor segmentation and multi-organ segmentation. These results indicate that the improved architecture is effective for segmenting rare pathological data and 3D medical images in multi-organ segmentation tasks.All experiments are conducted on an NVIDIA A100 (80GB) GPU with FP16 precision, using the BTCV dataset . The baseline model is SAM2-base (Hiera-B backbone,  216M parameters). In our comparative analysis, MedSAM2 and our proposed DSSAM2-LAPG demonstrate the following computational profiles: MedSAM2 requires approximately 431 GFLOPs, contains around 228.6M parameters, utilizes 14.2 GB of GPU memory, and achieves a per-volume inference time of 0.8 s. In comparison, DSSAM2-LAPG operates with about 464 GFLOPs, comprises roughly 255.6M parameters, consumes 22.3 GB of GPU memory, and completes inference in1.2 seconds per volume.

The quantitative results of the three datasets across the four segmentation networks are shown in Fig. 5. DSSAM2-LAPG significantly outperforms the other models in segmentation performance, achieving the highest Dice on all three datasets: 0.88 on the XYCH-cervical dataset, and 0.87 on the CCTH-cervical dataset, and 0.93 on the BTCV dataset. These values represent improvements of 0.07, 0.06, and 0.04, respectively, compared to the corresponding results from MedSAM2.

Fig. 5.

Learnable aAuto-prompt generation network.

To further validate our method, comprehensive experiments on the private XYCH-cervical dataset, comprising multiparametric Cervical Cancer on MRI (T1c, T2, DWI and ADC), systematically evaluated different prompting strategies with the quantified corresponding Dice and IoU metrics, including various 3D segmentation networks such as SAM-Med3D, AM-SAM, and SAM2-based networks integrating point prompts, bounding box prompts, automatic prompts, and other mechanisms. Benchmarking against state-of-the-art methods revealed DSSAM2-LAPG’s superior performance, where the auto-prompt mechanism eliminated manual annotation requirements while achieving significant Dice/IoU improvements of +0.07/+0.08 versus optimal manual prompting, as shown in Table 1.

Table 1.

Quantitative comparison of different prompting strategies across MRI from the private dataset (XYCH-cervical). The best Dice scores are highlighted in bold.

Method		T1c		T2		DWI
Model	Prompt	IoU	Dice	IoU	Dice	IoU	Dic
SAMMed3D27	Point	0.61	0.64	0.65	0.71	0.71	0.72
SAMMed3D	Box	0.65	0.68	0.68	0.72	0.79	0.76
MA-SAM26	Point	0.64	0.61	0.66	0.64	0.68	0.76
MA-SAM	Box	0.62	0.71	0.63	0.74	0.71	0.78
MedSAM231	Point	0.72	0.77	0.69	0.82	0.78	0.80
MedSAM2	Box	0.72	0.79	0.69	0.83	0.80	0.84
APMedSAM2	LAPG	0.58	0.65	0.62	0.68	0.61	0.72
SAM2-LAPG	LAPG	0.69	0.79	0.68	0.75	0.75	0.86
SAM2-LAPG	LAPG+box	0.70	0.80	0.69	0.77	0.76	0.88
DSSAM2-LAPG	LAPG	0.75	0.83	0.77	0.84	0.82	0.89
Ours(DSSAM2-LAPG)	LAPG+box	0.79	0.85	0.81	0.86	0.84	0.93

Method		ADC		Average
Model	Prompt	IoU	Dice	IoU	Dice
SAMMed3D27	Point	0.71	0.67	0.67	0.69
SAMMed3D	Box	0.69	0.73	0.70	0.72
MA-SAM26	Point	0.65	0.75	0.66	0.69
MA-SAM	Box	0.71	0.72	0.67	0.74
MedSAM231	Point	0.73	0.79	0.73	0.80
MedSAM2	Box	0.74	0.79	0.74	0.81
APMedSAM2	LAPG	0.69	0.76	0.63	0.70
SAM2-LAPG	LAPG	0.72	0.83	0.71	0.81
SAM2-LAPG	LAPG+box	0.73	0.83	0.72	0.83
DSSAM2-LAPG	LAPG	0.77	0.88	0.78	0.86
Ours(DSSAM2-LAPG)	LAPG+box	0.82	0.87	0.82	0.88

To further validate the method’s efficacy, supplementary qualitative analysis was conducted on the public CCTH-cervical dataset. Employing identical experimental configurations, DSSAM2-LAPG demonstrated superior Cervical Cancer segmentation performance, compared to the best MedSAM2 intervention, achieving Dice/IoU improvements of +0.06/+0.19 respectively versus optimal manual prompting, as shown in Table 2.

Table 2.

Quantitative comparison of different prompting strategies across MRI from the public CCTH-cervical dataset. The best Dice scores are highlighted in bold.

Method		SAG-T2		AX-T2		COR-T2
Model	Prompt	IoU	Dice	IoU	Dice	IoU	Dic
SAMMed3D27	Point	0.58	0.63	0.61	0.62	0.65	0.69
SAMMed3D	Box	0.62	0.64	0.67	0.72	0.72	0.76
MA-SAM26	Point	0.61	0.62	0.61	0.63	0.72	0.74
MA-SAM	Box	0.62	0.68	0.62	0.64	0.69	0.75
MedSAM231	Point	0.63	0.65	0.69	0.72	0.73	0.78
MedSAM2	Box	0.61	0.76	0.62	0.79	0.65	0.86
APMedSAM2	LAPG	0.58	0.53	0.59	0.62	0.55	0.68
SAM2-LAPG	LAPG	0.69	0.77	0.68	0.73	0.72	0.82
SAM2-LAPG	LAPG+box	0.71	0.83	0.71	0.81	0.77	0.85
DSSAM2-LAPG	LAPG	0.75	0.82	0.77	0.83	0.82	0.88
Ours(DSSAM2-LAPG)	LAPG+box	0.79	0.84	0.81	0.85	0.84	0.92

Method		SAG-DWI		Average
Model	Prompt	IoU	Dice	IoU	Dice
SAMMed3D27	Point	0.61	0.63	0.61	0.64
SAMMed3D	Box	0.71	0.73	0.68	0.71
MA-SAM26	Point	0.62	0.68	0.64	0.67
MA-SAM	Box	0.72	0.76	0.66	0.71
MedSAM231	Point	0.69	0.75	0.69	0.73
MedSAM2	Box	0.62	0.82	0.63	0.81
APMedSAM2	LAPG	0.52	0.62	0.56	0.61
SAM2-LAPG	LAPG	0.69	0.81	0.70	0.78
SAM2-LAPG	LAPG+box	0.72	0.83	0.73	0.83
DSSAM2-LAPG	LAPG	0.77	0.87	0.78	0.85
Ours(DSSAM2-LAPG)	LAPG+box	0.82	0.86	0.82	0.87

This enhancement primarily originates from the support memory -driven feature representation enrichment and synergistic dense prompt guidance, collectively optimizing predictive performance. Notably, DWI exhibited optimal segmentation accuracy attributable to its inherent high tissue contrast enhancing boundary delineation.Our DSSAM2-LPGA prompting protocol employs a standardized approach across all datasets. For point prompts, we use iterative clicking from 3 to 5 clicks per object, with points sampled from ground truth masks using center and boundary. Box prompts consist of single bounding boxes derived from ground truth minimum enclosing rectangles. “Optimal manual prompting” is defined as the prompt configuration that maximizes IoU with minimal clicks, specifically selecting prompts that yield the highest segmentation accuracy while mimicking realistic human interaction patterns. All prompts are normalized to relative coordinates and processed through a unified encoder to ensure cross-dataset consistency.

In the comparative experimental data analysis, for the MedSAM2 model, point prompts are less effective than bounding box prompts in terms of segmentation accuracy, because of bounding boxes provide sufficient information to identify the target and enhance semantic expression. To comprehensively address the effectiveness of prompts, we propose DSSAM2-LAPG, in the Preliminary prediction stage, it generates a mask and its corresponding bounding box as prompts for the refinement stage, achieving automatic prompting and effectively combining sparse prompts (points and bounding boxes) with dense prompts (masks), thereby enhancing adaptability and robustness.

For multi-organ segmentation, we conducted validation on the public BTCV dataset, which contains 14 abdominal organs. Semantic labels are assigned using object IDs that implicitly encode organ categories, where each organ is assigned a unique identifier that serves both as an instance marker and semantic category indicator. The model does not employ explicit class tokens or classifier heads for multi-class prediction but instead treats each organ as an independent object instance with its own distinct representation throughout the processing pipeline. The prediction method operates through per-object processing using separate object tokens for each organ. For every object ID, the model receives independent interactive prompts such as clicks or bounding boxes, then generates individual binary masks while maintaining organ-specific representations across video sequences. This approach utilizes memory-augmented tracking mechanisms including object pointer projections, memory encoders, and attention modules to preserve temporal consistency for each organ throughout the video frames, distinguishing it from simultaneous multi-class prediction systems. Experiments were performed using four designed segmentation networks for comparison. As shown in Table 3, the results demonstrate that DSSAM2-LAPG outperforms MedSAM2 by +0.04 in Dice, respectively. This improvement is attributed to the learnable automatic prompt generation, which produces target prompts and incorporates a memory attention mechanism to enhance the recognition and tracking of multiple targets across frames.For consistency with the BTCV experiments in Medical SAM 2³¹, we follow its memory attention parameter settings where k is 7 and is 1.

Table 3.

Quantitative comparison of different prompting strategies for MRI from the BTCV dataset. The table presents Dice scores for twelve organs across various segmentation networks. The best Dice scores are highlighted in bold.

Method	Spleen	R.Kid	L.Kid	Gall	Eso	Liver	Stomach
SAMMed3D27	0.87	0.88	0.93	0.79	0.79	0.94	0.89
MA-SAM26	0.97	0.95	0.95	0.68	0.82	0.97	0.93
MedSAM231	0.92	0.93	0.92	0.93	0.92	0.86	0.87
APMedSAM2	0.86	0.87	0.91	0.85	0.83	0.85	0.89
SAM2-LAPG	0.93	0.94	0.93	0.94	0.93	0.87	0.88
Ours	0.95	0.95	0.91	0.97	0.94	0.94	0.94

Method	Aorta	IVC	Veins	Pancreas	AG	Average
SAMMed3D27	0.87	0.80	0.81	0.78	0.79	0.85
MA-SAM26	0.91	0.87	0.80	0.87	0.74	0.87
MedSAM231	0.89	0.83	0.85	0.84	0.91	0.89
APMedSAM2	0.83	0.81	0.85	0.82	0.87	0.85
SAM2-LAPG	0.90	0.84	0.86	0.85	0.92	0.90
Ours	0.92	0.90	0.90	0.93	0.94	0.93

Visual comparison results and analysis

In addition to quantitative analysis, we conducted visual comparisons to further validate the performance of the proposed method. After obtaining institutional ethical approval, all images were processed using a unified protocolpatient images were randomly selected from a private dataset (XYCH-cervical), including four MRI modalities: T1c, T2, DWI and ADC. Each MRI scan displayed tumors with varying sizes, locations, and morphological characteristics. We performed a visual comparison of the segmentation results generated by SAM-Med3D²⁷, MA-SAM²⁶, MedSAM2³¹, and our improved methods—including APMedSAM2, SAM2-LAPG, and DSSAM-LAPG. Visual comparisons as show in Fig. 6. Visual analysis indicates that the MedSAM-LAPG segmentation network exhibits significant advantages in detail preservation and tumor boundary delineation. For segmenting cervical cancer tumors significantly improved segmentation accuracy. The left section demonstrates effectively addresses the challenge of low grayscale contrast between the tumor and adjacent tissues in segmenting tumors of varying shapes and sizes compared to other prompt-based strategies; the middle section shows that the model more accurately segments small-volume structures within discontinuous anatomical regions; the right section further confirms that MedSAM-LAPG is more effective and precise in identifying and outlining tumor boundaries within complex anatomical contexts. This performance improvement is attributed to the effective integration of LAPG and the memory attention mechanism: when encountering rare samples not seen during training, through memory attention for object tracking, the model adaptively captures tumor features across varying sizes, shapes, and spatial contexts.

Fig. 6.

Visual comparison of segmentation results from three sets of MRI DWI axial (AX) view images in a private dataset (XYCH-cervical). From top to bottom: segmentation outputs from various networks based on different prompt strategies. Green: ground truth (GT); red: predicted segmentation mask; yellow: overlap. From left to right: segmentation results for three different images.

This model exhibits strong generalization capability, performing well on cross-institutional data collected from different hospitals (e.g., the public dataset CCTH-cervical and the private dataset XYCH-cervical). For instance, Dice scores reach 0.87 on CCTH-cervical and 0.88 on XYCH-cervical, respectively. Similarly, several cervical cancer cases were randomly selected from a public dataset (CCTH-cervical) and analyzed using T2c MRI sequences. Segmentation maps generated by different segmentation networks and prompt strategies were compared, with emphasis on tumor morphology and spatial accuracy. The results show that the proposed DSSAM2-LAPG method for cervical cancer tumor segmentation outperforms other methods in terms of anatomical fidelity and boundary accuracy, producing results closer to the ground truth annotations. Visual comparisons as show in Fig. 7, further validating the performance enhancement of the DSSAM2-LAPG model. The LAPG mechanism does not simply replace manual prompts with automatic ones; rather, it enhances contextual semantic understanding of key regions in areas with structural discontinuities or low contrast, guiding the model to focus on subtle tumor boundaries or small lesions, thereby improving segmentation accuracy.

Fig. 7.

Visual comparison of segmentation results from three sets of MRI T2c coronal (COR) view images in a public dataset (CCTH-cervical). From top to bottom: segmentation outputs from various networks using different prompting methods. Green: GT; red: predicted mask; yellow: overlap. From left to right: segmentation results for three different images.

To demonstrate the effectiveness of the cross-frame association mechanism in multi-object segmentation, we conducted a visual analysis on the public dataset BTCV. Visual comparisons as show in Fig. 8, A single frame from a randomly selected MRI sequence was used, and three different organs were annotated within the same image to compare segmentation performance. The results indicate that DSSAM2-LAPG outperforms other methods in multi-organ recognition and segmentation accuracy. This advantage stems from the two-stage mask optimization mechanism in DSSAM2-LAPG: during the preliminary mask prediction stage, each organ target has independent object tokens and object pointers. The object pointer tokens participate in cross-frame attention computation, and the memory attention mechanism enables the model to effectively utilize historical information of relevant targets, thereby achieving robust multi-object tracking.

Fig. 8.

Visual comparison of segmentation results for three different organs from the same image in the public dataset BCTV. Columns 1–4: original image, predicted mask, ground truth (GT) mask, and overlap result. In the overlap result (Column 4), green represents GT, red indicates predicted masks, and yellow shows overlapping regions. The same representation applies to subsequent columns.

Ablation experiments

We conducted ablation studies to evaluate the contribution of the LAPG module in our proposed method, which integrates the output of LAPG into a memory attention mechanism to enhance overall segmentation performance. First, we compared the average Dice scores obtained using learnable prompts (with LAPG) versus manual prompts (without LAPG) in the SAM2 and MedSAM2 architectures to validate the effectiveness of the LAPG module. Furthermore, we replaced the LAPG module in the DSSAM2-LAPG model with manual prompts and compared the average Dice scores; the results are shown in Table 4. The key advantage of LAPG lies in its ability to map the visual features of the input image to a representation space aligned with SAM’s prompt embedding space, thereby learning a set of general and transferable capabilities for “visual concept decomposition and recombination.” When processing new or rare images, the image encoder extracts their visual features (such as shape, texture, edges, and contrast with the background). The prompt embedding generation then converts these features into effective object tokens and dense embeddings based on the learned general mapping rules. The generated prompts effectively provide the instruction “there is a significant foreground object here,” which the SAM decoder recognizes and uses to produce accurate segmentation masks.

Table 4.

Ablation results comparing models with LAPG (learnable prompts) and without LAPG (manual prompts) on three datasets: the Public datase CCTH-cervical, Private dataset Cervical-XYCH datase and BTCV dataset.

Model	Dice Score (with LAPG/without LAPG)
	CCTH-cervical	XYHC-cervical	BTCV
SAM2	0.76/0.71	0.75/0.71	0.86/0.80
MedSAM2	0.81/0.77	0.78/0.76	0.89/0.87
DSSAM2-LAPG	0.88/0.83	0.87/0.82	0.93/0.88

The results consistently demonstrate that integrating LAPG improves segmentation performance across all models and datasets. Specifically, DSSAM2-LAPG achieves the highest Dice scores, with improvements of 5–6 over its manual-prompt counterpart, confirming the critical role of LAPG in enabling robust and automatic prompt generation for 3D medical image segmentation.

Additionally, the ablation study assessed the contribution of the first-stage preliminary mask prediction module. We compared the performance of the following three configurations: using only the dense and sparse prompts generated in the second stage (without the first stage referred to as NF), using only the mask generated in the first stage (referred to as OM), and integrating the bounding box (bbox) obtained from the first stage into the sparse prompts( referred to as BM). The average Dice scores were compared, and the average Dice scores across three different datasets are shown in Table 5. The experimental results indicate that relying solely on second-stage prompts leads to poorer segmentation performance, demonstrating that MedSAM2 heavily depends on high-quality prompts, while the first stage provides more comprehensive prompt information. Incorporating bounding boxes further improves prompt quality by enhancing the representation of spatial location information.

Table 5.

The average result of correlated ablation experiments under using only the mask generated in the first stage (referred to as OM), and integrating the bounding box (referred to as BM) on the The Public datase CCTH-cervical, Private dataset Cervical-XYCH datase and BTCV dataset.

Dice(%)	w/o NF	w/o OM	w/o BM(Ours)
XYCH-cervical	0.82	0.86	0.88
CCTH-cervical	0.81	0.85	0.87
BTCV	0.89	0.91	0.93

Discussion and conclusion

To address the significant performance degradation of MedSAM2 under automatically generated prompts, particularly in rare lesion and multi-object segmentation scenarios, this paper proposes a dual-stage 3D Medical Image Segmentation network named DSSAM2-LAPG, which integrates learnable automatic prompt generation (LAPG) with memory attention. The method includes two core designs: first, the LAPG module generates spatially and semantically aligned coarse masks to replace manual prompts; second, the dual-stage architecture effectively combines LAPG with a memory attention mechanism, enriching spatiotemporal context and semantic information in 3D medical images and enhancing cross-frame segmentation consistency. This design significantly improves boundary accuracy and 3D structural continuity, increases the Dice coefficient on both private and public segmentation datasets, substantially reduces reliance on manual annotations, and demonstrates strong generalization capability across institutions. However, the additional parameters introduced by automatic prompt generation present a trade-off; therefore, lightweight implementation remains a promising direction for future research.

Author contributions

F.L. and J.Y. wrote the main manuscript text J.Y. and T.L. prepared figures 1–8; F.W. and L.D. conducted experiments and analyzed the experimental data; Y.D.Z. and H.X. provided guidance on the research methodology. Authors reviewed the manuscript.

Funding

YanDuo Zhang received funding from the National Natural Science Foundation of China under Grant ID 62072350; Hui Xing received funding from the National Natural Science Foundation of China under Grant ID 82472634; and the Key Research and Development Program of Hubei Province, Grant ID 2024BCB055.

Data availability

The experiments were conducted on two types of datasets for 3D medical image segmentation analysis. One is public datasets: the CCTH- cervical dataset, obtained from the public cervical cancer imaging archive, which provides multiparametric MRI sequences (T1-weighted, T2-weighted,and DWI), the The URL is https://www.cancerimagingarchive.net/collection/cc-tumor- heterogeneity; and the BCTV dataset, a public multi-organ segmentation dataset focused on the abdomen, the The URL is https://www.synapse.org/Synapse:syn3193805/files/. The other is a private dataset, named XYCH-cervical, derived from the cervical cancer MRI sequences at XiangYang Central Hospital. It comprises 8,586 expert-annotated MRI slices prospectively collected from 150 patients between June 2018 and August 2023. This study was approved by the Medical Ethics Committee of Xiangyang Central Hospital, Affiliated Hospital of Hubei University of Arts and Science. Informed consent was obtained from all participants. Approved by the institutional review board, This dataset includes four MRI modalities: diffusion-weighted imaging (DWI), contrast-enhanced T1-weighted (T1CE), T2-weighted (T2), and apparent diffusion coefficient (ADC).

Competing interests

The authors declare no competing interests.