⏶7
基于 GS-Jacobi 迭代加速 TarFlow 采样
发表
由 
zhenqin 提交
zhenqin 提交作者: 
Ben Liu, 
Zhen Qin
Ben Liu, Zhen Qin摘要
图像生成模型已获得广泛应用。例如，TarFlow 模型结合了 Transformer 架构和 Normalizing Flow 模型，在多个基准测试上取得了最先进的结果。然而，由于因果形式的注意力需要顺序计算，TarFlow 的采样过程非常缓慢。在本文中，我们证明通过一系列优化策略，可以使用 Gauss-Seidel-Jacobi（缩写为 GS-Jacobi）迭代方法极大地加速 TarFlow 的采样。具体来说，我们发现 TarFlow 模型中的块具有不同的重要性：少量块在图像生成任务中起主要作用，而其他块贡献相对较少；一些块对初始值敏感且容易发生数值溢出，而另一些则相对稳健。基于这两个特性，我们提出了收敛排名指标（CRM）和初始猜测指标（IGM）：CRM 用于识别 TarFlow 块是“简单”（在少数迭代内收敛）还是“困难”（需要更多迭代）；IGM 用于评估迭代的初始值是否良好。在四种 TarFlow 模型上的实验表明，GS-Jacobi 采样可以在显著提高采样效率的同时保持生成图像的质量（通过 FID 衡量），在 Img128cond 中提速 4.53 倍，在 AFHQ 中提速 5.32 倍，在 Img64uncond 中提速 2.96 倍，在 Img64cond 中提速 2.51 倍，而不会降低 FID 分数或样本质量。代码和检查点可在 https://github.com/encoreus/GS-Jacobi_for_TarFlow 访问。
图像生成模型已获得广泛应用。例如，TarFlow 模型结合了 Transformer 架构和 Normalizing Flow 模型，在多个基准测试中取得了最先进的结果。然而，由于注意力的因果形式需要按顺序计算，TarFlow 的采样过程极其缓慢。在本文中，我们证明通过一系列优化策略，通过使用高斯-赛德尔-雅可比（简称 GS-Jacobi）迭代方法，TarFlow 采样可以大大加速。具体来说，我们发现 TarFlow 模型中的块具有不同的重要性：少数块在图像生成任务中起主要作用，而其他块贡献相对较少；一些块对初始值敏感且容易发生数值溢出，而另一些块相对稳健。基于这两个特性，我们提出了收敛排序指标 (CRM) 和初始猜测指标 (IGM)：CRM 用于识别 TarFlow 块是“简单”（在很少的迭代中收敛）还是“困难”（需要更多迭代）；IGM 用于评估迭代的初始值是否良好。在四种 TarFlow 模型上的实验表明，GS-Jacobi 采样可以显著提高采样效率，同时保持生成图像的质量（通过 FID 衡量），在 Img128cond 中实现了 4.53 倍加速，在 AFHQ 中实现了 5.32 倍加速，在 Img64uncond 中实现了 2.96 倍加速，在 Img64cond 中实现了 2.51 倍加速，而没有降低 FID 分数或样本质量。代码和检查点可在[此]处获取