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

2026-03-06

AI Summary

This work addresses the limited generalization of medical image segmentation models across institutions, imaging devices, or patient populations, as well as the high communication overhead of conventional federated fine-tuning. To this end, the authors propose MemSeg-Agent, which pioneers a paradigm shift from parameter-based adaptation to memory-based operations. By leveraging static memory, few-shot memory, and test-time working memory—all while keeping the backbone network fixed—the method unifies few-shot learning, federated learning, and test-time adaptation without requiring model fine-tuning. This enables dynamic adaptation to new domains with substantially reduced communication costs. Experiments demonstrate that static memory alone matches or surpasses strong supervised baselines, and incorporating test-time memory further enhances both in-domain and cross-domain performance, achieving high parameter efficiency and robustness.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Multiagent Systems: Agent Communication

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• medical image segmentation
• domain generalization
• federated learning
• communication overhead
• model adaptation

Innovation

• memory-augmented agent
• weight-free adaptation
• federated learning
• test-time adaptation
• medical image segmentation