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

Mathias Unberath; Xinyuan Shao; Yiqing Shen

REVIEW 5 major objections 6 minor 19 references

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

T0 review · 5 major / 6 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Adding a retrievable memory of high-quality training representations to a 3D SAM raises its average Dice by 11.36% at a cost of 4.38 ms per image.

desk verdict A useful plug-in idea whose headline gain is still confounded by per-class specialist training versus the single-model baseline. read the letter →

arxiv 2412.13908 v1 pith:5HE3WHE4 submitted 2024-12-18 cs.CV

Xinyuan Shao , Yiqing Shen , Mathias Unberath This is my paper · ORCID

classification cs.CV

keywords medicalimagesegmentationSegmentAnythingModelMemorizingTransformer3DvolumetrickNNattentionTotalSegmentatorfoundationmodels

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper asks whether a memorizing Transformer—a plug-in that retrieves stored key-value representations of past inputs—can close the performance gap of 3D Segment Anything Models on volumetric medical segmentation. It introduces Memorizing SAM, which precomputes a memory bank from class-specific fine-tuned models on 210 training and tuning volumes, saves it to disk, and retrieves the top-k memories via kNN attention during inference. Evaluated on 33 TotalSegmentator classes and 500 test volumes, the method beats the FastSAM3D baseline by an average of 11.36% Dice while adding only 4.38 ms per image and no new parameters. The authors argue the benefit is largest when the SAM has not been task-specifically fine-tuned, which matters because well-annotated 3D medical data for fine-tuning is scarce.

What carries the argument

The central object is the Memorizing Transformer block, a plug-in that retains the standard dense self-attention and feed-forward layers and adds a kNN attention branch. For a query, it retrieves the top-$k$ (key, value) pairs from a precomputed external memory and computes attention ratios $R_i$ from the kNN distances $D_i$, combining them with the local attention output $A_L$ as $A_c = A_L \cdot R_L + \sum_{i=0}^k R_i \cdot A_i$, with the local ratio $R_L$ set to 0.3. The memory itself is generated offline: the dataset is split into $N$ one-class subsets, a class-specific SAM is fine-tuned on each, and the resulting internal representations are saved as (key, value, image) tuples to disk, loaded into cache only at retrieval time. This design keeps FLOPs at 23.14 G and parameters at 45.64 M, identical to the FastSAM3D baseline, while adding 4.38 ms per image for the retrieval and attention computation.

What would settle it

Run Memorizing SAM on a CT dataset from a different scanner vendor or patient population whose anatomical variants are absent from the 210 memorized volumes; if the average Dice against FastSAM3D falls to zero or negative, the memory-representativeness premise is false. A second check: replace the stored memory keys with random keys from the same volumes; if Dice does not drop substantially, the improvement comes from the kNN operation itself rather than the memory content.

Watch

Extended reading notes

Core claim

The central claim is that a Memorizing Transformer block, used as a plug-in in 3D SAM, improves segmentation accuracy by retrieving and attending to precomputed internal representations of high-quality training examples. The memory is built by training one-class SAM models on a small dataset (200 fine-tuning samples), generating key-value pairs for each class, and storing them in external memory alongside the images. During inference, the block performs an approximate kNN search over these keys and combines the retrieved attention outputs with local attention using a ratio-based weighting scheme. Measured on 33 classes of TotalSegmentator, Memorizing SAM outperforms FastSAM3D by 11.36% average Dice, with inference time rising from 2.51 ms to 6.89 ms per image while FLOPs, parameters, and GPU memory stay essentially unchanged. The authors interpret the gains as evidence that memory retrieval can substitute for part of the fine-tuning signal, especially for un-fine-tuned models.

Load-bearing premise

The precomputed memory bank, built from 210 training and tuning volumes, must contain key-value pairs representative enough that kNN retrieval helps on 500 unseen test volumes across 33 anatomical classes; if the stored keys fail to cover the anatomical variation in the test set, retrieval returns irrelevant values and the reported 11.36% Dice gain disappears.

Editorial extensions

If this is right

Memorizing SAM consistently beats FastSAM3D across most of the 33 anatomical classes, with the largest gains on challenging or underrepresented structures.
The memory mechanism is most valuable for models that have not been task-specifically fine-tuned, suggesting it can serve as a lightweight alternative to fine-tuning when annotated volumes are scarce.
The performance gain comes with negligible overhead: inference time increases by 4.38 ms per image, while FLOPs, memory usage, and parameter count stay essentially flat.
Varying the number of retrieved memories k has only a modest effect, with k=3 giving the best trade-off; later-ranked, less similar memories contribute little to the output.
Because the block is a plug-in and memory is precomputed, the approach can be attached to any SAM variant, not only FastSAM3D.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

If the memory bank is representative, the same retrieval mechanism could be extended to other foundation models in medical imaging, for instance to adapt a model to a new hospital's scan protocol by storing a small set of labeled cases from that site.
A direct ablation would be to corrupt the memory content, e.g., by replacing stored keys with random or wrong-class representations; if Dice does not drop substantially, the gain comes from the kNN operation itself rather than the memorized content.
The results hint at a 'memorize instead of fine-tune' trade-off, where retrieving expert-annotated slices approximates part of the benefit of gradient updates—potentially relevant in federated or privacy-constrained settings where raw training data cannot be shared but representations might be.
Because the reported timing is on an NVIDIA T4 GPU, the 4.38 ms overhead may scale differently on edge or CPU hardware; measuring retrieval latency on deployment hardware would clarify clinical feasibility.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Short version: this is a reasonable engineering paper with a confounded headline number. The idea—bolt a kNN-attention memory block onto a 3D SAM and fill it with precomputed key-value pairs from high-quality internal representations—is not a new mechanism (it follows Memorizing Transformers and MoViT), but packaging it as a plug-in for FastSAM3D and testing on 33 TotalSegmentator classes is a useful thing to have. The design choice of preparing memory before inference and loading it on demand is practical, and the 4.38 ms/image overhead is attractive if it holds up.

What the paper does not yet establish is where the 11.36% average Dice gain comes from. Section 2 says they train N separate Memorizing SAM models, one per class. The baseline FastSAM3D is described as a single model, but the paper never says whether it was also fine-tuned per class. If not, the comparison bundles two interventions: per-class specialization (a much easier binary problem) and the memory block. A per-class FastSAM3D without memory is the missing control, and without it the central causal claim is not supported. If the authors actually used per-class FastSAM3D, they need to say so; the text as written doesn't.

Other soft spots, in decreasing order. There are no error bars or significance tests across the 33 classes, so we are looking at point estimates. Only one baseline is used, from the same group. Table 1 reports identical FLOPs after adding kNN attention, which is at minimum underspecified; if FLOPs excludes the lookup and the extra attention, that should be stated. The phrase 'particularly marginal improvements in un-fine-tuned scenarios' contradicts the later claim that memory is especially beneficial without fine-tuning; probably a typo, but it should be fixed. The k hyperparameter is selected using the separate 10-sample tuning set, so that particular concern is less serious than it first looks—but the paper should be explicit that test data was not used for model selection.

On citations, the paper cites the right prior work and evaluates against an external public benchmark, so the result is falsifiable once the code is released. Formal evidence is limited to the dataset and promised code; there is no verified artifact.

Who is this for: applied medical imaging researchers who want a cheap retrofit to a SAM variant for volumetric segmentation. The method description is compact and reproducible in principle. I would send this to review, with the per-class baseline question as the main condition. If the gain survives that control, it is a solid incremental contribution. If it doesn't, the paper is still a useful description of a plug-in that may not be worth the complexity.

Referee Report

5 major / 6 minor

Summary. The paper proposes Memorizing SAM, a 3D medical image segmentation model that augments FastSAM3D with a memorizing transformer block. During training, class-specific models are trained and their internal key-value representations are saved as external memory; at inference, kNN search retrieves top-k pairs that are combined with local attention via a ratio. Evaluated on 33 TotalSegmentator classes (200 training, 10 tuning, 500 test), the method is reported to outperform FastSAM3D by 11.36% average Dice with only 4.38 ms additional inference time. The authors provide code and claim no extra FLOPs or parameters.

Significance. If the reported gains hold under controlled comparison, the memorizing plug-in would be a low-cost way to improve 3D SAM without retraining, especially for under-represented classes. The idea of using precomputed high-quality representations as memory is elegant, and the code release is a strength. However, the current experiments do not isolate the effect of memory from per-class specialization, so the central claim is not yet established; the paper needs additional control experiments and corrected efficiency reporting.

major comments (5)

[Section 2, Overview of Memorizing SAM] The method trains N separate Memorizing SAM models, one per class (N=33), and the headline average Dice improvement is computed against FastSAM3D, which is presented as a single multi-class model. This confounds the memory plug-in with per-class specialization, since each specialized model solves an easier binary segmentation task. Without a per-class FastSAM3D baseline (fine-tuned without the memory block) or a multi-class Memorizing SAM, the 11.36% gain cannot be attributed to the memorizing mechanism.
[Table 1] The reported FLOPs are identical (23.14 G) for FastSAM3D and Memorizing SAM, despite Memorizing SAM adding kNN searches and k attention computations in the memorizing block. This is not physically consistent unless the kNN operations are excluded from the measurement; please report the actual FLOPs including the memory block or clarify precisely which operations are counted.
[Section 3, Results] All performance claims rest on single point estimates with no error bars, standard deviations, or significance tests, and the only baseline is FastSAM3D. Please report repeated-run variability and add at least one additional 3D SAM-based baseline; otherwise the robustness of the 11.36% improvement cannot be assessed.
[Section 3, Ablation Study] The k hyperparameter is selected by ablation, but the text does not state whether the ablation is performed on the 10-sample tuning set or on the 500-sample test set; if test data are used for model selection, the reported Dice values are optimistically biased. Please specify the data split used for choosing k and RL, and avoid test-set-based hyperparameter selection.
[Sections 2-3, inference protocol] The paper does not explain how the 33 class-specific models are run at inference time or how the 6.89 ms per-image inference time is aggregated across classes. If all 33 models must process each volume to produce the final multi-class segmentation, the effective inference time is much larger than 6.89 ms; this directly affects the claim of only 4.38 ms additional cost.

minor comments (6)

[Throughout the text] The term 'un-fine-tuned' is used repeatedly; please replace with 'not fine-tuned' or 'zero-shot' for clarity.
[Abstract] The phrase 'FastSAM3Di.e., FastSAM3D' is missing a space after 'FastSAM3D'; the same typo appears in the abstract and in Section 2.
[Section 3, Results] The value '11.36 ˙%' contains a stray dot and should be '11.36%'.
[Section 2, Memorizing Transformer Block] The equation for Ri is not typeset clearly ('DiPk i=0 RL Di'); please rewrite the combination rule with clear notation, define all symbols, and specify how the distances Di are normalized.
[Figure 2] The subfigures are extremely small and the caption is dense; the per-class Dice comparison would be more readable as a table or box plot with class names.
[Section 2, Overview of Memorizing SAM] The paper does not describe the prompt setting used for all models; since SAM is prompt-based, specify whether the same prompts (e.g., bounding boxes) are provided to FastSAM3D and Memorizing SAM.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity; the claimed gains are evaluated on a held-out external benchmark.

full rationale

The paper's central claim is an empirical improvement in Dice on the TotalSegmentator test set (500 held-out samples) when a memorizing Transformer block retrieves key-value pairs saved during training. This is not circular: the memory bank is explicitly built from training volumes ('Using these trained class-specific SAM models, we process each one-class dataset to generate and save key-value pairs'), and the reported 11.36% average Dice increase is measured on a separate test split, not on the memory training data. The use of the authors' own FastSAM3D as baseline ([16]) is a self-citation, but the comparison is an external evaluation against a public benchmark rather than an argument that reduces to the citation; therefore it is not load-bearing in a circular sense. The per-class specialization of the 33 models is a possible experimental confound (correctness risk), but it does not make the derivation circular by construction. No equation in the paper defines the predicted quantity in terms of the input metric, and no fitted parameter is renamed as a prediction.

Assumptions & free parameters 2 free parameters · 4 assumptions · 0 invented entities

The central claim rests on two fitted hyperparameters (RL = 0.3 and k = 3), a coverage assumption for the small memory set, a comparability assumption for the feature space, and an ad hoc attention combination rule. No new physical entities are introduced.

free parameters (2)

k (number of kNN neighbors) = 3
Chosen via ablation on the test set (Fig. 2c) as the best balance, not on a held-out validation set.
RL (local attention combination ratio) = 0.3
Preset to 0.3 following prior Memorizing Transformer work, with no experiment justifying this value for 3D medical SAM.

assumptions (4)

domain assumption Key-value pairs generated from a small one-class training subset are representative of test anatomy.
Section 2 states that memory quality directly influences efficacy; without coverage of test anatomy, kNN retrieval returns irrelevant values.
domain assumption The same FastSAM3D feature space can compare training-time keys with test-time queries.
The kNN search assumes keys and queries remain comparable after training; no distribution-shift analysis is provided.
ad hoc to paper The combination rule Ac = AL * RL + sum Ri * Ai preserves useful information.
The ratio-based combination is an unproved heuristic defined in the Memorizing Transformer Block section.
domain assumption TotalSegmentator labels are high-quality ground truth for all 33 classes.
The memory source is derived from these labels, so label noise would degrade the memorized values.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Memorizing SAM: 3D Medical Segment Anything Model with Memorizing Transformer." pith.science (2026). https://pith.science/paper/5HE3WHE4

@misc{pith2026241213908,
  author       = {Pith},
  title        = {Pith review of: Memorizing SAM: 3D Medical Segment Anything Model with Memorizing Transformer},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/5HE3WHE4}},
  note         = {Machine review of arXiv:2412.13908}
}

read the original 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

Figures

Figures reproduced from arXiv: 2412.13908 by the authors.

**Figure 2.** Performance comparison between the baseline 3D SAM model [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

19 extracted references · 12 canonical work pages

[1]

State-of-the-art methods for brain tissue segmentation: A review,

Dora, L., Agrawal, S., Panda, R., et al., “State-of-the-art methods for brain tissue segmentation: A review,” IEEE reviews in biomedical engineering 10, 235–249 (2017)

work page 2017
[2]

A survey on deep learning for skin lesion segmentation,

Mirikharaji, Z., Abhishek, K., Bissoto, A., et al., “A survey on deep learning for skin lesion segmentation,” Medical Image Analysis , 102863 (2023)

work page 2023
[3]

Swin transformer improves the idh mutation status prediction of gliomas free of mri-based tumor segmentation,

Wu, J., Xu, Q., et al., “Swin transformer improves the idh mutation status prediction of gliomas free of mri-based tumor segmentation,” Journal of Clinical Medicine 11(15), 4625 (2022)

work page 2022
[4]

Deep learning techniques for tumor segmentation: a review,

Jiang, H., Diao, Z., and Yao, Y.-D., “Deep learning techniques for tumor segmentation: a review,” The Journal of Supercomputing 78(2), 1807–1851 (2022)

work page 2022
[5]

U-net: Convolutional networks for biomedical image segmen- tation,

Ronneberger, O., Fischer, P., and Brox, T., “U-net: Convolutional networks for biomedical image segmen- tation,” 234–241 (2015)

work page 2015
[6]

Clusterseg: A crowd cluster pinpointed nucleus segmentation framework with cross-modality datasets,

Ke, J., Lu, Y., Shen, Y., et al., “Clusterseg: A crowd cluster pinpointed nucleus segmentation framework with cross-modality datasets,” Medical Image Analysis 85, 102758 (2023)

work page 2023
[7]

Transnuseg: A lightweight multi-task transformer for nuclei segmen- tation,

He, Z., Unberath, M., Ke, J., et al., “Transnuseg: A lightweight multi-task transformer for nuclei segmen- tation,” in [ International Conference on Medical Image Computing and Computer-Assisted Intervention ], 206–215, Springer (2023)

work page 2023
[8]

Mednext: Transformer-driven scaling of convnets for medical image segmentation,

Roy, S., Koehler, G., Ulrich, C., Baumgartner, M., Petersen, J., Isensee, F., Jaeger, P. F., and Maier-Hein, K., “Mednext: Transformer-driven scaling of convnets for medical image segmentation,” arXiv preprint arXiv:2303.09975 (2024)

arXiv 2024

Show all 19 references

[9]

Segment anything,

Kirillov, A. et al., “Segment anything,” arXiv preprint arXiv:2304.02643 (2023)

2023 arXiv
[10]

An image is worth 16x16 words: Transformers for image recognition at scale,

Dosovitskiy, A. et al., “An image is worth 16x16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929 (2020)

2020 arXiv
[11]

Segment anything in medical images,

Ma, J., He, Y., Li, F., et al., “Segment anything in medical images,” Nature Communications 15(1), 654 (2024)

2024
[12]

Sam-med2d,

Cheng, J. et al., “Sam-med2d,” arXiv preprint arXiv:2308.16184 (2023)

2023 arXiv
[13]

Segment anything model for medical image analysis: an experimental study,

Mazurowski, M. A., Dong, H., Gu, H., et al., “Segment anything model for medical image analysis: an experimental study,” Medical Image Analysis 89, 102918 (2023)

2023
[14]

Sam3d: Segment anything model in volumetric medical images,

Bui, N.-T., Hoang, D.-H., Tran, M.-T., et al., “Sam3d: Segment anything model in volumetric medical images,” arXiv preprint arXiv:2309.03493 (2023)

2023 arXiv
[15]

Sam-med3d,

Wang, H. et al., “Sam-med3d,” arXiv preprint arXiv:2310.15161 (2023)

2023 arXiv
[16]

Fastsam3d: An efficient segment anything model for 3d volumetric medical images,

Shen, Y. et al., “Fastsam3d: An efficient segment anything model for 3d volumetric medical images,” arXiv preprint arXiv:2403.09827 (2024)

2024 arXiv
[17]

Memorizing transformers,

Wu, Y., Rabe, M. N., Hutchins, D., and Szegedy, C., “Memorizing transformers,” arXiv preprint arXiv:2203.08913 (2022)

2022 arXiv
[18]

Movit: Memorizing vision transformers for medical image analysis,

Shen, Y., Guo, P., and Wu, J. a., “Movit: Memorizing vision transformers for medical image analysis,” in [International Workshop on Machine Learning in Medical Imaging ], 205–213, Springer (2023)

2023
[19]

Totalsegmentator: Robust segmentation of 104 anatomic structures in ct images,

Wasserthal, J. et al., “Totalsegmentator: Robust segmentation of 104 anatomic structures in ct images,” Radiology: Artificial Intelligence 5(5) (2023)

2023

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.

[1] [1]

State-of-the-art methods for brain tissue segmentation: A review,

Dora, L., Agrawal, S., Panda, R., et al., “State-of-the-art methods for brain tissue segmentation: A review,” IEEE reviews in biomedical engineering 10, 235–249 (2017)

work page 2017

[2] [2]

A survey on deep learning for skin lesion segmentation,

Mirikharaji, Z., Abhishek, K., Bissoto, A., et al., “A survey on deep learning for skin lesion segmentation,” Medical Image Analysis , 102863 (2023)

work page 2023

[3] [3]

Swin transformer improves the idh mutation status prediction of gliomas free of mri-based tumor segmentation,

Wu, J., Xu, Q., et al., “Swin transformer improves the idh mutation status prediction of gliomas free of mri-based tumor segmentation,” Journal of Clinical Medicine 11(15), 4625 (2022)

work page 2022

[4] [4]

Deep learning techniques for tumor segmentation: a review,

Jiang, H., Diao, Z., and Yao, Y.-D., “Deep learning techniques for tumor segmentation: a review,” The Journal of Supercomputing 78(2), 1807–1851 (2022)

work page 2022

[5] [5]

U-net: Convolutional networks for biomedical image segmen- tation,

Ronneberger, O., Fischer, P., and Brox, T., “U-net: Convolutional networks for biomedical image segmen- tation,” 234–241 (2015)

work page 2015

[6] [6]

Clusterseg: A crowd cluster pinpointed nucleus segmentation framework with cross-modality datasets,

Ke, J., Lu, Y., Shen, Y., et al., “Clusterseg: A crowd cluster pinpointed nucleus segmentation framework with cross-modality datasets,” Medical Image Analysis 85, 102758 (2023)

work page 2023

[7] [7]

Transnuseg: A lightweight multi-task transformer for nuclei segmen- tation,

He, Z., Unberath, M., Ke, J., et al., “Transnuseg: A lightweight multi-task transformer for nuclei segmen- tation,” in [ International Conference on Medical Image Computing and Computer-Assisted Intervention ], 206–215, Springer (2023)

work page 2023

[8] [8]

Mednext: Transformer-driven scaling of convnets for medical image segmentation,

Roy, S., Koehler, G., Ulrich, C., Baumgartner, M., Petersen, J., Isensee, F., Jaeger, P. F., and Maier-Hein, K., “Mednext: Transformer-driven scaling of convnets for medical image segmentation,” arXiv preprint arXiv:2303.09975 (2024)

arXiv 2024