从‘椅子旋转’到代码：图解神经网络中的等变（Equivariant）与不变（Invariant），附向量神经元实例

从椅子旋转到向量变换：几何深度学习的等变与不变原理实战

想象你坐在一把可以360度旋转的办公椅上，左手拿着咖啡杯，右手握着手机。当你顺时针旋转30度后，咖啡杯和手机的位置关系保持不变——它们依然在你的左右手中，只是相对于房间的绝对方向改变了。这种空间关系的一致性，正是几何深度学习中**等变（Equivariant）与不变（Invariant）**概念的现实映射。本文将用日常物品的几何变换作为引子，逐步拆解这两个关键概念在神经网络中的实现原理，并通过PyTorch代码展示如何设计能"理解"空间关系的向量神经元。

1. 等变与不变的现实隐喻

1.1 旋转椅上的坐标系

让我们继续用旋转椅的比喻来建立直觉。假设你戴着一个智能眼镜，可以识别手中的物品：

• 等变场景：当椅子旋转时，眼镜检测到的"咖啡杯在左侧，手机在右侧"的空间关系会同步变化。如果旋转θ角度，检测结果坐标系也会旋转θ角度，但相对位置描述不变。

• 不变场景：同一系统中，物品分类结果（"咖啡杯"和"手机"的标签）不应因旋转而改变。无论怎么转，杯子不会变成手机。

这种区别在3D物体识别中至关重要。下表对比了两个概念的典型应用场景：

其中ρ(g)表示群g在输入空间的表示，ρ'(g)表示输出空间的群表示

1.2 从物理世界到向量空间

将这个概念延伸到神经网络，我们需要处理的不再是具体的咖啡杯，而是它们的向量表示。传统全连接层的一个根本局限在于，它处理的是孤立的标量值，完全忽略了输入元素可能存在的空间关系。这就是**向量神经元（Vector Neurons）**的设计动机——让网络能够原生处理具有几何意义的数据结构。

考虑一个简单的3D点坐标预测任务：

# 传统标量神经元处理3D点 class PointNet(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(3, 64) # 将xyz作为独立标量处理 def forward(self, points): return self.fc(points) # 丢失空间关系信息

这种处理方式的问题在于，当输入点云旋转时，网络需要重新学习旋转后的所有可能变体。而等变网络的设计目标是通过数学约束，让网络自动适应这种变换。

2. 向量神经元的实现解剖

2.1 核心设计原理

向量神经元层的核心思想是将每个神经元扩展为可以处理向量值的计算单元。与普通线性层不同，它的权重不再是简单的标量缩放系数，而是能保持向量空间关系的变换矩阵。以下是关键设计要素：

1. 输入输出结构：每个神经元处理的是向量而非标量，因此输入维度为(batch, channels, vector_dim)
2. 权重张量：权重变为四维张量(out_channels, in_channels, vector_dim, vector_dim)
3. 等变操作：使用矩阵乘法而非点积，保持向量空间关系

class VectorNeuronLayer(nn.Module): def __init__(self, in_channels, out_channels, dim=3): super().__init__() # 权重形状：(out_ch, in_ch, dim, dim) self.weight = nn.Parameter(torch.randn(out_channels, in_channels, dim, dim)) # 偏置形状：(out_ch, dim) self.bias = nn.Parameter(torch.randn(out_channels, dim)) def forward(self, x): # x形状：(batch, in_ch, dim) # einsum解释：对in_ch维度求和，对dim维度矩阵乘法 return torch.einsum('bic,ocde->boe', x, self.weight) + self.bias

2.2 等变性的数学验证

让我们验证这个设计如何满足等变性。假设输入x旋转矩阵R，根据等变定义应有：

VNLayer(x @ R) ≈ VNLayer(x) @ R'

对于我们的实现，当输入旋转时：

rotated_x = torch.einsum('bic,cd->bid', x, rotation_matrix) rotated_output = vn_layer(rotated_x) # 根据权重设计，应有： true_rotated = torch.einsum('boc,cd->bod', vn_layer(x), rotation_matrix)

当权重被正确初始化时（如使用正交矩阵），两者差异应该很小。这种性质使得网络无需见过所有可能的旋转变体，就能正确处理旋转后的输入。

3. 与传统网络的性能对比

3.1 旋转鲁棒性实验

为了直观展示等变网络的优势，我们设计一个简单的点云分类实验：

1. 数据集：包含50类基本几何形状（立方体、球体等）的1000个样本
2. 任务：识别旋转后的形状类别（不变性任务）
3. 对比模型：
   • 基准模型：普通PointNet
   • 等变模型：VectorNeuron网络

实验结果如下表所示：

注意：虽然等变网络在原始数据上表现略低，但其旋转鲁棒性显著优于传统网络

3.2 计算开销分析

等变性带来的性能提升并非没有代价。向量神经元层的主要计算瓶颈在于：

1. 内存占用：权重张量从二维扩展到四维，显著增加参数量
2. 矩阵乘法：einsum操作比普通matmul更耗资源

实际部署时需要权衡的考虑因素：

• 对于小规模几何数据（如分子结构），等变网络通常是优选
• 对绝对旋转不敏感的任务（如点云分割），传统网络可能更高效
• 可以使用群等变卷积等技巧降低计算复杂度

4. 前沿扩展：SE(3)-等变网络

4.1 从SO(3)到SE(3)

前述向量神经元主要处理旋转（SO(3)群），而现实应用常需要同时处理旋转和平移（SE(3)群）。最新的SE(3)-Transformer等架构通过以下创新扩展了等变性：

1. 位置感知注意力：将相对位置编码纳入注意力机制
2. 向量场消息传递：在特征更新时保持等变性质
3. 标量-向量混合表示：同时处理不变和等变特征

class SE3Layer(nn.Module): def __init__(self, channels): super().__init__() # 标量部分处理不变特征 self.scalar_proj = nn.Linear(channels, channels) # 向量部分处理等变特征 self.vector_proj = VectorNeuronLayer(channels, channels) def forward(self, scalar_feats, vector_feats): new_scalar = self.scalar_proj(scalar_feats) new_vector = self.vector_proj(vector_feats) return new_scalar, new_vector

4.2 实际应用案例

等变网络已在多个领域展现独特优势：

• 分子动力学：预测蛋白质3D结构时保持物理对称性
• 机器人抓取：不同视角下的抓取姿态估计
• 医学影像：旋转无关的器官分割

在AlphaFold2等突破性成果中，等变网络组件扮演了关键角色。它们使模型能够自然地处理蛋白质骨架的刚体运动，而无需昂贵的数据增强。