周弈帆的博客

最近我在看位置编码最新技术时，看到了一个叫做 “NTK-aware” 的词。我想：「”NTK”是什么？Next ToKen （下一个词元）吗？为什么要用这么时髦的缩写？」看着看着，我才发现不对劲。原来，NTK 是神经网络理论里的一个概念，它从 kernel regression 的角度解释了神经网络的学习方法。基于 NTK 理论，有人解释了位置编码的理论原理并将其归纳为一种特殊的 Fourier Feature （傅里叶特征）。这么多专有名词一下就把我绕晕了，我花了几天才把它们之间的关系搞懂。

在这篇文章里，我主要基于论文 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains （后文简称为「傅里叶特征论文」），介绍傅里叶特征这一概念。为了讲清这些理论的发展脉络，我会稍微讲一下 NTK 等理论概念。介绍完傅里叶特征后，我还会讲解它在其他方法中的应用。希望读完本文后，读者能够以这篇论文为基点，建立一个有关位置编码原理的知识网络，以从更深的层次来思考新的科研方向。

用 MLP 表示连续数据

我们先从一个具体的任务入手，直观体会傅里叶特征能够做些什么事。

我们知道，神经网络，哪怕是最简单的多层感知机（MLP），都有着很强的泛化能力：训练完毕后，对于训练集里完全没见过的输入，网络也能给出很正确的输出。特别地，如果新输入恰好和训练集的某个输入很近，那么它的输出也会和对应的训练集输出很近；随着新输出与训练集输入的距离不断增加，新输出也会逐渐变得不同。这反映了神经网络的连续性：如果输入的变化是连续的，那么输出的变化也是连续的。

基于神经网络的这一特性，有人想到：我们能不能用神经网络来表示连续数据呢？比如我想表达一张处处连续的图像，于是我令神经网络的输入是 (x, y) 表示的二维坐标，输出是 RGB 颜色。之后，我在单张图像上过拟合这个 MLP。这样，学会表示这张图像后，哪怕输入坐标是分数而不是整数，神经网络也能给出一个颜色输出。

这种连续数据有什么好处呢？我们知道，计算机都是以离散的形式来存储数据的。比如，我们会把图像拆成一个个像素，每个像素存在一块内存里。对于图像这种二维数据，计算机的存储空间还勉强够用。而如果想用密集的离散数据表达更复杂的数据，比如 3D 物体，计算机的容量就捉襟见肘了。但如果用一个 MLP 来表达 3D 物体的话，我们只需要存储 MLP 的参数，就能获取 3D 物体在任何位置的信息了。

这就是经典工作神经辐射场 (Neural Radiance Field, NeRF) 的设计初衷。NeRF 用一个 MLP 拟合 3D 物体的属性，其输入输出如下图所示。我们可以用 MLP 学习每个 3D 坐标的每个 2D 视角处的属性（这篇文章用的属性是颜色和密度）。根据这些信息，利用某些渲染算法，我们就能重建完整的 3D 物体。

上述过程看起来好像很简单直接。但在 NeRF 中，有一个重要的实现细节：必须给输入加上位置编码，MLP 才能很好地过拟合连续数据。这是为什么呢？让我们先用实验复现一下这个现象。

MLP 拟合连续图像实验

为了快速复现和位置编码相关的问题，我们简单地用一个 MLP 来表示图像：MLP 的输入是 2D 坐标，输出是此处的三通道 RGB 颜色。我为这篇博文创建一个 GitHub 文件夹 https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/FourierFeature ，该实验的 Notebook 代码在文件夹的 image_mlp.ipynb 中，欢迎大家 clone 项目并动手尝试。

一开始，我们先导入库并可视化要拟合的图片。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.io import read_image, ImageReadMode
from torchvision.transforms.functional import to_pil_image

from tqdm import tqdm
from einops import rearrange

def viz_image(pt_img: torch.Tensor):
    pil_img = to_pil_image(pt_img)
    display(pil_img)

    
input_image = read_image('misuzu.png', ImageReadMode.RGB)
input_image = input_image.to(torch.float32) / 255
input_image = input_image.unsqueeze(0)
input_image = F.interpolate(input_image, (256, 256), mode='bilinear')
viz_image(input_image[0])

我们再定义一个 MLP 类。稍后我们会并行地传入二维坐标。具体来说，我们会将输入定义为一个 [1, 2, H, W] 形状的数据，其中通道数 2 表示 (i, j) 格式的坐标。由于输入是以图像的形式并行输入的，我们可以用 $1 \times 1$ 的 2D 卷积来表示二维数据上的并行 MLP。所以在下面这个 MLP 里，我们只用到 $1 \times 1$ 卷积、激活函数、归一化三种层。按照傅里叶特征论文的官方示例，网络最后要用一个 Sigmoid 激活函数调整输出的范围。

1
2
3
4
5
6
7
8
9
10
11
12

class MLP(nn.Module):
    def __init__(self, in_c, out_c=3, hiden_states=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Conv2d(in_c, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),
            nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),
            nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),
            nn.Conv2d(hiden_states, out_c, 1), nn.Sigmoid()
        )

    def forward(self, x):
        return self.mlp(x)

之后我们来定义训练数据。在一般的任务中，输入输出都是从训练集获取的。而在这个任务中，输入是二维坐标，输出是图像的颜色值。输出图像 input_image 我们刚刚已经读取完毕了，现在只需要构建输入坐标即可。我们可以用下面的代码构建一个 [1, 2, H, W] 形状的二维网格，grid[0, :, i, j] 处的数据是其坐标 (i, j) 本身。当然，由于神经网络的输入一般要做归一化，所以我们会把原本 0~H 和 0~W 里的高宽坐标缩放都到 0~1。最终 grid[0, :, i, j]==(i/H, j/W)。

1
2
3
4
5

H, W = input_image.shape[2:]

h_coord = torch.linspace(0, 1, H)
w_coord = torch.linspace(0, 1, W)
grid = torch.stack(torch.meshgrid([h_coord, w_coord]), -1).permute(2, 0, 1).unsqueeze(0)

准备好一切后，我们就可以开始训练了。我们初始化模型 model 和优化器 optimizer，和往常一样训练这个 MLP。如前所述，这个任务的输入输出非常直接，输入就是坐标网格 grid，目标输出就是图片 input_image。每训练一段时间，我们就把当前 MLP 拟合出的图片和误差打印出来。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = MLP(2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
input_image = input_image.to(device)
grid = grid.to(device)
for epoch in tqdm(range(n_loops)):
    output = model(grid)
    loss = F.l1_loss(output, input_image)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 100 == 0 or epoch == n_loops - 1:
        viz_image(output[0])
        print(loss.item())

运行代码，大致能得到如下输出。可以看到，从一开始，图像就非常模糊。

不过，如果我们在把坐标输入进网络前先将其转换成位置编码——一种特殊的傅里叶特征，那么 MLP 就能清晰地拟合出原图片。这里我们暂时不去关注这段代码的实现细节。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

class FourierFeature(nn.Module):
    def __init__(self, in_c, out_c, scale):
        super().__init__()
        fourier_basis = torch.randn(in_c, out_c // 2) * scale
        self.register_buffer('_fourier_basis', fourier_basis)
        
    def forward(self, x):
        N, C, H, W = x.shape
        x = rearrange(x, 'n c h w -> (n h w) c')
        x = x @ self._fourier_basis
        x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)
            
        x = 2 * torch.pi * x
        x = torch.cat([torch.sin(x), torch.cos(x)], dim=1) 
        return x
        
feature_length = 256
model = MLP(feature_length).to(device)
fourier_feature = FourierFeature(2, feature_length, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
for epoch in tqdm(range(n_loops)):
    x = fourier_feature(grid)
    output = model(x)
    loss = F.l1_loss(output, input_image)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 100 == 0 or epoch == n_loops - 1:
        viz_image(output[0])
        print(loss.item())
prev_output = output

简单地对比一下，此前方法的主要问题是 MLP 无法拟合高频的信息（如图块边缘），只能生成模糊的图像。而使用位置编码后，MLP 从一开始就能较好地表示高频信息。可见，问题的关键在于如何让 MLP 更好地拟合数据的高频信息。

接下来，我们来从一个比较偏理论的角度看一看论文是怎么分析位置编码在拟合高频信息中的作用的。

核回归

傅里叶特征论文使用了神经正切核（Nerual Tangent Kernel, NTK）来分析 MLP 的学习规律，而 NTK 又是一种特殊的核回归 (Kernel Regression) 方法。在这一节里，我会通过代码来较为仔细地介绍核回归。下一节我会简单介绍 NTK。

和神经网络类似，核回归也是一种数学模型。给定训练集里的输入和输出，我们建立这样一个模型，用来拟合训练集表示的未知函数。相比之下，核回归的形式更加简单，我们有更多的数学工具来分析其性质。

核回归的设计思想来源于我们对于待拟合函数性质的观察：正如我们在前文的分析一样，要用模型拟合一个函数时，该模型在训练数据附近最好是连续变化的。离训练集输入越近，输出就要和其对应输出越近。基于这种想法，核回归直接利用和所有数据的相似度来建立模型：假设训练数据为 $(x_i, y_i), i \in [1, n]$，我们定义了一个计算两个输入相似度指标 $K(x_1, x_2)$，那么任意输入 $x$ 的输出为：

也就是说，对于一个新输入 $x$，我们算它和所有输入 $x_i$ 的相似度 $w_i$，并把相似度归一化。最后的输出 $f(x)$ 是现有 $y_i$ 的相似度加权和。

这样看来，只要有了相似度指标，最终模型的形式也就决定下来了。我们把这个相似度指标称为「核」。至于为什么要把它叫做核，是因为这个相似度指标必须满足一些性质，比如非负、对称。但我们这里不用管那么多，只需要知道核是一种衡量距离的指标，决定了核就决定了核回归的形式。

我们来通过一个简单的一维函数拟合实验来进一步熟悉核回归。该实验代码在项目文件夹下的 kernel_regression.ipynb 中。

先导入库。

1
2
3
4

%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

再创建一个简单的非线性函数，做为我们的拟合目标。这个函数就是一个简单的周期为 $2$ 的正弦函数乘上线性函数 $(1-x)$。我们可以简单可视化一下函数在 $[-1, 1]$ 之间的图像。

1
2
3
4
5
6
7

def func(x):
    return np.sin(np.pi * x) * (1 - x)

xs = np.linspace(-1, 1, 100)
ys = func(xs)
plt.plot(xs, ys)
plt.show()

基于这个函数，我们等间距地选一些点做为训练数据。

1
2
3
4

sample_x = np.linspace(-1, 1, 10)
sample_y = func(sample_x)
plt.scatter(sample_x, sample_y)
plt.show()

有了数据后，我们来用核回归根据数据拟合这个函数。在决定核回归时，最重要的是决定核的形式。这里我们用正态分布的概率密度函数来表示核，该核唯一的超参数是标准差，需要我们根据拟合结果手动调整。标准差为 1 的标准正态分布核的图像如下所示。由于最后要做归一化，正态分布密度函数的系数被省略掉了。

1
2
3
4
5
6
7

def kernel_func(x_ref, x_input, sigma=1):
    return np.exp(-(x_input-x_ref)**2 / (2 * sigma**2))

xs = np.linspace(-1, 1, 100)
ys = kernel_func(0, xs)
plt.plot(xs, ys)
plt.show()

可以从图像中看出，离某输入越近（假设该输入是 0），那么相似度就越高。这符合我们对于相似度函数的要求。

有了核函数后，我们就直接得到了模型。根据核回归模型计算结果的函数为 kernel_regression。函数参数 xs, ys 表示训练数据，x_input 表示测试时用的输入坐标，sigma 是核回归的超参数。

假设有 n 个训练样本，有 m 个测试输入，那么我们要计算每个测试输入对每个训练输入的 n * m 个相似度，这些相似度会存到矩阵 weight 里。为此，我们需要对 xs 和 x_input 做一些形状变换，再用上面定义的核函数 kernel_func 求出每对相似度。有了相似度后，我们根据公式计算点乘结果 weight_dot 及归一化系数 weight_sum，并最终计算出核回归的结果 res。

基于这个函数，我们可以将测试输入定义成 [-1, 1] 上一些更密集的坐标，并用上面定义好的 10 个样本做为训练集，得到核回归的结果。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

def kernel_regression(xs, ys, x_input, sigma=1):
    # xs: [n, ]
    # ys: [n, ]
    # x_input: [m, ]
    N = xs.shape[0]
    xs = np.expand_dims(xs, 1)
    ys = np.expand_dims(ys, 1)
    x_input = np.expand_dims(x_input, 0)
    x_input = np.repeat(x_input, N, 0)
    weight = kernel_func(xs, x_input, sigma) # [n, m]
    weight_sum = np.sum(weight, 0)
    weight_dot = weight.T @ ys
    weight_dot = np.squeeze(weight_dot, 1)
    res = weight_dot / weight_sum
    return res

sigma = 1
xs = np.linspace(-1, 1, 100)
ys = kernel_regression(sample_x, sample_y, xs, sigma)
plt.title(f'sigma = {sigma}')
plt.plot(xs, ys)
plt.show()

我们可以通过修改 sigma 来得到不同的拟合效果。以下是我的一些结果：

可以看出，标准差越小，模型倾向于过拟合；随着标准差变大，曲线会逐渐平缓。我们需要不断调整超参数，在过拟合和欠拟合之间找到一个平衡。这种现象很容易解释：正态分布核函数的标准差越小，意味着每个训练数据的影响范围较小，那么测试样本更容易受到少数样本的影响；标准差增大之后，各个训练样本的影响开始共同起作用，我们拟合出的函数也越来越靠近正确的函数；但如果标准差过大，每个训练样本的影响都差不多，那么模型就什么都拟合不了了。

从实验结果中，我们能大致感受到核回归和低通滤波很像，都是将已知数据的平均效果施加在未知数据上。因此，在分析核回归的时候，往往会从频域分析核函数。如果核函数所代表低通滤波器的带宽 （bandwidth）越大，那么剩下的高频信息就更多，核回归也更容易拟合高频信息较多的数据。

神经正切核

那么，核回归是怎么和神经网络关联起来的呢？有研究表明，在一些特殊条件下，MLP 的最终优化结果可以用一个简单的核回归来表示。这不仅意味着我们可以神奇地提前预测梯度下降的结果，还可以根据核回归的性质来分析神经网络的部分原理。这种能表示神经网络学习结果的核被称为神经正切核（NTK）。

这些特殊条件包括 MLP 无限宽、SGD 学习率的学习率趋近 0 等。由于这些条件和实际神经网络的配置相差较远，我们难以直接用核回归预测复杂神经网络的结果。不过，我们依然可以基于这些理论来分析和神经网络相关的问题。傅里叶特征的分析就是建立在 NTK 上的。

NTK 的形式为

其中，$f$ 是参数为 $\theta$ 的神经网络，$\langle\cdot,\cdot \rangle$为内积运算。简单来看，这个式子是说神经网络的核回归中，任意两个向量间的相似度等于网络对参数的偏导的内积的期望。基于 NTK，我们可以分析出很多神经网络的性质，比如出乎意料地，神经网络的结果和随机初始化的参数无关，仅和网络结构和训练数据有关。

在学习傅里叶特征时，我们不需要仔细研究这些这些理论，而只需要知道一个结论：一般上述 NTK 可以写成标量函数 $h_{NTK}(\mathbf{x}_i^T\mathbf{x}_j)$，也就是可以先算内积再求偏导。这意味用核回归表示神经网络时，真正要关心的是输入间的内积。别看 NTK 看起来那么复杂，傅里叶特征论文其实主要就用到了这一个性质。

为了从理论上讲清为什么 MLP 难以拟合高频，作者还提及了很多有关 NTK 的分析，包括一种叫做谱偏差（spectral bias）的现象：神经网络更容易学习到数据中的低频特征。可能作者默认读者已经熟悉了相关的理论背景，这部分论述经常会出现逻辑跳跃，很难读懂。当然，不懂这些理论不影响理解傅里叶特征。我建议不要去仔细阅读这篇文章有关谱偏差的那一部分。

正如我们在前文的核回归实验里观察到的，核回归模型能否学到高频取决于核函数的频域特征。因此，这部分分析和 NTK 的频域有关。对这部分内容感兴趣的话可以去阅读之前有关谱偏差的论文。

傅里叶特征的平移不变性

在上两节中，我们花了不少功夫去认识谱回归和 NTK。总结下来，其实我们只需要搞懂两件事：

• 神经网络最终的收敛效果可以由简单的核回归决定。而核回归重点是定义两个输入之间的相似度指标（核函数）。
• 表示神经网络的核回归相似度指标是 NTK，它其实又只取决于两个输入的内积 $\mathbf{x}_i^T\mathbf{x}_j$。

根据这一性质，我们可以部分解释为什么在文章开头那个 MLP 拟合连续图像的实验中，位置编码可以提升 MLP 拟合高频信息的能力了。这和位置输入的特性有关。

当 MLP 的输入表示位置时，我们希望模型对输入位置具有平移不变性。比如我们现在有一条三个样本组成的句子 $(1, A), (2, B), (3, C)$。当我们同时改变句子的位置信息时，比如将句子的位置改成 $(11, A), (12, B), (13, C)$时，网络能学出完全一样的东西。但显然不对输入位置做任何处理的话， $(1, 2, 3)$ 和 $(11, 12, 13)$ 对神经网络来说是完全不同的意思。

而使用位置编码的话，情况就完全不同了。假如输入数据是二维坐标 $\mathbf{v}\in [0, 1)^d$，我们可以用下面的式子建立一个维度为 $2m$ 的位置编码：

其中 $a_i$ 是系数， $b \in \mathbb{R}^{m \times 2}$ 是一个投影矩阵，用于把原来 2D 的位置变成一个更长的位置编码。当然，由于位置编码中既要有 $\sin$ 也要有 $\cos$，所以最终的位置编码长度为 $2m$。

根据我们之前的分析，NTK 只取决于输入间的内积。算上位置编码后，一对输入位置 $\mathbf{v}_1, \mathbf{v}_2$ 的内积为：

而根据三角函数和角公式可知：

这样，上面那个内积恰好可以写成：

上式完全由位置间的相对距离 $\mathbf{v_1}-\mathbf{v_2}$ 决定。上式决定了 NTK，NTK 又决定了神经网络的学习结果。所以，神经网络的收敛结果其实完全取决于输入间的相对距离，而不取决于它们的绝对距离。也因此，位置编码使得 MLP 对于输入位置有了平移不变性。

加入位置编码后，虽然 MLP 满足了平移不变性，但这并不代表 MLP 学习高频信息的能力就变强了。平移不变性能给我们带来什么好处呢？作者指出，当满足了平移不变性后，我们就能手动调整 NTK 的带宽了。回想一下我们上面做的核回归实验，如果我们能够调整核的带宽，就能决定函数是更加高频（尖锐）还是更加低频（平滑）。这里也是同理，如果我们能够调大 NTK 的带宽，让它保留更多高频信息，那么 MLP 也就能学到更多的高频信息。

作者在此处用信号处理的知识来分析平移不变性的好处，比如讲了新的 NTK 就像一个重建卷积核 （reconstruction filter），整个 MLP 就像是在做卷积。还是由于作者省略了很多推导细节，这部分逻辑很难读懂。我建议大家直接记住推理的结论：平移不变性使得我们能够调整 NTK 的带宽，从而调整 MLP 学习高频的能力。

那我们该怎么调整 NTK 的带宽呢？现在的新 NTK 由下面的式子决定：

为了方便分析，我们假设$\mathbf{v}$和$\mathbf{b_j}$都是一维实数。那么，如果我们令$b_j=j$的话：

这个式子能令你想到什么？没错，就是傅里叶变换。$j$ 较大的项就表示 NTK 的高频分量。我们可以通过修改前面的系数 $a_j$ 来手动调整 NTK 的频域特征。我们能看到，位置编码其实就是在模拟傅里叶变换，所以作者把位置编码总结为傅里叶特征。

作者通过实验证明我们可以手动修改 NTK 的频谱。实验中，作者令 $b_j=j, a_j=1/j^p$。$p=\infty$ 表示位置编码只有第一项：$\gamma(v)=[\cos 2\pi v, \sin 2\pi v]^T$。不同 $p$ 时 NTK 的空域和频域示意图如下所示。可以看出，令 $p=0$ 时，即傅里叶特征所有项的系数都为 $1$ 时，NTK 的高频分量不会衰减。这也意味着 MLP 学高频信息和低频信息的能力差不多。

随机傅里叶特征

现在我们已经知道傅里叶特征的公式是什么，并知道如何设置其中的参数 $a_j$, $\mathbf{b}_j$ 了。现在，还有一件事我们没有决定：该如何设置傅里叶特征的长度 $m$ 呢？

既然我们说傅里叶特征就是把输入的位置做了一次傅里叶变换，那么一般来讲，傅里叶特征的长度应该和原图像的像素数一样。比如我们要表示一个 $256 \times 256$ 的图像，那么我们就需要令 $m = 256 \times 256 / 2$ ，$\mathbf{b}$ 表示不同方向上的频率：$[(1, 1), (1, 2), …, (128, 256)]$。但这样的话，神经网络的参数就太多了。可不可以令 $m$ 更小一点呢？

根据之前的研究 Random features for large-scale kernel machines 表明，我们不需要密集地采样傅里叶特征，只需要稀疏地采样就行了。具体来说，我们可以从某个分布随机采样 $m$ 个频率 $\mathbf{b_j}$ 来，这样的学习结果和密集采样差不多。当然，根据前面的分析，我们还是令所有系数 $a_j=1$。在实验中，作者发现，$\mathbf{b_j}$ 从哪种分布里采样都无所谓，关键是 $\mathbf{b_j}$ 的采样分布的标准差，因为这个标准差决定了傅里叶特征的带宽，也决定了网络拟合高频信息的能力。实验的结果如下：

我们可以不管图片里 $1/f^x$ 是啥意思，只需要知道 a, b, c 是三组不同的实验就行。虚线是密集采样傅里叶特征的误差，它的结果反映了一个「较好」的误差值。令人惊讶的是，不管从哪种分布里采样 $\mathbf{b_j}$，最后学出来的网络误差都差不多。问题的关键在于采样分布的标准差。把标准差调得够好的话，模型的误差甚至低于密集采样的误差。

也就是说，虽然我们花半天分析了位置编码和傅里叶变换的关系，但我们没必要照着傅里叶变换那样密集地采样频率，只需要随机选一些频率即可。当然，这个结论只对 MLP 拟合连续数据的任务有效，和 Transformer 里的位置编码无关。

代码实现随机傅里叶特征

现在，我们可以回到博文开头的代码，看一下随机傅里叶特征是怎么实现的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

class FourierFeature(nn.Module):
    def __init__(self, in_c, out_c, scale):
        super().__init__()
        fourier_basis = torch.randn(in_c, out_c // 2) * scale
        self.register_buffer('_fourier_basis', fourier_basis)
        
    def forward(self, x):
        N, C, H, W = x.shape
        x = rearrange(x, 'n c h w -> (n h w) c')
        x = x @ self._fourier_basis
        x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)
            
        x = 2 * torch.pi * x
        x = torch.cat([torch.sin(x), torch.cos(x)], dim=1) 
        return x
        
feature_length = 256
model = MLP(feature_length).to(device)
fourier_feature = FourierFeature(2, feature_length, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
for epoch in tqdm(range(n_loops)):
    x = fourier_feature(grid)
    output = model(x)
    loss = F.l1_loss(output, input_image)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 100 == 0 or epoch == n_loops - 1:
        viz_image(output[0])
        print(loss.item())
prev_output = output

傅里叶特征通过类 FourierFeature 实现。其代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

class FourierFeature(nn.Module):
    def __init__(self, in_c, out_c, scale):
        super().__init__()
        fourier_basis = torch.randn(in_c, out_c // 2) * scale
        self.register_buffer('_fourier_basis', fourier_basis)
        
    def forward(self, x):
        N, C, H, W = x.shape
        x = rearrange(x, 'n c h w -> (n h w) c')
        x = x @ self._fourier_basis
        x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)
            
        x = 2 * torch.pi * x
        x = torch.cat([torch.sin(x), torch.cos(x)], dim=1) 
        return x

构造函数里的 fourier_basis 表示随机傅里叶特征的频率，对应论文公式里的$\mathbf{b}$，scale 表示采样的标准差。初始化好了随机频率后，对于输入位置 x，只要按照公式将其投影到长度为 out_c / 2 的向量上，再对向量的每一个分量求 sin, cos 即可。按照之前的分析，我们令所有系数 $a$ 为 $1$，所以不需要对输出向量乘系数。

傅里叶特征在 StyleGAN3 里的应用

傅里叶特征最经典的应用就是 NeRF 这类过拟合连续数据任务。除此之外，傅里叶特征另一次大展身手是在 StyleGAN3 中。

StyleGAN3 希望通过平滑地移动生成网络的输入来使输出图片也发生对应的移动。为此，StyleGAN3 将生成网络的输入定义为频域上的一个有限带宽图像信号：根据信号处理知识，我们能够将有限带宽信号转换成空域上无限连续的信号。也就是说，不管输入的分辨率（采样率）多低，我们都能够平滑地移动输入图片。StyleGAN3 借助随机傅里叶特征来实现这样一个频域图像。

以下代码选自 StyleGAN3 中傅里叶特征的构造函数。这个函数的关键是随机生成一些频率固定，但方向可以不同的傅里叶频率。函数先随机采样了一些频率，再将它们归一化，最后乘上指定的带宽 bandwidth，保证所有频率大小相等。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

class SynthesisInput(torch.nn.Module):
    def __init__(self,
        w_dim,          # Intermediate latent (W) dimensionality.
        channels,       # Number of output channels.
        size,           # Output spatial size: int or [width, height].
        sampling_rate,  # Output sampling rate.
        bandwidth,      # Output bandwidth.
    ):
        super().__init__()
        self.w_dim = w_dim
        self.channels = channels
        self.size = np.broadcast_to(np.asarray(size), [2])
        self.sampling_rate = sampling_rate
        self.bandwidth = bandwidth

        # Draw random frequencies from uniform 2D disc.
        freqs = torch.randn([self.channels, 2])
        radii = freqs.square().sum(dim=1, keepdim=True).sqrt()
        freqs /= radii * radii.square().exp().pow(0.25)
        freqs *= bandwidth
        phases = torch.rand([self.channels]) - 0.5

而在使用这个类获取网络输入时，和刚刚的 MLP 实现一样，我们会先生成一个二维坐标表格 grid 用于查询连续图片每一处的颜色值，再将其投影到各个频率上，并计算新向量的正弦函数。

这段代码中，有两块和我们自己的实现不太一样。第一，StyleGAN3 允许对输入坐标做仿射变换（平移和旋转）。仿射变换对坐标的影响最终会转化成对三角函数相位 phases 和频率 freqs 的影响。第二，在计算三角函数时，StyleGAN3 只用了正弦函数，没有用余弦函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

def forward(self, ...):
   ...

    # Transform frequencies.
    phases = ...
    freqs = ...

    # Construct sampling grid.
    theta = torch.eye(2, 3, device=w.device)
    theta[0, 0] = 0.5 * self.size[0] / self.sampling_rate
    theta[1, 1] = 0.5 * self.size[1] / self.sampling_rate
    grids = torch.nn.functional.affine_grid(theta.unsqueeze(0), [1, 1, self.size[1], self.size[0]], align_corners=False)

    # Compute Fourier features.
    x = (grids.unsqueeze(3) @ freqs.permute(0, 2, 1).unsqueeze(1).unsqueeze(2)).squeeze(3) # [batch, height, width, channel]
    x = x + phases.unsqueeze(1).unsqueeze(2)
    x = torch.sin(x * (np.pi * 2))
    x = x * amplitudes.unsqueeze(1).unsqueeze(2)

    ...

    # Ensure correct shape.
    x = x.permute(0, 3, 1, 2) # [batch, channel, height, width]
    return x

我们在 MLP 拟合连续图像的实验里复现一下这两个改动。首先是二维仿射变换。给定旋转角 theta 和两个方向的平移 tx, ty，我们能够构造出一个 $3 \times 3$ 的仿射变换矩阵。把它乘上坐标 [x, y, 1] 后，就能得到仿射变换的输出。我们对输入坐标 grid 做仿射变换后得到 grid_ext，再用 grid_ext 跑一遍傅里叶特征和 MLP。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

N, C, H, W = grid.shape
tx = 50 / H
ty = 0
theta = torch.tensor(torch.pi * 1 / 8)
affine_matrix = torch.tensor([
    [torch.cos(theta), -torch.sin(theta), tx],
    [torch.sin(theta), torch.cos(theta), ty],
    [0, 0, 1]
]
).to(device)
grid_ext = torch.ones(N, 3, H, W).to(device)
grid_ext[:, :2] = grid.clone()
grid_ext = grid_ext.permute(0, 2, 3, 1)
grid_ext = (grid_ext @ affine_matrix.T)
grid_ext = grid_ext.permute(0, 3, 1, 2)[:, :2]

x = fourier_feature(grid_ext)
output = model(x)
viz_image(output[0])

在示例代码中，我们可以得到旋转 45 度并向下平移 50 个像素的图片。可以看到，变换成功了。这体现了连续数据的好处：我们可以在任意位置对数据采样。当然，由于这种连续数据是通过过拟合实现的，在训练集没有覆盖的坐标处无法得到有意义的颜色值。

之后，我们来尝试在傅里叶特征中只用正弦函数。我们将投影矩阵的输出通道数从 out_c / 2 变成 out_c，再在 forward 里只用 sin 而不是同时用 sin, cos。经实验，这样改了后完全不影响重建质量，甚至由于通道数更多了，重建效果更好了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

class FourierFeature(nn.Module):
    def __init__(self, in_c, out_c, scale):
        super().__init__()
        fourier_basis = torch.randn(in_c, out_c) * scale
        self.register_buffer('_fourier_basis', fourier_basis)
        
    def forward(self, x):
        N, C, H, W = x.shape
        x = rearrange(x, 'n c h w -> (n h w) c')
        x = x @ self._fourier_basis
        x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)
            
        x = 2 * torch.pi * x
        x = torch.sin(x)
        return x

StyleGAN3 论文并没有讲为什么只用 sin，网上也很少有人讨论傅里叶特征的实现细节。我猜傅里叶特征并不是非得和傅里叶变换完全对应，毕竟它只是用来给神经网络提供更多信息，而没有什么严格的意义。只要把输入坐标分解成不同频率后，神经网络就能很好地学习了。

只用 sin 而不是同时用 sin, cos 后，似乎我们之前对 NTK 平移不变的推导完全失效了。但是，根据三角函数的周期性可知，只要是把输入映射到三角函数上后，网络主要是从位置间的相对关系学东西。绝对位置对网络来说没有那么重要，不同的绝对位置只是让所有三角函数差了一个相位而已。只用 sin 的神经网络似乎也对绝对位置不敏感。为了证明这一点，我把原来位于 [0, 1] 间的坐标做了一个幅度为 10 的平移。结果网络的误差几乎没变。

1
2
3
4
5
6
7

for epoch in tqdm(range(n_loops)):
    x = fourier_feature(grid + 10)
    output = model2(x)
    loss = F.l1_loss(output, input_image)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

根据这些实验结果，我感觉是不是从 NTK 的角度来分析傅里叶特征完全没有必要？是不是只要从直觉上理解傅里叶特征的作用就行了？按我的理解，傅里叶特征在真正意义在于显式把网络对于不同频率的关注度建模出来，从而辅助网络学习高频细节。

总结

在这篇博文中，我们学习了傅里叶特征及其应用，并顺带了解其背后有关核回归、NTK 的有关理论知识。这些知识很杂乱，我来按逻辑顺序把它们整理一下。

为了解释为什么 NeRF 中的位置编码有效，傅里叶特征论文研究了用 MLP 拟合连续数据这一类任务中如何让 MLP 更好地学到高频信息。论文有两大主要结论：

• 通过从 NTK 理论的分析，位置编码其实是一种特殊的傅里叶特征。这种特征具有平移不变性。因此，神经网络就像是在对某个输入信号做卷积。而我们可以通过调整傅里叶特征的参数来调整卷积的带宽，也就是调整网络对于不同频率的关注程度，从而使得网络不会忽略高频信息。
• 傅里叶特征的频率不需要密集采样，只需要从任意一个分布随机稀疏采样。影响效果的关键是采样分布的标准差，它决定了傅里叶特征的带宽，也就决定了网络是否能关注到高频信息。

由于这些结论比较抽象，我们可以通过一个简单的二维图像拟合实验来验证论文的结论。实验表明直接将坐标输入给 MLP 不太行，必须将输入转换成傅里叶特征才能有效让网络学到高频信息。这个傅里叶特征可以是随机、稀疏的。

除了过拟合连续数据外，傅里叶特征的另一个作用是直接表示带宽有限信号，以实现在空域上的连续采样。StyleGAN3 在用傅里叶特征时，允许对输入坐标进行仿射变换，并且计算特征时只用了正弦函数而不是同时用正弦、余弦函数。这表明有关 NTK 的理论分析可能是没有必要的，主要说明问题的还是实验结果。

傅里叶特征论文仅研究了拟合连续数据这一类问题，没有讨论 Transformer 中位置编码的作用。论文中的一些结论可能无法适用。比如在大模型的位置编码中，我们还是得用密集的 sin, cos 变换来表示位置编码。不过，我们可以依然借助该论文中提到的理论分析工具，来尝试分析所有位置编码的行为。

只通过文字理解可能还不太够，欢迎大家尝试我为这篇博客写的 Notebook，通过动手做实验来加深理解。 https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/FourierFeature

推荐参考资料：

用代码示例理解核回归：https://towardsdatascience.com/kernel-regression-made-easy-to-understand-86caf2d2b844

直观理解 NTK: https://www.inference.vc/neural-tangent-kernels-some-intuition-for-kernel-gradient-descent/

傅里叶特征官方示例 PyTorch 实现: https://github.com/ndahlquist/pytorch-fourier-feature-networks/tree/master

论文速览

平时我自己读新论文时，往往简单看一下摘要和图表就能差不多明白文章的大意。而要用一篇博文来详细介绍新论文的设计动机、实现方法，需要花多得多的时间来写作。为了能分享更多的新论文，并且让我自己能够更好地整理论文间的关系，从这篇文章开始，我会不定期发表「论文速览」类文章，以简明扼要的文字来帮助计算机视觉的同行们快速了解新论文。

和发文频繁的自媒体相比，我的论文速览文章几乎不会展示论文的结果，仅会从科研角度介绍文章的贡献及方法，并给出我的批判性分析。和我之前的论文详解文章相比，这个新系列的文章会用更少的文字介绍背景，几乎不介绍方法细节，力求用更少的文字来表达关键信息，提升文字沟通效率。不过另一方面，我会花更多心思在介绍新论文与过往论文的关系上，为不熟悉某些背景知识的读者提供学习途径，以填补背景介绍上的空缺。

Diffusion Forcing

视频扩散模型普遍存在视频质量随着视频长度增加不断下降的问题。为此，论文作者提出了 Diffusion Forcing 这一建模任何序列生成问题的新范式：在训练该扩散模型时，序列中的每个元素会独立地添加不同强度的噪声。作者在简单视频生成、决策 (decision-making) 任务上验证了这种生成范式的有效性。我将主要从视频生成的角度介绍本工作。

以往工作

稍后我们会了解到，Diffusion Forcing 与此前两种主流序列生成范式密切相关：自回归生成与全序列扩散模型。

在序列的自回归生成中，模型会不断根据序列前 $n-1$ 个元素，预测下一个元素（第 $n$ 个元素）。自回归生成在 NLP 中最为常见，RNN、Transformer 用的都是这种生成方法。

扩散模型可以直接生成任何形状的数据。如果我们不把视频视为一种由图像组成的序列数据，而是将其视为一种「三维图像」，那么我们可以直接将 2D 图像扩散模型简单地拓展成 3D 视频扩散模型。这种做法在这篇论文中被称为「全序列扩散模型」。使用这一方法的早期工作有 DDPM 的作者提出的 Video Diffusion Models。Stable Diffusion 的作者也基于 LDM 提出类似的 Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models（Video LDM）工作。

全序列扩散模型仅能生成固定长度的视频。为了将其拓展到长视频生成，还是得将其和自回归结合起来。但是，自回归视频生成存在着生成质量与训练集质量不匹配的问题。Stable Video Diffusion 等工作参考 Cascaded Diffusion Models 的做法，通过给约束图像/视频帧加噪声缓解此问题。

AR-Diffusion: Auto-Regressive Diffusion Model for Text Generation 进一步探讨了自回归生成与全序列扩散模型的结合方法：在生成文本时，不同时刻的文本噪声不同，越早的文本上的噪声越少。无独有偶，FIFO-Diffusion: Generating Infinite Videos from Text without Training 提到了如何在预训练视频扩散模型上，以不同的噪声强度来生成不同的视频帧。或许是受到这些工作的启发，Diffusion Forcing 系统探讨了如何在训练时独立地给序列元素添加噪声。

科研动机

本工作的作者发现了自回归与全序列扩散模型的不足之处，并认为这两种生成范式是互补的：

• 自回归不能在推理时加入新的优化目标，且存在着前述的训练采样质量不匹配导致的降质问题。
• 全序列扩散模型不能生成变长的序列。

反过来讲：

• 全序列扩散模型可以在推理加入 Classifier-guidance，且在固定序列长度内几乎不存在降质。
• 自回归可以生成变长序列。

那能不能把二者结合起来呢？在推理时，我们希望序列是按时间顺序先后生成的。同时，从噪声强度这一维度上，我们希望每个元素能够从完全带噪到完全清晰逐渐生成。结合二者的方法是：在采样时，较早的元素有较少的噪声，较新的元素有较多的噪声，不同噪声强度的元素在序列中同时生成。比如采样分3步完成，共生成3帧，那么一步去噪会让$[x_1^{1/3\cdot T}, x_2^{2/3\cdot T}, x_3^{T}]$变成$[x_1^0, x_2^{1/3\cdot T}, x_3^{2/3\cdot T}]$。

为了实现这样的采样，我们在训练时就要让每一元素能够在其他元素和自己的噪声强度不同时也能顺利去噪。和以往的工作不同，本工作的作者发现，训练时我们不必和推理时一样固定每一元素的噪声强度，而是可以独立地对每帧随机采样噪声强度。

简单视频生成模型

本文的想法十分简洁，只需要在 DDPM 序列生成模型上把训练和推理时的噪声强度换掉即可。为了进一步了解本工作的方法，我们来看一下论文中有关视频模型的方法与实验。

整体上看，训练时，方法和 $\epsilon$ 预测的普通 DDPM 一样，只是不同帧噪声强度不一样。

在帧间关系上，Diffusion Forcing 用因果（causal）关系建模，即当前帧只能看到之前帧的信息。

具体来说，本工作用 RNN （准确来说是 GRU）的隐变量来建模之前帧传过来的信息。加入了 RNN 后，论文把本来很简单的公式变得无比复杂，不建议读者深究论文中有关 RNN 的内容。

在采样时，由于不同帧的噪声强度不同，现在我们需要定义一个二维的噪声强度表，表示每一帧在不同位置及不同去噪时刻的噪声强度。为了让每一时刻的噪声强度不同，一开始较新帧的去噪时刻会停留在原地。作者在附录中介绍了同时刻去噪的详细设计。

作者发现，Diffusion Forcing 的这种序列生成方式可以自然地推广到无限长度的视频生成上：在生成下一个片段时，不用滑动窗口，直接把 RNN 的初始隐变量初始化为上一个片段的隐变量输出。

作者用同样的 RNN 结构训练了两个基准模型：一个自回归模型，一个全序列因果扩散模型。定性结果表明，不管是在预定视频长度内，还是超出原本长度的长视频生成，Diffusion Forcing 的结果均好于基准方法。结果可以在官方项目网站上查看：

https://boyuan.space/diffusion-forcing/

批判性分析

论文开头说自回归模型在推理时缺少添加约束的方法。但这对视频生成来说并不致命，因为一般可以在训练时加入约束，推理时用 Classifier-free Guidance 就行了。

作者说出于简洁，他们在论文中用 RNN 实现了 Diffusion Forcing。但很明显 3D U-Net 才应该是直观上最简单实用的实现方法，毕竟最早期的视频扩散模型就是拿 3D U-Net 做的。在官方仓库中，有本科生帮他们实现了一个 3D U-Net 加时间注意力的模型，比原来视频模型效果要好。

我认为本文的视频生成基准方法设置过低。对于自回归视频生成/图生成视频，现在大家都会参照 Cascaded Diffusion, 对约束图像加噪并把噪声强度做为额外约束传入当前帧的生成模型。这种设计和 Diffusion Forcing 原理相似。为了体现新方法的好处，有必要跟这个更加强大的基准自回归方法做对比。

作者对于全序列视频扩散模型的设计也很奇怪。全序列视频扩散模型的初衷就是把视频当成 3D 图像来看待，允许帧间两两交换信息，只保证预定长度内的视频是连贯的。作者现在用 RNN 实现了一个因果版全序列视频模型，这个模型的表现肯定是不如非因果版的。虽然作者说 Diffusion Forcing 在因果视频生成上总是比全序列扩散模型更加连贯，我很怀疑去掉了因果这个条件后 Diffusion Forcing 还能否比得过。

Diffusion Forcing 在视频生成的主要好处应该体现在超出预定长度的长视频生成。因此，哪怕在预定长度内比不过全序列扩散模型，也没有关系。作者应该比较结合自回归和全序列扩散模型的方法，比如用 Stable Video Diffusion 这种图生视频模型，把上一个视频的末帧当作下一个视频的首帧约束，证明 Diffusion Forcing 在长视频生成上的优越性。

综上，我认为作者在视频生成任务上的实验是不够充分的。也的确，这篇论文有一半篇幅是在决策任务上，没有只讲视频生成任务。我相信 Diffusion Forcing 的设计会在长视频生成上缓解降质问题，这需要后续大公司的工作来跟进。但是，长视频的根本问题是记忆缺失，这一本质问题是 Diffusion Forcing 这种方法难以做好的。

这篇工作对我最大的启发是，我们一直把视频当成完整的 3D 数据来看待，却忘了视频可以被看成是图像序列。如果把视频当成 3D 数据的话，不同帧只能通过时序注意力看到其他帧在当前去噪时刻的信息；而对于序列数据，我们可以在不同帧的依赖关系上做更多设计，比如这篇工作的不同去噪强度。我很早前就在构思一种依赖更加强的序列生成范式：我们可不可以把序列中其他元素的所有去噪时刻的所有信息（包括中间去噪结果及去噪网络的中间变量）做为当前元素的约束呢？这种强约束序列模型可能对多视角生成、视频片段生成等任务的一致性有很大帮助。由于生成是约束于另一个去噪过程的，我们对此去噪过程做的任何编辑，都可以自然地传播到当前元素上。比如在视频生成中，如果整个视频约束于首帧的去噪过程，那么我们用任意基于扩散模型的图像编辑方法来编辑首帧，都可以自然地修改后续帧。当然，我只是提供一个大致的想法，暂时没有考虑细节，欢迎大家往这个方向思考。

有人肯定也会想到能否把 Diffusion Forcing 拓展到图像像素间关系上。我认为要实现训练是完全没有问题的，问题出在推理上：Diffusion Forcing 在推理时需要预定义不同元素间的去噪强度。对于视频这种有先后顺序的数据，我们很自然地可以让越早的帧噪声强度越低。但是，对于图像来说，如何定义不同像素在不同去噪步数时的去噪强度并不是一个易解的问题。

近期，有两个大型多模态模型于同期公布：一个是来自 Meta 的 Transfusion，另一个是来自 Show Lab 和字节跳动的 Show-o 。好巧不巧，二者都宣称自己的模型是几乎最早将多模态任务用一个 Transformer 完成的，不需要借助额外的文本编码器实现图像生成，同时结合了自回归生成和扩散模型。我很好奇这两篇工作究竟有多少创新，于是快速扫完了这两篇论文，并简单给大家分享一下它们的核心内容。

在这篇文章中，我会快速介绍两篇工作的核心模型架构与部分实验结果。由于我仅对视觉任务比较熟悉，对语言和多模态没有那么了解，我的分析将主要围绕视觉任务。

论文 Arxiv 链接：

Transfusion: https://arxiv.org/pdf/2408.11039

Show-o: https://arxiv.org/pdf/2408.12528

读前准备

在阅读这两篇新工作时，建议大家先熟悉以 Transformer 为代表的自回归生成、以 DDPM、LDM、DiT 为代表的扩散模型、以 MaskGIT (Masked Generative Image Transformer), MAR (Masked autoregressive models, 于 Autoregressive Image Generation without Vector Quantization 论文中提出) 为代表的掩码自回归图像生成这三类生成模型，并简单了解此前较为先进的 Chameleon (Chameleon: Mixed-Modal Early-Fusion Foundation
Models) 多模态模型。本文不会对这些知识做深入回顾，如果读者遇到了不懂的旧概念，请先回顾有关论文后再来看这两篇新文章。

自回归模型

自回归模型用于生成形如 $\mathbf{x}=\{x_1, x_2, …, x_n\}$ 这样的有序序列。自回归算法会逐个生成序列中的元素。假设正在生成第 $i$ 个元素，则算法 $F$ 会参考之前所有的信息 $\{x_j \mid j < i\}$，得到 $x_i = F(\{x_j \mid j < i\})$。比如：

自回归任务最常见的应用场景是文本生成。给定第一个词，生成第二个词；给定前两个词，生成第三个词……。

为了训练实现这一算法的模型，一般我们需要假设每个元素的取值是有限的。比如我们要建模一个生成单词的模型，每个元素是小写字母，那么元素的取值只有 a, b, c, ..., z。满足这个假设后，我们就可以像分类任务一样，用神经网络模型预测的类别分布建模已知之前所有元素时，下一个元素的分布，并用交叉熵损失函数去优化模型。这种训练任务被称为下一个词元预测 (next token prediction, NTP)。

用自回归生成建模某类数据时，最重要的是定义好每个元素的先后顺序。对于文本数据，我们只需要把每个词元 (token) 按它们在句子里的先后顺序编号即可。而对于图像中的像素，则有多种编号方式了。最简单的一种方式是按从左到右、从上到下的顺序给像素编号。

掩码自回归模型

由于图像的像素数很多，用自回归模型一个一个去生成像素是很慢的；另外，按从左到右、从上到下的顺序给像素编号显然不会是最合理的。为了提升自回归模型的速度和表现，研究者提出了掩码自回归模型。它做了两个改进：

1） 相比按序号一个一个生成元素的经典自回归模型，这种模型在每轮生成时可以生成多个像素（下图中的橙色像素）。

2） 相比从左到右、从上到下的固定顺序，像素的先后顺序完全随机（下图中的 (b) 和 (c) ）。

由于这种方式下必须一次给模型输入所有像素，并用掩码剔除未使用的像素，所以这种自回归被叫做掩码自回归。

扩散模型

扩散模型将图像生成表示成一个噪声图像 $\mathbf{x}_T$ 从时刻 $T$ 开始随时间变化 $\mathbf{x}_t = F(\mathbf{x}_{t+1})$，最后得到目标图像 $\mathbf{x}_0$ 的过程。和输入输出为部分像素的自回归模型不同，扩散模型的输入输出自始至终都是完整图像。

为了减少扩散模型的计算量，参考 Latent Diffusion Model (LDM) 的做法，我们一般会先用一个自编码器压缩图像，再用扩散模型生成压缩过的小图像。正如 NLP 中将文本拆成单词、标点的「词元化」(tokenize) 操作一样，这一步操作可以被称为「图块化」(patchify)。当然，有些时候大家也会把图块叫做词元，把图块化叫做图像词元化。

严格来说，本文讲到的「像素」其实是代表一个图块的图像词元。用「像素」是为了强调图像元素的二维空间信息，用「图像词元」是强调图像元素在自回归模型中是以一维序列的形式处理的。

有人认为，掩码自回归模型是一种逐渐把纯掩码图像变成有意义图像的模型，它和逐渐把纯噪声图像变成有意义图像的扩散模型原理类似。因此，他们把掩码自回归模型称为离散扩散模型。还有人认为扩散模型也算一种更合理的自回归，每轮输入一个高噪声图像，输出一个噪声更少的图像。但这些观点仅仅是从称呼上统一两种模型，两种模型在实现上还是有不少差别的。

Chameleon

Chameleon 似乎是此前最为先进的多模态模型，它是这两篇新工作的主要比较对象。在语言模型的基础上，Chameleon 并没有对图像的处理多加设计，只是以离散自编码器（如 VQGAN）的编码器为图像词元化工具，以其解码器为图像词元化还原工具，让被词元化后的图像词元以同样的方式与文本词元混在一起处理。

功能与效果

看方法前，我们先明确一下两个多模态模型能做的任务，及各任务的输入输出。

Transfusion 是一个标准多模态模型，也就是一个输入输出可以有图像词元的语言模型。它输入已知文本和图像，输出后续文本和图像。

基于这个多模态模型，可以做文生图任务。

这个模型似乎没有为特定任务设置特殊词元，所有图像功能完全靠文本指定。因此，要做图像编辑任务的话，需要在带文本标注的图像编辑数据集上微调。文章指出，只需要在一个仅有 8000 项数据的数据集上微调就能让模型具有一定的编辑能力。

相比之下，Show-o 可以在序列前多输入一个区分任务的特殊词元。因此，Show-o 可以完成多模态理解（输入多模态，输出文本描述）、图像生成（根据文字生成图像或填补图像）、多模态生成（输入输出都包含图片、文本）等丰富的任务。似乎特殊词元仅有多模态理解 (MMU, Multi-modal Understanding MMU)、文生图 (T2I, Text to Image) 这两种。

Transfusion 的基础模型在微调后才能做根据文本提示来编辑图像的任务，而 Show-o 的基础模型默认是在此类带有文本提示的图像编辑数据集上微调的。

方法

对于熟悉此前图像生成模型、语言模型的研究者来说，这两篇工作都仅用了现有技术，核心方法非常易懂。这两篇工作并不是试图开发一种新的图像生成技术，而是在考虑如何更好地将现有图像模型融入多模态模型。

在读新技术之前，我们先以 Chameleon 为例，看一下之前的多模态模型是怎么做生成的。在我看来，之前的多模态模型不应该叫「多模态模型」，而应该叫「强行把图像当成词元的语言模型」。语言模型在处理文本时，文本中的词语、标点会根据规则被拆成「词元」，成为模型的最小处理单位。然而，要用同样的方式处理图像，就要定义什么是「图像词元」。受到之前图像离散压缩模型（以 VQGAN 为代表）的启发，大家用一个编码器把图像切成图块，每个图块可以用一个 1, 2, 3 这样的整数序号表示，再用一个解码器把用序号表示的图块翻译回真实图像。这里的带序号图块就和文本里的单词一样，可以用「词元」来表示。文本、图像都被统一成词元后，就能用标准 Transformer 的下一个词元预测任务来训练多模态模型。

如下图所示，训练时，文本基于程序规则变成词元，而图像经过一个编码器模型变成词元。两类词元被当成一类数据，以同样的方式按下一个词元预测任务训练。生成时，多模态模型自回归地生成所有词元，随后文本词元基于程序规则恢复成文本，而图像词元通过解码器模型恢复成图像。

这种多模态模型最大的问题是没有充分设计图像词元生成，还是暴力地用 Transformer 自回归那一套。虽然有些模型会多加入一些图像生成上的设计，比如 LaVIT (Unified Language-Vision Pretraining in LLM with Dynamic Discrete Visual Tokenization) 用扩散模型来做图像词元的解码，但核心的图像词元生成还是离不开标准自回归。

Transfusion 和 Show-o 的设计初衷都是引入更先进的图像生成技术来改进图像词元生成。先看 Show-o。要改进标准的一个一个按顺序生成图像词元的图像自回归模型，最容易想到的做法就是按照 MaskGIT, MAR 那一套，将标准自回归换成掩码自回归。在做掩码自回归时，像素的先后顺序完全随机，且一次可以生成多个像素。另外，图像词元之间可以两两互相做交叉注意力，而不用像文本词元一样只能后面的对前面的做交叉注意力。

Show-o 莫名其妙地把自己的图像生成模型称为离散扩散模型。如前文所述，叫离散扩散模型还是掩码自回归，只是一个称呼上的问题。由于问题建模上的重大差异，大家一般还是会把扩散模型和掩码自回归看成两类模型。

而 Transfusion 更加激进地在革新了多模态模型中的图像生成方式。现在最好的图像生成技术不是扩散模型吗？我们干脆直接把整个扩散模型搬过来。于是，在 Transfusion 生成多模态内容时，程序会交替执行两种模式：在语言模型模式下，程序按标准自回归逐个生成文本词元。一旦生成了特殊词元 BOI (begin of image)，就切换到扩散模式。在扩散模式下，程序按 DiT, SD 3 (Stable Diffusion 3) 那种标准扩散模型的方式，一次性生成所有图像词元。结束此模式后，程序往词元序列里填一个 EOI (end of image)，重返语言模型模式。

同理，在训练时，两种模态也用不同的任务来训练。语言模型老老实实地按下一个词元预测训练，而扩散模型那部分就按照训练扩散模型的标准方式，先给所有图像词元加噪，再预测噪声。因此，只看图像生成任务的话，Transfusion 更像 SD 3 这种文生图模型，而不像此前的基于语言模型的多模态模型。

Transfusion 和 SD 3 之间的最大区别在于，文本词元还是按照语言模型那一套，只能在交叉注意力看到之前的文本词元。而图像词元之间两两都能看到。这种交叉注意力的设计和 Show-o 是一模一样的。当然，由于现在文本也会在同一个 Transformer 里处理，所以 Transfusion 自己就扮演了解读文本的工作，而不像 SD 3 那样还需要单独的文本编码器。

定量评测结果

我们最后来看一下两篇文章展示的定量评测结果。

Transfusion 用了许多篇幅来展示它与 Chameleon 之间的对比。从数值指标上看，Transfusion 全面领先 Chameleon。明明没有对文本任务做特别的优化，Transfusion 却在文本任务超越了 Chameleon，这挺令人惊讶的。为了探究这一现象的原因，作者从同一个预训练语言模型开始，以同样的配置训练了 Transfusion, Chameleon。结果显示，相比加入图像扩散模型，加入 Chameleon 那种离散图像词元对文本指标的损害更大。作者猜测，这是因为扩散模型更加高效，模型能够把更多精力放在文本任务上。

而从图像生成模型的对比上看，Transfusion 比之前多数文生图模型都要好，只是比顶尖文生图模型 SD 3 要差一点。

再来看 Show-o 的评测结果。Show-o 在部分文本指标上超过了之前的语言模型。作者也坦言，这些指标仅表明 Show-o 有潜力在文本任务上做到顶尖。

Show-o 也展示了图像任务的指标。和 Transfusion 一样，Show-o 展示了表示图像质量的 COCO FID 指标以及评价文本图像匹配度的 GenEval 指标。Show-o 在图像指标上超越了此前多数多模态模型，且超越了 Stable Diffusion 2.1 等图像生成模型。但是其图像指标比 Transfusion 还是差了不少。Show-o 的最大优点是需要的图像训练数据远远少于其他模型。

总结与讨论

此前多模态模型都只是强行把图像变成离散图像词元，再用标准自回归来生成图像词元。为了改进这些多模态模型，无独有偶，Transfusion 和 Show-o 都用到了更先进的图像生成技术。Show-o 将标准自回归改成了更强大的掩码自回归，而 Transfusion 激进地引入了完整的图像扩散模型，并把文本生成和图像生成当成两个相对独立的任务。二者的相同之处如下：

• 两个多模态模型都用同一个 Transformer 来处理文本和图像。
• 两个模型的 Transformer 都使用了同样的交叉注意力机制。文本词元只能看到此前的图像、文本词元，而图像词元可以看到此前的文本词元和当前所有图像词元。

二者的不同之处在于：

• Transfusion 使用标准扩散模型实现图像生成，而 Show-o 使用掩码自回归实现图像生成。Show-o 强行将自己的图像生成模型称为「离散扩散模型」，有借扩散模型的名头宣传之嫌。
• Transfusion 没有用特殊词元来区分不同任务。要用 Transfusion 编辑图像，需要在基础模型上用图像编辑数据集微调。Show-o 用特殊词元来区分不同任务，默认支持文本理解、图像生成、图像编辑、多模态生成等多种任务。
• 二者在文本、图像指标上都超越了之前的多模态模型。但二者相互对比之下，Transfusion 的表现更好，而 Show-o 需要的训练资源少得多。

我再来谈一下我看完这两篇文章后的一些感想。此前我对多模态模型很不感兴趣，觉得它们就是在语言模型的基础上，强行加入图像信息，然后加数据、加显卡，大火乱炖，没有太大的创新。而 Transfusion 和 Show-o 的设计令我眼前一亮。我觉得这两篇文章的结果再次表明，图像和文本的处理方法不应该是一致的，不能强行把文本自回归那套方法直接搬到图像任务上。不管是换成 MaskGIT 掩码自回归那一套，还是完全用扩散模型，都比之前的方法要好。

而究竟是掩码自回归好还是扩散模型更好呢？在我看来，扩散模型是更好的。文本是离散的，而图像是连续的。这种连续性不仅体现在图像的颜色值上，还体现在图像像素间的位置关系上。强行用 Transformer 自回归生成图像，一下子把这两种连续信息都破坏了。而何恺明团队近期的 MAR 工作则试图找回图像在颜色值上的连续性，但依然无法充分利用空间上的连续性。相比之下，扩散模型每步迭代都是对整幅图像做操作，不会破坏图像的连续性。这两个多模态工作也反映了这一点，Transfusion 的表现要比 Show-o 好很多。

在生成图像时，Transfusion 的行为几乎和 SD 3, FLUX.1 这些文生图模型一样了。两类模型的区别在于，SD 3 它们得用一个预训练的语言处理模型。而 Transfusion 用同一个 Transformer 来处理文本、图像信息。尽管现在 Transfusion 还没有超过 SD 3，但我认为文生图任务本质上是一个多模态任务，这两类模型的后续发展路线很可能会交汇到一起。文生图也应该从多模态中汲取经验。比如我们在 SD 3 的基础上，加入一些语言任务上的学习，说不定能进一步提升文生图时的图文匹配度。

当然，仅根据我读完这两篇文章后有限的认知，我认为多模态并不是一个值得广大科研人员投入的研究方向，而只适合有足够资源的大公司来做。其根本原因是验证多模态设计的代价太大了。在图像生成领域，要验证一个生成模型好不好，我们拿 ImageNet，甚至拿只有几万张图像的人类数据集训练及评测，都很有说服力。而多模态任务必须在大量文本、图像数据集上训练，评测文本、图像上的多种指标，小一点的团队根本做不来。这样，各种小创新都难以验证。而没有各种各样的小创新，领域也就很难快速发展。所以多模态只能不断从纯语言生成、纯图像生成方向找灵感，很难有仅属于多模态任务的创新。

今天看到一则新闻，说 Google DeepMind 发布了一个叫做 GameNGen （英文读音 “game engine”，游戏引擎）的扩散模型，它可以仅用神经网络的生成结果，模拟经典射击游戏 DOOM。

项目网站：https://gamengen.github.io/

视频链接：https://gamengen.github.io/static/videos/e1m1_t.mp4

作为一名未来的游戏设计师，每次看到这类「今天 AI 又取代了创作者」的新闻，我的第一反应总会是愤怒：创作是人类智慧的最高结晶，能做到这种程度的 AI 必然是强人工智能。但显然现在 AI 的水平没有那么高，那么这类宣传完全是无稽之谈。我带着不满看完了论文，果然这个工作并没有在技术上有革命性的突破。不过，这篇论文还是提出了一个比较新颖的科研任务并漂亮地将其解决了的，算是一篇优秀的工作。除了不要脸地将自己的模型称为「游戏引擎」外，这篇工作在宣传时还算克制，对模型的能力没有太多言过其实的描述。

如果不懂相关技术的话，外行人很容易对这篇工作的应用前景产生一些不切实际的幻想。在这篇文章中，我将完整而清晰地介绍这篇工作的内容，再给出我从游戏开发方面和科研方面对这篇工作的评价。

问题定义

看完了这篇工作的展示视频，大家第一个想问的问题一定是：这个模型的输入是什么？是只能随机生成游戏视频，还是能够根据用户的输入来生成后续内容？

答案是，模型可以根据之前的游戏进度及当前的用户输入，输出下一帧的游戏图片。我们来详细看一下该工作对于「游戏」的定义。

论文的第二章详细定义了该工作要完成的任务。作者认为，一个游戏可以由游戏状态、游戏画面、玩家操作集这三类信息组成。游戏状态包括角色血量、装备、地图、敌人等所有影响游戏进程的信息；游戏画面就是游戏屏幕上显示的二维图片；玩家操作集就是移动、射击等玩家所有可能进行的操作。此外，为了让游戏运行，还需要两类游戏机制：如何根据游戏状态生成当前游戏画面的渲染机制；如何根据游戏状态和当前玩家操作更新下一时刻游戏状态的逻辑机制。

而用一个生成模型来模拟游戏时，我们不需要让模型学会游戏机制、状态，只需要让模型根据之前所有的画面和玩家操作，以及当前时刻的玩家操作，输出当前时刻的画面。这样，我们给第一个初始帧和操作，模型就能输出第二帧；给前两帧和之前及现在的操作，模型就能输出第三帧……。也就是说，模型以自回归的方式模拟游戏画面。

给想读这篇论文的读者一点阅读上的提示。论文以「交互环境」$\mathcal{E}$来指代我上文中的「游戏」，以 「交互世界模拟」$q$ 来指代生成模型。在定义 $\mathcal{E}$ 用到的几个字母对应我前文的粗体名词。

在我看来，这篇工作的主要贡献，就是把游戏模拟任务以如此简明的形式清楚定义了出来。怎么把任务做好，纯粹只是工程实现问题。明确了任务后，相关领域的科研人员基本上能猜出这篇工作是怎么实现的了。下面，我们就来看一下论文中分享的实现过程。

GameNGen

用强化学习造数据

本工作的任务是根据某些信息生成下一帧的画面，这种生成某类图像的任务用日益成熟的图像扩散模型技术就能解决了。但是，为了训练扩散模型，本工作有一道跨不过的坎——缺乏游戏画面数据。

为了生成足够多的图片，作者利用强化学习训练了一个玩游戏的 AI。在这一块，作者用了一个非常巧妙的设计：和其他强化学习任务不同，这个玩游戏的 AI 并不是为了将游戏漂亮地通关，而是造出尽可能多样的数据。因此，该强化学习的奖励函数包括了击中敌人、使用武器、探索地图等丰富内容，鼓励 AI 制造出不同的游戏画面。

带约束图像扩散模型

有了数据后，问题进一步缩小：现在该怎么用 Stable Diffussion 这个比较成熟的图像生成模型来根据之前的图片和之前及当前的操作来生成图片。

「之前的图片」和「之前及现在的操作」是输入给图像生成模型的两类额外信息。用专业术语来说，它们是给一个随机生成图像的模型的约束条件，用于让模型的输出不那么随机。而带约束图像生成也是一个被研究得比较透的任务了。

以下内容是写给相关科研者看的，看不懂可以跳过。

先谈怎么让模型约束于之前的图片。本文参考了经典的 Cascaded Diffusion 实现图像约束：将之前的图片过 VAE 编码器，与扩散模型原来的噪声输入拼接。然而，这种约束方式存在分布不匹配的问题：训练时，图像约束来自训练集；推理时，图像约束来自于模型之前自回归的生成结果。为了填平两类图像在分布上的差异，我们需要给所有约束图像加噪，并把加噪程度当成额外约束输入进模型。这种图像约束方法和图生视频的 Stable Video Diffusion 是一样的。

再看模型怎么约束于操作。每个操作都有独特的含义，它和语言中的单词是类似的。因此，我们可以把离散的操作变成嵌入向量，用 Stable Diffusion 处理单词的机制来处理操作。所以，在 GameNGen 的扩散模型中，文本约束没有了，被操作约束取代了。顺带一提，我们要输入之前及当前的操作，它们构成了一个操作序列。我们只要把这个操作序列当成由单词构成的文本，还是用原来 Stable Diffusion 那套处理文本的机制就行了。

此外，该工作还微调了 Stable Diffusion 的 VAE 的解码器，用以提升其在特定数据上的重建效果。这种操作也是比较常见的，比如 Stable Video Diffusion 将图像解码器微调成了视频解码器。

实验及结果

本工作仅对图像约束用了强度为 1.5 的 Classifier-Free Guidance (CFG)，没有对操作约束加 CFG。采样用的是 DDIM 采样器，实验表明 4 步采样的结果就足够好了。在单个 TPU-v5 上，模拟模型每秒能渲染 20 帧画面。

扩散模型基于 Stable Diffusion V1.4 训练。输入图像分辨率为 $320 \times 256$。之前图像和操作的上下文窗口长度为 $64$。训练集包含 900M 张图片，要用 128 块 TPU-v5e 训练 700,000 步。

论文还评价了模型的生成质量。由于训练集能够提供当前要预测的帧的真值，因此我们可以用重建误差来反映模型的质量。具体来说，论文展示了表示图像相似度的 PSNR 和图像感知误差的 LPIPS。 重建质量通过两个任务来反映：

• 图像质量：输入之前图像、操作均为数据集（按理说是测试集而不是训练集）里的真值，仅评测当前帧。此时重建质量较好，平均 PSNR 为 29.43，和一般图像经 JPEG 压缩产生的损耗相近。
• 视频质量：给定初始图像和数据集里的之前操作，让模型自回归地生成一段视频。这种情况下每一帧的质量会逐渐降低，如下所示。

除此之外，论文还展示了人类的评估结果：给出两段长度为 1.6 秒或 3.2 秒的游戏视频，分别来自真实游戏和 AI 生成，请人类分辨哪段视频是 AI 生成的。对于 1.6 秒的视频，正确率为 58%；对于 3.2 秒的视频，正确率为 60%。

此类评测的最低（最优）正确率是 50%，因为我们总是可以随便猜一个。

论文总结与评价

信息整理

到目前为止，我已经客观介绍了论文中展示的内容。让我用外行人也能看懂的语言总结一下：

本工作提出的 GameNGen 模型可以根据之前的游戏画面、用户的历史操作、当前操作，在完全不了解游戏机制的前提下，生成包含游戏逻辑（血量、弹药）的当前帧画面。生成画面的模型是一个深度学习模型。因此，需要一个包含了过往操作、真实游戏画面的大型数据集来训练该模型。本工作利用强化学习制造了大量数据。训练用了 128 张 TPU-v5e 计算卡。该生成模型能在单张 TPU-v5 (专门用来做深度学习的计算「显卡」，只租不卖，比价值 25000 美元的 H100 显卡要好) 上能以 20 帧每秒的速度生成画面。模型最多能记住 64 帧，即 3.2 秒内的游戏信息。若将 3.2 秒的生成视频与真实游戏视频比较，人类仅有 60% 的概率分辨出 AI 生成的视频。当然，利用自回归技术反复输入之前生成的画面，模型也能够生成更长的视频，正如本文开头所展示的那个视频。

论文中提及的缺陷

哪怕不懂深度学习，大家也可以根据上述信息，提出自己的看法。当然，在那之前，我们先看一下作者在论文里是怎么描述模型的缺陷的。

作者讲到，GameNGen 受制于有限的记忆。模型仅能获取三秒之短的历史信息，却能在极其长的时间里保持游戏逻辑稳定，这很了不起（怎么讲着讲着又开始夸起自己了？）。但由于模型学到了太多东西，模型会创造短期记忆处理不了的场景，比如角色在某处打倒了敌人，数分钟后又返回打倒了敌人的地方（文章讲了半天讲不出个东西，我替他们总结了一下）。对于这些问题，以当前的模型架构，再怎么加大记忆窗口也无济于事。

面向交互视频游戏的新范式

基于这篇工作的成果，作者还展望了未来：（以下是我对原文的精心翻译，欢迎对比）

如今，视频游戏是靠人类编程实现的。而 GameNGen 表明了以神经网络的权重来描述游戏这种新范式的部分可行性。GameNGen 展示了，在现有硬件上用神经网络模型来高效地交互运行一个复杂的游戏 （DOOM）不是奢望，这样的模型架构与模型权重是存在的。尽管还有很多重要问题要解决，我们希望这种范式能够造福众生。比如说，用这种新范式开发视频游戏的代价可能更小，门槛更低，借此，我们或许能仅通过编辑文本描述或示例图片来开发游戏。这个愿景的一小部分，即在现有游戏上略作调整或创造新行为，也许会在不久的将来就能实现。比如即使拿不到作者的源代码，我们也可能可以把几段游戏截图变成一个可玩的游戏关卡或者基于现有图像创建新角色。这个范式还能带来其他好处，比如严格控制帧率、内存占用不变。路漫漫其修远兮，尽管我们目前还没有探究这些方向，我们还是乐于尝试！希望我们这一小步，会在未来的某天，化为人们享受视频游戏的美好瞬间；甚至更进一步，化为人们使用交互软件系统的点点日常。

作为一个对这段抒情深有共鸣的人，在看论文前，我有千言万语想喷，却找不到切入点。看了这段话，我总算知道该怎么针对性地发表意见了。

锐评

总有人说，AI 要取代人类了。

之前是说 AI 可以代替画师，又是说 AI 能代替人写小说。现在，来说 AI 能够完全模拟游戏了。

我真的很不解：人类的作品怎么会沦落到和现有 AI 的对比了？

我想了很久，为什么我无法容忍「现有 AI 技术能代替人类创作」这种观点。我的核心论点是：1）现有基于深度学习的 AI 无法达到和人同等的智力水平；2）达到人类同等的智力水平，意味着能够理解人类的行为，进而在创作、编程、教学、心理咨询等所有现在看来比较困难的领域看齐甚至超越人类。

详细对深度学习了解了一段时间后，绝大多数人都能推理出深度学习的上限。通过对深度学习应用的种种观察，我们能够用我们脑中的那个「神经网络」，那种基于数据推理现实的能力，预知深度学习的能力上限。深度学习适用于一些数据定义良好，目标定义良好的任务。只要给了数据，给了目标，网络就能学习，甚至涌现出一些意想不到的强大理解、生成能力。

但是，数据不是全部，永远有从大量数据学不到的东西。

那就是人心。

就和无数探讨机器人的作品所展示的一样。

正因为我们是人，所以我们能感知我们生活在这个世界上。我们感受痛苦，所以我们思考，并追逐美好。

无论是穷尽多少数据，复读多少经书也体会不到的；无论是洞察多少规律，拟合多少逻辑也推理不出的，就是人心。

和人心等价的一切任务，是现在的深度学习 AI 做不到的。

有关深度学习是否能达到人类水平，那是技术讨论，我们在别处再谈。如果最终能认同目前深度学习无法达到人类水平，但目前认为「现有 AI 技术能代替人类创作」的话，简单来看有两种可能：

1. 不懂深度学习

2. 不懂创作的难度

而对于一个相关专业人士来说，只能是第二种可能了。所以听到懂深度学习的人讲出「AI 要代替人类」时，我都会下意识地认为他在贬低人类创作的含金量，自然是怒火中烧。不懂创作，就不要妄加评论。

除此之外，还有人明明知道当前 AI 的实力，还要违心地宣传 AI 如何如何，宣传自己的垃圾工作多么有价值。这种人只会为自己的利益考虑，纯纯的坏而已。这种无可救药的坏连讨论的价值也没有。

最后还有一种可能，很多人做研究时，并不会像我这样想这么多。在他们看来，科研就是从现有知识出发，朝外迈一小步。不管这一步是否方向正确，不管这一步有多么小，只要是拓宽人类的知识边界，那就是好的。他们的研究是纯洁的、无私的，不在意有生之年能否看到自己的成果被用上，甚至不在意自己的研究是否真的会被用到，只是为了科研增砖添瓦而已。我不得不承认，这是真正的、高尚的科研。

他们是幸福的，可以不在意眼前的得失。所以，他们可以望着远处高峰，轻松而豪放地说出：「希望我的研究，能化为人们享受视频游戏的美好瞬间」。

然而，这句话，对我而言，是沉重的。如果这种话是我说出来的，那么它不会是期盼，而是矢志不渝的誓言。不是拿着望远镜向远处眺望，不是用手指着地图挥斥方遒，而是用我的脚，一步，一步，踏出来的。

我是各种作品的鼓舞下走过来的，优秀的作品对我而言是神圣的。所以，我希望立刻，亲眼见到更多的好作品。没有对好作品急功近利的渴望，也就说明他们生活中有更加便捷的能量来源。所以我说，真正能以纯洁的心做科研的人，是幸福的人。

把我这些话总结一下，能认为深度学习能代替人类创作游戏，要么是深度学习的信徒，要么是不懂深度学习，要么是不懂创作，要么是坏，要么就是没想那么多觉得有新科研工作就是好事。

有人可能还会说：「我也同意深度学习代替不了人类，但也不能说这些技术就完全没用」。这我非常同意，我就认为大家应该把现在的 AI 当成一种全新的工具。基于这些新工具，我们把创新的重点放在如何适配这些工具上，辅助以前的应用，或者开发一些新的应用，而不是非得一步到位直接妄想着把人类取代了。比如，根据简笔画生成图片就是一个很好的新应用啊。

回到这篇文章的锐评上来。一上来，标题就写着《扩散模型是实时游戏引擎》。其实这是一个在顶级计算卡上每秒生成 20 张低分辨率图片的模型，这真的是我们认为的实时吗？作为一个游戏引擎，你能修改游戏机制吗？哦，我对引擎的理解有误，这不是游戏开发引擎，而是一个运行游戏的引擎。那踏踏实实叫做「游戏模拟器」不好吗？标题取得夸张一点，想吸引大家注意，能够理解，不多讲了。

文章主体部分都是客观陈述，写得非常清晰，我读起来也很舒服。本来都准备把「锐评」改成「简述」的，看完作者最后那段对未来的畅想后，一阵无名火在我心中燃起。通过「我们或许能仅通过编辑文本描述或示例图片来开发游戏」这段话，我感觉这是作者是几个对游戏制作质量没有那么高要求的人。可是，就是这样的人，却能写出「路漫漫其修远兮，尽管我们目前还没有探究这些方向，我们还是乐于尝试！」这说明他们可能是真心热爱游戏的玩家。那为什么，为什么只做出了这种程度的工作呢？900M 张图片，128 块卡，别说爱好者，就是一般的大学实验室，都难以跟进这篇工作，这是想要给设计师、爱好者开发新工具的态度吗？没有其他更加贴近用户的项目了吗？好，你说你们以长期的科研为主，这只是这个方向的初步尝试，你们重心在科研上而不是提供游戏开发工具上。那你们是抱着多大的觉悟说出「希望我们这一小步，会在未来的某天，化为人们享受视频游戏的美好瞬间」的？给人的感觉就是一群深居象牙塔的人，一辈子也不去了解业界真的需要什么，只是「正确地」做着科研而已。

非常抱歉，以上都是我的主观评价，请恕我对作者的想法妄加猜测。看完文章最后那段话后，我就有了这样一种矛盾的愤怒感。我喷了这么多，其实也不是想喷这篇工作的作者，更想批判的，是我长年以来在生活中的见闻。把气撒到这篇工作上，可能只是我嫉妒他们，没有 128 块卡去做想做的事情而已。

但我毫不怀疑地相信，我要有的资源最后都会有的。「游戏开发的新范式」、「造福众生」、「今天的一小步」……，如果有一天我说出了这些话，那必然不是在论文里，而是在我的产品得到了用户的充分肯定后，向世界吹嘘的胜利宣言吧。

新科研方向的讨论

先谈一下这篇工作在科研上给我们的启发。我认为有三点：

1. 对于深度学习应用来说，不要去在意功能有多么异想天开，只要把问题定义好，数据准备好，问题就可解。
2. 可以以用户操作为约束，用生成模型建模一个可交互世界。这个「用户操作」不一定局限于游戏玩家的操作。
3. 强化学习可以用来造大批数据。并且，我们需要精心设计模型的学习目标，使其造出多样的数据。

再来看顺着这篇文章的结果，我们能有怎样的新思考。拿图像生成模型这种结果极不稳定的东西做要求输入输出可控的游戏是不可能的。但是，我们应该把思路逆转过来：哪些任务可以以不确定的图像为输入？最容易想到的是缺数据的自动驾驶任务。也就是说，这篇工作实际上是提出了一种带交互的图像数据生成器。

这篇工作用图像模型学习了 3D 场景在移动后的变化。也就是说，模型「理解」了 3D 场景。那么，有没有办法从模型中抽取出相关的知识呢？按理说，能理解 3D，就能理解物体是有远近的。那么，深度估计、语义分割这种任务是不是可以直接用这种模型来做呢？以交互为约束的图像生成模型可能蕴含了比文生图模型更加丰富的图像知识。很可惜，不知道这篇工作最后会不会开源。

如果一个模型能够建模一个简单的世界，我们下一步要思考的是怎么编辑这个世界。就像有了图像生成模型，我们要给它加上文本约束一样。最容易想到的是以二维图片为约束，生成一个世界里的三维物体。但目前这个模型要的训练量太大了，做这种新实验的代价根本不敢想。

当前这个模型的结果还是比较弱的。别看这里用了扩散模型，模型的训练目标实际上是一个重建任务而非生成任务，最后的评测指标也是重建误差而没有考虑 FID 等图像生成质量指标。有没有办法让模型设法输出有多样性的新内容？还有，这个模型在设计上应该是一个自回归模型，只不过下一张图片是用扩散模型隐式建立了概率分布。但由于用户操作是在线而不是离线一次性给出，这种任务在时序上没办法用视频扩散模型建模。所以，从本质上看，要设计带用户操作的生成模型，其实是要一种时序信息在线输入的生成模型。除了自回归模型外，能不能用一种全新的生成范式呢？一旦有了这样一种新式模型，训练交互世界神经网络的代价将大幅降低。

除了时序信息不好输入外，这个任务还面临另一个问题：随着时间不断推移，模型会忘记以前的内容。NLP 领域通过 Transformer 全局信息交互暴力地解决了这类问题，但在长视频生成中，这种问题还是没能很好解决。或许需要其他领域为长时序建模提供了更好的工具后，才能考虑长时间的世界交互。但我们也可以从另一个角度思考：把要模拟的世界简化，不要去做模拟游戏那么难的任务，做一个时序依赖少的模拟任务。

结束科研新方向的思考前，我再对这个模型的训练量喷两句。搞大批数据，搞大模型，一堆 GPU 狂训，强行拟合一个任务，在我看来是非常不优雅的做法。得不到学术界广泛跟进的工作，是发展不下去的。在思考这种世界模拟模型的种种应用前，最重要的还是把训练量降下去。我认为优化在线时序约束生成模型是最有价值的方向。

最后我再从游戏开发的角度讲一下这篇工作的启发。我们只讨论如何用生成模型减少美术工作量，不去探究怎么让模型学习设计复杂的游戏机制。

真的想要为游戏开发提供工具的话，应该从 2D 游戏而不是 3D 游戏入手。拟合 2D 游戏画面的训练代价较小，且 2D 游戏对开发者来说更容易做。仅基于现在的 AI 技术，我们就已经能够做出一些简单的美术生成应用了。比如前段比较火的植物大战僵尸杂交版，我们完全可以想办法定义一个科研任务，训练一个融合两种植物的生成模型。明明有更切实际的 2D 不做去做 3D，这也是我为什么觉得这篇工作在为游戏开发提供工具这个层面上显得非常没有诚意。这篇文章建模的 3D 世界模拟器其实更适合前面提到的自动驾驶等真实场景，反而不是很适合游戏开发。

从这篇工作出发，或许可以联想出非常多游戏开发应用。这里我随便举一个 2D 网格地图生成任务。用同一种地板拼接地形时，根据地板的连通情况，游戏开发引擎会自动生成完整地形。但是，这种生成是基于写死的规则的，且最终每种地形还是以正方形网格为单位呈现。能不能用生成模型让这种地形生成更加多样呢？

从本文的启发 2 出发，我们或许可以有一些新想法：可以以操作为约束，用生成模型建模一个可交互世界。那么，我们将「操作」定义为设计师向地图上铺上一块地板，而不是定义为玩家的操作。这样，模型就能站在设计师的角度的学习生成地图了。

如果数据足够的话，这种定义方式一定能让神经网络学会地形生成的。问题的关键就在于如何获取数据。根据玩家的操作，我们能够自然生成大量游戏数据。而根据设计师的每一步想法，每一步都绘制合理且高质量的图片，可能要花费大量资源。因此，如果是从这个角度出发，需要大量美术、游戏设计、深度学习的相关人员参与进来。

总结

GameNGen 将模拟 3D 可交互场景的任务定义为根据历史画面、历史及当前操作生成当前画面的带约束图像生成任务。该工作用强化学习巧妙造出大量数据，用扩散模型实现带约束图像生成。结果表明，该模型不仅能自回归地生成连贯的游戏画面，还能学会子弹、血量等复杂交互信息。然而，受制于硬件及模型架构限制，模型要求的训练资源极大，且一次只能看到 3.2 秒内的信息。这种大量数据驱动的做法难以在学校级实验室里复刻，也不能够归纳至更一般的 3D 世界模拟任务上。

我个人认为，从科研的角度来看，这篇工作最大的贡献是提出了一种用带约束图像生成来描述 3D 世界模拟任务的问题建模方式。其次的贡献是确确实实通过长期的工程努力把这个想法做成功了，非常不容易。但从游戏开发的角度来看，这个工作现阶段没什么用处。

从科研启发的角度思考，这篇工作告诉我们，定义好交互世界里的操作，我们就能部分地用图像生成模型建模一个交互世界。从本质上来看，这是一个每一个时刻的约束都在线给出的视频生成任务。针对这个任务，我们既可以去思考能否用自回归以外的更高效的方式来实现它，也可以去思考是否可以修改对于「操作」的定义来实现模拟玩家操作以外的世界模拟任务。

前几个月，推出了著名文生图模型 Stable Diffusion 的 Stability AI 公司曝出了核心团队集体离职的消息。一时间，AI 从业者们议论纷纷，不知道这究竟是团队出现了矛盾，还是这些员工觉得文生图模型做下去没有前途了。而近期，该核心团队重新组建的创业团队 Black Forest Labs（黑森林实验室）带着名为 FLUX.1 的文生图模型「复仇归来」。FLUX.1 受到了用户的广泛好评，让人们期盼更强开源文生图模型的热情得以延续。

Black Forest Labs 的成员基本上都是 Stable Diffusion 3 的作者，其中三名元老级成员还是 Stable Diffusion 论文的作者。同时，FLUX.1 也是一个在 Stable Diffusion 3 架构上做改进的模型。不管从哪个角度，FLUX.1 都称得上是Stable Diffusion 3 的「精神续作」。秉承着此前的开源精神，FLUX.1 也在上线之始就为社区开放了源代码和模型权重。不过，配套的技术论文并没能及时发布，想要了解 FLUX.1 技术细节的用户恐怕还得等上一阵子。为了尽快搞清楚 FLUX.1 相较 Stable Diffusion 3 做了哪些改进，我直接去细读了 FLUX.1 的源码。在这篇文章中，按照惯例，我将主要从源码层面上分享 FLUX.1 中已知的科研创新，做一个官方论文发布前的前瞻解读，而不会评测 FLUX.1 的图像生成效果。

具体来说，我会介绍 FLUX.1 中的以下改动：

• 略微变动的图块化策略
• 不使用 Classifier-Free Guidance 的指引蒸馏
• 为不同分辨率图像调整流匹配噪声调度
• 用二维旋转式位置编码 (RoPE) 代替二维正弦位置编码
• 在原 Stable Diffusion 3 双流 Transformer 块后添加并行单流 Transformer 块

我会先简单介绍 FLUX.1 的官方公告及 Diffusers 版使用示例，再按照我读代码的逻辑，从熟悉整个代码框架，到深究每一处创新的代码细节，最后分享我对于 FLUX.1 科研改进上的分析。对源码不感兴趣的读者，可以跳过通读代码框架章节，或者直接阅读感兴趣的那部分改动。想看省流版文章的读者，可以直接跳到结尾看总结。

建议读者在学习 Flux.1 前熟悉 Stable Diffusion 3。欢迎参考我之前写的文章：Stable Diffusion 3 论文及源码概览。

模型简介与 Diffusers 示例脚本

在正式阅读源码前，我们先来看一下官方推文（https://blackforestlabs.ai/announcing-black-forest-labs/ ）中有关 FLUX.1 的简介，并在 Diffusers 中跑通 FLUX.1 的图像生成示例脚本。

据官方介绍，FLUX.1 是一套文生图模型。它有三个变体（variant，可以理解成结构相似或相同，但权重不同的几个模型）：

• FLUX.1 [pro]: FLUX.1 系列的最强模型，只能通过付费的 API 或者在线平台使用。
• FLUX.1 [dev]：FLUX.1 [pro] 的指引蒸馏（guidance-distilled）模型，质量与文本匹配度与原模型相近，运行时更高效。
• FLUX.1 [schnell]：为本地开发和个人使用而裁剪过的本系列最快模型。据 Diffusers 中的文档介绍，这是一个 Timestep-distilled（时间戳蒸馏）的模型，因此仅需 1~4 步就可以完成生成。无法设置指引强度。

官方对这些模型的详细介绍少之又少。FLUX.1 [dev] 用到的指引蒸馏技术似乎来自论文 On Distillation of Guided Diffusion Models，其目标是让模型直接学习 Classifier-Free Guidance (CFG) 的生成结果，使得模型一次输出之前要运行两次才能得到的指引生成结果，节约一半的运行时间。官方也没有讲 FLUX.1 [schnell] 的蒸馏细节，似乎它是从 FLUX.1 [dev] 中用扩散模型加速蒸馏手段得到的模型。因此，FLUX.1 [schnell] 不仅能一次输出有指引的生成结果，还能在极少的采样步数里完成生成。

官方推文中还说，FLUX.1 的生成神经网络基于 Stable Diffusion 3 的 MMDiT 架构和并行的 DiT 块，参数量扩大至 120 亿。生成模型是根据流匹配（flow matching）推导的扩散模型。为了提升性能与效率，模型新引入了旋转式位置编码 (RoPE) 和并行注意力层。

这段话这么长，还把并行注意力说了两遍，其实没有多少新东西。说白了，FLUX.1 就是在 Stable Diffusion 3 的基础上，加了 RoPE 和并行注意力层。官方推文到这里就没有其他有关模型细节的介绍了。FLUX.1 具体做了哪些改动，我们直接去源码里看。

FLUX.1 的官方仓库是 https://github.com/black-forest-labs/flux 。相比 Stable Diffusion 那个臃肿杂乱的 generative-models 仓库，这个仓库的代码要简洁很多。不过，我还是推荐使用 Diffusers 框架来运行 FLUX.1。

Diffusers 中运行 FLUX.1 的官方文档为 https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux 。目前（2024 年 8 月 11 日），相关代码还在 Diffusers 的在线主分支里进行开发，并没有集成进 pip 版的 Diffusers 里。因此，要在 Diffusers 中使用 FLUX，必须要从源码安装 Diffusers：

1
2
3

git clone https://github.com/huggingface/diffusers.git
cd diffusers
pip install -e .

安装完毕后，我们可以随便新建一个 python 脚本，填入以下的官方示例代码。在能够连通 Hugging Face 的环境中运行此脚本的话，脚本会自动下载模型并把生成结果保存在 image.png 中。注意，FLUX.1 的神经网络很大，显存占用极高，可能至少需要在 RTX 3090 同等级的显卡上运行。在示例代码中，我还改了一行，使用 pipe.enable_sequential_cpu_offload() 让模型把更多参数临时放到 CPU 上，避免显存不够。经测试，改了这一行后，FLUX.1 才勉强能在显存为 24G 的 RTX 3090 上运行。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
# pipe.enable_model_cpu_offload()
pipe.enable_sequential_cpu_offload()

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt,
    guidance_scale=0.0,
    num_inference_steps=4,
    max_sequence_length=256,
    height=1024,
    width=1024,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("image.png")

由于随机数是固定的，运行后，我们应该总能得到这样的图片：

通读代码框架

由于开发还没有结束，在当前 Diffusers 的 FLUX.1 源码中，我们能看到各种潦草的写法及残缺不全的文档，这让读源码变成了一项颇具趣味的挑战性任务。让我们先看一下代码的整体框架，找出 FLUX.1 相较 Stable Diffusioni 3 在代码上的改动，再来详细分析这些创新。

和 Diffusers 中的其他生成模型一样，FLUX.1 的采样算法写在一个采样流水线类里。我们可以通过示例脚本里的 FluxPipeline 类跳转到定义该类的文件 diffusers/pipelines/flux/pipeline_flux.py 里。这个文件是从 Stable Diffusion 3 的采样流水线文件 diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py 改过来的，大部分文档都没有更新。我们可以用肉眼对比两份文件的区别。

先看构造函数。Stable Diffusion 3 用了三个文本编码器，clip-vit-large-patch14, CLIP-ViT-bigG-14-laion2B-39B-b160k, t5-v1_1-xxl，而 FLUX.1 没有用第二个 CLIP 编码器，只用了另外两个文本编码器。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin):
def __init__(
        self,
        transformer: SD3Transformer2DModel,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: AutoencoderKL,
        text_encoder: CLIPTextModelWithProjection,
        tokenizer: CLIPTokenizer,
        text_encoder_2: CLIPTextModelWithProjection,
        tokenizer_2: CLIPTokenizer,
        text_encoder_3: T5EncoderModel,
        tokenizer_3: T5TokenizerFast,
    ):

class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
def __init__(
        self,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: AutoencoderKL,
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        text_encoder_2: T5EncoderModel,
        tokenizer_2: T5TokenizerFast,
        transformer: FluxTransformer2DModel,
    ):

再往下翻，我们能用火眼金睛发现 FLUX.1 的 VAE 压缩比是 16，是所有版本的 Stable Diffusion VAE 压缩比的两倍。这是为什么呢？不是增加压缩比会让 VAE 重建效果下降吗？

1
2
3
4
5
6
7
8
9
10
11

# SD3
self.vae_scale_factor = (
    2 ** (len(self.vae.config.block_out_channels) - 1) 
    if hasattr(self, "vae") and self.vae is not None else 8
)

# FLUX.1
self.vae_scale_factor = (
    2 ** (len(self.vae.config.block_out_channels)) 
    if hasattr(self, "vae") and self.vae is not None else 16
)

查看周围其他代码，我们能找到 _pack_latents，_unpack_latents 这两个方法。_pack_latents 其实就是一个图块化操作，它能把 $2 \times 2$ 个像素在通道维度上拼接到一起，而 _unpack_latents 是该操作的逆操作。原来，代码把图块化的两倍压缩比也算进 VAE 里了。这里直接把 vae_scale_factor 乘个 2 是一种非常差，歧义性极强的写法。

1
2
3
4
5
6
7

@staticmethod
def _pack_latents(latents, batch_size, num_channels_latents, height, width):
    ...

@staticmethod
def _unpack_latents(latents, height, width, vae_scale_factor):
    ...

相比 SD3, FLUX.1 将图块化操作写在了去噪网络外面。因此，SD3 的去噪网络的输入通道数是 16，和 VAE 的隐空间通道数相同；而 FLUX.1 由于把 $2 \times 2$ 个像素在通道上拼接到了一起，其去噪网络的输入通道数是 64。

1
2
3
4
5
6
7
8
9

{
  "_class_name": "SD3Transformer2DModel",
  "in_channels": 16,
}
{
  "_class_name": "FluxTransformer2DModel",
  "in_channels": 64,
}

再来看采样主方法 __call__。先看一下它的主要参数。相比之下，FLUX.1 少了一组提示词，且没有负面提示词。少一组提示词是因为少用了一个文本编码器。而没有负面提示词是因为该模型是指引蒸馏过的，在文本指引上没那么灵活。我们稍后还会看到 FLUX.1 具体是怎么利用文本指引的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

# SD3
def __call__(
    self,
    prompt: Union[str, List[str]] = None,
    prompt_2: Optional[Union[str, List[str]]] = None,
    prompt_3: Optional[Union[str, List[str]]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 28,
    timesteps: List[int] = None,
    guidance_scale: float = 7.0,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    negative_prompt_2: Optional[Union[str, List[str]]] = None,
    negative_prompt_3: Optional[Union[str, List[str]]] = None,
)

# FLUX.1
def __call__(
    self,
    prompt: Union[str, List[str]] = None,
    prompt_2: Optional[Union[str, List[str]]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 28,
    timesteps: List[int] = None,
    guidance_scale: float = 7.0,
)

之后的内容都与其他扩散模型流水线一样，代码会判断输入是否合法、给输入文本编码、随机生成初始化噪声。值得关注的是初始化噪声采样器前的一段新内容：代码会算一个 mu，并传进 retrieve_timesteps 里。这个变量最后会传到流匹配采样算法里。我们先把该改动记在心里，不看细节。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

mu = calculate_shift(
    image_seq_len,
    self.scheduler.config.base_image_seq_len,
    self.scheduler.config.max_image_seq_len,
    self.scheduler.config.base_shift,
    self.scheduler.config.max_shift,
)
timesteps, num_inference_steps = retrieve_timesteps(
    self.scheduler,
    num_inference_steps,
    device,
    timesteps,
    sigmas,
    mu=mu,
)

在去噪循环部分，FLUX.1 没有做 Classifier-Free Guidance (CFG)，而是把指引强度 guidance 当成了一个和时刻 t 一样的约束信息，传入去噪模型 transformer 中。CFG 的本意是过两遍去噪模型，一次输入为空文本，另一次输入为给定文本，让模型的输出远离空文本，靠近给定文本。而负面提示词只是一种基于 CFG 的技巧。把 CFG 里的空文本换成负面文本，就能让结果背离负面文本。但现在这个模型是一个指引蒸馏模型，指引强度会作为一个变量输入模型，固定地表示输入文本和空文本间的差距。因此，我们就不能在这个模型里把空文本换成负面文本了。

除了指引方式上的变动外，FLUX.1 的去噪网络还多了 txt_ids 和 img_ids 这两个输入。我们待会来看它们的细节。

FLUX.1 的去噪网络和 SD3 的一样，除了输入完整文本嵌入 prompt_embeds 外，依然会将池化过的短文本嵌入 pooled_prompt_embeds 输入进模型。我们现在可以猜测，FLUX.1 使用了和 SD3 类似的文本约束机制，输入了两类文本约束信息。

代码里的 /1000 是临时代码。之后所有涉及乘除 1000 的代码全可以忽略。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

for i, t in enumerate(timesteps):
    timestep = t.expand(latents.shape[0]).to(latents.dtype)

    # handle guidance
    if self.transformer.config.guidance_embeds:
        guidance = torch.tensor([guidance_scale], device=device)
        guidance = guidance.expand(latents.shape[0])
    else:
        guidance = None

    noise_pred = self.transformer(
        hidden_states=latents,
        timestep=timestep / 1000,
        guidance=guidance,
        pooled_projections=pooled_prompt_embeds,
        encoder_hidden_states=prompt_embeds,
        txt_ids=text_ids,
        img_ids=latent_image_ids,
        joint_attention_kwargs=self.joint_attention_kwargs,
        return_dict=False,
    )[0]

    # compute the previous noisy sample x_t -> x_t-1
    latents_dtype = latents.dtype
    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]

采样流水线最后会将隐空间图片解码。如前所述，由于现在图块化和反图块化是在去噪网络外面做的，这里隐空间图片在过 VAE 解码之前做了一次反图块化操作 _unpack_latents。对应的图块化操作是在之前随机生成初始噪声的 prepare_latents 方法里做的，为了节约时间我们就不去看了。

1
2
3
4
5
6
7

if output_type == "latent":
    image = latents
else:
    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
    latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
    image = self.vae.decode(latents, return_dict=False)[0]
    image = self.image_processor.postprocess(image, output_type=output_type)

接下来，我们再简单看一下去噪网络的结构。在采样流水线里找到对应类 FluxTransformer2DModel，我们能用代码跳转功能定位到文件 diffusers/models/transformers/transformer_flux.py。SD3 去噪网络类是 SD3Transformer2DModel，它位于文件 diffusers/models/transformers/transformer_sd3.py。

同样，我们先对比类的构造函数。构造函数的新参数我们暂时读不懂，所以直接跳到构造函数内部。

在使用位置编码时，SD3 用了二维位置编码类 PatchEmbed。该类会先对图像做图块化，再设置位置编码。 而 FLUX.1 的位置编码类叫 EmbedND。从官方简介以及参数里的单词 rope 中，我们能猜出这是一个旋转式位置编码 (RoPE)。

1
2
3
4
5
6
7
8
9
10
11
12

# SD3
self.pos_embed = PatchEmbed(
    height=self.config.sample_size,
    width=self.config.sample_size,
    patch_size=self.config.patch_size,
    in_channels=self.config.in_channels,
    embed_dim=self.inner_dim,
    pos_embed_max_size=pos_embed_max_size,  # hard-code for now.
)

# FLUX.1
self.pos_embed = EmbedND(dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope)

再往下看，FLUX.1 的文本嵌入类有两种选择。不设置 guidance_embeds 的话，这个类就是 CombinedTimestepTextProjEmbeddings，和 SD3 的一样。这说明正如我们前面猜想的，FLUX.1 用了和 SD3 一样的额外文本约束机制，将一个池化过的文本嵌入约束加到了文本嵌入上。

设置 guidance_embeds 的话，CombinedTimestepGuidanceTextProjEmbeddings 类应该就得额外处理指引强度了。我们待会来看这个类是怎么工作的。

1
2
3
4
5
6

text_time_guidance_cls = (
    CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
)
self.time_text_embed = text_time_guidance_cls(
    embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
)

之后函数定义了两个线性层。context_embedder 在 SD3 里也有，是用来处理文本嵌入的。但神秘的 x_embedder 又是什么呢？可能得在其他函数里才能知道了。

1
2

self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)

函数的末尾定义了两个模块列表。相比只有一种 Transformer 块的 SD3，FLUX.1 用了两种结构不同的 Transformer 块。

1
2
3
4
5
6
7
8
9
10
11
12
13

self.transformer_blocks = nn.ModuleList(
    [
        FluxTransformerBlock(...)
        for i in range(self.config.num_layers)
    ]
)

self.single_transformer_blocks = nn.ModuleList(
    [
        FluxSingleTransformerBlock(...)
        for i in range(self.config.num_single_layers)
    ]
)

我们再来看 forward 方法，看看之前看构造函数时留下的问题能不能得到解答。

forward 里首先是用 x_embedder 处理了一下输入。原本在 SD3 中，输入图像会在 pos_embed 里过一个下采样两倍的卷积层，同时完成图块化和修改通道数两件事。而现在 FLUX.1 的图块化写在外面了，所以这里只需要用一个普通线性层 x_embedder 处理一下输入通道数就行了。这样说来，变量名有个 x 估计是因为神经网络的输入名通常叫做 x。既然这样，把它叫做 input_embedder 不好吗？

1
2
3
4
5

# SD3
hidden_states = self.pos_embed(hidden_states)

# FLUX.1
hidden_states = self.x_embedder(hidden_states)

下一步是求时刻编码。这段逻辑是说，如果模型输入了指引强度，就把指引强度当成一个额外的实数约束，将其编码加到时刻编码上。具体细节都在 time_text_embed 的类里。

1
2
3
4
5
6
7
8
9
10

timestep = timestep.to(hidden_states.dtype) * 1000
if guidance is not None:
    guidance = guidance.to(hidden_states.dtype) * 1000
else:
    guidance = None
temb = (
    self.time_text_embed(timestep, pooled_projections)
    if guidance is None
    else self.time_text_embed(timestep, guidance, pooled_projections)
)

下一行是常规的修改约束文本嵌入。

1

encoder_hidden_states = self.context_embedder(encoder_hidden_states)

再之后的两行出现了一个新操作。输入的 txt_ids 和 img_ids 拼接到了一起，构成了 ids，作为旋转式位置编码的输入。

1
2

ids = torch.cat((txt_ids, img_ids), dim=1)
image_rotary_emb = self.pos_embed(ids)

此后图像信息 hidden_states 和文本信息 encoder_hidden_states 会反复输入进第一类 Transformer 块里。和之前相比，模块多了一个旋转式位置编码输入 image_rotary_emb。

1
2
3
4
5
6

encoder_hidden_states, hidden_states = block(
    hidden_states=hidden_states,
    encoder_hidden_states=encoder_hidden_states,
    temb=temb,
    image_rotary_emb=image_rotary_emb,
)

本来过了这些块后，SD3 会直接会直接返回 hidden_states 经后处理后的信息。而 FLUX.1 在过完第一类 Transformer 块后，将图像和文本信息拼接，又输入了第二类 Transformer 块中。第二类 Transformer 块的输出才是最终输出。

1
2
3
4
5
6
7
8
9

hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
for:
    hidden_states = block(
        hidden_states=hidden_states,
        temb=temb,
        image_rotary_emb=image_rotary_emb,
    )

hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]

到这里，我们就把 FLUX.1 的代码结构过了一遍。我们发现，FLUX.1 是一个基于 SD3 开发的模型。它在图块化策略、噪声调度器输入、位置编码类型、Transformer 块类型上略有改动。且由于开源的 FLUX.1 是指引蒸馏过的，该模型无法使用 CFG。[dev] 版可以以实数约束的方式设置指引强度，而 [schnell] 版无法设置指引强度。

在这次阅读中，我们已经弄懂了以下细节：

• 采样流水线会在去噪网络外面以通道堆叠的方式实现图块化。
• 指引强度不是以 CFG 的形式写在流水线里，而是以约束的形式输入进了去噪网络。

我们还留下了一些未解之谜：

• 输入进噪声采样器的 mu 是什么？
• 决定旋转式位置编码的 txt_ids 和 img_ids 是什么？
• 旋转式位置编码在网络里的实现细节是什么？
• 新的那种 Transformer 块的结构是怎么样的？

针对这些问题，我们来细读代码。

调整流匹配标准差

在采样流水线里，我们见到了这样一个神秘变量 mu。从名字中，我们猜测这是一个表示正态分布均值的变量，用来平移 (shift) 某些量的值。

1
2
3
4
5

mu = calculate_shift(...)
timesteps, num_inference_steps = retrieve_timesteps(
    ...
    mu=mu,
)

我们先看 calculate_shift 做了什么。第一个参数 image_seq_len 表示图像 token 数，可以认为是函数的自变量 x。后面四个参数其实定义了一条直线。我们可以认为 base_seq_len 是 x1, max_seq_len 是 x2，base_shift 是 y1，max_shift 是 y2。根据这两个点的坐标就可以解出一条直线方程出来。也就是说，calculate_shift 会根据模型允许的最大 token 数 4096 ($64 \times 64$) 和最小 token 数 256 ($16 \times 16$)，把当前的输入 token 数线性映射到 0.5 ~ 1.16 之间。但我们暂时不知道输出 mu 的意义是什么。

1
2
3
4
5
6
7
8
9
10
11

def calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.16,
):
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    mu = image_seq_len * m + b
    return mu

再追踪进调用了 mu 的 retrieve_timesteps 函数里，我们发现 mu 并不在参数表中，而是在 kwargs 里被传递给了噪声迭代器的 set_timesteps 方法。

1
2
3
4
5
6
7
8
9
10

def retrieve_timesteps(
    scheduler,
    num_inference_steps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
    timesteps: Optional[List[int]] = None,
    sigmas: Optional[List[float]] = None,
    **kwargs,
):
...
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)

根据流水线构造函数里的类名，我们能找到位于 diffusers/schedulers/scheduling_flow_match_euler_discrete.py 调度器类 FlowMatchEulerDiscreteScheduler。

1
2
3
4

def __init__(
    self,
    scheduler: FlowMatchEulerDiscreteScheduler,
    ...)

再找到类的 set_timesteps 方法。set_timesteps 一般是用来设置推理步数 num_inference_steps 的。有些调度器还会在总推理步数确定后，初始化一些其他变量。比如这里的流匹配调度器，会在这个方法里初始化变量 sigmas。我们可以忽略这背后的原理，仅从代码上看，输入 mu 会通过 time_shift 修改 sigmas 的值。

这里的变量命名又乱七八糟，输入 time_shift 的 sigmas 是第三个参数，而在 time_shift 里的 sigmas 是除了 self 以外的第二个参数。这是因为 Diffusers 在移植官方代码时没有取好变量名。

1
2
3
4
5
6
7
8
9
10
11
12

def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)

def set_timesteps(
    self,
    num_inference_steps: int = None,
    device: Union[str, torch.device] = None,
    sigmas: Optional[List[float]] = None,
    mu: Optional[float] = None,
):
    if self.config.use_dynamic_shifting:
        sigmas = self.time_shift(mu, 1.0, sigmas)

我们再跑出去看一下流水线里输入的 sigmas 是什么。假设总采样步数为 $T$，则 sigmas 是 $1$ 到 $\frac{1}{T}$ 间均匀采样的 $T$ 个实数。

1

sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)

现在要解读 mu 的作用就很容易了。假设 sigmas 是下标和值构成的点，我们可以测试 mu 不同的情况下， sigmas 经过 time_shift 函数形成的曲线图。

可以看出，mu=0则不修改曲线。随着 mu 增大，曲线逐渐上凸。

我对流匹配的具体细节不是很懂，只能大概猜测 mu 的作用。流匹配中，图像沿着某条路线从纯噪声运动到训练集中，标准差 sigma 用于控制不同时刻图像的不确定性。时刻为 0 时，图像为纯噪声，标准差为 1； 时刻为 1 时，图像为生成集合中的图像，标准差要尽可能趋于 0。对于中间时刻，标准差默认按照时刻线性变化。而 mu 是一个 0.5 ~ 1.16 之间的数，可能控制的是中间时刻的噪声均值。图像分辨率越大，token 越多，mu 越大，要加的噪声越重。这也符合之前 Stable Diffusion 3 论文在 Resolution-dependent shifting of timestep schedules 小节里的设计，对于分辨率越高的图像，需要加更多噪声来摧毁原图像的信号。总之，这个 mu 可能是训练的时候加的，用于给高分辨率图像加更多噪声，推理时也不得不带上这个变量。

FLUX.1 官方仓库对应部分是这样写的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

def get_lin_function(
    x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
) -> Callable[[float], float]:
    m = (y2 - y1) / (x2 - x1)
    b = y1 - m * x1
    return lambda x: m * x + b


def get_schedule(
    num_steps: int,
    image_seq_len: int,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
    shift: bool = True,
) -> list[float]:
    # extra step for zero
    timesteps = torch.linspace(1, 0, num_steps + 1)

    # shifting the schedule to favor high timesteps for higher signal images
    if shift:
        # eastimate mu based on linear estimation between two points
        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
        timesteps = time_shift(mu, 1.0, timesteps)

    return timesteps.tolist()

mu 的作用确实和高信号图片有关。但他们的设计初衷是偏移时间戳，而不是根据某种公式修改 sigma。比如原来去噪迭代 0~500 步就表示 t=0 到 t=0.5，偏移时间戳后，0~500 步就变成了 t=0 到 t=0.3。偏移时间戳使得模型能够把更多精力学习对如何对高噪声的图像去噪。

使用单流并行注意力层的 Transformer 架构

接下来的问题都和 FLUX.1 的新 Transformer 架构相关。我们先把整个网络架构弄懂，再去看旋转式位置编码的细节。

为了理清网络架构，我们来根据已知信息，逐步完善网络的模块图。首先，我们先粗略地画一个 Transformer 结构，定义好输入输出。相比 SD3，FLUX.1 多了指引强度和编号集 txt_ids, img_ids这两类输入。

接下来，我们把和 SD3 相似的结构画进来。所有 Transformer 块都是那种同时处理两类 token 的双流注意力块。输入文本的 T5 嵌入会作为文本支流进入主模型。输入文本的 CLIP 嵌入会经池化与MLP，与经过了位置编码和 MLP 的时刻编码加到一起。时刻编码会以 AdaLayerNorm 的方式修改所有层的数据规模，以及数据在输出前的尺寸与均值。

在 CombinedTimestepGuidanceTextProjEmbeddings 类中，我们能知道小文本嵌入、时刻嵌入、指引嵌入是怎么加到一起的。我们主要关心指引嵌入的有关操作。由于指引强度 guidance 和时刻 timestep 都是实数，所以 guidance_emb 的处理方式与 timesteps_emb 一模一样。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
    def __init__(self, embedding_dim, pooled_projection_dim):
        super().__init__()

        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
        self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
        self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu")

    def forward(self, timestep, guidance, pooled_projection):
        timesteps_proj = self.time_proj(timestep)
        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))  # (N, D)

        guidance_proj = self.time_proj(guidance)
        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))  # (N, D)

        time_guidance_emb = timesteps_emb + guidance_emb

        pooled_projections = self.text_embedder(pooled_projection)
        conditioning = time_guidance_emb + pooled_projections

        return conditioning

在去噪模型 FluxTransformer2DModel 的 forward 方法中，原先的图块化及二维位置编码模块被一个简单的线性层 x_embedder 取代了，现在的位置编码 image_rotary_emb 会输入进所有层中，而不是一开始和输入加在一起。

1
2
3
4

def forward(hidden_states, ...):
    hidden_states = self.x_embedder(hidden_states)
    ids = torch.cat((txt_ids, img_ids), dim=1)
    image_rotary_emb = self.pos_embed(ids)

之后，除了过 MM-DiT 块以外，文本信息还会和图像信息融合在一起，过若干个单流 Transformer 块。过了这些模块后，原来文本 token 那部分会被丢弃。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

for index_block, block in enumerate(self.transformer_blocks):
    encoder_hidden_states, hidden_states = block(
        hidden_states=hidden_states,
        encoder_hidden_states=encoder_hidden_states,
        temb=temb,
        image_rotary_emb=image_rotary_emb,
    )

hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

for index_block, block in enumerate(self.single_transformer_blocks):
    hidden_states = block(
        hidden_states=hidden_states,
        temb=temb,
        image_rotary_emb=image_rotary_emb,
    )

hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]

我们已经画完了去噪模型的结构，最后把 VAE 部分加上就比较完美了。

多数模块的细节都可以在 SD3 的论文里找到，除了 RoPE 和单流 DiT 块。我们在这一节里再仔细学习一下单流 DiT 块的结构。

根据官方介绍，FLUX.1 的 Transformer 里用到了并行 Transformer。准确来说，FLUX.1 仅在最后的单流 DiT 块里用到了并行注意力层。并行注意力层是在文章 Scaling Vision Transformers to 22 Billion Parameters 中提出的。如下图所示，这项技术很好理解，只不过是把注意力和线性层之间的串联结构变成并联结构。这样的好处是，由于数据在过注意力层前后本身就要各过一次线性层，在并联后，这些线性层和 MLP 可以融合。这样的话，由于计算的并行度更高，模型的运行效率会高上一些。

顺带一提，在 Q, K 后做归一化以提升训练稳定性也是在这篇文章里提出的。SD3 和 FLUX.1 同样用了这一设计，但用的是 RMSNorm 而不是 LayerNorm。

我们可以在 FluxSingleTransformerBlock 类里找到相关实现。代码不长，我们可以一次性读完。相比上面的示意图，Q, K, V 的投影操作被单独放进了 Attention 类里，并没有和第一个线性层融合。而做了注意力操作后，Att-out 和 MLP-out 确实是放在一起做的。attn_output 和 mlp_hidden_states 拼接了起来，一起过了 proj_out。此外，这里的归一化层还是 DiT 里的 AdaLN，模块能接收时刻编码的输入。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class FluxSingleTransformerBlock(nn.Module):
    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
        super().__init__()
        self.mlp_hidden_dim = int(dim * mlp_ratio)

        self.norm = AdaLayerNormZeroSingle(dim)
        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
        self.act_mlp = nn.GELU(approximate="tanh")
        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)

        processor = FluxSingleAttnProcessor2_0()
        self.attn = Attention(...)

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        temb: torch.FloatTensor,
        image_rotary_emb=None,
    ):
        residual = hidden_states
        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))

        attn_output = self.attn(
            hidden_states=norm_hidden_states,
            image_rotary_emb=image_rotary_emb,
        )

        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
        gate = gate.unsqueeze(1)
        hidden_states = gate * self.proj_out(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states

此处具体的注意力运算写在 FluxSingleAttnProcessor2_0 类里。跳过前面繁杂的形状变换操作，我们来看该注意力运算的关键部分。在做完了标准注意力运算 scaled_dot_product_attention 后，一般要调用 attn.to_out[0](hidden_states) 对数据做一次投影变换。但是，在这个注意力运算中，并没有对应的操作。这表明该模块确实是照着并行注意力层设计的，离开注意力的投影与 MLP 的第二个线性层融合到了一起。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

def __call__(
    self,
    attn: Attention,
    hidden_states: torch.Tensor,
    encoder_hidden_states: Optional[torch.Tensor] = None,
    attention_mask: Optional[torch.FloatTensor] = None,
    image_rotary_emb: Optional[torch.Tensor] = None,
) -> torch.Tensor:

    ...

    if image_rotary_emb is not None:
        query, key = apply_rope(query, key, image_rotary_emb)

    hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)

    hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
    hidden_states = hidden_states.to(query.dtype)

    ...
    
    return hidden_states

旋转式位置编码思想及 FLUX.1 实现

旋转式位置编码是苏剑林在 RoFormer: Enhanced Transformer with Rotary Position Embedding 中提出的一种专门为注意力计算设计的位置编码。在这篇文章中，我们来简单地了解一下旋转式位置编码的设计思想，为学习 FLUX.1 的结构做准备。

想深究旋转式位置编码的读者可以去阅读苏剑林的博文，先阅读《让研究人员绞尽脑汁的Transformer位置编码》（https://kexue.fm/archives/8130） 了解该怎么设计位置编码，再阅读《Transformer升级之路：2、博采众长的旋转式式位置编码》（https://kexue.fm/archives/8265） 了解旋转式位置编码的细节。

Transformer 仅包括注意力和全连接两种运算，这两种运算都是和位置无关的。为了让 Transformer 知道词语的前后关系，或者像素间的空间关系，就要给 Transformer 中的 token 注入某种位置信息。然而，仅仅告诉每个 token 它的绝对位置是不够好的，这样做最明显的缺点是模型无法处理训练时没有见过的长序列。比如训练集里最长的句子是 512 个 token，如果输入 600 个 token，由于模型没有见过编号超过 512 的位置编码，就不能很好地处理 512 号以后的 token。因此，我们不仅希望每个 token 知道自己的绝对位置，还希望 token 能从位置编码里知道相对位置的信息。

在提出 Transfomer 的论文中，作者给出了如下的一套正弦位置编码方案。这也是多数工作默认使用的位置编码方式。为了简化表示，我们假设输入 token 是一个二维向量，这样，每个 token 需要的位置编码也是一个二维向量。

其中，$k$ 表示第 $k$ 个 token。这样做的好处是，根据三角函数和角公式，位置编码之间可以用线性组合来表示，这种编码蕴含了一定的相对位置信息。

当我们要把二维向量拓展成 $d$ 维向量时，只需要把 $d$ 维两两打包成一组，每组用不同周期的正弦函数即可。因此，在后文中，我们也不讨论 $d$ 维的 token，只需要搞明白二维的 token 该怎么编码就行。

尽管正弦编码能表示一定的相对信息，但是，由于位置编码之间是线性关系，经过了 Transformer 中最重要的操作——注意力操作后，这种相对位置信息几乎就消失了。有没有一种位置编码方式能够让注意力计算也能知道 token 间的相对位置关系呢？

经苏剑林设计，假设每个 token 的二维位置编码是一个复数，如果用以下的公式来定义绝对位置编码，那么经过注意力计算里的求内积操作后，结果里恰好会出现相对位置关系。设两个 token 分别位于位置 $m$ 和 $n$，令给位置为 $j$ 的注意力输入 Q, K $q_j, k_j$ 右乘上 $e^{ij/10000}$的位置编码，则求 Q, K 内积的结果为：

其中，$i$ 为虚数单位，$*$ 为共轭复数，$Re$ 为取复数实部。只是为了理解方法的思想的话，我们不需要仔细研究这个公式，只需要注意到输入的 Q, K 位置编码分别由位置 $m$, $n$ 决定，而输出的位置编码由相对位置 $m-n$ 决定。这种位置编码既能给输入提供绝对位置关系，又能让注意力输出有相对位置关系，非常巧妙。

根据欧拉公式，我们可以把 $e^i$ 用一个含 $sin$ 和 $cos$ 的向量表示。由于该变换对应向量的旋转，所以这种位置编码被称为「旋转式位置编码」。在实际实现时，我们不需要复数库，只需要用两个分别含 $sin$ 和 $cos$ 的数来表示一个位置编码。也就是说，原来正弦位置编码中每个位置的编码只有一个实数，现在需要两个实数，或者说要一个二维向量。

总结一下用旋转式位置编码替换正弦位置编码后，我们在实现时应该做的改动。现在，我们不是提前算好位置编码，再加到输入上，而是先预处理好位置编码，在每次注意力 Q，K 求内积前给输入乘上。和正弦编码一样，我们会把特征长度为 $d$ 的 token 向量的分量两两分组，分别维护位置关系。但是，现在每个分量的编码由两个而不是一个实数表示。所以，在之后的代码中，我们会看到生成位置编码时，会先把 token 特征向量长度除以二，再给每组 token 生成 $2 \times 2$ 个编码，对应每组两个编码，每个编码长度为二。

我们来看一下 FLUX.1 的 Transformer 是怎么处理位置编码的。在 FluxTransformer2DModel 的 forward 方法中，我们能看到输入的 0, 1, 2, 3 这样的整数位置编码 ids 被传入了位置编码层 pos_embed 中。

1
2

ids = torch.cat((txt_ids, img_ids), dim=1)
image_rotary_emb = self.pos_embed(ids)

位置编码层类 EmbedND 定义了位置编码的具体计算方式。这个类的逻辑我们暂时跳过，直接看最后真正在算旋转式位置编码的 rope 函数。函数中，输入参数 pos 是一个 0, 1, 2, 3 这样的整数序号张量，dim 表示希望生成多长的位置编码，其值应该等于 token 的特征长度，theta 用来控制三角函数的周期，一般都是取常数 10000。我们能看到，rope 计算了输入的三角函数值，并把长度为 dim 的编码两两分组，每组有 (2, 2) 个位置编码值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
    assert dim % 2 == 0, "The dimension must be even."

    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
    omega = 1.0 / (theta**scale)

    batch_size, seq_length = pos.shape
    out = torch.einsum("...n,d->...nd", pos, omega)
    cos_out = torch.cos(out)
    sin_out = torch.sin(out)

    stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
    out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
    return out.float()


class EmbedND(nn.Module):
    def __init__(self, dim: int, theta: int, axes_dim: List[int]):
        super().__init__()
        self.dim = dim
        self.theta = theta
        self.axes_dim = axes_dim

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        n_axes = ids.shape[-1]
        emb = torch.cat(
            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
            dim=-3,
        )
        return emb.unsqueeze(1)

我们来看一下位置编码是怎么传入 Transformer 块的注意力计算的。在预处理完位置编码后，image_rotary_emb 会作为输入参数传入所有 Transformer 块，包括前面的双流块和后面的单流块。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

def forward(...):
    ids = torch.cat((txt_ids, img_ids), dim=1)
    image_rotary_emb = self.pos_embed(ids)

    ...

    for index_block, block in enumerate(self.transformer_blocks):
        encoder_hidden_states, hidden_states = block(
                        ...
                        image_rotary_emb=image_rotary_emb,
                    )

    for index_block, block in enumerate(self.single_transformer_blocks):
        hidden_states = block(
            ...
            image_rotary_emb=image_rotary_emb,
        )

位置编码 image_rotary_emb 最后会传入双流注意力计算类 FluxAttnProcessor2_0 和单流注意力计算类 FluxSingleAttnProcessor2_0。由于位置编码在这两个类中的用法都相同，我们就找 FluxSingleAttnProcessor2_0 的代码来看一看。在其 __call__ 方法中，可以看到，在做完了 Q, K 的投影变换、形状变换、归一化后，方法调用了 apply_rope 来执行旋转式位置编码的计算。而 apply_rope 会把 Q, K 特征向量的分量两两分组，根据之前的公式，模拟与位置编码的复数乘法运算。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

class FluxSingleAttnProcessor2_0:

    def __call__(
        self,
        ...
        image_rotary_emb: Optional[torch.Tensor] = None,
    ):
        query = attn.to_q(hidden_states)
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)

        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)
        if attn.norm_k is not None:
            key = attn.norm_k(key)

        if image_rotary_emb is not None:
            query, key = apply_rope(query, key, image_rotary_emb)
        
def apply_rope(xq, xk, freqs_cis):
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

这样，我们就看完了旋转式位置编码在 FLUX.1 里的实现。但是，我们还遗留了一个重要问题：在 NLP 中，句子天然有前后关系，我们按照 0, 1, 2, 3 给 token 编号就行了。而在这个模型中，既有图像 token，又有文本 token，该怎么给 token 编号呢？

图像及文本 token 的位置编号

现在，我们把目光倒回到流水线类。输入给去噪模型的序号变量有两个：text_ids，latent_image_ids。它们是怎么得到的？

1
2
3
4
5
6

noise_pred = self.transformer(
    ...
    txt_ids=text_ids,
    img_ids=latent_image_ids,
    ...
)[0]

在文本编码方法中，我们看到，text_ids 竟然只是一个全零张量。它的第一维表示 batch 大小，第二维序列长度等于文本编码 prompt_embeds 的长度，第三维序号长度为 3。也就是说，对于每一个文本 token 的每一个位置，都用 (0, 0, 0) 来表示它的位置编号。这也暗示在 FLUX.1 中，token 的位置是三维的。

1
2
3
4
5
6
7
8

def encode_prompt(
    ...
):
    ...
    text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
    text_ids = text_ids.repeat(num_images_per_prompt, 1, 1)

    return prompt_embeds, pooled_prompt_embeds, text_ids

而 latent_image_ids 主要是在 _prepare_latent_image_ids 函数里生成的。这个函数的主要输入参数是图像的高宽。根据高宽，函数会生成 (0, 0) ~ (height, width) 的二维位置坐标表格，作为位置坐标 latent_image_ids 的第二、第三维。而位置坐标的第一维全是 0。也就是说，位置为 (i, j) 的像素的位置编号为 (0, i, j)。代码里给高宽除以 2 是因为输入没有考虑 2 倍的图块化，这写得真够乱的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
    latent_image_ids = torch.zeros(height // 2, width // 2, 3)
    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]

    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape

    latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
    latent_image_ids = latent_image_ids.reshape(
        batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
    )

    return latent_image_ids.to(device=device, dtype=dtype)

def prepare_latents(...):
    height = 2 * (int(height) // self.vae_scale_factor)
    width = 2 * (int(width) // self.vae_scale_factor)

    ...

    latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
    return latents, latent_image_ids

文本位置编号 txt_ids 和 img_ids 会在第二维，也就是序列长度那一维拼接成 ids。ids 会输入给 EmbedND 类的实例 pos_embed。EmbedND 的构造函数参数中，dim 完全没有被用到，theta 控制编码的三角函数周期，axes_dim 表示位置坐标每一维的编码长度。比如 FLUX.1 的位置坐标是三维的， axes_dim 是 [16, 56, 56]，那么它就表示第一个维度用长度 16 的位置编码，后两维用长度 56 的位置编码。位置编号经 rope 函数计算得到旋转式位置编码后，会拼接到一起，最后形成 128 维的位置编码。注意，所有 Transformer 块每个头的特征数 attention_head_dim 也是 128。这两个值必须相等。

「头」指的是「多头注意力」里的「头」。头数乘上每次参与注意力运算的特征长度才等于总特征长度。由于位置编码是给 Q, K 准备的，所以位置编码的长度应该与参与注意力运算的特征长度相同。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class FluxTransformer2DModel():
def __init__(
        self,
        patch_size: int = 1,
        in_channels: int = 64,
        num_layers: int = 19,
        num_single_layers: int = 38,
        attention_head_dim: int = 128,
        num_attention_heads: int = 24,
        joint_attention_dim: int = 4096,
        pooled_projection_dim: int = 768,
        guidance_embeds: bool = False,
        axes_dims_rope: List[int] = [16, 56, 56],
    ):
    self.pos_embed = EmbedND(dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope)

def forward(...):
    ids = torch.cat((txt_ids, img_ids), dim=1)
    image_rotary_emb = self.pos_embed(ids)

class EmbedND(nn.Module):
    def __init__(self, dim: int, theta: int, axes_dim: List[int]):
        super().__init__()
        self.dim = dim
        self.theta = theta
        self.axes_dim = axes_dim

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        n_axes = ids.shape[-1]
        emb = torch.cat(
            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
            dim=-3,
        )
        return emb.unsqueeze(1)

我们来整理一下 FLUX.1 的位置编码机制。每个文本 token 的位置编号都是 (0, 0, 0)。位于 (i, j) 的像素的位置编号是 (0, i, j)。它们会生成 128 维的位置编码。编码前 16 个通道是第一维位置编号的位置编码，后面两组 56 个通道分别是第二维、第三位位置编号的位置编码。也就是说，在每个头做多头注意力运算时，特征的前 16 个通道不知道位置信息，中间 56 个通道知道垂直的位置信息，最后 56 个通道知道水平的位置信息。

乍看下来，这种位置编号方式还是非常奇怪的。所有 token 的第一维位置编号都是 0，这一维岂不是什么用都没有？

FLUX.1 旋转式位置编码原理猜测与实验

在这一节中，我将主观分析 FLUX.1 的现有源码，猜测 FLUX.1 未开源的 [pro] 版本中旋转式位置编码是怎么设置的。此外，我还会分享一些简单的相关实验结果。

已开源的 FLUX.1 为什么会出现 (0, 0, 0), (0, i, j) 这样奇怪的位置编号呢？由于现在已开源的两版模型是在 FLUX.1 [pro] 上指引蒸馏的结果，很可能原模型在指引机制，也就是和文本相关的处理机制上与现有模型不同。因此，我使用我独创的代码心理学，对现有源码进行了分析。

首先，令我感到疑惑的是采样流水线里生成位置编号的代码。latent_image_ids 一开始是一个全零张量，你写它加一个数，和直接赋值的结果不是一样的吗？为什么要浪费时间多写一个加法呢？

1
2
3

latent_image_ids = torch.zeros(height // 2, width // 2, 3)
latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]

为了确认这段代码不是 Diffusers 的开发者写的，我去看了 FLUX.1 的官方代码，发现他们的写法是一样的。在看 Diffusers 源码时，我们还看到了其他一些写得很差的代码，这些代码其实也都是从官方仓库里搬过来的。

1
2
3
4
5
6
7

def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
    ...

    img_ids = torch.zeros(h // 2, w // 2, 3)
    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)

从这些代码中，我们不难猜出开发者的心理。FLUX.1 的开发者想，我们要赶快搞一个大新闻，论文也不写了，直接加班加点准备开源。Diffusers 的开发者一看，你们这么急，我们也得搞快一点。于是他们先把 SD3 的代码复制了一遍，然后又照搬了 FLUX.1 官方仓库里的一些逻辑，直接修改 SD3 的代码。

相信大家都有这样的代码重构经历：把自己写的个人开发代码，急忙删删改改，变成能给别人看的代码。能少改一点，就少改一点。上面的代码用加法而不是赋值，就是重构的时候代码没删干净的痕迹。这说明，一开始的 img_ids 很可能不是一个全零张量，而是写了一些东西在里面。

而另一边，设置文本位置编号的官方源码里，非常干脆地写着一个全零向量。我倾向于这部分代码没有在开源时改过。

1

txt_ids = torch.zeros(bs, txt.shape[1], 3)

那么，问题就来了，这个看似全零的图像位置编号一开始是什么？它对整个位置编码的设计有什么影响？

1

img_ids = torch.zeros(h // 2, w // 2, 3)

我猜开发者设置这个变量的目的是为了区分文本和图像 token。目前，所有文本 token 的位置编号是 (0, 0, 0)，这其实不太合理，因为这种做法实际上是把所有文本 token 都默认当成位置为 (0, 0) 图像 token。为了区分文本和图像 token，应该还有其他设计。我猜最简单的方法是在第一维上做一些改动，比如令所有图像 token 的第一维都是 1。但看起来更合理的做法是对三个维度的编号都一些更改，比如给所有图像位置编号都加上一个常量 (a, b, c)。这样，图像 token 间的相对位置并不会变，而图像和文本 token 的相对位置就不同了，文本就不会默认在图像 (0, 0) 处了。从代码里的加法来看，我更倾向于认为 img_ids 原来是一个三个维度都有值的常量，且这个量或许是可以学习的。而在指引蒸馏时，位置编号的设计被简化了。

网上有人说文本位置编码全零是因为 t5 编码器自带位置编码。而在我看来，过了一个文本编码器后，文本的每个 token 已经包含所有文本的全局信息，文本 token 之间的位置编码在这里已经不重要了。重要的是文本 token 和图像 token 之间的「位置」关系，这并不能通过 t5 的位置编码来反映。

为了验证位置编码的作用，我尝试修改了图像位置编号的定义，还是跑本文开头那个测试示例。

如果把图像位置编号全置零，会得到下面的结果。这说明位置编码对结果的影响还是很大的，模型只能从位置编码处获取 token 间的相对关系。

如果把位置编号除以二，会得到下面的结果。我们能发现，图像好像变模糊了一点，且像素有锯齿化的倾向。这非常合理，因为位置编号除以二后，模型实际上被要求生成分辨率低一倍的结果。但突然又多了一些距离为 0.5 的像素，模型突然就不知道怎么处理了，最终勉强生成了这种略显模糊，锯齿现象明显的图片。注意哦，这里虽然像素间的关系不对，但图中的文字很努力地想要变得正常一点。

位置编号乘二的结果如下所示。可能模型并没有见过没有距离为 1 的图像 token 的情况，结果全乱套了。但尽管是这样，我们依然能看到图中的 “Hello World”。结合上面的结果，这说明文本指引对结果的影响还是很大的，正常的文本 token 在努力矫正图像生成结果。

位置编号乘 1.2 的结果如下所示。图像的结果还是比较正常的。这说明这套位置编码允许位置编号发生小的扰动，且模型能认识非整数的位置编号，即在模型看来，位置编号是连续的。

原图片和将位置编号第一维全置 1 的结果如下所示。如我所料，位置编号的第一维几乎没什么作用。图片只是某些地方发生了更改，整体的画面结构没有变化。

目前看下来，由于现在我们有了显式定义 token 相对位置关系的方法，要在 FLUX.1 上做一些图像编辑任务的科研，最容易想到地方就是位置编码这一块。我目前随便能想到的做法有两个：

• 直接基于位置编号做超分辨率。想办法修改位置编码的机制，使得所有图像 token 距离 2 个单位时也能某种程度上正常输出图片。以此配置反演一张低分辨率图片，得到纯噪声，重新以图像 token 距离 1 单位的正常配置来生成图片，但旧像素不对新像素做注意力，再想一些办法控制文本那部分，尽量保持旧像素输出不变，最后就能得到两倍超分辨率的结果了。inpainting 似乎也能拿类似的思路来做。
• 目前所有文本 token 的位置默认是 (0, 0)，改变文本 token 的位置编号或许能让我们精确控制文本指定的生成区域。当然，这个任务在之前的 Stable Diffusion 里好像已经被做滥了。

总结

在这篇文章中，我们围绕 FLUX.1 相对 Stable Diffusion 3 的改动，仔细阅读了 FLUX.1 在 Diffusers 中的源码。这些改动具体总结如下：

• SD3 是在去噪网络里用下采样 2 倍的卷积实现图块化，而 FLUX.1 通过把 $2 \times 2$ 个图像 token 在通道上堆叠直接实现图块化。
• FLUX.1 目前公布的两个模型都是指引蒸馏过的。我们无需使用 Classifier-Free Guidance，只要把指引强度当成一个约束条件输出进模型，就能在一次推理中得到带指定指引强度的输出。
• FLUX.1 遵照 Stable Diffusion 3 的噪声调度机制，对于分辨率越高的图像，把越多的去噪迭代放在了高噪声的时刻上。但相较 Stable Diffusion 3，似乎不仅训练时有这种设计，采样时也需要用到这种设计。
• FLUX.1 将文本的位置编号设为 (0, 0, 0)，图像的位置编号设为 (0, i, j)，之后用标准的旋转式位置编码对三个维度的编号编码，再把三组编码拼接。这种看似不太合理的位置编号设计方式或许是指引蒸馏导致的，目前从源代码中看不出原 FLUX.1 模型的位置编号设计方式。
• 在原 Stable Diffusion 的 MM-DiT 块之后，FLUX.1 将文本和图像 token 拼接，输入进了一个单流的 Transformer 块。该 Transformer 块遵照之前并行注意力层的设计，注意力层和 MLP 并联执行，在执行速度上有所提升。

FLUX.1 的总模型结构图如下所示。

作为最强开源 DiT 文生图模型，FLUX.1 狠狠打脸了拖拖拉拉刚开源没多久的 Stable Diffusion 3。可以预见，之后大家会把开发图像编辑工作的基础模型从 U-Net 版 Stable Diffusion 逐渐换成 FLUX.1。这方面的研究目前还是蓝海，值得大家投入精力研究。

FLUX.1 还是在科研上能给我们一些启示的。RoPE 都是 NLP 那边已经出了很久的工作了，直到现在才搬到图像生成这边来。我们或许能够把 NLP 或者其他视觉任务中使用的神经网络技术搬到图像生成这边来，不费什么力气地改进现有的图像生成模型。

但是，在搬运 NLP 技术中，我们也要思考如何更合理地在视觉应用中使用这些技术。文本和图像存在本质上的区别：文本是离散的，而图像是连续的。这种连续性不仅体现在图像的颜色值上，还体现在图像像素间的位置关系上。就以这里的旋转式位置编码为例，NLP 中，token 间的距离就得是整数。而在 CV 中，如果我们认为图像是一种连续信号，那么非整数的 token 距离或许也是有意义的。从文本和图像的本质区别出发，我们或许能够把 NLP 的技术更好地适配到 CV 上，而不是把 Transformer 搬过来，然后加数据一把梭。

自回归是一种根据之前已生成内容，不断递归预测下一项要生成的内容的生成模型。这种生成方式十分易懂，符合我们对生活的观察。比如我们希望模型生成一句话，第一个是「今」字，那么第二个字很可能就是「天」字。如果前三个字是「今天早」，那么第四个字就很可能是「上」。

1
2
3
4

（空）  -> 今
今      -> 天
今天    -> 早
今天早  -> 上

为了解决此问题，何恺明团队公布了论文 Autoregressive Image Generation without Vector Quantization。作者分析了目前最常见的自回归图像生成模型后，发现模型中的向量离散化 (Vector Quantization, VQ) 是拖累模型能力的罪魁祸首。作者用一些巧妙的方法绕过了 VQ，最终设计出了一种新式自回归模型。该模型在图像生成任务上表现出色，在 ImageNet 图像生成指标上不逊于最先进的图像扩散模型。在这篇博文中，我们就来学习一下这种新颖的无 VQ 自回归图像生成模型。

建议读者在阅读本文前熟悉 VQ-VAE、Transformer、DDPM 等经典工作，了解 NLP 和图像生成中连续值和离散值的概念。可以参考我之前写的文章：

轻松理解 VQ-VAE：首个提出 codebook 机制的生成模型

VQGAN 论文与源码解读：前Diffusion时代的高清图像生成模型

Transformer 论文精读

扩散模型(Diffusion Model)详解：直观理解、数学原理、PyTorch 实现

Stable Diffusion 解读（一）：回顾早期工作

知识回顾

连续值与离散值

在计算机科学中，我们既会用到连续值，也会用到离散值。比如颜色就是一个常见的连续值，我们用 0~1 之间的实数表示灰度从全黑到全白。而词元 (token) 需要用离散值表示，比如我们用 “0” 表示字母 “A”，”1” 表示 “B”, “2” 表示 “C”，并不代表 “B” 是 「’A’ 和 ‘C’ 的平均值」。离散值的数值只是用来区分不同概念的。

神经网络默认输入是连续变化的。因此，一个连续值可以直接输入进网络。而代表离散值的整数不能直接输入网络，需要先过一个嵌入层，再正常输入进网络。

自回归与类别分布

在自回归文本生成模型中，为了不断预测下一个词元，通常的做法是用一个神经网络建模下一个词元的类别分布（categorical distribution）。如下面的例子所示，所谓类别分布，就是下一步选择每一个词元的概率。有了概率分布后，我们就能用采样算法采样出下一个词元。

要训练这个预测模型也很简单。每次预测下一个词元的类别分布，其实就是一个分类任务。我们直接照着分类任务的做法，以数据集里现有句子为真值，用交叉熵损失函数就能训练这个预测模型了。

自回归图像生成

由于 Transformer 在 NLP 中的成功，大家也想用 Transformer 做图像生成。在用自回归模型生成图像时，需要考虑图像和文本的两个区别：

1. 文本是一维的，天然有先后顺序以供自回归生成。而图像是二维的，没有先后顺序。
2. 图像的颜色值是连续而非离散的。而只有离散值才能用类别分布表示。

解决问题 1 的方法很简单：没有先后顺序，我们就人工定义一个先后顺序就好了，比如从左上到右下给图像编号。

而对于问题 2，一种最简单的方式是把连续的颜色值离散化。比如将原来 0 ~ 1 的灰度值转换为「0 号灰度」、「1 号灰度」、…… 「7号灰度」。神经网络像对待词元一样对待这些灰度值，不知道它们之间的大小关系，只知道生成图像的颜色只能由这 8 种「颜色词语」构成。

向量离散化

把颜色值离散化后，我们的确可以用自回归做图像生成了。但是，由于图像的像素数比文章的词元数要多很多，这种逐像素生成方式会非常慢。为了加速自回归生成，VQ-VAE, VQGAN 等工作借由向量离散化自编码器（VQ 自编码器）实现了一个两阶段的图像生成方法：

• 训练时，先训练一个包括编码器 (encoder) 和解码器 (decoder) 两个子模型的 VQ 自编码器，再训练一个生成压缩图像的自回归模型。
• 生成时，先用自回归模型生成出一个压缩图像，再用 VQ 自编码器将其复原成真实图像。

相比普通的自编码器，VQ 自编码器有一项特点：它生成的压缩图像仅由离散值组成。这样，它就同时完成了两项任务，使得自回归模型能够高效地实现图像生成：1）将连续图像变成离散图像；2）减少要生成的像素数。

如果你还是不太理解 VQ 的作用，请先回顾 VQ-VAE 工作，再来学习这篇工作。

抛弃 VQ，拥抱扩散模型

我们来总结一下为什么要使用基于 VQ 的自回归图像生成：大家想用基于 Transformer 的自回归模型做图像生成。自回归模型在预测下一个词元/像素时，通常会用一个类别分布来建模下一项数据。由于类别分布只能描述离散数据，而图像又是连续数据，我们需要把连续像素值变成离散值。一种常用的将连续图像变成离散图像的方法是 VQ 自编码器，它既能减少图像尺寸以提高生成效率，又能将连续图像变成离散图像。

但相比普通的自编码器，如 VAE，VQ 自编码器有着一些缺点：

• VQ 自编码器很难训练
• VQ 自编码器的重建效果没有 VAE 好。比如在 Stable Diffusion 中，开发者选择了用 VAE 而不是 VQ-VAE 作为自编码器

出于抛弃 VQ 的想法，论文的作者发问道：「自回归图像生成真的需要和 VQ 绑定起来吗？」注意到，在我们刚刚阐述使用 VQ 自回归生成的动机时，用了几个「通常」、「常用」这样的非肯定词。这表明我们的这条推理链不是必然的。要取代 VQ，我们可以从两个方面入手：

1. 换一种更强力的把连续图像变成离散图像的方法
2. 从更根本处入手，不用类别分布来建模下一项数据

论文的作者选择了第二种做法：不就是建模一个像素值的分布吗？我们为什么要用死板的类别分布呢？既然扩散模型如此强大，能够拟合复杂的图像分布，那用它来拟合一个像素值的分布还不是轻轻松松？论文的核心思想也就呼之欲出了：用扩散模型而不是类别分布来建模自回归模型中下一个像素值的分布，从而抛弃自编码器里的 VQ 操作，提升模型能力。

可能读者第一次看到这个想法时会有些疑惑：扩散模型不是用来生成一整张图像的吗？它怎么建模一个像素值的分布？它和自回归模型又有什么关系？我们来多花点时间深入理解这个想法。

在文本自回归生成中，输入是已生成文本，输出是下一个词元的类别分布。

而在图像自回归生成中，输入是已生成像素，输出是下一个像素的类别分布。现在，我们希望不用类别分布，而用另一种方式，根据之前的像素生成出下一个像素。

论文作者从扩散模型中获取了灵感。扩散模型是一种强力的生成模型，它可以不根据任何信息，或根据类别、文本等信息，隐式建模训练集的图像分布，从而生成符合训练集分布的图像。既然扩散模型能够建模复杂的图像分布，那它也可以根据之前像素的信息，建模下一个像素的分布。

那么，在这种新式自回归模型里，我们可以用约束于 Transformer 输出的上下文信息的扩散模型来建模下一个像素的分布，尽管现在我们并不知道每种颜色出现的概率。

这样做的好处是，以前我们只能用离散的有限类型的颜色（准确来说是图像词元）来表示图像，现在我们能够用连续值来表示图像。模型能够更加轻松地生成内容丰富的图像。

当然，抛弃了 VQ 后，自回归模型确实不需要 VQ 自编码器来把连续图像变成离散图像了。但是，我们依然需要用自编码器来压缩图像，减少要生成的像素数。本工作依然采取了 VQ-VAE、VQGAN 那种两阶段的生成方式，只不过把 VQ 自编码器换成了用 KL loss 约束的 VAE。

训练这种扩散模型的方法很简单。在每一步训练时，我们知道上下文像素是什么，也知道当前像素的真值是什么。那么，只要以上下文像素为约束，用当前像素的真值去训练一个带约束扩散模型就行了。作者把训练这种隐式描述下一个像素值分布的误差函数称为 Diffusion Loss。

具体来说，本工作使用了最基础的带约束 DDPM 扩散模型。它和标准 DDPM 的唯一区别在于误差函数多了一个约束信息 $z$，该信息是上下文像素过 Transformer 的输出。

$t$ 时刻的噪声图像 $x_t$也是由 DDPM 加噪公式得来的。

Diffusion Loss 不仅可以用来训练表示分布的扩散模型，还可以训练前面提取上下文信息的 Transformer。由于约束信息 $z$ 来自 Transformer，可以把 Diffusion Loss 的梯度通过 $z$ 回传到 Transformer 的参数里。

扩散模型的采样公式也和 DDPM 的一样，这里不再赘述。特别地，以前的自回归模型在使用类别分布时，会用温度来控制采样的多样性。为了在扩散模型中也加入类似的温度参数，本工作参考了 Diffusion models beat GANs on image synthesis 论文的有关设计。

在具体模型超参数上，本工作的 DDPM 训练时有 1000 步，采样时有 100 步。乍看之下，DDPM 会为整个生成模型增加许多计算量，但由于只需要建模一个像素的分布，这套模型的 DDPM 可以用非常轻量级的结构。默认配置下，这套模型的 DDPM 的去噪模型是一个由 3 个残差块组成小型 MLP。每个残差块由 LayerNorm、线性层、SiLU、线性层组成。约束信息 $z$ 会和时刻 $t$ 的编码加在一起，用 DiT (Scalable diffusion models with Transformers) 里的 AdaLN 约束机制输入进 LayerNorm 层里。

套用更先进的自回归模型

仅是去掉 VQ，把 Diffusion Loss 加进标准自回归模型，并不能得到一个很好的图像生成模型。于是，作者用更加先进的一些自回归模型（掩码生成模型 Masked Gernerative Models，如 MaskGIT: Masked generative image Transformer、MAGE: Masked generative encoder to unify representation learning and image synthesis）代替标准自回归模型，极大提升了模型的生成能力。

双向注意力

在标准 Transformer 中（如下图 (a) causal 所示），每一个词元只能看到自己及之前词元的信息。这样做的好处是模型能够并行训练，串行推理。训练和推理的速度都会比较快。但是，由于每个词元看不到后面词元的信息，Transformer 提取整个句子（图像）特征的能力会下降。

而 MAE (Masked autoencoders are scalable vision learners) 论文提出了一种双向注意力机制，它可以让词元两两之间都传递信息。但是，这样模型就不能用同一个句子并行训练了，也失去了 KV cache 加速推理的手段。

如果你不太了解 Transformer 为什么是并行训练，请仔细回顾 Transformer 论文中有关自回归机制的描述。

广义自回归模型

除了双向注意力外，作者还将一些掩码生成模型的设计融合进标准自回归模型。这种广义上的自回归模型效果更好，且能缓解双向注意力导致的推理速度慢的问题。

一般来说，用图像自回归模型时，我们都是按从左到右，从上到下的顺序生成词元，如下图 (a) 所示。但是，这种顺序不一定是最合理的。

按理来说，模型应该可以通过任何顺序生成词元，这样模型学到的生成方式更加多样。更合理的生成方式应该如下图 (b) 所示，不是从左到右，从上到下给词元编号，而是随机选择一个排列给图像编号。这样就能按照随机的顺序生成图像的词元了。

而在掩码自回归生成中，模型可以一次性生成任意一个集合的词元。因此，为了加速 (b) 模型，我们可以如下图 (c) 所示，在随机给词元编号后一次生成多个词元。(b) 可以看成是 (c) 一次只预测下一个词元的特例。

Transformer 模型配置

本工作并没有给 Transformer 加入新设计，我们来确认一遍论文中介绍的 Transformer 配置。

本工作依然采取了两阶段的生成方法。第一个阶段的自编码器（又可以理解成 NLP 中的 tokenizer）来自 LDM 工作官方仓库的 VQ-16 和 KL-16 模型。前者是 VQ 自编码器（VQGAN），后者是一个加强版的 VAE。

本工作用的 Transformer 和 ViT 一样。得到图像词元后，词元会加上位置编码，且词元序列开头会附加一个 [cls] 词元，用以在类别约束生成任务里输入类别。

基于这个类别词元，本工作使用了一种特别的 Classifier-free guidance (CFG) 机制：模型用一个假类别词元来表示「类别不明」。训练时，10% 的正确类别词元被替换成了假类别词元。这样，在用扩散模型时，就可以根据标准 CFG 的做法，用正确类别和假类别实现 CFG。详情请参见论文附录 B。

在训练掩码自回归模型时，70%~100% 的词元是未知的。由于采样序列可能会很短，作者在输入序列前附加了 64 个 [cls] 词元。掩码自回归模型的其他主要设计都与 MAE 相同。

实验结果

本工作面向的是图像生成任务，主要评估 ImageNet 数据集上按类别生成的 FID 和 IS 指标。FID 越低越好，IS 越高越好。这篇工作的实验结果中有许多信息，让我们来仔细看一看这份结果。

Diffusion Loss 与广义自回归模型

论文首先展示了 Diffusion Loss、广义自回归模型这两项主要设计的优越性，如下表所示。由于图像是按类别生成的，可以用 CFG 提升模型的生成效果。为了公平比较，模型使用的 VQ 自编码器和 KL 自编码器都来自 LDM 仓库。

表格的 4 大行展示了改进自回归模型的影响，每一大行里不同 loss 的对比体现了 Diffusioin Loss 的影响。

从第一大行可以看出，Diffusion Loss 似乎对标准自回归的改进不是很明显，且这一套方法的生成能力并不出色。只有把自回归模型逐渐改进后，Diffusion Loss 的效果才能逐渐体现出来。在后几行掩码自回归模型中，Diffusion Loss 的作用还是很大的。

而对比前三大行，我们可以发现自回归模型的架构极大地提升了生成效果，且似乎将 Transformer 由 causal 改成 bidirect 的提升更加显著。

第四大行相比第三大行，提升了每次预测的词元数，主要是为了加速。这两行的对比结果表明，做了这个加速操作后，模型生成能力并没有下降多少。后续实验都是基于第四行的配置。

Diffusion Loss 适配不同的自编码器

相比原来类别分布，用 Diffusion Loss 解除了自编码器必须输出离散图像的限制。因此，目前的模型能够适配多种自编码器，如下表所示。图中 rFID 指的是图像重建任务的 FID，越低越好。这里的 VQ-16 指的是将 VQGAN 的 VQ 层当作解码器的一部分，这样 VQGAN 的编码器输出也可以看成是连续图像，和 LDM 里的做法一样。最后一行的 KL-16 是作者重新重新在 ImageNet 上训练的 VAE，而前两行的 VQ-16 和 KL-16 是在 OpenImages 上训练的。由于后文的实验都基于 ImageNet，所以后文都会用第五行那个 VAE。

首先对比一下这里 VQ-16 w/o CFG 的 FID 和上表里最后一大行 CrossEnt 的 FID。这两组实验的自编码器相同，仅有误差函数不同。将误差函数从交叉熵换成了 Diffusion Loss 后，FID 从 8.79 变成了 7.82。这一项直接对比的实验证明了不考虑自编码器的改进时，Diffusion Loss 本身的优越性。

再对比前两行，KL 的自编码器无论是图像恢复指标还是最后的生成指标都优于 VQ 的自编码器。这印证了论文开头想要抛弃 VQ 自编码器的动机：VQ 自编码器逊于 KL 自编码器。

第三、第四行展示了方法也可以兼容下采样 8 倍的自编码器。本来测试用的 ImageNet 是 $256 \times 256$ 大小的，按照一开始下采样 16 倍的配置，能得到 $16 \times 16$ 的压缩图像，即输入 Transformer 的词元序列长度为 $16 \times 16$。现在改成了下采样 8 倍后，为了兼容之前 $16 \times 16$ 的序列长度，作者把 $2 \times 2$ 个像素打包成一个词元。论文里没讲是怎么打包的，我猜测是在通道上拼接。Consistency 是另一套自编码器，作者展示这个估计是为了说明这套方法兼容性很强。

和 SOTA 图像生成模型对比

为了证明方法的优越性，论文还展示了本工作与其他 SOTA 工作在 ImageNet 图像生成任务上的定量对比结果。下表是 ImageNet $256 \times 256$ 的结果。为了方便对比，我还贴出了 DiT 论文里展示的表格（左表）。本文的模型在表里被称作 MAR。

下表是 ImageNet $512 \times 512$ 的结果。左边那张表是 EDM2 展示的结果。

从表里可以看出，本工作在 ImageNet 图像生成任务上表现很不错，超越了绝大多数模型。

图像生成速度对比

下面是不同生成模型的速度对比结果。第一张图是本论文展示的和 DiT 的对比结果。DIT 采用的扩散模型采样步数是 (50 ,75, 150, 250)。由于本工作的性能瓶颈在自回归模型而不在扩散模型上，所以本工作展示的不同采样步数由自回归步数决定。图中的自回归步数是 (8, 16, 32, 64, 128)。中间的图是 LDM 的结果，同模型不同点表示的是采样步数为 (10, 20, 50, 100, 200) 的结果。右边的表是 EDM2 的采样速度等指标。左边两张图是 ImageNet $256 \times 256$ 上的，最右边的表是 ImageNet $512 \times 512$ 上的。

由于不同图表的采样速度指标不太一样，我们将指标统一成每秒生成的图像。从第一张图的对比可以看出，DiT 最快也是一秒 2.5 张图像左右，而 MAR 又快又好，默认（自回归步数 64）一秒生成 3 张图左右。同时，通过 MAR 和有 kv cache 加速的标准 AR 的对比，我们能发现 MAR 在默认自回归步数下还是比标准 AR 慢了不少。

我们再看中间 LDM 的速度。我们观察一下最常使用的 LDM-8。如果是令 DDIM 步数为 20 （第二快的结果）的话，LDM-8 的生成速度在一秒 16 张图像左右，还是比 MAR 快很多。DDIM 步数取 50 时也会比 MAR 快一些。

最后看右边较新的图像扩散模型 EDM2 的速度。由于这个是在 $512 \times 512$ 的图片上测试的，和前面的速度相比时大概要乘个 4。哪怕是最大的 XXL 模型，在有 guidance 时，生成速度也是 2 张图片每秒。换算到 $256 \times 256$ 上约 8 张图片每秒，还是比 MAR 快。

总结

自回归图像生成中的向量离散化和类别分布必须同时使用。为了去除表现较差的向量离散化操作，本工作的作者重新用扩散模型建模了自回归中下一个图像词元的分布，从而提升了模型的生成能力。由于标准自回归模型生成能力有限，为了进一步提升模型，作者又引入了最新的掩码自回归模型。最终的模型在 ImageNet 图像生成指标上取得了几乎最顶尖的结果。

以上是论文的叙述逻辑。但掩码自回归那一块应该是之前工作的研究成果，这篇文章实际上就是把新提出的 diffusion loss 用到了掩码自回归上，把本来在 ImageNet 上生成能力尚可的掩码自回归推到了最前列。

这篇文章在科研上的最大创新是打破了大家在图像自回归上的固有思维，认为必须用离散词元，必须用类别分布。但仔细一想，建模一个分布的方法其实许许多多。随便把另一种生成完整图像的模型用到生成一个像素上，就能取代之前的类别分布，得到更好的图像生成结果。这篇文章用简单的 DDPM 只是为了验证这个想法的可行性，用更复杂的模型或许能有更好的结果，但用 DDPM 做验证就足够了。之后肯定会有各种后续工作，研究如何用更好的模型来建模本框架中一个像素值的分布。

反过来想，这篇文章也在提醒我们，扩散模型并不只是可以用来生成图像，它的本质是建模一个分布。如果某个模型中间需要建模一个简单的分布的话，都可以尝试用 DDPM。

相比其科研创新，这篇文章在 ImageNet 图像生成指标的成就反而没有那么耀眼了。本工作在 ImageNet 的 FID 等指标上取得了几乎最优的结果，战胜了多数最强的扩散模型，有望将大家的科研眼光从扩散模型移到自回归上。但由于自回归本身步数较多，且每一步要在 Transformer 里做完整的注意力操作，这种方法的速度还是比扩散模型要慢一点。

目前 GitHub 上已有本工作的复现：https://github.com/lucidrains/autoregressive-diffusion-pytorch 。

Diffusers 库为社区用户提供了多种扩散模型任务的训练脚本。每个脚本都平铺直叙，没有多余的封装，把训练的绝大多数细节都写在了一个脚本里。这种设计既能让入门用户在不阅读源码的前提下直接用脚本训练，又方便高级用户直接修改脚本。

可是，这种设计就是最好的吗？关于训练脚本的最佳设计风格，社区用户们往往各执一词。有人更喜欢更贴近 PyTorch 官方示例的写法，而有人会喜欢用 PyTorch Lightning 等封装度高、重复代码少的库。而在我看来，选择哪种风格的训练脚本，确实是个人喜好问题。但是，在开始使用训练脚本之前，我们要从细节入手，理解训练脚本到底要做哪些事。学懂了之后，不管是用别人的训练库，还是定制适合自己的训练脚本，都是很轻松的。不管怎么说，Diffusers 的这种训练脚本是一份很好的学习素材。

当然，我在用 Diffusers 的训练脚本时，发现一旦涉及多类任务的训练，比如既要能训练 Stable Diffusion，又要能训练 VAE，那么这份脚本就会用起来比较困难，而写两份训练脚本又会有很大的冗余。Diffusers 的训练脚本依然有改进的空间。

在这篇文章中，我会主要面向想系统性学习扩散模型训练框架的读者，先详细介绍 Diffusers 官方训练脚本，再分享我重构训练脚本的过程，使得脚本能够更好地兼容多类模型的训练。文章的末尾，我会展示几个简单的扩散模型训练实例。

在阅读本文时，建议大家用电脑端，一边看源代码一边读文章。「官方训练脚本细读」一节细节较多，初次阅读时可以快速浏览，看完「训练脚本内容总结」中的流程图，再回头仔细看一遍。

准备源代码

我们将以最简单的 DDPM 官方训练脚本 examples/unconditional_image_generation/train_unconditional.py 为例，学习训练脚本的通用写法。examples 文件夹在位于 Diffusers 官方 GitHub 仓库中，用 pip 安装的 Diffusers 可能没有这个文件夹，最好是手动 clone 官方仓库，再在本地查看这个文件夹。使用 Diffusers 训练时，可能还要安装其他库。官方在不同的训练教程里给了不同的安装指令，建议大家都安装上。

1
2
3

cd examples/text_to_image
pip install -r requirements.txt
pip install diffusers[training]

我为本教程准备的脚本在仓库 https://github.com/SingleZombie/DiffusersExample 中。请 clone 这个仓库，再切换到 TrainingScript 目录下。train_official.py 是原官方训练脚本 train_unconditional.py，train_0.py 是第一次修改后的训练脚本
，train_1.py 是第二次修改后的训练脚本。

官方训练脚本细读

先拉到文件的最底部，我们能在这找到程序的入口。在 parse_args 函数中，脚本会用 argparse 库解析命令行参数，并将所有参数保存在 args 里。args 会传进 main 函数里。稍后我们看到所有 args. 打头的变量调用，都表明该变量来自于命令行参数。

1
2
3

if __name__ == "__main__":
    args = parse_args()
    main(args)

接着，我们正式开始学习训练主函数。一开始，函数会配置 accelerate 库及日志记录器。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

logging_dir = os.path.join(args.output_dir, args.logging_dir)
accelerator_project_config = ProjectConfiguration(
    project_dir=args.output_dir, logging_dir=logging_dir)

# a big number for high resolution or big dataset
kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=7200))
accelerator = Accelerator(...)

if args.logger == "tensorboard":
    if not is_tensorboard_available():
        ...

elif args.logger == "wandb":
    if not is_wandb_available():
        ...
    import wandb

在配置日志的中途，函数插入了一段修改模型存取逻辑的代码。为了让我们阅读代码的顺序与实际运行顺序一致，我们等待会用到了这段代码时再回头来读。

1
2
3
4
5
6

# `accelerate` 0.16.0 will have better support for customized saving
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
    def save_model_hook(models, weights, output_dir):
        ...
    def load_model_hook(models, input_dir):
        ...

跳过上面的代码，还是日志配置。

1
2
3
4
5
6
7
8
9
10
11
12
13

# Make one log on every process with the configuration for debugging.
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
logger.info(accelerator.state, main_process_only=False)
if accelerator.is_local_main_process:
    datasets.utils.logging.set_verbosity_warning()
    diffusers.utils.logging.set_verbosity_info()
else:
    datasets.utils.logging.set_verbosity_error()
    diffusers.utils.logging.set_verbosity_error()

之后其他版本的训练脚本会有一段设置随机种子的代码，我们给这份脚本补上。

1
2
3

# If passed along, set the training seed now.
if args.seed is not None:
    set_seed(args.seed)

接着，函数会创建输出文件夹。如果我们想把模型推送到在线仓库上，函数还会创建一个仓库。

这段代码还出现了一行比较重要的判断语句：if accelerator.is_main_process:。在多卡训练时，只有主进程会执行这个条件语句块里的内容。该判断在并行编程中十分重要。很多时候，比如在输出、存取模型时，我们只需要让一个进程执行操作就行了。这个时候就要用到这行判断语句。

1
2
3
4
5
6
7

# Handle the repository creation
if accelerator.is_main_process:
    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)

    if args.push_to_hub:
        repo_id = create_repo(...).repo_id

准备完辅助工具后，函数开始准备模型。输入参数里的 model_config_name_or_path 表示预定义的模型配置文件。如果该配置文件不存在，则函数会用默认的配置创建一个 DDPM 的 U-Net 模型。在写我们自己的训练脚本时，我们需要在这个地方初始化我们需要的所有模型。比如训练 Stable Diffusion 时，除了 U-Net，需要在此处准备 VAE、CLIP 文本编码器。

1
2
3
4
5
6

# Initialize the model
if args.model_config_name_or_path is None:
    model = UNet2DModel(...)
else:
    config = UNet2DModel.load_config(args.model_config_name_or_path)
    model = UNet2DModel.from_config(config)

这份脚本还帮我们写好了维护 EMA（指数移动平均）模型的功能。EMA 模型用于存储模型可学习的参数的局部平均值。有时 EMA 模型的效果会比原模型要好。

1
2
3
4
5
6
7

# Create EMA for the model.
if args.use_ema:
    ema_model = EMAModel(
        model.parameters(),
        model_cls=UNet2DModel,
        model_config=model.config,
        ...)

此处函数还会根据 accelerate 配置自动设置模型的精度。

1
2
3
4
5
6
7

weight_dtype = torch.float32
if accelerator.mixed_precision == "fp16":
    weight_dtype = torch.float16
    args.mixed_precision = accelerator.mixed_precision
elif accelerator.mixed_precision == "bf16":
    weight_dtype = torch.bfloat16
    args.mixed_precision = accelerator.mixed_precision

函数还会尝试启用 xformers 来提升 Attention 的效率。PyTorch 在 2.0 版本也加入了类似的 Attention 优化技术。如果你的显卡性能有限，且 PyTorch 版本小于 2.0，可以考虑使用 xformers。

1
2
3

if args.enable_xformers_memory_efficient_attention:
    if is_xformers_available():
        ...

准备了 U-Net 后，函数会准备噪声调度器，即定义扩散模型的细节。

注意，扩散模型不是一个神经网络，而是一套定义了加噪、去噪公式的模型。扩散模型中需要一个去噪模型来去噪，去噪模型一般是一个神经网络。

1
2
3
4
5
6
7

# Initialize the scheduler
accepts_prediction_type = "prediction_type" in set(
    inspect.signature(DDPMScheduler.__init__).parameters.keys())
if accepts_prediction_type:
    noise_scheduler = DDPMScheduler(...)
else:
    noise_scheduler = DDPMScheduler(...)

准备完所有扩散模型组件后，函数开始准备其他和训练相关的模块。其他版本的训练脚本会在这个地方加一段缓存梯度和自动放缩学习率的代码，我们给这份脚本补上。

1
2
3
4
5
6
7

if args.gradient_checkpointing:
    unet.enable_gradient_checkpointing()

if args.scale_lr:
    args.learning_rate = (
        args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
    )

函数先准备的训练模块是优化器。这里默认使用的优化器是 AdamW。

1
2
3
4
5
6
7

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

函数随后会准备训练集。这个脚本用 HuggingFace 的 datasets 库来管理数据集。我们既可以读取在线数据集，也可以读取本地的图片文件夹数据集。自定义数据集的方法可以参考 https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder 。

1
2
3
4
5
6
7
8
9
10
11
12

if args.dataset_name is not None:
    dataset = load_dataset(
        args.dataset_name,
        args.dataset_config_name,
        cache_dir=args.cache_dir,
        split="train",
    )
else:
    dataset = load_dataset(
        "imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")
    # See more about loading custom images at
    # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder

有了数据集后，函数会继续准备 PyTorch 的 DataLoader。在这一步中，除了定义 DataLoader 外，我们还要编写数据预处理的方法。下面这段代码的编写顺序和执行顺序不同，我们按执行顺序来整理一遍下面的代码：

1. 将预定义的预处理函数传给数据集对象 dataset.set_transform(transform_images)。在使用数据集里的数据时，才会调用这个函数预处理图像。
2. 使用 PyTorch API 定义 DataLoader。train_dataloader = ...
3. 每次用 DataLoader 获取数据时，一个数据词典 examples 会被传入预处理函数 transform_images。examples 里既包含了图像数据，也包含了数据的各种标签。而对于无约束图像生成任务，我们只需要图像数据，因此可以直接通过词典的 "image" 键得到 PIL 格式的图像数据。用 convert("RGB") 把图像转成三通道后，该 PIL 图像会被传入预处理流水线。
4. 图像预处理流水线 augmentations 是用 Torchvision 里的 transform API 定义的。默认的流水线包括短边缩放至指定分辨率、按分辨率裁剪、随机反转、归一化。
5. 处理过的数据会被存到词典的 "input" 键里。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# Preprocessing the datasets and DataLoaders creation.
augmentations = transforms.Compose(
    [
        transforms.Resize(
            args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(
            args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
        transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)

def transform_images(examples):
    images = [augmentations(image.convert("RGB"))
                for image in examples["image"]]
    return {"input": images}

logger.info(f"Dataset size: {len(dataset)}")

dataset.set_transform(transform_images)
train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)

在准备工作的最后，函数会准备学习率调度器。

1
2
3
4
5
6
7

# Initialize the learning rate scheduler
lr_scheduler = get_scheduler(
    args.lr_scheduler,
    optimizer=optimizer,
    num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
    num_training_steps=(len(train_dataloader) * args.num_epochs),
)

准备完了所有模块，函数会调用 accelerate 库来把所有模块变成适合并行训练的模块。

1
2
3
4
5
6

model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, lr_scheduler
)

if args.use_ema:
    ema_model.to(accelerator.device)

之后函数还会用 accelerate 库配置训练日志。默认情况下日志名 run 由当前脚本名决定。如果不想让之前的日志被覆盖的话，可以让日志名 run 由当前的时间决定。

1
2
3

if accelerator.is_main_process:
    run = os.path.split(__file__)[-1].split(".")[0]
    accelerator.init_trackers(run)

马上就要开始训练了。在此之前，函数会准备全局变量并记录日志。注意，这里函数会算一次总的 batch 数，它由输入 batch 数、进程数（显卡数）、梯度累计步数共同决定。梯度累计是一种用较少的显存实现大 batch 训练的技术。使用这项技术时，训练梯度不会每步优化，而是累计了若干步后再优化。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

total_batch_size = args.train_batch_size * \
    accelerator.num_processes * args.gradient_accumulation_steps
num_update_steps_per_epoch = math.ceil(
    len(train_dataloader) / args.gradient_accumulation_steps)
max_train_steps = args.num_epochs * num_update_steps_per_epoch

logger.info("***** Running training *****")
logger.info(f"  Num examples = {len(dataset)}")
logger.info(f"  Num Epochs = {args.num_epochs}")
logger.info(
    f"  Instantaneous batch size per device = {args.train_batch_size}")
logger.info(
    f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
logger.info(
    f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
logger.info(f"  Total optimization steps = {max_train_steps}")

global_step = 0
first_epoch = 0

在开始训练前，如果设置了 args.resume_from_checkpoint，则函数会读取之前训练过的权重。负责读取训练权重的函数是 load_state。

1
2
3
4
5
6
7
8
9
10
11
12
13

if args.resume_from_checkpoint:
    if args.resume_from_checkpoint != "latest":
        path = ..
    else:
        # Get the most recent checkpoint
        ...

    if path is None:
        ...
    else:
        accelerator.load_state(os.path.join(args.output_dir, path))
        accelerator.print(f"Resuming from checkpoint {path}")
        ...

在每个 epoch 中，函数会重置进度条。接着，函数会进入每一个 batch 的训练迭代。

1
2
3
4
5
6
7

# Train!
for epoch in range(first_epoch, args.num_epochs):
    model.train()
    progress_bar = tqdm(total=num_update_steps_per_epoch,
                        disable=not accelerator.is_local_main_process)
    progress_bar.set_description(f"Epoch {epoch}")
    for step, batch in enumerate(train_dataloader):

如果是继续训练的话，训练开始之前会更新当前的步数 step。

1
2
3
4
5

# Skip steps until we reach the resumed step
if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
    if step % args.gradient_accumulation_steps == 0:
        progress_bar.update(1)
    continue

训练的一开始，函数会从数据的 "input" 键里取出图像数据。此处的键名是我们之前在数据预处理函数 transform_images 里写的。

1

clean_images = batch["input"].to(weight_dtype)

之后函数会设置扩散模型训练中的其他变量，包含随机噪声、时刻。由于本文的重点并不是介绍扩散模型的原理，这段代码我们就快速略过。

1
2
3
4

noise = torch.randn(...)
timesteps =...
noisy_images = noise_scheduler.add_noise(
    clean_images, noise, timesteps)

接下来，函数会用去噪网络做前向传播。为了让模型能正确累计梯度，我们要用 with accelerator.accumulate(model): 把模型调用与反向传播的逻辑包起来。在这段代码中，我们会先得到模型的输出 model_output，再根据扩散模型得到损失函数 loss，最后用 accelerate 库的 API accelerator 代替原来 PyTorch API 来完成反向传播、梯度裁剪，并完成参数更新、学习率调度器更新、优化器更新。

1
2
3
4
5
6
7
8
9
10
11
12
13

with accelerator.accumulate(model):
    # Predict the noise residual
    model_output = model(noisy_images, timesteps).sample

    loss = ...

    accelerator.backward(loss)

    if accelerator.sync_gradients:
        accelerator.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    lr_scheduler.step()
    optimizer.zero_grad()

确保一步训练结束后，函数会更新和步数相关的变量。

1
2
3
4
5

if accelerator.sync_gradients:
    if args.use_ema:
        ema_model.step(model.parameters())
    progress_bar.update(1)
    global_step += 1

在这个地方，函数还会尝试保存模型。默认情况下，每 args.checkpointing_steps 步保存一次中间结果。确认要保存后，函数会算出当前的保存点名称，并根据最大保存点数 checkpoints_total_limit 决定是否要删除以前的保存点。做完准备后，函数会调用 save_state 保存当前训练时的所有中间变量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

f accelerator.is_main_process:
    if global_step % args.checkpointing_steps == 0:
        if args.checkpoints_total_limit is not None:
            checkpoints = os.listdir(args.output_dir)
            checkpoints = [
                d for d in checkpoints if d.startswith("checkpoint")]
            checkpoints = sorted(
                checkpoints, key=lambda x: int(x.split("-")[1]))

            if len(checkpoints) >= args.checkpoints_total_limit:
                ...

            save_path = os.path.join(
            args.output_dir, f"checkpoint-{global_step}")
            accelerator.save_state(save_path)
            logger.info(f"Saved state to {save_path}")

在这个地方，主函数开头设置的存取模型回调函数终于派上用场了。在调用 save_state 时，会自动触发下面的回调函数来保存模型。如果不加下面的代码，所有模型默认会以 .safetensor 的形式存下来。而用了下面的代码后，模型能够被 save_pretrained 存进一个文件夹里，就像其他标准 Diffusers 模型一样。

这里的输入参数 models 来自于之前的 accelerator.prepare，感兴趣可以去阅读文档或源码。

1
2
3
4
5
6
7
8
9
10
11

def save_model_hook(models, weights, output_dir):
    if accelerator.is_main_process:
        if args.use_ema:
            ema_model.save_pretrained(
                os.path.join(output_dir, "unet_ema"))

        for i, model in enumerate(models):
            model.save_pretrained(os.path.join(output_dir, "unet"))

            # make sure to pop weight so that corresponding model is not saved again
            weights.pop()

与上面的这段代码对应，脚本还提供了读取文件的回调函数。它会在继续中断的训练后调用 load_state 时被调用。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

def load_model_hook(models, input_dir):
    if args.use_ema:
        load_model = EMAModel.from_pretrained(
            os.path.join(input_dir, "unet_ema"), UNet2DModel)
        ema_model.load_state_dict(load_model.state_dict())
        ema_model.to(accelerator.device)
        del load_model

    for i in range(len(models)):
        # pop models so that they are not loaded again
        model = models.pop()

        # load diffusers style into model
        load_model = UNet2DModel.from_pretrained(
            input_dir, subfolder="unet")
        model.register_to_config(**load_model.config)

        model.load_state_dict(load_model.state_dict())
        del load_model

两个回调函数需要用下面的代码来设置。

1
2

accelerator.register_save_state_pre_hook(save_model_hook)
accelerator.register_load_state_pre_hook(load_model_hook)

回到最新的代码处。训练迭代的末尾，脚本会记录当前步的日志。

1
2
3
4
5

logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
if args.use_ema:
    logs["ema_decay"] = ema_model.cur_decay_value
progress_bar.set_postfix(**logs)
accelerator.log(logs, step=global_step)

执行完了一个 epoch 后，脚本调用 accelerate API 保证所有进程均训练完毕。

1
2

progress_bar.close()
accelerator.wait_for_everyone()

此处脚本可能会在主进程中验证模型或保存模型。如果当前是最后一个 epoch，或者达到了配置指定的验证/保存时刻，脚本就会执行验证/保存。

1
2
3
4
5
6

if accelerator.is_main_process:
    if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
        ...

    if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
        ...

脚本默认的验证方法是随机生成图片，并用日志库保存图片。生成图片的方法是使用标准 Diffusers 采样流水线 DDPMPipeline。由于此时模型 model 可能被包裹成了一个用于多卡训练的 PyTorch 模块，需要用相关 API 把 model 解包成普通 PyTorch 模块 unet。如果使用了 EMA 模型，为了避免对 EMA 模型的干扰，此处需要先保存 EMA 模型参数，采样结束再还原参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
    unet = accelerator.unwrap_model(model)
    if args.use_ema:
        ema_model.store(unet.parameters())
        ema_model.copy_to(unet.parameters())

    pipeline = DDPMPipeline(
        unet=unet,
        scheduler=noise_scheduler,
    )

    generator = torch.Generator(device=pipeline.device).manual_seed(0)
    # run pipeline in inference (sample random noise and denoise)
    images = pipeline(...).images

    if args.use_ema:
        ema_model.restore(unet.parameters())

    # denormalize the images and save to tensorboard
    images_processed = (images * 255).round().astype("uint8")

    if args.logger == "tensorboard":
        ...
    elif args.logger == "wandb":
        ...