【扩散模型】11、FreeU | 推理的时候只需加两个缩放因子就能免费提升图像生成效果！

论文：FreeU: Free Lunch in Diffusion U-Net

代码：https://chenyangsi.top/FreeU/

出处：2023.09

效果：

• 无需训练，只需要在推理的时候调节 U-Net 特征 concat 时的两个因子（分别控制 backbone 特征和 skip connection 特征的贡献权重），就能获得 “免费的午餐”，也就是提升模型生成图片的细节性和保真性。
• FreeU 可以轻松集成到现有的扩散模型中，如 Stable Diffusion, DreamBooth, ModelScope，Rerender 和 ReVersion

一、背景

本文中，作者发现了 U-Net 的一个未被利用的特性，能看做“免费的午餐”，也就是能够在不需要任何额外开销的情况下即时显著提升生成质量

作者是怎么发现的呢：

• 现有 U-Net 可能存在的问题：U-Net 结构的 backbone 的作用在于去噪，跳连结构主要作用是将高频特征从 encoder 送到 decoder 中，直接 concat 的操作可能会导致高频特征影响 backbone 的去噪能力，选择性提高基础连接权重并降低跳跃连接中低频权重可能会提高生成质量
• 利用这一发现，本文提出了一种简单而有效的方法——称为 “FreeU” ——它可以在不需要额外训练或微调的情况下增强生成质量。
• 主要操作：重新加权来自U-Net跳跃连接和骨干特征图的贡献，以利用U-Net架构的两个组成部分的优势。在图像和视频生成任务上的有希望的结果表明，
• 可扩展性：FreeU 可以轻松集成到现有的扩散模型中，如 Stable Diffusion, DreamBooth, ModelScope，Rerender 和 ReVersion

U-Net：

• 扩散模型一般都使用 U-Net 的结构来作为去噪模型，U-Net 网络训练的目标呢就是逐步的预测每步的噪声来实现去噪
• 现有的工作一般都聚焦于使用预训练好的 U-Net 来进行下游工作，很少有再去探究优化 U-Net 结构的

作者画出了去噪的中间过程生成的图片的高频和低频特征，如图 2 所示：

• 第一行是生成的图片
• 第二行是经过傅里叶变换后得到的低频图
• 第三行是经过傅里叶变化后得到的高频图
• 可以看出随着去噪过程的进行，低频图的变化比较小，高频图的变化比较大

图 3 中的曲线更是证明了这一点：

• 低频特征包含很多平滑特征和全局特征，在去噪一开始就能很好的被捕捉到
• 高频特征包含图像中变化剧烈的特征，如边缘和纹理，这些细节信息更容易被噪声干扰
• 所以，去噪的过程需要再能够去除噪声的同时保存细节信息才能更好的生成图片

作者根据这个观察到的观点，又去查看了 U-Net 的结构和高频、低频信号的关系，发现：

• 每个 step 的 U-Net 结构，encoder 和 decoder 都有跳连，跳连过去的特征都是和 backbone 原本的特征直接 concat 然后输入下一层的。
• 作者认为 U-Net 的 backbone 主要作用是用于去噪，跳连的作用是将高频信息传递到 decoder 中，这样的跳连能够传递更多的细节特征，更容易复原出输入的样子
• 但是这样直接 concat 会有问题，就是高频信息可能会影响 inference 过程中 backbone 的去噪能力，导致模型产生一些异常的细节，如图 1 所示

作者如何改进这个问题的呢：FreeU

• 在推理阶段，作者使用了两个调制因子来平衡 U-Net 的 backbone 和 skip connections 的贡献程度
• backbone 权重：如果该权重很大，则有利于去噪过程，但会让生成的图片过于平滑模糊
• skip connections 权重：如果该权重很大，则可能会产生一些不想要的细节信息
• 所以作者提出在 U-Net 的结构中，动态控制这两个因子的大小，能够无任何额外开销的情况下提升模型生成图片的细节丰富性和保真性

二、方法

2.1 扩散过程

扩散模型如 DDPM，一般都是先对图片加噪，然后逐步学习噪声然后去噪的过程：

加噪过程：

去噪过程：

2.2 U-Net 如何实现去噪

去噪的过程一般是使用以时间为条件的 U-Net 来实现，每一步的 U-Net 都是学习对应该步的噪声，然后实现去噪

U-Net 的结构如图 4，包括一个 backbone 和很多 skip connections，为了度量这个两个部件在去噪过程中的作用，作者使用了控制实验，分别使用了两个尺度因子来控制 concat 之前这两者的权重：

• b：控制 backbone 生成的特征图的权重
• s：控制 skip 连接的特征图的权重

如图 5 所示：

• 提升 backbone 的权重 b 能够直接提升生成的图片的质量
• 调节 skip connection 的权重 s，对生成的图片的质量几乎没有影响

观察到了这个现象后，作者就进一步探究了为什么出现这个现象，结果如图 6 所示：

• 分析得到，图片质量的提升和 U-Net 的backbone的去噪能力有很强的关系
• 如图 6，当 b 的改变步长相同时（0.6到0.8,0.8到1.0 等），对高频信息的抑制作用的改变是近似的，就是高频区域不同线条的距离都差不多，这说明提升 backbone 特征的贡献，能够提升 U-Net 的去噪能力，在保真和细节上都有提升

那么 skip connection 的作用是什么呢：为解码器提供信息，而不是为了改善网络的去噪性能

• skip connection 在结构上来说，就是将 encoder 中的特征传递给 decoder 中
• 如图 7 所示，这些传递过去的特征主要包含的是高频特征，作者推测，在U-Net训练过程中，这些高频特征可能会导致解码器过快地学习到噪声，而不是所需的有用信号，这会加速网络对噪声的预测。
• 而且，前面图 5 中说明了调整跳跃连接传递的特征似乎对解码器的输出影响不大。这表明跳跃连接主要是为解码器提供信息，而不是为了改善网络的去噪性能，这可能会减少网络本身去除噪声的能力。

这个观察结果引发了一个问题：

在 U-Net 架构中，编码器的去噪作用和跳跃连接之间应该如何平衡，以便在整个网络中实现最佳的去噪效果。简而言之，这是关于如何最佳地利用网络的不同部分来提高图像去噪性能的问题。

2.3 在 U-Net 中如何使用 Free lunch

在 U-Net 解码器的第 $l$ 个块中，有两种特征图：

• 一种是来自前一个块主干网络的特征图，记为 $x_l$
• 另一种是通过对应的跳跃连接传播的特征图，记为 $h_l$
• 为了调节这些特征图，引入了两个标量因子：一个是针对 $x_l$ 的主干特征缩放因子 $b_l$ ，另一个是针对 $h_l$ 的跳跃特征缩放因子 $s_l$
• 具体来说，$b_l$ 旨在放大主干特征图 $x_l$ ，而 $s_l$ 设计用来衰减跳跃连接的特征图 $h_l$ 。

如何放大主干网络特征图（$x_l$）的作用呢：

作者引入了一种与图像结构相关缩放的方法，可以为每步采样动态的调节对应的 scaling，不是固定的 scaling

首先，沿着通道维度计算平均特征图，就是每个 block $l$ 会得到一个对应的平均特征图，也就是该 block 上所有通道的特征图的平均，如图 8 所示：

• ${\bar{x}}_l$ 是沿着通道维度的平均特征图
• $C$ 是通道的数量
• $x_{l,i}$ 是第 $l$ 个块的第 $i$ 个通道的特征图。

然后，会得到每个 block 上对应的 backbone factor map ${\alpha }_l$，就是和特征图大小相同的一个缩放特征图，后面会乘到原始特征图上去做增强：

• ${\alpha }_l$：该 block 的backbone 缩放特征图，该特征图相当于对每个像素位置会有不同的缩放尺度，${\alpha }_l$ 某个像素位置的值越大，则 backbone 特征图上对应位置的值的变化就越大。
• $b_l$：标量常数，就是设定的尺度因子

通过前面的实验发现，如果不加选择地放大 $x_l$ 的所有通道，即每个 block 的所有通道都和 ${\alpha }_l$ 相乘来进行增强，会导致合成图像中的纹理过于平滑。这是因为增强的 U-Net 在去噪的同时，会损害图像的高频细节。因此，作者只对 $x_l$ 的一半通道进行增强，另外一半不进行增强：

为什么要使用整个 block 所有通道的平均特征图来指导加强呢：

如图 8，average feature map ${\bar{x}}_l$ 主要包括结构信息，所以 backbone factor map ${\alpha }_l$ 能够让 backbone feature map $x_l$ 以结构信息对齐。这种策略的目的是为了减轻过平滑的问题。首先，它增强了主干特征图的去噪能力，使其能够更有效地过滤掉噪声。其次，它避免了在整个特征图上无差别应用缩放的负面效果，从而在降噪和保留纹理之间达到了更细致的平衡。

对跳跃连接传播的特征图（$h_l$）如何削弱呢：

为了进一步减轻由于增强去噪能力而导致的纹理过度平滑问题，作者在傅里叶域对跳连特征图（$h_l$）进一步采用光谱调制，以选择性地减少其低频分量，也就是将低于某个阈值的低频信息全部变为 $s_l$，这样低频信息就减少了，保留下来的都是高频的信息，更有细节。

• 首先，对特征图（$h_l$）进行傅里叶变换（FFT）
• 然后，对经过傅里叶变化的特征图和 ${\beta }_{l,i}$ 进行点乘（$\odot$），来改变其傅里叶系数模值
• 最后，对经过模值缩放的特征进行傅里叶反变换（IFFT），得到经过频域缩放后的跳连特征 $h^{\prime }$

${\beta }_{l,i}$ 是傅里叶频谱的掩码 mask，这个掩码用于调节低频信息，以解决过度平滑的问题，怎么解决的呢：

• 如果频谱图半径 r（傅里叶域内的一个点到中心的距离，可以理解为频率）小于某个阈值 $r_{t}hresh$，也就是频率小于该阈值的信息，全部乘以 $s_l$，模值会改变
• 如果频谱图半径大于 r，则全部乘以 1，即模值不变

是与傅里叶系数的模值相关的一个函数，是与 frequency-dependent scaling factor $s_l$ 有如下关系，r 是半径，$r_{thresh}$ 是频率阈值

这个过程可以帮助在去噪的同时保持图片的纹理细节，因为去噪往往会平滑掉图片的高频细节。通过在傅里叶域内调节频率成分，可以更精细地控制去噪过程，避免过度平滑的结果。

三、效果

3.1 text-to-image

3.2 text-to-video

3.3 下游任务