Composition of Memory Experts for Diffusion World Models
Authors: Sebastian Stapf
Pablo Acuaviva
Aram Davtyan
Paolo Favaro
Venue: ICLR 2026 Poster (2026)
Keywords: World Model
Diffusion Model
Memory
Generative Models
Video Generation
Primary area: generative models

Abstract: World models aim to predict plausible futures consistent with past observations, a
capability central to planning and decision-making in reinforcement learning. Yet,
existing architectures face a fundamental memory trade-off: transformers preserve
local detail but are bottlenecked by quadratic attention, while recurrent and state-
space models scale more efficiently but compress history at the cost of fidelity. To
overcome this trade-off, we suggest decoupling future–past consistency from any
single architecture and instead leveraging a set of specialized experts. We introduce
a diffusion-based framework that integrates heterogeneous memory models through
a contrastive product-of-experts formulation. Our approach instantiates three
complementary roles: a short-term memory expert that captures fine local dynamics,
a long-term memory expert that stores episodic history in external diffusion weights
via lightweight test-time finetuning, and a spatial long-term memory expert that
enforces geometric and spatial coherence. This compositional design avoids mode
collapse and scales to long contexts without incurring a quadratic cost. Across
simulated and real-world benchmarks, our method improves temporal consistency,
recall of past observations, and navigation performance, establishing a novel
paradigm for building and operating memory-augmented diffusion world models.

Published as a conference paper at ICLR 2026
COMPOSITION OF MEMORY EXPERTS FOR DIFFUSION WORLD MODELS
Sebastian Stapf & Pablo Acuaviva Huertos & Aram Davtyan & Paolo Favaro Computer Vision Group
Department of Computer Science
University of Bern
(sebastiane stapf, pablo. acuavivahuertos aram. davytan, paolo. favarof@unibe.co
ABSTRACT
World models aim predict plausible futures consistent with past observations, capability central to planning and decision- decision-making in reinforcement learning. Yet, existing architectures face fundamental memory trade-off:: transformers preserve local detail but are bottlenecked by quadratic attention, while recurrent and state- space models scale more efficiently but compress history at the cost of fidelity. To overcome this trade-off, we suggest decoupling future-past consistency from any single architecture and instead leveraging set specialized experts. We introduce liffusion-based framework that integrates heterogeneous memory models through a contrastive product-of-eppert formulation. Our approach instantiates three complementary roles: short-term memory expert that captures fine local dynamics, a long-term memory expert that stores episodic history in external diffusion weights via lightweight test-time finetuning, and spatial long-term memory expert that enforces geometric and spatial coherence. This compositional design avoids mode collapse and scales long contexts without incurring quadratic cost. Across simulated and real-world benchmarks, our method improves temporal consistency, recall of past observations, and navigation performance, establishing novel paradigm for building and operating vemory-aaugmetee diffusion world models.
INTRODUCTION
World models (WMs) are powerful tools designed to predict plausible future states of the world based on past observations Ha Schmidhuber (2018b ay; Hainer et al. (2019); Hu al. (2023): [Buuee et al.) (2024). By learning to model the distribution of the observed environment, WMs implicitly capture its underlying rules and dynamics. This capability makes them especially attractive for decision-making in downstream tasks Alonso al. (2024); Ha & Schmidhuber (201 8a): Hafner et al. (2019). Recent advances, particularly in diffusion based world models, have demonstrated remarkable progress in generating high-quality future trajectories Hassan al. (2024). Crucially, this predictive ability makes WMs well suited for navigation and interaction, where agents must anticipate future states in order plan and effectively. In such settings, an agent explores, perceives objects, and selects trajectories under uncertainty by simulating candidate futures and evaluating their outcomes before acting. These imagined rollouts are only useful if they remain consistent with prior observations; for example, previously visited room should not spontaneously change its contents on later visit. Maintaining such cross-time consistency is the crux of reliable prediction and decision decision-making.
Unfortunately, predictive WMs face structural trade-off: richer temporal context improves fidelity but explodes compute. Transformer architectures can produce high-quality rollouts, yet their quadratic attention scales poorly and caps context length Vaswani et al. (001)). Recurrent networks and state- space models scale more gracefully in context, but compress history into hidden states, inevitably discarding detail over time. There is silver bullet here: each family wins in one regime and loses in another. Simply "turning up" context is not a sustainable solution, because training becomes unstable and expensive, and inference costs quickly blow past practical budgets.

---

Published as a conference paper at ICLR 2026
We advocate different stance: memory should be distributed across system rather than confined to a single architecture. Human cognition illustrates this principle by separating fast, sapacity-limiter short-term memory (STM) from slower but durable long-term memory (LTM). These forms of memory differ in both mechanism and purpose, and precisely this division that makes them effective together Liu et al. (2025). We adopt the same philosophy for world models (WMs)): instead of burdening a single backbone with every temporal demand, we structure memory as a composition of specialized experts. This distributed approach enables WMs balance efficiency with fidelity, supporting scalable and reliable use of past experience. Beyond composition, we address the hard reality that scaling context in training and inference is costly and brittle. Even with linear-time time designs, pushing horizons indefinitely is not viable. To keep costs flat, we add long-term memory channel that stores episodic knowledge directly the weights of an external diffusion expert. A small number of targeted "memorization" updates amortizes future recall, enabling constant-time reuse of past experience without dragging full histories through the core model at every step.
Diffusion models are particularly good fit for this idea because they permit principled inference- time composition of heterogeneous experts without retraining Du et al. (023). We leverage this to integrate large pretrained backbones with lightweight adapters and auxiliary specialists into unified framework. Naïve composition, however, not enough: when experts still share modes that are not consistent with the past, standard product-of-expprrs,, regions and collapse consistency. We therefore introduce a contrastive mechanism that factors redundant modes during composition, preserving agreement where it matters while avoiding over-confidence where experts are merely echoing each other.
we show, tends to over-sharpen common
Finally, for navigation the world is not just temporal, it is spatial. Where the agent is, and how observations tie to place, is decisive for consistency. We hypothesize that anchoring memory to spatial priors (e.g.. pose/topological cues or map-aligned keys) improves retrieval and composition, especially in RL-style tasks that revisit locations under changing viewpoints. Accordingly, we ground our memory modules in spatial structure, so imagined futures respect both what was seen where it was seen.
To summarize, our main contributions are:
1. We aim solve the memory problem, by using multiple expert models that excel modeling memory at different scales and formulate it as a product-of-expert (PoE) problem, enabling a principled probabilistic view of memory integration video world models;
2. We introduce product of contrastive experts (PoCE) a novel strategy tailored for this setting, ensuring that heterogeneous memory experts can be fused together for consistent prediction. We provide both thcoretical and empirical evidence supporting its necessity;
3. We propose leveraging an external diffusion model as long-term memory (LTM) with finetuning strategy that adapts pretrained priors while retaining domain-genedt capabilities, and extend this framework with spatial long-term memory models (SLTM) to further improve accuracy and spatial coherence in generated sequences.
PRIOR WORK
Video Generation and World Models Recent video diffusion models have demonstrated strong capabilities in modeling temporally coherent video sequences, though ensuring long-term consistency remains an open challenge He et al. (2022); Henschel et al. (2024); (Bar-Tal et al.) [
(2024): Ma et al. (2024); [in et al. (2024); Davtyan et al. (023).. This has motivated several lines of work, including architectural modifications such as transformer variants designed for temporal stability Tay et al. (2020); Lu et al. (2024) or different modeling frameworks Fuest et al. 2025): Ge et al. (2022). Alternative sampling strategies have also been explored, with methods that guide generation toward temporally aligned predictions based on prior context Yin al. (2023); Chen (et al.] (2024): [Song et al.] (2025); Davtyan et al. (2023). In world modeling, early work learned approaches incorporate generative memory to extend reasoning
dynamics from pixels;
2

---

Published as a conference paper at ICLR 2026
over longer horizons Feng et al. (2023) Deng et al. (2023): Samsami al. (0224 Alonso et al. (2024)
Product of Experts and Compositional Adaptation Adapting large pretrained diffusion models to new domains is challenging due to their size and domain specificity. One approach, probabilistic adaptation Yang et al. (2024a), trains smaller diffusion model alongside frozen pretrained one, combining their scores to flexibly adapt to new domains while retaining the robustness of the original.
Combining multiple conditional experts Du et al. (2023) has been applied for image generation. Beyond images, compositionality has been explored for task and visual planning Liu et al. (2022). as well as for factorizing language instructions to guide video generation models Ajay et al. (2023)
Memory for Video Models SlowFast Hong et al. (2025) introduces a two-stage memory mechanism, applying LoRA adapters Hu et al. (2022) a pretrained diffusion model for fast adaptation to recent sequences (fast learning), while maintaining a separate slower loop to consolidate changes for future predictions (slow learning). Unlike our approach, this method follows meta-learning-styll memory paradigm and requires additional training of the potentially large pretrained network. Other approaches implement memory by explicitly storing relevant past frames and learning how to retrieve and inject them into the model's context window Xiao et al. (2025). However, it introduces challenges related to increasing storage demands and retrieval complexity. Other lines of work utilize spatial memory models that retrieve relevant past context based on camera position Xiao et al. (2025 ; Yu et al. (2025); ; Wu et al. (025). And lastly approaches that implement recurrent networks such as state-space models into their world model [Savov et al. (2025): Po al. (2025).
3 MEMORY ADAPTATION
3.1 BACKGROUND
Denoising Diffusion Probabilistic Models Diffusion models Ho et al. (2020) aim learn data distribution ((xo) by introducing a sequence of ''' progressively noisier versions {xe [ of the data Xg. This is done via a forward Markov chain where each transition probability q(x, xx-1) N(xx; VI 3 x1-1, 3,I) is Gaussian with a predefined noise schedule {.3: ) where < 3r S 1. The generative process learns to reverse this diffusion, modeling po px(xx-1|x0 Ni t), 3,I), which approximates the reverse conditional (xx-1 /xx). The mean μ4 (x,, t)
is predicted using neural network that outputs noise estimate (x t). This model is trained by minimizing the simplified objective Ho et al. (2020)
[()) Ex E.x. ||ee eø (x 1)
3 which implicitly learns the score function of the data distribution Xx, lo VI-a. cp (x t). After training, sampling iterates Langevin-style updates
1
Be
X Var
xx
VT co (x), + om,, ~~N(0.I)
(2)
with Ox,, ãz defined by the noise schedule.
Compositional Diffusion via Product of Experts In many generative settings, is desirable to combine the knowledge or constraints encoded by multiple models. A principled way to achieve this is through the Product of Experts (PoE) framework [Hinoon| (2002), which defines composite distribution as the (normalized) product of individual distributions. Each model can be viewed as an expert px(x), and the query itself may define an additional constraint Pquery(x), The PoE then computes joint distribution p(x) X Pquery (x) II, Pt (x), sharply focusing on latent states consistent with all sources of evidence. This allows for selective, content-addresssalee generation, retrieving or synthesizing patterns that satisfy all contributing distributions. In the context of diffusion models, this corresponds to composing multiple pretrained models, po, (x+) = px(x,) into a unified generative
process Liu al. 2022): Du al. (2023).
3.2 COMPOSITION OF MEMORY EXPERTS
Our goal is to augment video world models with the ability to efficiently condition generation on specific past histories, while preserving the diversity and generalization of a pretrained prior. We
3

---

Published as a conference paper at ICLR 2026
denote a video sequence of T frames E RT* with spatial resolution H x W. The context M represents the set of all past frames, which could be of arbitrary length. In most modern video world models, T is fixed and limited by design Lin et al. (2024); Ma et al. (2024): Yang et al. (2024b); Polyak et al. (024)). Our objective is to model the conditional distribution pl x M) for coherent generation that reflects both recent and long-term context.
Different architectures offer trade-offs between consistency with the immediate past, long-term memory, generation quality, and inference speed. There no universally optimal solution. Scaling transformers with longer context windows quickly incurs quadratic attention costs and often leads to unstable training Song et al. (025).. Sparsified attention or purely state-space models reduce memory usage but lose fidelity and long-rang consisteney Gu & Dao (2024); Behrouz et al. (2024).
Motivated by this, we propose compose specialized experts, each responsible for a particular memory role, and fuse them probabilistically at sampling time. Diffusion models naturally support such composition through a PoE formulation, without requiring retraining Du et al. (2023). Concretely, we assume the past history M can be decomposed into multiple (possibly overlapping) subsets of frames {c: [= with e; C M and _ = M. We assume that a diffusion model px(x) has been trained to predict future frames conditioned on the corresponding context c.. To approximate p(x|MM, we then use
K
p(x M) = p(x . .........)).... ) x II (x(
which balances contributions from all experts.
Product of Contrastive Experts. The PoE framework amplifies regions where all experts assign high likelihood and suppresses others. While this sharpens consensus, in practice it can also lead to over-confident distributions and reduced diversity, particularly as the number of experts increases. To address this, we introduce a Product of Contrastive Experts (PoCM) formulation that explicitly factors out spurious distribution modes.
Formally, (x)) denotes set of heterogeneous experts. We define their composition as
K
p(x M)
II [x),, where px(x) pi(x)** 1=xx)== 1=1
and where a; > are contrasting coefficients, and (x) denotes the expert's unconditional baseline distribution (i.e., without conditioning the context cj), and p; the contrastive experts. The contrastive product is motivated by the fact that each expert is only an approximation of true show below, PoCE provides way suppress undesired spurious modes.
distribution, and,
Taming spurious modes. A first impulse suppress spurious modes of approximate expert is to temper it, i.e., use (x)" with a > 1. This indeed downweights lighter modes more than heavier ones. However, tempering alters not only the magnitude of the modes but also their spread: for standard kernels (e.g., Gaussians),
E)) x M(x u4 E.),
so the variance shrinks by a factor 1/a. Thus, naive tempering narrows kernels and changes local geometry, which reduces diversity when sampling from it.
Contrastive experts. To suppress spurious modes while preserving local shapes, we combine each conditional expert p; with its unconditional baseline P, via contrastive mixture: for a; > 1,
c (x) x (x) (x(x) 1-a.
This operation is log- linear (additive in log density), ), but, unlike pf. does not, in the KDE setting below, tighten kernels; it reweights the existing components and leaves their shapes intact.
Proposition 1 (KDE reweighting under disjoint support). Assume a KDE for (x(x) = EM TH (x(x) and (x() EM, wt he (x) with Ex के. o. 2 0. Ex wt 1, w 2 and hx dx==
4

---

Published as a conference paper at ICLR 2026
Suppose each kernel hy has limited bandwidth and the supports of {hk 1 disjoint. Then the contrastive expert satisfies
M (m1)) (w+) 1-o, (x).).
(x(
k=1
In particular, if if = M (uniform baseline), ), the factor (wi))=- is constant in and is absorbed into normalization, yielding p:(x) x (T()) ( (x).
Proof. On the support of hg, the disjointness assumption implies pr(x) = of (x)xx)) and px( x w (x). Hence pi(xx (x) xx (Tt) (rr)) (wif) 111-a-hh (x))))xtt-110 (mp ) (wr (1=a) (x). Summing over
k and renormalizing gives the claim.
Interpretation and trade-offs. By Proposition [ the transformation p; + p; leaves the kernels h& untouched and acts only the mixture weights, via oíद. + # x (mz))<(w()) (wk))_a, Hence lighter modes are disproportionately downweighted when a; > 1. This suppresses spurious modes without distorting local shapes. The downside is that genuinely valid but secondary modes (with smaller weights) are also reduced, which may lower diversity; og must therefore balance artifact removal diversity preservation. The KDE modeling assumptions p. and 7., are also valid in very general settings. One can always approximate both distributions with the same basis functions (non- overlapping kernels) but different coefficients. As further discussed in Sec. A the main requirement is instead that p; has larger magnitudes for the valid modes than for the spurious ones compared to Ps. which should be natural outcome of the training of these experts. Additional analysis and a simplified illustrative example are provided in Sec.
3.3 INSTANTIATING MULTIPLE MEMORY EXPERTS
We instantiate three complementary experts that address different aspects of temporal conditioning. Together, they form compositional memory system that balances fidelity to recent frames, long-term consistency, and spatial disambiguation. In practice, we build upon pretrained diffusion models that generate videos from small number of input frames. For first memory expert, which denote as po (x c), we use a standard image-to-video (12V) diffusion model. It excels at producing high-quality frames conditioned on the immediate past (in the range of 1-3 images).
Short-Terme Memory (STM). The STM expert is a diffusion model that directly attends to short-term recent context window esT (in the range of 10-100 images). This expert is effective at capturing local dynamics and fine details, though its temporal consistency is limited by the finite context size. We denote as (x | OST). implemented using an architecture with moderately extended context length, motivated by fast weight learning SSchlag et al. (202T). The contrastive expert is the unconditional model (x) Px 2²). The notation Ø indicates the absence of conditioning.
Long-Term Memory (LTM). While STM effectively models limited history, many applications require consistency across hundreds of frames. Scaling the context length of existing models quickly becomes computationally prohibitive, motivating a different strategy. We therefore introduce a separate diffusion model that be conditioned through two channels: 1) by stores episodic knowledge in its weights z) via test-time finetuning on the long-term history Cur (in the range of 100-1000 images) and 2) through the standard context conditioning. To clarify the notation, let Peter) (x| c) denote a base diffusion model conditioned on a context c and with weights conditioned on separate context c'. Using Ø as one of the contexts indicates that no conditioning has been provided or used to finetune the weights vb. We then introduce Peter)) (x c), where the parameters (C(LT) incorporate long-term context through finetuning cur using the diffusion loss (eq. (.). The contrastive expert is then the model without any context P.. (x) = Pe(e) (x Ø). where Pe(z) = Pa-
A key challenge is catastrophic forgetting in the adapted model. While explicit regularization strategies such as KL penalties between base and finetuned models are possible, we find that finetuning set of LoRA adapters Hu al. (2022) provides sufficient implicit regularization. This approach avoids overfitting to the past, reduces computational cost, and preserves the generalization capacity of the original model.
5

---

Published as a conference paper at ICLR 2026
Spatial Long-Term Memory (SLTM). LTM can be viewed associative memory that retrieves long-range knowledge conditioned on immediate visual cues c e(Schaeffer et al. (2024). However, when cues are ambiguous (e.g., looping paths or visually similar but distinct locations) additional spatial priors are essential.
To address this, extend the context M with auxiliary spatial signals S (e.g., camera poses point maps) extracted from the history M.. We define a spatial prior Px (x S), where Enc (M) encodes the memory into a spatial representation. Here, Encx(-) may correspond from-motion algorithm (e.g., SLAM [Smihh et al. (1866) or a learned encoder Wang et al. (202)). This spatial prior improves long-range consistency by disambiguating locations and reducing drift in repetitive or visually ambiguous environments. The contrastive expert, as in STM, is obtained by dropping the spatial prior: px(x (0))
structure-
Composition of Memory Experts. We now combine the individual experts into unified distribu- tion. Specializing eq. (3) to the STM ps. LTM pv. and SLTM px, together with the pretrained prior Po: we obtain
(C CST. CLT, S)
pe(( 12) 131-aopaxx (x c)
Po (x 1)1-a1pp(66 (CT) [xx(x 2) px (x ()) 8))0)
(x (x 1 p) 0))²
Equation (4) shows the product-of-expert formulation where each memory model is contrastive expert. The coefficients a; 2 balance unconditional and conditional contributions, ensuring that recent context, long-term consistency, and spatial priors integrated in principled manner. For the integration this approach into diffusion models, we refer to Sec.
4 EXPERIMENTS
In this section, we evaluate our proposed method through comparisons with baselines on synthetic and woal-wrldd datasets. We describe the datasets, baseline models, and evaluation metrics before presenting quantitative and qualitative results. Additional details, including implementation settings, computational cost analysis, extended ablation studies, and further visualizations, are provided in the appendix.
Datasets. We assess Composition of Memory Experts (CoME) on synthetic and real-world datasets. Memory Maze Pasukonis et al. 2022 consists of 30k offline trajectories (lk frames each) of agents navigating 3D mazes, with absolute and relative position signals, making suitable for evaluating long-term memory learning. RE/OK Zhou et al. (2018) provides indoor scenes with camera pose annotations and serves as challenging real-world benchmark. RECON [Shah et al. (202T) contains over 5k trajectories of real-world outdoor navigation and provides absolute camera trajectories, making it suitable for navigation planning tasks.
Evaluation Metries. For high-fidelity prediction tasks (i.e., low uncertainty), we report frame-wise Learned Perceptual Image Patch Similarity (LPIPS) Zhang et al. (2018). Structural Similarity Index (SSIM), and Peak Signal-to-Noise Ratio (PSNR). For navigation planning, evaluate trajectory alignment using Absolute Trajectory Error (ATH) and Relative Pose Error (RPE), where ATE measures the global deviation of the predicted trajectory from ground truth and RPE captures the local consistency of relative motion.
4.1 CONSISTENCY WITH PAST OBSERVED FRAMES
Baseline Models. Our primary baseline is Diffusion Transformer (DiT) Peebles also following prior work Chen et al.) (2024); Song et al. (025)). The model uses three context frames with relative positional information to predict the next seventeen frames.
(Dellle & Xie (202hh),
We further evaluate specialized memory architectures. The short-term memory (STM) expert is DiT(-S) with sliding window attention chunk size of 20, with full attention layers at the and layers, taking 33 frames conditioning and predicting the next 17 frames. The long-tem tnn-term memory (LTM) module adapts a pretrained DiT with LoRA finetuning. The spatial long-temn memory (SLTM) variant. doubles the patch size in the input projection and conditions absolute camera poses. To

---

Published as a conference paper at ICLR 2026
Table 1: Comparison of different memory con- figurations on reconstruction metrics with two steps per frame. Best values highlighted
enter on
0.17
50M
0.36
0.55
Method LPIPS 1
SSIM
PSNR
Base
0.209
0.771
19.16
0.3
STM
0.156
0.820
21.29
0.2
LTM
0.171
0.805
19.98
031
SLTM
0,150
0.833
20.65
Context Frames
STM+LTM
0.114
0.862
22.32
Figure 1: (LPIPS 1) as function of context
CoME
0.097
0.892
23.07
length on the Memory Maze dataset. Longer contexts lead to more faithful reconstructions, with no saturation observed up to 480 frames.
Sliding
0.183
0.753
19.02
SSM
0.158
0.828
20.62
Full
0.113
0.859
22.78
align relative trajectories with absolute coordinates, we introduce an auxiliary encoder that maps relative positions and current clean frame estimates into absolute camera positions.
Although our approach in principle allows training- free composition of different pretrained models, still compare with specialized architectures that represent memory-oriented alternatives: (i) full-attention transformer as the absolute baseline, (ii) chunked transformer inference with interleaved Mamba blocks Gu & Dao (2024) for SSM-style memory, and (iii) sliding-window attention with interleaved full-attention blocks. All baselines have comparable parameter budget and are trained for 150k steps.
We report quantitative results across SSIM, PSNR, and LPIPS (Table D on Memory Maze. Both LTM and SLTM are finetuned with two gradient steps per context frame. Each additional memory component consistently improves temporal consistency over the base model. In particular, augment- ing with long-term memory, and especially combining STM and LTM, yields significant gains in consistency, beating all provided baselines.
Training cost is also noteworthy. Full-attention training required roughly 60x the compute of STM and 12x that of the base model, with also more than 3x the amount of training steps to converge with the same base settings, see Section E.T Moreover, full attention with the standard diffusion regime was unstable, requiring modified training scheme for convergence SereenceeSnggaaaaa et al. (2025). In contrast, our memory composition approach achieves comparable or better consistency at a fraction of the training and inference cost, demonstrating the efficiency of leveraging specialized memory experts over brute-force scaling.
Effect of Context Length. We further investigate the effectiveness of LTM by varying the number of context frames used for adaptation, ranging from 50 to 480 with 3 update steps per frame. Increasing the context length consistently improves perceptual quality, as measured by LPIPS (Figure Longer contexts allow the model to leverage more temporal information and generate reconstructions that are more faithful the ground truth. Notably, consistency continues to improve up 500 frames, with no saturation observed (limited by dataset length). These findings highlight the complementary roles of STM and LTM in balancing short- and long-range temporal dependencies.
Effect on Planning. Consistency is particularly critical for planning tasks. As testbed for visual planning, we adopt the Navigation WM framework [Bar al.) (2024), where the model applied to goal-condee goal-conditioned navigation. Given past observations and goal image, candidate sequences of trajectories are scored by executing NWM rollouts of length eight and computing the I.PIS distance between the final frame and the goal image, averaged over three runs. Our STM is implemented as a finetuned DiT-B on twelve context frames, while the LTM is instantiated by adapting a frozen pretrained world model with LoRA finetuning.
As shown in Table 2, adding memory components consistently improves planning accuracy. Even lightweight additions such smaller STM yield tangible gains, supporting the principle that structured memory enhances temporal consistency and thereby benefits downstream tasks such as navigation.
7

---

Published as a conference paper at ICLR 2026
Table 2: Planning accuracy on the RECON benchmark (100 sampled trajectories). We We report Absolute Trajectory Error (ATE) and Relative Pose Error (RPE). Comparison of CoME with GNM Shah et al. 2022). NOMAD Sridhar et (2023) and NWM Bar al. 2024).
GNM
NOMAD
NWM
CoME
STM
LTM
NWM+STM
NWM+LTM
ATE (L)
1.87
1.93
1.13
0.96
1.05
1.10
0.98
1.07
RPE (1)
0.73
0.52
0.35
0.28
0.32
0.32
0.30
0.33
Table 3:
Comparison of evaluation metrics across different methods and rollout lengths on RealEstate10K. PSNR and SSIM (1) indi- cate higher is better; LPIPS (1) lower is better. Highlighting marks the best results.
0.22
Dest
Base
0.20
0.18
0.16
0.14
Method
LPIPS 4
PSNR 1
SSIM t
0.12
Base
0,405
19.8
0.794
611
CoME
0.359
21.3
0.83
Chunk
HG-v HG-t
0,414
19.2
0.764
0.400
19.7
0.788
Figure 2: Evaluation on Memory Maze with 10 chunks of 20 frames each. Our method incre- mentally memorizes to maintain consistency.
4.2 CONSISTENCY WITH A CONTINUOUS STREAM OF OBSERVATIONS
Streaming Evaluation. In reinforcement learning settings. agents receive continuous stream of observations and must remain consistent with them over time. To evaluate whether our memory composition extends beyond fixed contexts at the beginning of generation, experiment: at each step, the model memorizes the last 50 frames for 25 gradient updates, predicts the next frames, and then appends the ground-truth frames back into the context. This procedure is repeated for 11 iterations. As shown in Figure 2 on Memory Maze the STM already improves consistency across iterations by reinforeing short-term dynamics. Adding the LTM yields com- pounding effect: consistency continues to improve new observations are integrated, demonstrating that the LTM successfully stores additional information in its weights and leverages it in subsequent generations. This highlights the importance of online memory adaptation for handling continuous observation streams.
design the following
Recall Ability. We evaluate the recall capabilities of CoME on the RE10K dataset using large pretrained network Song et al. (2025). In this experiment, the model generates frames and memorizes them, and then follows the reversed trajectory to reconstruct previously visited states. We test sequences of forward rollouts, each followed by backward rollouts. To measure recall accuracy, compare the frames of the backward rollout with the forward rollout using LPIPS for perceptual similarity. As reported Table 3 CoME consistently outperforms baseline methods semantic recall. In Figure 3 we provide qualitative results, where we see that CoME accurately is able to retrieve the exact frames, preserving a consistent scene view. In contrast, the baseline model without memory generates plausible but inconsistent frames, failing to recall the original state.
Notes on Compute. In the online setting, we perform two memorization steps per frame, resulting in approximately 100 steps total. When using the LTM with LoRA at rank r 64, this introduces about 4x compute overhead due to memorization, in addition to the overhead from the added experts. With r 16, the overhead decreases roughly 2x increase in sampling time. Further details are provided in Section
4.3 ABLATION
Contrastive Expert Composition. We first ablate the role of the contrastive formulation Table 4 Applying the contrastive PoE approach improves the performance of each individual memory expert
8

---

Published as a conference paper at ICLR 2026
(a)
(b)
Figure 3: Qualitative results the RealEstate108 dataset. We generate six forward rollouts and then reverse the camera trajectory. Unlike the base model without long-term memory, Composition of Memory Experts correctly recalls the initial frame in both examples, ensuring scene consistency. (The first frame and last frame of a sequence should be equal.)
Table 4: LPIPS results with and without the ad- dition of the contrastive experts.
Table 5: Effect of LTM adaptation capacity and context length given in LPIPS.
Model
w/o Contrastive
w/ Contrastive
Context Frames
Base
0.203
0.200
Rank
50
150
450
Params
STM
0.175
0.156
0.193
0.161
0.146
LTM
0.188
0.171
32
0.175
0.169
0.128
4%
SLTM
0.178
0.150
64
0.161
0.158
0.124
7%
STM+LTM
0.170
0.114
256
0.162
0.142
0.118
25%
Full
0.221
0.188
0.125
100%
All
0.192
0.097
with the Memory Maze setting from Section 4 We can see that the naive product-of-experts quickly collapses or cancels out consistency gains, whereas our contrastive composition is essential for stable improvements. Without contrastive weighting, stacking multiple experts fails to outperform the base model and in some cases even reduces consistency.
LTM Architecture and Context Frames. We next analyze the effect of LTM adaptation capacity and context length in Table Increasing LoRA rank consistently improves performance, while full fine-tuning tends to overfit unless the available context is highly diverse. This suggests that LoRA provides implicit regularization, enabling sufficient adaptation while retaining general priors. Longer contexts further amplify these benefits, with no saturation observed even at 450 frames, underscoring the scalability our long-term memory design.
CONCLUSIONS
In this work, we explored how different memory mechanisms can be combined to improve the consis- tency and reliability of video world models. First, we demonstrated that models with complementary strengths in memory modeling can be effectively combined to yield more consistent generations. We showed analytically and experimentally that our proposed Product of Contrastive Experts can successfully eliminate spurious modes without introducing side-effects in naive Product of Experts formulation. We further introduced novel long-term memory module that allows to store efficiently information across more than 500 frames. enabling substantially improved temporal consistency. These contributions were validated across both synthetic and real-world datasets, as well in reinforcement learning navigation planning tasks, where reliable memory is crucial for

---

Published as a conference paper at ICLR 2026
effective decision decision-making. Looking ahead, future work will explore extending long-term memory mechanisms to explicitly capture temporal dependencies, incorporating richer conditioning strategies, and developing more exhaustive forms of spatial conditioning to further enhance generative fidelity and downstream task performance.
ACKNOWLEDGEMENTS
This work supported by grant from the Swiss National Supercomputing Centre (CSCS) under project ID a144 on Alps as part of the Swiss Al Initiative, and by SNSF Grant 10001278. Additional computations carried on UBELIX (hnps://www.dduniiecchhhh) the high-performance computing cluster at the University of Bern.
REFERENCES
Anurag Ajay, Seungwook Han, Yilun Du, Shuang Li, Abhi Gupta, Tommi Jaakkola, Joshua B. Tenenbaum, Leslie Pack Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional founda- tion models for hierarchical planning. NeurlPS, 2023.
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos J. Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter atari. In NeurlPS, 2024. URL
nips ec httppppprr...llssppppr/////0//0p 2024/hash/ -
Amir Bar, Gaoyue Zhou, Danny Tran, Trevor Darrell, and Yann LeCun. Navigation world models. 2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15791- 15801, 2024.
Omer Bar-Tal, Hila Chefer, Omer Tov, Charles Herrmann, Roni Paiss, Shiran Zada, Ariel Ephrat, Junhwa Hur, Yuanzhen Li, Tomer Michaeli, Oliver Wang, Deqing Sun, Tali Dekel, and Inbar Mosseri. Lumiere: A space-time diffusion model for video generation. CoRR. abs/2401.12945 2024. URLpttps: http://oooo.o.o./////oo//o/o 48550/ ar..50/arxiv... 2401 12945
Charles Beattie, Joel Z. Leibo, Denis Teplyashin, Tom Ward, Mareus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, Julian Schrittwieser, Keith Anderson, Sarah York, Max Cant, Adam Cain, Adrian Bolton, Stephen Gaffney, Helen King, Demis Hassabis, Shane Legg, and Stig Petersen. Deepmind lab. CoRR, abs/1612.03801, 2016. URL http:
//axxiv.org/abs//1
03801
Ali Behrouz, Peilin Zhong, and Vahab S. Mirrokni. Titans: Learning memorize at test time. ArXiv, abs/2501.00663 2024.
Arwen Bradley and Preetum Nakkiran. Classifier-free guidance is a predictor-correctt ArXiv, abs/2408.09900) 2024.
Jake Bruce, Michael Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, Sarah Bechtle, Feryal M. P. Behbahani, Stephanic Chan, Nicolas Manfred Otto Heess, Lucy Gonzalez, Simon Osindero. Sherjil Ozair, Scott Reed, Jingwei Zhang, Konrad Zolna. Jeff Clune, Nando de Freitas, Satinder Singh, and Tim Rocktaschel. Genie: Generative interactive environments. ArXiv, abs/2402.15331. 2024.
Boyuan Chen, Diego Martí Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets full-sequence diffusion. In The eight Annual Conference on Neural Information Processing Systems, 2024. URL
Aram Davtyan, Sepehr Sameni, and Paolo Favaro. Efficient video prediction via sparsely conditioned
flow matching. In ICCV, pp.
2023. URL
TCCV510777 2023. 02126
10

---

Published as a conference paper at ICLR 2026
Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: RNNs, transformers, and s4. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: http:s//ppenreeeee...o net.///wweeeeeeeeeeeeeeeeeeeeeee how..foooummii-00XXXXX..
Yilun Du, Conor Durkan, Robin Strudel, Joshua B. Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Narain Sohl-Dickstein, A. Doucet, and Will Grathwohl. Reduce, reuse, recycle: Composi- tional generation with mergy-based diffusion models and meme. In International Conference on Machine Learning, 2023.
Yunhai Feng, Nicklas Hansen, Ziyan Xiong, Chandramouli Rajagopalan, and Xiaolong Wang. Finetuning offline world models in the real world. In CoRL, pp. 425-445, 2023. URL https www.e //proceedings. mlr press /tcees/vv2290 v229 (wt.222/eng22a.mmmi html
Michael Fuest, Vincent Tao Hu, and Bjor Ommer. Maskflow: Discrete flows for flexible and efficient long video generation. ArXiv, abs/2502. 1234, 2025.
Songwei Ge, Thomas Hayes, Harry Yang. Xi Yin, Guan Pang, David Jacobs, Jia-Bin Huang, and Devi Parikh. Long video generation with time-agnostic vqgan and time-sensitive transformer. In ECCV (17), pp. 102-118, 2022. URL http::ttp///// //cci. wtt.///....00011...00001000
0---031-997 1_7
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference Language Modeling, 2024.
EEYskw1VE VY2
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. CoRR. abs/1 1899009.0110 0199, 2018a.
01999
David R Ha and Jürgen Schmidhuber. World models. ArXiv, abs/1 10122, 2018b.
Danijar Hafner, Timothy P. Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control:
Learning behaviors by latent imagination. CoRR, abs/1912.01603, 2019. URL httt:://axxxi
org/abs/ /192. 912. 01 61633
Mariam Hassan, Sebastian Stapf, Ahmad Rahimi, Pedro M. B. Rezende, Yasaman Haghighi. David Brüggemann, Isinsu Katircioglu, Lin Zhang, Xiaoran Chen, Suman Saha, Marco Cannici, Elie Aljalbout, Botao Ye, Xi Wang, Aram Davtyan, Mathieu Salzmann, Davide Scaramuzza, Marc Pollefeys, Paolo Favaro, and Alexandre Alahi. Gem: generalizable ego-vision multimodal world model for fine-gra ine-grained ego-motion, object dynamics, and scene composition control. CoRR, 11198
abs/2412. 111998 2024. URL httph:: http://iioor/o/...././oo/ /10.
08.111...855000 48550
Yingqing He, Tianyu Yang, Yong Zhang, Ying Shan, and Qifeng Chen. Latent video diffusion models for high-fidelity video generation with arbitrary lengths. CoRR, abs/2211.132211, 2022. URL
https://doi..om 0ttp://11...1/0//
48550/550//00 arwiv..ivvv.m 2211 13221
Roberto Henschel, Levon Khachatryan, Daniil Hayrapetyan, Hayk Poghosyan. Vahram Tadevosyan, Zhangyang Wang, Shant Navasardyan, and Humphrey Shi. Streamingt2v: Consistent, dynamic, and extendable long video generation from text. CoRR, Abs/2403.144733 2024. htt.pt
//doi. 48550/arXiv 2403.1477
Geoffrey E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14:1771-1800, 2002.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusson probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN
Yining Hong, Beide Liu, Maxine Wu, Yuanhao Zhai, Kai-Wei Wei Chang, Linjie Li, Kevin Lin, Chung- Ching Lin, Jianfeng Wang, Zhengyuan Yang, Ying Nian Wu, and Lijuan Wang. Slowfast- VGen: Slow-fast learning for action-driven long video generation. In The Thirteenth International Conference on Learning Representations, 2025. /ttpps://ppeeeeeeeww http:/wwwwweee..tteeeeeee
id=01.8b54P96c
11

---

Published as a conference paper at ICLR 2026
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. CoRR, abs/2301.11099, 2023. URL htt.attst///c..o/////00
arXiv. 2301 11093
Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shotton, and Gianluca Corrado. Gaia-l: A generative world model for autonomous driving, ArXiv, abs/2309.17080, 2023.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL htRp:tt::: wwwpeeoeeeeeeeeeeeeeeeeeeeee..ee net/ forum?
id=nZevKeeFYf9
Bin Lin, Yunyang Ge, Xinhua Cheng, Zongjaan Li, Bin Zhu, Shaodong Wang, Xianyi He, Yang Ye, Shenghai Yuan, Liuhan Chen, Tanghui Jia, Junwu Zhang, Zhenyu Tang, Yatian Pang, Bin She, Cen Yan, Zhiheng Hu, Xiaoyi Dong, Lin Chen, Zhang Pan, Xing Zhou, Shaoling Dong, Yonghong Tian, and Yuan. Open-sora plan: Open Open-source large video generation model. CoRR, abs/2412.00131,, 2024. URL https: / doi 100 / arXiv. 2412 00131
Bangbang Liu, Xinfeng Li, Jiayi Zhang, Jinlin Wang, Tanjin He, Sirui Hong, Hongzhang Liu, Shaokun Zhang, Kaitao Song, Kunlun Zhu, Yuheng Cheng, Suyuchen Wang. Xiaoqiang Wang, Yuyu Luo, Haibo Jin, Peiyan Zhang, Ollie Liu, Jiaqi Chen, Huan-Rui Zhang, Zhaoyang Yu, Hao Shi, Boyan Li. De Ming Wu, Fengwei Teng, Xiao Jia, Jiawei Xu, Jinyu Xiang. Yizhang Lin, Tianming Liu, Tongliang Liu, Yu Su, Huan Sun, Glen Berseth, Jianyun Nie, Ian T. Foster, Logan T. Ward, Qingyun Wu, Yu Gu, Mingchen Zhuge, Xiangru Tang, Haohan Wang, Jiaxuan You, Chi Wang, Jian Pei, Qiang Yang, Xiaoliang Qi, and Chenglin Wu. Advances and challenges intelligence to evolutionary, collaborative, and safe
in foundation agents: From
systems, ArXiv, 0bs/2550400199) 2025.
Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B. Tenenbaum. Compositional visual generation with composable diffusion models. CoRR, abs/2206.01714, 2022. URL https
//doi.org/11 10. 100.0.8855500 48550/8555//0 arviv http:/aaxxvv.......r.. 01714)
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL htl.:ttps:
/ttl.tts////ppeeeiiiie /wpp///ppnrvvvveeeeeeeeeeee net ftr./formmm di
Bkg6RiCqY7
Yu Lu, Yuanzhi Liang, Linchao Zhu, and Yi Yang. Freelong: Training-free long video generation with spectralblend temporal attention. In The Thihty- Annual Conference on Neural Information Processing Systems, 2024. URL
htt.:eewwwwwwweeeeeeeeeeeeeeeeee wwr.:foooo forum d= ga50005500 5200v
Xin Ma, Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. CoRR, abs/2401 1.00048 2024. URL
wwwooogg...... 10. 48550/ arxiv 2401.
Jurgis Pasukonis, Timothy P. Lillicrap, and Danijar Hafner. Evaluating long- term memory in 3d mazes. CoRR, 0tss/2110.13882 2022. htt.attp:e//da...r//10044445
13:3
William Peebles and Saining Xie. Scalable diffusion models with transformers. CoRR, abs/2212.007448, 2022a.
2212.09745
William S. Peebles and Saining Xie. Scalable diffusion models with transformers. 2023 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 4182-4182, 2022b.
Ryan Po, Yotam Nitzan, Richard Zhang, Berlin Chen, Tri Dao, Eli Shechtman, Gordon Wetzstein, and Xun Huang. Long-context state-space video world models. ArXiv, abs/2505.20171, 2025.
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih- Yao Ma, Ching Yao Chuang, David Yan, Dhruv Choudhary, Dingkang Wang, Geet Sethi, Guan Pang, Haoyu Ma, Ishan Misra, Ji Hou, Jialiang Wang, Ki ran Jagadeesh, Kunpeng Li, Luxin Zhang, Mannat Singh, Mary Williamson, Matt Le, Matthew Yu, Mitesh Kumar Singh, Peizhao Zhang, Peter Vajda, Quentin Duval, Rohit Girdhar, Roshan Sumbaly, Sai Saketh
12

---

Published as a conference paper at ICLR 2026
Rambhatla, Sam S. Tsai, Samaneh Azadi, Samyak Datta, Sanyuan Chen, Sean Bell, Sharadh Ramaswamy, Shelly Sheynin, Siddharth Bhattacharya, Simran Motwani, Tao Xu, Tianhe Li, Tingbo Hou, Wei-Ning Hsu, Xi Yin, Xiaoliang Dai, Yaniv Taigman, Yaqiao Luo, Yen-Cheng Liu, Yi-Chiao Wu, Yue Zhao, Yuval Kirstain, Zecheng He, Zijian He, Albert Pumarola, Ali K. Thabet, Artsiom Sanakoyeu, Arun Mallya, Baishan Guo, Boris Araya, Breena Kerr, Carleigh Wood, Ce Liu, Cen Peng, Dimitry Vengertsev, Edgar Schonfeld, Elliot Blanchard, Felix Juefei-Xu, Fraylie Nord, Jeff Liang, John Hoffman, Jonas Kohler, Kaolin Fire, Karthik Sivakumar, Lawrence Chen, Licheng Yu, Luya Gao, Markos Georgopoulos, Rashel Moritz, Sara K. Sampson, Shikai Li, Simone Parmeggiani, Steve Fine, Tara Fowler, Vladan Petrovic, and Yuming Du. Movie gen: A cast of media foundation models. ArXiv, abs/2410. 0att22400.13700, 13720, 2024.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. CoRR, 12./2112.10552, 2021. URL
/wwp///a////g httpss:/aaxii>: org/ abs /2tp/2112... 10752
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar, Master- ing memory tasks with world models. In The Twelfth International Conference Learning Representations, 2024. URL penttts::/openreeviii
net/
Nedko Savov, Naser Kazemi, Deheng Zhang, Danda Pani Paudel, Xi Wang, and Luc Van Gool. Statespacediffuser: Bringing long context to diffusion world models. CoRR, abs/2505.22246, May 2025. URL http::tts:://ott. http://i..oorrg////// 48550/ 20r/5550//iii.... arxiv 2005.220223
Rylan Schaeffer, Nika Zahedi, Mikail Khona, Dhruv Pai, Sang T. Truong, Yilun Du, Mitchell Ostrow, Sarthak Chandra, Andres Carranza, Ila Rani Fiete, Andrey Gromov, and Sanmi Koyejo, Bridging associative memory and probabilistic modeling. CoRR, abs/2402.10202, 2024. URL
https: /trggg/1 10 48550/ arwiv. 2402 10202
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 Proceedings of Machine Learning Researck, pp.
9355-9366 PMLR, 18-24 Jul 2021. URL https: httpps://procedinggg mlr.. wwt..................... schlag2la.htm.
Dhruv Shah, Benjamin Eysenbach, Nicholas Rhinehart, and Sergey Levine. Rapid exploration for open-world navigation with latent goal models. In CoRL., pp. 674-684, 2021. URL https
//proceeding ing. ml www.ggg...............m.mm.emm.. press / v1 /1/1144///aahaaa shah22a html
Dhruv Shah, Ajay Kumar Sridhar, Arjun Bhorkar, Noriaki Hirose, and Sergey Levine. Gnm: A general navigation model to drive any robot. 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 7226-7233, 2022.
Pete Shinners and the Pygame community. Pygame -python game development. pygame . org/ 2000-. Version X.Y.Z.
https:/ //bw::///wwwwww// /www
Randall Smith, Matthew Self, and Peter C. Cheeseman. Estimating uncertain spatial relationships in robotics. Proceedings. 1987 IEEE International Conference on Robotics and Automation, 4: 850-850, 1986.
Kiwhan Song, Boyuan Chen, Max Simchowitz, Yilun Du, Russ Tedrake, and Vincent Sitzmann. History-guided video diffusion. ArXiv, abs/2502.06664, 2025.
Ajay Kumar Sridhar, Dhruv Shah, Catherine Glossop, and Sergey Levine. Nomad: diffusion policies for navigation and exploration. 2024 IEEE International Conference on Roboties and Automation (ICRA), pp. 63-70, 2023.
Goal masked
Chengkun Sun, Jinqian Pan, Russell Terry, Jiang Bian, and Jie Xu. Bgdb: Bernoulli-gaussian decision block with improved denoising diffusion probabilistic models. CoRR, abs/2409. 2024. URL https://doi.oze
48550/ 10550/a00iii...// 2409. 13116
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, CoRR, abs/2009.06732. 2020. URL /ttp:///aaiiviii http:/...s///aat..ooo//a////////
wwwpaawwwwwwww./////////// 06732
13

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pp. 100006010, 2017. URL httt:://papae nips htce/paper/1155 wtw.:iiaatteiiooooo --
Jianyuan Wang, Minghao Chen, Nikita Karaev, Andrea Vedaldi, Christian Rupprecht, and David Novotný. Vggt Vggt: Visual geometry grounded transformer. 2025 IEEF/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 5294-5306, 2025.
Tong Wu, Shuai Yang, Ryan Po, Yinghao Xu, Ziwei Liu, Dahua Lin, and Gordon Wetzstein. Video world models with long-term spatial memory. ArXiv, abs/2506.05284, 2025.
Zeqi Xiao, Yushi Lan, Yifan Zhou, Wenqi Ouyang, Shuai Yang, Yanhong Zeng, and Xingang Pan. Worldmem: Long-term consistent world simulation with memory. 2025.
Wilson Yan, Danijar Hafner, Stephen James, and Pieter Abbeel. Temporally consistent transformers for video generation. In ICML, pp. 19962-39998, 2023. URL SRRLttts/////ceeediiii mlr
Dress/v202/yanee
html
Sherry Yang, Yilun Du Bo Dai, Dale Schuurmans, Joshua Tenenbaum, and Picter Abbeel. Probabilistic adaptation of black-box text-to-video models. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL https: http://s///peneevee.... net wwwntee/ffouurrr forum?id=
PitIEgSCE3
Zhuoyi Yang, Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng. Da Yin, Xiaotao Gu, Yuxuan Zhang, Weihan Wang, Yean Cheng, Ting Liu, Bin Xu, Yuxiao Dong. and Jie Tang. Cogvideox: Text-to-video diffusion models with expert transformer. ArXiv, 0bs/2408.060177 2024b.
Shengming Yin, Chenfei Wu, Huan Yang, Jianfeng Wang, Xiaodong Wang, Minheng Ni, Zhengyuan Yang, Linjie Li, Shuguang Liu, Fan Yang. Jianlong Fu, Ming Gong, Lijuan Wang, Zicheng Liu, Houqiang Li, and Nan Duan. Nuwa-xl: Diffusion over diffusion for extremely long video generation. In ACL (1), pp. 1309-1320, 2023. URL /ttps:://oi..rrgg /t5.//.00..rr1111.111333/// 186533vv1 2023. acl- 73
Jiwen Yu, Jianhong Bai, Yiran Qin, Quande Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Xihui Liu. Context as memory: Scene-consistent interactive long video generation with memory retrieval. CoRR, abs/2506.03141, June 2025. URL
48550/ 2506
03141
Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features perceptual metric. 2018 IEEE/CVF Conference Computer Vision and Putlern Recognition, pp. 586-595, 2018.
Tinghui Zhou, Richard Tucker, John Flynn, Graham Fyffe, and Noah Snavely. Stereo magnification: learning view synthesis using multiplane images. ACM Trans. Graph., 37(4):65, 2018. URL https://doi org/10 11.....11311919 3201323
A FURTHER ANALYSIS FOR PRODUCT OF CONTRASTIVE EXPERTS
Non-Uniform Weights and Mode Selectivity. If the weights wf are not uniform, Proposition yields
# x (IT) (wf) (wf) a
where सं. are the KDE weights for gi. iee,, px(x) = DM (x). Let Ps - mL/wf. denote the contrast ratio of conditional to baseline weights. Then /// x Ph which is (strictly) increasing in Pk for a; 1. Therefore: (i) modes with Pk 1 (more emphasized by the conditional than the baseline) are amplified relative to baseline, (ii) modes with Pk < are suppressed, and the degree of separation grows with Os.
14

---

which balances experts and Published as a conference paper at ICLR 2026
Experts
Cantrastive
Composttons (normalized integral)
1.0
p10x)
610x
P_
p26x)
6201
Ruk
p304
0300
MTr_kh (Lknn(
density Gcaa 00
678
p4ix)
Cerstyy (Sce c0)
donsity
Figure 4: Illustration of Mixture of Contrastive Experts. (a) Individual experts, modeled as Gaussian mixtures, modes are decreasing geometrically (from left to right). (b) Individual contrastive experts, with uniform modes. (c) Product of Experts, Exponentially scales PoE, Product of Contrastive Experts. PoCE suppresses inconsistent modes (e.g., the four rightmost peaks) while preserving the dominant left kernel. The vertical line indicates the center of probability mass for the PoE and PoCE.
Spurious vs. Secondary Modes. To reliably remove spurious modes while retaining secondary but genuine modes, one needs margin in the contrast ratios:
3>>1 s.t. Pk 2 T for genuine modes and Pk STT for spurious modes.
Under such gap, choosing Oş that T is large cleanly separates the two groups after renormaliza- tion. secondary modes have AI only slightly above (or even below), aggressive a, will attenuate them too; in that case one (a) keep as close reflect "background" mass (e.g., flatten clearly spurious regions).
(b) baseline model with w's that better
Toy Example and Visualization To illustrate the effect of the Mixture of Contrastive Experts, construct a toy example where each expert (I) is modeled as mixture of M Gaussians,
(x) Wk.m N(x ( po kmm))
(5)
with weights chosen either to decay geometrically (emphasizing a few dominant modes) or to equalize peak heights across all components. The latter case defines the contrastive distributions (x))
We then study their interaction by composing multiple experts using different rules:
the product of experts (PoE) II Pk(x), which amplifies consensus and yields various spurious likelihood modes;
product of contrastive experts (PoCE)
their contrastive countcrparts.
All products are renormalized to define valid probability densities. The results highlight that stan- dard PoE compositions concentrate probability mass around shared modes, while contrastive and interpolated forms provide control over the degree of sharpening and mitigate trivial collapse.
As shown in Figure the standard PoE yields nearly equal likelihood for the two leftmost modes, centering the probability mass close to the second peak. In contrast, the contrastive composition removes the four right-hand modes almost entirely, while maintaining the structure left kernel. This demonstrates how contrastive weighting prevents spurious consensus and yields sharper, more interpretable mixtures.
dominant
B SAMPLING AND TEST- TIME TRAINING
Sampling. Given the compositional formulation in eq. 44) sampling proceeds by replacing the
single-model score function the reverse diffusion update (eq. with the joint score of all experts. At noise level 1, compute
V x, log PCoM (x:) = Σ an Vx, log (x.) + (1 v log (x))
(6)
15

---

Published as a conference paper at ICLR 2026
where the unconditional counterpart of each expert acts contrastive baseline. This additive score naturally integrates into the denoising step (eq. (2) by substituting cø with the weighted combination of expert noise predictors.
Time Adaptation of LTM. Among the experts, the Long-Term Memory (LTM) uniquely adapts online to the current episode. At the beginning of rollout, the base diffusion prior (e( is finetuned on the long- gingteemm context C.T. yielding Pe(er)) We perform Tomem gradient steps on mini-batches sampled from the memory M using the diffusion loss (eq. (.). To mitigate catastrophic forgetting, updates are restricted to low-rank adapters (LoRA layers), while the backbone remains frozen. This strategy injects episodic knowledge into the weights at low computational cost without degrading the pretrained prior.
After adaptation, sampling resumes with the composite score (eq. 66 Newly observed clips are appended to the memory M,, and the LTM may be updated again for longer rollouts. This establishes a recurrent loop: the STM ensures short-term fidelity, the LTM maintains evolving long-term consistency, and the spatial prior resolves global ambiguities. Together, these components yield stable online generation process grounded in both immediate dynamics and extended history.
C DATASETS
Memory Maze The Memory Maze environment Pasukonis et al. (2022) provides trajectories of an agent navigating through a 3D maze. However, also provides both absolute and relative positional information of the agent within the maze. Designed for reinforcement learning, the agent's weents task is to navigate toward colored balls placed in the maze, with the correct color cue provided via a colored border around the frame. Relative position is represented by 2D displacement vectors and two rotation angles, while absolute position includes 2D coordinates and an orientation angle (defined with respect to the maze center). The environment supports four maze sizes. In our experiments, we use an offline dataset comprising 30,000 trajectories of 1,000 frames each, collected using a stochastic navigation policy targeting random locations under action noise. This policy ensures diverse trajectory paths that often form loops, facilitating long-term memory learning.
The RECON Shah al. (202I) dataset contains over 5k trajectories of real-worke outdoor navigation with Clearpath Jackal robot, providing absolute camera trajectories along with RGB, stereo, thermal, LiDAR, GPS, and IMU data, making a comprehensive benchmark for studying latent goal models and topological memory in navigation planning.
RECON
RealEstate10K RealEstate 10K Zhou et al. (2018) is a video dataset captured in real-world real estate environments, annotated with accurate 3D camera poses and trajectories. This rich spatial information supports training video models with explicit control over camera motion and scene geometry. The dataset is partieularly valuable for memory and consistency evaluations, it allows testing whether models can retain and recall visual information about different rooms or viewpoints within the same property. Unlike synthetic datasets, RealEstate 10K provides real-world complexity and diversity, helping to validate model generalization to non- non-simulated settings. In our experiments, we use videos resized 256 256.
DeepMind Lab (DMLab-40K) The DMLab-40K dataset consists of 64 x 64 resolution videos depicting agent trajectories in a 3D maze environment Yan et al.) (2023); Beattie et al. (2016). The DeepMind Lab simulator procedurally generates random mazes with varied floor and wall textures. Each trajectory is 300 frames long and involves an agent traversing 7 x 7 mazes by selecting random target points and navigating to them using the shortest path. In total, 40,000 such action-conditioneed videos are provided. The maze layout is static throughout each episode, making the environment particularly suitable for evaluating recall capabilities.
Minecraft-2009 Minecraft-200k Yan et al. (2023) is large-scale dataset comprising 200,000 videos of Minecraft gameplay. where the player navigates using three discrete actions: forword left. and right. Each video is 300 frames long at resolution of 128 128 pixels, with corresponding action labels for each frame. The gameplay occurs Minecraft's 's marsh biome, a procedurally generated 3D world featuring complex terrain such hills, forests, rivers, and lakes. The player
16

---

Published as a conference paper at ICLR 2026
alternates between walking forward for random number of steps and rotating randomly, causing parts of the environment to disappear and reappear in the field of view. This dynamic makes the dataset particularly well-suited for evaluating memory and temporal consistency, especially under changing viewpoints. Following the protocol of Song et al. (2025), we upsample the original frames from 128 x 128 to 256 x 256 resolution enable higher-quality video generation.
Memory Cards Dataset To provide simplified and interpretable testbed for discrete recall, introduce the Memory Cards dataset. Each sample is a 4 x 4 grid containing eight unique objects. The player move up, down, left, right, and may static during an episode. A competent world model should be able to reconstruct the complete object configuration after observing all tiles at least once. We generate 100,000 sequences of 250 frames using random actions from the pygame library [Shinners & the Pygame community (2000 Actions are sampled such that covering/uncovering occurs with 0.4 probability, while each directional movement (up/down/leff/rr
uncover tiles. The object layout remains
occurs with 0.15 probability, promoting frequent tile changes. The dataset is split into 90% training and 10% test sets. Additionally, we provide specialized test set where sequences start either fully covered or uncovered and then transition through zig-7ag uncovering/covering patterns.
D ADDITIONAL EXPERIMENTS
D.1 MEMORY-CARDE DISCRETE RECALL EVALUATION
This experiment evaluates the model's ability to recall individual objects that have been concealed or occluded over time. We use the discrete Memory Cardy dataset, specifically constructed for this purpose.
We train standard diffusion U-Net [Snn et al. 2024) for 20k steps, described in Section F. The evaluation input consists of video sequences showing 4 x 4 grid of cards, which are sequentially revealed and concealed line by line. To measure object-level recall, we compute the tile-wise mean squared error (MSE) between the predicted and ground-truth frames at the end of each sequence. A grid tile considered correctly reconstructed if its MSE falls below a threshold of 2 x 10-5, Our method achieves recall accuracy of 79, .2% over a sample of 50 sequences with 20 memorization steps. In contrast, the standard diffusion (SD) baseline performs at chance level, achieving only 13%. See Section H for qualitative results.
D.2 DMLAB: DETERMINISTIC ENVIRONMENT ANALYSIS
We replicate the Memory Maze experiment in DMLab, fixed environment that allows for determin- istic evaluation of recall capabilities. The agent initially explores the maze to collect observations, after which we generate 50 future frames and compute L.IPS scores against ground-truth frames. We use a UVIT model, as described in Section E.1
The baseline SD model achieves an LPIPS of 0.558, while our method improves this to 0.456, marking an 18% improvement. Notably, no clear overfitting ceiling is observed: increasing the number of memorization steps from 1000 to 2000 yields a modest further improvement of 0.028 in LPIPS. This suggests that, for deterministic tasks, performance may continue to benefit from additional memorization steps. Visual examples are provided in Section H
D.3 ABLATION: LANGEVIN CORRECTION STEPS
Motivated by techniques from Du al. (202)) and predictor-correcto sampling in diffusion mod- els BBradley & Nakkiran (2024), we examine whether inserting Langevin dynamics steps during sampling improves generation quality and memory recall.
We conduct this ablation on the Memory Maze dataset, using 200 context frames and generating 17 future frames. We vary both the number and scale of Langevin steps, testing fixed step sizes well as step sizes proportional to the noise schedule 3,. Across all configurations, no significant improvement over the baseline, indicating that these modifications do not provide a favorable tradeoff between computational cost and generation quality. Furthermore using unadjusted
observe
17

---

Published as a conference paper at ICLR 2026
Hamiltonian Monte Carlo, did not yield better results, with the same settings detailed in Du et al. )003
D.4 EXPERIMENT ON MINECRAFT (OASIS-STYLE SETTING).
CoME is intended general framework rather than standalone world model: whenever diffusion- based world models are used, CoME can be layered on top. Beyond the diverse architectures evaluated in Tables 2-3, additionally apply CoME to Oasis/Diamond-styll setup on the simpler Minecraft Marsh dataset. We train a diffusion model that takes input context frames and predicts 17 future frames. The STM variant again uses doubled patch size for efficiency and 33 context frames to predict 17 future frames, while all other settings follow Section
Table Minecraft Marsh results (3-+17 prediction). CoME consistently improves perceptual and reconstruction metrics over the base model and single-cxpert variants.
Model
LPIPS +
SSIM f
PSNR 1
Base
0,408
0,450
16.19
+STM
0.375
0.549
17.03
+LTM
0.389
0.552
17.28
CoME
0.369
0.574
17.78
These results demonstrate that CoME transfers effectively to Oasis-style diffusion world models in Minecraft, yiclding consistent gains across perceptual similarity (LPIPS) and reconstruction metrics (SSIM, PSNR).
E
IMPLEMENTATION DETAILS
This section outlines the experimental configurations used across datasets to ensure reproducibility of Composition of Memory Experts. For consistency, all memorization steps, i.e., the gradient updates on past trajectories, are performed using the Adam W optimizer Loshchilov & Hutter (2019) with the same learning rate and hyperparameters (2025), we retain the conditioning during memorization, mirroring the training setup. Regarding hyperparameters of Composition of Memory Experts, found that simple and stable settings (e.g. the same a within a narrow, fixed range [2, for all experts) worked robustly across tasks.
during training. Unlike prior work Hong et al.
E. EXPERIMENTAL DETAILS
Memory-Maze We evaluate approach on the Memory-Mace 9x9 environment using the DiT(- B) configuration Peebles & Xie (2022a). Input videos are 64 x 64 resolution and are subsampled every second frame. We train VAE with 2x spatial downsampling. The model takes in total of 20 frames, 3 context frames and predicts the next 17 frames. For relative positional encoding, we concatenate translational and rotational changes and stack conditioning channels for skipped frames. In the case of relative position conditioning, we omit conditioning on skipped frames. The model is trained for 150k steps with a batch size of 16. We generate 512 videos for evaluation and exclude the ground truth context frames from metric computations.
We implement the STM as DiT(-S) with sliding window attention chunk size of 20, with 2 full attention layers at the and 6 layers, taking 33 frames as conditioning and predicting the next 17 frames. The LTM is simply the base model but with LoRA applied to query, key, value projections, MLPs, and out projections.
The SLTM is implemented as a DiT(-B) transformer with patch size 4 (instead 2). It takes absolute 2D position and angle (2+2 dimensions) as direct conditioning. To infer absolute position from relative motion alone, we train an additional encoder composed of 12 3D convolutional layers. This encoder embeddings formed by combining the current prediction of clean frames with the relative positioning information. In our experiments we infuse the SLTM with the information by test-time training similar to the L.T..
18

---

Published as a conference paper at ICLR 2026
Our baselines take 200 context frames and predict the next 17 frames. The full attention baseline consists a DiT(-B) transformer with patch size 4. The SSM baseline employs a DiT(-B) transformer that processes 20 frames at time. Every third layer integrates Mamba Gu & Dao (2024) layer, which causally propagates information to subsequent chunks. For the sliding window baseline, we replace the Mamba global layers with full attention layers, which mimics our STM architecture.
RealEstate For RealEsttate 10K, we follow Song et al. 2025). reusing the pretrained model and associated hyperparameters. All methods use 4 context frames and 4 future frames within the attention window. We compare against DFoT [Song et al. (2025), evaluating both History Guidance variants: vanilla (HG-v) and temporal (HG-t). Each method generates 256 videos for evaluation, with 4 clean ground truth frames provided as context. Note that this setting differs from the interpolation setup in Song et al. (2025), so reused hyperparameters may not be optimal.
UVIT backbone Hoogeboom al. (203)), detailed in Table 7 We follow the same diffusion and sampling procedure as in Memory Maze, but employ a continuous noise schedule and operate in pixel space. Training uses a linear learning rate scheduler with 1000 warmup steps, a learning rate of 2 x 10-1, weight decay of 0.001, and Adam optimizer. Videos resized to 64 64 and subsampled with stride of 2. We train for 50k steps with batch size of 256. Evaluation involves generating 512 videos using 50 DDIM steps. The SD baseline mirrors the sampling algorithm used for Memory Maze.
DMLab For DMLab,
Table 7: Configuration of the UViT3D model.
Component
Setting
Encoder Channels
[128, 128, 256, 256]
Block Types
2xResBlock, 2xTransformerBlock
Up/Down Blocks
[3. 3]
Embedding Dimension
1024
Patch Size
2
Mid Transformer Blocks
16
Attention Heads
diffusion framework, and training hyperparameters from Song et al. (2025). We train latent diffusion model using precomputed latents obtained from a pretrained image VAË Rombach et al.) (021). The model attends to 20-frame window with clean context frames. For the STM, we double the patch size for efficiency and use 33 context frames to predict 17 future frames. Other settings follow Section E.
Minecraft We adopt the DiT architecture
Petpcctutureeeeeeeeeeeeeeeee Peebles & (2c22022a)
Memory-Cards We use a UNet backbone Sun et al. (2024), detailed in Table g Input resolution 48 48. The model is conditioned on frames to predict future frames. Diffusion and training settings follow those used for DMLab. We train for 20k steps with a batch size of 256 until convergence. For evaluation, we sample 100 videos using 20 DDIM steps.
Table 8: Configuration of the UNet3D backbone.
Component
Setting
Encoder Channels
[64, 128, 256, 512]
Number of Resnets in Block
8
Attention Resolutions
[8, 16. 32, 64] 4
Attention Heads
F COMPUTE ANALYSIS
For our compute analysis, we employ the DiT configuration used in the Memory Maze experiments, retaining most settings described in Sections F.1 Inference with memory adaptation generally
19

---

Published as a conference paper at ICLR 2026
compute compared the base model, where is the additional number of experts we chose due to the need for multiple forward passes, one through the expert model and another through the contrastive expert. However, the actual overhead depends heavily the size of the expert networks. In settings where the experts have significantly fewer parameters than the base model (e.g., 45% of the pretrained model's size in the STM, SLTM), the additional compute can be substantially reduced. In our Memory Maze setting, this results in only 2x increase with the addition of LTM and 3x increase in forward compute for using all experts in CoME.
requires up to 2
Table 9: Average runtime (in seconds) for memorization and sampling. For 50 memorization and 50 denoising steps with different LTM, such as different LoRA ranks and smaller adapter network.
Adatter Config
Memorization Time
Sample Time
64
3.33
1.88
32
2.42
1.88
r" 16
2.07
1.87
LTM (45%)
4.01
1.38
We compare the memorization and sampling time the streaming setting from Section 4 we evaluate the runtime on a Nvidia RTX 4090 by sampling 64 videos with activated TorchScript compilation. Each video involves 50 sampling steps interleaved with 50 memorization steps. Average runtimes for different adapter configurations presented in Table 9
Full FLOPs and parameter tables across benchmarks. For completeness, we report the same back-of-thee -nevelleee FLOPs and parameter analysis for all major experimental settings. Unless otherwise stated, FLOPs are computed for 50 forward denoising steps. For configurations involving test-time memorization (LTM/SLTM), we additionally report the backward-pass cost incurred during adaptation.
Table 10: Memory Maze (DiT backbone). Reference corresponds 50 forward passes of the base DiT.
Config
Total Params
Trainable
GFLOPs (fwd)
GFLOPs (bwd)
Base
57,918,992
0
585.89
STM
32,662,720
206.32
LTM
59,500,304
1,581,312
602.83
419.45
SLTM
59,550,272
1,582,080
150.92
104.92
SSM
58,114,833
0
6,293.09
Sliding
57,918,992
11,198.53
Full
57,918,992
0
5,858.95
DiT (Match)
172,645.649
0
1,752.87
Table 11: Navigation World Models (NWM). Reference corresponds to 50 forward passes. CoME configurations include 25 backward passes for memorization.
Config
Total Params
Trainable
GFLOPs (fwd) GFLOPs (bwd)
CDiT-XL CDiT-L-LTM
675M 464M
1,719.87
6.3M
1,182.43
805.95
CDiT-B-STM
231M
1,405.33
These numbers highlight two practical points: (i) the added cost of CoME comes primarily from extra forward passes rather than from a large increase in trainable parameters (LoRA adapters are ~ ~.6. parameters), and (ii) scaling via heterogeneous experts be competitive with common single-model alternatives such as sliding-windon or SSM-style approaches under comparable rollout settings.
20

---

Published as a conference paper at ICLR 2026
Table 12: RealEstateI0K (R10k). Reference HG-t uses classifier-free guidance (CFG) evaluations step. CoME includes 100 backward passes for memorization.
Config
Total Params
Trainable
GFLOPs (fwd)
GFLOPs (bwd)
UViT-Base
288M
0
2,420.06
UViT-LTM
293M
4.6M
2,493. 46
1,161.93
UViT-HG-v
288M
0
2 x 2,420.06
UViT-HG-t
288M
X 2,420.06
G LIMITATIONS
While CoME provides a flexible and effective framework for composing heterogeneous memory experts, CoME relies the assumption that each expert provides plausible and reasonably calibrated predictions. In practice, we observed failure cases primarily when one expert produced outputs that the remaining experts considered implausible. This typically occurred when an expert's s expressiveness was severely constrained (e.g., due to aggressive parameter reduction), preventing it from reproducing basic scene structure. In such cases, the fused prediction can degrade in quality. While did not observe instability under normal settings, CoME remains sensitive to the quality and capacity of its constituent experts.
H ADDITIONAL VISUAL EXAMPLES
This section provides qualitative results that complement the main experiments described in Sec- tion E For each dataset, we visualize the generations produced by model after receiving fixed context window, The context frames are highlighted red. The sampling configuration matches that used in Section E.
We present results across the Memory Maze (Figures and 6), RealEstate10KK (Figures and 8). DMLab (Figure(9), and Memory Cards (Figure 10p environments, highlighting spatial and temporal coherence in the generated rollouts. For DMLab datasets, we pick videos where the memory is required, when the actions lead to revisit of the scene. For the other datasets, randomly pick videos for visualization.
I LARGE LANGUAGE MODEL USAGE
Large language models were used solely to assist in revising and improving the clarity, structure, and grammar of the camera-ready manuscript. No LLMs were used for generating experimental results, modifying data, designing algorithms, or conducting analyses. All technical content, experiments, and conclusions were developed and verified by the authors.
21

---

Published as a conference paper at ICLR 2026
GT
Ours
SD
Figure 5: Memory Maze CoME. After memorizing trajectory through the maze, and only given three context frame, the model generates frames that accurately reflect the correct turns and re-entries.
22

---

Published as a conference paper at ICLR 2026
GT
Ours
SD
Figure 6: Memory Maze COME After memorizing a trajectory through the maze, and only given three context frame, the model generates frames that accurately reflect the correct turns and re-entries.
23

---

Published as a conference paper at ICLR 2026
GT
Ours
SD
DEOT-HGT
DFoT-HG-L
Figure 7: RealEstate10K- COME, We provide 4 ground truth context frames, here highlighted with red border. We generate forward rollouts and 3 rollouts with the reversed trajectory camera position.
24

---

Published as a conference paper at ICLR 2026
Ours
SD
DFOT-HG-V
DFoT-HG-t
Figure 8: RealEstate10K COME, We provide 4 ground truth context frames, here highlighted with red border. We generate forward rollouts and 3 rollouts with the reversed trajectory camera position.
25

---

Published as a conference paper at ICLR 2026
GT
Ours
SD
Figure DMLab COME. Given 50 context frames, here highlighted with a red border, we visualize the next 4 rollouts. Our method produces coherent forward trajectories that reflect consistent agent movement and scene structure.
26

---

Published as a conference paper at ICLR 2026
QuiS
Target
65
OuT
Target
6S
OUIS
Target
Figure 10: Memory Cards Discrete object recall. At the end of the sequence, the model is evaluated on its ability to regenerate occluded tiles. After being shown sequence of uncovering and covering actions, such that all tiles were visible at some point, our method more accurately recalls the occluded tiles, here given 'target', demonstrating effective memory recall.
27


Reviews:
[{"review_id": "w0GUkdvYmS", "reviewer": "Reviewer_5qnA", "summary": "This paper proposes CoME, which effectively combines long-term and short-term memory in Diffusion-based World Models to address memory limitations. Based on the existing limitations and concerns, I recommend **weak reject**. However, I am willing to raise my score if the authors can address these limitations and concerns.", "questions": "* How does CoME world model condition on action input?\n* Can the authors provide additional experiments on Minecraft, since many existing works experiment on this environment (Dreamer4 [1], Oasis [2], MineWorld [3])? Or compare with other existing World Model methods on standard datasets (StateSpaceDiffuser [4], RLVR-World [5], GameNGen [6])?\n* What is the difference between SLTM and LTM besides patch size and conditioning? Do they both use the same test-time finetuning method? Section E.1 states \"similar to LTM\" but provides no details.\n* If the test set does not include relative and absolute position information (i.e., only contains images and action sequences) would this method still be effective?\n* Can the authors provide detailed computational efficiency comparisons with baselines, including end-to-end latency and throughput?\n* What is the total parameter count of CoME, can author provides parameter count comparison with existing baselines.\n* Why does SLTM use doubled patch size?\n\n[1] Training Agents Inside of Scalable World Models\n\n[2] Oasis: A Universe in a Transformer\n\n[3] MineWorld: a Real-Time and Open-Source Interactive World Model on Minecraft\n\n[4] StateSpaceDiffuser: Bringing Long Context to Diffusion World Models\n\n[5] RLVR-World: Training World Models with Reinforcement Learning\n\n[6] Diffusion Models Are Real-Time Game Engines", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Well-motivated problem: The fundamental trade-off between context length and computational cost in world models is clearly articulated. The human cognition analogy (STM/LTM) provides good intuition.\n* Effective and elegant solution: Using test-time finetuning to overcome context length limitations in the DiT architecture is an effective and concise approach.\n* Thorough ablations: Table 4 convincingly demonstrates the necessity of the contrastive formulation; Table 5 analyzes LTM capacity trade-offs.", "weaknesses": "* Insufficient baselines: Missing comparisons with key related works (WorldMem [1], StateSpaceDiffuser [2]).\n* Low computational efficiency.\n* CoME uses 4 models vs single models for baselines, is this a fair comparison? \n* No failure case analysis.\n\n[1] WORLDMEM: Long-term Consistent World Simulation with Memory\n\n[2] StateSpaceDiffuser: Bringing Long Context to Diffusion World Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ngY7sD0Qtp", "reviewer": "Reviewer_LCCy", "summary": "- The authors propose approaches for augmenting diffusion world models with external memories (a short, long, and short-long term memory) that enable diffusion world models to have better long-term rollouts. The authors propose a product of contrastive experts approach that improves performance, and demonstrate broadly that their proposed approach has good properties empirically and loosely theoretically.", "questions": "- Overall, I'm open to revising my score, but that primarily relies on addressing my concerns/questions given in the weaknesses.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The results are generally very strong, with the proposed CoME generally outperforming all existing models in LPIPS, SSIM, and PSNR.\n- The writing and motivation are good, as I agree that incorporating short/long/short+long term memory into existing models would be very impactful.\n- The ablations are generally very strong in terms of supporting the author's claims.", "weaknesses": "- It's not clear that a fair comparison was done to me, which is my biggest concern. The authors don't include rich details regarding the parameters/FLOPs of models in Table 1, Table 2, and table 3. This makes it hard to determine whether or not the memories are working as intended and really providing strong performance for their size, or are just giving the model more capacity. \n    - I see this as the largest issue. Is it possible to reveal this information as well as do a comparison between a baseline model with more FLOPs/parameters to match the size of the baseline with the memory experts?\n    - Further, is it possible to do comparisons at multiple different model sizes, to see if the approach scales well?\n- Visual perceptive metrics are the primarily method for evaluating the model. While this is common, pixel level reconstruction has been shown to be less effective than modeling in a learned latent space [1], meaning the world models may produce higher fidelity rollouts but not actually be better world models. What do the authors think of this?\n\n[1] https://arxiv.org/abs/2506.09985", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ov4Xq8rpYy", "reviewer": "Reviewer_DBh7", "summary": "To leverage the benefits from different memory models, the proposed approach (CoME) utilises a diffusion framework with a contrastive mechanism to compose information from memories. Experiments across different settings showed that the strategy improves performance.", "questions": "The author could highlight more in the baseline, which one is the baseline that contains the simpler way of combining knowledge from memory models.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The mechanism to compose heterogeneous experts is well-presented in Section 3.2, and different memory models (long-term, short-term, and Spatial Long-term Memory) are detailed clearly in Section 3.3.\n- Experiments were conducted on a range of datasets, including Memory Maze, RealEstate10K, DeepMind Lab, Memory Cards Dataset, and RECON.  The comparisons of different memory combinations are also highlighted in Table 1. The setting in section 4.2 shows that the composed mechanism can handle a stream of observations.", "weaknesses": "Since the approach proposed to use all memory models, it will be more expensive than using one memory. However, the authors also discussed and provided compute analysis in Appendix F.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t7kEGiewQZ", "reviewer": "Reviewer_9rob", "summary": "In this paper, the authors propose a novel diffusion world model that employs a composition of multiple types of memory experts to overcome the *memory-compute trade-off* inherent in existing architectures (e.g., Transformers, State Space Models (SSMs)). This model generates the next frame by probabilistically fusing the predicted probability distributions from each expert into a single, unified distribution.\n\nThe paper introduces three types of experts: a **Short-Term Memory (STM)** expert that predicts based on recent context, a **Long-Term Memory (LTM)** expert that stores long-term context in its weights via test-time fine-tuning, and a **Spatial Long-Term Memory (SLTM)** expert that utilizes spatial knowledge for prediction . In their experiments, the authors validate the positive impact of each memory component on model performance. They also demonstrate that performance **monotonically increases** as more context is provided to the LTM.\n\nFurthermore, the authors propose a novel formulation, the **Product of Contrastive Experts (PoCE)**, to ensure that each expert focuses more intently on the given context . This method selectively amplifies predictions that have a higher relative probability when conditioned on the context, compared to an unconditional baseline . Through experiments, they show that this approach mitigates the **consistency collapse** phenomenon—often observed in the conventional Product of Experts (PoE) method—and achieves superior performance .\n\nFinally, the authors experimentally validate that their full model, combining the STM, LTM, and SLTM experts with the PoCE unification formulation, outperforms existing memory-augmented world models. In the appendix, they provide a detailed description of their experimental setup and demonstrate the effectiveness of PoCE through a toy example . They also validate the superior performance of their world model through a discrete recall evaluation (Memory-Cards environment) and a deterministic environment analysis (DMLab) . Notably, in the discrete recall evaluation, their model improves recall accuracy by approximately **66%** compared to a standard diffusion baseline (79.2% vs. 13%).", "questions": "- Could you please share a more detailed experimental analysis for the Short-Term Memory (STM) and Spatial Long-Term Memory (SLTM) experts? For instance, it would be insightful to see an experiment where, in the absence of SLTM, the model's performance on a dedicated spatial reasoning task (e.g., as explored in [1]) is shown to decrease. This would provide stronger evidence that the SLTM is indeed contributing the specific spatial capabilities as intended.\n\n[1] Cho, Junmo, Jaesik Yoon, and Sungjin Ahn. \"Spatially-aware transformer for embodied agents.\" _arXiv preprint arXiv:2402.15160_ (2024).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **Addresses a significant research question and proposes a viable alternative:** The authors tackle the critical _memory-compute tradeoff_  in world models, a problem where models are either limited to using short-term memory or incur excessive computational overhead when employing long-term memory . To address this, they introduce the concept of a Product of Experts (PoE), proposing a model where specialized memory experts operate compositionally, and they experimentally validate its effectiveness through significant performance improvements.\n    \n- **Novel formulation for a contrastive Product of Experts:** The authors do not simply adopt the standard PoE framework. They identify a potential issue of consistency collapse when applying it to memory experts and propose a novel contrastive formulation to mitigate this problem . The superiority of this method is convincingly demonstrated through both an intuitive toy visualization and an empirical ablation study .\n    \n- **Thorough experimental validation and detailed methodology:** The authors demonstrate the effectiveness of their proposed compositional memory experts through a comprehensive set of ablation studies. For the Long-Term Memory (LTM) expert, in particular, they show a monotonic performance improvement as the context length increases . Furthermore, they strengthen their claims by including additional insightful experiments in the appendix, such as the PoCE toy visualization and a discrete recall evaluation, which experimentally prove the model's efficacy . The paper is also well-written, with the methodology and experimental settings described in sufficient detail in both the main paper and the appendix to ensure reproducibility.", "weaknesses": "- **Insufficient discussion on whether the proposed method truly resolves the memory-compute tradeoff:** The authors frame their work as a solution to the memory-compute tradeoff. However, their proposed method, Composition of Memory Experts (CoME), requires forward passes for each expert and its contrastive counterpart, leading to a computational complexity increase of approximately $2\\times n$ (where $n$ is the number of experts), as discussed in Appendix F. This raises doubts as to whether CoME fundamentally solves the issue, as requiring more memory (i.e., more experts) still leads to a direct increase in computation. While the authors mention in Appendix F that this overhead can be reduced by using smaller expert models, this point is critical to the paper's central research question and warrants a more in-depth discussion with concrete experimental results in the main paper, not just the appendix .\n    \n- **Lack of discussion on limitations:** The paper lacks a thorough discussion of the limitations of the CoME framework. For instance, as the number of experts with diverse characteristics increases, the computational overhead will grow, and the process of fusing their predicted distributions could become more challenging (e.g., tuning the hyperparameters $\\alpha_i$ ​may become prohibitively complex). Moreover, while the framework is effective if the chosen memory experts are well-suited for the task, inappropriate experts could have a detrimental impact, potentially leading to performance worse than the base model. A discussion of these potential failure points and scalability challenges should be included.\n    \n- **Insufficient experimental analysis for Short-Term Memory and Spatial Long-Term Memory:** While the paper demonstrates a monotonic performance improvement for Long-Term Memory (LTM) with increased context, similar in-depth analyses for Short-Term Memory (STM) and Spatial Long-Term Memory (SLTM) are missing. The ablation studies confirm that using these modules is beneficial, but they do not provide clear evidence that the performance gains are achieved for the intended reasons—i.e., that STM excels by focusing on recent observations and SLTM by leveraging spatial knowledge. This makes it difficult to verify whether these components are functioning as hypothesized.\n\n- **The Statement about the Use of Large Language Model:** The authors doesn’t state explicitly how they used LLM in their submission, which is guided in https://iclr.cc/Conferences/2026/AuthorGuide.\n    \n- **(Minor) Comments on Presentation:**\n    - In line 173, explaining the \"subsets of frames\" with a concrete example would help prevent reader confusion about how these subsets are determined.\n    - In line 221, there appears to be a small formatting artifact at the end of the line.\n    - In Section 3.3, the description of SLTM is ambiguous about its use of fine-tuning. This should be clarified in the main text, as it is mentioned in Appendix E.1.\n    - The discussion on training cost (lines 359-364) could be more detailed with the training time for each model.\n    - In Section 4.3, the dataset used for the ablation studies should be explicitly mentioned.\n    - In Figure 3, the meaning of each row and what the reader should observe (e.g., which frames are forward vs. reverse) is not clearly explained in the caption.\n    - In line 975, it should be explicitly stated that the \"$2\\times$\" increase is in comparison to the base model.\n    - In Figures 5-10, the red box annotations appear misaligned or incorrectly drawn in several places (e.g., all frames are boxed in \"Ours\" in Figure 5; boxes are unclear in Figure 10).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bF1bjg0yif", "reviewer": "Reviewer_hHXL", "summary": "This paper proposes a memory-compositional diffusion world model that integrates multiple specialized “memory experts”—short-term memory (STM), long-term LoRA-based memory (LTM), and spatial long-term memory (SLTM)—into a unified model via a **Product of Contrastive Experts (PoCE)** mechanism. The PoCE combines multiple memory outputs while reducing redundant features through contrastive weighting, encouraging each expert to specialize. Experiments on navigation and reconstruction datasets show clear improvements in long-horizon consistency, memory recall, and generalization to novel sequences.", "questions": "1. How sensitive is performance to the number of experts or their relative weights?\n\n2. Could you provide compute and runtime comparisons to conventional diffusion world models?\n\n3. Are there stability issues when experts disagree strongly during sampling?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* The conceptual idea of composing multiple memory experts through a contrastive product is innovative and well-articulated.\n\n* The use of test-time LoRA as long-term memory is practical, allowing scalable specialization without retraining the core model.\n\n* Experiments convincingly demonstrate that the composition improves both reconstruction quality and navigation performance.\n\n* The paper is clearly written, with mathematical formulation and visual illustrations that make the mechanism understandable.\n\n* The approach connects well to emerging work on compositional generative models and memory-based reasoning.", "weaknesses": "* The assumptions behind the contrastive composition (such as the independence of expert outputs) are not empirically validated.\n\n* Computational overhead and runtime comparisons to single-memory baselines are missing.\n\n* While results are strong, the evaluation scope is somewhat limited to a few datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)