08 · Sana 高效高分辨率生成

Sana 用 LLM text encoder、linear attention 和 SVDQuant 量化降低 DiT 的显存与计算开销，目标是在消费级 GPU 上生成高分辨率图像，包括 4K 档位。

1. 高效高分辨率问题的本质

高分辨率下的显存压力

扩散模型在高分辨率下的显存压力来自三个层面：

2. Sana 的三项取舍

设计 1：LLM as Text Encoder（用 Gemma-2B 替代 CLIP+T5）

为什么小 LLM 更好：LLM（Gemma-2B）的语言理解能力远强于 CLIP（后者只能做图文匹配）。2B 的小 LLM 比 11B 的 T5-XXL 小 5×，但文本理解不输。Sana 用 cross-attention（text tokens 作为 K/V，image tokens 作为 Q）注入文本条件，而非 MMDiT 的 joint attention。这意味着 text tokens 不参与自注意力，进一步降低注意力复杂度。

设计 2：Linear Attention（O(n) 替代 O(n²)）

O(n²) Softmax Attention 的问题：

# 标准 softmax attention
A = softmax(Q · K^T / √d) · V    # Q,K 都是 n×d
# 中间矩阵 QK^T 是 n×n → O(n²) 内存

# 对 4096px 图像：n = 65536
# QK^T = 65536 × 65536 × 2 bytes (fp16) ≈ 8.6 GB（不可行）

Linear Attention 的原理：

# Linear attention（kernel trick）
φ = ReLU  # 或 ELU+1, cos, 等非线性核

# 计算 K^T·V：d×d 矩阵，与 n 无关
KV = φ(K)^T · V         # shape: (d, d)，不依赖 n
A  = φ(Q) · KV          # shape: (n, d)

# 中间矩阵 KV 是 d×d
# 对任意分辨率：KV = 1536 × 1536 × 2 bytes ≈ 4.7 MB（可行）

内存复杂度：Linear attention 的中间 KV 矩阵是 d×d（d 是 hidden dim，通常 1536），与 token 数 n 无关。它去掉了随分辨率按 O(n²) 增长的 QK^T 矩阵；Q、K、V 和输出激活仍按 O(nd) 增长。

设计 3：SVDQuant 4-bit 权重量化

Sana 配合 MIT HAN Lab 的 SVDQuant（基于奇异值分解的量化）将 DiT 权重从 fp16 压缩到 int4：

SVDQuant 的质量损失 < 1%（按官方 benchmark），在推理中几乎不可见。

3. Sana 的推理数据流

prompt
  │
  └─→ Gemma-2B tokenizer → Gemma-2B (2304d) → text embeddings
       pooled: last token hidden state (B, 2304) → adaLN
  │
  ▼
noise latent z₁ ~ N(0, I)  shape: (1, C, H/8, W/8)
  │
  ▼
denoising loop: t = 1 → 0 (20~40 步, Sprint 仅 2 步)
  ├─ patchify(p=2) → (B, N, D)
  ├─ Linear self-attention: φ(Q) · (φ(K)^T · V)
  ├─ Cross-attention: Q=image tokens, K/V=text tokens
  ├─ FFN + AdaLN modulation
  ├─ CFG: v_cfg = v_uncond + s · (v_cond − v_uncond)
  └─ Euler step
  │
  ▼
AE decoder → pixel image

4. 显存友好路径

# 受限显存示例：SVDQuant int4 Sana-1.6B
pipe = SanaPipeline.from_pretrained(
    'mit-han-lab/SVDQuant-int4-Sana-1.6B',  # 量化版
    torch_dtype=torch.float16
)
pipe = pipe.to('cuda')
# 此配置不使用 offload
image = pipe('一只柴犬在樱花树下', num_inference_steps=20, guidance_scale=4.5).images[0]
# VRAM < 8 GB（1024px）；可进一步测试 4096px

5. Sana 与其他图像模型的对比