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

arXiv: https://arxiv.org/abs/2603.05873 · HTML: https://mlsanity.com/paper/2603.05873/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical and scientific domain agents, one of 16 directions in LLM agents and tool use.
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 6 Mar 2026.

## Placed in

- LLM agents and tool use / [Medical and scientific domain agents](https://mlsanity.com/agents/medical-and-scientific-domain-agents/)

## Builds on (most cited first, 25 of 25)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [SAM 2: Segment Anything in Images and Videos](https://mlsanity.com/paper/2408.00714/) (2024)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [DINOv3](https://mlsanity.com/paper/2508.10104/) (2025)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [A-MEM: Agentic Memory for LLM Agents](https://mlsanity.com/paper/2502.12110/) (2025)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [CHAOS Challenge -- Combined (CT-MR) Healthy Abdominal Organ Segmentation](https://mlsanity.com/paper/2001.06535/) (2020)
- [Deep Learning for Segmentation using an Open Large-Scale Dataset in 2D Echocardiography](https://mlsanity.com/paper/1908.06948/) (2019)
- [Medical Image Segmentation Review: The success of U-Net](https://mlsanity.com/paper/2211.14830/) (2022)
- [Foundational Models Defining a New Era in Vision: A Survey and Outlook](https://mlsanity.com/paper/2307.13721/) (2023)
- [MedSAM2: Segment Anything in 3D Medical Images and Videos](https://mlsanity.com/paper/2504.03600/) (2025)
- [Agentic Memory: Learning Unified Long-Term and Short-Term Memory Management for Large Language Model Agents](https://mlsanity.com/paper/2601.01885/) (2026)
- [MedRAX: Medical Reasoning Agent for Chest X-ray](https://mlsanity.com/paper/2502.02673/) (2025)
- [Parameter-Efficient Fine-Tuning for Foundation Models](https://mlsanity.com/paper/2501.13787/) (2025)
- [GraphEcho: Graph-Driven Unsupervised Domain Adaptation for Echocardiogram Video Segmentation](https://mlsanity.com/paper/2309.11145/) (2023)
- [Retrieval-augmented Few-shot Medical Image Segmentation with Foundation Models](https://mlsanity.com/paper/2408.08813/) (2024)
- [Adapting Vision Foundation Models for Real-time Ultrasound Image Segmentation](https://mlsanity.com/paper/2503.24368/) (2025)
- [K-DeCore: Facilitating Knowledge Transfer in Continual Structured Knowledge Reasoning via Knowledge Decoupling](https://mlsanity.com/paper/2509.16929/) (2025)

## Abstract (arXiv)

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.