Block-Causal Diffusion Transformer (DiT)

• Block-Causal DiT is a novel design that decomposes data into blocks and enforces causal dependencies to scale diffusion models for long videos and high-resolution images.
• It leverages linearized attention kernels and block-wise autoregressive inference, as demonstrated in SANA-Video and Inf-DiT, to drastically reduce memory and computation costs.
• Empirical results show significant speedup and memory efficiency improvements, enabling state-of-the-art visual synthesis with linear or constant memory requirements.

The Block-Causal Diffusion Transformer (DiT) encompasses a class of efficiency-driven attention designs for diffusion models, enabling tractable, high-quality generation over extremely long temporal sequences (videos) or ultra-high-resolution spatial grids (images) by leveraging block decomposition, linearized attention kernels, and block-wise causal dependency graphs. Its key instantiations—through SANA-Video’s block-wise linear causal attention for video generation (Chen et al., 29 Sep 2025) and Inf-DiT’s unidirectional block attention for image super-resolution (Yang et al., 2024)—demonstrate block-causal DiTs as state-of-the-art solutions for large-scale visual synthesis using constant or linear memory.

1. Foundational Principles

Block-causal DiTs address the quadratic memory bottlenecks of conventional transformer self-attention, which impede scaling diffusion models to large video and image domains. The core principle is to tile the data domain (temporal or spatial) into blocks, and enforce causality constraints at the attention/conditioning level: each block can (directly or via attention) see only a bounded window of previously generated blocks, while information flows globally through stacking multiple attention/FFN layers.

The causality and block structure enable storing only a constant (or linear) number of compressed key/value caches. This design grounds both SANA-Video (Chen et al., 29 Sep 2025) (for long-form video) and Inf-DiT (Yang et al., 2024) (for image upsampling and super-resolution), while retaining global context and supporting large-scale autoregressive or diffusion sampling.

2. Block-Causal Attention Mechanisms

SANA-Video: Block Linear Causal Attention

SANA-Video replaces softmax-based attention with a linear kernel and block-wise autoregressive inference. The linear kernel is

$\phi(x) = \mathrm{ReLU}(x) + \epsilon$

applied to queries ($Q$) and keys ($K$). Rotary position encoding (3D RoPE) is applied after $\phi$ to ensure numerical stability, yielding representations such as $\mathrm{RoPE}(\phi(Q))$.

For a sequence of $N$ tokens, let $Q_i, K_i, V_i \in \mathbb{R}^{1\times D}$. Linear attention computes: $$O_i = \frac{ \mathrm{RoPE}(\phi(Q_i)) \left(\sum_{j=1}^{N} \mathrm{RoPE}(\phi(K_j))^{T} V_j \right) }{ \phi(Q_i) \left(\sum_{j=1}^{N} \phi(K_j)^{T}\right) }$$ The block-wise strategy chunks tokens into $B$ blocks of $M$ tokens, running diffusion per block sequentially and maintaining only two small buffer matrices for all prior context:

• $Q$0
• $Q$1

This constant-memory cache supports minute-long, globally coherent video synthesis.

Inf-DiT: Unidirectional Block Attention (UniBA)

Inf-DiT decomposes $Q$2 images into $Q$3 blocks of $Q$4 pixels. The unidirectional attention assigns for each target block $Q$5, direct attention to its “self,” “above,” “left,” and “above-left” neighbors, imposing a directed acyclic dependency graph (DAG) for block-wise autoregressive generation.

Formally, for a layer $Q$6:

• $Q$7 (hidden state per patch in block)
• Attention source blocks $Q$8 are concatenations of the four neighbor blocks, each with learnable block-level positional embeddings. Attention computation is softmax over all patches in these four blocks, optionally with QK-normalization.

Each layer’s restricted connectivity is $Q$9, yet stacking $K$0 such layers yields a global receptive field. Only $K$1 (number of column blocks plus current inference batch $K$2) blocks' KV-caches must be stored at any time, yielding linear, rather than quadratic, scaling in inference memory.

3. Theoretical Analysis of Memory and Computation

Block-causal DiTs achieve substantial reductions in memory and compute complexity compared to full self-attention:

Attention Scheme	Memory per token	Compute per token	Global Context
Full causal softmax	$K$3	$K$4	Yes (exact)
Block-local window	$K$5	$K$6	Limited ($K$7)
Causal linear (SANA-Video)	$K$8	$K$9	Yes (linear, ReLU+RoPE)
UniBA (Inf-DiT)	$\phi$0	$\phi$1	Yes (via DAG, depth)

• SANA-Video keeps attention memory at a few hundred megabytes even for minute-long sequences, since only $\phi$2 and $\phi$3 matrices are retained independent of sequence length.
• Inf-DiT with UniBA reduces memory for $\phi$4 images from $\phi$580 GB (U-Net) to $\phi$615 GB, a $\phi$75x reduction.

Both designs guarantee that global dependencies are conveyed either exactly (SANA-Video, via associative kernel) or approximately but reliably (Inf-DiT, via deep stacking over the block DAG).

4. Integration with DiT and Diffusion Pipelines

Block-causal attention is integrated in transformer-based diffusion models as follows:

• SANA-Video utilizes a Diffusion Transformer backbone initially designed for text-to-image, modified as:
  • All softmax attentions replaced by linear attention with ReLU kernel.
  • 3D RoPE for spatiotemporal position encoding.
  • Feed-forward MLP replaced by a Mix-FFN incorporating a 1D temporal convolution (kernel size 3) for local motion continuity.
  • Constant-memory block-wise autoregressive diffusion sampling.
• Inf-DiT upsampler processes images as:
  • Channel-wise concatenation of bicubic-upsampled LR image and noisy diffusion input, partitioned into $\phi$8 blocks, each further split into $\phi$9 patches.
  • Per-block linear embedding, 2D RoPE, and small learnable block-relative positional embeddings.
  • In the initial transformer layer, blocks use multi-head cross-attention with adjacent LR patches to enforce local sharpness.
  • Inference proceeds in $\mathrm{RoPE}(\phi(Q))$0 block batches, discarding unneeded KV-caches after each batch.

Hyperparameters in Inf-DiT (e.g., $\mathrm{RoPE}(\phi(Q))$1, $\mathrm{RoPE}(\phi(Q))$2, $\mathrm{RoPE}(\phi(Q))$3) optimize the memory/runtime trade-off; ablations confirm block-causal attention is essential for artifact-free upsampling at ultra-high resolutions.

5. Empirical Results and Performance Metrics

Block-causal DiTs demonstrate state-of-the-art performance on both video and image domains at unprecedented scale:

SANA-Video

• 5 s, $\mathrm{RoPE}(\phi(Q))$4 video generation in 36 s (BF16 on H100; 29 s with NVFP4 on RTX 5090).
• 4×–16× speedup over DiT baselines such as Wan-2.1-1.3B or Wan-2.1-14B.
• End-to-end memory $\mathrm{RoPE}(\phi(Q))$5 GB (H100 BF16), $\mathrm{RoPE}(\phi(Q))$6 GB (NVFP4).
• Quality: VBench total score 84.05 (vs. 83.28 for Wan-2.2-5B); FID low 20s, CLIP similarity $\mathrm{RoPE}(\phi(Q))$70.28.

Inf-DiT

• $\mathrm{RoPE}(\phi(Q))$8 image generation at $\mathrm{RoPE}(\phi(Q))$915 GB inference memory vs. $N$0 GB for SDXL or U-Net.
• Ultra-high-resolution FID (HPDv2, $N$1): 66.0 (Ours) vs. 66.3 (SDXL+BSRGAN), 67.5 (DemoFusion).
• DIV2K $N$2 super-resolution: FID = 20.2, PSNR = 26.3, SSIM = 0.74 (substantially improves over BSRGAN, Real-ESRGAN).
• Human evaluation: Inf-DiT ranked first on authenticity, global coherence, and LR consistency.

6. Approximation Properties and Practical Trade-Offs

Block-causal attention yields global context exposure, with primary deviation from traditional softmax attention arising from the kernel choice and block-level compression:

• SANA-Video’s ReLU kernel discards softmax normalization, incurring minor smoothing of attention peaks but preserving global semantics and temporal coherence.
• Inf-DiT’s blockwise DAG induces indirect (multi-layer) dependency between distant pixels; ablation removing cross-block attention reveals block seam artifacts and sharp FID/quality loss.
• For Inf-DiT, block size $N$3 governs a tradeoff between compute/memory per block and total number of cache entries: larger $N$4 reduces total block count but raises per-block memory and vice versa.

Empirical ablations confirm the necessity of nearby LR cross-attention and block-level CLIP guidance for both objective (FID) and subjective (human rating) improvement (Yang et al., 2024).

7. Influence, Extensions, and Future Directions

Block-causal DiT architectures as exemplified by SANA-Video and Inf-DiT are now primary blueprints for efficient transformer-based visual generative modeling at scale, enabling deployment of multi-minute video synthesis and gigapixel super-resolution within commercially viable memory/computation footprints.

Potential avenues for further research include:

• Exploration of alternative kernelized attention functions or hybrid block/local attention for improved sharpness or long-range alignment.
• Generalization to non-rectangular domains and multi-modal tasks.
• Dynamic block partitioning or learned block graphs to optimize context exposure adaptively during inference.
• Compression/quantization advances, such as NVFP4 precision, to further reduce inference costs.

The block-causal DiT paradigm constitutes a foundational advance in transformer-based generative models, directly enabling practical large-scale applications and setting new standards across video and super-resolution benchmarks (Chen et al., 29 Sep 2025, Yang et al., 2024).