在这篇博文中,我会先回顾与 VAR 密切相关的早期工作 VQVAE 和 VQGAN,再介绍论文的方法细节与实验结果,最后分享我对该工作的测试结果与原理探究。在读 VAR 论文时,我发现有个地方的设计存在缺陷。相关实验结果表明, VAR 论文并没有完整地分析出这套方法有效的原因。欢迎大家仔细阅读这一部分并提出自己的思考与见解。
VAR 用了一种更加高级的做法:用残差金字塔来表示这些隐空间特征。我们先来回顾一下拉普拉斯金字塔这一经典图像处理算法。我们知道,图像每次下采样的时候,都会损失一些信息。既然如此,我们可以将一张高分辨率的图像表示为一张低分辨率的图像及其在各个分辨率下采样后的信息损失。如下图所示,最右侧的一列表示拉普拉斯金字塔的输出。
VAR 的方法部分我们看得差不多了,现在来简单看一下实验部分。论文宣称 VAR 在图像生成实验和参数扩增实验上都取得了不错的成果。特别地,VAR 的拟合能力胜过了 DiT,生成速度是 DiT 的 45 倍以上。我们就主要看一下 VAR 在 ImageNet $256 \times 256$ 图像生成上的实验结果。以下是论文中的表格。我同时还附上了何恺明团队的 MAR 工作(Autoregressive Image Generation without Vector Quantization)的实验结果。
先比一下 DiT 和 VAR。先看速度,不管是多大的模型,DiT 的速度都远远慢于 VAR。再看以 FID 为代表的图像拟合指标。VAR 在参数量为 600M 左右时并没有 DiT 效果好。但继续增加参数量后,DiT 的 FID 没有变好的趋势,而 VAR 的 FID 一直在降。最终 VAR 的 FID 甚至超过了 ImageNet 的验证集,可以认为 FID 再低的也意义不大了。
再比一下 MAR 和 VAR。MAR 的刷指标能力更加恐怖,943M 的模型就能有 1.55 的 FID。但根据 MAR 论文,其速度是 DiT-XL 的 5 倍左右,也就是说 VAR 还是比 MAR 快,是 MAR 速度的 9 倍左右。
ImageNet 图像生成已经被各个模型刷到头了。FID 结果能说明 VAR 的拟合能力很强,最起码不逊于 DiT。但在更有挑战性的文生图任务上,VAR 的效果还有待验证。另外,虽然刷指标的时候 DiT 用了 250 步采样,但实际用起来的时候一般就是采样 20 步。如果算上蒸馏的话,采样步数能缩小到 4 步。加上这些加速技巧的话,VAR 不见得会比 DiT 快。
VAR 各尺度生成结果
看完了论文的主要内容,我来分享一下我对 VAR 的一些理论分析与实验结果。
先看一下随机采样结果。我用的是最大的 d=30 的 VAR 模型。在官方采样脚本的默认配置下,两个随机种子 (0, 15) 的输出如下所示。用到的图像类别为火山、灯塔、老鹰、喷泉,每个类别的图各生成了两张。图像的生成速度很快,一秒就生成了全部 8 张图片。
此前,以经典工作 VQGAN 为代表的图像自回归生成模型无论在速度上还是图像质量上都不尽如人意。究其原因,下一个图像词元预测的建模方式既不够合理,也拖慢了生成速度。为此,VAR 提出一种新式自回归策略:将词元图像拆分成多个尺度,通过下一尺度预测实现图像生成。为了兼容这一设计,VAR 对 VQGAN 的自编码器和 Transformer 都进行了修改:自编码器能够将图像编码成多尺度的残差词元图像,而 Transformer 同时输出同一尺度每个词元的独立分布。实验表明,VAR 在 ImageNet 图像生成指标上超越了以 DiT 为代表的扩散模型,且生成速度至少比 DiT 快 45 倍。另外,还有实验表明 VAR 符合扩增定律:增加参数量即可提升模型性能。
我个人认为,和其他前沿生成模型一样,VAR 在 ImageNet 上的表现已经满分了。它能否完成更困难的图像生成认为还有待验证。最近字节跳动发布了 VAR 的文生图版本:Infinity,但这个模型还没有开源。我们可以持续关注 VAR 的各个后续工作。VAR 的生成速度也没有比 DiT 快上那么多,通过减小采样步数,再加上模型蒸馏,DiT 不会比 VAR 慢。当然,VAR 或许也存在进一步加速的可能,只是相关研究暂时没有扩散模型那么多。
VAR 的数学模型是存在缺陷的:词元图的分布不应该等于词元间的独立分布的乘积。最起码论文里没有任何相关分析(用了类似做法的 MAR 论文也没有分析)。通过一些简单的生成实验,我们发现由于 VAR 在其他设计上提升了输出图像的连续性,哪怕同一尺度的词元间是独立采样,甚至是随机均匀采样,模型的输出质量也不会太差。我们需要通过更深入的实验来挖掘 VAR 的生效原理。
我觉得如果一个科研工作能够解释清楚 VAR 中哪些模块起到了最主要的作用,并取其精华,去其糟粕,提出一个更好的生成模型,那这会是一个很不错的工作。我觉得能够探索的方向有:
VAR 的前几个尺度的词元图是最重要的。能不能用更好的方式,比如用扩散模型,来生成前几个尺度的图像,而更大尺度的词元图用一个比 Transformer 更高效的模型来生成。这样模型的质量和效率能进一步提升。
VAR 还是用了 VQ 自编码器。无论怎么样,VQ 操作都会降低模型的重建质量。但另一方面,VQ 也能起到规范解码器输入的作用。究竟我们能不能把 VQ 自编码器换成精度更高的 VAE 呢?换了之后怎么设计多尺度编码呢?
将图像生成拆解成从低分辨率到高分辨率是一种很常见的思想。基于扩散模型,有多种方式来应用这种思想。一种比较直接的方式是显式将图像生成分解成生成最低分辨率的图像和多轮超分辨率,代表工作是 Cascaded Diffusion Models for High Fidelity Image Generation;另一种更加巧妙的方式是将图像上采样和扩散模型的去噪同时进行,代表工作是 f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation。本文的多尺度设计和 f-DM 非常相似,我会在文末详细分析二者的区别。
和这篇工作非常相关的早期工作是苹果在 2022 发表的 f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation。f-DM 将扩散模型的加噪推广到了降采样、模糊等其他退化策略上。降采样版的 f-DM 有非常多的设计和本工作很像。苹果该团队次年发表的 Matryoshka Diffusion Models 也用到了按分辨率逐次生成的设计。
将拉普拉斯金字塔融入扩散模型
拉普拉斯金字塔是一种图像表示方法,它把图像按频率成分拆成几张分辨率不同的图像,分辨率越低的图像表示频率越低的图像成分。我们直接通过下面的例子学习它的原理。假如 x 是原图,那么 x(3)=down(down(x)),x(2)=down(x)-up(x(3)), x(1)=x-up(down(x))。对 x(1), x(2), x(3) 求加权和就可以还原输入图像。
Lumina-Next,前沿扩散 Transformer (Diffusion Transformer, DIT) 模型,采用了长度外推技术:Lumina-Next : Making Lumina-T2X Stronger and Faster with Next-DiT (https://arxiv.org/abs/2406.18583)
近几年和 NTK 理论比较相关的论文叫做 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains。这篇论文用 NTK 理论分析了 NeRF 这种以位置坐标为输入的 MLP 需要位置编码的原因,并将这类位置编码归纳为「傅里叶特征」。
在这篇文章里,我主要基于论文 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains (后文简称为「傅里叶特征论文」),介绍傅里叶特征这一概念。为了讲清这些理论的发展脉络,我会稍微讲一下 NTK 等理论概念。介绍完傅里叶特征后,我还会讲解它在其他方法中的应用。希望读完本文后,读者能够以这篇论文为基点,建立一个有关位置编码原理的知识网络,以从更深的层次来思考新的科研方向。
这种连续数据有什么好处呢?我们知道,计算机都是以离散的形式来存储数据的。比如,我们会把图像拆成一个个像素,每个像素存在一块内存里。对于图像这种二维数据,计算机的存储空间还勉强够用。而如果想用密集的离散数据表达更复杂的数据,比如 3D 物体,计算机的容量就捉襟见肘了。但如果用一个 MLP 来表达 3D 物体的话,我们只需要存储 MLP 的参数,就能获取 3D 物体在任何位置的信息了。
这就是经典工作神经辐射场 (Neural Radiance Field, NeRF) 的设计初衷。NeRF 用一个 MLP 拟合 3D 物体的属性,其输入输出如下图所示。我们可以用 MLP 学习每个 3D 坐标的每个 2D 视角处的属性(这篇文章用的属性是颜色和密度)。根据这些信息,利用某些渲染算法,我们就能重建完整的 3D 物体。
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
classFourierFeature(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) defforward(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 == 0or epoch == n_loops - 1: viz_image(output[0]) print(loss.item()) prev_output = output
根据之前的研究 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}$,最后学出来的网络误差都差不多。问题的关键在于采样分布的标准差。把标准差调得够好的话,模型的误差甚至低于密集采样的误差。
classFourierFeature(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) defforward(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 == 0or 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
classFourierFeature(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) defforward(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$,所以不需要对输出向量乘系数。
之后,我们来尝试在傅里叶特征中只用正弦函数。我们将投影矩阵的输出通道数从 out_c / 2 变成 out_c,再在 forward 里只用 sin 而不是同时用 sin, cos。经实验,这样改了后完全不影响重建质量,甚至由于通道数更多了,重建效果更好了。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
classFourierFeature(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) defforward(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
只用 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()
扩散模型可以直接生成任何形状的数据。如果我们不把视频视为一种由图像组成的序列数据,而是将其视为一种「三维图像」,那么我们可以直接将 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 系统探讨了如何在训练时独立地给序列元素添加噪声。
作者说出于简洁,他们在论文中用 RNN 实现了 Diffusion Forcing。但很明显 3D U-Net 才应该是直观上最简单实用的实现方法,毕竟最早期的视频扩散模型就是拿 3D U-Net 做的。在官方仓库中,有本科生帮他们实现了一个 3D U-Net 加时间注意力的模型,比原来视频模型效果要好。
这篇工作对我最大的启发是,我们一直把视频当成完整的 3D 数据来看待,却忘了视频可以被看成是图像序列。如果把视频当成 3D 数据的话,不同帧只能通过时序注意力看到其他帧在当前去噪时刻的信息;而对于序列数据,我们可以在不同帧的依赖关系上做更多设计,比如这篇工作的不同去噪强度。我很早前就在构思一种依赖更加强的序列生成范式:我们可不可以把序列中其他元素的所有去噪时刻的所有信息(包括中间去噪结果及去噪网络的中间变量)做为当前元素的约束呢?这种强约束序列模型可能对多视角生成、视频片段生成等任务的一致性有很大帮助。由于生成是约束于另一个去噪过程的,我们对此去噪过程做的任何编辑,都可以自然地传播到当前元素上。比如在视频生成中,如果整个视频约束于首帧的去噪过程,那么我们用任意基于扩散模型的图像编辑方法来编辑首帧,都可以自然地修改后续帧。当然,我只是提供一个大致的想法,暂时没有考虑细节,欢迎大家往这个方向思考。
作为一名未来的游戏设计师,每次看到这类「今天 AI 又取代了创作者」的新闻,我的第一反应总会是愤怒:创作是人类智慧的最高结晶,能做到这种程度的 AI 必然是强人工智能。但显然现在 AI 的水平没有那么高,那么这类宣传完全是无稽之谈。我带着不满看完了论文,果然这个工作并没有在技术上有革命性的突破。不过,这篇论文还是提出了一个比较新颖的科研任务并漂亮地将其解决了的,算是一篇优秀的工作。除了不要脸地将自己的模型称为「游戏引擎」外,这篇工作在宣传时还算克制,对模型的能力没有太多言过其实的描述。
为了生成足够多的图片,作者利用强化学习训练了一个玩游戏的 AI。在这一块,作者用了一个非常巧妙的设计:和其他强化学习任务不同,这个玩游戏的 AI 并不是为了将游戏漂亮地通关,而是造出尽可能多样的数据。因此,该强化学习的奖励函数包括了击中敌人、使用武器、探索地图等丰富内容,鼓励 AI 制造出不同的游戏画面。
有人可能还会说:「我也同意深度学习代替不了人类,但也不能说这些技术就完全没用」。这我非常同意,我就认为大家应该把现在的 AI 当成一种全新的工具。基于这些新工具,我们把创新的重点放在如何适配这些工具上,辅助以前的应用,或者开发一些新的应用,而不是非得一步到位直接妄想着把人类取代了。比如,根据简笔画生成图片就是一个很好的新应用啊。
这篇工作用图像模型学习了 3D 场景在移动后的变化。也就是说,模型「理解」了 3D 场景。那么,有没有办法从模型中抽取出相关的知识呢?按理说,能理解 3D,就能理解物体是有远近的。那么,深度估计、语义分割这种任务是不是可以直接用这种模型来做呢?以交互为约束的图像生成模型可能蕴含了比文生图模型更加丰富的图像知识。很可惜,不知道这篇工作最后会不会开源。
GameNGen 将模拟 3D 可交互场景的任务定义为根据历史画面、历史及当前操作生成当前画面的带约束图像生成任务。该工作用强化学习巧妙造出大量数据,用扩散模型实现带约束图像生成。结果表明,该模型不仅能自回归地生成连贯的游戏画面,还能学会子弹、血量等复杂交互信息。然而,受制于硬件及模型架构限制,模型要求的训练资源极大,且一次只能看到 3.2 秒内的信息。这种大量数据驱动的做法难以在学校级实验室里复刻,也不能够归纳至更一般的 3D 世界模拟任务上。
我个人认为,从科研的角度来看,这篇工作最大的贡献是提出了一种用带约束图像生成来描述 3D 世界模拟任务的问题建模方式。其次的贡献是确确实实通过长期的工程努力把这个想法做成功了,非常不容易。但从游戏开发的角度来看,这个工作现阶段没什么用处。
defcalculate_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 方法。
# 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)
经苏剑林设计,假设每个 token 的二维位置编码是一个复数,如果用以下的公式来定义绝对位置编码,那么经过注意力计算里的求内积操作后,结果里恰好会出现相对位置关系。设两个 token 分别位于位置 $m$ 和 $n$,令给位置为 $j$ 的注意力输入 Q, K $q_j, k_j$ 右乘上 $e^{ij/10000}$的位置编码,则求 Q, K 内积的结果为:
其中,$i$ 为虚数单位,$*$ 为共轭复数,$Re$ 为取复数实部。只是为了理解方法的思想的话,我们不需要仔细研究这个公式,只需要注意到输入的 Q, K 位置编码分别由位置 $m$, $n$ 决定,而输出的位置编码由相对位置 $m-n$ 决定。这种位置编码既能给输入提供绝对位置关系,又能让注意力输出有相对位置关系,非常巧妙。
for index_block, block inenumerate(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 特征向量的分量两两分组,根据之前的公式,模拟与位置编码的复数乘法运算。
由于不同图表的采样速度指标不太一样,我们将指标统一成每秒生成的图像。从第一张图的对比可以看出,DiT 最快也是一秒 2.5 张图像左右,而 MAR 又快又好,默认(自回归步数 64)一秒生成 3 张图左右。同时,通过 MAR 和有 kv cache 加速的标准 AR 的对比,我们能发现 MAR 在默认自回归步数下还是比标准 AR 慢了不少。
我们再看中间 LDM 的速度。我们观察一下最常使用的 LDM-8。如果是令 DDIM 步数为 20 (第二快的结果)的话,LDM-8 的生成速度在一秒 16 张图像左右,还是比 MAR 快很多。DDIM 步数取 50 时也会比 MAR 快一些。