机器学习中的旋转等变性：原理、与不变性的区别及PyTorch实现

做图像分类、目标检测或者三维点云任务时很多开发者都会遇到一个很朴素的烦恼同一个物体只要旋转了一个角度网络输出的结果就可能完全跑偏。为了缓解这种问题最常见的做法是随机旋转增强把不同角度的样本“喂”给网络。可问题往往没有这么简单——增强可以提升训练集的覆盖度却不能从结构上保证模型对旋转真的有稳定的响应能力。这篇文章围绕Rotational Equivariance in Machine Learning展开也就是机器学习和深度学习里的“旋转等变性”。我会先解释旋转等变和旋转不变这两个容易被混淆的概念再分析为什么普通卷积神经网络先天缺乏旋转等变能力接着介绍主流的实现思路并提供一个基于 PyTorch 的小实验来验证普通卷积的旋转敏感性。文章适合刚接触几何深度学习的新手也适合希望在图像、点云、分子等任务中引入旋转等变结构的开发者。1. 一个让普通卷积失效的问题先来看一个非常常见的场景。假设我们有一个训练好的图像分类模型训练集里的“汽车”基本都是水平摆放车头朝左。测试时出现同一辆汽车但视角被旋转了 30 度。对人类的视觉系统来说这仍然是同一辆车分类结果不应该改变。但传统卷积神经网络不一定能给出正确答案因为卷积核在二维网格上学习到的特征例如车轮和车身的相对位置会随着旋转发生明显变化。你可能会说这个问题好解决只要在训练时加入随机旋转增强把数据样本做得足够多模型不就能适应了吗数据增强确实有效但它的本质是在“数据分布”层面做补偿不是让网络结构本身掌握旋转对称性。它更像用大量样本来掩盖结构上的不足。那么是否存在一种网络结构能够让模型天然地“理解”旋转而不是靠大量数据去硬学这正是旋转等变性要讨论的问题。如果把“物体在图像中旋转”看成是对输入施加了一个旋转变换就很自然地会问当输入发生了旋转模型内部的特征和输出应该发生什么变化这个问题的答案分为两种输出保持完全不变例如物体分类结果不会因为图像旋转而改变。输出跟着旋转例如检测模型产生的检测框、分割模型产生的像素级掩码应该随图像同步旋转。第二种需求就对应了旋转等变性。它并不是一个只在论文里出现的高级概念而是真实工程任务中非常普遍的需求。2. 什么是机器学习中的旋转等变性2.1 旋转不变性与旋转等变性的区别先给一个直观定义。如果一个模型记为函数 (f)现在输入是一张图像 (x)我们对图像做一个旋转变换 (R)得到 (R(x))。如果模型在旋转前后给出的结果完全一样也就是[ f(R(x)) f(x) ]那么把这种性质称为旋转不变性Rotation Invariance。图像分类是一个典型任务因为“旋转后是否还是同一只猫”这个语义不取决于角度分类标签通常不会因为旋转而改变。但是如果模型输出本身带有空间结构比如分割掩码、检测框、像素级关键点那么旋转后输出也应该跟着旋转。假设对输出特征执行同一个旋转操作 (R)期望满足[ f(R(x)) R(f(x)) ]这就是旋转等变性Rotation Equivariance。为了减少符号上的压力可以把这两个概念用一个简单例子理解把一张“向右箭头”图片输入网络。如果网络输出结果是“右”那么旋转后图像还应该输出“右”这是旋转不变性因为类别没有方向。如果网络输出的是一个箭头位置图原图箭头向右旋转后箭头指向下那么输出特征图里的箭头也应该指向下这就是旋转等变性。两者容易混淆一个简单判断方法是看任务输出是否携带空间方向信息。只关心“有没有某类物体”的分类任务更多需要旋转不变关心“物体在哪里、朝向哪里”的任务通常会需要旋转等变。用一张表来总结区别性质数学表达典型任务旋转不变性(f(R(x)) f(x))图像分类、全局回归旋转等变性(f(R(x)) R(f(x)))语义分割、目标检测、位姿估计混合形式特征层等变、顶层不变先提取旋转等变特征再做全局池化完成分类实际网络中等变和不变并不互斥。很多旋转等变网络先把底层特征做成对旋转同步变化的表示到了最上层需要全局标签时再利用旋转维度的池化把等变特征转换为不变特征。2.2 普通CNN为什么先天不具备旋转等变性这个问题的根源在于卷积网络的设计原则。卷积神经网络之所以对“平移”有很强的适配能力是因为同一个卷积核会在输入的所有空间位置重复滑动这种参数共享机制使它天然具有平移等变性。把一个物体在图像中移动几个像素网络的卷积响应也会对应移动因此平移等变在普通CNN里几乎是“免费”的。但旋转变换不是这样。CNN的卷积核形状通常是矩形例如 (3\times3) 或 (5\times5)这些卷积核本身带有明确的“方向偏好”。一个专门提取横向边缘的卷积核很难响应纵向边缘除非网络在别的卷积核里重新学一套纵向模板。你可能会说数据量足够大时网络可以在同一层中学习到多个方向的卷积核。这种说法有一定道理深度学习模型确实能通过海量数据隐式覆盖多个角度但请注意这属于数据驱动的近似行为不是模型结构对旋转的数学承诺。网络需要额外参数来记住不同方向的模式而不是利用旋转对称性来共享参数。对训练集里出现频率低的方向比如某些特殊旋转角度模型泛化效果通常较差。当数据中存在噪声、遮挡或样本量受限时这种近似覆盖往往不可靠。所以普通CNN虽然可以实现平移等变却没有先天旋转等变能力。要让网络真正对旋转保持稳定的结构诱导偏置就需要显式设计包含旋转对称性的模型结构。2.3 旋转等变性在哪些场景更关键旋转等变性并不是所有任务都需要但在以下场景中会非常关键。第一类是图像方向不固定的视觉任务。比如遥感图像中的建筑物、船舶由于卫星拍摄方向和目标摆放没有统一朝向目标在图像中可以呈现任意旋转角度医学病理切片、显微镜图像也常出现细胞或组织的任意朝向工业缺陷检测中产品可能在传送带上发生不同角度翻转。这些数据如果直接交给普通CNN往往需要做大量增强。第二类是三维几何与点云任务。三维点云里的物体可以在空间中任意旋转例如用扫地机器人扫描到的椅子、用自动驾驶激光雷达捕获到的车辆都不存在整齐统一的“世界方向”。模型如果把一个方向的椅子记住了换一个方向就容易失效。点云分类和部件分割任务中引入旋转等变结构能够显著降低对视角多样性的依赖。第三类是分子性质预测和物理模拟。分子的几何结构在三维空间中做任意旋转都不会改变它的化学性质因此分子的能量、属性预测函数应该满足旋转不变性而力场中的每个原子受力方向会随分子旋转而旋转所以力预测等任务需要旋转等变性。这类数据通常样本量小、标注成本高结构对称性约束非常有用。第四类是目标检测和位姿估计。当模型不仅要识别物体还要输出带方向的边界框、抓取位姿、车辆朝向时普通分类网络常用的全局池化会丢掉方向信息。旋转等变特征可以将“响应位置”与“响应方向”同时保存在特征里方便后续网络回归位姿。3. 实现旋转等变性的主要技术路线旋转等变性听起来很高深但工程上并不是只能靠从头写论文模型来实现。下面几种路线各有优缺点从最简单的数据层处理到严格的结构化卷积横向对比会更容易理解。3.1 数据增强最朴素但最不彻底最常见的方案是在训练时对输入做随机旋转。PyTorch 里通常这样写train_transform transforms.Compose([ transforms.RandomRotation(30), transforms.ToTensor(), # 其他标准化 ])这个方案的优点是实现成本极低对已有模型结构没有侵入性也不需要改网络。它能让网络在一定程度上适应旋转后的输入分布。缺点也很明显网络参数依然没有共享不同旋转方向的模式。模型在角度采样足够密集、增强强度足够大时表现可能变好但需要更多训练数据与训练时间。如果旋转角度超出了训练时设置的范围模型依然容易崩溃。它不能保证输出的结构跟随输入旋转例如分割掩码的边界仍需要解码层来对齐方向。因此数据增强适合作为一种通用兜底手段但很难用“让模型学会旋转”来描述它的作用更准确的说法是“让模型见过更多旋转样本”。3.2 主方向对齐用预处理规避旋转另一种思路是在输入进入网络前先估计物体的主方向