CVPR 2024 | FreeU：扩散U-Net的“免费午餐“，无需训练即可提升生成质量！

这篇论文最有趣的地方不是又加了一个模块而是它直接告诉你你已经在用的U-Net结构里就藏着一个可以即插即用、零成本、零参数提升效果的秘密武器——只需要两行代码。论文: FreeU: Free Lunch in Diffusion U-Net作者: Chenyang Si, Ziqi Huang, Yuming Jiang, Ziwei Liu发表单位: S-Lab, Nanyang Technological University发表: CVPR 2024论文链接: CVPR 2024 PDF | arXiv:2309.11497代码链接: GitHub - ChenyangSi/FreeU一、引言扩散模型Diffusion Model在过去几年彻底改变了图像生成领域Stable Diffusion、DALL·E 等产品已经走向大众。然而如何在推理阶段进一步提升生成质量而不需要重新训练、不需要增加参数、不需要额外的计算开销始终是一个极具吸引力的研究方向。FreeUFree Lunch in Diffusion U-Net正是在这个背景下提出的。它发现扩散U-Net架构中存在一个被忽视的潜力点——通过对主干特征和跳跃连接特征进行巧妙的重加权可以在零成本下显著提升生成质量。核心贡献揭示U-Net内部工作机制发现主干网络主要负责去噪而跳跃连接主要引入高频特征导致网络忽视主干语义提出FreeU方法通过两个简单的缩放因子在推理阶段即可提升生成质量广泛的适用性可与Stable Diffusion、DreamBooth、ModelScope、Rerender、ReVersion等多种扩散模型即插即用二、核心动机扩散U-Net的核心设计是编码器-解码器结构 跳跃连接skip connection。编码器负责提取特征解码器负责去噪重建跳跃连接则将编码器的特征直接传递给解码器的对应层级。然而作者通过深入的实验分析发现了一个关键问题跳跃连接引入的高频特征淹没了主干网络的语义信息。解码器在拼接跳跃特征后更多地依赖这些细节信息而忽略了主干网络中更重要的去噪语义。这是一个非常反直觉的发现——在CV领域跳跃连接通常被视为一种好东西如ResNet但在这个特定的扩散去噪场景中跳跃连接带来的高频细节反而干扰了去噪过程。基于这一发现FreeU的解决思路非常直接放大主干特征的权重同时衰减跳跃连接中的高频噪声。三、方法3.1 模块整体设计图1FreeU整体框架FreeU的核心思想可以概括为两个操作均在扩散U-Net的解码器decoder阶段执行Backbone Rescale主干重缩放对解码器主干特征的前半部分通道进行放大Skip Fourier Filter跳跃傅里叶滤波对跳跃连接传入的特征进行频域低通滤波这两个操作分别在解码器的前两个阶段即通道数为1280和640的层级执行。前向过程简述解码器每层先弹出一个跳跃连接特征res_hidden_states根据当前层级通道数对主干特征hidden_states的前一半通道乘以缩放因子b1或b2对跳跃特征res_hidden_states做傅里叶低通滤波scale参数为s1或s2拼接后进入ResNet处理3.2 核心操作Fourier Filter傅里叶滤波傅里叶滤波是FreeU中处理跳跃连接高频特征的核心操作。其数学原理如下x filtered F − 1 ( F ( x ) ⊙ M ) x_{\text{filtered}} \mathcal{F}^{-1}(\mathcal{F}(x) \odot M)xfiltered​F−1(F(x)⊙M)其中F \mathcal{F}F为二维FFTF − 1 \mathcal{F}^{-1}F−1为逆FFTM MM为频域掩码。掩码的构建方式为在频谱中心区域低频设置缩放系数s ss其余区域保持为1。M c , j { s , ∣ j − H / 2 ∣ ≤ τ ∧ ∣ c − W / 2 ∣ ≤ τ 1 , otherwise M_{c,j} \begin{cases} s, |j - H/2| \leq \tau \land |c - W/2| \leq \tau \\ 1, \text{otherwise} \end{cases}Mc,j​{s,1,​∣j−H/2∣≤τ∧∣c−W/2∣≤τotherwise​其中τ \tauτ为频率阈值论文中取1s ss为衰减系数1。importtorchimporttorch.fftasfftdefFourier_filter(x,threshold,scale): 对特征图进行频域低通滤波 :param x: 输入特征图 [B, C, H, W] :param threshold: 频域滤波阈值保留中心的频率范围 :param scale: 低频区域缩放系数1 表示衰减 :return: 滤波后的特征图 dtypex.dtype xx.type(torch.float32)# 正向FFT将空间域转换到频域x_freqfft.fftn(x,dim(-2,-1))x_freqfft.fftshift(x_freq,dim(-2,-1))# 将零频率移到中心B,C,H,Wx_freq.shape masktorch.ones((B,C,H,W)).to(x.device)# 构建频域掩码中心区域设置为scale其余保持为1crow,ccolH//2,W//2mask[...,crow-threshold:crowthreshold,ccol-threshold:ccolthreshold]scale x_freqx_freq*mask# 逆FFT从频域回到空间域x_freqfft.ifftshift(x_freq,dim(-2,-1))x_filteredfft.ifftn(x_freq,dim(-2,-1)).real x_filteredx_filtered.type(dtype)returnx_filtered关键理解当scale 1时中心低频区域被衰减这意味着保留更多高频细节的同时整体降低了跳跃连接特征的幅度从而让主干语义更有话语权。3.3 核心操作Backbone Rescale主干重缩放主干重缩放的操作更为简单直接对解码器特征的前一半通道即来自主干网络的特征乘以缩放因子。以通道数为1280的第一阶段为例# 对前640个通道主干特征进行缩放hidden_states[:,:640]hidden_states[:,:640]*b1其中b1 1论文推荐值为1.2-1.6用于增强主干特征的贡献。3.4 完整FreeU解码器前向过程FreeU通过monkey-patch的方式注入到解码器的UpBlock中在跳跃连接拼接前执行两个操作defregister_free_upblock2d(model,b11.2,b21.4,s10.9,s20.2): 将FreeU注入到Unet的UpBlock2D中 :param model: diffusers模型 :param b1: 第一阶段主干缩放因子通道数1280 :param b2: 第二阶段主干缩放因子通道数640 :param s1: 第一阶段跳跃连接滤波系数 :param s2: 第二阶段跳跃连接滤波系数 defup_forward(self):defforward(hidden_states,res_hidden_states_tuple,tembNone,upsample_sizeNone):forresnetinself.resnets:# 弹出跳跃连接特征res_hidden_statesres_hidden_states_tuple[-1]res_hidden_states_tupleres_hidden_states_tuple[:-1]# ---------- FreeU核心操作 ----------# 仅在前两个解码器阶段执行ifhidden_states.shape[1]1280:# 放大主干特征前640通道hidden_states[:,:640]hidden_states[:,:640]*self.b1# 对跳跃特征做傅里叶滤波res_hidden_statesFourier_filter(res_hidden_states,threshold1,scaleself.s1)ifhidden_states.shape[1]640:# 放大主干特征前320通道hidden_states[:,:320]hidden_states[:,:320]*self.b2# 对跳跃特征做傅里叶滤波res_hidden_statesFourier_filter(res_hidden_states,threshold1,scaleself.s2)# ------------------------------------# 拼接并送入ResNethidden_statestorch.cat([hidden_states,res_hidden_states],dim1)hidden_statesresnet(hidden_states,temb)ifself.upsamplersisnotNone:forupsamplerinself.upsamplers:hidden_statesupsampler(hidden_states,upsample_size)returnhidden_statesreturnforward# 遍历所有up_blocks替换forward函数fori,upsample_blockinenumerate(model.unet.up_blocks):ifisinstance_str(upsample_block,UpBlock2D):upsample_block.forwardup_forward(upsample_block)setattr(upsample_block,b1,b1)setattr(upsample_block,b2,b2)setattr(upsample_block,s1,s1)setattr(upsample_block,s2,s2)3.5 即插即用模块说明属性说明模块名称FreeUFree Lunch in Diffusion U-Net核心创新通过Backbone Rescale和Skip Fourier Filter两个操作在零成本下重加权U-Net的主干与跳跃连接特征提出动机U-Net的跳跃连接引入高频特征导致网络忽视主干语义影响去噪质量适用任务文本到图像生成、文本到视频生成、个性化生成DreamBooth、图像修复、视频重渲染等所有基于扩散U-Net的生成任务缝合位置扩散模型U-Net解码器的UpBlockUpBlock2D和CrossAttnUpBlock2D前向过程中跳跃连接拼接之前参数量0额外参数——仅通过运行时缩放因子调整不引入任何可学习参数内存开销0额外内存——不增加显存占用推理时间几乎无增加——仅增加少量FFT计算3.6 推荐参数作者为不同的Stable Diffusion版本提供了推荐参数模型版本b1b2s1s2SD 1.41.31.40.90.2SD 1.51.51.60.90.2SD 2.11.41.60.90.2SDXL1.31.40.90.2参数调节范围建议b1: 1 ≤ b1 ≤ 1.2或更高视效果调整b2: 1.2 ≤ b2 ≤ 1.6s1: s1 ≤ 1s2: s2 ≤ 1四、实验结果4.1 图像生成质量对比FreeU在多种扩散模型上均取得了显著的生成质量提升。以下是Stable Diffusion 2.1的对比图2SD-XL 生成结果对比无FreeU vs 有FreeU图3ModelScope 生成结果对比无FreeU vs 有FreeU从对比图中可以明显看出使用FreeU后的生成结果在以下方面得到了改善细节丰富度纹理和边缘细节更加清晰色彩饱和度颜色更加鲜艳自然构图合理性减少了解码器因忽略主干语义而产生的结构性错误4.2 多模型泛化性FreeU的优势不仅限于Stable Diffusion还在以下模型中得到了验证模型任务类型效果Stable Diffusion文本到图像显著提升SD-XL高分辨率文本到图像4096×4096显著提升DreamBooth个性化文本到图像显著提升ModelScope文本到视频显著提升Rerender视频到视频显著提升ReVersion关系反转生成显著提升ControlNet条件控制生成显著提升LCMLatent Consistency Model快速采样显著提升4.3 消融实验FreeU的两项操作各自独立有效且联合使用时效果最佳配置效果描述仅Backbone Rescale改善生成质量增强主干语义贡献仅Skip Fourier Filter改善生成质量减少跳跃连接噪声干扰FreeU两者结合效果最佳两项操作协同增强五、总结与思考FreeU的核心价值在于它揭示了一个简单但深刻的事实扩散U-Net的设计并非完美其内部特征流动存在可优化的空间。核心贡献总结揭示了U-Net内部机制的盲区跳跃连接的高频特征可能淹没主干语义零成本即插即用无需训练、无需参数、无需额外显存仅靠两个缩放因子广泛的模型兼容性可无缝集成到几乎所有基于U-Net的扩散模型中个人评价FreeU的设计哲学值得每个扩散模型从业者学习——在做加法之前先看看已有的结构里有没有被浪费的潜力。它用最简单的操作乘法 FFT解决了生成质量提升这一难题体现了大道至简的工程美学。适用任务所有基于扩散U-Net的生成任务文本到图像、图像到图像、文本到视频、个性化生成等在解码器的跳跃连接拼接前插入Backbone Rescale和Skip Fourier Filter操作即可零成本提升生成质量。六、快速上手6.1 基础依赖层Fourier Filterimporttorchimporttorch.fftasfftdefFourier_filter(x,threshold,scale):频域低通滤波——衰减跳跃连接中的高频特征dtypex.dtype xx.type(torch.float32)# 正向FFT 移位x_freqfft.fftn(x,dim(-2,-1))x_freqfft.fftshift(x_freq,dim(-2,-1))B,C,H,Wx_freq.shape masktorch.ones((B,C,H,W)).to(x.device)# 中心低频区域设为scalecrow,ccolH//2,W//2mask[...,crow-threshold:crowthreshold,ccol-threshold:ccolthreshold]scale x_freqx_freq*mask# 逆移位 逆FFTx_freqfft.ifftshift(x_freq,dim(-2,-1))x_filteredfft.ifftn(x_freq,dim(-2,-1)).realreturnx_filtered.type(dtype)6.2 工具函数类型检查defisinstance_str(x:object,cls_name:str): 检查对象x的继承链中是否有类名为cls_name的类 用于monkey-patch时不依赖具体的类实现 for_clsinx.__class__.__mro__:if_cls.__name__cls_name:returnTruereturnFalse6.3 核心模块注入FreeU到UpBlock2Ddefregister_free_upblock2d(model,b11.2,b21.4,s10.9,s20.2): 将FreeU注入到扩散模型U-Net的UpBlock2D中 参数: model: diffusers的StableDiffusionPipeline对象 b1: 第一阶段1280通道主干缩放因子 b2: 第二阶段640通道主干缩放因子 s1: 第一阶段跳跃连接滤波系数 s2: 第二阶段跳跃连接滤波系数 defup_forward(self):defforward(hidden_states,res_hidden_states_tuple,tembNone,upsample_sizeNone):forresnetinself.resnets:# 弹出跳跃连接特征res_hidden_statesres_hidden_states_tuple[-1]res_hidden_states_tupleres_hidden_states_tuple[:-1]# ---- FreeU核心操作仅前两个解码器阶段----ifhidden_states.shape[1]1280:hidden_states[:,:640]hidden_states[:,:640]*self.b1 res_hidden_statesFourier_filter(res_hidden_states,threshold1,scaleself.s1)ifhidden_states.shape[1]640:hidden_states[:,:320]hidden_states[:,:320]*self.b2 res_hidden_statesFourier_filter(res_hidden_states,threshold1,scaleself.s2)# --------------------------------------------hidden_statestorch.cat([hidden_states,res_hidden_states],dim1)hidden_statesresnet(hidden_states,temb)ifself.upsamplersisnotNone:forupsamplerinself.upsamplers:hidden_statesupsampler(hidden_states,upsample_size)returnhidden_statesreturnforwardfori,upsample_blockinenumerate(model.unet.up_blocks):ifisinstance_str(upsample_block,UpBlock2D):upsample_block.forwardup_forward(upsample_block)setattr(upsample_block,b1,b1)setattr(upsample_block,b2,b2)setattr(upsample_block,s1,s1)setattr(upsample_block,s2,s2)6.4 核心模块注入FreeU到CrossAttnUpBlock2Ddefregister_free_crossattn_upblock2d(model,b11.2,b21.4,s10.9,s20.2): 将FreeU注入到扩散模型U-Net的CrossAttnUpBlock2D中 适用于带有交叉注意力的Upsampling块 defup_forward(self):defforward(hidden_states,res_hidden_states_tuple,tembNone,encoder_hidden_statesNone,cross_attention_kwargsNone,upsample_sizeNone,attention_maskNone,encoder_attention_maskNone,):forresnet,attninzip(self.resnets,self.attentions):res_hidden_statesres_hidden_states_tuple[-1]res_hidden_states_tupleres_hidden_states_tuple[:-1]# ---- FreeU核心操作 ----ifhidden_states.shape[1]1280:hidden_states[:,:640]hidden_states[:,:640]*self.b1 res_hidden_statesFourier_filter(res_hidden_states,threshold1,scaleself.s1)ifhidden_states.shape[1]640:hidden_states[:,:320]hidden_states[:,:320]*self.b2 res_hidden_statesFourier_filter(res_hidden_states,threshold1,scaleself.s2)# -----------------------hidden_statestorch.cat([hidden_states,res_hidden_states],dim1)hidden_statesresnet(hidden_states,temb)hidden_statesattn(hidden_states,encoder_hidden_statesencoder_hidden_states,cross_attention_kwargscross_attention_kwargs,)[0]ifself.upsamplersisnotNone:forupsamplerinself.upsamplers:hidden_statesupsampler(hidden_states,upsample_size)returnhidden_statesreturnforwardfori,upsample_blockinenumerate(model.unet.up_blocks):ifisinstance_str(upsample_block,CrossAttnUpBlock2D):upsample_block.forwardup_forward(upsample_block)setattr(upsample_block,b1,b1)setattr(upsample_block,b2,b2)setattr(upsample_block,s1,s1)setattr(upsample_block,s2,s2)6.5 完整使用示例importtorchfromdiffusersimportStableDiffusionPipeline# 1. 加载模型model_idstabilityai/stable-diffusion-2-1pipelineStableDiffusionPipeline.from_pretrained(model_id,torch_dtypetorch.float16).to(cuda)# 2. 注册FreeUSD 2.1推荐参数register_free_upblock2d(pipeline,b11.4,b21.6,s10.9,s20.2)register_free_crossattn_upblock2d(pipeline,b11.4,b21.6,s10.9,s20.2)# 3. 生成图像promptA fantasy landscape, trending on artstationtorch.manual_seed(42)imagepipeline(prompt,num_inference_steps25).images[0]image.save(freeu_result.png)注FreeU通过monkey-patch方式注入不影响原始模型权重随时可通过重新加载模型恢复基线行为。如需对比效果先用b11.0, b21.0, s11.0, s21.0生成基线再用推荐参数生成FreeU版本。