SE(3)-Transformer 解析：用等变注意力解决三维旋转崩溃问题

点云也好，分子也好，蛋白质结构也好，只要你的数据活在三维空间里，就迟早会被同一个问题恶心到：模型在标准坐标系下训练得好好的，准确率也漂亮，可一旦把输入整体旋转个 90 度，或者换个原点重新归一化，预测结果就像被重新投胎了一样，完全不是那么回事。

我最早是在做一个点云分类项目时撞上这堵墙的。当时用图神经网络喂绝对坐标，训练精度很高，我用旋转增强做了一轮测试，准确率直接从 90% 掉到 60% 出头。那一刻我突然意识到，数据增强只是把一个错误的架构"打补丁"：你教会了网络应付训练时见过的那些旋转角度，但三维空间的旋转是连续无穷多个的，永远补不齐。

后来我认真读完了 SE(3)-Transformer：三维旋转-平移等变注意力网络，才真正理解了这个问题该从架构层面解决。这篇论文提出的是一个对 SE(3) 群（三维刚性变换群）天然等变的注意力网络：输入点云怎么旋转、怎么平移，输出特征就跟着怎么旋转、怎么平移，整个网络的预测规律完全不依赖坐标系的选择。它能处理点云分类、分子性质预测、蛋白质结构建模这类任务，也是后来一大批等变图神经网络（EGNN、Equiformer、MACE 等）的思想源头。

这篇文章我不打算复述论文完事，而是从一个动手实践者的角度，把等变性的动机、群论与不可约表示的数学地基、注意力机制的拆解、基于 e3nn 的复现细节，以及我踩过的几个坑全部摊开来讲。你不需要学过群论，我尽量用生活化的类比把公式讲成人话，但核心推导和代码片段我会保留——出了 bug 你才知道去哪里排查。

1. "转一下就不认识"：三维深度学习的坐标系诅咒

1.1 等变性是什么：从最直观的旋转实验说起

先看一张想象中的实验图：你有一个椅子的点云，网络输出一个特征向量 v，用来表示"椅背指向"这个方向。如果椅子绕竖直轴旋转了 30 度，椅背指向也旋转了 30 度，那么理想的网络输出应该是 R·v，也就是把原来的向量也跟着转 30 度。这种"输入怎么变，输出就怎么跟着变"的性质，数学上叫等变性（equivariance）：

f(Rx + t) = R·f(x) + t

这里 R 是旋转矩阵，t 是平移向量。如果任务只是给椅子分类，你其实只需要不变性（invariance）：f(Rx + t) = f(x)，旋转之后类别标签不变。但如果你要预测的是一把力的方向、一个原子的受力、一个残基的朝向，那么输出本身就是向量，它必须跟着输入一起旋转——这时候只有等变网络才能做到"物理上正确"。

请注意这里的区别：不变性是等变性的一种特例，你只需要把等变输出中的标量通道（l=0 特征）取出来，再做一层全局池化，就能得到严格旋转不变的预测。所以等变网络是一个更通用的框架，它既能做分类，也能做需要方向感知的回归。

1.2 数据增强救不了这个问题

有人会问：那我训练时把所有物体都旋转若干角度，做成增强数据，不就行了吗？我试过，结论是：它只能缓解症状，治不了病根。

第一个问题是旋转群是连续群。三维空间中旋转角度是无穷多个，你不可能枚举所有方向。常见做法是采样几十个甚至几百个角度，但无论多少，测试时总会出现一个没见过的角度，网络在这些角度上的表现完全不可控。

第二个问题更本质：增强训练只能让网络"统计学上"对旋转鲁棒，但无法保证"结构上"对旋转正确。对于分类任务，这通常够用；但对于力、速度、方向这类向量型输出，标准网络根本无法同时做到"把输入旋转 30 度后，输出的向量也跟着旋转 30 度"。你可以旋转输入，也必须旋转标签去训练，但一个用绝对坐标建模的 MLP 它的中间特征在旋转下会乱成一团，你没法期待它学会这么精确的变换规律。

换句话说，数据增强是在逼网络用有限样本去逼近一个连续的对称性，而等变网络是把对称性直接写进了网络结构里，让它在所有角度下都天生正确。这就像你教小孩认字，与其让他死记硬背一万种字体的"一"，不如直接告诉他"一"这个字的本质特征——后者的泛化能力是结构性的。

1.3 一个插曲：CSS 的 translate 和 SE(3) 的"平移"压根是两码事

写这篇文章前我看到热搜里挂着"css平移"和"图片平移会卡帧么"，觉得挺有意思，就多说一句。做前端的同学都知道，CSS 里用transform: translate()去移动元素，走的是合成器，流畅不卡帧；而用left/top去改位置，会触发 layout 重排，元素多了就容易掉帧。这里面的关键差异在于：transform告诉你"把元素在空间里挪个位置，这个元素本身的渲染结果不变"，而left/top是"重新计算一遍布局，所有依赖它的兄弟节点全部跟着变"。

SE(3) 里的"平移"其实也是这个味道：它是三维空间里的一种对称变换 x → x + t。一个等变网络在这种变换下，内部计算结构保持"流畅"——你平移整个输入，不过是把坐标系换了个参考点，网络的行为完全一致。而一个依赖绝对坐标的非等变网络，坐标系一换，就像 CSS 里改了left/top触发了全局重排，输出立刻"卡帧"式地乱跳。图片平移卡不卡帧取决于浏览器走不走合成层，三维网络平不平移供参考地"卡帧"则取决于它是不是等变的。当然，这只是一个直觉类比，SE(3) 的平移是数学上的群元素，不是渲染属性，但理解了这个类比，你就能明白为什么等变性对三维深度学习这么重要。

2. 数学地基：SO(3)、SE(3) 群与不可约表示

2.1 群与表示：网络如何"理解"旋转和平移

先过一遍最基本的术语。三维旋转的集合构成一个群，叫 SO(3)：所有满足 R^T R = I 且行列式为 1 的 3×3 矩阵。它是个李群，意思是旋转可以连续地从小到大变化。把旋转和平移放在一起，构成更大的群 SE(3)：(R, t)，作用在一个三维点上就是 x → Rx + t。

一个值得强调的事实：SE(3) 不是"旋转和翻译各管各的"那么简单。旋转与平移的组合顺序会影响结果——先旋转再平移，和先平移再旋转，得到的位置通常不一样。这是一个非交换群。这意味着你在设计网络时，不能简单地把"旋转处理"和"平移处理"两个模块串起来就完事，必须让所有层在完整的 (R, t) 变换下保持行为一致。

"表示"（representation）是群论里的核心概念，简单说就是：群元素怎么作用在一个向量空间上。对网络来说，每一层特征都可以看作某个表示空间里的向量。一个等变网络层本质上是一个满足约束的线性映射 W，它必须满足：

W(ρ_in(g)·x) = ρ_out(g)·W(x)

这里的 ρ_in、ρ_out 是输入和输出特征各自遵循的表示规则。如果你设计的所有层都满足这个约束，那么整个网络自动就是等变的。这就把一个"网络架构设计问题"转化成了"如何构造满足约束的算子"的问题，SE(3)-Transformer 所有技术细节都在围绕这个约束展开。

2.2 球谐函数：构造等变核的万能积木

有了表示的概念，下面要解决一个具体问题：怎么构造一个以向量为输入、以向量为输出、且满足等变约束的函数？

答案是球谐函数。我们可以把它类比成傅里叶展开：任何一个周期信号都能分解成不同频率的正弦波叠加；类似地，任何一个在单位球面上的函数，都能分解成不同度数 l 的球谐函数 Y_lm 的线性组合。球谐函数最关键的性质是：在旋转 R 下，度数为 l 的一组球谐函数只会在这组函数内部互相线性组合，不会泄漏到其他度数。这个变换矩阵叫 Wigner D 矩阵 D^l(R)。

这个性质太有用了。它意味着我们可以把度数为 l 的空间当作一个"等变的积木块"：网络上任意一层特征，都可以拆成若干个度数为 l 的不可约表示（irrep）的直和。l=0 是标量（旋转不变），l=1 是三维向量（跟着旋转），l=2 是五维的对称无迹二阶张量（能描述曲率类的局部形状），l 越大描述的形状越精细。每个 l 的维度是 2l+1。

在 e3nn 库中，你会经常看到这样的类型定义：

irreps = o3.Irreps("16x0e + 8x1o + 4x2e")

这表示特征里有 16 个标量通道（偶宇称 0e）、8 个向量通道（奇宇称 1o）、4 个二阶张量通道（偶宇称 2e）。这里的 e/o 是宇称（parity），描述特征在镜像翻转下如何变换，对于分子这类有镜像对称性的数据很重要。

2.3 张量积与 CG 系数：等变特征的组合规则

光有积木还不够，网络还需要把不同度数的特征组合起来。比如你要把 l=0 的原子类型信息和 l=1 的方向信息融合成一个 l=1 的输出，这就要用到张量积与 Clebsch-Gordan（CG）系数。

两个等变特征做张量积，结果通常不是单一度数的特征，而是多个度数的直和。规则很简单：l_1 和 l_2 做张量积，得到的所有度数范围在 |l_1 - l_2| 到 l_1 + l_2 之间。CG 系数就是这套分解中的固定权重，它告诉你哪些分量属于哪个 l_3。好消息是这些系数不需要你自己算，e3nn 里全部预计算好了，你只需要指定输入和输出的 irreps。

这整套机制实际上继承自 Tensor Field Network（TFN）。TFN 提出了一个核心思想：一个等变卷积核可以被分解为径向函数乘以球谐函数：

W(x) = ∑_l w_l(‖x‖) · Y_l(x̂)

其中 x̂ 是相对方向的单位向量，‖x‖ 是距离，w_l 是只用距离作为输入的 MLP。因为距离是平移不变的标量，方向经过球谐展开后又具备正确的旋转变换性质，所以整个卷积核天然是 SE(3) 等变的。SE(3)-Transformer 的注意力消息机制，本质就是在这个框架上叠加了注意力权重。

3. 机制拆解：SE(3)-Transformer 的等变注意力到底怎么算

3.1 标准 Transformer 在三维空间遇到了什么麻烦

先回忆一下标准 Transformer 的注意力：attention = softmax(QK^T / √d)·V。这个东西放在自然语言处理里没问题，但放到三维空间，第一个要解决的问题就是位置编码。

如果你直接拿绝对坐标当位置编码，旋转一下输入，所有位置编码全变了，注意力权重也会变，网络输出自然跟着乱套。即便你用相对坐标，也要小心：相对向量既包含距离也包含方向，方向在旋转下会变化，如果你不对它做特殊处理，注意力权重同样不具备旋转不变性。

SE(3)-Transformer 的解决思路非常清醒，它把问题拆成了两半：注意力权重只依赖旋转不变量，消息函数用等变机制构造。这样旋转根本不会影响"谁注意谁"的权重分布，而消息本身又携带了方向信息，在旋转下能正确地跟着转。

3.2 注意力权重：全部由旋转不变量算出

具体来说，论文把注意力权重的计算限制在了"类型-0"特征上，也就是那些旋转不变的标量通道。过程大概是这样的：

• 对节点 i，用一个线性层从它的不变特征算出 query：φ_i；
• 对节点 j，用线性层从它的不变特征、再叠加上一个只依赖两节点距离 r_ij 的径向 MLP 输出，算出 key：ψ_ij；
• 注意力分数就是两者的点积再除以 √d，然后对邻居做 softmax。

这个设计最妙的地方是：点积运算是旋转不变的，距离也一样。所以注意力权重 α_ij 在输入整体旋转、平移后完全不变。换句话说，旋转一个分子，注意力热力图不会有一丝变化——"谁关注谁"这件事与坐标系无关。

你可能会问：为什么不直接对向量特征做内积？因为两个向量做内积确实也是旋转不变的，但问题在于生成这两个向量的线性层必须满足等变约束，而且用向量特征算注意力会让实现变复杂、数值上也不够稳定。论文选择了只在标量通道上算注意力，既保证了不变量，又大幅降低了实现难度。这个设计决策在我看来非常务实。另外再补一句，注意力权重保持"不变量"还有一个额外收益：你可以直接把它当成可解释性工具，比如在蛋白质任务里看哪些残基在互相注意，这个热力图不会因为结构被旋转而改变。

3.3 消息传递：球谐张量积如何保证等变

注意力告诉你"该多重视哪些邻居"，但真正把邻居信息传过来的是消息函数。SE(3)-Transformer 的消息构造沿用了 TFN 的卷积思路：

1. 对源节点 j 的特征做一个线性变换：v = W_V·h_j；
2. 计算边缘方向向量 x̂_ij 的球谐展开 Y_l(x̂_ij)；
3. 用张量积把 v 和 Y 组合成消息，组合权重由只依赖距离 r_ij 的径向 MLP 给出；
4. 把注意力权重 α_ij 乘到消息上，再按目标节点聚合，最后残差连接到本身的特征上。

这个流程每一部都满足等变约束：方向向量经过球谐展开后按 Wigner D 矩阵变换，张量积用 CG 系数做组合，径向权重是旋转不变量，平移则根本不起作用——因为所有输入都是相对量。你把整个点云平移一段距离，边缘向量和距离完全不变，注意力和消息都不变，节点特征当然也不变；你旋转点云，节点特征就会按照其所属度数正确地旋转。

值得注意的另一个细节是：论文在实际模型中，除了注意力消息路径，有时还会加一条不带注意力的等变卷积路径。我个人的理解是，这相当于给模型一个"保底"的信息通路——注意力在训练初期可能非常不平滑，保底卷积能维持稳定的梯度流。在多个任务上，加这条路径都能提升稳定性，下面的踩坑部分我会再展开。

3.4 等变非线性与归一化：被多数教程忽略的细节

如果你只照抄"张量积 + 注意力"的思路，很快就会发现一个尴尬问题：注意力层的输出要经过非线性激活函数，但 ReLU 用在向量组件上是错的。

原因很简单：ReLU 是逐元素操作，它对特征向量的每个分量做 max(x, 0)。但度数为 l 的等变特征，它的各个分量在旋转下会互相混合（通过 Wigner D 矩阵）。你单独给某个分量加 ReLU，等于在一个"非自然坐标系"上做了截断，旋转对称性就被破坏了。训练时你甚至能观察到网络"学聪明了"——它会把向量通道全部训练成 0，绕开这个错误的非线性，但这样等变特征就完全退化了。

正确做法有两种。最常用的是门控非线性（gated nonlinearity）：先用标量通道（0e）过一遍 sigmoid，得到每个通道的门控值，再把这个门控值乘到对应的高阶特征上。e3nn 里有现成的Gate模块，它会自动处理哪些通道当门控、哪些通道被门控，甚至会把奇宇称特征用奇偶规则分开处理，避免镜像对称性被破坏。另一种方案是基于范数的非线性：每个 irrep 的范数是旋转不变的，你可以对这个范数做任意 MLP，再把缩放因子乘回原特征。两种我都用过，门控在大多数任务上表现更稳定。

归一化同样是个坑。BatchNorm 是逐通道统计的，它无法感知"哪些分量属于同一个 irrep"，直接用在等变特征上会引入坐标依赖。我自己一般要么跳过归一化，要么用 e3nn 提供的NormActivation这类基于范数的方案。这个细节很多人第一次写等变网络时都会掉进去，后面我详细讲。

4. 基于 e3nn 的复现：从 Irreps 到完整网络

4.1 准备工作与库选型

动手实现 SE(3)-Transformer，最省力的路径是用 e3nn。这是一个基于 PyTorch 的等变神经网络库，o3.Irreps管理特征类型，o3.TensorProduct处理张量积和 CG 系数，o3.SphericalHarmonics计算球谐展开，你不需要自己写任何 Wigner D 矩阵的逻辑。

官方参考实现是 Fabian Fuchs 等人的 JAX 版本，github 上也有 lucidrains 的 PyTorch 移植版。我个人的建议是：你可以参考官方代码理解细节，但要在自己任务上做定制，最好还是基于 e3nn 自己搭。原因是这两个仓库都绑定了一些实验设置，比如特征初始化、数据增强策略，直接移植到新任务反而束手束脚。邻域图构建可以直接用 torch_geometric，radius_graph一步搞定。

4.2 Irreps 设计和输入特征处理

这是最容易迷茫的一步：我的特征该设成 16x0e 还是 32x0e + 8x1o + 4x2e？

我的经验法则是：标量通道管"这里大概是什么"，向量通道管"这里的朝向是什么"，二阶张量通道管"这里的局部形状是弯曲还是平直"。对于大多数任务，从16x0e + 8x1o + 4x2e起步是稳妥的。如果数据量小，可以换成8x0e + 4x1o + 2x2e。如果任务是性质预测（只需要标量标签），甚至可以只保留16x0e + 8x1o，l_max 设为 1。

输入特征上有个细节需要注意：绝对坐标永远不要直接作为节点特征喂进去。在 SE(3)-Transformer 里，坐标只是用来计算相对位置、距离和方向的，节点本身输入的应该是语义特征（原子类型、点云的法向量、类别嵌入等）。论文里的惯例是先把语义特征映射成标量嵌入，然后让第一层的张量积自动生成高阶等变特征。向量通道可以初始化为 0，网络会自己学会让它们"长出方向"来。我第一次写的时候把坐标直接拼进特征里，结果等变性被一个不该出现的绝对坐标彻底破坏，这个低级错误希望大家别犯。

4.3 邻域图和边缘特征：半径选择与球谐预计算

构建邻域图是决定计算量的关键环节。对蛋白质 Cα 结构，我一般用 10-12 Å 的半径截断，折叠蛋白中这个范围内大约有 20-30 个邻居残基，信息量很充分。对坐标已经归一化到 [0,1] 区间的点云，半径取 0.1-0.2，或者直接用 k-NN（k=16-32）。这里有一个权衡：半径太小会失去全局上下文，网络成了"盲人摸象"；半径太大，远距离节点的注意力会被稀释，而且边数爆炸，显存直接告急。

边缘特征有两个：距离 r_ij 和单位方向向量 x̂_ij。注意，方向向量不需要存成三维坐标，你要做的是把 x̂_ij 代入球谐函数，得到它在 l=0 到 l_max 下的所有球谐分量，维度是 (l_max+1)^2。这个值只依赖于方向和 l_max，跟层数无关，所以强烈建议在数据预处理阶段就把所有边的球谐展开算好存起来，训练时直接查表，能省下一大块时间。我遇到过有人在每个前向过程里重新算一遍球谐，速度慢了近一倍，完全没必要。

4.4 等变注意力层的核心代码骨架

下面给一个简化但结构完整的代码骨架，去掉了多头拆分的细节，方便你抓到核心算子：

import torch from torch import nn from torch_geometric.utils import scatter_softmax, scatter_sum from e3nn import o3 from e3nn.nn import FullyConnectedNet class SE3AttentionLayer(nn.Module): def __init__(self, irreps_in, irreps_out, l_max, num_heads=8): super().__init__() self.num_heads = num_heads sh_irreps = o3.Irreps.spherical_harmonics(l_max) # 1) query/key 只使用不变（0e）通道 self.query_net = nn.Sequential( o3.Linear(irreps_in, o3.Irreps("32x0e")), nn.SiLU(), o3.Linear(o3.Irreps("32x0e"), o3.Irreps(f"{num_heads}x0e"))) self.key_net = nn.Sequential( o3.Linear(irreps_in, o3.Irreps("32x0e")), nn.SiLU(), o3.Linear(o3.Irreps("32x0e"), o3.Irreps(f"{num_heads}x0e"))) # 2) 径向网络：距离 -> 张量积权重 self.radial_net = FullyConnectedNet( [1, 64, 64, sh_irreps.dim * irreps_in.num_irreps], act=nn.SiLU()) # 3) 消息张量积：源特征 ⊗ 球谐 -> 输出特征 self.tp = o3.TensorProduct( irreps_in, sh_irreps, irreps_out, shared_weights=False, internal_weights=False) self.out_linear = o3.Linear(irreps_out, irreps_out) def forward(self, x, pos, edge_index): src, dst = edge_index vec = pos[dst] - pos[src] # 相对位置，平移不变 dist = vec.norm(dim=-1, keepdim=True) + 1e-8 direction = vec / dist Y = o3.SphericalHarmonics( o3.Irreps.spherical_harmonics(self.l_max), direction, normalize=True) # 注意力权重（旋转不变量） q = self.query_net(x)[src] # (E, H) k = self.key_net(x)[dst] # (E, H) attn = (q * k).sum(-1, keepdim=True) / (self.num_heads ** 0.5) attn = scatter_softmax(attn, dst) # 按目标节点 softmax # 带权重张量积消息 w = self.radial_net(dist) # (E, C1*C2) msg = self.tp(x[dst], Y, w) # (E, irreps_out.dim) msg = msg * attn # 聚合与残差 out = scatter_sum(msg, dst, dim=0, dim_size=x.size(0)) return x + self.out_linear(out)

有几个地方要提醒。第一，o3.TensorProduct在internal_weights=False时，权重参数需要从外部传入，具体形状取决于指令数量，我这里用radial_net生成的通道数对应了输入特征类型数与球谐展开维度数的乘积，实际使用时要按 e3nn 版本的 API 核对一下维数对不上常见报错，不要慌，看错误信息里提示的 shape 去调最后那个64就行。第二，scatter_softmax和scatter_sum来自 torch_geometric，如果你不想引入这个依赖，自己写 for 循环聚合也可以，但会慢不少。第三，我这里省略了等变非线性层和层归一化，真实网络里通常在每层之间加Gate。整套完整的 SE(3)-Transformer 就是堆叠若干个这样的层，最后根据任务决定输出方式：分类取 0e 通道做全局池化，回归力则直接取 1o 通道。

4.5 超参数、显存估算与训练配置

下面是我在不同任务上常用的参数范围，直接抄作业基本不会翻车：

显存估算有个很实用的公式：消息张量是最大的内存大户，它的大小近似等于 边数 × 输出 irreps 维度 × 4 字节 × 多头系数。举个例子，1024 个点，平均 20 个邻居，边数约 2 万，输出特征是 16x0e + 8x1o + 4x2e，维度是 16×1 + 8×3 + 4×5 = 60，那么消息张量约 2 万 × 60 × 4 字节 ≈ 4.8 MB。看起来不大，但张量积中间过程会为每条指令额外分配缓冲区，再加回传图的存储，实际占用往往是这个数字的 5-10 倍。所以如果你的边数上万、l_max 又高，OOM 几乎是必然的。下面我的几个踩坑记录就是围绕这些问题展开的。

5. 复现实战：我踩过的四个坑

5.1 高 l 阶特征与显存爆炸

我第一次在自己数据集上复现时，贪心地把 irreps 设成了64x0e + 32x1o + 16x2e + 8x3o，l_max=3，还天真地以为"更高阶=更强表达力"。结果模型在 batch size=2 的时候就 OOM 了，NVIDIA 直接报 "CUDA out of memory"，我一度以为是代码有幽灵泄漏。

排查过程的结论是：张量积的指令数是按输入特征类型和球谐类型乘积增长的，l_max 从 2 升到 3，球谐通道数从 9 变成 16，接近翻倍；再加上更多的指令路径，每一层的中间缓冲区数量暴增。显存爆炸不是线性的，是指数级逼近的。

解决办法我按优先级排序：第一，把 l_max 降到 2，这一步能立刻省掉一大半显存；第二，减少通道数，从 64x0e+32x1o 缩到 16x0e+8x1o，很多任务根本用不了那么多通道；第三，开混合精度训练（AMP），等变层的张量积是标准的 CUDA 算子，AMP 收益很直接；第四，batch size 降到 1 配合梯度累积。另外别忘了一件事：球谐展开一定要预计算存起来，别在前向里现算，否则不仅费显存还费时间。

5.2 注意力退化成 one-hot，训练死水一潭

第二个坑发生在我把注意力权重算错之后。大约第 5 个 epoch 左右，我发现验证集指标停滞不前，训练集 loss 也在一个高位上震荡。我把注意力矩阵打出来看了一眼，好家伙，几乎每个节点都把 90% 以上的权重压在了同一个邻居身上，注意力熵不到 0.3 bit。这种退化让网络失去了聚合多样性信息的能力，等于退化成只有一跳的消息传递。

我用工具箱从三个方向修。一是检查 query/key 的尺度：如果点积结果绝对值太大，softmax 就会饱和，这时给注意力分数乘一个温度系数（比如 0.5）能立竿见影。二是检查邻居数量：如果半径太小，每个节点只有 3-4 个邻居，softmax 本来就容易变陡，适当扩大半径让每个节点有 15 个以上的邻居会平滑很多。三是加 attention dropout，在注意力的分布上加一点随机扰动，迫使网络不要过度依赖单一路径。我个人的建议是把这个指标监控起来：训练时周期性打印平均注意力熵，一旦发现快速下降，优先调温度和邻居数量，而不是盲目加层。

5.3 尺度与坐标归一化的隐藏影响

SE(3)-Transformer 只对旋转和平移等变，对缩放可不等变。这个特性经常被忽略，但它直接决定了你的模型能不能在新数据上工作。径向 MLP 的输入是绝对距离，如果你训练时用的是埃（Å）作为单位，到了另一个数据集换成了纳米（nm），所有距离输入都变了，模型预测基本报废。这跟标准图神经网络还得重新做归一化是一个道理，但等变网络给了你一个错觉："反正我对平移旋转免疫，坐标怎么放都行"——对平移确实免疫，对尺度并不免疫。

更隐蔽的问题是坐标数值范围过大时的精度问题。float32 下，如果坐标量级到 1e6，vec.norm()时相对方向的分量精度会被截断，球谐展开的输出变成噪声，轻则 loss 波动，重则出现 NaN。我的习惯是：训练前统一把数据居中（减去质心），再按一个合理的物理尺度缩放。居中不影响等变性，纯粹是为了数值稳定；缩放则需要考虑任务本身的语义——比如对分子结构，强行把所有键长缩放到单位长度会破坏化学含义，正确做法是保持键长分布不变，只做全局一致的尺度对齐。

5.4 等变非线性用错，向量通道悄悄"退化"

这个坑我前面在机制部分预告过：在等变特征上直接使用逐元素 ReLU 会破坏等变性。但这里有个更隐蔽的现象值得单独说——网络不会直接报错，它只是默默把向量通道训练成 0，让你发现不了问题。

我当时的复现是：每层张量积之后接nn.ReLU()，训练指标居然还不错，没有崩溃。我一度以为"啊，原来 ReLU 也没事嘛"。直到我检查了各层特征的范数统计，才发现 1o 和 2e 通道的范数在几轮迭代后趋于 0，整个模型事实上退化成了纯标量网络。方向信息根本没被利用，只是被"安全地"忽略了。

正确做法是用 e3nn 的Gate模块。它在内部会把特征分成三部分：标量部分用非线性激活后再作为门控值，门控值乘到高阶特征上。注意一个细节：门控值必须来自偶宇称的标量（0e），不能来自奇宇称标量，否则在镜像对称下门的符号会翻转，整个特征的宇称性质就错了。如果你不想用门控，也可以用范数激活：对每个 irrep 计算旋转不变的范数，用 MLP 生成缩放因子，再乘回去。两种都试过后，门控在我的任务上收敛更快。设置好之后，一个简单的验证方法是把输入整体旋转 30 度，看模型的向量输出是不是也旋转了 30 度——如果只有标量输出正确、向量输出对不上，多半是某个非线性或归一化环节破坏了等变性。

6. 应用场景与选型判断：什么时候非它不可

6.1 蛋白质结构：等变性的天然主场

蛋白质结构是最适合 SE(3)-Transformer 的场景之一，原因很朴素：蛋白质在自然界中不存在"标准朝向"，一个蛋白质从哪个方向看过去都是同一个蛋白质。如果你用依赖绝对坐标的模型，就必须做旋转增强，而且对于需要输出向量方向的任务（比如结合位点的朝向、配体进入通道的方向），纯增强根本无能为力。

在实际使用中，你可以把蛋白的 Cα 原子位置作为点云输入，每个残基的氨基酸类型映射成标量嵌入，然后堆叠几层 SE(3)-Transformer 层。次结构预测这类任务，取 0e 通道做分类头即可；如果要做残基接触界面预测，注意力热力图本身就是非常好的中间表示。我做过的实验里，SE(3)-Transformer 在残基级方向预测上的优势非常明显，因为它的 1o 通道天然输出一个随结构旋转的向量，而普通 GNN 需要额外加一个 MLP 从标量特征"硬猜"方向，泛化能力差很多。不过要提醒一句：蛋白质数据规模通常不小，Cα 数量上千时，边的规模会非常可观，注意控制半径截断和批次大小。

6.2 点云分类与分子性质预测

在点云分类上，论文在 ModelNet40 上做了实验，用远小于其他点云模型的参数量达到了有竞争力的准确率。我自己在室内点云分类上的体验是：SE(3)-Transformer 的优势在于不需要做旋转增强，训练数据少时泛化更稳。但也有局限——它不做全局形状归纳，纯靠局部等变消息传递，对细粒度类别区分有时不如专门设计的下采样结构，所以实际工程中我会把它和 PointNet++ 式的层级下采样结合用。

分子性质预测是另一个典型场景。对 QM9 上的能量预测，SE(3)-Transformer 完全可以用；但如果只是预测标量性质，其实 EGNN 这类更轻量的模型往往性价比更高。SE(3)-Transformer 真正的不可替代场景，是你需要预测力、加速度、偶极方向这类向量型物理量的时候——比如分子动力学中的原子受力，它输出的 1o 通道可以直接学习力的方向，物理上天然正确。

6.3 选型对照：SE(3)-Transformer vs EGNN vs Equiformer

很多人问我："那我是不是无脑用 SE(3)-Transformer 就行？"当然不是。我把常见几个模型的选型逻辑整理成一张表：

EGNN 的思路比 SE(3)-Transformer 更激进：它干脆不维护高阶等变特征，只在标量特征之外额外维护一个"坐标向量"，消息传递时同时更新坐标，从而以极低的成本获得 E(n) 等变性。它的代价是表达能力上限不如 SE(3)-Transformer——你无法用它直接表达二阶张量级的几何信息。我的经验是：如果数据量大、维度高、任务以分类和标量回归为主，EGNN 是更务实的选择；如果数据量中等、需要精细的方向推理或者模型需要输出等变向量场，SE(3)-Transformer 的高阶特征不可替代。Equiformer 则可以看成"SE(3)-Transformer 的升级版"，它在注意力计算中更充分地使用等变特征，同时用了一些训练技巧，在分子基准上通常比原版更准，但实现复杂度也更高。

6.4 我个人的经验总结

做了几轮对比实验之后，我现在的选型习惯是这样的：先问自己两个问题——输出是否需要等变？数据规模能否承受高阶特征的代价？如果两个答案都是否，直接上轻量级模型；如果输出需要方向或张量，或者你明确需要网络的注意力机制能对齐到物理结构，SE(3)-Transformer 这条路线就是最顺的。它虽然计算量大，却把"旋转鲁棒"从数据增强的玄学变成了结构的必然，这种确定性在工程上是巨大的安心感。

最后再分享一个调试小技巧：在训练第一个 epoch 时，取一个输入样本，分别做三次变换——不旋转、旋转 45 度、平移 (0.5, -1, 2)，把模型的中间特征输出存下来，检查它们是否满足预期的等变关系。这一步能在几分钟内暴露出 90% 的等变性 bug，比训练完再看验证集崩溃高效太多了。我每次搭新的等变网络，都会先跑这个"三连测试"，它已经帮我抓出过不止一次低级错误。