CVPR 2022 | 应用于旋转回归的流形投影梯度层

论文地址：https://arxiv.org/abs/2110.11657

该工作代码已开源，欢迎大家关注~

项目主页：https://jychen18.github.io/RPMG/

代码地址：https://github.com/JYChen18/RPMG

三维旋转（朝向）估计旨在预测三维空间中三自由度的旋转。这一问题广泛存在于计算机视觉领域中的众多任务中，如物体/人体/手的位姿估计、相机重定位等。近年来，随着神经网络的兴起，人们开始尝试用神经网络来解决三维旋转的预测任务。其中，直接回归（regression）三维旋转具有速度快、端到端的优势，吸引了很多研究者的关注。

我们知道三维平移和三维旋转一同构成了六维位姿。然而与三维平移不同的是，三维旋转属于非欧几里得变量，其非欧的性质对神经网络直接进行回归带来了重大的挑战。具体来说，全部三维旋转构成了一个连续群 SO(3)，这一连续群是一种三维黎曼流形（Riemannian manifold），而并非传统的欧几里得空间。我们知道神经网络的输出在线性激活的情况下是处在欧式空间 上的，这种空间性质上的差异为旋转预测带了不小的困难，如何弥合这一鸿沟就成为了一个重要的研究问题。

现有的工作主要围绕三维旋转的表示（representation）。我们知道三维旋转只有三个独立自由度，但是学界为其发展出了多种不同的表示方法，包括三维的欧拉角（Euler angle）表示、三维的轴角（axis-angle）表示、四维的四元数（quaternion）表示和九维的旋转矩阵表示。虽然旋转的欧拉角和轴角表示是最小表示且张成了一个欧式空间，但是已有的工作发现使用神经网络直接回归欧拉角和轴角表示的准确度偏低，因此长久以来四元数成为了比较主流的神经网络的预测表示。

2019年 Zhou et al. [1] 在 CVPR 的工作第一提出了一个理论框架，对不同旋转表示与神经网络的契合程度进行了分析并提出了一种六维的冗余表示。在此之后学者们相继提出了基于 SVD 正交化的九维旋转矩阵表示和十维表示。这些表示的一个共性问题是他们的维数 都高于实际的三维自由度，构成了在高维欧几里得空间中嵌入的三维黎曼流形。因此神经网络的输出需要经过一个流形映射（manifold mapping），将 维欧几里得空间的输出投影到流形上才能转化为有效的流形上表示，并最终转化成一个旋转矩阵作为输出并施加监督。

在本项工作中，我们第一次关注到了这样常常是多对一的流形映射在梯度回传中构成的特有问题，并针对这样一类从欧式空间到黎曼流形的映射设计了一个新的梯度反传层——针对流形的正则投影梯度层（regularized projective manifold gradient layer），来取代简单的基于链式法则的梯度回传，以促进网络的优化。我们的方法不仅能广泛适用于之前工作提出的各种非欧表示（四元数/六维/九维/十维)，还能推广到其他的黎曼流形，如回归 上的二维方向向量等。通过大量的实验，我们展示了我们的方法能够在不带来额外时空开销的情况下，在旋转矩阵的各种非欧表示（四元数/六维/九维/十维）及多项任务（如基于点云的三维物体位姿估计、基于图像的三维物体位姿估计、相机重定位等）中取得一致且显著的提升。

图1. 方法完整流程图。

为了便于理解，我们用一个简单的例子，预测三维空间中的单位球面 （也是一个黎曼流形）上的一个单位方向向量来介绍我们的方法。如果我们要用神经网络去回归 上的一个元素 （满足 ），那么最常用的做法是首先用神经网络去预测一个 中的元素 ，然后再通过一个流形映射 来得到 。这就是神经网络正向传播的过程。在训练的时候，人们会对 施加一个损失函数 （常见的选择是 ），然后用基于链式法则来计算 梯度并进一步反向传播给网络，用于更新网络的参数。

我们发现这样的反向传播过程并不是最优的。我们注意到，由于 属于一个嵌入在三维空间中的二维流形（单位球面），我们希望网络利用梯度更新权重后能够使得 更加靠近同在流形上的 ，可以理解为可以得到一个特殊的梯度（黎曼梯度）使得 可以在流形上更新。然而，直接根据 计算梯度（欧几里得梯度）然后加到 上，常常会导致 离开球面上。从这个意义上讲，这样的梯度含有噪声，继续反传这一梯度给网络权重会使得更新后的网络未必能回归出一个很好的 使得其经过流形映射 后得到的 更接近 。

因此，我们首先采用流形上的优化技术——黎曼优化（Riemannian optimization）来得到一个在黎曼流形 上的更新 ，之后进一步利用它来计算 的梯度：我们希望达到的目的是让 沿梯度方向做一次更新后能满足 。最直接的解决办法是令 ，但是由于映射 常常是多对一的映射， 得到的往往是一个集合而非单一元素（对二维球面来说，这个元素集合就是一条从原点指出通过 的射线），因此我们需要从这个集合中选一个最优的 来计算 。

我们进一步意识到，这个寻找最优的过程其实可以类比于多正确答案（multi-ground truth）的监督问题。He et al. [4] 在监督网络四元数的输出时曾遇到过这样一个问题：由于四元数 和 都代表了同一个旋转矩阵，相当于存在两个正确答案，但是网络的输出并不是一个分布而是一个特定向量，它只能往 或者 其中一个靠近，因此在计算损失函数时必须从 和 这两个答案中选取其一。考虑到希望网络能以尽可能小的变动去实现所需要的更新，他们提出的解决办法是选取离当前网络输出值最近的那个答案。受他们的启发，我们也采用了同样的办法去选取最优的 ： 。针对旋转矩阵的不同表示，我们都进行了推导，求出了这一步的解析解，因此不需要采用优化的办法去找 ，不会引入额外的时间开销。至此，我们可以得到 的 PMG（projective manifold gradient，图2左蓝色箭头） 。

在实验中，我们发现由于 ，以上述方式计算出来的梯度 PMG 会持续地给 的模长 一个变小的趋势（如图2中），在不同流形、不同表示上均有此现象。这会使得网络的输出变得特别的小，相当于等效地增加了相对学习率（learning rate），造成了训练的不稳定和网络收敛的困难。为了解决这一问题，我们需要引入一项额外的正则项来稳定输出的模长，最终得到 的 RPMG（regularized projective manifold gradient，图2左紫色箭头） 。如图2右所示，我们在实验中发现只需要取一个非常小的 就可以避免模长消失的问题。

图2. 方法直观示意图（左）；不加正则项时网络直接输出 x 的模长关于训练过程的变化趋势图（中）；加正则项时网络直接输出 x 的模长关于训练过程的变化趋势图（右）。

虽然上面我们只介绍了 的情形，但我们的方法可以推广到旋转矩阵的各种表示（representation），如四元数（quaternion）、6D [1]、9D [2]、10D [3]。通过大量实验，我们验证了在旋转矩阵的各种表示、各类任务中，使用我们提出的梯度反传层得到的方法 RPMG-Quat、RPMG-6D、RPMG-9D、RPMG-10D，相较于各自的 baseline 都能有极其显著的性能提升。

我们的主实验是基于三维点云的类别级物体位姿估计。我们用一个物体的三维点云作为输入，希望网络去预测输入的点云 相对于该类物体预定义的物体坐标系的旋转 。我们采用和 ground truth 旋转矩阵的 L2 loss 来作为损失函数。实验结果如图3、图4所示，可以看出我们的梯度层能在不同类别的物体、不同的旋转矩阵表示下（图3），在各个误差区间都能获得显著的提升（图4右）。另外我们的方法能收敛得更快更好（图4左）。

图3. 基于 ModelNet [5] 三维点云的类别级物体位姿估计。Mn，Md，Acc 分别代表平均角度误差（Mean）、角度误差中位数（Median）、角度误差小于5°的百分比（5°Accuracy）。

下左图是对于飞机这个类的训练过程中平均角度误差的变化情况，可以看出我们的方法能收敛的更快更好。下右图是角度误差的分布情况，可以看出我们的方法在各个误差区间都能显著优于 baseline。

图4. 训练过程中训练集上平均角度误差的变化情况（左），训练完成后在测试集上角度误差的分布情况（右）。物体类别均为飞机。

此外，我们的方法在基于图片的类别级物体位姿估计的实验中（图5），也能有显著、一致的提升。实验的设置和第一个实验完全一致，唯一的区别是给网络的输入变成了一张包含特定类别物体的二维图片。

图5. 基于 ModelNet [5] 二维图片的类别级物体位姿估计。Mn，Md，Acc 分别代表平均角度误差（Mean）、角度误差中位数（Median）、角度误差小于5°的百分比（5°Accuracy）。

我们还利用对于瓶子这一类的三维点云位姿估计实验，在其他流形（ ）上验证了我们提出的算法（图6）。由于瓶子具有对称性，它的位姿被长轴方向唯一确定，因此我们只需要预测一个三维空间中模长为1的单位向量即可（也是我们在介绍方法时所举的简单的例子）。网络的输入依然是完整的、随机旋转后的物体点云，输出是该瓶子点云的长轴方向。我们的 RPMG 依然可以显著提升该任务的表现。

图6. 其他流形（ ）上元素的回归实验。

更多其他实验，如自监督的三维点云位姿估计、相机重定位的实验，请详见原论文。

参考文献：

[1] Zhou, Yi, et al. "On the continuity of rotation representations in neural networks." CVPR 2019.

[2] Levinson, Jake, et al. "An analysis of svd for deep rotation estimation." NeurIPS 2020.

[3] Peretroukhin, Valentin, et al. "A smooth representation of belief over so (3) for deep rotation learning with uncertainty." Robotics: Science and Systems 2020.

[4] Wang, He, et al. "Normalized object coordinate space for category-level 6d object pose and size estimation." CVPR 2019.

[5] Wu, Zhirong, et al. "3d shapenets: A deep representation for volumetric shapes." CVPR 2015.