作为一名未来的游戏设计师,每次看到这类「今天 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 快一些。
# `accelerate` 0.16.0 will have better support for customized saving if version.parse(accelerate.__version__) >= version.parse("0.16.0"): defsave_model_hook(models, weights, output_dir): ... defload_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 isnotNone: set_seed(args.seed)
# Initialize the model if args.model_config_name_or_path isNone: 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, ...)
if args.resume_from_checkpoint: if args.resume_from_checkpoint != "latest": path = .. else: # Get the most recent checkpoint ...
if path isNone: ... 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 inrange(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 inenumerate(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
接下来,函数会用去噪网络做前向传播。为了让模型能正确累计梯度,我们要用 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
f accelerator.is_main_process: if global_step % args.checkpointing_steps == 0: if args.checkpoints_total_limit isnotNone: 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]))
defsave_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 inenumerate(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()
脚本默认的验证方法是随机生成图片,并用日志库保存图片。生成图片的方法是使用标准 Diffusers 采样流水线 DDPMPipeline。由于此时模型 model 可能被包裹成了一个用于多卡训练的 PyTorch 模块,需要用相关 API 把 model 解包成普通 PyTorch 模块 unet。如果使用了 EMA 模型,为了避免对 EMA 模型的干扰,此处需要先保存 EMA 模型参数,采样结束再还原参数。
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": ...
在保存模型时,脚本同样会先用去噪模型 model 构建一个流水线,再调用流水线的保存方法 save_pretrained 将扩散模型的所有组件(去噪模型、噪声调度器)保存下来。
# old if cfg.model_config isNone: model = UNet2DModel(...) else: config = UNet2DModel.load_config(cfg.model_config) model = UNet2DModel.from_config(config)
# Create EMA for the model. if cfg.use_ema: ema_model = EMAModel(...) ...
# new trainer.init_modules(enable_xformers, cfg.gradient_checkpointing)
# The config must have a "base" key base_cfg_dict = data_dict.pop('base')
# The config must have one another model config assertlen(data_dict) == 1 model_key = next(iter(data_dict)) model_cfg_dict = data_dict[model_key] model_cfg_cls = __TYPE_CLS_DICT[model_key]
虽然整流模型是这样宣传的,但实际上 SD3 还是默认用了 28 步来生成图像。单看这篇文章,原整流论文里的很多设计并没有用上。对整流感兴趣的话,可以去阅读原论文 Flow straight and fast: Learning to generate and transfer data with rectified flow
按照之前高分辨率文生图模型的训练方法,SD3 会先在 $256^2$ 的图片上训练,再在高分辨率图片上微调。然而,开发者发现,开始微调后,混合精度训练常常会训崩。根据之前工作的经验,这是由于注意力输入的熵会不受控制地增长。解决方法也很简单,只要在做注意力计算之前对 Q, K 做一次归一化就行,具体做计算的位置可以参考上文模块图中的 “RMSNorm”。不过,开发者也承认,这个技巧并不是一个长久之策,得具体问题具体分析。看来这种 DiT 模型在大规模训练时还是会碰到许多训练不稳定的问题,且这些问题没有一个通用解。
from ..unet_3d_blocks import MidBlockTemporalDecoder, UpBlockTemporalDecoder ...
classUpBlockTemporalDecoder(nn.Module): def__init__(...): super().__init__() for i inrange(num_layers): ... resnets.append(SpatioTemporalResBlock(...))
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order self._num_timesteps = len(timesteps) with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t inenumerate(timesteps):
# expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
Video LDM 在 U-Net 中加入时序层的方法与多数同期方法相同,是在每个原来处理图像的空间层后面加上处理视频的时序层。Video LDM 加入的时序层包括 3D 卷积层与时序注意力层。这些新模块本身不难理解,但我们需要着重关注这些新模块是怎么与原模型兼容的。
要兼容各个模块,其实就是要兼容数据的形状。本来,图像生成模型的 U-Net 的输入形状为 B C H W,分别表示图像数、通道数、高、宽。而视频数据的形状是 B T C H W,即视频数、视频长度、通道数、高、宽。要让视频数据复用之前的图像模型的结构,只要把数据前两维合并,变成 (B T) C H W 即可。这种做法就是把 B 组长度为 T 的视频看成了 $B \cdot T$ 张图片。
对于之前已有的空间层,只要把数据形状变成 (B T) C H W 就没问题了。而 SVD 又新加入了两种时序层:3D 卷积和时序注意力。我们来看一下数据是怎么经过这些新的时序层的。
2D 卷积会对 B C H W 的数据的后两个高、宽维度做卷积。类似地,3D 卷积会对数据最后三个时间、高、宽维度做卷积。所以,过 3D 卷积前,要把形状从 (B T) C H W 变成 B C T H W,做完卷积再还原。
接下来我们来看新的时序注意力。这个地方稍微有点难理解,我们从最简单的注意力开始一点一点学习。最早的 NLP 中的注意力层的输入形状为 B L C,表示数据数、token 长度、token 通道数。L 这一维最为重要,它表示了 L个 token 之间互相交换信息。如果把其拓展成图像空间注意力,则 token 表示图像的每一个像素。在这种注意力层中,L 是 (H W),B C H W 的数据会被转换成 B (H W) C 输入进注意力层。这表示同一组图像中,每个像素两两之间交换信息。而让视频数据过空间注意力层时,只需要把 B 换成 (B T) 即可,即把数据形状从 (B T) C H W 变为 (B T) (H W) C。这表示同一组、同一帧的图像的每个像素之间,两两交换信息。
在 SVD 新加入的时序注意力层中,token 依旧指代是某一组、某一帧上的一个像素。然而,这次我们不是让同一张图像的像素互相交换信息,而是让不同时刻的像素互相交换信息。因此,这次 token 长度 L 是 T,它表示要像素在时间维度上交换信息。这样,在视频数据过时序层里的自注意力层时,要把数据形状从 (B T) C H W 变成 (B H W) T C。这表示每一组、图像每一处的像素独立处理,它们仅与同一位置不同时间的像素进行信息交换。
此处如果你没有理解注意力层的形状变换也不要紧,它只是一个实现细节,不影响后面的阅读。如果感兴趣的话,可以回顾一下 Transformer 论文的代码,看一下注意力运算为什么是对 B L C 的数据做操作的。
微调 VAE 解码器
Video LDM 的另一项改动是修改了图像压缩模型 VAE 的解码器。具体来说,方法先在 VAE 的解码器中加入类似的时序层,并在 VAE 配套的 GAN 的判别器里也加入了时序层,随后开始微调。在微调时,编码器不变,仅训练解码器和判别器。
如果你没听过这套 VAE + GAN 的架构的话,请回顾 Stable Diffusion 论文及与其紧密相关的 VQGAN 论文。
以上就是 Video LDM 的模型结构。SVD 对其没有做任何更改,所以也没有在论文里对模型结构做详细介绍。稍有不同的是,Video LDM 仅微调了新加入的模块,而 SVD 在加入新模块后对模型的所有参数都进行了重新训练。
Video LDM 曾提出了一种把基础视频模型变成视频插帧模型方法。该方法以视频片段的首末帧为额外约束,在此新约束下把视频生成模型微调成了预测中间帧的视频预测模型。SVD 以同样方式实现了这一应用。
多视角生成
多视角生成是计算机视觉中另一类重要的任务:给定 3D 物体某个视角的图片,需要算法生成物体另外视角的图片,从而还原 3D 物体的原貌。而视频生成模型从数据中学到了物体的平滑变换规律,恰好能帮助到多视角生成任务。SVD 论文用不少篇幅介绍了如何在 3D 数据集上生成视频并微调基础模型,从而得到一个能生成环绕物体旋转的视频的模型。
结语
Stable Video Diffusion 是在文生图模型 Stable Diffusion 2.1 的基础上添加了和 Video LDM 相同的视频模块微调而成的一套视频生成模型。SVD 的论文主要介绍了其精制数据集的细节,并展示了几个微调基础模型能实现的应用。通过微调基础低分辨率文生视频模型,SVD 可以用于高分辨率文生视频、高分辨率图生视频、视频插帧、多视角生成。
模型输出不同,最容易想到的是模型权重不同。于是,我尝试交换使用二者的模型权重,再比较输出的 FID。两个库的模型定义不太一样,不能直接换模型文件名。我用强大的代码魔改实力强行让新权重分别都跑起来了。结果非常神奇,算上之前的两个 FID,我一共得到了 4 个不一样的 FID 结果。也就是说,A 库 A 模型、B 库 B 模型、A 库 B 模型,B 库 A 模型,结果均不一样。