张芷铭的个人博客

扩散模型与 Transformer 的融合通过全局建模能力和灵活条件注入显著提升生成质量与可控性。

结合机制

主干网络替换(DiT)

DiT(Diffusion Transformer)直接用 Transformer 替代传统 U-Net,形成纯 Transformer 架构:

  1. 图像分块与嵌入:将输入分割为 patch,通过线性投影转化为 token 序列
  2. 位置编码:正弦-余弦位置编码赋予空间位置信息
  3. 时间步与条件嵌入:通过 MLP 或交叉注意力注入 Transformer 层
  4. 多层 Transformer 块:多头自注意力 + FFN 捕捉全局依赖

混合架构(U-Net + Transformer)

在 Stable Diffusion 等模型中协同工作:

  • U-Net 处理局部特征:卷积层提取多尺度细节
  • Transformer 增强全局建模:瓶颈层插入 Transformer 模块处理全局上下文
  • 交叉注意力注入条件:文本编码器通过交叉注意力与中间特征融合

训练流程

输入预处理

  • 图像分块:256×256 图像分割为 16×16 patch,得到 256 个 token
  • 时间步编码:正弦嵌入转换为连续向量
  • 条件编码:文本通过 CLIP 编码为语义向量

损失函数

噪声预测损失:

$$\mathcal{L} = \mathbb{E}{x_0,\epsilon,t} \left[ | \epsilon - \epsilon\theta(x_t, t) |^2 \right]$$

条件注入方式

方式机制特点
交叉注意力Query 来自图像,Key/Value 来自文本全局上下文建模
adaLN条件向量动态调整归一化参数轻量级注入
adaLN-Zero初始化为恒等函数训练稳定

推理优化

采样策略

  • DDIM 采样:非马尔可夫过程减少采样步数
  • 渐进式蒸馏:每步仅更新部分 token

并行推理

  • PipeFusion:图像分块分配到多设备,流水线并行
  • 模型压缩:知识蒸馏或剪枝减少参数

架构对比

模型主干网络条件注入应用场景
DiT纯 TransformeradaLN、交叉注意力高分辨率图像生成
Stable DiffusionU-Net + Transformer交叉注意力文本到图像生成
Sora时空 Transformer交叉注意力 + 时空自注意力长视频生成

挑战与方向

  1. 计算成本:Transformer 二次复杂度限制高分辨率生成,需稀疏注意力或硬件优化
  2. 多模态对齐:跨模态语义对齐需改进交叉注意力机制
  3. 实时交互:结合脑机接口或手势控制需优化推理延迟

Comments