Linear DiT / SANA: Diffusion Transformer with Linear Attention.

Implements DiT (Diffusion Transformer) architecture with linear attention replacing the quadratic softmax attention. This achieves comparable image quality at dramatically reduced computational cost.

Key Innovation: Linear Attention in Diffusion

Standard DiT uses O(N²) softmax attention, which becomes prohibitive for high-resolution images. Linear DiT replaces this with O(N) linear attention using kernel feature maps, enabling:

• 100x speedup for high-resolution generation
• Same quality as quadratic DiT
• Scalable to 4K+ resolution images

Architecture

Input [batch, num_patches, patch_dim]
      |
      v
+---------------------------+
| Patchify + Position Embed |
+---------------------------+
      |
      v
+---------------------------+
| Linear DiT Block x depth  |
|  AdaLN-Zero(condition)    |
|  Linear Attention         |  <- O(N) instead of O(N²)
|  Residual                 |
|  AdaLN-Zero(condition)    |
|  MLP                      |
|  Residual                 |
+---------------------------+
      |
      v
| Final AdaLN + Linear     |
      |
      v
Output [batch, num_patches, patch_dim]

Linear Attention Mechanism

Standard: Attn(Q,K,V) = softmax(QK^T/sqrt(d)) * V [O(N²)]

Linear: Attn(Q,K,V) = phi(Q) * (phi(K)^T * V) / (phi(Q) * sum(phi(K))) [O(N)]

Where phi(x) = ELU(x) + 1 ensures non-negative attention weights.

Usage

model = LinearDiT.build(
  input_dim: 64,
  hidden_size: 512,
  num_layers: 12,
  num_heads: 8,
  patch_size: 2
)

References

• SANA: "Efficient High-Resolution Image Synthesis with Linear Diffusion Transformer" (2024)
• DiT: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023)
• Linear Attention: "Transformers are RNNs" (Katharopoulos et al., 2020)