🤖 machine learning

Composition of Memory Experts for Diffusion World Models

This paper introduces a diffusion-based world model that overcomes the memory-fidelity trade-off of existing architectures by composing specialized short-term, episodic, and spatial memory experts through a contrastive product-of-experts formulation, thereby achieving scalable, high-fidelity future prediction without quadratic computational costs.

Original authors: Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to tell a story about a long journey you took through a maze. You want the story to be perfect: the turns should make sense, the rooms you visited hours ago should look exactly the same when you return, and the details of the scenery should be sharp.

The problem with current AI "storytellers" (called World Models) is that they have a memory problem.

The "Short-Term" AI is like a person with a great memory for the last few sentences they heard, but they forget what happened an hour ago. They can describe the immediate next step perfectly, but they lose the plot of the whole journey.
The "Long-Term" AI is like a historian who remembers the whole trip but gets the details of the immediate next step fuzzy. They know the general shape of the maze but might mix up the colors of the walls in the room you are standing in right now.

Trying to build one giant AI that remembers everything perfectly is like trying to carry a library in your backpack while running a marathon; it's too heavy, too slow, and eventually, you trip over the books.

The Solution: A Team of Specialists

The authors of this paper propose a new way to build these AI storytellers. Instead of one giant brain, they use a team of specialized experts who work together to tell the story. They call this Composition of Memory Experts (CoME).

Think of it like a production team for a movie:

The "Short-Term" Expert (The Scriptwriter): This expert focuses on the immediate scene. They look at the last few frames of video and say, "Okay, the character just turned left, so the next frame should show a wall on the right." They are great at fine details and immediate movement but can't remember what happened 100 frames ago.
The "Long-Term" Expert (The Archivist): This expert has a massive library of the entire journey. They don't look at every single frame in real-time (which would be too slow). Instead, they "study" the history of the trip and update their own internal notes (their "weights") to remember the big picture. If you ask, "What did that room look like 5 minutes ago?" they can recall it perfectly because they've memorized the episode.
The "Spatial" Expert (The Mapmaker): This expert keeps track of where things are. If the AI is walking in a loop (like a hallway that looks the same everywhere), the other experts might get confused. The Mapmaker says, "Wait, we are at the North entrance, not the South one," ensuring the story stays in the right place.

How They Work Together: The "Product of Experts"

Usually, if you ask three people for advice, you might just take the average. But in AI, averaging can sometimes create a muddy, confused result.

The authors use a clever trick called a Product of Experts. Imagine the three experts are all holding up signs with their predictions.

If the Scriptwriter, the Archivist, and the Mapmaker all agree on what the next frame should look like, the AI is very confident and draws that frame.
If they disagree, the AI doesn't just guess; it looks for the "common ground."

However, there's a catch. Sometimes experts agree on the wrong things just because they are echoing each other. To fix this, the authors invented a "Contrastive" method.

The Analogy: Imagine the experts are singing a song. Sometimes they all hum the same wrong note because they are listening to each other too much. The "Contrastive" method acts like a conductor who says, "Stop humming the background noise! Focus only on the unique notes that are different from your usual habits." This ensures the AI doesn't get stuck in a loop of repeating mistakes and keeps the story fresh and accurate.

The Results

The paper tested this team approach on:

Virtual Mazes: Where an agent has to navigate a 3D maze and remember where it's been.
Real-World Videos: Like walking through a house or driving a robot.

They found that this team approach was much better than trying to use one giant AI.

Better Memory: The AI could remember details from hundreds of frames ago without getting confused.
Faster: It didn't need to carry a massive "library" of data in its head at every single step; it just consulted its specialists.
More Consistent: When the AI walked back into a room it had visited earlier, the room looked exactly the same, not a blurry mess.

In Summary

The paper argues that instead of trying to build one super-brain that does everything, we should build a committee of specialized brains. By letting a short-term expert handle the details, a long-term expert handle the history, and a spatial expert handle the location, and then having them vote on the future using a special "contrastive" rule, we can create AI that remembers the past perfectly while predicting the future accurately—all without needing a supercomputer to do it.

Technical Summary: Composition of Memory Experts for Diffusion World Models

Problem Statement

World models (WMs) are essential for reinforcement learning (RL) and planning, as they predict plausible future states consistent with past observations. However, existing architectures face a fundamental trade-off between memory fidelity and computational efficiency. Transformer-based models preserve local detail but suffer from quadratic attention costs, limiting context length. Conversely, recurrent networks and state-space models (SSMs) scale efficiently but compress history into hidden states, inevitably discarding fine-grained details over time. Simply increasing context length in a single architecture leads to unstable training and prohibitive inference costs. Furthermore, naive composition of multiple models often results in "mode collapse," where the combined distribution over-sharpens common regions and loses diversity.

Methodology

The authors propose Composition of Memory Experts (CoME), a diffusion-based framework that decouples future–past consistency from a single monolithic architecture. Instead, CoME leverages a set of specialized experts, each handling a specific memory role, and fuses them probabilistically at inference time.

1. Product of Contrastive Experts (PoCE)

To address the limitations of standard Product of Experts (PoE), which can suppress diversity and amplify spurious modes, the authors introduce a Product of Contrastive Experts formulation.

Mechanism: For a set of heterogeneous experts $p_i(x)$ , the composition is defined as $p(x) \propto \prod \tilde{p}_i(x)$ , where $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$ .
Contrastive Baseline: $\bar{p}_i(x)$ represents the expert's unconditional baseline distribution. The coefficient $\alpha_i > 1$ acts as a contrasting factor.
Theoretical Justification: Unlike naive tempering (which shrinks kernel variance and distorts local geometry), this contrastive mixture reweights the mixture components while preserving the shape of the kernels. It effectively suppresses spurious modes (where an expert assigns high probability without strong evidence) while maintaining valid secondary modes, provided the conditional expert assigns higher probability to valid modes than the baseline.

2. Specialized Memory Experts

The framework instantiates three complementary experts:

Short-Term Memory (STM): A diffusion model conditioned on a recent context window ( $10\text{--}100$ frames). It captures fine local dynamics and immediate temporal dependencies using an architecture with extended context length (e.g., sliding window attention).
Long-Term Memory (LTM): A separate diffusion model that stores episodic history ( $100\text{--}1000$ frames) directly in its weights. This is achieved via test-time finetuning of lightweight LoRA adapters on the long-term context. This approach amortizes the cost of recall, allowing constant-time reuse of past experience without feeding the entire history through the model at every step. It avoids catastrophic forgetting by freezing the backbone and only updating adapters.
Spatial Long-Term Memory (SLTM): An extension of LTM that incorporates spatial priors (e.g., camera poses, point maps) extracted from the history. This expert enforces geometric and spatial coherence, disambiguating locations in repetitive or visually similar environments.

3. Unified Generation

The final distribution is a product of the pretrained prior and the three contrastive experts (STM, LTM, SLTM). During sampling, the score functions of these experts are combined additively (in log-space) to guide the denoising process, ensuring the generated future respects recent dynamics, long-term history, and spatial consistency.

Key Contributions

Decoupled Memory Architecture: The paper proposes solving the memory trade-off by distributing memory across specialized experts rather than a single backbone, formulated as a Product of Experts problem.
Product of Contrastive Experts (PoCE): A novel strategy to fuse heterogeneous experts. The authors provide theoretical and empirical evidence that this method prevents mode collapse and over-confidence by factoring out redundant modes, unlike standard PoE.
External Diffusion Memory: The introduction of an LTM module that uses test-time finetuning of LoRA adapters to store episodic knowledge in weights. This is extended to SLTM to incorporate spatial structure, improving accuracy in navigation tasks.

Experimental Results

The method was evaluated on synthetic (Memory Maze, DMLab, Minecraft) and real-world (RealEstate10K, RECON) benchmarks.

Reconstruction Quality: On the Memory Maze dataset, CoME achieved state-of-the-art performance in LPIPS (0.097), SSIM (0.892), and PSNR (23.07), significantly outperforming baselines including full-attention transformers, SSMs, and sliding-window models.
Scalability: Unlike full-attention models which require $60\times$ more compute than the STM baseline, CoME achieves superior consistency with a fraction of the training and inference cost.
Context Length: The model showed no saturation in performance consistency up to 480 frames of context, demonstrating the efficacy of the LTM component.
Planning and Navigation: In goal-conditioned navigation tasks (RECON benchmark), CoME reduced Absolute Trajectory Error (ATE) to 0.96, outperforming specialized navigation world models (NWM) and other baselines.
Recall Ability: In a reverse-trajectory experiment on RealEstate10K, CoME successfully recalled initial frames after forward rollouts, whereas baselines without memory generated plausible but inconsistent scenes.
Ablation Studies: Removing the contrastive formulation caused performance to collapse, confirming the necessity of the PoCE mechanism. Increasing LoRA rank and context length for LTM consistently improved results without overfitting.

Significance and Claims

The paper claims to establish a novel paradigm for building memory-augmented diffusion world models. By distributing memory across specialized experts and using a contrastive composition strategy, CoME balances efficiency with high-fidelity temporal consistency. The authors argue that this approach allows world models to scale to long contexts without incurring quadratic costs or suffering from the detail loss inherent in recurrent compression. The method is presented as a general framework applicable to various diffusion-based world models, enhancing their reliability for planning and decision-making in reinforcement learning settings. The work emphasizes that memory should be a distributed system property rather than a constraint of a single architecture.