Normalizing Flows with Iterative Denoising

这是一份关于论文《Normalizing Flows with Iterative Denoising》（具有迭代去噪的归一化流）的详细技术总结。

1. 研究背景与问题 (Problem)

背景：
图像生成领域主要由两类范式主导：

1. 扩散模型 (Diffusion Models)： 通过连续空间去噪生成图像，质量极高，但推理时需要数十到数百次迭代，计算开销大，且扩展性（Scaling Law）不如大语言模型（LLM）可预测。
2. 自回归模型 (Autoregressive Models)： 将数据视为离散 Token 序列（如 LLM），扩展性极佳且采样快，但量化过程会引入信息瓶颈，降低保真度。早期尝试在连续像素空间直接进行自回归建模（如 PixelCNN），但因序列过长导致训练和采样效率低下。

核心问题：
归一化流（Normalizing Flows, NFs）作为一种基于似然的经典方法，近期通过 TARFlow 等模型在图像建模上展现了潜力。然而，NFs 在训练中面临一个关键的**“噪声困境” (Noise Dilemma)**：

• 噪声过小： 模型倾向于生成纹理丰富但全局结构混乱的图像，且训练不稳定（对数似然在数据流形附近无界）。
• 噪声过大： 模型能捕捉良好的全局结构，但生成的图像细节模糊、平滑，且存在伪影（受 Tweedie 引理限制，后验均值随噪声增加而平滑化）。

现有的 NF 方法（如 TARFlow）通常使用单一噪声水平或简单的自去噪器，难以同时兼顾全局结构的连贯性和局部细节的清晰度。

2. 方法论 (Methodology)

本文提出了 iTARFlow (Iterative TARFlow)，一种结合了自回归生成与迭代去噪的归一化流框架。

核心思想

iTARFlow 的核心在于在多个噪声尺度上进行训练，并在推理时执行迭代去噪。这类似于扩散模型的训练方式，但保留了 NF 的端到端似然目标。

具体技术细节：

1. 多噪声水平训练 (Multi-noise Training)：

   • 模型被训练以最大化不同噪声水平 $t \in [t_{min}, t_{max}]$ 下的对数似然 $\log p_\theta(x_t, t)$。
   • 输入数据 $x$ 被添加高斯噪声 $x_t = x + t\epsilon$。
   • 通过最小化期望负对数似然，模型学习从不同噪声水平到数据的映射。
   • 超参数设置：$t_{min}=0.01$ 以保证训练稳定性，$t_{max}$ 根据数据集分辨率 empirically 设定（如 ImageNet-256 设为 0.7）。

2. 两阶段采样过程 (Two-Phase Sampling)：

   • 阶段一：自回归生成 (Autoregressive Generation)：
     • 模型首先从先验分布（高斯噪声）采样，利用 TARFlow 的因果 Transformer 架构，以自回归方式（Patch-by-Patch）生成一个高噪声的中间图像 $x_{t_{max}}$。
     • 此阶段利用 CFG (Classifier-Free Guidance) 来增强生成质量，特别是在大 Patch 尺寸下。
   • 阶段二：迭代去噪 (Iterative Denoising)：
     • 利用训练好的模型计算得分函数（Score Function）$\nabla_x \log p_\theta(x_t, t)$，该得分函数可通过自动微分直接获得。
     • 执行基于 ODE 的迭代去噪过程：$dx_t = -t \nabla_x \log p_\theta(x_t, t) dt$。
     • 该过程是完全并行的（非自回归），仅需少量步骤（如 5-10 步）即可将高噪声样本转化为高质量图像。

3. 架构设计：

   • 基于 Transformer 的自回归流（Autoregressive Flow），由堆叠的因果 Transformer 块组成。
   • 在层与层之间应用置换（Permutation）以增强表达能力。

3. 关键贡献 (Key Contributions)

1. 揭示并解决“噪声困境”： 系统性地研究了 NF 训练中的噪声水平选择问题，发现较高的噪声水平有助于生成更好的全局结构，但会牺牲细节。
2. 提出 iTARFlow 框架： 设计了一种结合多噪声训练和迭代去噪推理的新范式。该方法利用自回归生成全局结构，再通过并行迭代去噪恢复高频细节，成功融合了自回归模型的扩展性优势和连续表示的高保真度。
3. 性能突破： 在 ImageNet 64、128 和 256 分辨率上，iTARFlow 显著优于之前的 TARFlow 变体，并大幅缩小了归一化流与当前最先进（SOTA）扩散模型及自回归模型之间的性能差距。
4. 兼容性分析： 发现 iTARFlow 估计的得分函数与扩散模型（如 DiT）的得分函数具有可比性，且能无缝结合预训练的 DiT 去噪器，为构建统一的混合生成框架提供了可能。

4. 实验结果 (Results)

• ImageNet-64：
  • iTARFlow-B (7.7 亿参数) 达到 FID 1.68，优于 TARFlow (2.66)，并接近扩散模型（如 EDM-SDE 的 1.55），且参数量更少。
• ImageNet-128：
  • iTARFlow-L (12 亿参数) 达到 FID 2.44，显著优于 TARFlow (5.03)，并与扩散模型（如 Simple Diff 的 1.94）竞争。
• ImageNet-256：
  • iTARFlow-XL (22 亿参数) 达到 FID 3.32，优于之前的 NF 模型（TARFlow: 5.56, STARFlow: 4.69）和部分自回归模型。
  • 虽然与 SOTA 扩散模型（如 SiD2 patch 1 的 1.38）仍有差距，但差距已显著缩小。
• 消融实验：
  • 证明了迭代去噪（仅需 5-10 步）比单次自去噪器能显著提升图像质量。
  • 证明了 iTARFlow 的得分函数可以替代 DiT 作为去噪器，且训练效率更高（收敛更快）。

5. 局限性与未来方向 (Limitations & Significance)

局限性：

• 显存消耗： 迭代去噪过程中的自动微分需要对整个网络进行反向传播，导致 GPU 显存占用较高。
• 生成伪影： 在 256 分辨率下，观察到两类主要失败案例：
  1. 背景坍塌 (Collapsed Backgrounds)： 整个背景变黑，可能与大 Patch 尺寸下过高的 CFG 引导强度有关。
  2. 低频模糊 (Low-frequency Blur)： 尤其是图像右下角（自回归生成的第一个 Token 区域）出现失真和模糊，推测是因为第一个 Token 缺乏上下文条件，导致分布外预测并传播误差。
• 超参数理论缺失： 目前缺乏理论指导来选择最优的噪声范围 $[t_{min}, t_{max}]$，主要依赖经验调整。

意义：

• 范式融合： iTARFlow 证明了归一化流可以通过引入扩散式的迭代去噪机制，克服传统 NF 的局限性，成为扩散模型和离散自回归模型的强有力替代方案。
• 扩展性潜力： 保留了 NF 的精确似然训练和自回归架构的扩展性优势，为未来构建可扩展、高质量的连续图像生成模型指明了方向。
• 混合框架： 展示了 NF 与预训练扩散模型（DiT）结合的可能性，为构建更灵活的混合生成架构奠定了基础。

总结：
iTARFlow 通过“多噪声训练 + 迭代去噪”的策略，成功解决了归一化流在图像生成中的噪声困境，在保持训练稳定性和扩展性的同时，大幅提升了生成图像的全局结构和局部细节质量，是归一化流领域的一次重要进展。