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

arXiv:2603.05873 · cs.CV · Submitted 2026-03-06 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Shifting Adaptation from Weight Space to Memory Space".

Jane: Here's the gist of the paper. The gist: This work proposes 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.

Jane: MemSeg-Agent conditions a fixed backbone with lightweight static, few-shot, and test-time working memories The agent relies only on lightweight memories (static, few-shot, and test-time working), which are dynamically selected and composed by an agentic controller.

Paper summary: Tom: So, today we’re talking about this paper called "Shifting Adaptation from Weight Space to Memory Space: A Memory-Augmented Agent for Medical Image Segmentation." It sounds technical, but basically, it tackles a big problem in medical AI where models trained on one hospital's data just don't work well at another.

Jane: Exactly. The core idea here is that instead of constantly retraining or fine-tuning the massive foundation models every time you move to a new scanner or institution, this agent shifts the adaptation from those huge model weights into a specialized memory space. It’s about making the model smarter without changing its main structure.

Lu: What they claim is they introduce test-time working memory and static and few-shot memories that let them do things like few-shot learning and federated supervised learning all within one unified architecture. That’s the big promise of this work in terms of how it handles the heterogeneity of medical images, which is a huge issue <ref:2603.05873#pg2>.

Meng: From a deployment standpoint, I’m curious about what they mean by "memory-augmented agent." Are we talking about adding a massive new part to the pipeline that just gets in the way of inference speed? Because at scale, any added component has to be efficient.

Lalam: The paper says the agent conditions a fixed backbone with lightweight static, few-shot, and test-time working memories that are dynamically composed by an agentic controller <ref:2603.05873#pg3>. This controller decides which memory units are needed for any given input.

Tom: That sounds like a smart way to manage resources. So if we’re in a federated learning setting, what does that actually change? Does it make communication cheaper?

Jane: It does significantly. In federated settings, they update compact memory units instead of the whole model parameters <ref:2603.05873#pg3>. They claim this reduces per-round communication by about seventy-four point three times, which is around ninety-eight point six five percent over updating the SAM2-tiny backbone <ref:2603.05873#pg3>.

Lu: And they show that static memory alone, without any extra working memory, matches or even surpasses strong supervised baselines with high parameter efficiency <ref:2603.05873#pg1>. That suggests the initial learned priors are quite good for capturing those segmentation details.

Meng: High parameter efficiency is great because it means less storage and faster loading times on the edge devices where we actually run these things in a hospital setting. But what about when the data truly shifts, like moving from one patient population to another?

Jane: That’s where the test-time working memory comes in, which allows for continual feedback-driven updates during inference <ref:2603.05873#pg3>. If a human reviewer corrects a prediction, that corrected observation gets encoded into memory via a gated delta-update strategy <ref:2603.05873#pg4>.

Paper summary: Tom: So it’s not just static knowledge; it’s dynamic knowledge being absorbed on the fly based on real-time feedback from users. That sounds like something that could actually improve clinical workflows if implemented correctly.

Lu: The agentic memory controller orchestrates this by first checking static memory similarity using DINOv3 embeddings <ref:2603.05873#pg4>. If it doesn't find a good match, it requests few-shot examples to create new priors <ref:2603.05873#pg4>.

Jane: It’s a loop of checking static memory first, then grabbing few-shot examples if needed, and then composing the static and working memory for the target domain shift <ref:2603.05873#pg1>. This allows it to handle those cross-domain discrepancies without needing any backbone fine-tuning.

Meng: I’m looking at the results on four public datasets—CHAOS abdominal MRI, ACDC cardiac MRI, and CAMUS 2D echocardiography—and the performance jump is substantial <ref:2603.05873#pg1>. But what are the limitations? What does this system not do?

Tom: The authors themselves flag that they focus on using memory to unlock foundation models rather than suggesting a replacement for them entirely. They restrict cross-site updates to these lightweight memories, which is key because it keeps communication costs down <ref:2603.05873#pg1>.

Lu: They also noted in the ablation study that when static memory is already well-aligned with the target distribution, performance peaks at small k and then gradually degrades as memory size increases because of redundant aggregation <ref:2603.05873#pg1>.

Jane: So, for a perfectly aligned dataset, you don't need a huge working memory; you just need the right static priors. But under cross-domain conditions, increasing the working memory size consistently improves segmentation accuracy during inference <ref:2603.05873#pg1>.

Tom: That’s a crucial distinction—the trade-off between having enough knowledge versus having too much redundant information when things get messy. It sounds like this agent is designed to be adaptive without being overly complex in its storage requirements.

Meng: From a production view, if we deploy this, the complexity shifts from managing model weights to managing the memory unit composition and ensuring that the agentic controller doesn't waste computation time deciding what memory to pull for every single image. That decision-making process has to be super fast.

Lalam: The paper confirms that test-time adaptation via working memory improves performance to seventy-seven point thirty percent average Dice, which is a plus of forty-six point fifty-six percent absolute gain over the static-only variant <ref:2603.05873#pg1>. This shows the working memory component is genuinely useful for cross-domain adaptation.

Jane: It really does show that test-time feedback can compensate for distribution shift without ever touching those heavy backbone parameters <ref:2603.05873#pg1>. So, if you’re dealing with a new scanner, this framework lets you adapt quickly just by feeding it corrected samples.

Paper summary: Tom: I think the main implication here is that we might not always need to spend huge amounts of compute or massive parameter updates just to make a segmentation model work everywhere. This memory-space approach is a much more scalable way to deploy these complex vision models in real clinical settings <ref:2603.05873#pg1>.

Lu: It validates the idea that we can leverage the intrinsic capabilities of foundation models by augmenting them with explicit, task-specific memories instead of relying solely on parameter-efficient fine-tuning methods like PEFT <ref:2603.05873#pg3>. This is about unlocking what the models are already capable of doing.

Jane: So, to wrap up this discussion on "Shifting Adaptation from Weight Space to Memory Space: A Memory-Augmented Agent for Medical Image Segmentation," we see a system that uses memory—static, few-shot, and working—to handle learning across different sites and domains without needing constant model retraining.

Tom: That’s the main point. The authors show that updating compact memory units instead of the massive backbone parameters reduces communication by about seventy-four point three times in federated scenarios <ref:2603.05873#pg1>.

Meng: It means we can run these powerful segmentation models across many institutions with much less data transfer overhead than we currently estimate for full model updates <ref:2603.05873#pg1>. That’s a big win for real-world deployment efficiency.

Lu: The framework successfully unifies few-shot learning, federated supervised learning, and test-time adaptation in a single architecture by shifting adaptation from weights to memory <ref:2603.05873#pg1>.

Jane: It really shows that we don't always have to fine-tune the entire model just because the data is new or different. Instead, you can just update your specialized memory units with new information <ref:2603.05873#pg1>.

Tom: This whole concept of the Agentic Memory Controller dynamically composing these memories based on similarity checks sounds like a very robust way to manage adaptation at inference time <ref:2603.05873#pg4>.

Meng: I’m just hoping that when we try to deploy this, the memory retrieval process itself doesn't introduce too much latency in the actual segmentation step <ref:2603.05873#pg4>. That’s where my biggest concern lies with any added inference complexity.

Lu: The research suggests that static memory alone is a strong baseline, but adding that test-time working memory component is what truly improves cross-domain performance during inference <ref:2603.05873#pg1>.

Jane: So the big picture here is that this memory-augmented agent offers a new path for medical image segmentation where adaptation is just as important as initial training, and it does it by being very smart about where it stores its learned knowledge <ref:2603.05873#pg1>.

Conclusion: Tom: So, to wrap up this look at "Shifting Adaptation from Weight Space to Memory Space: A Memory-Augmented Agent for Medical Image Segmentation," we’re talking about how they move the adaptation process out of the model's main weights and into a separate memory space.

Jane: That sounds like they’re trying to solve that communication problem in federated learning, right? It suggests that instead of sending massive model updates between hospitals, you just send small memory units.

Lu: Exactly, it shifts the focus from updating the one hundred million weight parameters to updating these compact memory units. The authors show this reduces communication by about seventy-four point three times in a simulated four-site federated setting.

Meng: Seventy-four times is a huge reduction in data transfer, which means we could actually scale this type of model across many institutions without the network becoming the biggest bottleneck. But what about the performance trade-offs?

Tom: The results on static memory alone show it can match or even beat strong supervised baselines with high parameter efficiency. That’s a solid starting point before you even add any working memory stuff.

Jane: And when they bring in that test-time working memory, it boosts performance to seventy-seven point thirty percent average Dice, which is a gain of about forty-six point five six percent over the static-only version. That’s significant for handling data shift.

Lu: The method works by letting an agentic controller dynamically select between static memory and working memory based on similarity checks with embeddings, like DINOv3. It builds a unified architecture where adaptation happens in this memory layer instead of the backbone weights themselves.

Meng: I'm still thinking about the real-world deployment cost. How fast does that agentic controller have to decide which piece of memory to use for every single image during inference? That decision-making process has to be constant time, otherwise it eats up all our GPU cycles.

Tom: The authors note that they focused on using memory augmentation to unlock the foundation models' existing capabilities, rather than trying to replace them entirely with something smaller. They kept the backbone fixed.

Jane: So the big implication is that we can use these large vision models for complex tasks without needing constant, expensive fine-tuning every time we get new data. It makes deployment much more plug-and-play for clinical settings.

Lu: The method successfully unifies few-shot learning, federated supervised learning, and test-time adaptation into one structure by shifting the adaptation mechanism from weights to memory. That’s the core contribution here.

Meng: Okay, so we have a way to be more communication efficient and more adaptive without retraining the entire model structure. I wonder how they plan to handle cases where the static memory is completely useless for a new domain?

Tom: The setup suggests that if static memory similarity is low, the agent requests few-shot examples to generate new priors, which then get encoded into working memory. It’s an iterative process of building context.

Jane: That iterative building process sounds like it's what allows it to adapt during inference when the target domain deviates from the training distribution, using human feedback to guide those updates.

Lu: The paper confirms that static memory alone degrades under distribution shift, but the Agent-driven test-time working memory progressively improves segmentation accuracy as it absorbs target-specific representations.

Meng: So for a listener who only cares about running this on a hospital server, the message is: you can keep your massive model frozen and just manage these lightweight memories for adaptation instead of constantly pushing heavy parameter updates.

Tom: This approach offers a new way to handle cross-domain robustness in medical imaging by treating knowledge as something you can dynamically load and compose from memory units.

Jane: It’s a shift in thinking from parameter-space adaptation to an agentic memory space, which is really interesting for the future of applied AI.

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

University of California, Santa Barbara Department of Electrical and Computer Engineering Department of Biomedical Engineering New Jersey Institute of Technology United Imaging Intelligence Mayo Clinic Massachusetts General Hospital Harvard Medical School University of Georgia School of Computing

cs.CV

Submitted: 2026-03-06

Updated: 2026-10-18

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 89/100

The gist: The gist: This work proposes a memory-augmented segmentation agent (MemSeg-Agent) that shifts adaptation from weight space to memory space, enabling few-shot learning, federated supervised learning,

Key concepts

Memory-Augmented Agent (MemSeg-Agent): A unified architecture where a fixed neural network backbone is augmented with three types of lightweight memories: static (for general priors), few-shot (for limited labeled data), and test-time working memory (for online adaptation). An agentic controller dynamically selects and composes these memories during inference to handle different scenarios.
Shifting Adaptation from Weight Space to Memory Space: Instead of fine-tuning the massive parameters of a foundation model (weight space) when encountering new data or domains, this method adapts by modifying small, compact memory units. This allows for rapid adaptation—like learning from a few examples or correcting errors during testing—without altering the main model structure.

Terminology

Adaptation — 3× in this paper · explained in 1 episode
Communication Overhead — 2× in this paper · explained in 2 episodes
Fine-Tuning — 2× in this paper · explained in 5 episodes
Agentic AI — 1× in this paper · explained in 2 episodes
Domain Shift — 1× in this paper · explained in 7 episodes
Federated Learning — 1× in this paper · explained in 17 episodes
Few-Shot Learning — 1× in this paper · explained in 5 episodes
Foundation Model — 1× in this paper · explained in 3 episodes
Foundation Models — 1× in this paper · explained in 7 episodes
Parameter Efficiency — 1× in this paper · explained in 3 episodes
Robustness — 1× in this paper · explained in 5 episodes

Summary

The gist: This work proposes 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.

How it works

MemSeg-Agent conditions a fixed backbone with lightweight static, few-shot, and test-time working memories The agent relies only on lightweight memories (static, few-shot, and test-time working), which are dynamically selected and composed by an agentic controller.
In federated settings, the system updates compact memory units instead of model parameters, substantially reducing communication overhead.
The framework is built upon a frozen SAM2 backbone where task adaptation is achieved through explicitly parameterized memories, including static memory (Section 3.2), few-shot memory (Section 3.3), and test-time working memory (Section 3.4).

Static Memory via Supervision

SAM2’s memory mechanism implicitly defines a representation space M that encodes segmentation priors.
To construct static memory, the authors optimize only the pseudoobservations zi N i=1 while keeping SAM2 backbone fθ as well as the memory encoder Em fixed by solving min Lsegfθ x; for Em(zi) N i=1, y.
The loss function used is L = Ldice(ˆy, y) + 0.3Lobj on slice-wise binary segmentation using Dice Loss plus an auxiliary objectness BCE on foreground-presence logits.

Few-shot Memory

When labeled data is scarce, retrieval augmented few-shot segmentation can be seamlessly integrated into the framework to support few-shot scenarios.
Given a support set S = (xi, yi) K i=1, each pair is encoded into memory mi = Em(xi, yi) ∈ M as non-parametric priors.
At inference, relevant memories are retrieved by similarity and condition the frozen backbone for the query xq, i.e., yˆq = fθ(xq; m i∗).

Test-time Working Memory

A test-time working memory adapts segmentation priors online during inference when the target domain deviates from the training distribution.
For each test sample xt, a human-in-the-loop may provide a corrected annotation y∗ t when the model prediction is unsatisfactory.
The corrected observation (xt, y∗ t) is encoded into memory mt = Em(xt, y∗ t) ∈ M and written into a working memory buffer via a gated delta-update strategy.
Updates are triggered only when Dice(y∗ t, yt) < τdice, and then DINOv3 embedding similarity is computed to existing entries to decide whether to add or update a memory entry.

Agentic Memory Controller

The Agentic Memory Controller orchestrates memory usage at inference time by querying static memory first through DINOv3 similarity.
If a sufficiently similar prototype exists, the agent directly conditions the frozen SAM2 backbone with static memory.
Otherwise, the controller requests few-shot examples and encodes them as few-shot memory priors.
During deployment, the controller composes static and working memory to handle domain shift, increasing the working memory contribution and update frequency when similarity to static memory is low.

The proposed method demonstrates strong performance and robustness to domain shift on four public datasets. Specifically, the static-only variant reaches 30.74% average Dice and outperforms nnU-Net, SwinUNETR, and U-Net. Furthermore, enabling test-time adaptation through working memory improves performance to 77.30% average Dice, a +46.56% absolute gain over static-only and nearly doubling MedSAM2. In a simulated four-site FedAvg setting, updating a 2M-parameter memory instead of the 148.63M SAM2-tiny backbone reduces per-round communication by ∼74.3× (∼98.65%) while maintaining competitive performance. This indicates that working memory can dynamically absorb domain-specific information from test samples and compensate for cross-domain discrepancies without updating backbone parameters. The framework introduces a new paradigm for scalable and adaptive medical image segmentation in the era of agentic AI. The overall contribution is that MemSeg-Agent unifies few-shot learning, federated supervised learning, and test-time adaptation in a single architecture by shifting adaptation from weights to memory. The results show that static memory alone matches or outperforms strong baselines with high parameter efficiency. This approach provides an effective and scalable alternative to parameter updates for cross-domain adaption.

The ablation study confirms the necessity of both components, showing that removing either static or test-time working memory leads to performance degradation across all datasets. The effect of working memory size shows that when static memory is well-aligned with the target distribution (in-domain), performance peaks at small k and gradually degrades with larger memory sizes due to redundant aggregation. Conversely, under cross-domain settings, increasing k consistently improves performance, suggesting that larger working memory capacity facilitates progressive domain adaptation by absorbing target-specific representations. The framework successfully achieves strong performance with minimal additional parameters.

The authors evaluate the method on three public datasets: CHAOS abdominal MRI, ACDC cardiac MRI, and CAMUS 2D echocardiography. Cross-domain testing is performed using CardiacUDA as an external benchmark. The evaluation confirms that conventional baselines suffer substantial performance drops under domain shift. The results indicate that static memory alone leads to severely degraded performance under distribution shift, whereas enabling the Agent-driven test-time working memory progressively improves segmentation accuracy during inference. This demonstrates that working memory can dynamically absorb domain-specific information from test samples and compensate for cross-domain discrepancies without updating backbone parameters. The framework achieves performance largely comparable to fully supervised baselines across datasets and supervision levels.

The paper discusses the limitations of foundation models, noting that most works still require task-specific finetuning of foundation models. The authors leverage the memory mechanism to unlock and fully utilize the intrinsic capabilities of these foundation models. They restrict cross-site updates to lightweight memories, significantly reducing communication costs while preserving strong performance. This approach validates the effectiveness of static memory for encoding segmentation priors. The framework achieves comparable performance with substantially fewer learnable parameters, enabling scalable plug-and-play deployment without retraining the segmentation model. The final conclusion is that MemSeg-Agent improves cross-domain robustness and reduces FL communication by updating compact memory units instead of model parameters.

The paper evaluates the performance of static memory in Table 2 and qualitative results in Fig. 2. The cross-domain evaluation shows that only static memory leads to severely degraded performance under distribution shift, whereas enabling the Agent-driven test-time working memory progressively improves segmentation accuracy during inference.

Improvements for AI systems

Shifting Adaptation Paradigm to Memory Space: Few-Shot Learning and Test-Time Adaptation in a Unified Architecture. The MemSeg-Agent architecture shifts adaptation from weight space to memory space, enabling few-shot learning, federated supervised learning, and test-time adaptation within a unified architecture. This allows the model to adapt to new domains or conditions without requiring task-specific fine-tuning.
Communication Efficient Federated Learning via Memory Updates. By updating compact memory units instead of model parameters, the system achieves communication efficient federated learning via memory-only updates, reducing per-round communication by ∼74.3× (∼98.65%) over updating the SAM2-tiny backbone while preserving segmentation accuracy. This enables collaborative training across institutions under strict data privacy constraints with minimal network overhead.
Scalable and Task/Site-Specific Prior Maintenance via Static Memory. The framework utilizes static memory to compactly maintain task-/site-specific segmentation priors to plug-and-play condition a frozen backbone, achieving strong performance with minimal additional parameters. This allows the system to rapidly initialize segmentation tasks by retrieving relevant prior memories based on similarity.
Continual Feedback-Driven Adaptation via Test-Time Working Memory. The test-time working memory enables continual feedback-driven updates, where human corrections are encoded into memory via a gated delta-update strategy [29] to adapt segmentation priors while avoiding uncontrolled memory growth and noisy accumulation. This allows the system to improve performance on unseen test samples online during inference.
Agentic Inference Orchestration for Dynamic Memory Composition. The Agentic Memory Controller dynamically selects and composes memories, such as request[ing] few-shot examples, to condition the backbone, ensuring that when static memory similarity is low, the system increase[s] the working memory contribution and update frequency. This enables adaptive behavior during inference by balancing pre-learned priors with on-the-fly adaptation.