Fractal Generative Models论文阅读笔记与代码分析

何恺明分型模型这篇文章在二月底上传到arXiv预出版网站到现在已经过了三个月，当时我也听说这篇文章时感觉是大有可为，但是几个月不知道忙啥了，可能错过很多机会，但是亡羊补牢嘛，而且截至目前，该文章应该也还是生成模型领域最重大的理论突破。

原文阅读吃力的也有将该文章翻译成中文的博客链接可供参考，原文及代码链接均在该博客中。

论文笔记

该论文首先展示了一张分型模型的示意图，并说明在这个实例中，我们使用自回归模型作为分形生成器。通过在自回归模型中递归调用自回归模型，我们构建了一个在不同级别上具有自相似性的分形状框架。
分型模型示意图
对该部分有两点要补充：
自回归模型：：Autoregressive model，简称AR模型，是一种通过历史数据预测未来值的统计模型，其核心思想是当前值与过去若干时刻的值存在线性关系。
工程上可以理解为，该模型训练好后，可根据已有序列不断生成后续序列，甚至只要有一个 $x_0$ 就可以生成任意长度的序列。

分形？递归？：从图中来看，模型结构更像是递归生成。最上层的自回归模型调用下一层自回归模型，不断深入直到像素级生成，在该过程同时使用分治并行计算，提高效率。

摘要与引言

作者首先提出模块化是计算机科学的基石，它将复杂的功能抽象为原子构建块故本文也将生成模型抽象为原子生成模块，引入了一个新的模块化水平。我们的方法通过递归调用原子生成模块来构建一种新型的生成模型，从而产生自相似的分形架构，我们称之为分形生成模型。使用自回归作为原子生成模块，在似然估计和生成质量上都取得了不错的效果，并且认为可能在生成领域打开新范式。

作者的灵感来自自然界中的图像，如雪花、树枝和DNA等生物结构，这些结构通常表现出分形或近分形结构，这些图像由子图像构成，而子图像本身也是图像，重点在于用何种规则生成子图像。

故论文设计的关键也是定义递归生成规则的生成器，在该基础上进一步实例化为逐像素生成器。目前直接对像素序列进行建模的方法效果都不太好，但逐像素建模仍然对于分子结构，蛋白质和生物神经网络十分重要，作者提出的分形框架在这项任务表现出色。

随后作者介绍了分形、分层和分层生成模型，比如latent diffusion通过将图像压缩到潜空间，也是分层生成理念的应用，还有MegaByte和尺度空间自回归方法等。

模块化单元在神经架构的首次应用是FractalNet，它通过递归调用简单的扩展规则来构建非常深入的神经网络。该方法虽然思想类似，但有两个关键方面的方法不同。
首先，FractalNet使用一小块卷积层作为其模块化单元，而我们使用整个生成模型，代表不同级别的模块化。
其次，FractalNet主要是为分类任务设计的，因此只输出低维logits。
相比之下，我们的方法利用分形图案的指数缩放行为来生成大量输出（例如，数百万个图像像素），展示了分形启发设计在分类之外更复杂的任务中的潜力。

分型模型

分形模型的关键是通过简单模型构造高级的生成模型，但实现上重点在于两部分：
1，如何定义分形规则
2，如何组合分形模型的结果

分形过程用数学语言描述为： $x_{i+1}=g_i(x_i)$ ，其中 $g_i$ 分形生成器，因为生成器可以从单个输入中同时生成多个输出，所以只需要线性时间就能生成指数结果，生成效率天然较高。

模型设计

作者使用原子生成模块作为参数分形生成器，通过这种方式，神经网络直接从数据中“学习”递归规则，在执行过程中采用分治策略。

其设计结构如下：
Fractal生成模型结构
每个自回归模型的输入来自上一层生成器，并将图像拆分为多个片patches输入Transformer，随后Transformer将计算结果放在patches前，一起给下一层生成器，使得生成器能获得已有像素及序列内部信息。
随后更下一级生成器对patches进行建模，进一步细分为更小的patches并递归进行上述步骤。在该过程中生成模型与Transformer不断减小规模，最后一层使用轻量Transformer对像素的RGB通道进行自回归建模，并在预测中应用 256位cross-entropy交叉熵损失。

时间复杂度分析

使用该方法生成256x256的图像成本仅为64x64的两倍，各级生成器计算成本如下：
Fractal生成器计算成本
论文设计了两种变体，类似GPT的因果单向顺序建模Transformer（AR）和类似BERT的双向随机建模Transformer（MAR），二者都遵循预测自回归原则，示意图如下：
AR与MAR
AR使用键值缓存加速生成，MAR使用双向注意力支持并行预测的同时还能使图像更一致，总得来说还是MAR性能更好一些。

与现有尺度模型预测的区别主要就是在于快。
现有方法使用单个回归模型预测，该方法采用分而治之的策略，使用生成子模块递归建模原始像素。
另外一个区别就是计算复杂度，我们前面学习的Transformer中提到，越往后计算需要关注的序列就越多，以生成一张256x256的图像为例，注意力矩阵维度为 $(256\times256)^2=4,294,967,296$ ，而我们的方法因为只对很小的一个区域关注如(4x4)，故计算复杂度仅为 $(4\times4)^2=256$ ，总的注意力矩阵大小为 $(64\times64)\times(4\times4)^2=1,048,576$ ，在该计算下可以认为计算效率提升了四千倍。

实施细节

训练
广度优先的方式遍历分形架构，在原始图像像素上端到端地训练分形模型，训练期间，每个自回归模型都从前一个自回归模型中获取一个输入，并生成一组输出，供下一级自回归模型作为输入，该过程递归进行直到最终像素级别，最终的自回归模型预测出RGB通道上的值。
计算交叉熵损失并反馈到所有级别的自回归模型，实现对整个模型的训练。

作者在训练中发现使用“引导像素”对训练更有帮助，具体实施为：模型使用前一层的输出来预测当前输入的平均像素值，将该值用于Transformer转换的附加条件。

因为该过程是分块递归进行的，彼此间独立可能导致patches块之间的边界不一致。作者在训练中不仅提供了当前 patch 的输出，还提供了周围四个 patchs 的输出。加入额外输入会提高序列长度，但也能显著减少边界伪影。

默认情况下模型训练使用AdamW优化器，训练了800轮，AdamW 的权重衰减和动量分别为0.05和（0.9，0.95）。对ImageNet 64×64 使用批处理大小2048，对 ImageNet 256×256 使用批处理大小1024，并将基本学习率（lr）设置为5e-5。

生成
生成过程逐像素，采用深度优先方法，以MAR随机顺序生成为例，第一级自回归捕获16x16图像块patches之间的关系，每一步根据已有patches生成下一步输出。
第二级根据这16x16patches中4x4的图像块依赖进行建模，第三级再分为4x4，最后一级从自回归预测的RGB中采样获得预测值。

无分类器指导CFG生成的训练过程中，10%的样本类条件被替换为虚拟类标记。推理时，模型同时使用给定的类标记和虚拟标记运行，为每个像素通道生成两组逻辑 $l_c$ 和 $l_u$ ，根据公式 $l=l_u+w(l_c-l_u)$ 调整预测逻辑，第一级线性推理中采用线性CFG。
当预测概率较小时CFG可能出现数值不稳定性，为了缓解此问题，我们在应用 CFG 之前对条件 logit 应用阈值为 0.0001 的 top-p 采样。

实验

该部分主要通过图像和表格来展示，图像可以从人眼主观感受鉴别，但是需要一些指标对结果进行定量分析，所以首先介绍几个指标。
FID
FID是生成图像评价指标，通过将生成图像的特征与真实图像特征相比判断真实图像的可能性。提取特征的过程通过其预训练的深度神经网络（如Inception V3）实现，比较这些特征在高维空间中的分布差异来评估生成质量，该值通常越小越好。
有关该指标的具体计算过程可见深入解析FID。
IS
nception Score是在GAN网络背景下提出，判别器只能鉴别图像是否为生成图像，但对图像的多样性和类别清晰质量上缺少关注，故提出了IS指标，衡量图像的多样性和生成质量。
该值越大越好，使用时多和FID结合以弥补多样性不足。
Precision and Recall
pre精度和rec召回率，精度表示预测为正确的元素真为正确的比例，召回率表示样本中正确的元素被正确预测的比例，详细解释可见召回率和精度。

逐像素生成的部分结果展示如下，256x256的图像，作者说可以实现平均1.29秒生成一张，比起diffusion的速度还是快了不少，但是人家用的H100，复现之前对效率暂时不做评价。

逐像素生成结果
量化标准的评价参数如下：
Fractal量化参数
从表里看出大部分指标其实是不占优的，现在唯一的优势就是速度快，这还需要后续复现实验检验。
并且作者认为，随着参数增大效果是可能更好甚至无上限的。

另外还有裁剪修复及语义替换效果展示如下：

对于鸡头的修复甚至比原图还好，就是最后换的狗头有点违和。

小结

作者提出了一种分治递归解决图像生成的方法，主要通过训练多层预测器，最后一层预测RGB通道值并反向传播给各层级预测器实现。因为多个预测器可并行进行，时间和计算复杂度上性能都表现较好，为生成模型的基础研究提供了新方向。

我觉得同时也还有很多不足，可能是这篇文章暂时没发表的原因？
首先大部分量化指标比起其他生成模型都落后，提出新的生成模型如果效果不佳也不占优势。
其次生成虽然能实现并行，但目前主流的扩散生成也有潜空间优化策略，二者效率究竟如何还需比较。
最后逐像素生成相比扩散或GAN具有什么特殊优势吗？在处理具有自相似结构如雪花、树枝、DNA等可能训练时会比较轻松，但日常图像处理中这种逐像素生成的算法应用效果究竟如何还有待检验。

实验分析

复现代码第一步也是最重要的一步就是把代码跑起来，能运行了以后心里踏实才能继续深入研究，所以本文也先记录环境配置及使用情况，再分析具体代码结构。

环境配置

Readme文档中给出的方法是进入到代码目录后，直接使用conda env create -f environment.yaml创建虚拟环境，配置文件中需要的环境如下：

name: fractalgen
channels:
  - pytorch
  - defaults
  - nvidia
dependencies:
  - python=3.8.5
  - pip=20.3
  - pytorch-cuda=11.8
  - pytorch=2.2.2
  - torchvision=0.17.2
  - numpy=1.22
  - pip:
    - opencv-python==4.1.2.30
    - timm==0.9.12
    - tensorboard==2.10.0
    - scipy==1.9.1
    - gdown==5.2.0
    - -e git+https://github.com/LTH14/torch-fidelity.git@master#egg=torch-fidelity

直接使用命令要等个半天，最后到GitHub下载的文件可能因为网络问题下不下来报错中断，笔者使用配置文件安装后发现到pip以后的包都没装好，别的都很简单，pip install一下就没问题，主要是最后安装github项目没搞过。

-e参数表示以可编辑模式安装，后续的链接直接复制无法访问，因为其中有git的专用指令，@表示从master分支下载，最后#egg表示要安装的包。

执行操作为：首先复制urlhttps://github.com/LTH14/torch-fidelity.git到浏览器打开，下载代码的zip文件，移动到项目目录下解压，conda activate fractalgen进入虚拟环境，使用pip install -e .安装当前目录的包。
安装后可通过在python中使用如下命令检验：

import torch_fidelity
print(torch_fidelity.__version__)
# 出现如下输出说明安装成功
0.4.0-beta

使用该代码需要预训练模型，后续会通过代码下载，但是网络可能有问题，为了方便还是推荐手动下载，点击链接跳转网站自己下载即可，要注意的是所有模型名字都叫checkpoint，下载要自己标注好对应类型。

生成图像

这部分自己研究半天，后来才发现Readme里早就写了，只不过是通过链接给出的。
用法中点击Demo即可打开图像生成代码演示（该步需要魔法），示例给出的代码还是很详细的，从包的引入，模型下载，到生成图像都有。

引入包及设置代码如下，注释部分是源代码有的，但是尝试觉得没什么用的：

# import os
# !git clone https://github.com/LTH14/fractalgen.git
# os.chdir('fractalgen')
# os.environ['PYTHONPATH'] = '/env/python:/content/fractalgen'
# !pip install timm==0.9.12
import torch
import numpy as np
from models import fractalgen
from torchvision.utils import save_image
from util import download
from PIL import Image
# from IPython.display import display
# IPython的display专门用于Jupyter Notebook的相关展示，网站内还需要谷歌账号
# 没法注册只能本地跑，所以这个库暂时没用
torch.set_grad_enabled(False)
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    print("GPU not found. Using CPU instead.")

下载预训练模型64x64，网站中可以更改模型为mar，其实就是把model_type换一下：

model_type = "fractalmar_in64" #@param ["fractalar_in64", "fractalmar_in64"]
if model_type == "fractalar_in64":
  download.download_pretrained_fractalar_in64(overwrite=False)
  num_conds=1
elif model_type == "fractalmar_in64":
  download.download_pretrained_fractalmar_in64(overwrite=False)
  num_conds=5
else:
  raise NotImplementedError
model = fractalgen.__dict__[model_type](
    guiding_pixel=False,
    num_conds=num_conds,
).to(device)
state_dict = torch.load("pretrained_models/{}/checkpoint-last.pth".format(model_type))["model"]
model.load_state_dict(state_dict)
model.eval() # important!

使用该预训练模型生成图片代码如下：

# Set user inputs
seed = 0 #@param {type:"number"}
torch.manual_seed(seed)
np.random.seed(seed)
class_labels = 207, 360, 388, 113, 355, 980, 323, 979 #@param {type:"raw"}
num_iter_list = 64, 16 #@param {type:"raw"}
cfg_scale = 4 #@param {type:"slider", min:1, max:10, step:0.5}
cfg_schedule = "constant" #@param ["linear", "constant"]
temperature = 0.93 #@param {type:"slider", min:0.9, max:1.2, step:0.01}
filter_threshold = 1e-4
samples_per_row = 4 #@param {type:"number"}

label_gen = torch.Tensor(class_labels).long().cuda()
class_embedding = model.class_emb(label_gen)
if not cfg_scale == 1.0:
  class_embedding = torch.cat([class_embedding, model.fake_latent.repeat(label_gen.size(0), 1)], dim=0)

with torch.no_grad():
  with torch.cuda.amp.autocast():
    sampled_images = model.sample(
      cond_list=[class_embedding for _ in range(num_conds)],
      num_iter_list=num_iter_list,
      cfg=cfg_scale, cfg_schedule=cfg_schedule,
      temperature=temperature,
      filter_threshold=filter_threshold,
      fractal_level=0,
      visualize=True)

# Denormalize images.
pix_mean = torch.Tensor([0.485, 0.456, 0.406]).cuda().view(1, -1, 1, 1)
pix_std = torch.Tensor([0.229, 0.224, 0.225]).cuda().view(1, -1, 1, 1)
sampled_images = sampled_images * pix_std + pix_mean
sampled_images = torch.nn.functional.interpolate(sampled_images, scale_factor=4, mode="nearest")
sampled_images = sampled_images.detach().cpu()

# Save & display images
save_image(sampled_images, "samples.png", nrow=int(samples_per_row), normalize=True, value_range=(0, 1))
samples = Image.open("samples.png")

# from IPython.display import display, clear_output
# clear_output(wait=True)
# display(samples)
# 上述代码用于Notebook，修改为如下代码可实现本地展示
import matplotlib.pyplot as plt
plt.imshow(samples)
plt.show()

后续还有256的使用示例，大差不差，使用上述代码生成8张不同类型的64x64的图像展示如下：
分形模型生成图像
不算模型加载时间，单纯梯度内生成图像花了131s，这速度比起参数量更大的stable diffusion也慢了不少，何况sd生成的还是高分辨率640x640的图像。
只生成两张图像也花了65s，可见图像生成数量和时间并不是简单的线性关系，所以时间可能有压缩空间，还需进一步探究。

模型训练

文档还说要下载ImageNet数据集，可到官网下载，但不管下载官网中的哪个版本都会发现——数据集实在是太太太大了。
以2010的版本为例，训练，检验加测试集有一百多万张图像，大约一百五十G。 ImageNet数据集
尝试了一个更轻量级的数据集Tiny ImageNet-200，该数据集图像分辨率64x64，200类，每类500张图像，验证集测试集分别包含10000张图像，下载的压缩包就200多MB，比满血版的小太多了，而且结构也差不多，内部都分train、val和test类型。

后续使用文档命令即可，这部分需要相当的硬件条件支持，本文暂不研究。

代码解析

本文只关注生成过程，故只解析生成部分代码，其中前半部分都是用于模型加载和初始化，用于核心生成部分的代码如下：

with torch.no_grad():
  with torch.cuda.amp.autocast():
    sampled_images = model.sample( # 核心函数sample
      cond_list=[class_embedding for _ in range(num_conds)],
      num_iter_list=num_iter_list,
      cfg=cfg_scale, cfg_schedule=cfg_schedule,
      temperature=temperature,
      filter_threshold=filter_threshold,
      fractal_level=0,
      visualize=True)

在torch环境下仅调用了一个方法，就是模型的sample，该负责在无梯度和混合精度模式下，通过分类引导扩散模型生成指定类别的高质量图像，并支持多层级迭代优化，相关参数说明如下：

cond_list: 用于生成图像的条件向量列表，这里将 class_embedding 复制 num_conds 次作为多个层次的条件输入。
num_iter_list: 控制不同层次的迭代次数，例如 (64, 16) 表示粗粒度和细粒度层次的迭代次数。
cfg: 分类引导系数，值越大生成结果越贴近指定类别标签。
cfg_schedule: 分类引导调度策略，可选 "linear" 或 "constant"。
temperature: 控制采样的随机性，值越高随机性越大。
filter_threshold: 噪声过滤阈值，用于控制生成图像的质量。
fractal_level: 分形层次级别，0 表示最低层级（整体结构生成）。
visualize: 是否在生成过程中可视化中间结果。

以在Mar中的Sample为例，其功能实现步骤主要分为以下几步：
1，初始化batch_size，根据是否使用cfg判断赋值；

if cfg == 1.0:
  bsz = cond_list[0].size(0)
else:
  bsz = cond_list[0].size(0) // 2

2，生成引导像素，根据文中描述，来自上一层输出的平均值；

# sample the guiding pixel
if self.guiding_pixel:
   sampled_pixels = self.guiding_pixel_loss.sample(cond_list, temperature, 
   cfg, filter_threshold)
if not cfg == 1.0:
   sampled_pixels = torch.cat([sampled_pixels, sampled_pixels], dim=0)
cond_list.append(sampled_pixels)

3，初始化掩码；

# init token mask
mask = torch.ones(bsz, self.seq_len).cuda()
patches = torch.zeros(bsz, self.seq_len, 3 * self.patch_size**2).cuda()
orders = self.sample_orders(bsz)
num_iter = min(self.seq_len, num_iter)

4，迭代生成图像，调用predict获取下一层条件，更新掩码，next_level_sample_function计算新的patches，该部分代码较长，需要单独解析。
5，组合所有patches序列为图像。

图像生成

生成图像sample

sample的核心步骤是在迭代逐步预测和填充新的patches，各部分含义直接通过注释给出：

for step in range(num_iter):
		# 克隆当前patches集合，防止修改原始数据
        cur_patches = patches.clone()
		# 处理分类器自由引导，如果不是设置为1，
		# 则 patches 和 mask 进行复制并拼接，用于同时运行有引导和无引导的生成样本。
        if not cfg == 1.0:
            patches = torch.cat([patches, patches], dim=0)
            mask = torch.cat([mask, mask], dim=0)

        # get next level conditions
        # 调用 predict 方法，根据当前的 patches、mask 和 cond_list（条件列表）
        # 生成下一级的条件信息，用于指导后续生成
        cond_list_next = self.predict(patches, mask, cond_list)

        # mask ratio for the next round, following MAR.
        # 根据余弦衰减策略计算当前步的掩码比例 mask_ratio，
        # 然后计算需要掩盖的 token 数量 mask_len
        mask_ratio = np.cos(math.pi / 2. * (step + 1) / num_iter)
        mask_len = torch.Tensor([np.floor(self.seq_len * mask_ratio)]).cuda()

        # masks out at least one for the next iteration
        # 确保下一步要掩盖的 token 数量不会超过当前剩余未掩盖的数量，并且至少保留一个 token。
        mask_len = torch.maximum(torch.Tensor([1]).cuda(),
                                 torch.minimum(torch.sum(mask, dim=-1, keepdims=True) - 1, mask_len))

        # get masking for next iteration and locations to be predicted in this iteration
        # 生成新的掩码
        # 根据 orders 和 mask_len 构建新的掩码 mask_next。
		# 计算当前步骤中需要预测的区域 mask_to_pred，即在当前掩码与新掩码之间发生变化的位置。
		# 更新 mask 为 mask_next，进入下一步迭代。
        mask_next = mask_by_order(mask_len[0], orders, bsz, self.seq_len)
        if step >= num_iter - 1:
            mask_to_pred = mask[:bsz].bool()
        else:
            mask_to_pred = torch.logical_xor(mask[:bsz].bool(), mask_next.bool())
        mask = mask_next
        # 调整预测掩码，将 mask_to_pred 扩展以匹配复制后的 patches 和 mask
        if not cfg == 1.0:
            mask_to_pred = torch.cat([mask_to_pred, mask_to_pred], dim=0)

        # sample token latents for this step
        # 提取 cond_list_next 中与 mask_to_pred 对应的部分
        for cond_idx in range(len(cond_list_next)):
            cond_list_next[cond_idx] = cond_list_next[cond_idx][mask_to_pred.nonzero(as_tuple=True)]

        # cfg schedule
        #  动态调整 CFG 参数，
        if cfg_schedule == "linear":
            cfg_iter = 1 + (cfg - 1) * (self.seq_len - mask_len[0]) / self.seq_len
        else:
            cfg_iter = cfg
        # 使用外部提供的 next_level_sample_function 函数，根据当前条件生成新的 patch
        sampled_patches = next_level_sample_function(cond_list=cond_list_next, cfg=cfg_iter,
                                                     temperature=temperature, filter_threshold=filter_threshold)
        sampled_patches = sampled_patches.reshape(sampled_patches.size(0), -1)

        if not cfg == 1.0:
            mask_to_pred, _ = mask_to_pred.chunk(2, dim=0)
		# 将生成的新 patch 插入到 cur_patches 中对应位置
		# 更新 patches 为最新的 cur_patches
        cur_patches[mask_to_pred.nonzero(as_tuple=True)] = sampled_patches.to(cur_patches.dtype)
        patches = cur_patches.clone()

条件信息生成predict

该函数用以生成下一级条件信息，以列表形式返回给cond_list_next，解析同样通过注释给出如下：

def predict(self, x, mask, cond_list):
	# 将每个 patch 映射到嵌入空间中，得到 x 的形状为 (bsz, seq_len, embed_dim)
    x = self.patch_emb(x)
    
    # prepend conditions from prev generator
    # 将外部传入的多个条件（如前一级生成器的输出）映射后拼接到序列前面。
	# 形状变为 (bsz, num_conds + seq_len, embed_dim)
    for i in range(self.num_conds):
        x = torch.cat([self.cond_emb(cond_list[i]).unsqueeze(1), x], dim=1)

    # prepend guiding pixel
    # 如果启用了 guiding_pixel，将最后一个条件作为“全局颜色”特征插入序列最前面。
	# 形状变为 (bsz, num_conds + guiding_pixel + seq_len, embed_dim)
    if self.guiding_pixel:
        x = torch.cat([self.pix_proj(cond_list[-1]).unsqueeze(1), x], dim=1)

    # masking
    # 构建完整的掩码（包括条件部分），并用 mask_token 替换被遮挡的 patch。
    mask_with_cond = torch.cat([torch.zeros(x.size(0), self.num_conds+self.guiding_pixel, device=x.device), mask], dim=1).bool()
    x = torch.where(mask_with_cond.unsqueeze(-1), self.mask_token.to(x.dtype), x)

    # position embedding
    # 添加可学习的位置编码,对输入进行归一化处理。
    x = x + self.pos_embed_learned
    x = self.patch_emb_ln(x)

    # apply Transformer blocks
    # 通过多个 Transformer Block 处理输入序列，最终输出经过 LayerNorm。
    if self.grad_checkpointing and not torch.jit.is_scripting() and self.training:
        for block in self.blocks:
            x = checkpoint(block, x)
    else:
        for block in self.blocks:
            x = block(x)
    x = self.norm(x)

    # return 5 conditions: middle, top, right, bottom, left
    # 提取中间结果并构建 5 种上下文条件，分别代表上下左右和当前patch的特征
    middle_cond = x[:, self.num_conds+self.guiding_pixel:]
    bsz, seq_len, c = middle_cond.size()
    h = int(np.sqrt(seq_len))
    w = int(np.sqrt(seq_len))
    top_cond = middle_cond.reshape(bsz, h, w, c)
    top_cond = torch.cat([torch.zeros(bsz, 1, w, c, device=top_cond.device), top_cond[:, :-1]], dim=1)
    top_cond = top_cond.reshape(bsz, seq_len, c)

    right_cond = middle_cond.reshape(bsz, h, w, c)
    right_cond = torch.cat([right_cond[:, :, 1:], torch.zeros(bsz, h, 1, c, device=right_cond.device)], dim=2)
    right_cond = right_cond.reshape(bsz, seq_len, c)

    bottom_cond = middle_cond.reshape(bsz, h, w, c)
    bottom_cond = torch.cat([bottom_cond[:, 1:], torch.zeros(bsz, 1, w, c, device=bottom_cond.device)], dim=1)
    bottom_cond = bottom_cond.reshape(bsz, seq_len, c)

    left_cond = middle_cond.reshape(bsz, h, w, c)
    left_cond = torch.cat([torch.zeros(bsz, h, 1, c, device=left_cond.device), left_cond[:, :, :-1]], dim=2)
    left_cond = left_cond.reshape(bsz, seq_len, c)

    return [middle_cond, top_cond, right_cond, bottom_cond, left_cond]

生成新patch的next_level

if fractal_level < self.num_fractal_levels - 2:
	# 不是最后一层，使用functools.partial固定参数，避免递归调用重复传参，
    next_level_sample_function = partial(
        self.next_fractal.sample,
        num_iter_list=num_iter_list,
        cfg_schedule="constant",
        fractal_level=fractal_level + 1
    )
else:
	# 是倒数第二层或更深层级，调用自身
    next_level_sample_function = self.next_fractal.sample

# Recursively sample using the current generator.
# 当前生成器的sample方法
return self.generator.sample(
    cond_list, num_iter_list[fractal_level], cfg, cfg_schedule,
    temperature, filter_threshold, next_level_sample_function, visualize
)

next_level的预测是一个递归结构，直到遇到最后一层才调用PixelLoss方法，最终层的生成方法如下。

递归底层生成RGB

def sample(self, cond_list, temperature, cfg, filter_threshold=0):
    """ generation """
    if cfg == 1.0:
        bsz = cond_list[0].size(0)
    else:
        bsz = cond_list[0].size(0) // 2
    pixel_values = torch.zeros(bsz, 3).cuda()

    for i in range(3):
        if cfg == 1.0:
            logits, _ = self.predict(pixel_values, cond_list)
        else:
            logits, _ = self.predict(torch.cat([pixel_values, pixel_values], dim=0), cond_list)
        logits = logits[:, i]
        logits = logits * temperature

        if not cfg == 1.0:
            cond_logits = logits[:bsz]
            uncond_logits = logits[bsz:]

            # very unlikely conditional logits will be suppressed
            cond_probs = torch.softmax(cond_logits, dim=-1)
            mask = cond_probs < filter_threshold
            uncond_logits[mask] = torch.max(
                uncond_logits,
                cond_logits - torch.max(cond_logits, dim=-1, keepdim=True)[0] + torch.max(uncond_logits, dim=-1, keepdim=True)[0]
            )[mask]

            logits = uncond_logits + cfg * (cond_logits - uncond_logits)

        # get token prediction
        probs = torch.softmax(logits, dim=-1)
        sampled_ids = torch.multinomial(probs, num_samples=1).reshape(-1)
        pixel_values[:, i] = (sampled_ids.float() / 255 - self.pix_mean[i]) / self.pix_std[i]

    # back to [0, 1]
    return pixel_values