流匹配与扩散模型-05：可扩展的模型架构

流模型和扩散模型需要用神经网络参数化向量场 $u_t^\theta(x|y)$ 。该函数包含三个输入和一个输出：

输入：空间向量 $x \in \mathbb{R}^d$ （图像像素或潜变量）；引导变量 $y \in \mathcal{Y}$ （如文本、类别等）；时间步长 $t \in [0, 1]$ 。
输出：预测的向量场 $u_t^\theta(x|y) \in \mathbb{R}^d$ 。

低维任务可直接拼接 $(x,y,t)$ 并使用 MLP。图像和视频生成维度高，需要能同时处理带噪视觉状态、提示词和时间信息的结构。

嵌入条件变量

简单模型可直接输入时间步 $t$ ；实际生成模型通常用傅立叶特征将 $t$ 嵌入高维空间，以表达不同噪声水平下的时间依赖：

$TimeEmb(t) = \sqrt{\frac{2}{d}} [ \cos(2\pi w_1 t) \dots \cos(2\pi w_{d/2} t), \sin(2\pi w_1 t) \dots \sin(2\pi w_{d/2} t) ]^T$

其中频率系数 $\omega_i$ 表示为：

$w_i = w_{min} \left( \frac{w_{max}}{w_{min}} \right)^{\frac{i-1}{d/2-1}}, \quad i=1, \dots, d/2$

该形式不是唯一选择，但能将时间映射到 $d$ 维隐空间，并保持 $||TimeEmb(t)||=1 \quad (sin^2 + cos^2 = 1)$ 。

若 $y_{raw} \in \{0, \dots, N\}$ 是类别标签，可为每个类别学习一个嵌入向量 $y$ ，并与向量场参数一起训练。

若 $y_{raw}$ 是文本提示词，通常用冻结的预训练文本模型将文本嵌入连续向量空间。常见选择：

CLIP：生成全局语义嵌入 $y = CLIP(y_{raw}) \in \mathbb{R}^{d_{CLIP}}$ 。
T5：提供细粒度的序列嵌入，形如 $PromptEmbed(y_{raw}) \in \mathbb{R}^{S \times k}$ ，允许模型通过注意力机制关注特定单词。

图像可表示为张量 $x \in \mathbb{R}^{C_{image} \times H \times W}$ 。DiT 使用注意力机制参数化向量场，设隐藏维度为 $d$ ，层数为 $L$ ，注意力头数为 $h$ 。

与 Vision Transformer 类似，DiT 将图像切成 Patch，映射为 token 序列，经 Transformer 处理后再 Depatchify 回原图尺寸。

结构流程：

Patch 化：输入一张图片张量 $x \in \mathbb{R}^{C \times H \times W}$ ，以 $P \times P$ 的 Patch 大小，产生 $N=(H/P) \cdot (W/P)$ 个 Patch；每个 Patch 做 Patchify 之后的的维度 $C' = CP^2$ 。表示为 $Patchify(x) \in \mathbb{R}^{N \times C'}$ 。
Patch 嵌入：学习一个矩阵 $W \in \mathbb{R}^{C' \times d}$ 来将每个 Patch Token 潜入到隐藏空间 $d$ 中。表示为 $\tilde{x}_0=PatchEmb(x)=Patchify(x)W \in \mathbb{R}^{N \times d}$ 。
时间/提示词嵌入： $\tilde{t} = TimeEmb(t) \in \mathbb{R}^d$ 和 $\tilde{y} = PromptEmbed(y) \in \mathbb{R}^{S \times d}$
输入 DiT：每层接收 $\tilde{x}_i, \tilde{t}, \tilde{y}$ ，并计算 $\tilde{x}_{i+1} = DiTBlock(\tilde{x}_i, \tilde{t}, \tilde{y}) \in \mathbb{R}^{N \times d}, \quad (i=0, \dots, L-1)$ 。Block 主要包含 Patch 自注意力、提示词交叉注意力和时间控制的 AdaLN。
解 Patch：学习一个矩阵 $\tilde{W} \in \mathbb{R}^{d \times C'}$ 来将 DiT 的输出映射回到一张图片，即： $u = Depatchify(\tilde{x}_N \tilde{W}) \in \mathbb{R}^{C \times H \times W}$ 。这个 $u$ 就是模型的输出，也就是我们需要预测的向量场 $u_t^\theta(x|y)$ 。

DiT Block 的核心组件：

缩放点积注意力：

$Attn(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_h}}\right)V \in \mathbb{R}^{N \times d_h}$
多头注意力：学习投影矩阵 $W_Q^{(h)}, W_K^{(h)}, W_V^{(h)} \in \mathbb{R}^{d \times d_h}$ 。定义 $head_h(x, z) = Attn(xW_Q^{(h)}, zW_K^{(h)}, zW_V^{(h)})$ ，其中 $z=x$ 为自注意力， $z=y$ 为交叉注意力。最后拼接并投影：

$\text{MultiHeadAttention}(x, z) = \text{Concat}(head_1, \dots, head_h)W_O \in \mathbb{R}^{N \times d}$
AdaLN 自适应层归一化：用 MLP $g: \mathbb{R}^d \rightarrow \mathbb{R}^{2d}$ 从时间嵌入 $\tilde{t}$ 预测 $(\gamma, \beta) = g(\tilde{t})$ 。初始化时令 $g$ 输出零，使 Block 初始接近恒等映射，提高训练稳定性。

$AdaNorm_{\tilde{t}}(x) = (1 + \gamma) \odot \text{Norm}(x) + \beta$

计算流程：

自注意力： $x \leftarrow x + g_{self}(\tilde{t}) \cdot \text{MultiHeadAttention}(AdaNorm_{\tilde{t}}(x), AdaNorm_{\tilde{t}}(x))$
交叉注意力： $x \leftarrow x + g_{cross}(\tilde{t}) \cdot \text{MultiHeadAttention}(AdaNorm_{\tilde{t}}(x), \tilde{y})$
前馈网络： $x \leftarrow x + g_{MLP}(\tilde{t}) \cdot MLP(AdaNorm_{\tilde{t}}(x))$ (注： $g_{\dots}(\tilde{t})$ 为学习到的门控参数)。

U-Net 是扩散模型中的另一类主流架构，本质是卷积神经网络，特点是输入与输出形状一致。

在扩散模型中，我们需要构建一个参数化的向量场：

$x \mapsto u_t^\theta(x|y)$

在固定 $y,t$ 时，输入 $x$ 和输出 $u$ 都是图像形状，因此 U-Net 适合参数化该向量场。

典型 U-Net 包含：

以 $256 \times 256$ RGB 图像为例，输入维度为 $(3, 256, 256)$ ：

输入阶段： $x_t^{input} \in \mathbb{R}^{3 \times 256 \times 256}$ ，表示带噪图像或潜变量。
编码压缩阶段： $x_t^{latent} = \mathcal{E}(x_t^{input}) \in \mathbb{R}^{512 \times 32 \times 32}$ ，分辨率降低、通道数增加。
中间处理阶段： $x_t^{latent} = \mathcal{M}(x_t^{latent}) \in \mathbb{R}^{512 \times 32 \times 32}$ ，维度保持不变。
解码还原阶段： $x_t^{output} = \mathcal{D}(x_t^{latent}) \in \mathbb{R}^{3 \times 256 \times 256}$ ，上采样恢复空间分辨率和通道数。