这篇论文做的不是推理阶段减少采样步数而是训练阶段决定应该训练哪个时间步。引言普通扩散模型训练时一般从 1∼T 中均匀随机抽取时间步 t。作者认为这样不合理因为不同时间步的训练难度不同有些时间步梯度波动大、收敛慢却和简单时间步得到相同的训练次数。梯度方差不同梯度就是优化方向不同图片在这一步的优化方向波动很大不容易收敛某个时间步的梯度方差越大它达到相同精度所需要的迭代次数越多另外adam对梯度方差大时候学习率也会降低而均匀采样给的迭代次数一致需要使用不用的迭代次数。是否可以按梯度方差来采样不可以因为不同时间步之间存在依赖作者只采样1-200T发现为采样部分loss变大梯度方差只解释这一步难不难训练所以作者想找哪一个时间步对loss有用于是作者额外训练一个时间步采样器动态判断当前阶段训练哪个 t最能降低整个扩散模型的目标函数。贡献点1、从梯度方差解释非均匀训练为什么有效2、提出基于学习的自适应时间步采样方法3、验证方法的鲁棒性相关工作扩散模型的非均匀时间步训练1Weighting methods损失加权方法时间步采样概率相同权重不同2Sampling 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也就是偏向后期 timestepalpha 小, beta 大更偏向 0也就是偏向前期 timestepalpha ≈ beta比较居中alpha beta 1接近均匀分布在时间步 t 上更新一次后整个VLB下降多少对所有时间步求平均。让采样器更倾向于选择高收益时间步。但是极端成本太高使用下面的方法收益的计算依然开销很大使用5.2的近似。5.2 Approximation of Δkt​怎么低成本估计一个时间步的训练价值作者认为相邻时间步的 loss 变化存在较强相关性所以没有必要每次都检查1000个。只用3个代表性时间步近似1000个时间步的整体变化但是需要选择有代表性的三步。代码中用SelectKBest(f_regression, k3)代码流程train.py文件1、解析命令行参数2、创建 replay bufferreplay_buffer SharedMemoryManager(capacityargs.capacity, world_sizeargs.num_gpus)多个 GPU 各自计算 timestep 的 KL 改变量- 放进共享 replay buffer- 用于更新策略网络3、判断是否多 GPU 启动运行train() 函数train() 函数1、读取配置文件和基本信息、设置随机种子、读取训练参数配置文件中包含datasetdiffusionmodeltrainActorNetworkValueNetwork参数 todo2、读取 diffusion 配置并创建 DDPM1生成 beta schedulebetas get_beta_schedule(...) 确定前向加噪公式的系数2创建扩散模型过程diffusion GaussianDiffusion(betasbetas, **diffusion_config)里面写了q_sample()训练时加噪train_losses()计算训练 lossp_sample()生成图片时逐步去噪q_posterior_mean_var()/p_mean_var()计算真实后验和模型后验用于 KL、采样和评估3、创建 UNet 主模型_model UNet(**model_config)4、创建 ValueNetwork 和 ActorNetworkValueNetwork用于价值估计/辅助策略训练但是后面没有用到ActorNetwork根据图片生成 timestep 采样分布5、创建三个优化器optimizer Adam(model.parameters(), lrtrain_config.model_lr, ...) optimizer_pi Adam(policy.parameters(), lrtrain_config.policy_lr, ...) optimizer_v Adam(value_function.parameters(), lrtrain_config.value_lr, ...)三个优化器分别更新optimizerUNetoptimizer_piActorNetworkoptimizer_vValueNetwork6、创建学习率调度器scheduler LambdaLR(...)scheduler_v LambdaLR(...)scheduler_pi LambdaLR(...)用于学习率 warmup。7、创建trainer Trainer8、创建评估器 EvaluatorFID 越低一般说明生成质量越好。Evaluator.eval()↓清空 Inception 统计器↓循环生成 eval_total_size 张图片↓每批图片送入 Inception 网络提特征↓累计生成图片特征的均值和协方差生成训练集风格的图片↓读取真实数据集预计算的均值和协方差直接下载↓计算 FID↓返回 {fid: fid}9、开始训练trainer.train负责按 epoch/batch 调用step()训练模型并定期保存生成样例图、评估 FID、记录日志和同步多 GPU。train()|v设置采样图片网格|vfor 每个 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 e1632 or e 2039:是 -计算fid|v记录 wandb 日志|v多 GPU 同步for结束def step(self, x, e, i, global_steps1, loggerNone):作用给当前 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计算2vActorNetwork 根据 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、根据第五步就是rewardreward 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散度的计算1def compute_singlestep_KL(self, x, sampled_t)算出一张图片在1000个时间步的每个像素的kl散度每个像素得到1000个kl对每个像素取平均1、self.diffusion.train_losses用ddpm正向加噪获得加噪图片sampled_xt并可计算多种loss2、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 值2kl_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|vActorNetwork(x) 输出 alpha, beta|vBeta(alpha, beta)创建beta分布|v从 Beta 分布中采样 dist_sampled ∈ [0, 1]|v映射到 [0, self.timesteps - 1]|vround 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个时间步的选择可能不准确13这个数字2代表时间步选择方法文章做了消融实验对比1、3、10但依据不足。可以迁移的训练一个时间步选择器选择某个时间步后整体loss下降。采样器参数ϕ→选择 t→更新 SDθ→评价 Δkt​→更新 ϕ迁移小集合 S降低时间步评价成本选择损失计算的时间步记录每次时间步更新后的损失根据此损失选择3个代表时间步去计算reward不是只根据这一步的数值防止偶然性影响假设20次是有相关性根据unet来定的。DDPM 训练完整地记成DDPM采样/生成阶段也有前向传播每一次用unet预测噪声没有反向。