跪拜 Guibai
← Back to the summary

Nano World Models Puts a Reproducible Baseline Under World-Model Research

【World Models】Nano World Models (NanoWM) Analysis Notes

Author: 罗西的思考

0x00 Overview

World models (ha2018recurrent; dawid2023introduction) have become the foundation of spatial intelligence (yang2025cambrian; wang2026vagen) and real-world decision making (richens2025general; guo2025ctrl) - they generate high-fidelity future predictions conditioned on an agent's historical observations and actions.

Nano World Models is a minimalist, out-of-the-box codebase that aims to advance a cautious and scientific approach to world-model design.

0x01 Motivation

The motivation behind the Nano World Model project is simple: although industrial-grade world models achieve impressive visual effects, they ultimately rely on a small set of simple, mature techniques: video diffusion (ho2022video; svd), diffusion forcing (chen2024diffusion), and consistency distillation (song2023consistency).

The authors of Nano World Models argue that as world-model algorithms gradually stabilize, research focus will shift from "inventing new techniques" toward "developing deeper understanding of more subtle scientific design decisions" - including architecture choice, training objectives, and of course data composition and scaling behavior. However, it is impossible to truly understand the science without being able to easily experiment with models. Especially in the current highly fragmented landscape of world-model research, different datasets (brohan2022rt; pearce2022counter; zhou2024dino), training schemes (yang2023diffusion; lipman2024flow; li2025back), evaluation protocols (vafa2024world; zhang2026worldinworld), and downstream tasks (alonso2024diffusion; quevedo2026worldgym; guo2025ctrl) are scattered across a large number of independent sources, making rigorous scientific research extremely difficult.

Nano World Models serves as a usable, minimalist, and out-of-the-box world-model implementation in this fragmented landscape. Its codebase is specifically built around the diffusion forcing (chen2024diffusion) paradigm. It supports a Hydra-based configuration system (Yadan2019Hydra), modular code for data loading, model design, training, evaluation, and downstream tasks, and carefully curated documentation. Specifically, it supports the following features:

0x02 What is a world model, and what are its necessary components

2.1 Explain to a twelve-year-old

What is a world model?

When you play a game, like Super Mario, and you see a pit ahead, your brain immediately "imagines": if I jump over it, I'll land safely; if I walk straight in, I'll fall. You can "imagine" this outcome because your brain already contains a "small model" of the world - you know how gravity works, and how far Mario can jump.

A World Model teaches a computer to have this same "imagination." You show the computer lots of videos (e.g., a robot pushing things, a person playing games), and the computer learns "how the world works." Afterwards, you ask it "if I do this action, what will happen next?" and it can imagine the result in its "mind."

What must a world model contain?

Just like you need three things to play a game well:

  1. Eyes to see frames (encoder) - turns "what the world looks like" into data the computer can understand
  2. Brain to learn patterns (prediction model) - learns "if action X is taken, the frame becomes Y"
  3. Mouth to speak results (decoder) - turns the ideas in the brain back into videos or images you can understand

Without any one of these, a world model cannot function.

2.2 Explain to a semi-professional

What is a world model?

A world model is a generative model that learns to simulate environment dynamics by learning from environment interaction data. It takes an agent's historical observation sequence and action sequence as input and learns a conditional probability distribution:

p(o_{t+1}, o_{t+2}, ..., o_{t+H} | o_{1:t}, a_{1:t+H-1})

That is, given the observed historical frames and the future action sequence to be executed, it predicts the future observation sequence. The core value of a world model is that it encodes implicit environmental rules (physical laws, object interactions, scene evolution) into a queryable parameterized function - the agent does not need to trial-and-error in the real environment, but can reason about the outcomes of different actions in the model's "imagination."

The necessary components of a world model:

From a functional architecture perspective, a complete world model requires four core components:

Component Function Analogy Implementation in NanoWM
Perception Encoder Compresses high-dimensional observations (RGB images) into a compact latent space Human visual cortex - turns retinal signals into "object" concepts SD-VAE / Web-DINO / V-JEPA 2.1
Dynamics Predictor Models temporal state transitions in latent space, conditioned on actions Human cerebellum - predicts "what happens after moving" NanoWM (DiT + Diffusion Forcing)
Objective Defines "good" vs "bad" criteria for guiding planning Human prefrontal cortex - judges "is that good to do" MSE / Cosine distance vs goal latent
Planner / Decoder Uses the predictor to search for optimal action paths, or decodes latent variables back to understandable observations Human motor cortex - turns plans into actual actions CEM Planner / VAE Decoder

All four components are indispensable: without an encoder, the model cannot process raw pixels; without a predictor, the model has no "imagination"; without an objective, it cannot judge which future is better; without a planner/decoder, "imagination" cannot translate into action.

2.3 Explain to a professional

2.3.1 What is a world model

A world model is a generative framework that encodes environment interaction dynamics into a differentiable parameterized function. Formally, given a historical observation sequence of length T, o_{1:T}, and an action sequence a_{1:T}, the world model learns the conditional distribution of future observations o_{T+1:T+H}:

p_θ(o_{T+1:T+H} | o_{1:T}, a_{1:T+H-1})

Differences from traditional dynamics models (e.g., Kalman filters, Gaussian processes):

  1. High-dimensional observation space - handles RGB images (not low-dimensional state vectors)
  2. Generative uncertainty modeling - does not output a deterministic o_{t+1} = f(o_t, a_t), but models the full distribution p(o_{t+1} | ...)
  3. Long-range temporal dependencies - learns causal structure between frames, not single-step transitions under a Markov assumption

The role of a world model in a decision system is as a differentiable, queryable simulator - the agent substitutes expensive real-environment interactions with "trial-and-error in imagination."

2.3.2 Classification of world models

From the perspective of learning paradigms, modern world models can be divided into three categories:

Paradigm Representative work Core idea Objective
Reconstruction-based DreamerV3, IRIS, NanoWM w/ SD-VAE Compress observations into a decodable latent space, predict dynamics in that space Pixel-level fidelity + planning
Representation-based JEPA, DINO-WM Predict in semantic/geometric feature space, not requiring decodability Planning-oriented abstract representation
Pixel-based Genie, GameNGen Predict directly in pixel space (diffusion/autoregressive) Extreme visual fidelity (costly)

2.3.3 Necessary components

Necessary components of a world model (system architecture perspective):

Reasoning from first principles, a world model usable for decision making must include the following five subsystems, which have strict dependencies among them:

1. Observation Encoder enc: O → Z, where O is the observation space (usually a high-dimensional continuous space like RGB images, or discrete space like tokens), and Z is the latent space.

Necessity: the dimension of o_t is usually [3, H, W] ≈ 10^5-10^6; modeling such a high-dimensional distribution p(o_{t+1} | o_t, a_t) directly is statistically infeasible. The encoder reduces the problem to a tractable latent space (e.g., VAE's 4×32×32 = 4096 dimensions, a compression ratio of about 200×).
Additionally, the encoder can introduce inductive biases (e.g., VAE's reconstruction orientation, DINO's semantic orientation).

2. Temporal Transition Model trans: Z_t × A_t → Z_{t+1}, generalized as a conditional generation process p(z_{t+1:t+H} | z_{1:t}, a_{t:t+H-1}).

Necessity: this is the "core intelligence" of the world model. Without a transition model, the model can only memorize seen trajectories, and cannot generalize to unseen situations. NanoWM uses diffusion forcing (Diffusion Forcing) to implement this - through a scheduling matrix, it controls the noise level of each frame, and unifies teacher-forcing, autoregressive generation, and full-sequence generation in a single framework.

3. Action Conditioning Interface inject: (Z_t, A_t) → Z_t', which injects action information into the state representation.

Necessity: without action conditioning, the model can only do unconditional video prediction ("the most likely future"), and cannot do counterfactual reasoning ("what if I do action A?"). Counterfactual reasoning is the basis of planning.
NanoWM supports 5 injection mechanisms, with complexity ranging from additive (zero extra parameters) to cross-attention (fully parameterized interaction), providing a continuous spectrum from "weak conditioning" to "strong conditioning."

4. Generator / Decoder dec: Z → O or gen: Z → p(O), which maps latent space predictions back to observation space.

Necessity: latent space predictions may be sufficient for downstream tasks (like planning), but not for human understanding and visualization.
The decoder ensures the world model's output is interpretable. Moreover, in MPC planning, the objective is usually defined in latent space (cosine distance to goal latent), so the decoder is not strictly necessary for planning, but it is necessary for interpretability and verification.

5. Inference / Planning Engine plan: (Z_init, Z_goal, trans) → A*, which searches for the optimal action sequence on the learned dynamics model.

Necessity: a world model itself is just a generator, not a decision-maker.
To make it produce action value, an optimization algorithm must run on top of the model. CEM is the most common approach:
  1) Sample N action sequences
  2) For each rollout, obtain a future trajectory
  3) Score with the objective function
  4) Select elites to update the sampling distribution
  5) Execute the first action of the optimal sequence → re-plan (MPC)

Formal dependency relations between components:

Component relationships

Component relationships

Data flow:

Data flow

Data flow

Action conditioning is an input interface to the transition model, not an independent component; the whole system has only 4 necessary components - encoder, transition model, decoder, planner.

2.3.4 Open questions

Open questions in current world-model architectures:

0x03 World model vs video generation model

Looking at the backbone alone, the difference between a world model and a video generation model is indeed small. The core difference is not "which skeleton is used," but four fundamental differences:

3.1 Different conditioning signals: video generation vs world model

Video generation models (SVD, Genie, etc.) World models (NanoWM, etc.)
Input conditioning Text/image Action sequence (continuous vectors) + historical frames
Output semantics "The most likely future" "The future that will happen if this action is taken"
Counterfactual reasoning ❌ Not present ✅ Queryable: what if action X? What if action Y?

Video generation models learn p(video | text) - given text, they generate one video. World models learn p(o_{t+1:t+H} | o_{1:t}, a_{t:t+H-1}) - given history and actions, they generate a future controlled by actions.

This is the basis of counterfactual reasoning. For example:

The "generalization" of a world model is not "generating pretty videos," but rather:
- During training, only sees left turn → frame moves left, right turn → frame moves right
- At test time, asked: what if you push the stick up?
- Correct answer: the frame moves up (even if the exact "push up" scenario was never seen in the data)

A video generation model cannot do this - it has no "action" channel to accept such queries.

3.2 Rollout error accumulation - a challenge unique to world models

Video generation models only need a single-pass generation (text → video). World models need to repeatedly use their own output as input for the next step:

Video generation (single-pass):
  text → [o₁, o₂, ..., o₁₆]     ← generate 16 frames at once, errors do not cross steps

World model (auto-regressive rollout):
  1: o₁, o₂ → predict o₃          ← step 1
  2: o₂, o₃ → predict o₄          ← o₃ already contains step-1 error
  3: o₃, o₄ → predict o₅          ← o₄ contains cumulative error from first two steps...
  ...
  50: o₄₉, o₅₀ → predict o₅₁     ← error accumulates exponentially

This is error accumulation. Video generation models do not have this problem because their generation does not depend on their own output as subsequent input.

This directly dictates that the architecture design of a world model must include:

3.3 Latent space interface and codec protocol

Video generation models are usually fixed in pixel space or a single latent space. World models must expose a multifunctional interface:

# The world model's latent_codec must support all of:
z = latent_codec.encode(o)           # encode: o → latent space (for prediction)
ô = latent_codec.decode(z)           # decode: latent space → pixels (for visualization)
dist = cosine(z_pred, z_goal)        # compare in latent space (for planning objective)

This means the encoder design of a world model must consider the trade-off across three uses:

Finding #4 of NanoWM exposes the difficulty of this trade-off - DINO/V-JEPA latent spaces, despite being semantically good, cannot learn controllable dynamics under the diffusion framework.

3.4 Inference-time interaction paradigm: MPC closed loop

Video generation is open-loop - input conditioning, output generation, done.

World models are closed-loop - they are embedded in a decision loop:

Each decision loop:
  1. Get the latest observation o_t from the environment
  2. Use the world model to rollout multiple candidate action sequences
  3. Compare each rollout result with the goal
  4. Select the best action → execute the first step in the environment
  5. Environment returns new observation o_{t+1}
  6. Back to step 1 (re-planning)

This is called Model Predictive Control (MPC) - the world model acts as a "queryable simulator."

This loop imposes additional engineering requirements on the world model:

3.5 Summary: why a world model cannot be equated with a DiT video backbone

Comparison

Comparison

The backbone is the same, but the surrounding systems, usage, and engineering challenges are completely different. To use an analogy: a truck and a race car might both have V8 engines, but you wouldn't say "race cars and trucks are no different" - the chassis, transmission, suspension, and purpose are all different.

But your skepticism is entirely correct on one point: once backbones converge, the "moat" of world-model research is not in the model architecture itself, but in how to make that backbone learn controllable, planning-usable dynamics - and that is exactly the question NanoWM's paper tries to answer with its six findings.

0x04 Understanding Nano World Models from three perspectives

4.1 Perspective one: if I were a twelve-year-old

Imagine you are playing a video game, like Minecraft. You walk around in this world, push obstacles aside, pick up items... Your brain unconsciously learns to "predict what will happen next" - for example, if you push a rock, it will roll forward.

Nano World Models (NanoWM for short) is about teaching computers this "predict the future" ability.

Specifically, we show the computer a bunch of videos, like a robot pushing things, or a person playing CS:GO. Then we ask the computer: "If you do this action now, what will the next frame look like?" By learning the patterns in the videos, the computer can "imagine" future frames just like a person.

What makes this project special:

In short: NanoWM is a toolbox that lets computers "imagine the future," and this toolbox is both simple and powerful.

4.2 Perspective two: if I were a semi-professional

Nano World Models is a world-model codebase based on the Diffusion Forcing paradigm. A "world model" learns the dynamic evolution rules of an environment by studying large amounts of sequential observation data (video frame sequences), so it can predict future observations given current states and action sequences.

Core capabilities at a glance:

Capability Description
Future video prediction Given several historical frames + action sequence, generate subsequent frames
Long-range autoregressive generation Use sliding windows to exceed training length, generating videos of 50+ frames
Action-conditioned modeling Supports 5 action injection mechanisms (additive, AdaLN, AdaLN-Fuse, FiLM, cross-attention)
Multiple diffusion targets Supports x-prediction, ε-prediction, v-prediction, and Flow Matching
Multi-scale architectures NanoWM-S (40M) / B (160M) / L (600M) / XL (830M)
Multiple latent encoders SD-VAE (decodable to pixels), Web-DINO (semantic features), V-JEPA 2.1 (video-pretrained features)
Downstream applications MPC planning (CEM algorithm), video-to-3D point cloud, long rollouts
Unified configuration system Hydra-based modular YAML config, all components pluggable

Tech stack: PyTorch + PyTorch Lightning + Hydra + OmegaConf, with a custom implementation of GaussianDiffusion and FlowMatching for the diffusion process; the core model is a DiT (Diffusion Transformer) architecture.

4.3 Perspective three: if I were a professional

Nano World Models is a modular, reproducible research framework for future video prediction, with core contributions:

  1. Unified interface design: With diffusion forcing (Chen et al. 2024) as the core abstraction, it unifies different prediction targets (x/ε/v/flow), architecture scales (S/B/L/XL), action injection mechanisms (additive/adaLN/FiLM/cross-attention), latent encoders (SD-VAE/Web-DINO/V-JEPA), and dataset formats (D4RL/HDF5/LeRobot) into a single training-sampling pipeline - a rarity among existing open-source implementations.
  2. Scalable DiT video backbone: Uses alternating spatial-temporal attention (Spatial-Temporal Transformer) instead of the traditional UNet, with 2D + 1D Rotary Position Embedding for position encoding, and causal masking for autoregressive generation. Patch size (1/2/4/8) controls the number of latent tokens, naturally supporting variable-length sequences.
  3. Diffusion Forcing scheduling matrix: A custom scheduling matrix supports three sampling modes - full_sequence (traditional full-sequence denoising), pyramid (staggered initial noise across frames), and sequential (frame-by-frame autoregressive denoising) - for the first time unifying teacher-forcing and autoregressive generation in a single framework.
  4. MPC (model predictive control) planning interface: Wraps the world model as a batched rollout function, decoupled from the CEM (Cross-Entropy Method) planner, supporting direct latent-space objective computation (MSE / cosine) for zero-shot planning.
  5. Scaled ablation experiments: The paper systematically compares prediction targets, model size, action injection methods, and latent space types across 6+ datasets (RT-1, PushT, CS:GO, Point Maze, etc.) for video quality and planning performance, providing reproducible baselines.

4.4 Core findings of the paper

The paper's six core findings:

# Finding Experimental basis
#1 v-prediction and x-prediction beat ε-prediction. v-prediction gives the best FID (42.27), x-prediction the best PSNR (23.37). ε-prediction is worst (FID=48.86) RT-1 50K-step ablation, cosine + ZTSNR scheduling
#2 Larger models perform better. Scaling from NanoWM-S/2(40M)→B/2(160M)→L/2(~460M) monotonically improves PSNR, SSIM, LPIPS, and FID RT-1 ablation, all hyperparameters fixed
#3 Action injection effectiveness is task-dependent. FiLM gives the best FID on RT-1 (40.62), but additive is best on PushT (FID=23.89) and gives the best parameter-quality trade-off RT-1 + PushT dual-task ablation, 5 injection mechanisms compared
#4 Semantic latent spaces (Web-DINO, V-JEPA 2.1) do not automatically suit planning. Action embedding RMS is near 0 (0.002 vs SD-VAE's 0.112), leading to 0% MPC planning success. Reason: the diffusion objective plus semantic latent space cannot force the model to learn controllable dynamics PushT goal-reaching planning experiment
#5 Long rollouts accumulate error, but increasing DDIM sampling steps mitigates it. In 50-frame CS:GO generation, LPIPS grows monotonically with frame number; 250-step DDIM beats 50-step CS:GO long rollout experiment
#6 The same training-evaluation recipe transfers across domains, but simple domains (Point Maze FID=9.66) are much better than complex domains (Granular FID=40.05, RT-1 FID=35.08) Standardized evaluation across 6 datasets

0x05 Related background knowledge

5.1 Diffusion Forcing

5.1.1 Explain Diffusion Forcing to a twelve-year-old

Imagine completing a jigsaw puzzle, but not all at once - you have 4 puzzle areas on the table, corresponding to the 4 frames of a video:

Even better, you can also switch play styles:

Key difference: the traditional method is like "all students start running from the same starting line at the same time"; Diffusion Forcing is like "each student can start from a different starting line and at a different time" - this lets the model handle more flexibly the question of "which frames are already determined and which need to be generated."

5.1.2 Explain Diffusion Forcing to a semi-professional

Standard diffusion models (e.g., DDPM, DDIM) apply the same noise level to all frames during sampling and then gradually denoise synchronously. This has a fundamental limitation for conditional video generation - context frames (already observed) and future frames (to be generated) should be at different noise levels, but standard diffusion forces them to share the same noise schedule.

The core improvement of Diffusion Forcing (DF) is that each frame at each timestep can have an independent noise level. Specifically:

Traditional video diffusion:                 Diffusion Forcing:
step:  [250, 250, 250, 250]   step: [ -1, 250, 200, 150]
                                  (-1 = clean, already the final result)

The scheduling matrix is the mathematical abstraction of DF - it is a [T×F] matrix, where T is the number of sampling steps and F the number of frames. Each element represents the noise level of the corresponding frame at that sampling step (-1 means already complete). By changing this matrix, DF unifies three generation modes:

Mode Scheduling matrix characteristics Semantics
full_sequence All frames share the same noise level, decreasing synchronously Traditional unconditional video generation
sequential Each frame goes from high noise to clean one by one, the next frame waits for the previous to finish Autoregressive generation (teacher-forcing)
pyramid Noise start times are staggered across frames (overlap), multiple frames are parallel but ordered Balances parallel efficiency and causality

During training: each frame's t in DF is also sampled independently (t ∈ [B, F] or [B]), meaning the model must learn to make correct predictions even when adjacent frames have completely different noise levels. This forces the model to learn more robust temporal representations.

Comparison between standard Diffusion Forcing (original Chen et al. 2024) and NanoWM's implementation: NanoWM's DF implementation simplifies some aspects - it mainly uses the scheduling matrix to control the sampling process, while the training loss still uses the standard MSE of GaussianDiffusion. Its core value is supporting both teacher-forcing and autoregressive generation in a single model without modifying the model architecture.

5.1.3 Explain Diffusion Forcing to a professional

Problem background

Standard diffusion models (DDPM/DDIM) apply the same noise corruption and denoising schedule to all data dimensions (e.g., all pixels of an image, all frames of a video). Formally:

For a video x_{1:F}, the forward process is q(x_t | x_0) = N(√ᾱ(t)·x_0, (1-ᾱ(t))·I), with the same ᾱ(t) applied to all frames.

This creates a contradiction in practice: in conditional video generation (like world models), context frames (known) and future frames (unknown) should be at different noise states - context is known (should be near clean), future is unknown (should be highly noisy). But standard diffusion forces both to share the same noise level.

Formalization of Diffusion Forcing

Diffusion Forcing extends diffusion to frame-wise noise scheduling:

Define a set of noise indices K ⊂ R. For a trajectory x_{1:F}, each frame x_f is assigned an independent noise index k_f ∈ K. The model is trained on input noisy trajectories x_{k_1}, ..., x_{k_F} (where x_{k_f} = √ᾱ(k_f)·x_f + √(1-ᾱ(k_f))·ε_f) and learns to predict clean data or noise.

The scheduling matrix S ∈ R^{T×F} is the core abstraction of DFoT (Diffusion Forcing of Time):

# generate_scheduling_matrix(full_sequence, T=250, F=4):
S_full = [[250, 250, 250, 250],   # step 0: all high noise
          ...,
          [  0,   0,   0,   0]]   # step 250: all clean

# generate_scheduling_matrix(sequential, T=250, F=4):
S_seq = [[250,  -1,  -1,  -1],    # step 0: only frame 1 denoises
         ...
         [  0,  -1,  -1,  -1],    # step 250: frame 1 complete
         [  0, 250,  -1,  -1],    # step 251: start frame 2
         ...
         [  0,   0,   0,   0]]    # step 1000: all complete

where -1 indicates the frame is "frozen" at that step (already generated, no longer participates in denoising).

Relationship to existing frameworks
Perspective Equivalence
Teacher Forcing Sequential mode: each frame is generated conditioned on previously generated frames
Masked Autoregressive Noise level controls visibility/masking degree (not binary mask, but continuous noise level)
DDPM Full-sequence mode is strictly equivalent to standard DDPM/DDIM
Cold Diffusion DF is a generalization of cold diffusion - allowing different dimensions to have different degradation degrees

5.1.4 Pseudocode implementation

Based on NanoWM's implementation, here is a complete two-stage pseudocode: training (diffusion forcing) and sampling (DFoT).

1. Training stage (independent timesteps per frame)
# ──────────────────────────────────────────────────────────────
# Training: Diffusion Forcing single step
# Core difference: t is [B, F] not [B], each frame has its own noise level
# ──────────────────────────────────────────────────────────────
def train_step(model, diffusion, optimizer, batch):
    # batch.video: [B, F, 3, H, W]   raw pixel video
    # batch.action: [B, F, act_dim]  action sequence

    # (1) Encode to latent space (frozen latent_codec)
    with torch.no_grad():
        x = rearrange(batch.video, 'b f c h w -> (b f) c h w')
        z = latent_codec.encode(x)             # [B*F, C_lat, H_lat, W_lat]
        z = rearrange(z, '(b f) c h w -> b f c h w', b=B)

    # (2) Sample timestep independently per frame  ←←← core of Diffusion Forcing
    t = sample_timesteps(shape=[B, F],        # shape is [B, F] not [B]!
                         max=diffusion.num_timesteps)   # each frame a random t_f ~ U[0,1000)

    # (3) Forward noising: each frame noised by its own t
    noise = randn_like(z)                      # [B, F, C, H, W]
    z_t = q_sample(z, t, noise)                # z_t = √ᾱ(t)·z + √(1-ᾱ(t))·ε
                                               # note ᾱ is taken per-frame!

    # (4) Model forward: input noisy z_t and per-frame t
    pred = model(z_t, t=t, action=batch.action)
    #      model internally: t is [B,F] → flatten → t_embedder → per-frame conditioning
    #      output: [B, F, C, H, W]

    # (5) Compute target (v-prediction default)
    target = v(z, t, noise)                    # v = √ᾱ·ε - √(1-ᾱ)·z

    # (6) Loss: weighted with Min-SNR-gamma
    mse = mean((pred - target)², dim=[-3,-2,-1])      # [B, F]
    weight = clamp(SNR(t), γ) / SNR(t)                # lower weight for low-noise steps
    loss = mean(mse * weight)

    # (7) Backpropagation
    loss.backward()
    clip_grad_norm_(model.parameters(), 0.1)          # enabled after step > 20000
    optimizer.step()
2. Scheduling matrix generation
# ──────────────────────────────────────────────────────────────
# Scheduling matrix: the "score" for sampling
# Shape [num_steps, F], element = noise level of that frame at that sampling step
# Special value -1 = that frame is frozen (does not participate in denoising)
# ──────────────────────────────────────────────────────────────
def generate_scheduling_matrix(F, num_timesteps, sampling_steps, mode, n_context):
    # Only generate for frames that need generation (context frames are set to -1 at the end)
    G = num_gen_frames = F - n_context

    if mode == 'full_sequence':           # traditional: all frames denoise synchronously
        schedule = np.arange(sampling_steps, -1, -1)      # [steps+1, 1]
        schedule = schedule.repeat(G, axis=1)              # [steps+1, G]
        # shape: (sampling_steps+1, G)

    elif mode == 'pyramid':               # staggered: 1-step delay between frames
        schedule = np.zeros((sampling_steps + G, G), dtype=int)
        for step in range(sampling_steps + G):
            for f in range(G):
                eff = step - f                          # frame f delayed by f steps
                if   eff < 0:             schedule[step,f] = sampling_steps     # not started → pure noise
                elif eff >= sampling_steps: schedule[step,f] = 0                # complete
                else:                     schedule[step,f] = sampling_steps - eff
        # shape: (sampling_steps + G, G)

    elif mode == 'sequential':            # frame-by-frame: complete one frame before starting the next
        schedule = np.full((sampling_steps * G + 1, G), sampling_steps)
        for f in range(G):
            start = f * sampling_steps
            for s in range(sampling_steps + 1):
                schedule[start+s:, f] = sampling_steps - s
        # shape: (sampling_steps * G + 1, G)

    # Context frames: set to -1 (frozen)
    if n_context > 0:
        schedule = concat(zeros(steps, n_context), schedule, axis=1)
        # context column = 0 → mapped to -1

    # DDIM index → real timestep mapping (0→-1, sampling_steps→999)
    return ddim_idx_to_timestep(schedule, num_timesteps, sampling_steps)

Comparison of three modes (assuming sampling_steps=4, G=3):

full_sequence:  [4 4 4] → [3 4 4] → [2 3 4] → [1 2 3] → [0 1 2] → [-1 0 1] → [-1 -1 0] → [-1 -1 -1]   (7 steps)
pyramid:        [4 4 4] → [3 4 4] → [2 3 4] → [1 2 3] → [0 1 2] → [-1 0 1] → [-1 -1 0] → [-1 -1 -1]   (7 steps)
sequential:     [4 4 4] → [3 4 4] → [2 4 4] → [1 4 4] → [0 4 4] → [-1 3 4] → ... → [-1 -1 -1]   (13 steps)
3. DFoT sampling (multi-step denoising)
# ──────────────────────────────────────────────────────────────
# DFoT sampling: denoise frame-by-frame according to the scheduling matrix
# ──────────────────────────────────────────────────────────────
def dfot_sample(diffusion, model, shape, context, n_context,
                scheduling_mode, num_sampling_steps, model_kwargs):
    B, F = shape[:2]

    # (1) Generate scheduling matrix
    S = generate_scheduling_matrix(F, diffusion.num_timesteps,
                                   num_sampling_steps, scheduling_mode, n_context)
    # S: [num_steps, B, F] (expanded to batch)

    # (2) Initialize: all noise
    z = randn(*shape)                     # [B, F, C, H, W]

    # (3) Replace context frames with real frames
    if n_context > 0:
        z[:, :n_context] = context        # first n_context frames = latent of real observations

    # (4) (Optional) History stabilization: add a small noise to context frames
    if history_stabilization > 0:
        t_stab = int(history_stabilization * (num_timesteps - 1))
        z[:, :n_context] = q_sample(z[:, :n_context], t_stab)

    # (5) Main denoising loop: call the model in batches at each step
    for step in range(num_steps - 1):
        curr_t = S[step]      # [B, F] current noise level per frame
        next_t = S[step+1]    # [B, F] next noise level per frame

        z_prev = z.clone()    # backup (to restore context frames)

        # Single-step DDIM: one model inference, all frames updated simultaneously
        out = dfot_ddim_step(model, z, curr_t, next_t, model_kwargs)
        z = out['sample']

        # Restore context frames: use mask to keep frozen frames unchanged
        z = where(context_mask_expanded, z_prev, z)

    # (6) Return fully denoised latent
    return z                  # [B, F, C_lat, H_lat, W_lat]
4. Single-step DDIM (with per-frame noise levels)
# ──────────────────────────────────────────────────────────────
# DFoT DDIM single step: one model inference, handles all frames (regardless of noise level)
# ──────────────────────────────────────────────────────────────
def dfot_ddim_step(model, z, curr_t, next_t, model_kwargs):
    # z: [B, F, C, H, W]
    # curr_t, next_t: [B, F]   current/target noise level per frame
    #                      -1 means frame is frozen

    # Mask: only update frames where curr_t > 0 and not -1
    mask = (curr_t != -1) & (next_t < curr_t)      # [B, F]

    # Model inference (once, not per-frame!)
    pred = model(z, t=curr_t, **model_kwargs)      # predicts target (e.g., v)
    #      ↑ model must handle inputs where each frame has a different noise level
    #        (implemented via per-frame timestep embedding)

    # Recover x0 from prediction
    x0 = v_predict_to_x0(z, pred, curr_t)          # or corresponding ε/x₀ recovery

    # Compute denoised z (DDIM deterministic, eta=0)
    z_next = z + (sqrt(ᾱ(next_t)) - sqrt(ᾱ(curr_t))) * x0

    # Only update frames that need updating
    z = where(mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1), z_next, z)
    return {'sample': z}
5. Summary: correspondence between DF training and sampling
Dimension Training (single step) Sampling (multi-step)
Timestep shape t: [B, F] randomly sampled S matrix [steps, B, F] predefined
Number of model calls 1 sampling_steps times
Inter-frame noise relation Independent random Constrained by scheduling matrix
Context frames None (all are training data) Frozen (-1 or history stabilization)
Gradient Yes No (no_grad)
Noising direction Clean → noisy (forward) Noisy → clean (reverse)

Unified core idea: training and sampling share the same key mechanism - each frame has an independent noise level t_f. During training, t_f is randomly sampled; during sampling, it is determined by the scheduling matrix. The model learns to "given any sequence of frames with arbitrary noise levels, predict the denoising target," and this capability is fully reused in both stages.

5.1.5 Implementation details in NanoWM

Independence manifestation

NanoWM's DF implementation embodies independence at two levels:

1. Training: timesteps sampled independently

# train_experiment.py, training_step()
if diffusion_mode == "diffusion_forcing":
    t_shape = (x.shape[0], x.shape[1])  # [B, F] - one timestep per frame
elif diffusion_mode == "full_seq_diffusion":
    t_shape = (x.shape[0],)             # [B] - all frames share one timestep

When t is 2D [B, F], each frame receives an independent timestep embedding in the forward pass:

# nanowm.py forward()
if t.ndim == 2:
    t_flat = t.flatten()               # [B*F]
    t_emb_flat = self.t_embedder(t_flat)  # [B*F, D]
    timestep_spatial = t_emb_flat       # [B*F, D] - used by SpatialBlock
    timestep_temp = repeat(t_emb, ...)   # [B*P, F, D] - used by TemporalBlock

This means the SpatialBlock's adaLN_modulation receives a frame-specific conditioning vector - frame 1 (clean) and frame 4 (high noise) go through different adaLN paths.

TemporalBlock: after SpatialBlock processes intra-frame patch relations, rearrange changes the shape from [(B*F), P, D] to [(B*P), F, D] - i.e., each patch position observes its evolution across F frames. TemporalBlock performs 1D RoPE self-attention on this, allowing frames to communicate. If causal=True, frame t can only look at frames ≤ t, enabling autoregression. In one sentence: SpatialBlock handles "within a frame," TemporalBlock handles "across frames."

2. Sampling: driven by the scheduling matrix

dfot_sample() takes scheduling_mode and n_context_frames:

z_sample = dfot_sample(
    diffusion=self.diffusion,
    model=eval_model,
    shape=z.shape,                      # [B, F, C, H, W]
    context=z[:, :n_context],           # first n_context frames as clean context
    n_context_frames=n_context,
    scheduling_mode=scheduling_mode,    # "sequential" / "pyramid" / "full_sequence"
    num_sampling_steps=num_sampling_steps,
    model_kwargs=model_kwargs,
    ...
)

Internal logic is:

  1. Initialize z = torch.randn(shape), but replace z[:, :n_context] with real context frames
  2. Call generate_scheduling_matrix() to generate the denoising path matrix
  3. For each step, according to the per-frame noise level in the matrix:
    • If curr_t == -1: skip (frame already complete)
    • Otherwise: perform one DDIM sampling step on frame f
Key advantages
  1. Unified framework: no need to maintain three separate code paths for teacher-forcing, autoregressive, and full-sequence generation
  2. Smooth interpolation: noise levels are continuous, allowing partial frames to be in a "partially noisy" state - more flexible than binary masking in masked language models
  3. Sliding-window long-range generation: sequential mode naturally supports generating sequences far beyond training length - generate one frame at a time, then slide it into the context window
  4. History stabilization: paper experiments find that in sequential mode, applying a small noise to already-generated frames and re-denoising mitigates error accumulation
Limitations and open questions

5.2 Three diffusion prediction targets (x / ε / v) plus Flow Matching

First, recall the core formula of the diffusion forward process:

x_t = √ᾱ·x₀ + √(1-ᾱ)·ε    （Gaussian diffusion）

where x₀ is the clean data, ε ~ N(0,I) is the noise, and ᾱ is determined by the noise schedule (larger t means ᾱ closer to 0).

The model receives x_t and timestep t during training and must predict some target. Different targets are like teaching the model to understand the denoising process from different "perspectives":

5.2.1 ε-prediction (predicting noise)

5.2.2 x₀-prediction (predicting clean data)

5.2.3 v-prediction (predicting velocity, default)

5.2.4 Flow Matching

Unlike the three Gaussian diffusions above, Flow Matching changes the underlying logic:

x_t = (1-τ)·x₀ + τ·ε       ← linear interpolation, not exponential decay
τ = 1 - (t+1)/T             ← time parameter uniformly 0→1

5.2.5 Key differences at a glance

Key differences at a glance

Key differences at a glance

5.3 Five action injection mechanisms

NanoWM's core challenge is: how to make video prediction conditioned on actions? That is, given past frames plus an action sequence, the model must correctly predict the future frames corresponding to those actions.

Action injection occurs in each TransformerBlock, with 5 methods listed in increasing order of complexity:

① Additive (simplest)

How: add the action embedding directly to patch tokens

# Spatial block: action_emb [B*F, 1, D]
x = x + action_emb              # added to all patch tokens
# Temporal block: action_emb [B*P, F, D]
x = x + action_emb              # added to all frame tokens

Advantages: zero extra parameters, simple and efficient. Disadvantages: limited expressiveness - action is just a "additive bias" that cannot dynamically adjust feature processing. Paper result: best on PushT (PSNR=26.20), best parameter-quality trade-off

② AdaLN (Adaptive Layer Normalization)

adaLN_modulation: it is an nn.Sequential(SiLU, Linear(hidden_size, 6*hidden_size)). Its role is to dynamically generate 6 parameters based on the conditioning vector c (timestep or fused timestep+action) to control the shift, scale, and gate of the two sublayers (Self-Attention and MLP). This lets the model adjust feature processing according to "how far denoising has progressed."

Background: NanoWM's backbone uses adaLN-Zero - the timestep embedding passes through an MLP to generate 6 parameters (shift/scale/gate × 2 sublayers) used to modulate LayerNorm and gated residuals.

AdaLN variant: give actions a separate parallel adaLN branch, and add its generated shift/scale/gate to the timestep's corresponding parameters:

t_params = adaLN_modulation(c)        # [6*D] from timestep
a_params = action_adaLN_modulation(a) # [6*D] from action
shift = t_shift + a_shift，scale = t_scale + a_scale，gate = t_gate + a_gate

Advantages: actions can dynamically adjust the activation distribution of each layer (which features matter and which don't). Disadvantages: each Block adds a 6×D→D linear layer, increasing parameter count.

③ AdaLN-Fuse (fused injection)

Difference from AdaLN: action does not go through a separate adaLN branch; instead, the action embedding is added directly to the timestep conditioning vector c, so the same adaLN modulator handles both timestep and action information:

c = c + action_emb     # fusion
shift, scale, gate = adaLN_modulation(c)  # single branch handles fused conditioning

Advantages: no extra adaLN parameters; action and timestep interact through the same nonlinear mapping. Disadvantages: the mixing of action and time is fixed (addition), slightly less flexible than separate branches.

④ FiLM (Feature-wise Linear Modulation)

FiLM applies an affine transformation to features after Self-Attention and before MLP, using two vectors γ and β generated from the action embedding:

# After self-attention, before MLP:
gamma, beta = film_modulation(action_emb)  # shape [D]
x = (1 + gamma) * x + beta                # feature modulation

Difference from adaLN:

Advantages: intuitive "gating" semantics - "if action is turn left, enhance feature channels representing left-turn." Paper result: best on RT-1 (FID=40.62, PSNR=23.20)

⑤ Cross-Attention

How: at the end of the block, use video tokens as Q and action as K/V; cross-attention lets each visual position query relevant information from the action sequence:

x = x + cross_gate * cross_attn(norm_cross(x), action_emb)
# Q from video tokens, K/V from action embeddings

Advantages: most flexible - actions can influence each visual token in a non-local, content-adaptive way. Different patches can learn different information from the same action. Disadvantages:

5.4 Three latent space encoders

5.4.1 SD-VAE (Stable Diffusion VAE)

Property Value
Source stabilityai/sd-vae-ft-mse
Architecture Convolutional VAE (encoder + decoder)
Input size 256×256 RGB
Latent shape [4, 32, 32]
Downsampling factor 8×
Has decoder? ✅ Can reconstruct RGB
Latent characteristics Reconstruction-oriented - few channels (4), strong "material" quality, good pixel alignment

How it works: standard VAE structure - the encoder downsamples via convolutions, learning a compact Gaussian representation (mean + variance) in low-dimensional space; the decoder reconstructs the original image from that distribution. The latent is trained with KL regularization and MSE reconstruction loss.

Why used as world-model latent space? The 4-channel × 32×32 is a very small "token space" (1024 values total) but contains enough information to reconstruct a 256×256 image. The model performs diffusion prediction in this compact space, which is far more computationally efficient than pixel space.

Key result: SD-VAE is the only latent space that successfully supports MPC planning (PushT success 25%), because the model learns to use action conditioning.

5.4.2 Web-DINO (semantic feature encoder)

Property Value
Source facebook/webssl-dino300m-full2b-224
Architecture ViT-L/14 (Vision Transformer)
Input size 224×224 RGB
Latent shape [1024, 16, 16]
Downsampling factor 14× (patch size)
Has decoder? ❌ Encoder-only
Latent characteristics Semantic-oriented - many channels (1024), emphasizes object-level invariance

How it works: DINOv2 uses self-supervised learning (without labels) to make the model focus on the semantic content of images rather than pixel details. The same object under different views, lighting, etc. gets similar DINO features. Web-DINO is a large-scale variant of DINO.

Why used as world-model latent space? In theory, semantic features are better for planning - "I don't need to know the exact RGB value of each pixel of a chair; I only need to know there is a chair there." But experiments show that when using DINO features for diffusion, the model gets "lazy" - because DINO features of adjacent frames are already very similar, the model finds it can make acceptable predictions without using actions, so the action branch's weights converge to 0 (RMS=0.002).

Key result: planning success is 0%. The model learns to ignore actions.

5.4.3 V-JEPA 2.1 (video-pretrained feature encoder)

Property Value
Source Video-JEPA 2.1
Architecture ViT-L/16
Input size 384×384 RGB
Latent shape [1024, 16, 16]
Downsampling factor 16× (patch size)
Has decoder? ❌ Encoder-only
Latent characteristics Video-prediction-oriented - pretrained on large video data to learn video dynamics

How it works: V-JEPA (Video Joint Embedding Predictive Architecture) is a self-supervised video learning paradigm: given part of a video (visible blocks), it predicts the feature representations of masked parts. It is pretrained at the video level, so theoretically better at capturing temporal dynamics.

Why used as world-model latent space? Although V-JEPA is pretrained on video prediction, when used as a world-model latent space it performs similarly to Web-DINO - action embedding RMS is only 0.00129, and planning success is 0%. The paper's analysis suggests that the diffusion loss combined with a semantic latent space is insufficient to force the model to learn controllable dynamics. Possible reasons:

  1. Semantic feature space is inherently "smooth" - small action changes produce too little difference in semantic space
  2. Diffusion models tend to "fill in averages" during decoding rather than precise action conditioning

5.4.4 Systematic comparison of three encoders in action utilization ability

Encoder Action embedding RMS MSE difference (GT action vs random action) Planning success
SD-VAE 0.112 0.015 vs 0.081 (6x difference) 25%
Web-DINO 0.002 0.834 vs 0.834 (almost no difference) 0%
V-JEPA 2.1 0.001 0.584 vs 0.584 (almost no difference) 0%

This shows that it's not that semantic latent spaces are inherently bad, but rather that under the diffusion training framework, semantic latent spaces have difficulty learning action-conditioned dynamics. This is an important open problem in current world-model research.

Transformer-cover

Transformer-cover

0xFF References

This article is typeset using markdown.com.cn.

Comments

Top 1 from juejin.cn, machine-translated. The original thread is authoritative.

掘一二三桶金

This learning content is quite a lot, thanks for sharing.