EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling

We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.

Q1. Limited Architectural Innovation and Potential Redundancy.

We appreciate the reviewer's observation. We would like to bring to the reviewer's attention that this represents significant architectural innovation, as evidenced by the superior performance compared to prior works in this domain. Our design choice to maintain both components stems from the complementary roles they serve in the diffusion process. The hidden embedding captures long-term temporal dependencies and global context across the entire sequence history, while the explicit frame stacking provides pixel-based spatial details from recent observations that are crucial for pixel-level prediction accuracy. The diffusion model benefits from both the abstract temporal representations (via cross-attention with hidden embedding) and direct access to recent visual features (via frame stacking) for generating high-fidelity predictions.

For empirical validation to address this concern, we conducted an ablation study comparing: (1) EDELINE with both components, (2) Hidden embedding only (without frame stacking), and (3) DIAMOND. As shown in Table R4, the results demonstrate that in Atari Breakout, which requires minimal long-term memory while demanding high visual detail processing, EDELINE without frame stacking performs slightly higher than DIAMOND, though still significantly worse than full EDELINE. In MiniGrid-MemoryS9, which requires substantial memory capabilities while having simpler visual complexity, EDELINE without frame stacking can still successfully learn to solve the task and outperforms DIAMOND. Note that our analysis reveals that removing either component leads to degradation in either temporal consistency or visual fidelity.

We appreciate the reviewer's suggestion and consider new architectural design or more elegant integration mechanisms as interesting avenues for future research directions. However, we would like to respectfully emphasize that our current framework achieves the primary objective of combining long-term memory with high-quality visual prediction while maintaining computational efficiency, which is the main theme of this manuscript. The experimental results faithfully validate our contributions, and we have experimental evidence demonstrating that these components are complementary in achieving superior performance.

Table R4: Frame Stacking Ablation Study. Performance comparison demonstrating the necessity of both SSM hidden states and explicit frame stacking across environments with different memory and visual complexity requirements.

Q2. Typo in Line 18 and Line 328.

We thank the reviewer for catching these typographical errors. We will correct these issues in the revised manuscript.

Q3. In Figure 3, the labels are too small.

We appreciate the reviewer's feedback regarding figure readability. We will increase the label sizes in Figure 3 to enhance clarity in the revised version.

We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.

Q1. What's the training time and inference time of EDELINE compared to other baselines?

We thank the reviewer for this important question. We provide the total training time for DreamerV3 [1], IRIS [2], DIAMOND [3], and EDELINE on a single RTX 4090 GPU for Atari 100k benchmark training. This time encompasses both world model training and policy training using the world model and data collection.

In addition, we provide the single-step imagination time (i.e., world model inference) for these baselines. While EDELINE's inference time is slightly higher than DIAMOND's, the adoption of Mamba's parallel scan algorithm and the random target timestep selection technique during world model training makes EDELINE more efficient in world model training compared to DIAMOND. This enables EDELINE to maintain computational efficiency comparable to DIAMOND while addressing DIAMOND's memory capacity limitations. The detailed comparison is presented in Table R2.

Table R2: Training and Inference Time Comparison. Comparison of training efficiency, inference speed, and performance across different world model architectures on Atari 100k benchmark, including Atari 100k Human Normalized Score (HNS).

[1] Hafner, Danijar, et al. "Mastering diverse domains through world models." Nature (2025).

[2] Micheli, Vincent, Eloi Alonso, and François Fleuret. "Transformers are sample-efficient world models." ICLR (2023).

[3] Alonso, Eloi, et al. "Diffusion for world modeling: Visual details matter in atari." NeurIPS (2024).

Q2. Could the authors provide additional experiments on challenging visual RL?

We thank the reviewer for this suggestion. We conducted additional evaluations on the Battle scenario in ViZDoom, a 3D FPS game scenario proposed in [1]. This scenario features a 3D enclosed maze with multiple ammunition pickups, health packs, and continuously moving monsters that attack the player. The evaluation metric is the number of monsters killed by the player. We employed a sparse reward function where the agent receives a reward of 1 only upon successfully killing a monster, with the agent limited to 1M environment interactions. This setting simultaneously challenges the agent's capabilities in handling complex visual inputs, sparse rewards, 3D FPS gameplay, long-term memory, and sample efficiency. As demonstrated in Table R3 below, EDELINE significantly outperforms DIAMOND [2] due to its memory advantages and unified architecture.

Table R3: ViZDoom Battle Scenario Results. Performance comparison on a challenging 3D FPS environment requiring long-term memory and sparse reward handling.

[1] Dosovitskiy, Alexey, and Vladlen Koltun. "Learning to act by predicting the future." ICLR (2017).

[2] Alonso, Eloi, et al. "Diffusion for world modeling: Visual details matter in atari." NeurIPS (2024).

We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.

Q1. It is unclear how EDELINE differs from and improves upon existing works that also integrate SSMs with diffusion models or design world models to predict multiple environment elements. The authors may want to clarify the key novelty of their method.

We thank the reviewer for the question. Our primary contribution lies in designing a unified architecture that leverages SSMs to address the memory limitations of existing diffusion-based world models within the Model-based Reinforcement Learning domain, while maintaining computational efficiency through Mamba's parallel scan algorithm during training.

We would like to clarify that most current works combining diffusion models with Mamba, such as DiM [1], DiMSUM [2], Matten [3], and other recent approaches, primarily utilize Mamba's linear time complexity and unbounded context advantages as a replacement for traditional attention-based architectures in image generation and video synthesis tasks. These works, however, concentrate primarily on architectural fusion rather than addressing the specific challenges of diffusion-based world models. Regarding world models that predict multiple environment elements like Dynalang [4], those world models focus on grounding diverse language types through GRU-based sequence modeling, while our proposed EDELINE specifically targets the memory constraints of diffusion-based world models like DIAMOND (e.g., fixed 4-frame context windows).

Our key novelty lies in integrating Mamba SSMs to enable unbounded temporal consistency in diffusion world modeling while maintaining superior visual fidelity. This approach addresses a fundamentally different challenge from existing approaches.

[1] Teng, Yao, et al. "DiM: Diffusion mamba for efficient high-resolution image synthesis." arXiv preprint arXiv:2405.14224 (2024).

[2] Phung, Hao, et al. "DiMSUM: Diffusion Mamba--A Scalable and Unified Spatial-Frequency Method for Image Generation." NeurIPS (2024).

[3] Gao, Yu, et al. "Matten: Video generation with mamba-attention." arXiv preprint arXiv:2405.03025 (2024).

[4] Lin, Jessy, et al. "Learning to model the world with language." ICML (2024).

Q2. There exists symbol abuse and inconsistency in the paper.

We appreciate the reviewer's attention to this detail. We will address these notation issues in our revised manuscript by replacing the observation space symbol $O$ with $\mathcal{O}$ and changing the burn-in length symbol $B$ to $\mathcal{B}$ to eliminate the overloading of symbols.

Q3. Why is the input $o_{t+1}$ required for predicting $\hat{o}^0_{t+1}$ in the next-frame predictor module of Figure 2?

We thank the reviewer for this question. Figure 2 illustrates the training process of EDELINE's world model. During training, the noised version of $o_{t+1}$ ​ serves as the input to the diffusion process, where the model performs denoising based on the current conditioning information and noise level. This follows the standard training procedure for conditional diffusion models. During inference, no ground truth frame is available. Instead, the model generates the next frame $\hat{o}\_{t+1}$ by starting from pure Gaussian noise and progressively denoising through multiple steps without requiring $o_{t+1}$ as input. This approach is consistent with previous diffusion-based world models such as DIAMOND and follows established practices in conditional diffusion literature. We will incorporate this clarification in the revised manuscript.

Q4. Can the author provide the complete mathematical formulation showing how $h_t$ is derived from $h_{t-1}$ using matrices $A, B, C,$ and $\Delta$ in Figure 2?

We thank the reviewer for this question. Our implementation follows the established Mamba literature. We concatenate the current timestep's observation $o_t$ and action $a_t$ ​as input, denoted as $u_t = (o_t, a_t)$. We project $u_t$ to initialize the input-dependent matrices $B, C$, and $\Delta$, while matrix $A$ remains input-independent. The discretization is then performed using $\Delta$ to obtain the discrete matrices $\tilde{A}$ and $\tilde{B}$. Let $x_t$ denote the Mamba latent state at timestep $t$. The Mamba architecture updates the latent state through the following recurrence:

$x_t = \tilde{A} x_{t-1} + \tilde{B} u_t$

The hidden embedding $h_t$ in the EDELINE architecture is obtained by applying matrix $C$ to the mamba latent state $x_t$:

$h_t = C x_t$

This architecture enables Mamba to store unbounded contextual information within the latent state and compute the current timestep's hidden embedding $h_t$ through the selective mechanism. As a result, $x_{t-1}$ contains the information from $h_{t-1}$ ​since $h_{t-1} = C x_{t-1}$, establishing the complete derivation chain from $h_{t-1}$​ to $h_t$​ through the latent state recurrence. To enhance the clarity of Figure 2, we will include the Mamba latent states $x_{t-1}$ and $x_t$​ in the revised version of the diagram.

We appreciate the reviewer for this follow-up question, and we would be pleased to clarify the differences.

In EDELINE's architecture, Mamba serves as an independent sequence processing module (REM) that specifically handles temporal sequences of (observation, action) pairs to establish causal understanding of environment dynamics. We preserve the original diffusion model architecture while we inject Mamba's hidden embeddings as conditioning information into the diffusion process through cross-attention and adaptive group normalization (AGN). EDELINE is applied in the Model-based RL domain, with the primary goal of addressing memory limitations of previous diffusion-based world models and improving RL agent sample efficiency.

We would like to bring to the reviewer's kind attention that, in contrast, works like DiM [1] primarily address computational efficiency issues of Diffusion Transformers (DiT) in high-resolution image generation. DiM tackles the fundamental challenge of adapting Mamba's 1D sequential modeling to 2D image data by introducing multi-directional scanning patterns that alternate between row-major and column-major traversals to provide global receptive fields. The method inserts learnable padding tokens at spatial boundaries to preserve image structure during sequence flattening and incorporates lightweight depth-wise convolutions for local feature enhancement. This approach enables linear computational complexity while maintaining generation quality for high-resolution image synthesis tasks. Similarly, DiMSUM [2] integrates Mamba directly into the diffusion model's core architecture with a focus on solving the spatial scanning order challenges inherent to 2D image processing. DiMSUM introduces a spatial-frequency fusion approach using wavelet transforms to capture both spatial and frequency domain information simultaneously. The architecture replaces traditional attention mechanisms within the diffusion model with Mamba-based spatial-frequency processing blocks, specifically designed to address the inductive bias limitations of Mamba when processing 2D visual data for high-quality static image generation tasks.

On the other hand, Matten [3] adopts a hybrid architectural approach and directly integrates both Mamba and Attention mechanisms within the diffusion model for video generation. Matten serially combines Spatial Attention, Temporal Attention, and Global-Sequence Mamba scanning to form unified processing blocks that capture local and global spatiotemporal relationships in videos. We note that the key difference from EDELINE is that Matten embeds Mamba as a component of the diffusion model architecture directly into the generation process with a focus on spatiotemporal continuity modeling between video frames for unconditional video generation tasks.

These architectural differences reflect fundamentally different design objectives. DiM and DiMSUM pursue efficient static image generation, Matten focuses on high-quality video synthesis, while EDELINE is specifically designed for causal dynamics modeling in reinforcement learning environments. Specifically, our approach requires the maintenance of memory across extended observation-action trajectories to predict future states, rewards, and terminations in partially observable environments.

It is important to note that EDELINE's contribution in world modeling goes beyond visual prediction by utilizing the same Mamba hidden embeddings for multi-task learning. It simultaneously predicts next-frame observations, reward signals, and termination signals. This unified representation learning is essential for world modeling in RL, where accurate environment modeling directly impacts agent sample efficiency. In contrast, DiM and DiMSUM are specifically designed for independent image generation tasks and are not designed for world modeling scenarios that require modeling sequential state transitions and reward prediction. Similarly, Matten focuses on unconditional video generation and does not include mechanisms for modeling the causal relationships between actions and environmental outcomes necessary for world modeling.

Finally, EDELINE's REM module design enables it to process unbounded sequence lengths while maintaining causal understanding of environment dynamics, which differs fundamentally from video generation tasks that primarily focus on visual coherence and quality. Therefore, EDELINE's contribution lies in specifically applying Mamba's sequence modeling capabilities to address the memory limitations of diffusion-based world models, while simultaneously enabling multi-task prediction for complete environment modeling, rather than primarily focusing on generation quality or computational efficiency.

We appreciate the reviewer's thoughtful engagement and welcome further discussion to address any additional questions that may arise.

We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.

Q1. The comparison between the other REM modules should be placed in the main body.

We thank the reviewer for this valuable suggestion. We acknowledge the reviewer's recommendation and will incorporate the comparison between different REM modules into the main body of the revised manuscript to better support our architectural design choices.

Q2. This particular statement is not sufficient to rule out the choice of Transformers.

We thank the reviewer for this important observation. We agree with the reviewer that Transformer represents a viable and potentially promising architectural choice for our framework. Our initial experiments employed a window size of 16 to maintain computational feasibility and fair comparison with existing baselines during our preliminary investigations. However, we acknowledge this may have inadvertently disadvantaged Transformer architectures. When we removed this limitation and extended the window size to match Crafter's maximum episode steps (10,000), we found that Transformer performance became comparable to Mamba, as shown in Table R1.

Our experimental results demonstrate that both Transformer and Mamba architectures achieve strong performance in our framework. We selected Mamba not due to its superiority over Transformer, but rather as we found that the simpler linear-time complexity solution already achieves satisfactory performance for our objectives. Given that both architectures perform well, and considering that the primary theme and most important aspect of this work is efficiency, we opted for the more computationally straightforward approach while acknowledging that Transformer remains an equally valid choice. We will update our discussion to incorporate the perspective of Transformer's capabilities in the revised manuscript.

Table R1: Recurrent Embedding Module (REM) Architecture Comparison on Crafter Environment. We conducted additional comparisons between Transformer and Mamba architectures as the REM component in EDELINE. The evaluation includes performance (average return), computational efficiency (i.e., REM single-step inference FLOPs measured at context length = 32), and model complexity (i.e., number of REM parameters). The results demonstrate that both architectures are viable choices, with Transformer achieving competitive performance when the window size limitation is removed.

Q3. I would suggest sharing stats of the context lengths used during training, and also practically useful context length in inference.

We thank the reviewer for the suggestion. During world model training, we employ different sequence lengths depending on the environment: 19 steps for Atari, MiniGrid, and ViZDoom, and 32 steps for Crafter. For the imagination phase when training the Actor-Critic component, we use four steps of world model burn-in followed by 15 steps of imagination horizon for Atari, MiniGrid, and ViZDoom, while Crafter employs 32 steps of burn-in with a 15-step imagination horizon. Please note that due to Mamba's characteristics, there are no specific context length limitations.

Q4. The training details are not sufficiently presented.

We thank the reviewer for pointing this out. We have provided comprehensive implementation details in our supplementary materials. Algorithm 1 in Appendix E presents our overall training procedure, while Appendix D.12 contains hyperparameters used during training. In the revised version, we will include a detailed table of contents and reorganize the supplementary materials to avoid potential confusion and improve accessibility.

Q5. What are practically sufficient context lengths needed here? Do the authors have insights on true long-horizon tasks in the real world?

We thank the reviewer for the question. As detailed in our response to Q3 above, the sequence lengths vary by environment and training phase. While we agree with the reviewer regarding real-world applicability opportunities, please note that the primary theme, focus, and contributions of this work concern addressing the memory limitations of existing diffusion-based world models. Our contributions enhance the potential for deploying these models in real applications, especially when long-term memory capabilities are essential for effective decision-making.

Q6. Could the authors provide detailed specifications of the Mamba architecture used as well as the training sequence lengths utilized across the different environments?

We thank the reviewer for the question regarding reproducibility. We have provided Mamba-related configurations in Section D.12, including the number of Mamba layers, Mamba's latent state size, and expansion ratio. Additional settings follow standard Mamba defaults as established in the original implementation. The training sequence lengths across different environments are as described in our response to Q3 above. To further improve clarity and accessibility, we will reorganize these technical specifications for better accessibility in the revised manuscript.