Memorizing SAM: 3D Medical Segment Anything Model with Memorizing Transformer

TL;DR

This paper introduces Memorizing SAM, a 3D medical image segmentation model with a memory mechanism that enhances accuracy by recalling internal representations, outperforming existing models with minimal additional inference time.

Contribution

It proposes a novel memorizing Transformer plug-in for 3D SAM, improving segmentation accuracy in medical images with limited computational overhead.

Findings

Outperforms FastSAM3D with 11.36% higher Dice score

Increases inference time by only 4.38 milliseconds

Demonstrates effectiveness across 33 medical image categories

Abstract

Segment Anything Models (SAMs) have gained increasing attention in medical image analysis due to their zero-shot generalization capability in segmenting objects of unseen classes and domains when provided with appropriate user prompts. Addressing this performance gap is important to fully leverage the pre-trained weights of SAMs, particularly in the domain of volumetric medical image segmentation, where accuracy is important but well-annotated 3D medical data for fine-tuning is limited. In this work, we investigate whether introducing the memory mechanism as a plug-in, specifically the ability to memorize and recall internal representations of past inputs, can improve the performance of SAM with limited computation cost. To this end, we propose Memorizing SAM, a novel 3D SAM architecture incorporating a memory Transformer as a plug-in. Unlike conventional memorizing Transformers that save the internal representation during training or inference, our Memorizing SAM utilizes existing highly accurate internal representation as the memory source to ensure the quality of memory. We evaluate the performance of Memorizing SAM in 33 categories from the TotalSegmentator dataset, which indicates that Memorizing SAM can outperform state-of-the-art 3D SAM variant i.e., FastSAM3D with an average Dice increase of 11.36% at the cost of only 4.38 millisecond increase in inference time. The source code is publicly available at https://github.com/swedfr/memorizingSAM