阅读量:0
MobileDiffusion：移动设备上即时文本到图像生成
Paper Title：MobileDiffusion: Instant Text-to-Image Generation on Mobile Devices
Paper是谷歌出品
Paper地址
图 1：MobileDiffusion 用于 (a) 文本到图像的生成。(b) Canny 边缘到图像、风格 LoRA 和修复。所有样本均在一步内生成。
Abstract.
大规模文本转图像扩散模型在移动设备上的部署受到模型规模大和高延迟的阻碍。在本文中，我们介绍了 MobileDiffusion，这是一种超高效的文本转图像扩散模型，通过对架构和采样技术进行大量优化而获得。我们对模型架构设计进行了全面检查，以最大限度地减少模型大小和 FLOPs，同时保持图像生成质量。此外，我们重新审视了扩散 GAN 的高级采样技术，并使一步采样与在基础模型上训练的下游应用程序兼容。定量和定性进行的实证研究证明了我们提出的技术的有效性。借助它们，MobileDiffusion 实现了移动设备上的即时文本转图像生成，建立了新的领先水平。
1 Introduction
文本到图像的扩散模型 [1, 40, 42, 44, 45, 47] 在以文本为条件生成高质量图像方面具有出色的能力。这些模型是各种应用的基础，包括图像编辑 [15,55]、受控生成 [39, 65]、个性化内容生成 [13, 46]、视频合成 [4, 14] 和底层视觉任务 [26, 61]。这些大规模模型本质上被认为必须在具有强大神经计算单元的服务器上运行。只有少数工作 [25, 27] 几乎没有涉及在移动设备上运行扩散模型，这仍然是一个悬而未决的挑战。
我们发现导致文本到图像传播模型效率低下的两个主要因素。首先，文本到图像传播模型的网络架构复杂，涉及大量参数，通常达到数十亿个，导致计算成本高昂。其次，传播模型的固有设计需要迭代去噪来生成图像，因此需要对模型进行多次评估 [17, 53]。这些效率低下的挑战对在资源受限的环境（例如移动设备）中部署构成了重大障碍。因此，尽管具有增强用户体验、解决新出现的隐私问题和节省成本等潜在好处，但这一方面在当前文献中仍然相对未被探索。
在本文中，我们的目标是开发一种适用于一系列设备应用的超高效扩散模型，这需要一种轻量、快速且功能多样的设计，以便高效处理各种下游生成任务，包括修复、可控生成和个性化生成。我们采用分而治之的策略来应对已确定的挑战（架构和采样效率），分别解决每个问题。
现有文献很少涉及扩散模型的架构效率挑战。先前的尝试主要集中在消除冗余的神经网络块 [25,27] 或重新组织它们以提高效率 [18]，但缺乏对模型架构组件的全面分析。我们的研究通过深入审查扩散网络填补了这一空白，从而精心优化了模型架构。我们的模型参数不到 4 亿，实现了高质量的图像生成，在效率上超越了以前的努力。
第二个挑战，即扩散模型的采样效率，一直是一个活跃的研究领域。借助先进的数值求解器 [2,23,33,34,51] 或蒸馏技术 [27, 35, 38, 48, 52]，采样所需的网络评估次数已从几百次显著减少到不到 10 次。尽管取得了这些进展，但减少评估步骤仍然会对移动设备构成挑战，因为即使是最低限度的处理需求也可能过于繁重。因此，我们的注意力转向开创性的单步扩散生成方法 [49,62,63]。这些进步有可能将扩散模型转变为高效的单步文本到图像生成器。然而，这些模型仍需要通过集成单独训练的附加模块来灵活地适应各种条件生成任务。为了弥补这一差距，我们探索了扩散-GAN 混合模型的细微设计空间 [28、49、60、62]，这是越来越受欢迎的一类前沿一步式扩散模型。具体来说，我们的研究集中在 UFOGen 框架 [62] 上，这是该领域的一项最新创新。我们的工作通过深入分析 UFOGen 的目标函数和训练方法来推进这一论述，最终形成了一种开发高效一步式扩散 GAN 模型的改进方法。这项工作旨在实现双重目标：促进高质量文本到图像合成的创建，并确保模型对各种下游应用的无缝适应性，从而标志着朝着实现超快速、多功能的一步式扩散模型迈出了重要一步。
我们在架构设计和采样效率方面的共同努力使得在移动设备上即时生成高质量的 512 × 512 512 \times 512 512×512图像成为可能，尤其是在 iPhone 15 Pro 上实现了 0.2 秒，这大约是人类对视觉刺激的平均响应时间。这一成就标志着在设备上文本到图像生成方面较之前最先进的技术有了重大飞跃 [27]。我们将我们的模型称为 MobileDiffusion，强调其作为边缘设备基础生成模型的潜力。
我们的论文对设备上生成建模领域做出了几项重要贡献，具体如下：
- 我们对文本到图像扩散模型的架构效率进行了全面探索。我们的工作引入了一种精致的扩散模型架构，它不仅高效轻量，而且还在移动设备上表现出色。
- 我们研究一步扩散 GAN 模型的设计空间，以提出训练此类模型的最佳方案，从而通过一步实现超越文本到图像应用的高生成能力。据我们所知，我们的工作开创了一步扩散模型适用于广泛应用的先河。
- 综合我们的努力，我们推出了 MobileDiffusion，这是一种专为设备上部署而设计的超高效扩散模型框架。
2 Related works
如第 1 节所述，为了提高文本到图像扩散模型的推理效率并最终实现其在移动设备上的部署，主要关注两个领域：架构效率和采样效率。我们简要回顾一下前期的工作。
架构效率 之前有少量研究讨论了扩散模型的架构效率。早期的扩散模型 [10,45,47] 采用了带有 transformer 块的 UNet 结构。最近的研究 [3,41] 开始将视觉 transformer 引入扩散。[18] 提出了 UViT，一种带有 transformer 主干的 UNet 结构。随着 transformer 的不断进步，这些方法的不同之处在于通过 UNet 或 patchify 来平衡计算和成本的方法。[25] 介绍了一种方法，通过有选择地从教师模型中删除特定的块，将较大的扩散模型蒸馏成较小的学生模型。同时，[27] 提出了一种有效的架构搜索方法。他们使用鲁棒的训练技术 [21,64] 训练了一个带有冗余块的 UNet，然后根据指标修剪某些块，从而得到一个适合蒸馏的架构。在我们的工作中，我们与 [18] 中关于 UNet 中的 Transformer 组件的见解保持一致，同时引入了第 3 节中阐述的独特观点。与 [25, 27] 等方法相比，我们的重点不仅限于单纯的块移除。相反，我们对架构进行了更细致的分析和修改。此外，与 [25, 27] 不同，我们选择不使用来自较大模型的知识蒸馏。我们的观察表明，从头开始训练我们的模型会产生令人满意的结果。
采样效率已出现了多种减少扩散模型所需采样步骤的策略，大致分为两类。第一种方法涉及开发快速求解器以更有效地求解与去噪过程相关的微分方程，从而减少必要的离散化步骤 [2, 12, 23, 33, 33]。第二种方法利用知识蒸馏技术来压缩采样轨迹 [27, 35, 37, 48, 52]。其中，一致性蒸馏 [35, 36, 52] 已显示出将采样步骤数减少到 4-8 的前景，同时仍促进下游生成任务，如 [59] 所强调的。尽管如此，进一步减少步骤数的尝试往往会导致输出质量下降。最近的创新旨在通过实现单步生成来彻底改变这一领域 [30,49,62,63]。一步扩散的一个重要研究方向是扩散-GAN 混合 [60]，即具有对抗性目标的微调扩散模型 [49, 62]。尽管取得了这些进展，但将一步扩散模型应用于更广泛的下游任务仍然是一个尚未探索的领域。一项同时开展的工作 [28] 采用了参数高效的对抗性微调。然而，他们得到的参数高效的微调模型不利于一步生成。
3 Designing a mobile-friendly diffusion architecture
在本节中，我们将介绍如何制作高效的文本到图像扩散模型，最终实现移动设备上的亚秒级生成。遵循 [42, 45]，我们采用潜在扩散，因为它可以在潜在空间中高效地学习文本引导生成。我们的模型设计在于高效的扩散网络和轻量级 VAE 解码器。我们使用 CLIP-ViT/L14 [43] 作为文本编码器，它可以在设备上运行几毫秒。
3.1 Diffusion network
早期的扩散模型 [10, 45, 47] 采用了具有卷积和注意机制的 UNet 架构。最近的研究 [3, 41] 将视觉变换器 (ViT) 主干引入到扩散中，在扩展模型方面具有巨大潜力。我们注意到，对于较长的序列长度，变换器的成本很高。在设计预算受限的模型时需要考虑这一点。SnapFusion [27] 以最高分辨率移除潜在空间中的变换器块。ViT [41] 将空间输入修补为具有缩短长度的标记序列。
我们遵循 UViT [18] 来设计带有文本指导的扩散网络。如图 2 所示，我们利用 UNet 来减少序列长度并高效计算变压器。我们对 UViT 架构进行了新颖的更改，对两个基本构建块进行了全面研究：变压器和卷积。在整个研究过程中，我们控制训练管道（例如数据、优化器）以研究不同架构的影响。我们优化网络以减少参数和浮点运算（FLOPs）数量，同时保持质量。我们利用 MS-COCO 2014 30K 上的两个指标，即 Fréchet Inception Distance (FID) 和 CLIP-ViT B/32，并辅以视觉检查来监控每个模型变体的质量。
图2：MobileDiffusion的模型示意图。（Conv：卷积。SA：自注意力（可选）。CA：交叉注意力。 ϕ \phi ϕ：用于计算注意力权重的非线性函数。 e e e：门控线性层中的扩展因子[50]。DW：深度卷积。）
3.2 Optimize transformer
扩展主干网络。**模型优化并不总是意味着缩小规模。**通过重新分配模型参数，我们可以实现更高效的计算。我们注意到变压器在低分辨率下是高效的。因此，我们通过将更多的变压器层移动到瓶颈中，同时保持参数总数不变，来扩展 UViT 的变压器主干。此外，自注意力机制的计算复杂度为 O ( n d 2 + n 2 d ) O\left(n d^2+n^2 d\right) O(nd2+n2d)，其中 n n n 是序列长度， d d d 是通道维度。在低分辨率（例如 16 × 16 16 \times 16 16×16）下，通道维度对计算的贡献更大。我们的经验观察表明，略微减少瓶颈中的通道维度不会对生成样本的定量指标或视觉质量产生不利影响。另一方面，尝试在大幅减少通道维度的情况下堆叠更多的变压器块会产生不利影响，导致视觉质量明显下降，表现为对象组合不佳和复杂的伪影。我们发现 1024 是在效率和质量之间取得良好平衡的最佳通道维度。
在文本到图像的扩散模型中，自注意力在捕捉长距离依赖关系方面发挥了关键作用，尽管在较高分辨率下会带来显著的计算成本。例如，在 32 × 32 32 \times 32 32×32 的分辨率下，自注意力的序列长度为 1024。在之前的工作中 [ 18 , 42 ] [18,42] [18,42]，自注意力和交叉注意力层在接近高效设计时一起移动。经过进一步调查，我们发现保留交叉注意力层而仅在高分辨率下丢弃自注意力层并不会导致性能下降。我们推测交叉注意力在不同分辨率下都是有用的，因为文本指导对图像的全局布局和局部纹理都至关重要。值得注意的是，由于文本嵌入的序列长度较小（例如，SD 的文本嵌入长度为 77），因此在高分辨率下交叉注意力的计算成本显著低于自注意力。因此，仅删除自注意力层可以显著提高效率。根据这些见解，我们采用了一种在提高效率的同时保持性能的设计：在最高分辨率 ( 64 × 64 ) (64 \times 64) (64×64) 下完全移除变压器块；消除 32 × 32 32 \times 32 32×32 分辨率和外部 16 × 16 16 \times 16 16×16 堆栈中的变压器块中的自注意力层；保留内部 16 × 16 16 \times 16 16×16 堆栈和最内层瓶颈堆栈中的完整变压器块。
在注意力层中，键和值都是从相同的输入 x x x 投影的，表示为 K = x ⋅ W K K=x \cdot W_K K=x⋅WK 和 V = x ⋅ W V V=x \cdot W_V V=x⋅WV。我们的实验表明，采用参数共享方案，特别是对于自注意力层设置 W K = W V W_K=W_V WK=WV 并不会对模型性能产生不利影响。因此，我们选择实施此参数共享策略，从而大约减少了 5 % 5 \% 5% 的参数量。
将gelu替换为swish。GLU（门控线性单元）[50]采用了gelu激活函数，这一选择在移动设备上的float16或int8推理时，由于涉及近似中的三次操作而不幸引入了数值不稳定性问题 [8]。此外，gelu激活的计算速度较慢，需要依赖特定的硬件优化 [ 54 , 56 ] [54,56] [54,56]。因此，我们建议将gelu替换为swish，swish保持了相似的形状，但更具成本效益和计算效率。重要的是，实验证明这种替换不会导致指标或感知质量下降。
将softmax微调为relu。给定key K K K，query Q Q Q 和value V V V，注意力计算为 x = ϕ ( K ⊤ Q ) V \mathbf{x}=\phi\left(K^{\top} Q\right) V x=ϕ(K⊤Q)V，其中函数 ϕ ( ⋅ ) \phi(\cdot) ϕ(⋅)为softmax。然而，由于在序列长度上进行指数运算和求和的非高效并行化，softmax函数，表示为softmax ( x ) = e x / ∑ j = 1 N e x j (\mathbf{x})=e^{\mathbf{x}} / \sum_{j=1}^N e^{x_j} (x)=ex/∑j=1Nexj和 x = ( x 0 , x 1 , … , x N ) \mathbf{x}=\left(x_0, x_1, \ldots, x_N\right) x=(x0,x1,…,xN)，在计算上是昂贵的。相反，像relu这样的逐点激活提供了一种更快的替代方案，不依赖于特定的硬件优化，并且可以作为一种可行的替代品 [57]。因此，我们建议在注意力计算中使用relu。一个有趣的发现是，不需要从头训练一个relu注意力模型；相反，从一个预训练的softmax注意力模型进行微调就足够了，这一微调过程可以快速完成，例如，在10,000次迭代内完成。有关此修改的视觉消融结果，请参见附录。
修剪前馈层。在 Transformer 的前馈层中，扩展率默认为 4，使用门控单元 [50] 后，该比例进一步翻倍。这显著增加了参数数量。例如，通道维度 1280 在投影后激增至 10240。如此高的维度可能会限制资源受限的移动应用程序。经过彻底的消融，我们发现将扩展率修剪为 3 可获得几乎相同的性能。通过这种调整，FID 分数仅略微增加 0.27，同时导致参数数量减少 10%。
3.3 Optimize convolution
可分离卷积。在经典卷积层中的残差块通常涉及大量参数，因此有多种努力旨在提高卷积层的参数效率。在这一背景下，一种行之有效的方法是可分离卷积 [19]。我们观察到，将UNet中较深层次的经典卷积层替换为轻量级可分离卷积层可以获得类似的性能。因此，我们将UNet中的所有卷积层替换为可分离卷积层，最外层除外。我们采用的可分离卷积块如图2所示，与ConvNeXt [31]有相似之处，但采用了较小的 3 × 3 3 \times 3 3×3内核尺寸。尽管我们试验了更大的内核尺寸，例如 7 × 7 7 \times 7 7×7和 9 × 9 9 \times 9 9×9，但发现它们并未提供额外的改进。
裁剪冗余的残差块。对高分辨率特征图进行卷积运算尤其费力，而裁剪是提高模型效率的一种简单方法。通过全面的网络搜索，我们将所需残差块的数量从22（在SD中）减少到更高效、更精简的12。更具体地说，除最内层块外，我们为每个块设置1层而不是SD中的2层，以在资源消耗和模型性能之间保持平衡。
3.4 Model details
利用前面讨论的优化技术，我们通过对参数和 FLOPs 的数量施加上限来精心优化我们对可行架构候选的选择。我们的目标是 4 亿个参数和 200 GFLOPs，以实现设备上的即时生成。值得注意的是，为了便于模型搜索，我们将下层堆栈、上层堆栈和中间瓶颈的最内层组合成一个统一的模块，因为它们共享相同的特征图维度。这对应于图 2 所示的架构中间的 5 层。在这个受限的搜索空间内，我们系统地探索了变压器层数和最内层通道维度的变化，专注于显著影响模型性能的关键架构参数。由于与训练扩散模型相关的计算需求，我们优先优化这些关键参数以提高模型的效率和有效性。在表 1 中，我们对 MobileDiffusion 与 SD UNets、针对设备优化的 SnapFusion UNet [27] 和 DiT XL/2 [41] 进行了比较分析。它强调 MobileDiffusion 表现出最高的效率。与 UNets 相比，MD 将更多的计算从卷积转移到注意力，并且比纯 Transformer 架构包含更少的 FLOPs。
3.5 Optimize VAE
在潜变量扩散中，VAE将一个 H × W × 3 H \times W \times 3 H×W×3的图像转换为 H f × W f × c \frac{H}{f} \times \frac{W}{f} \times c fH×fW×c的潜变量。SD [45] 使用 f = 8 f=8 f=8 和 c = 4 c=4 c=4，而 EMU [9] 使用 f = 8 f=8 f=8 和 c = 16 c=16 c=16。使用较小的通道数 c c c可以实现更好的压缩效果，但重建质量较差。经过实证研究，我们选择 f = 8 f=8 f=8 和 c = 8 c=8 c=8 来训练我们的VAE。这一修改提升了图像重建的质量，如图3所示。与用于潜变量扩散模型的其他VAEs相似，我们的VAE使用多种损失函数的组合进行训练，包括 L 2 \mathcal{L}_2 L2重建损失、正则化的KL散度、感知损失和对抗损失。我们使用与扩散模型训练相同的数据集，采用256的批量大小进行200万次迭代训练VAE。正如[27]中所指出的，当使用较少的扩散步骤进行评估时，图像解码器成为一个不可忽略的组件。为了进一步提高效率，我们通过修剪模型的宽度和深度设计了一个轻量级的解码器架构。我们在冻结编码器的情况下训练轻量级解码器40万次迭代。在目标中去掉了不必要的KL正则化，并减少了对抗损失的权重。蒸馏的解码器带来了显著的性能提升，比SD的解码器快3倍，如表4所示。
图 3：SD 解码器、我们的解码器和我们的蒸馏解码器之间的 VAE 重建比较。
4 Elucidating the design space of UFOGen
在本节中，我们详细介绍了我们对超快速一步式扩散模型的研究，重点是提高图像质量，并且至关重要的是，能够将广泛的下游应用集成到一步式生成范式中。我们的研究主要集中在 UFOGen [62] 上，这是一种创新的对抗性微调方法，已显示出在一步式文本到图像生成方面的巨大潜力。我们首先简要概述 UFOGen，然后讨论我们为改进其训练方法而进行的一系列消融研究
4.1 Overview of UFOGen
UFOGen 通过对现有扩散模型的对抗性微调实现了一步生成。这个过程包括一步生成器 G θ G_\theta Gθ 与判别器 D ϕ D_\phi Dϕ 的联合训练，目标如下：
min θ max ϕ E q ( x 0 ) q ( x t − 1 ∣ x 0 ) , p θ ( x 0 ′ ) p θ ( x t − 1 ′ ∣ x 0 ′ ) [ [ log ( D ϕ ( x t − 1 , t ) ) ] + [ log ( 1 − D ϕ ( x t − 1 ′ , t ) ) ] ⏟ adversarial loss + λ γ t ∥ x 0 − x 0 ′ ∥ 2 ⏟ diffusion loss ] , ( 1 ) \begin{aligned} & \underset{\theta}{\min } \max _\phi \mathbb{E}_{q\left(x_0\right) q\left(x_{t-1} \mid x_0\right), p_\theta\left(x_0^{\prime}\right) p_\theta\left(x_{t-1}^{\prime} \mid x_0^{\prime}\right)}[ \\ & \underbrace{\left[\log \left(D_\phi\left(x_{t-1}, t\right)\right)\right]+\left[\log \left(1-D_\phi\left(x_{t-1}^{\prime}, t\right)\right)\right]}_{\text {adversarial loss }}+\underbrace{\lambda \gamma_t\left\|x_0-x_0^{\prime}\right\|^2}_{\text {diffusion loss }}],\quad(1) \end{aligned} θminϕmaxEq(x0)q(xt−1∣x0),pθ(x0′)pθ(xt−1′∣x0′)[adversarial loss [log(Dϕ(xt−1,t))]+[log(1−Dϕ(xt−1′,t))]+diffusion loss λγt∥x0−x0′∥2],(1)
其中 q ( x 0 ) q\left(x_0\right) q(x0) 是训练数据分布， q ( x t − 1 ∣ x 0 ) q\left(x_{t-1} \mid x_0\right) q(xt−1∣x0) 是前向扩散过程。生成器通过生成一个干净的图像 x 0 ′ = G θ ( x t , t ) ∼ p θ ( x 0 ′ ) x_0^{\prime}=G_\theta\left(x_t, t\right) \sim p_\theta\left(x_0^{\prime}\right) x0′=Gθ(xt,t)∼pθ(x0′) 来参数化， x t − 1 ′ ∼ p θ ( x t − 1 ′ ∣ x 0 ′ ) x_{t-1}^{\prime} \sim p_\theta\left(x_{t-1}^{\prime} \mid x_0^{\prime}\right) xt−1′∼pθ(xt−1′∣x0′) 是在生成器输出 x 0 ′ x_0^{\prime} x0′ 上应用前向扩散过程。除了对抗性损失之外，UFOGen 的目标中还包含原始的扩散损失，这对于稳定训练被证明是至关重要的。生成器和判别器都具有相同的结构，并通过预训练的扩散模型进行初始化。UNet 判别器聚合输出位置的对数以做出决策。
UFOGen 已证明其能够通过一步从文本提示生成高质量图像，从而显著提高扩散模型中的采样效率。然而，UFOGen 的训练策略尚未得到很好的研究。此外，将 LoRA [20]、ControlNet [65] 和 T2I 适配器 [39] 等单独训练的模块与一步采样无缝结合仍未得到探索。虽然最初的 UFOGen 研究 [62] 展示了从深度图和精明边缘成功生成图像，但这些结果是通过特定于任务的对抗性训练实现的，因此需要为每个新的条件任务单独复制一份副本——这远非理想情况。为此，我们仔细研究了 UFOGen 的设计空间，并从相关方法中汲取灵感。
我们最终的设计选择是通过比较文本到图像生成和下游条件生成任务的不同变体获得的。
4.2 Reconstruction term in the training objective
UFOGen 训练目标中值得探索的第一个设计方面是重建项的作用。如公式 1 所述，UFOGen 的公式整合了对抗性损失和重建损失。对抗性损失的目的是提高通过单一扩散步骤生成的图像的视觉准确性。传统上，扩散模型的单步生成过程旨在预测给定噪声输入的预期原始图像，用 E [ x 0 ∣ x t ] \mathbb{E}\left[x_0 \mid x_t\right] E[x0∣xt] 表示。因此，由于数据分布内在的平均化，通常会生成一些模糊的图像，缺乏细节。加入对抗性损失项使模型能够生成更清晰的图像，尤其是在噪声水平较高时。相比之下，训练过程中重建损失的功能更为复杂。它被认为在对抗性训练阶段提供稳定作用，提供简单的学习指导。除此之外，我们认为它在维护预训练扩散模型特征的完整性方面也发挥了重要作用。这种保留是至关重要的；在其缺失的情况下，单步模型可能会明显偏离预训练模型的原始特征空间。这样的偏离可能会严重限制与专门为原始扩散框架量身定制的下游模块的兼容性。
我们的猜测表明，为了鼓励保留预训练模型的特征，可以提供增强的正则化，确保单步模型与基础扩散模型的特征保持紧密一致。为此，我们研究了一些替代方案以替代公式 1 中的重建损失。第一个是蒸馏损失，它明确将 UFOGen 的生成器的特征与教师模型对齐：
L distill = ∥ sg ( G teacher ( x t , t ) ) − x 0 ′ ∥ 2 , ( 2 ) \mathcal{L}_{\text {distill }}=\left\|\operatorname{sg}\left(G^{\text {teacher }}\left(x_t, t\right)\right)-x_0^{\prime}\right\|^2, \quad(2) Ldistill =sg(Gteacher (xt,t))−x0′2,(2)
其中 G θ teacher G_\theta^{\text {teacher }} Gθteacher 是预训练的教师模型， sg ( ⋅ ) \operatorname{sg}(\cdot) sg(⋅) 表示停止梯度。注意，蒸馏损失类似于对抗性扩散蒸馏（ADD）[49] 目标中的正则化项。
更一般地，我们还测试了 EMA 蒸馏目标，形式化如下：
L ema = ∥ sg ( G θ E M A ( x t , t ) ) − x 0 ′ ∥ 2 , ( 3 ) \mathcal{L}_{\text {ema }}=\left\|\operatorname{sg}\left(G_\theta^{\mathrm{EMA}}\left(x_t, t\right)\right)-x_0^{\prime}\right\|^2,\quad(3) Lema =sg(GθEMA(xt,t))−x0′2,(3)
其中 G θ EMA G_\theta^{\text {EMA }} GθEMA 是 G θ G_\theta Gθ 的指数移动平均。由于 G θ G_\theta Gθ 源自预训练模型，EMA 机制有效地保留了预训练模型的大部分信息。这种损失也可以确保保留预训练模型的本质，同时提供更多的灵活性。请注意，公式 2 中的蒸馏损失是我们损失的特例，其中 EMA 衰减参数设置为 1。我们在第 5.3 节对这些损失项进行了全面的消融研究。
4.3 Parameter-efficient adversarial fine-tuning
以 UFOGen [62] 和 ADD [49] 为例，当前的对抗性扩散微调方法会调整整个模型。这种广泛的参数修改可能会导致内部特征表示发生重大变化，从而可能削弱模型有效处理下游任务的能力。参数高效微调是一种引人注目的解决方案，它借鉴了视觉和语言领域的成功应用 [11,20]。在这些方法中，LoRA（低秩自适应）微调 [20] 因其改进扩散模型的能力而脱颖而出 [36]，它提供了一种平衡的方法，既支持快速采样，又保持任务适应性。
受 LCM-LoRA 成功的启发，我们探索了将 LoRA 微调集成到 UFOGen 训练框架中的潜力。我们将 LoRA 专门应用于生成器，从预先训练的扩散框架初始化模型权重，同时将 LoRA 层设置为以零效应初始化。鉴别器仍然从预先训练的扩散模型初始化。
总结：高效扩散模型（MobileDiffusion）从模型架构和采样效率两方面进行了优化。模型架构方面：1. 模型优化并不总是意味着缩小规模。即将transformer block移动到bottleneck实现高效计算。2. 移除高分辨率下的self-attention 3.key value投影层权重共享 4. gelu替换为swish 5. softmax替换为relu 6. 减少通道数 7. 采用深度可分离卷积 8. 移除冗余残差块 9. ve采用轻量化decoder
采样效率方面：使用一步式蒸馏技术UFOGen（加了权重正则化和LoRA微调）
