MSDiT: Masked Spatial-Temporal Diffusion Transformer

• MSDiT is a unified video diffusion model that employs spatial-temporal masks and a pure Transformer backbone to tackle tasks like unconditional generation, prediction, interpolation, and completion.
• It tokenizes videos using a frozen VAE and processes them via alternating temporal and spatial self-attention layers with adaptive LayerNorm for effective diffusion time-conditioning.
• Its dual-mask extension, Mask²DiT, enhances fine-grained, multi-scene video synthesis and achieves superior performance metrics such as lower FVD and improved consistency on benchmark datasets.

The Masked Spatial-Temporal Diffusion Transformer (MSDiT) is a class of video generative models that integrate modularized spatial-temporal Transformer architectures with a unified mask-based conditioning mechanism to support general-purpose video generation tasks. Crucially, MSDiT leverages a binary spatial-temporal mask to seamlessly address tasks such as unconditional video generation, prediction, interpolation, and spatio-temporal completion within a single framework. Recent extensions also employ dual-mask designs to enable fine-grained segment-level alignment in multi-scene video generation.

1. Architectural Foundations

MSDiT is grounded in a pure Transformer-based backbone, where each video is first tokenized via a frozen VAE. For a video $$\mathbf{V}\in\mathbb{R}^{F\times H\times W\times 3}$$, the tokenizer produces a latent $$\mathcal{F}\in\mathbb{R}^{F\times(H/8)\times(W/8)\times C}$$. Spatial $N\times N$ patches are flattened to $C$-dimensional tokens and stacked temporally, yielding a sequence of $L=F\cdot\frac{H/8\times W/8}{N^2}$ tokens. Each token is enriched by learnable linear projections and sinusoidal spatial and temporal positional encodings.

MSDiT blocks alternate modular Multi-Head Self-Attention (MHSA) modules—temporal and spatial—plus an MLP, employing adaptive LayerNorm (adaLN) to inject diffusion time-conditioning:

• Temporal MHSA: $$\mathbf{Z}_{\mathrm{temp}} = \mathrm{MHSA}_{\mathrm{temp}}(\mathrm{adaLN}(\mathbf{H}, t))$$
• Spatial MHSA: $$\mathbf{Z}_{\mathrm{spat}} = \mathrm{MHSA}_{\mathrm{spat}}(\mathrm{adaLN}(\mathbf{Z}_{\mathrm{temp}}, t))$$
• Feed-forward: $$\mathbf{H}' = \mathrm{MLP}(\mathrm{adaLN}(\mathbf{Z}_{\mathrm{spat}}, t))$$ where $$\mathrm{adaLN}(\mathbf{x}, t) = \gamma(t)\,\mathrm{LayerNorm}(\mathbf{x}) + \beta(t)$$.

Architectural variants include MSDiT-Small (12 layers, hidden=384, 6 heads, $\sim$28M parameters) and MSDiT-Large (28 layers, hidden=1152, 16 heads, $\sim$596M parameters) (Lu et al., 2023).

2. Unified Mask Modeling and Task Generalization

At the core of MSDiT is a unified spatial-temporal binary mask $\mathcal{M}\in\{0,1\}^L$ specifying whether each token is a “real” conditional token or a pure noise token. For input $\mathcal{I}$:

$$\mathcal{I} = \mathcal{F}\odot (1-\mathcal{M}) + \mathcal{C}\odot \mathcal{M}$$

where $\mathcal{F}$ are noise tokens, $\mathcal{C}$ are conditional tokens, and $\odot$ denotes element-wise multiplication.

Masking regimes flexibly encode a wide set of tasks:

• Unconditional generation: $\mathcal{M}\equiv 0$
• Video prediction: $\mathcal{M}=1$ on first $k$ frames, $0$ elsewhere
• Interpolation: Arbitrary interior frames set to $0$
• Image$\to$Video: Only a single frame revealed, rest zeroed
• Spatial-temporal completion: Block-wise random masking (analogous to BEiT pretext)

In training, $\mathcal{M}$ is sampled from these regimes, so the model learns to denoise under any arbitrary spatio-temporal missing pattern.

3. Diffusion Process and Training Objective

MSDiT follows the discrete-time diffusion framework. At step $t$:

• Forward process: $$q(\mathbf{x}_t|\mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{\alpha_t}\mathbf{x}_{t-1}, (1-\alpha_t)\mathbf{I})$$
• Reverse process: $$p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t, \mathcal{I}), \Sigma_\theta(t))$$

The noise prediction parameterization is optimized using MSE:

$$\mathcal{L}_{\mathrm{diff}} = \mathbb{E}_{t,\mathbf{x}_0,\epsilon}\left\| \epsilon - \epsilon_\theta(\sqrt{\bar\alpha_t}\mathbf{x}_0 + \sqrt{1-\bar\alpha_t}\epsilon, t, \mathcal{I}) \right\|^2$$

where $\bar\alpha_t = \prod_{s=1}^t\alpha_s$ (Lu et al., 2023).

4. Conditioning and Zero-Shot Generalization

MSDiT implements token-level concatenation: conditional tokens are concatenated with noise tokens to form a single sequence, processed jointly by the spatio-temporal Transformer backbone. Sinusoidal positional encodings permit zero-shot generalization to variable-length conditioning; e.g., an MSDiT trained on 8 conditioning frames generalizes to longer contexts at test time without retraining.

No extra loss is applied for the mask itself; the binary mask only determines which tokens contribute to the diffusion loss, further unifying the training process.

5. Dual Masked Approaches for Multi-Scene Generation

The Mask$^2$DiT extension incorporates two explicit masking mechanisms to address multi-scene video generation (Qi et al., 25 Mar 2025):

• Symmetric Binary Mask ($M^s_l$): Inserted at each attention layer to block cross-segment text-video interactions, ensuring each text prompt attends only to its respective video segment and audiovisual tokens freely interact.
• Segment-Level Conditional Mask ($M^c_l$): Applied during staged decoding (auto-regressive scene extension), allowing a segment to attend to all prior segments while preventing lookahead.

This dual masking provides one-to-one textual-visual alignment per scene and supports both fixed-length and variable-length (via auto-regressive extension) multi-scene video synthesis.

Dual Mask Attention Schema (summarized)

Attention Query \ Target	Same Scene Tokens	Visual Tokens (all)	Cross-Scene (Text/Video)
Text Query	Attend (0)	No	Blocked ($-\infty$)
Visual Query	Attend (0)	Attend (0)	Attend (0)

6. Empirical Evaluations and Benchmarks

MSDiT achieves strong performance on unconditional and conditional generation benchmarks:

• UCF-101 (unconditional, 16$\times$64$\times$64): MSDiT-L FVD = 225.7 versus VDM (295.0) and MCVD (1143.0)
• Cityscapes (cond 2$\to$pred 28): MSDiT-L FVD = 142.3, SSIM = 0.880; MCVD-concat FVD = 141.4, SSIM = 0.690
• Physion video prediction (cond 8$\to$pred 8): VQA accuracy 65.3% (scene-centric baseline 63.1%) (Lu et al., 2023)

Ablations indicate token concatenation conditioning outperforms cross-attention and adaLN-only strategies (FVD = 129.1 vs. 134.9 and 270.8, respectively). Spatial-only pretraining followed by joint tuning is preferred to direct spatio-temporal training.

In the multi-scene context, Mask$^2$DiT attains the best Visual Consistency (70.95%), strong Sequence Consistency (47.45%), and lowest FVD (720.01) on a 5,371-video test set compared to DiT and U-Net-based architectures (Qi et al., 25 Mar 2025).

7. Task Scope, Applicability, and Extensions

MSDiT and its extensions offer a unified interface for:

• Unconditional video synthesis
• Video prediction and extrapolation
• Arbitrary spatio-temporal interpolation
• Spatio-temporal completion
• Multi-modal and flexible conditioning, including variable-length and zero-shot settings

Qualitative evidence demonstrates spatially and temporally coherent outputs, accurate physics/dynamics simulation, and robust generalization, particularly under the dual-mask regime for multi-scene long-form video. The mask modeling mechanism is readily extensible to more complex data modalities (e.g., text-video).

Summary Table: Notable MSDiT and Mask$^2$DiT Results

Model	Scenario	Visual Con.	Seq. Con.	FVD
MSDiT-L	UCF-101, unconditional	—	—	225.7
MSDiT-L	Cityscapes, cond 2$\to$pred 28	—	—	142.3
Mask$^2$DiT	Multi-scene test set (5,371 vid)	70.95%	47.45%	720.01

• "VDT: General-purpose Video Diffusion Transformers via Mask Modeling" (Lu et al., 2023)
• "Mask$^2$DiT: Dual Mask-based Diffusion Transformer for Multi-Scene Long Video Generation" (Qi et al., 25 Mar 2025)