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

TL;DR

This paper introduces MemSeg-Agent, a memory-augmented model that shifts adaptation from model weights to memory units, enabling efficient, scalable, and robust medical image segmentation across diverse datasets and domain shifts.

Contribution

It proposes a novel memory-based framework for medical image segmentation that reduces communication overhead and supports few-shot, federated, and test-time adaptation within a unified architecture.

Findings

Memory units match or outperform supervised baselines.

Test-time working memory enhances cross-domain performance.

Reduces communication overhead in federated learning.

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.