Research map / Architectures and efficiency / Attention and transformers
EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes EDELINE, a diffusion world model that uses a Mamba-based recurrent embedding to condition next-frame prediction on full history and jointly predicts rewards/terminations, achieving SOTA model-based Atari 100k HNS without look-ahead search.
Problem
Diffusion-based world models such as DIAMOND condition on a fixed short window of frames, limiting memory in partially observable environments, and use separate networks for reward/termination prediction.
Approach
Integrate a Mamba selective SSM as a recurrent embedding module whose hidden state conditions a diffusion next-frame predictor (via adaptive group norm and cross-attention) and directly feeds reward and termination MLP heads, with HarmonyDream-style dynamic loss weighting and parallel-scan training on sampled trajectory segments.
Key results
- On Atari 100k, EDELINE reaches mean HNS 1.866, median 0.817, IQM 0.940, with 13 superhuman games versus 11 for DIAMOND (mean 1.459).
- On Crafter with a 1M step budget, EDELINE achieves 11.5 ± 0.9 average return with an 11M-parameter world model, versus 2.8 for DIAMOND (10.4M params) and 9.2 for DreamerV3 XL (200M params).
- Linear probing shows observation-prediction loss averaged over environments is reduced by 57.3% (0.3610 vs 0.8462) relative to DIAMOND, with comparable reward-prediction loss (0.0246 vs 0.0258).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace DIAMOND's fixed four-frame context with EDELINE's recurrent Mamba memory and unified world model | Diffusion-based world model · Atari 100k | Mean HNS: 1.459 -> 1.866 | Helped |
| Replace DIAMOND's fixed four-frame context with EDELINE's recurrent Mamba memory and unified world model | Diffusion-based world model · Crafter | Avg Return: 2.8 ± 0.5 -> 11.5 ± 0.9 | Helped |
| Use EDELINE's unified hidden representation instead of DIAMOND's hidden states | World model hidden representation · Atari 100k (26 games) | Average observation probing loss: 0.8462 -> 0.3610 | Helped |
| Use EDELINE's unified hidden representation instead of DIAMOND's hidden states | World model hidden representation · Atari 100k (26 games) | Average reward probing loss: 0.0258 -> 0.0246 | No clear effect |
| Use EDELINE's recurrent Mamba memory and conditioning instead of DIAMOND's four-frame context | Diffusion-based world model · Atari 100k (26 games) and ViZDoom-DeadlyCorridor | Mean normalized pixel MSE: 1.000 -> 0.918 | Helped |
Surprising findings
- EDELINE's world-model update is about 26.8% faster than DIAMOND, yet its actor-critic updates take about 17.7% more time, so total training time is roughly comparable.
- EDELINE underperforms DIAMOND on some games such as Assault, Frostbite, and MsPacman, and on pixel MSE for PrivateEye and Asterix.
Limitations
- Diffusion-based world models still require multiple forward passes, so computational overhead remains high.
- Performance is not uniformly better: EDELINE trails DIAMOND on several Atari games (e.g., Assault, Frostbite, MsPacman).
Experimental setup
Datasets: Atari 100k Crafter ViZDoom MiniGrid-Memory
Benchmarks: Atari 100k Crafter ViZDoom MiniGrid-Memory
Metrics: human-normalized score mean HNS median HNS IQM optimality gap average return pixel-wise MSE linear probe loss
Baselines: DIAMOND STORM DreamerV3 IRIS TWM Drama SimPLe Delta-IRIS
Builds on: DIAMOND Mamba HarmonyDream EDM DreamerV3
Closest papers by meaning
- Diffusion for World Modeling: Visual Details Matter in Atari (NeurIPS 2024 · similarity 0.90)
- Zero-shot World Models via Search in Memory (NeurIPS 2025 · similarity 0.87)
- DyMoDreamer: World Modeling with Dynamic Modulation (NeurIPS 2025 · similarity 0.87)
- StateSpaceDiffuser: Bringing Long Context to Diffusion World Models (NeurIPS 2025 · similarity 0.86)
- Learning 3D Persistent Embodied World Models (NeurIPS 2025 · similarity 0.86)
- Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.