扩散模型实现约定

前置

• 离散时间扩散
• EDM

1. 分离状态层级、转移步与数组下标

设离散扩散共有 $N$ 次转移，对应 $N+1$ 个状态层级：

$$0=t_0<t_1<\cdots <t_N=T.$$

数学状态下标为

$$i\in \{0,\dots ,N\},$$

其中 $i=0$ 表示干净数据，$i=N$ 表示终端噪声。第 $i$ 次前向转移是

$${\mathbf{x}}_{i-1}⟶{\mathbf{x}}_i,i\in \{1,\dots ,N\}.$$

实现中必须把两类数组分开：

1. 状态级数组长度为 $N+1$，代码下标 state_idx = i，可直接表示
   $t_i$、${\bar{\alpha }}_i$、${\sigma }_i$ 等状态量；
2. 转移级数组长度为 $N$，代码下标
   step_idx = i - 1，对应数学转移 $i\in \{1,\dots ,N\}$ 和参数
   ${\beta }_i$。

因此：

• step_idx = 0 对应从 $t_0$ 到 $t_1$ 的第一步，而不是状态 $t_0$ 或 $t_1$ 本身；
• step_idx = N - 1 对应从 $t_{N-1}$ 到 $t_N$ 的第 $N$ 步；
• state_idx = 0 对应 $t_0$；
• state_idx = N 对应 $t_N$。

原有“idx = 0 既对应 $t_1$ 又改成对应 $t_0$”的矛盾，来自把转移数组下标和状态数组下标共用同一个 idx。本文禁止这种未标类型的写法。

2. DDPM 数组约定

数学上定义

$${\beta }_i\in (0,1),i=1,\dots ,N,$$ $${\alpha }_i:=1-{\beta }_i,$$ $${\bar{\alpha }}_0:=1,{\bar{\alpha }}_i:=\prod _{j=1}^i{\alpha }_j.$$

推荐实现：

• betas.shape == (N,)，betas[step_idx] 表示
  ${\beta }_i$，其中 i=\texttt{step_idx}+1；
• alphas.shape == (N,)，使用同一转移下标；
• alpha_bars.shape == (N + 1,)；
• alpha_bars[0] = 1；
• alpha_bars[i] = prod(alphas[:i])，$i=1,\dots ,N$。

这样任意状态都可直接写成

$${\mathbf{x}}_i=\sqrt{{\bar{\alpha }}_i}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_i}\mathbf{ϵ}.$$

特别地，$i=0$ 自动给出 ${\mathbf{x}}_0$，不需要伪造
${\beta }_0=0$。若某个旧接口坚持使用长度 $N+1$ 的 betas_state，必须明确
betas_state[0] 只是 padding，不是一次实际转移。

2.1 任意两状态之间

对 $0\le a<b\le N$，

$${\mathbf{x}}_b\mid {\mathbf{x}}_a\sim \mathcal{N}\left(\sqrt{\frac{{\bar{\alpha }}_b}{{\bar{\alpha }}_a}}{\mathbf{x}}_a,\left[1-\frac{{\bar{\alpha }}_b}{{\bar{\alpha }}_a}\right]\mathbf{I}\right).$$

这里 $a,b$ 都是 state_idx，不是 step_idx。当 $a=0$ 时
${\bar{\alpha }}_a=1$，第一步和任意时刻扰动都无需特殊边界补丁。

3. 连续时间与 EDM 网格

连续时间模型的训练时间 $t$ 可以直接采样为实数，不应强行解释成 DDPM 的转移数组下标。若数值求解使用网格

$$t_0<t_1<\cdots <t_N,$$

则网格仍有 $N+1$ 个状态点和 $N$ 个积分区间。

EDM 采样通常按噪声从大到小存储：

$${\sigma }_0={\sigma }_{\max }>{\sigma }_1>\cdots >{\sigma }_{N-1}={\sigma }_{\min }>{\sigma }_N=0.$$

此时：

• sigma_idx = 0 是采样起点的最大噪声；
• sigma_idx = N 是最终零噪声；
• 第 solver_idx = i 个求解区间是
  ${\sigma }_i\to {\sigma }_{i+1}$；
• 该数组方向与正向物理时间是否递增无关，变量名应写
  sigma_idx 或 solver_idx，不要复用 DDPM 的 step_idx。

若前向 VE 过程在连续时间 $t_a<t_b$ 上满足

$${\mathbf{X}}_{t_b}={\mathbf{X}}_{t_a}+\sqrt{\sigma (t_b{)}^2-\sigma (t_a{)}^2}\mathbf{ϵ},$$

需确保根号内非负，并明确 $\sigma (0)$ 是否为零。若
$\sigma (0)>0$，从干净数据到首个噪声层的方差约定不能被数组下标掩盖。

4. 宏观形状与微观形状

张量完整形状分为：

• Macro shape：批次和样本组织维度；
• Micro shape：单个样本内部的特征维度；
• Full shape：Macro + Micro。

常规固定长度批次中：

• 文本表示形状 $(B,L,D)$：
  macro 为 $(B,)$，micro 为 $(L,D)$；
• 图像形状 $(B,C,H,W)$：
  macro 为 $(B,)$，micro 为 $(C,H,W)$；
• 单个向量批次 $(B,D)$：
  macro 为 $(B,)$，micro 为 $(D,)$。

所有按样本采样的时间、噪声标准差、损失权重等最初形状通常都是
$(B,)$。在与数据做逐元素运算前，应扩展为

$$(B,\underset{\text{micro ndim}}{\underbrace{1,\dots ,1}}).$$

例如图像时间从 $(B,)$ 变为 $(B,1,1,1)$，向量批次从
$(B,)$ 变为 $(B,1)$。

5. Batch flattening

变长数据常把批次内各样本的首个微观维度拼接。设批次有 $B$ 个样本，第
$i$ 个样本长度为 $n_i$，其余微观形状为
$(d_1,\dots ,d_k)$。令

$$M:=\sum _{i=1}^B{n}_i.$$

拼接后数据形状为

$$(M,d_1,\dots ,d_k).$$

本文约定：

• 逻辑 macro shape 仍为 $(B,)$；
• 被拼接的变长轴由 $M$ 表示；
• flatten 后供基础网络处理的逐元素 batch 轴为 $(M,)$；
• 剩余 micro shape 为 $(d_1,\dots ,d_k)$。

需要一个归属映射

$$batch\_index\in \{0,\dots ,B-1{\}}^M,$$

其中每个样本编号 $i$ 重复 $n_i$ 次。若逐样本时间为

$$\mathbf{t}\in {\mathbb{R}}^B,$$

先映射为

$${\mathbf{t}}_{\mathrm{flat}}=\mathbf{t}[batch\_index]\in {\mathbb{R}}^M,$$

再扩展为

$$(M,\underset{k}{\underbrace{1,\dots ,1}}).$$

不能仅凭 reshape 把 $(B,)$ 变成 $(M,)$；必须根据
batch_index、lengths 或等价的 segment 元数据重复各样本值。

6. 时间条件的形状变换

6.1 常规 batch

输入 $\mathbf{x}$ 形状为 $(B,C,H,W)$：

1. 初始时间 t.shape == (B,)；
2. 网络条件可保持 (B,) 交给时间嵌入层；
3. 若要与 $\mathbf{x}$ 广播运算，扩展为 (B,1,1,1)。

不要为了基础网络的时间嵌入接口而永久覆盖广播后的张量；推荐分别命名
t_condition 与 t_broadcast。

6.2 flattening batch

输入 $\mathbf{x}$ 形状为 $(M,d_1,d_2)$：

1. 初始逐样本时间 t_sample.shape == (B,)；
2. 通过 batch_index 映射为
   t_flat.shape == (M,)；
3. 传给基础网络的时间条件通常是 (M,)；
4. 与输入逐元素运算时扩展为 (M,1,1)。

如果 hook 负责第 2 步，接口契约必须注明输入是逐样本还是逐 flatten 元素，避免上层和 hook 各重复一次。

7. 广播辅助函数的契约

一个可靠的 append_dims(value, target_ndim) 应满足：

1. 只在末尾添加 singleton 维；
2. 若 value.ndim > target_ndim，直接报错；
3. 不改变 batch/flatten 轴长度；
4. 不负责把 $(B,)$ 映射为 $(M,)$；该步骤由显式
   batch_index 完成。

数学上，

$$append\_dims:(m_1,\dots ,m_r)\mapsto (m_1,\dots ,m_r,\underset{q-r}{\underbrace{1,\dots ,1}}),$$

其中 $q$ 是目标维数。

8. 最小命名规范

• state_idx：离散状态层级，范围 $0,\dots ,N$；
• step_idx：DDPM 转移数组，范围 $0,\dots ,N-1$，对应数学步
  i=\texttt{step_idx}+1；
• solver_idx：数值积分区间；
• sigma_idx：EDM 噪声级数组；
• t_continuous：连续时间实数；
• t_sample：形状 $(B,)$ 的逐样本时间；
• t_flat：形状 $(M,)$ 的 flatten 后时间；
• t_broadcast：追加 singleton 维后的广播张量。

只要这些对象不都缩写成 t 或 idx，$t_0$ 边界、第一步噪声、累计系数和 flattening 时间对齐就不会再互相矛盾。