Shifting Adaptation from Weight Space to Memory Space: A Memory-Augmented Agent for Medical Image Segmentation

Bowen Chen; Qiaohui Gao; Shaowen Wan; Shanhui Sun; Wei Liu; Xiang Li; Tianming Liu; Lin Zhao

cs.CVMar 6, 2026

Shifting Adaptation from Weight Space to Memory Space: A Memory-Augmented Agent for Medical Image Segmentation

Authors: Bowen Chen, Qiaohui Gao, Shaowen Wan, Shanhui Sun, Wei Liu, Xiang Li, Tianming Liu, Lin Zhao

Organizations: Department of Electrical and Computer Engineering, University of California, Santa Barbara, Santa Barbara, CA 93106, USA · College of Engineering, Northeastern University, Boston, MA 02115, USA · Department of Biomedical Engineering, New Jersey Institute of Technology, Newark, NJ 07102, USA · United Imaging Intelligence, Burlington, MA 01803, USA · Department of Radiation Oncology, Mayo Clinic, Scottsdale, AZ 85259, USA · Department of Radiology, Massachusetts General Hospital and Harvard Medical School, Boston, MA 02115, USA · School of Computing, University of Georgia, Athens, GA 30602, USA

Abstract

Medical image segmentation is fundamental to clinical workflows, yet models trained on a single dataset often fail to generalize across institutions, scanners, or patient populations. While vision foundation models have shown great promise in addressing this challenge, their deployment typically requires task-specific fine-tuning, which introduces substantial communication overhead in federated learning and prevents continuous knowledge evolution during deployment. In this work, we propose a memory-augmented segmentation agent (MemSeg-Agent) that shifts adaptation from weight space to memory space, enabling few-shot learning, federated supervised learning, and test-time adaptation within a unified architecture. MemSeg-Agent conditions a fixed backbone with lightweight static, few-shot, and test-time working memories, which are dynamically composed by an agentic controller. In federated settings, we update compact memory units instead of model parameters, substantially reducing communication overhead. Experiments on four public datasets demonstrate strong performance and robustness to domain shift: Static memory alone matches or surpasses strong supervised baselines with high parameter efficiency, and test-time working memory further improves in-domain and cross-domain performance without fine-tuning. Overall, MemSeg-Agent introduces a new paradigm for scalable and adaptive medical image segmentation in the era of agentic AI.

Figures & tables

Figure 1: MemSeg-Agent framework. The agent retrieves and composes static, few-shot, and test-time working memories to condition a frozen backbone for segmentation. Human corrections are encoded into the working-memory buffer via gated delta updates, while static memory can be federated for communication-efficient adaptation.

	nnU-Net [ 10 ]	SwinUNETR [ 9 ]	U-Net [ 23 ]	MedSAM2 [ 19 ]	Ours(Static)	Ours(Static+WM)
LV	14.65	33.57	20.30	47.26	39.75	81.53
LA	30.64	0.07	0.03	38.31	21.72	73.07
Avg.	22.65	16.82	10.15	42.75	30.74	77.30

Table 1: Cross-domain evaluation (Train: CAMUS [ 15 ] , Test: CardiacUDA [ 28 ] ) on left ventricle (LV) and left atrium (LA). Dice (%).

		CHAOS (ABD-MRI) [ 13 ]								ACDC [ 4 ]						CAMUS [ 15 ]
		LK		RK		Spl.		Liv.		RV		Myo		LV		LV		Myo		LA
Sup.	Method	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$
100%	UNet [ 23 ]	76.67	43.42	71.37	90.12	67.91	35.03	87.35	13.37	82.03	6.54	86.83	3.05	92.18	4.28	93.36	3.03	86.97	5.00	90.48	4.53
	SwinUNETR [ 9 ]	75.34	94.75	80.98	124.78	72.32	111.60	83.15	54.21	71.06	49.87	79.67	51.16	80.69	46.57	93.32	2.91	94.76	1.62	96.22	1.64
	nnUNet [ 10 ]	89.15	14.57	84.91	15.73	82.82	11.36	89.26	10.76	90.53	1.62	90.36	1.12	94.94	1.20	94.61	1.97	89.94	3.01	92.67	3.01
	MedSAM2 [ 19 ]	85.42	23.05	87.49	18.90	77.45	37.08	88.65	15.62	87.89	6.27	88.03	1.37	92.05	4.13	93.62	2.34	87.75	4.05	90.39	4.59

Table 2: Quantitative evaluation of single-class segmentation results under different supervision levels on CHAOS (ABD-MRI), ACDC, and CAMUS.

Figure 2: Qualitative segmentation results on CHAOS, ACDC, and CAMUS under 100% and 30% supervision. Average Dice scores are shown below each prediction.

Explore similar work

Memory-Supported Synergistic Adaptation for Training-Free Test-Time Medical Image Segmentation
Jul 20, 2026Lingrui Li, Nan Pu, Dong Zhao +4Memory-Augmented VLMs VLM Adaptation
School of Computer Science, University of Nottingham, UK · School of Computer Science and Information Engineering, Hefei University of Technology, China · Information Systems Technology and Design Pillar, Singapore University of Technology and Design, Singapore
Fed-ADApt: Federated Anytime Depth Adaptation for Resource-Aware Medical Image Segmentation
Oct 2, 2026Abhijeet Parida, Zhifan Jiang, Pooneh Roshanitabrizi +6Medical Image Segmentation Heterogeneous Federated Learning
Sheikh Zayed Institute for Pediatric Surgical Innovation, Children’s National Hospital, Washington, DC, United States · Escuela Técnica Superior de Ingenieros de Telecomunicación, Universidad Politécnica de Madrid, Madrid, Spain · CIBER-BBN, Instituto de Salud Carlos III, Madrid, Spain +2
ADP-FL-MedSeg: Adaptive Differential Privacy for Federated Medical Segmentation Across Diverse Modalities
Apr 7, 2026Puja Saha, Eranga UkwattaDifferentially Private Federated Learning Medical Image Segmentation
College of Engineering, University of Guelph, 50 Stone Road E, Guelph, ON, Canada

Jul 20, 2026J/K move · Enter open · S save

Memory-Supported Synergistic Adaptation for Training-Free Test-Time Medical Image Segmentation

Lingrui Li, Nan Pu, Dong Zhao, Wenjing Li+3

School of Computer Science, University of Nottingham, UK · School of Computer Science and Information Engineering, Hefei University of Technology, China · Information Systems Technology and Design Pillar, Singapore University of Technology and Design, Singapore

Memory-Augmented VLMs VLM Adaptation Text-Guided Image Segmentation Medical Image Segmentation

Open paper AskPDF arXiv

Test-time adaptation (TTA) aims to mitigate distribution shifts by adapting models with unlabeled target data at inference time. While TTA with vision-language models (VLMs) has shown promising results in classification, extending it to medical image segmentation remains challenging. In this setting, the adaptation gains from optimizing on VLM-generated predictions are often outweighed by the degradation to the VLM's strong pretrained features caused by noisy, update-driven learning, resulting in limited and unstable improvements. We therefore propose Memory-Supported Synergistic Adaptation (MSSA), a novel training-free TTA framework for medical image segmentation. Without updating model parameters, MSSA dynamically selects reliable image-text predictions to construct an online memory, uses them as text-guided semantic priors, and couples them with cross-image structural alignment for robust adaptation. Specifically, MSSA consists of (i) a noise-aware memory construction module that filters and stabilizes cross-modal predictions, and (ii) a relevance-driven prototype alignment module that aligns the target sample with structurally consistent memory samples and their reliable predictions to improve adaptation. Extensive experiments on multiple medical segmentation benchmarks demonstrate that MSSA consistently improves VLM-based segmentation models and outperforms existing fine-tuning-based TTA methods by a clear margin, with gains of up to 12.2% DSC and 11.7% mIoU. Project page: https://lingrayy.github.io/MSSA/ .

Abstract

Explore similar work

Memory-Supported Synergistic Adaptation for Training-Free Test-Time Medical Image Segmentation

Fed-ADApt: Federated Anytime Depth Adaptation for Resource-Aware Medical Image Segmentation

ADP-FL-MedSeg: Adaptive Differential Privacy for Federated Medical Segmentation Across Diverse Modalities