回复
4倍内存效率，生成和超分双SOTA！清华&智谱AI发布最新Inf-DiT模型
文章链接：https://arxiv.org/pdf/2405.04312
github 链接：https://github.com/THUDM/Inf-DiT
扩散模型在近年来的图像生成中表现出了显著的性能。然而，由于生成超高分辨率图像（如 4096 × 4096）时内存需求呈二次方增加，生成图像的分辨率通常限制在 1024 × 1024。
本文提出了一种单向块注意力机制，可以在推理过程中自适应地调整内存开销并处理全局依赖关系。基于这个模块，本文采用 DiT 结构进行上采样，并开发了一种能够对各种形状和分辨率的图像进行上采样的无限超分辨率模型。综合实验表明，本文的模型在机器和人类评估中都达到了生成超高分辨率图像的最新技术水平。与常用的 UNet 结构相比，本文的模型在生成 4096 × 4096 图像时可以节省超过 5 倍的内存。
介绍
近年来，扩散模型取得了迅速进展，显著推动了图像生成和编辑领域的发展。尽管有这些进展，但仍然存在一个关键限制：现有图像扩散模型生成的图像分辨率通常限制在 1024×1024 像素或更低，这在生成超高分辨率图像时构成了重大挑战。而超高分辨率图像在各种实际应用中是不可或缺的，包括复杂的设计项目、广告、海报和壁纸的制作等。
一种常用的生成高分辨率图像的方法是级联生成，即首先生成低分辨率图像，然后应用多个上采样模型逐步提高图像的分辨率。这种方法将高分辨率图像的生成分解为多个任务。基于前一阶段生成的结果，后一阶段的模型只需进行局部生成。在级联结构的基础上，DALL-E2 和 Imagen 都能有效生成分辨率为 1024 的图像。
对于上采样到更高分辨率图像的最大挑战是显著的 GPU 内存需求。例如，如果在图像推理中使用广泛采用的 U-Net 架构（如 SDXL，见下图 2），观察到随着分辨率的增加，内存消耗急剧上升。具体而言，生成一个 4096×4096 分辨率的图像（包含超过 1600 万个像素）需要超过 80GB 的内存，这超出了标准的 RTX 4090 或 A100 显卡的容量。此外，高分辨率图像生成模型的训练过程加剧了这些需求，因为它需要额外的内存来存储梯度、优化器状态等。
LDM 通过利用变分自编码器（VAE）来压缩图像并在较小的潜在空间中生成图像，从而减少了内存消耗。然而，文中也强调，过高的压缩比会显著降低生成质量，严重限制了内存消耗的减少。
基于这一算法，本文优化了扩散 Transformer（DiT），并训练了一个名为 Inf-DiT 的模型，该模型能够对不同分辨率和形状的图像进行上采样。此外，设计了几种技术，包括提供全局图像 embedding 以增强全局语义一致性，并提供 zero-shot 文本控制能力，以及通过交叉注意力机制提供所有相邻的低分辨率（LR）块以进一步增强局部一致性。评估结果表明，Inf-DiT 在机器和人类评估中均显著优于其他高分辨率生成模型。
主要贡献如下：
- 基于这些方法，训练了一个图像上采样扩散模型 Inf-DiT，这是一种 700M 的模型，能够对不同分辨率和形状的图像进行上采样。Inf-DiT 在机器评估（HPDv2 和 DIV2K 数据集）和人类评估中均达到了最新技术水平。
- 设计了多种技术来进一步增强局部和全局一致性，并提供灵活的文本控制的 zero-shot 能力。
方法
单向块注意力 (UniBA)
生成超高分辨率图像的关键障碍是内存限制
随着图像分辨率的增加，网络中对应的隐藏状态的大小呈二次方增长。例如，仅一层中形状为 2048 × 2048 × 1280 的单个隐藏状态就需要 20GB 的内存，使得生成非常大的图像变得异常艰难。如何避免在内存中存储整个图像的隐藏状态成为关键问题。
- 块之间的生成依赖关系是单向的，并且可以形成一个有向无环图（DAG）。
- 每个块对其他块只有少量的直接（一级）依赖关系，因为块及其直接依赖块的隐藏状态需要同时保存在内存中。
此外，为了确保整个图像的一致性，还需要确保每个块具有足够大的感受野，以处理长程依赖关系。
根据上述条件和分析，本文选择了一种高效的实现方式，即下图 3 所示的单向块注意力（UniBA）。对于每一层，每个块直接依赖于三个一阶相邻块：顶部的块、左侧的块和左上角的块。例如，如果本文采用了 Inf-DiT 的基础架构 Diffusion Transformer（DiT）架构，则块之间的依赖关系是注意力操作，其中每个块的查询向量与其左上角和本身的四个块的键值向量进行交互，如下图 3 所示。
形式上，Transformer 中的 UniBA 过程可以表示为：
需要注意的是，尽管每个块在每层中只关注少量相邻块，但随着特征层层传播，块可以间接与远处的块交互，从而捕捉长短距离关系。本文的设计与自然语言模型 Transformer-XL 具有相似的精神，可以看作是本文的一维情况的特例。
使用 O(N)内存消耗的推理过程
尽管本文的方法可以顺序生成每个块，但它不同于自回归生成模型，在自回归生成模型中，下一个块依赖于前一个块的最终输出。在本文的模型中，只要它们的依赖块的集合已经生成，就可以并行生成任意数量的块。基于这一特性，本文实现了一个简单但有效的推理过程。如上面图 3 所示，本文一次生成 n×n 个块，从左上到右下。生成一组块后，本文丢弃不再使用的隐藏状态（即 KV 缓存），并将新生成的 KV 缓存附加到内存中。
在实际操作中，尽管对于不同的 n，图像生成的总 FLOPs 保持不变，但由于操作初始化时间和内存分配时间等开销，当 n 增加时，生成时间会减少。因此，在内存限制允许的情况下，选择最大的 n 是最优的。
基本模型架构
下图 4 概述了本文模型 Inf-DiT 的架构。该模型使用了类似 DiT 的主干结构，DiT 将视觉 Transformer（ViT）应用于扩散模型，并证明了其有效性和可扩展性。与基于卷积的架构（如 UNet）相比，DiT 仅使用注意力作为块之间的交互机制，这使得单向块注意力的实现变得方便。为了适应单向块注意力并增强上采样性能，本文进行了如下的几项修改和优化。
模型输入
Inf-DiT 首先将输入图像划分为多个不重叠的块，然后将这些块进一步划分为边长等于 patch 大小的 patch。与 DiT 不同，考虑到压缩损失（如颜色偏移和细节损失），Inf-DiT 在 RGB 像素空间中进行 patch 划分，而不是在潜在空间中。在超分辨率因子为 f 的情况下，Inf-DiT 首先将低分辨率 RGB 图像条件上采样 f 倍，然后在特征维度上将其与扩散的噪声输入连接起来，再输入到模型中。
位置编码
与可以通过卷积操作感知位置关系的基于 UNet 的扩散模型不同，Transformer 中的所有操作（包括自注意力和前馈神经网络）都是置换不变函数。因此，基于 Transformer 的模型需要辅助输入显式位置信息以学习 patch 之间的关系。正如最近在大型语言模型中的研究所示，相对位置编码在捕捉单词位置相关性方面比绝对位置编码更有效，本文参考了旋转位置编码（RoPE）的设计，该设计在长上下文生成中表现良好，并将其适配为二维形式用于图像生成。具体来说，本文将隐藏状态的通道分成两半，一半用于编码 x 坐标，另一半用于编码 y 坐标，并在这两个部分中应用 RoPE。
本文创建了一个足够大的 RoPE 位置编码表，以确保在生成过程中满足需求。为了确保模型在训练期间能看到位置编码表的所有部分，本文采用了随机起始点：对于每个训练图像，本文随机分配一个位置（x,y）作为图像的左上角，而不是默认的（0,0）。
全局和局部一致性
使用 CLIP 图像 embedding 实现全局一致性
低分辨率（LR）图像中的全局语义信息（如艺术风格和物体材质）在上采样过程中起着至关重要的作用。然而，与文本生成图像模型相比，上采样模型有一个额外的任务：理解和分析低分辨率图像的语义信息，这大大增加了模型的负担。这在没有文本数据进行训练时尤其具有挑战性，因为高分辨率图像很少有高质量的配对文本，使得这些方面对模型来说很难处理。
使用邻近 LR 交叉注意力实现局部一致性尽管将 LR 图像与噪声输入连接起来已经为模型学习 LR 和 HR 图像之间的局部对应关系提供了良好的归纳偏置，但仍然可能存在连续性问题。原因在于，对于给定的 LR 块，有多种上采样的可能性，这需要结合多个邻近 LR 块进行分析以选择一个解决方案。假设上采样仅基于其左侧和上方的 LR 块进行，它可能会选择一个与右侧和下方的 LR 块冲突的 HR 生成方案。那么，当上采样右侧的 LR 块时，如果模型认为符合其对应的 LR 块比与左侧块连续更重要，则会生成与先前块不连续的 HR 块。一个简单的解决方案是将整个 LR 图像输入到每个块中，但当 LR 图像的分辨率也很高时，这样做成本太高。
为了解决这个问题，本文引入了邻近 LR 交叉注意力。在 Transformer 的第一层，每个块对周围的 3×3 LR 块进行交叉注意力，以捕捉附近的 LR 信息。本文的实验表明，这种方法显著降低了生成不连续图像的概率。值得注意的是，这个操作不会改变本文的推理过程，因为在生成之前整个 LR 图像是已知的。
本文进一步设计了包括无类别指导的连续性、基于 LR 的噪声初始化、QK 规范化等技术。
实验
在本节中，本文首先介绍 Inf-DiT 的详细训练过程，然后通过机器和人类评价全面评估 Inf-DiT 的性能。结果表明，Inf-DiT 在超高分辨率图像生成和上采样任务中均优于所有基线模型。最后，本文进行消融研究，以验证本文设计的有效性。
训练细节
数据集
数据集由分辨率高于 1024×1024 且美学评分高于 5 的 LAION-5B 子集和来自互联网的 10 万张高分辨率壁纸组成。与之前的工作[20,23,30]相同，本文在训练期间使用固定大小的 512×512 分辨率的图像裁剪。由于上采样可以仅使用局部信息进行，因此在推理时可以直接在更高分辨率下进行，这对大多数生成模型来说并不容易。
数据处理
由于扩散模型生成的图像通常包含残留噪声和各种细节不准确性，因此增强上采样模型的鲁棒性以解决这些问题变得至关重要。本文采用类似于 Real-ESRGAN 的方法，对训练数据中的低分辨率输入图像进行各种降质处理。
在处理分辨率高于 512 的训练图像时，有两种替代方法：直接执行随机裁剪，或将较短的一侧调整为 512 后再进行随机裁剪。直接裁剪方法保留了高分辨率图像中的高频特征，而调整后裁剪方法则避免了频繁裁剪出单色背景区域，这对模型的收敛性不利。因此，在实践中，本文随机选择这两种处理方法中的一种来裁剪训练图像。
训练设置
在训练期间，本文设置块大小为 128，patch 大小为 4，这意味着每个训练图像被分成 4×4 个块，每个块有 32×32 个 patch。本文采用 EDM 框架进行训练，并将上采样因子设置为 4 倍。由于上采样任务更关注图像的高频细节，本文调整了训练噪声分布的均值和标准差为-1.0 和 1.4。为了解决训练期间的溢出问题，本文使用了具有更大数值范围的 BF16 格式。本文的 CLIP 模型是一个在 Datacomp 数据集上预训练的 ViT-L/16。由于 CLIP 只能处理 224×224 分辨率的图像，本文首先将低分辨率图像调整为 224×224，然后将其输入到 CLIP 中。
机器评价
在这部分中，本文对 Inf-DiT 在超高分辨率图像生成任务上与最先进方法进行定量比较。基线包括两大类高分辨率生成：1. 直接高分辨率图像生成，包括 SDXL 的直接推理、MultiDiffusion、ScaleCrafte 等；2. 基于超分辨率技术的高分辨率图像生成，包括 BSRGAN、DemoFusion 等。本文采用 FID（Fréchet Inception Distance）来评估超高分辨率生成的质量，这在图像生成任务中广泛用于评估图像的感知质量。为了进一步验证本文模型的超分辨率能力，本文还将其与经典的超分辨率模型在典型超分辨率任务上进行了基准测试。
超高分辨率生成本文使用 HPDv2 的测试集进行评估。它包含 3200 个提示，并分为四个类别：“动画”、“概念艺术”、“绘画”和“照片”。这允许对模型在各个领域和风格中的生成能力进行全面评估。本文在两个分辨率上进行测试：2048x2048 和 4096x4096。对于基于超分辨率的模型，本文首先使用 SDXL 生成 1024x1024 分辨率的图像，然后在没有文本的情况下对其进行上采样。本文使用 BSRGAN 的 2× 和 4× 版本分别进行 2048x2048 和 4096x4096 的生成。尽管 Inf-DiT 是在 4× 上采样的设置下进行训练的，但本文发现它在较低的上采样倍数下也能很好地泛化。因此，对于 2048x2048 的生成，本文直接将 LR 图像从 1024x1024 调整为 2048x2048，并与噪声输入连接起来。本文从 LAION-5B 中随机选择了 3200 张 2048x2048 和 4096x4096 的图像作为真实图像的分布。
如下表 2 所示，本文的模型在所有指标上均达到了最先进水平。这表明，作为一个超分辨率模型，本文的模型不仅在任意尺度上表现出色，而且在最大限度地保留全局和详细信息的同时，还能恢复与原始图像非常接近的结果。
人类评价
为了进一步评估 Inf-DiT 并更准确地从人类视角反映其生成质量，本文进行了人类评价。比较设置与上节中相同，不过本文排除了 MultiDiffusion 和 Direct Inference 因其非竞争性的结果。对于每个类别，本文随机选择了十组比较集，每组集合包含了四个模型的输出，共计 40 组形成了人类评价数据集。为了保证公平性，在每个比较集中本文对模型输出的顺序进行了随机化。人类评估者被要求根据三个标准评估模型：细节真实性、全局连贯性和与原始低分辨率输入的一致性。每位评估者平均收到 20 组图像。在每个集合中，评估者需要根据三个标准对由四个模型生成的图像进行从高到低的排名。
本文最终收集了 3600 组比较。如下图 7 所示，本文的模型在所有三个标准中均优于其他三种方法。值得注意的是，其他三种模型中的每一种在至少一个评估标准上排名相对较低，而 Inf-DiT 在所有三个标准上的得分都最高：细节真实性、全局连贯性和与低分辨率输入的一致性。这表明本文的模型是唯一能够在高分辨率生成和超分辨率任务中同时表现出色的模型。
迭代上采样
由于本文的模型可以对任意分辨率的图像进行上采样，测试模型是否能够迭代上采样自身生成的图像是一个自然的想法。在这项研究中，本文在一张 322 分辨率的图像上进行实验，通过三次迭代上采样，将其生成一张 2048x2048 分辨率的图像，即 64 倍放大。下图 8 展示了这个过程的两个案例。在第一个案例中，模型成功地在三个阶段的上采样后生成了一张高分辨率图像。它在不同分辨率的上采样中生成了不同频率的细节：脸部轮廓、眼球形状和个别睫毛。然而，模型很难纠正在早期阶段生成的不准确性，导致错误的积累。第二个样本展示了这个问题。本文将这个问题留给未来的工作。
消融研究
相关工作
扩散图像生成
扩散模型已经成为图像生成领域的焦点，近年来取得了一系列突破性进展。最初于 2015 年引入，并通过诸如 DDPM 和 DDIM 等工作进一步发展，这些模型利用随机扩散过程，概念化为马尔可夫链，将简单的先验分布（如高斯噪声）转化为复杂的数据分布。这一方法在生成的图像质量和多样性方面取得了令人印象深刻的成果。
近期的增强显著提升了扩散模型的生成能力。CDM 创建了一个级联生成 pipeline，其中包括多阶段的超分辨率模型，可应用于大型预训练模型。引入潜在扩散模型（LDMs）代表了一个重要的扩展，它结合了潜在空间，提升了效率和可扩展性。除此之外，网络架构的优化也取得了显著进展。扩散 Transformer（DiT）的出现取代了 U-Net，使用 ViT 进行噪声预测。
图像超分辨率
这里 D 和 F 分别表示退化过程和超分辨率模型。δ和θ 代表参数。
近年来，盲目 SR 一直是主要关注的焦点：其中退化过程是未知的且可学习的。这一视角导致了有效的建模技术的发展，例如 BSRGAN 和 Real-ESRGAN。最近，基于扩散的 SR 方法取得了令人兴奋的结果。这些工作专注于对预训练的文本到图像扩散模型进行微调，以利用其优秀的生成能力。具体来说，DiffBir 在预训练的稳定扩散模型上使用了 ControlNet，而 PASD 通过执行像素感知的交叉注意力来增强它。这两种方法在固定分辨率超分辨率方面取得了相当大的成功，但不能直接用于更高的分辨率。
超高分辨率图像上采样器
目前，图像生成方法在生成超高分辨率图像方面存在着内存限制和训练效率问题。在这种情况下，MultiDiffusion 和 Mixture of Diffusers 将多个扩散生成过程绑定在一起，通过将图像划分为重叠的块，分别处理每个块，然后将它们拼接在一起，旨在保持块之间的连续性。然而，由于它们仅使用局部加权平均进行聚合，导致了交互效率低下，使得很难确保图像的全局一致性。
鉴于这一问题，DemoFusion 和 ScaleCrafter 采用了扩张策略，包括扩张采样和扩张卷积核，旨在获取更多的全局信息。这些方法确实在全局语义水平上取得了改进，而无需额外的训练。然而，训练和生成之间的巨大差异导致这些方法很容易产生不合逻辑的图像。
Inf-DiT 能够对任何生成模型生成的图像执行上采样，在这里展示了更多的情况。
结论
在这项工作中，本文观察到生成超高分辨率图像的主要障碍是模型隐藏状态占用了大量内存。基于此，本文提出了一种单向块注意力机制（UniBA），它可以通过在块之间进行批量生成来降低空间复杂度。利用 UniBA，本文训练了 Inf-DiT，这是一种 4 倍内存效率的图像上采样器，在生成和超分辨率任务中均取得了最先进的性能。
本文转自 AI生成未来 ，作者：Zhuoyi Yang等
赞
收藏
回复
相关推荐
