Sora和快手可灵背后的核心技术 | 3DVAE:通过小批量特征交换实现身体和面部的三维形状变分自动编码器

【摘要】学习3D脸部和身体生成模型中一个解开的、可解释的和结构化的潜在表示仍然是一个开放的问题。当需要控制身份特征时,这个问题尤其突出。在本文中,论文提出了一种直观而有效的自监督方法来训练一个3D形状变分自动编码器(VAE),以鼓励身份特征的解开潜在表示。通过交换不同形状间的任意特征来管理迷你批次生成,这允许定义一个利用已知的潜在表示中的差异和相似性的损失函数。在3D网格上进行的实验结果表明,用于潜在解开的最先进的方法无法解开脸部和身体的身份特征。论文提出的方法正确地解耦了这些特征的生成,同时保持了良好的表示和重建能力。

原文：3D Shape Variational Autoencoder Latent Disentanglement via Mini-Batch Feature Swapping for Bodies and Faces
地址：https://arxiv.org/abs/2111.12448
代码：https://github.com/simofoti/3DVAE-SwapDisentangled
出版：CVPR
机构: 伦敦大学学院

1 研究问题

本文研究的核心问题是: 如何设计一种自监督的方法,在3D人脸和人体形状的生成模型中实现身份特征解开的潜在表示。

::: block-1
想象论文要开发一个3D游戏或虚拟试衣间,需要程序自动生成各种逼真的人脸和人体模型。论文希望这个生成系统有以下特点:1)能够通过调节少量的控制参数,就生成千变万化的脸部和身体;2)这些控制参数应该是可解释的,比如对应眼睛、鼻子、嘴巴等具体的面部器官。这样的生成工具不仅使用方便,而且便于实现局部的、细粒度的形状控制。
:::

本文研究问题的特点和现有方法面临的挑战主要体现在以下几个方面:

身份特征具有明显的局部性,而现有的3D形状生成模型往往只学习到全局的latent表示,缺乏对局部身份特征的显式建模。
缺乏可用于身份特征解开学习的监督信号。现有的latent解开方法主要针对人脸表情和姿态的解耦,而缺少面向身份特征的benchmark和评估标准。
身份特征之间往往存在复杂的相关性,简单地对latent空间做划分,无法学习到语义清晰的特征表示。

针对这些挑战,本文提出了一种基于"迷你批次特征交换"的VAE训练范式:

::: block-1
核心思想是在组建每个训练批次时,随机地在不同形状实例之间交换局部的身份特征(如鼻子、眼睛、下巴等),从而得到一批在特定身份特征上存在差异,而在其他特征上保持一致的"合成"形状。然后要求VAE学习到的latent表示能够反映这种差异和一致性,即属于同一身份特征的潜变量应该在批内保持不变,而控制不同身份特征的潜变量应该各不相同。这其实是构建了一种弱监督信号,引导VAE去学习身份特征解开的表示（Sora和可灵等应该构建了很多弱监督信号）。同时,本文还提出了一个新的latent一致性损失函数,进一步利用批内形状的相似性约束,来强化解开学习的效果。实验表明,本文方法能在不影响生成质量的情况下,大幅提升latent表示的可解释性,实现对身份特征的有效解耦和操控。
:::

2 研究方法

本文提出了一种新颖的方法来训练3D形状变分自编码器(3D-VAE),以学习人脸和人体3D网格的解耦、可解释和结构化的潜在表示。该方法包括三个关键部分:3D-VAE的基本结构、特征交换mini-batch的生成,以及潜在一致性损失函数的构建。接下来将详细介绍每一部分。

2.1 3D形状变分自编码器

论文采用文献[16]中提出的网格卷积VAE作为基本架构。编码器 $E$ 将输入网格 $X$ 映射到潜在分布 $Z$ , 定义为近似后验 $q (z ∣ X)$ 。解码器(也称生成器) $G$ 根据似然 $p (X ∣ z)$ 将潜在向量 $z$ 映射回网格。编码器使用spiral++卷积、ELU激活和池化层,而解码器使用spiral++卷积、ELU激活和反池化层。此外还有三个全连接层:两个预测潜在分布均值和对角协方差,一个将 $z$ 转换回网格。

训练时最小化如下损失函数:

$L_{VAE} = L_R + αL_L + βL_{KL}$

其中 $L_R$ 是重建损失, $L_L$ 是基于网格拉普拉斯矩阵的平滑正则项, $L_{KL}$ 是变分分布 $q (z ∣ X)$ 与先验 $p (z)$ 间的KL散度。 $α$ 和 $β$ 是平衡系数。

2.2 特征交换mini-batch生成

核心创新点之一是特征交换方法。具体来说,论文先在一个模板网格上定义 $F$ 个特征区域,每个区域对应一组顶点。由于数据集中的网格具有顶点对应关系,可以很容易地识别出每个网格的各个特征区域。举个例子,对于人脸,论文可以定义眼睛、鼻子、嘴巴等不同的特征(图1左)。

基于此,论文精心设计了mini-batch的生成过程。一个大小为 $B$ 的mini-batch可以看作一个 $\sqrt{B} \times \sqrt{B}$ 的矩阵,其中对角线上的元素从数据集中采样,其余元素通过特征交换得到。每次生成一个mini-batch时,论文随机选择一个特征并交换不同网格间该特征对应的顶点坐标(图1左)。这样,矩阵的每一行包含相同的人脸/人体,但具有不同的特征;而每一列包含不同的人脸/人体,但具有相同的特征。

2.3 潜在一致性损失函数

mini-batch的精心设计为构建潜在一致性损失函数奠定了基础。当一个mini-batch输入到编码器后,论文得到大小为 $\sqrt{B} \times \sqrt{B}$ 的潜在向量矩阵,其中每个向量 $z_{ij}$ 被均匀地划分为 $F$ 个子集,分别控制一个特征(图1中间)。

假设当前mini-batch交换了特征 $f$ , 则 $z_{ij} = (z^f_{ij}|z^c_{ij})$ 。其中 $z^f_{ij}$ 控制特征 $f$ , 而 $z^c_{ij}$ 控制其他特征。直觉上,同一列的 $z^f_{ij}$ 应该接近,因为它们对应相同的特征 $f$ ;而同一行的 $z^f_{ij}$ 应该远离。反之,同一行的 $z^c_{ij}$ 应该接近,而同一列的 $z^c_{ij}$ 应该远离。这就好比对角线上的人脸彼此不同,但具有相同的特征 $f$ 。

基于这一思想,论文构建潜在一致性损失函数如下:

$\begin{aligned} L_c = \frac{γ}{\sqrt{B}} \sum_{s,p,q=1}^{\sqrt{B}} \mathbb{I}_{p \neq q} \Big(&\max[0, \|z^f_{ps} - z^f_{qs}\|^2_2 - \|z^f_{sp} - z^f_{sq}\|^2_2 + η_1] \\ +&\max[0, \|z^c_{sp} - z^c_{sq}\|^2_2 - \|z^c_{ps} - z^c_{qs}\|^2_2 + η_2] \Big) \end{aligned}$

其中 $γ$ 是归一化系数, $\mathbb{I}$ 是指示函数, $η_1$ 和 $η_2$ 是容差。该损失函数遍历mini-batch矩阵的对角线,并要求匹配的潜在变量对之间的距离比不匹配的对之间的距离小一个容差 $η$ 。这鼓励 VAE 学习到的潜在表示能够解耦地控制每个特征的生成。

最后,结合原有的VAE损失(1)和潜在一致性损失(2),论文得到总的损失函数:

$L = L_{VAE} + κL_c = L_R + αL_L + βL_{KL} + κL_c$

其中 $κ$ 控制潜在一致性损失的权重。

综上所述,本文提出了一种通过特征交换生成mini-batch、并构建潜在一致性损失函数的方法,使得3D-VAE能够学习到更加解耦、可解释和结构化的3D人脸和人体的潜在表示。该方法只需要对标准的VAE进行很小的改动,却能学习到可以独立控制每个特征生成的潜在空间。这极大地提高了生成模型的可解释性和可用性。

3 实验

3.1 实验场景介绍

本文提出了一种新的VAE框架,通过批特征交换与潜在一致性损失,使学习到的潜在表征具有更好的可解释性、结构性和解纠缠特性。实验旨在验证该方法在人脸与人体3D网格生成任务中的有效性。

3.2 实验设置

Datasets：从UHM和STAR线性模型中随机生成人脸与人体3D网格数据集,各1万个,划分为训练集(90%)、验证集(5%)和测试集(5%)
Baseline：VAE、DIP-VAE-I/II、Factor VAE等最先进的自监督潜在解纠缠模型
Implementation details：
- 骨干网络:SpiralNet++,人脸模型4层卷积,人体模型3层卷积
- 超参数:40个epoch,ADAM优化器,学习率1e-4,批大小16,潜变量维度(人脸60维,人体33维)等
metric：
- 重建误差:平均/最大每顶点误差
- 生成多样性:随机生成的样本对之间的平均每顶点距离
- 分布距离:JSD、MMD、COV等
- 潜变量解纠缠:观察潜变量遍历对生成特征的影响

3.3 实验结果

3.3.1 实验一、与其他潜在解纠缠方法的性能对比

目的：评估所提出的模型与其他自监督潜在解纠缠方法在生成任务上的性能表现
涉及图表：表1
实验细节概述：在人脸数据集上,比较了VAE、DIP-VAE-I/II、Factor VAE等模型的重建误差、生成多样性、分布距离等指标
结果：

本文模型的重建误差与普通VAE相当,优于其他方法,说明批特征交换和一致性损失并未降低重建能力
本文模型的生成多样性与大多数方法相近,Factor VAE多样性更高但生成样本真实性较差
其他评估生成性能的指标表明,本文方法与其他模型性能相当

3.3.2 实验二、潜在表征的解纠缠效果评估

目的：定性评估学习到的潜在表征对控制不同面部/身体特征的解纠缠能力
涉及图表：图3、图4、图5、图10、图11
实验细节概述：遍历每个潜变量,生成两个改变该变量至最小值(-3)和最大值(+3)的3D网格,观察顶点位移,分析其对不同特征的影响
结果：

本文模型的潜在表征展现出明显的结构性和可解释性,不同潜变量对应控制了特定的面部/身体特征
相比之下,其他方法学到的潜在表征难以解释,大多数变量均表现出纠缠的特性

3.3.3 实验三、消融研究

目的：分析模型关键组件(批特征交换、一致性损失、拉普拉斯正则化、归一化)的作用
涉及图表：图13
实验细节概述：分别移除模型的潜在一致性损失、拉普拉斯损失和输入归一化,重新训练并对比性能
结果：

移除潜在一致性损失后,仅保留批特征交换无法获得解纠缠的潜在表征
移除拉普拉斯损失会导致生成网格表面不规则,出现不连续
移除输入归一化则降低了生成人脸的真实性

3.3.4 实验四、泛化能力评估

目的：验证模型对于训练分布外样本的拟合能力
涉及图表：图14
实验细节概述：在CoMA人脸数据集的中性表情上,通过优化潜码最小化与目标网格的重建误差,测试模型的泛化性能
结果：

重建误差随着施加的噪声线性增长,但整体仍保持较低水平,说明模型泛化能力较好
本文方法的重建误差在所有模型中最低,证明其更强的泛化性

4 总结后记

本论文针对3D生成模型中潜在表示解耦的问题,提出了一种基于Mini-Batch特征交换和潜在一致性损失的自监督方法。通过在Mini-Batch内的样本间交换预定义的网格特征并约束对应潜变量的相似性,实现了对身体和面部3D网格的特征级解耦生成控制。实验结果表明,所提方法能够学习到更加解释性强、结构化的潜在表示,生成效果优于现有方法。

::: block-2
疑惑和想法：