搞三维深度学习的人多少都遇到过这种糟心事一个点云分类模型训练时F1刷到0.92结果部署时点云被旋转了45度指标直接掉到0.7或者做分子性质预测同一个分子换了个朝向预测能量完全对不上。数据增强能缓解一部分但解决不了根本问题。SE(3)-Transformer就是冲着这个痛点来的——它让网络自带三维旋转和平移的等变性喂进去的物体怎么转、怎么挪输出都跟着变逻辑完全自洽。这篇文章我会从原理讲到实现再给出应用场景和排查经验适合正在做点云、分子建模、机器人抓取位姿估计这些方向的朋友参考。SE(3)-Transformer全称是SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Network发表于NeurIPS 2020是继Tensor Field Networks之后又一个把等变性和注意力机制结合得很漂亮的框架。它解决的核心问题可以概括成一句话如何让Transformer在三维空间中感知朝向同时输出随刚体变换协同变化。下面我按自己的理解把这个东西拆开讲透。1. 为什么三维网络必须考虑等变性1.1 对称性归纳偏置的核心价值先解释下SE(3)这个名字。SE(3)是Special Euclidean Group的缩写指三维空间中的旋转和平移组合总共包含6个自由度——3个旋转自由度加3个平移自由度。一个刚体在三维空间中的任意运动都可以拆成一个旋转加一个平移这就是SE(3)变换。等变equivariance这个词有些朋友可能听着陌生我举个例子。假设有个函数f输入是一个点云输出是每个点的特征向量。如果输入点云经过旋转R和平移t变成(Rxt)输出也能同步旋转并叠加平移也就是f(Rxt) R·f(x) t这个函数就对这个变换等变。注意区分一下不变性输入怎么折腾输出完全不变比如“点云属于飞机这个类别”的结论。等变性输入怎么折腾输出跟着变但变换关系是已知的、可预测的比如每点的法向量估计、每点的位移向量、抓取姿态的回归。为什么等变性重要核心原因是对称性归纳偏置。三维空间的物理规律本身就不依赖坐标系朝向——你把一个分子转180度它的能量物理上不应该变。如果网络结构本身不保证这一点模型就得靠海量数据去“硬学”这个规律。而带等变性约束的网络天然就把这个先验写进了结构里数据效率高得多泛化能力也强得多尤其是在训练数据覆盖不到的姿态上。1.2 数据增强为什么治标不治本有人会问那我用旋转增强不就行了把所有训练样本随机旋转让网络见得多一点。这个思路没错但有几个现实问题。第一数据增强本质上是在“打补丁”补得再密集也总有漏洞。比如你的训练分布里物体主要朝上放置增强虽然能补全姿态但网络容量的一部分永远在“对抗”这些随机变换而不是专注学习几何结构本身的判别特征。第二某些任务的输出本身带方向比如预测物体表面每个点受力的方向增强的时候你必须同步变换标签这没问题但网络如果结构不保证等变它对不同朝向的相同局部结构会学出不同的特征表示这会严重拖慢收敛速度。第三在精细任务上比如蛋白质口袋的配体结合姿态预测输入稍微偏一点输出可能就完全错了靠增强很难让模型达到“饿汉式”的精度。SE(3)-Transformer这类等变网络解决的就是直接从结构上根除这个问题——输出按照数学推导确定性的变换不需要学也不可能错。2. SE(3)-Transformer架构拆解2.1 从标准Transformer到几何Transformer标准Transformer的核心是自注意力。给定输入特征x_i计算query q_i、key k_j、value v_j然后打分权重a_ij softmax(q_i^T·k_j)输出y_i Σ_j a_ij v_j。问题在于这套机制完全没考虑几何关系。输入点的顺序置换下输出是会变的如果不加位置编码更不用说旋转平移了。SE(3)-Transformer的改造思路很直接在注意力机制中引入几何先验让权重计算与特征聚合方式都对SE(3)协同变换。具体来说模型的每个输入节点不仅有一个标量或向量的特征还有一个明确的空间坐标x_i。在一层内对于点i先找到它给定半径r内的邻居节点j把邻居的相对位置r_ij x_j - x_i作为几何线索。这个相对位置天然具有平移不变性——整体平移时所有点的x都减去同一个t相对位置完全不变。旋转时r_ij会跟着旋转。因此只要把相对位置以合适的方式编码进注意力就能保证注意力对平移不变对旋转等变。这里有个关键点注意力权重本身是旋转不变的。因为物理规则不应该依赖于朝向所以“点j相对于点i有多重要”这个决策本身不该随旋转而变。但聚合出来的特征必须随旋转同步旋转。这两个需求看似矛盾解决办法是各管各的权重是标量用旋转不变的特征计算特征是几何向量用旋转等变的线性变换来聚合。2.2 球谐函数与不可约表示要把旋转等变性揉进神经网络需要用到的数学工具是李群的不可约表示irreducible representations简称irreps。这个概念听着吓人用大白话说就是三维旋转群SO(3)的每一个作用可以被分解成一系列独立的“模块”标量l0、向量l1、二阶张量l2等等每一类都在SO(3)作用下独立旋转互不干扰。这就是张量场网络TFN的方法论把一个节点的特征分成若干个“类型”type每种类型带一个阶数l在旋转时按对应阶数的Wigner-D矩阵进行变换。球谐函数在这里扮演的角色相当于把空间方向信息编码成一组基。比如相对位置r_ij它可以分解成径向距离r和单位方向向量\hat{r}。方向向量按照不同阶数的球谐函数展开就得到了一个在各阶irrep类型间变换的向量。这样任意一个方向敏感的函数都可以表示成“径向函数 × 球谐函数”的线性组合形式而且组合系数只依赖距离是旋转不变的——这正好满足了上面说的“权重旋转不变”的需求。SE(3)-Transformer沿用了TFN的思路但把注意力机制的query-key-value都扩展到了irrep特征上对每个邻居j计算q_i和k_j之间的“几何感知”注意力打分打分函数里加入了基于相对位置的球谐嵌入。通过一个可学习的径向网络把相对距离r_ij映射成一组系数。输出特征的每个irrep块都通过一组“张量积”操作从邻居的值向量v_j中聚合而来而张量积本身有明确的旋转等变性质。这样一层层叠加整个网络的每一层都保持SE(3)等变最后的输出可以是标量比如每个节点的分类分数也可以是向量或高阶特征比如位移场、力矩预测视下游任务而定。2.3 从TFN到注意力聚合权力的再分配TFN的聚合方式和卷积类似邻域内每个邻居的权重由径向函数决定这种权重和方向无关在均匀点云上没问题但在结构不均匀、存在明显方向依赖的场景下就吃亏了。SE(3)-Transformer把邻居权重改成了数据依赖的注意力分数让网络自己决定“应该关注哪个方向的邻居”这就在保留等变性的前提下大大增强了表达能力。我个人的理解是TFN像是给每个邻居固定的“话语权”SE(3)-Transformer让网络学会根据局部几何动态分配话语权。这在分子体系中尤其重要——比如判断某个原子是否为反应中心注意力机制可以让模型更关注特定的配位方向而不是对周围所有原子一视同仁。3. 核心实现细节与实操要点3.1 网络结构的一层层拆解整个前向过程大致分几步我结合e3nn库的实现来梳理。第一步构建输入特征和irrep类型。输入的每个原子或点通常给一个初始标量特征比如原子序数的embedding再加上坐标向量x_i。在e3nn里这对应于一种“scalar vector”的mixed irrep特征表示。这里的重点是确定每个特征通道的阶数l的选择通常l_max取1或2再高的高阶特征参数量会爆炸而且小数据集上很难带来收益。第二步构造邻域图。以每个点为中心取半径r内或k近邻的点作为邻居。这一步和PointNet、DGCNN等点云方法的构图方式类似但注意半径的选择对模型性能影响很大——半径太小注意力没有足够的感受野等变特征学到的局部结构有限半径太大计算复杂度和内存消耗急剧上升。一般建议先用可视化工具观察数据集的平均点间距再确定半径的基数。第三步计算几何嵌入。对每个邻居对(i,j)计算相对向量r_ij分解出方向和距离。方向用球谐函数展开到l_max阶距离送到一个MLP里和注意力权重计算整合在一起。第四步注意力计算。对节点i和邻居j计算query特征q_i和key特征k_j的“等变点积”再加上基于相对距离和方向嵌入的偏置项得到每个邻居的未归一化注意力分数经过softmax得到权重。第五步特征聚合。每个邻居的value特征v_j先通过一个依赖注意力权重与几何嵌入的线性组合操作可理解为gated的等变线性层然后加权求和得到输出特征。需要特别注意的是特征类型和权重绑定方式。在e3nn的实操中每层都会定义一组线性权重权重数量由irrep的size和输入feature channels共同决定模型学习和训练参数量在100万以内是相对轻量级的结构。3.2 超参数选择与初始设定我跑过不少等变模型结论是SE(3)-Transformer对超参数不算特别敏感但有几个注意点特征通道数简单任务比如分子性质预测或小规模点云分类hidden features从32到64起步。通道数直接决定模型大小和训练速度在等变模型上通道数加倍计算量不止加倍因为irrep分块和tensor product操作是排列组合级的所以别一上来就128、256。l_max通常从1开始因为l1的向量特征已经能描述方向信息很多任务够用。l_max2在某些任务上有提升比如带曲率或二阶几何信息的场景但计算代价明显更高。最稳妥的做法是固定其它变量先对比l_max1和l_max2。层数3到5层足够再深收益递减。因为每一层的感受野通过邻居聚合扩展深层会面临信息过平滑的问题这和标准Transformer的层数经验类似。归一化等变网络里不能用标准的BatchNorm它只对标量特征有效而且平移不变性下统计量会有问题。实践中常用的是每层的权重norm和输出norm的组合或者用基于节点数目的归一化。e3nn里的normalization一般选norm效果稳定。3.3 数值稳定性问题等变网络有个经常踩的坑高阶irrep特征的值域控制不好训练一两个epoch就NaN。原因是不可约表示的特征范数会随着阶数增大而快速增长如果没有合适的normalization反向传播时梯度直接爆炸。我的做法是第一每层输出后做一次per-feature的范数裁剪或scale第二初始化时把最后一层的线性权重设得很小确保初始预测在零附近第三学习率别裸用1e-3起步尤其是分子数据上先用1e-4跑几百步看看loss曲线是否稳定再决定要不要调高。另外一点做SE(3)等变时平移一般是通过相对位置自动满足的不需要额外处理。但如果你的输入坐标数值很大比如以纳米为单位坐标值在几十量级建议先归一化坐标到单位尺度不然后期注意力计算的距离偏差会非常大softmax趋于饱和梯度也就消失了。4. 应用场景与实验验证4.1 分子建模与蛋白质结构SE(3)-Transformer在科学计算领域最受欢迎。分子数据天然具备SE(3)对称性一个分子在真空中的能量和原子坐标的绝对朝向无关只依赖相对几何构型。所以等变网络做分子性质预测、势能面拟合几乎是对口克制。在QM9数据集约13万个有机小分子上SE(3)-Transformer在能量和力预测任务上达到了当时State-of-the-art的水平尤其在势能预测任务上不需要给模型提供任何方向的先验输出的力自动满足旋转等变。这对分子动力学的模拟特别有价值——训练好一个等变势能网络就可以直接用它做任意朝向构型的能量和力的预测不用一遍遍做坐标对齐预处理。蛋白质结构领域SE(3)-Transformer后续衍生出SE(3) Diffusion、ProteinMPNN的变体、以及Flow Matching相关工作都离不开它“几何特征 注意力机制”这个范式。在配体结合位点预测、侧链打包、构象生成等任务上它的等变特征天然匹配了三维残基之间的相对朝向约束。4.2 点云处理与抓取位姿估计另一大类应用场景在机器人操作。比如从点云估计物体抓取位姿6D pose estimation网络输出每个接触点的位移方向或姿态四元数这些输出天然是SE(3)等变的——物体旋转多少预测出来的抓取坐标和方向都应该跟着旋转多少。用SE(3)-Transformer做特征提取骨干在训练数据量只有几万级别的场景下比普通PointNet和DGCNN能更快收敛泛化到未见过的物体姿态的稳定性也更好。点云分割任务上虽然不是最强选择这个任务对局部精细几何依赖很强共享权重的注意力和卷积差别不大但它在点云配准、关键点检测这类需要输出方向信息的回归任务上优势明显。我做过的经验是做一个局部表面法向量的预测任务时等变模型预测出的法向量在旋转后仍然与真值法向量对齐而普通MLP在不同旋转程度上的误差随角度增加线性恶化。4.3 什么时候不该用SE(3)-Transformer有必要泼一盆冷水不是所有三维任务都适合直接套SE(3)-Transformer。如果你的数据本身就是对齐好的比如图像深度图直接输入、姿态已经固定完全没有旋转偏移等变性带来的收益微乎其微反而会增加计算开销。另外训练数据极其充足千万级点云时普通网络靠数据增强也能学到近似等变这时等变网络的性能优势会被大幅稀释。所以选型前先问自己两个问题第一测试时输入是否可能出现任意朝向第二训练数据量是否不足以覆盖所有朝向两个至少占一个再考虑用等变网络。5. 常见问题与排查技巧实录5.1 等变性验证写个简单的测试脚本很多朋友训练完模型根本不确定自己的网络是不是真等变的——因为实现细节里有个小bug就可能导致等变性被破坏。强烈建议训练前先写一个等变性测试脚本import torch def check_equivariance(model, coords, features, rot, trans): # 原始输入输出 out1 model(coords, features) # 对输入做旋转和平移 coords2 coords rot.T trans out2 model(coords2, features) # 期望的输出也应做相同旋转变换 expected out1 rot.T err (out2 - expected).norm(dim-1).max().item() return err具体做的时候attention权重产生的是标量标量部分应该保持严格不变误差在1e-5量级由浮点误差决定l 1的特征部分应该按对应Wigner-D矩阵变换。如果误差是1e-2以上多半是坐标没有从相对输入中正确提取或者某个线性层忘了加上等变约束。5.2 训练不稳定的排查步骤我先说结论等变网络训练不稳定80%的情况出在数值范围控制上。如果在loss曲线上看到突然的spike赶紧做下面几步检查每一层输出特征的L2范数是不是随层数指数增长。检查softmax前的注意力logits是否出现极大值如果是说明几何嵌入尺度太大。把学习率降到1e-4跑100步看梯度范数是否为稳定的有界值。检查是否有NaN出现在特定l阶特征上——如果是说明该阶irrep的特征在tensor product传播过程中出现了数值溢出可以针对性做norm控制。还有一个反直觉的经验把batch size调小有时候反而更稳定。因为等变网络对每个样本独立作用batch内部没有共享统计量除了部分标量分支可能会经过batch norm大批量并没有分担梯度的好处反而同等学习率下梯度范数更大。5.3 内存与速度优化技巧SE(3)-Transformer的GPU显存消耗大头在注意力计算和tensor product上。最直接有效的提速方法有三个第一限制邻居数量。用半径构图比k近邻图好——半径构图可以在稀疏区域自动减少邻居数节省大量计算。如果必须用k近邻k别超过16。第二降低l_max。把l_max从2降到1显存和速度都能优化超过一半。如果l1的特征向量加上l0的标量已经能满足物理需求就不要追求高阶张量。第三部分任务可以把坐标预处理到固定中心。虽然平移等变性是内置的但把所有点云的质心对齐到原点可以减少相对坐标的数值范围让注意力logits的分布在更早的阶段变得可控。这只影响数值效率不影响等变性质。6. 写在最后的一点心得最后一个话题聊下选型与扩展现。SE(3)-Transformer不是万能的但它是理解“几何深度学习”这个方向的一把钥匙。我自己最开始接触的时候花了好几天才把irrep、球谐、Wigner-D这一串数学概念捋顺真正上手写代码后才发现这些概念在e3nn里都被封装成了现成模块核心还是要理解“什么量是旋转不变的什么量是旋转等变的”这一对关系。如果用一句话概括这篇的经验等变网络解决的不是“拟合精度”问题而是“泛化到未见朝向”的鲁棒性问题。在数据量有限、朝向敏感的三维任务中它的收益非常明显如果数据量和算力都充裕它也能显著驯化模型的优化曲面让训练来得更稳。对我个人而言最值钱的经验是理解了“先等变验证再训练调参”这套流程——它能帮你把网络结构的bug和优化层面的问题分离开少走很多弯路。后面如果大家有兴趣我可以再写一篇基于e3nn实现SE(3)-Transformer的最小可跑代码解析把层与层之间的张量维度怎么对齐、径向MLP怎么接这些细节逐一展开。