MIT何恺明再次突破传统：分形递归架构引爆生成模型新纪元！
论文链接：https://arxiv.org/pdf/2502.17437
代码链接：https://github.com/LTH14/fractalgen
亮点直击
分形生成模型：首次将分形理论引入生成模型，提出了一种具有自相似性的递归生成框架。
递归模块化：通过递归调用生成模块，构建了复杂的生成系统，突破了传统模块化的限制。
高性能像素级生成：在像素级图像生成任务中，展示了优异的似然估计和生成质量，解决了现有方法的不足。
广泛应用潜力：该框架不仅适用于图像生成，还可应用于分子结构、蛋白质和生物神经网络等非顺序数据的生成，具有广泛的应用前景。
生成效果
总结速览
解决的问题
-
复杂生成模型的模块化：现有的生成模型（如扩散模型和自回归模型）虽然已经模块化，但缺乏更高层次的模块化抽象，难以构建更复杂的生成系统。
-
非顺序数据的生成：现有方法在处理像素级图像生成等非顺序数据时表现不佳，尤其是在似然估计和生成质量方面。
-
自然数据的分形特性：自然数据（如图像、蛋白质结构等）通常具有分形或近分形特性，现有生成模型未能充分利用这一特性。
提出的方案
-
分形生成模型：提出了一种新的生成模型框架，通过递归调用相同类型的生成模块来构建分形生成模型。这种模型具有自相似性，能够在不同层次上表现出复杂的结构。
-
自回归模型作为生成器：以自回归模型为例，展示了如何将分形生成模型应用于像素级图像生成任务。每个自回归模块由多个子自回归模块组成，形成分形结构。
应用的技术
-
递归模块化：通过递归调用生成模块，构建具有自相似性的分形生成模型。
-
自回归模型：使用自回归模型作为生成器，递归生成像素级图像。
-
分形理论：借鉴数学中的分形概念，设计生成模型的递归结构。
达到的效果
-
高性能生成：在像素级图像生成任务中，分形生成模型在似然估计和生成质量方面表现出色。
-
复杂结构建模：能够有效建模具有内在结构的非顺序数据，如图像、分子结构和蛋白质。
-
新研究范式：为生成模型的设计和应用提供了新的研究方向，有望推动生成模型领域的进一步发展。
分形生成模型
分形生成模型的核心思想是通过递归地从现有的原子生成模块构建更高级的生成模型。在本节中，首先介绍分形生成模型背后的高层次动机和直觉。然后，以自回归模型为例，展示如何实例化分形生成模型并将其用于建模非常高维的数据分布。
动机
分形是由简单的递归规则生成的复杂模式。在分形几何中，这些规则通常被称为“生成器”。通过不同的生成器，分形方法可以构造许多自然模式，例如云、山脉、雪花和树枝，并且与更复杂的系统相关联，例如生物神经网络的结构、非线性动力学和混沌系统。
形式上，一个分形生成器 指定了如何基于上一级生成器的一个输出 生成下一级生成器的新数据集 ：
例如，如图1所示，生成器可以通过在每个灰色框内递归调用类似的生成器来构造分形。
由于每个生成器级别可以从单个输入生成多个输出，分形框架可以在仅需要线性数量的递归级别的情况下实现生成输出的指数级增长。这一特性使其特别适合用相对较少的生成器级别来建模高维数据。
具体来说，引入了一种分形生成模型，该模型使用原子生成模块作为参数化分形生成器。通过这种方式，神经网络直接从数据中“学习”递归规则。通过将分形输出的指数增长与神经生成模块相结合，我们的分形框架能够建模高维非顺序数据。接下来，将展示如何以自回归模型作为分形生成器来实例化这一思想。
自回归模型作为分形生成器
在本节中，我们展示了如何使用自回归模型作为分形生成器来构建分形生成模型。我们的目标是建模一大组随机变量 的联合分布，但直接用单一自回归模型建模在计算上是不可行的。为了解决这个问题，我们采用了一种分而治之的策略。关键的模块化是将自回归模型抽象为一个模块化单元，用于建模概率分布 。通过这种模块化，我们可以在多个下一级自回归模型的基础上构建一个更强大的自回归模型。
假设每个自回归模型中的序列长度是一个可管理的常数 ，并且随机变量的总数为 ，其中 表示我们分形框架中的递归级别数。分形框架的第一级自回归模型将联合分布划分为 个子集，每个子集包含 个变量。形式上，我们将联合分布分解为：
每个包含 个变量的条件分布 由第二级递归的自回归模型建模，以此类推。通过递归调用这种分而治之的过程，我们的分形框架可以使用 级自回归模型高效地处理 个变量的联合分布，每级自回归模型操作在可管理的序列长度 上。
这种递归过程代表了一种标准的分而治之策略。通过递归分解联合分布，我们的分形自回归架构不仅显著降低了计算成本（与单一大型自回归模型相比），还捕捉了数据中的内在层次结构。
从概念上讲，只要数据表现出可以以分而治之方式组织的结构，就可以在我们的分形框架中自然建模。为了提供一个更具体的例子，在下一节中，我们将实现这种方法来解决像素级图像生成这一具有挑战性的任务。
图像生成实例
我们现在展示分形生成模型的一个具体实现，并将其应用于像素级图像生成这一具有挑战性的任务。尽管本文以图像生成作为测试平台，但相同的分而治之架构也可以潜在地适应其他数据领域。接下来，我们首先讨论像素级图像生成的挑战和重要性。
像素级图像生成
由于原始图像数据的高维性和复杂性，像素级图像生成仍然是生成建模中的一个重要挑战。该任务要求模型能够高效处理大量像素，同时有效学习像素之间丰富的结构模式和相互依赖性。因此，像素级图像生成已成为一个具有挑战性的基准任务，大多数现有方法仍然局限于似然估计，并且无法生成令人满意的图像，比如（Child et al., 2019; Hawthorne et al., 2022; Yu et al., 2023）。
尽管具有挑战性，像素级生成代表了一类更广泛的重要高维生成问题。这些问题旨在逐元素生成数据，但与长序列建模不同，它们通常涉及非顺序数据。例如，许多结构（如分子构型、蛋白质和生物神经网络）并不表现出顺序架构，但却包含非常高维和结构化的数据分布。通过选择像素级图像生成作为我们分形框架的实例，我们的目标不仅是解决计算机视觉中的一个关键挑战，还展示了分形方法在建模具有内在结构的高维非顺序数据方面的潜力。
架构
如图3所示，每个自回归模型将上一级生成器的输出作为其输入，并为下一级生成器生成多个输出。它还接收一张图像（可以是原始图像的一个 patch），将其分割成多个 patch，并将它们嵌入以形成 transformer 模型的输入序列。这些 patch 也被馈送到相应的下一级生成器。transformer 将上一级生成器的输出作为一个单独的 token，放置在图像 token 之前。基于这个组合序列，transformer 为下一级生成器生成多个输出。
遵循视觉transformer和图像生成模型的常见实践，将第一级生成器 的序列长度设置为 256，将原始图像划分为 16×16 的 patch。第二级生成器随后对每个 patch 进行建模，并将其进一步细分为更小的 patch，递归地继续这一过程。为了管理计算成本，我们逐步减少较小 patch 的宽度和 transformer 块的数量，因为建模较小的 patch 通常比建模较大的 patch 更容易。在最后一级，我们使用一个非常轻量级的 transformer 对每个像素的 RGB 通道进行自回归建模，并在预测上应用 256 路交叉熵损失。表1详细列出了不同递归级别和分辨率下每个 transformer 的配置和计算成本。值得注意的是，通过我们的分形设计，建模 256×256 图像的计算成本仅为建模 64×64 图像的两倍。
根据（Li et al., 2024），本文方法支持不同的自回归设计。在本工作中，主要考虑两种变体：一种是以光栅顺序、类似 GPT 的因果 transformer（AR），另一种是以随机顺序、类似 BERT 的双向 transformer（MAR）（图6）。两种设计都遵循自回归的下一个 token 预测原则，各有优缺点，详见附录 B。将使用 AR 变体的分形框架命名为 FractalAR，将使用 MAR 变体的分形框架命名为 FractalMAR。
与尺度空间自回归模型的关系
最近，一些模型被提出用于执行尺度间的自回归图像生成。这些尺度空间自回归模型与我们提出的方法之间的一个主要区别在于，它们使用单一的自回归模型逐尺度预测 token。相比之下，我们的分形框架采用分而治之的策略，通过生成子模块递归建模原始像素。另一个关键区别在于计算复杂度：尺度空间自回归模型在生成下一尺度的 token 时需要对整个序列执行全注意力机制，这导致计算复杂度显著增加。
例如，在生成分辨率为 256×256 的图像时，在最后一个尺度中，尺度空间自回归模型的每个注意力块中的注意力矩阵大小为 。相比之下，我们的方法在建模像素的相互依赖性时对非常小的 patch（4×4）执行注意力机制，每个 patch 的注意力矩阵仅为 ，总的注意力矩阵大小为 次操作。这种减少使我们的方法在最精细分辨率下的计算效率提高了 4000 倍，从而首次实现了逐像素建模高分辨率图像的能力。
与长序列建模的关系
大多数先前关于像素级生成的工作将问题表述为长序列建模，并利用语言建模方法来解决它，比如（Child et al., 2019; Roy et al., 2021; Ren et al., 2021; Hawthorne et al., 2022; Yu et al., 2023）。然而，许多数据类型（包括但不限于图像）的内在结构超越了一维序列。与这些方法不同，我们将此类数据视为由多个元素组成的集合（而非序列），并采用分而治之的策略递归建模具有较少元素的较小子集。这种方法的动机是观察到许多此类数据表现出近分形结构：图像由子图像组成，分子由子分子组成，生物神经网络由子网络组成。因此，设计用于处理此类数据的生成模型应由本身就是生成模型的子模块组成。
实现
我们简要描述了分形图像生成框架的训练和生成过程。更多细节和超参数可以在附录 A 中找到。
训练：通过以广度优先的方式遍历分形架构，在原始图像像素上端到端地训练分形模型。在训练过程中，每个自回归模型接收来自前一个自回归模型的输入，并为下一级自回归模型生成一组输出作为输入。这一过程一直持续到最后一级，其中图像被表示为像素序列。最终的自回归模型使用每个像素的输出以自回归方式预测 RGB 通道。在预测的 logits 上计算交叉熵损失（将 RGB 值视为 0 到 255 的离散整数），并通过所有级别的自回归模型反向传播此损失，从而端到端地训练整个分形框架。
生成：我们的分形模型以逐像素的方式生成图像，按照深度优先顺序遍历分形架构，如图2所示。这里我们以 MAR（Li et al., 2024）的随机顺序生成方案为例。第一级自回归模型捕捉 16×16 图像 patch 之间的相互依赖性，并在每一步基于已知 patch 生成下一级的输出。第二级模型随后利用这些输出来建模每个 16×16 patch 内 4×4 patch 之间的相互依赖性。类似地，第三级自回归模型建模每个 4×4 patch 内单个像素之间的相互依赖性。最后，最后一级自回归模型从自回归预测的 RGB logits 中采样实际的 RGB 值。
实验
在 ImageNet 数据集上进行了广泛的实验，分辨率为 64×64 和 256×256。评估包括无条件图像生成和类别条件图像生成，涵盖了模型的多个方面，如似然估计、保真度、多样性和生成质量。因此，报告了负对数似然（NLL）、Frechet Inception Distance（FID）、Inception Score（IS）、Precision 和 Recall以及可视化结果，以全面评估我们的分形框架。
似然估计
首先在无条件 ImageNet 64×64 生成任务上评估我们的方法，以检验其似然估计能力。为了检验分形框架的有效性，我们比较了不同分形级别下框架的似然估计性能，如表2所示。使用单一自回归模型建模整个 64×64×3=12,288 像素序列会导致极高的计算成本，使得训练不可行。此外，两级分形框架（首先建模整个像素序列，然后建模 RGB 通道）需要的计算量是我们三级分形模型的十倍以上。使用更多分形级别不仅计算效率更高，还提高了似然估计性能，这可能是因为它更好地捕捉了图像的内在层次结构。这些结果证明了我们分形框架的高效性和有效性。
进一步在表5中将我们的方法与其他基于似然的模型进行了比较。本文的分形生成模型（使用因果和掩码自回归分形生成器实例化）在似然性能上表现强劲。特别是，它实现了每维度 3.14 比特的负对数似然，显著优于之前最好的自回归模型（每维度 3.40 比特），并与基于扩散的先进方法保持竞争力。这些发现证明了我们的分形框架在像素级图像生成这一挑战性任务中的有效性，突显了其在建模高维非顺序数据分布方面的潜力。
生成质量
在类别条件图像生成任务上评估了 FractalMAR，分辨率为 256×256，使用了四个分形级别。按照标准实践报告了 FID、Inception Score、Precision 和 Recall 等标准指标，以评估其生成质量，如表4所示。具体而言，FractalMAR-H 实现了 6.15 的 FID 和 348.9 的 Inception Score，平均吞吐量为每张图像 1.29 秒（在单个 Nvidia H100 PCIe GPU 上以批量大小 1,024 进行评估）。
值得注意的是，本文方法在 Inception Score 和 Precision 上表现强劲，表明其能够生成具有高保真度和精细细节的图像，如图4所示。然而，其 FID 和 Recall 相对较弱，表明生成的样本多样性较低。推测这是由于逐像素建模近 200,000 个像素的巨大挑战所致。尽管如此，这些结果突显了我们的方法不仅在准确似然估计方面有效，还能生成高质量图像。
进一步观察到了一个有希望的扩展趋势：将模型大小从 186M 参数增加到 848M 参数，显著改善了 FID（从 11.80 降至 6.15）和 Recall（从 0.29 提升至 0.46）。预计进一步增加参数规模可以进一步缩小 FID 和 Recall 的差距。与依赖分词器的模型不同，我们的方法不受分词引入的重建误差影响，表明在更大模型容量下可能实现无上限的性能提升。
条件逐像素预测
进一步通过图像编辑中的常规任务检验了我们方法在条件逐像素预测方面的性能。图5提供了几个示例，包括修复、外推、去裁剪和类别条件编辑。如图所示，本文方法能够基于未掩码区域准确预测被掩码的像素。此外，它还能有效地从类别标签中捕捉高级语义，并在预测的像素中反映出来。这在类别条件编辑示例中得到了展示，模型通过以狗的类别标签为条件，将猫的脸替换为狗的脸。这些结果证明了我们的方法在给定已知条件下预测未知数据方面的有效性。
更广泛地说，通过逐元素生成数据，本文方法提供了一种比扩散模型或在潜在空间中操作的生成模型更易于人类理解的生成过程。这种可解释的生成过程不仅使我们能够更好地理解数据是如何生成的，还提供了一种控制和与生成过程交互的方式。这种能力在视觉内容创作、建筑设计和药物发现等应用中尤为重要。本文有竞争力的结果突显了我们方法在可控和交互式生成方面的潜力，为未来在这一方向上的探索铺平了道路。
讨论与结论
本文提出的分形生成模型的有效性通过像素级生成这一挑战性任务得到了验证，这为设计生成模型提供了新的机会。它突显了将复杂数据分布分解为可管理的子问题，并通过将现有生成模型抽象为模块化单元来解决这些问题的潜力。我们相信，分形生成模型特别适合建模具有超越一维顺序的内在结构的数据。希望本方法的简单性和有效性能够激发研究社区探索分形生成模型的新设计和应用。
参考文献
[1] Fractal Generative Models
