企业官网建设流程全解析

这篇论文做的不是推理阶段减少采样步数，而是训练阶段决定应该训练哪个时间步。

引言：

普通扩散模型训练时，一般从 1∼T 中均匀随机抽取时间步 t。作者认为这样不合理，因为不同时间步的训练难度不同：有些时间步梯度波动大、收敛慢，却和简单时间步得到相同的训练次数。

梯度方差不同：梯度就是优化方向，不同图片在这一步的优化方向波动很大，不容易收敛（某个时间步的梯度方差越大，它达到相同精度所需要的迭代次数越多，另外adam对梯度方差大时候学习率也会降低），而均匀采样给的迭代次数一致，需要使用不用的迭代次数。

是否可以按梯度方差来采样？不可以，因为不同时间步之间存在依赖（作者只采样1-200T，发现为采样部分loss变大），梯度方差只解释这一步难不难训练，所以作者想找哪一个时间步对loss有用

于是作者额外训练一个时间步采样器，动态判断当前阶段训练哪个 t，最能降低整个扩散模型的目标函数。

贡献点：1、从梯度方差解释非均匀训练为什么有效

2、提出基于学习的自适应时间步采样方法

3、验证方法的鲁棒性

相关工作：

• 扩散模型的非均匀时间步训练

（1）Weighting methods损失加权方法

时间步采样概率相同，权重不同

（2）Sampling methods 采样方法

改变采样概率

以上都是启发式方法，基于经验给一个规则，但不一定在所有情况下最合适，因此作者希望用学习方法动态找到更合适的时间步分配策略。

• Learning to Optimize，学习优化器

Learning to Optimize则希望训练一个模型，让它自己学习：给定当前梯度、损失和训练状态，下一步应该怎样更新参数。adam

损失函数loss：反向分布 q 和模型学习分布 pθ​ 之间的差异

正文：

5.1 Non-uniform Timestep Sampler：怎么定义和训练时间步采样器

使用beta分布进行采样，虽然只有两个参数，但能表达多种平滑采样模式

dist_sampled = dist.sample()#从 Beta 分布中采样

Beta 分布由 alpha 和 beta 控制形状：
alpha 大, beta 小：更偏向 1，也就是偏向后期 timestep
alpha 小, beta 大：更偏向 0，也就是偏向前期 timestep
alpha ≈ beta：比较居中
alpha = beta = 1：接近均匀分布

在时间步 t 上更新一次后，整个VLB下降多少，对所有时间步求平均。让采样器更倾向于选择高收益时间步。但是极端成本太高，使用下面的方法：

收益的计算依然开销很大，使用5.2的近似。

5.2 Approximation of Δkt​：怎么低成本估计一个时间步的训练价值

作者认为相邻时间步的 loss 变化存在较强相关性，所以没有必要每次都检查1000个。

只用3个代表性时间步，近似1000个时间步的整体变化，但是需要选择有代表性的三步。代码中用SelectKBest(f_regression, k=3)

代码流程：

train.py文件

1、解析命令行参数

2、创建 replay buffer

replay_buffer = SharedMemoryManager(capacity=args.capacity, world_size=args.num_gpus)

多个 GPU 各自计算 timestep 的 KL 改变量
-> 放进共享 replay buffer
-> 用于更新策略网络

3、判断是否多 GPU 启动，运行train() 函数

train() 函数

1、读取配置文件和基本信息、设置随机种子、读取训练参数

配置文件中包含：
dataset
diffusion
model
train
ActorNetwork
ValueNetwork参数 todo

2、读取 diffusion 配置并创建 DDPM

（1）生成 beta schedule：

betas = get_beta_schedule(...) 确定前向加噪公式的系数

（2）创建扩散模型过程：

diffusion = GaussianDiffusion(betas=betas, **diffusion_config)

里面写了：

• q_sample()：训练时加噪
• train_losses()：计算训练 loss
• p_sample()：生成图片时逐步去噪
• q_posterior_mean_var()/p_mean_var()：计算真实后验和模型后验，用于 KL、采样和评估

3、创建 UNet 主模型

_model = UNet(**model_config)

4、创建 ValueNetwork 和 ActorNetwork

ValueNetwork：用于价值估计/辅助策略训练，但是后面没有用到

ActorNetwork：根据图片生成 timestep 采样分布

5、创建三个优化器

optimizer = Adam(model.parameters(), lr=train_config.model_lr, ...) optimizer_pi = Adam(policy.parameters(), lr=train_config.policy_lr, ...) optimizer_v = Adam(value_function.parameters(), lr=train_config.value_lr, ...)

三个优化器分别更新：

• optimizer：UNet
• optimizer_pi：ActorNetwork
• optimizer_v：ValueNetwork

6、创建学习率调度器

scheduler = LambdaLR(...)
scheduler_v = LambdaLR(...)
scheduler_pi = LambdaLR(...)
用于学习率 warmup。

7、创建trainer = Trainer（）

8、创建评估器 Evaluator

FID 越低，一般说明生成质量越好。

Evaluator.eval()
↓
清空 Inception 统计器
↓
循环生成 eval_total_size 张图片
↓
每批图片送入 Inception 网络提特征
↓
累计生成图片特征的均值和协方差，生成训练集风格的图片
↓
读取真实数据集预计算的均值和协方差，直接下载
↓
计算 FID
↓
返回 {"fid": fid}

9、开始训练

trainer.train

负责按 epoch/batch 调用step()训练模型，并定期保存生成样例图、评估 FID、记录日志和同步多 GPU。

train()
|
v
设置采样图片网格
|
v
for 每个 epoch:
|
v
清空统计，模型进入 train 模式
|
v
遍历 trainloader:
取出 batch 图片 x
调用 step(x)主要是这里
更新进度条 loss
|
v
是否到保存图片的 epoch?
是 -> sample_fn() 生成图片并保存
|
v
是否到 FID 评估 epoch?e == 408 or e == 816 or e == 1224 or e==1632 or e == 2039:
是 ->计算fid
|
v
记录 wandb 日志
|
v
多 GPU 同步

for结束

def step(self, x, e, i, global_steps=1, logger=None):

作用：给当前 batch 用 ActorNetwork 选择 timestep
-> 用这些 timestep 训练 UNet 扩散模型
-> 计算训练前后 KL 变化
-> 把 KL 变化作为 reward
-> 更新 ActorNetwork 策略网络

获取 x 的batch size B 和 timestep 总数 T
|
v
每隔 update_policy ，论文给出40步，计算训练前的 KL
|
|-- 是：

| 1、按 GPU 数量划分 timestep，例如0-499,500-999，得到range_T
| 2、从主进程和x中随机选一张样本图，复制range_T份
| 3、计算1000时间步样本图训练前 KL before （每块gpu计算自己range_T的kl）

| kl_before_for_lasso一张图1000的kl取平均，用来更新3个时间步怎么选择，见下面kl计算（1）
| 4、计算当前 batch所有图片 在指定 timestep 上的 KL before kl_divergence_tensor_before用来计算reward 见下面kl计算（2）
v
ActorNetwork 根据 x 采样 timestep sampled_tself.sample_timesteps(x) 见下面
|
v
用 sampled_t 训练 UNet（根据 sampled_t 对图片加噪、UNet 预测噪声、计算 MSE loss、反向传播计算梯度）
|
v
更新 UNet
|
v
记录loss
|
v
每隔 update_policy 步，计算训练后的 KL
|
|-- 是：
| 1、计算样本图训练后这1000个时间步计算 KL after kl_after_for_lasso
| 2、KL diff lasso = KL before - KL after，在把1000求和
| 3、把KL diff lasso写入 replay buffer用于共享

| 4、计算当前 batch所有图片 在指定 timestep 上的 KL after

| 5、kl_diff = kl_divergence_tensor_before - kl_divergence_tensor_after，再求和[图1kl,...图128kl]
| 6、根据第五步就是reward，reward += self.ent_coef * entropy，鼓励保持一定随机性不要过早固定采样策略，reward是长度为 batch_size 的 tensor
| 7、根据第三步历史数据Q更新重要 timestep：哪些 timestep 的 KL 变化更能代表整体训练收益就把这些 timestep 选出来，利用 SelectKBest算法，用f_regression某个 timestep 的 KL 改变量越能解释整体 KL 改善，它的分数越高。论文给的Q是20。
| 8、reward怎么更新采样器：actor_loss = -log_prob * reward 采样概率*kl改善量，再求平均
| 9、更新 ActorNetwork （反向传播）

高reward->采样概率高

kl散度的计算：

（1）def compute_singlestep_KL(self, x, sampled_t)

算出一张图片在1000个时间步的每个像素的kl散度，每个像素得到1000个kl，对每个像素取平均

1、self.diffusion.train_losses

用ddpm正向加噪获得加噪图片sampled_xt，并可计算多种loss

2、self.compute_kl_divergence：

计算真实后验 q(x_{t-1}|x_t,x_0)：已知 x_t 和 x_0，推断上一时刻 x_{t-1}

用 UNet 预测模型后验 p_theta(x_{t-1}|x_t)：反向过程

计算 KL(q || p_theta)：得到差值

每个像素取平均计算每张图片的 KL 值

（2）kl_divergence_tensor_before = self.calculate_kl_for_all_x0_at_t(self.non_zero_coef_timesteps, x)

对 batch 里的每一张图片分别在指定的 timestep 上计算 KL divergence

初始化是012，后面会对non_zero_coef_timesteps进行更新

[
[图1t0kl, 图1t1kl, 图1t2kl],
[图2t0kl, 图2t1kl, 图2t2kl],
[图3t0kl, 图3t1kl, 图3t2kl],
...
]

def sample_timesteps(self, x):

batch x 里有几张图，timestep 里就有几个数字。

作用：用策略网络 ActorNetwork 为当前 batch 的图片自适应采样 timestep。

输入图片 x
|
v
ActorNetwork(x) 输出 alpha, beta
|
v
Beta(alpha, beta)创建beta分布
|
v
从 Beta 分布中采样 dist_sampled ∈ [0, 1]
|
v
映射到 [0, self.timesteps - 1]
|
v
round + long 得到整数 timestep
|
v
计算 log_prob 和 entropy

当前采样动作的 log 概率，策略网络采样出这个 timestep 的概率有多大

分布的熵，表示采样分布有多“分散”。
|
v
返回 timestep、log_prob、entropy 等

文章的问题：

1、Beta分布表达能力有限

2、模型最后是让fid小，采样是让vlb小，指标不一致

3、reward方差可能很大这么避免，加baseline或者正则化？计算reward的时候只是根据当前 step 中 UNet 更新前后在 3 个代表性 timestep 上的 KL 差异，应该加入之前的历史记录，但是不能太多记录unet一直更新，采用短期历史平滑加自适应遗忘。

4、3个时间步的选择可能不准确（1）3这个数字（2）代表时间步选择方法

文章做了消融实验对比1、3、10，但依据不足。

可以迁移的：

训练一个时间步选择器，选择某个时间步后整体loss下降。

采样器参数ϕ→选择 t→更新 SDθ→评价 Δkt​→更新 ϕ

迁移小集合 S：降低时间步评价成本

选择损失计算的时间步：记录每次时间步更新后的损失，根据此损失选择3个代表时间步去计算reward，不是只根据这一步的数值，防止偶然性影响（假设20次是有相关性，根据unet来定的。

DDPM 训练完整地记成

DDPM采样/生成阶段，也有前向传播，每一次用unet预测噪声，没有反向。