做分子表示学习的朋友一定绕不开“旋转”这两个字。一个分子在三维空间里不管怎么转、怎么平移它的化学性质不变但它的原子坐标全变了。如果你的神经网络连这种最基本的物理事实都学不会那后面无论是预测结合能、生成构象还是做力场拟合都会非常吃力。这些年搞3D等变几何深度学习本质上就是在解决一个问题怎么让网络既吃三维几何信息又对坐标系的选取无动于衷。而注意力机制作为当下特征交互的主流算子在三维空间里能不能做到“等变”就直接决定了整个模型的地基牢不牢。这篇文章聚焦等变注意力机制也就是1.1.3节这个位置的内容。我会把标准注意力为什么在三维空间里失效、等变注意力怎么设计、里面那些球谐函数和Clebsch-Gordan乘积到底在干什么以及我自己在复现和调试中踩过的坑一次讲清楚。适合刚入坑等变图神经网络的研究生也适合想把分子表征模型做得更稳的工程师。1. 等变性的本质先搞懂“等变”和“不变”的区别1.1 什么叫等变和“不变”差在哪很多新手容易把等变equivariant和不变invariant混为一谈这俩概念差出去十万八千里。用大白话讲不变是说“你输入转个角度输出结果完全不变”比如分子的总能量、偶极矩的模长这些是标量你分子怎么转它数值不变。等变是说“你输入转个角度输出跟着转同样的角度”比如每个原子上的受力方向、各个原子坐标组成的向量场分子整体旋转后每个向量也必须跟着转。用一个带方向性的例子来想假设你要预测一个分子体系里每个原子的偶极矩贡献方向。你把分子在空间里旋转90度那个原子的偶极矩方向必须也旋转90度不能还是原来的方向。这就是等变要求。如果你只要求预测标量性质比如logP、溶解度那不变就够了但如果你要做原子级别的向量输出力、偶极矩片段、构象生成等变特征就是刚需。所以在网络设计里等变和不变并不是非此即彼的关系。实际上特征里既要有不变的标量部分也要有等变的向量/张量部分。前者告诉你“这个原子周围长什么样”后者告诉你“这个方向上的信息往哪里去”。等变注意力机制干的事就是让标量和等变张量在每一层都能正确交互同时保证整体输出遵守坐标变换规则。1.2 为什么三维分子表示里“必须”等变分子表示学习的输入无非是原子种类、坐标可能还有电荷、自旋等属性。坐标这个输入非常特殊它本身就是一个三维欧氏空间里的向量。你把它喂给普通神经网络相当于默认了输入坐标的“绝对数值”是有意义的。但物理世界不这么想问题分子不会因为你把坐标系原点挪了个位置或者把z轴转了30度就变成另一个分子。那能不能不做等变用数据增强来硬扛比如训练的时候随机旋转分子让网络“见过”各种朝向最后近似学会旋转不变性。这种做法在少量数据、简单任务上确实能撑一撑但问题在于数据增强只能让网络在采样过的旋转角度上表现好不可能穷尽SO(3)群里所有连续旋转。而且分子构象本身是有手性的随意旋转数据增强如果处理不好还可能破坏手性信息。更根本的问题在泛化性。实际应用中测试集的分子构象常常是能量最小化之后的构象和训练集的坐标分布有偏差。等变网络因为内部特征本身就携带变换规则遇到没见过的朝向也能立刻给出物理正确的响应而不是靠“见过类似的”。这也是为什么近几年的SOTA分子表征模型清一色地往等变结构上靠。等变注意力就是这其中的核心交互层。2. 标准注意力机制为何在三维空间“失灵”2.1 在旋转坐标系里做点积会发生什么先回忆一个最朴素的注意力公式给定查询 (Q) 和键 (K)相关性打分 (a_{ij} Q_i \cdot K_j)再通过softmax归一化得到权重最后用权重去加权求和对应的值 (V)。这套机制在序列、图像里表现很好因为输入特征本身是定义在固定的坐标系或像素网格里的。换句话说图像里的“向右”就是像素坐标轴的“向右”这个方向全局一致。但把分子放进去就麻烦了。假设每个原子的特征里包含一个向量分量比如局部坐标系下的某个方向两个原子各自有一个局部坐标朝向。当你把整个分子旋转一个角度后每个原子的向量特征都会跟着旋转。此时你去计算两个向量之间的点积会发现这个值在旋转前后是保持不变的因为点积对旋转是保持不变的。听起来不错但问题在于标准注意力完全没有利用这个几何关系。它只计算了标量相关性然后对值向量做普通的线性组合组合系数是标量注意力权重。如果你更进一步把边特征里同时塞进相对位移向量 (\vec{r}_{ij} \vec{x}_j - \vec{x}i)然后直接用 (\vec{r}{ij}) 和某个节点特征做点积来算注意力这个点积在旋转下是不变的所以标量权重旋转不变没问题。但你最终更新节点特征时如果还是像Transformer那样把邻居的原始向量特征直接加权求和那这个新得到的方向就和全局坐标系绑定死了旋转后它不会跟着分子一起转。等变性就没了。2.2 几何信息缺失与SO(3)群的困境标准注意力在三维分子数据上的第二个大问题是它默认“特征顺序”就是“邻居顺序”没有显式建模原子对之间的几何关系。序列模型里位置编码是把token顺序塞进网络的注意力可以靠位置编码学到相邻关系。但分子的原子没有天然顺序你把原子编号打乱分子不变。于是图神经网络用边连接把相对位置向量编码进边特征这方向对了可如果只用标量距离 (||\vec{r}_{ij}||) 当边特征分子旋转后这距离确实不变然而原子对的方向信息全丢了。单纯丢掉方向信息有时候是故意的如果目标性质是标量那确实够用。但一旦你需要在中间层维护等变张量比如用向量特征表示键的方向、局部偶极方向或者输出原子受力那么注意力就必须知道“这个向量特征相对于某个邻居的方向关系”。可方向是定义在SO(3)群作用下的非标量对象网络层在数学上就要具备相应的群表示变换能力。这里点出了本质困境SO(3)群不是一个小离散群它是连续的李群。你没法用“枚举旋转角度”的方式去让网络内化这个对称性必须在算子本身的结构里注入对称性。注意力机制这种“加权求和”的线性操作本身在群作用下有良好的变换性质——只要权重标量是旋转不变的且值特征是等变的加权求和之后就还是等变的。这其实就是设计等变注意力的理论支点。3. 等变注意力机制的设计核心从标量权重到等变更新3.1 构造等变注意力权重的常见做法既然注意力权重的标量性质天然旋转不变核心问题就变成了怎么让权重更充分地利用三维几何信息。最简单的方案是只用距离(a_{ij} \phi(d_{ij}))其中 (\phi) 是一个标量函数常用的是径向基函数展开比如高斯核加一个MLP。距离是旋转不变的所以权重合法。但这样太浪费了原子对之间的方向信息完全被丢弃。稍微进阶一点可以在计算权重时额外引入“等变向量之间的点积”作为几何先验。比如节点特征里有等变向量型特征 (\vec{h}i)键方向是单位向量 (\hat{r}{ij})那 (\vec{h}i \cdot \hat{r}{ij}) 就是一个旋转不变的标量它反映的是“节点局部方向与这对原子连线方向的对齐程度”。把这样的标量拼接上距离编码再过一个MLP就能得到信息量更丰富的注意力权重。这类做法在文献里很常见也和化学直觉吻合拿一个原子的p轨道方向向量与其邻居原子的键方向作比较这个重叠关系就是化学键形成的本质。我在实际复现中发现把这类方向对齐特征加进注意力权重模型对分子构象的判别能力提升明显尤其是对涉及方向性相互作用的性质例如氢键倾向比纯粹用距离编码稳定很多。3.2 值的等变线性组合等变性如何被保证注意力权重建好之后值的更新方式才是决定模型等变性的关键。这里有两种路线。第一种是纯向量路线假设节点特征里有向量分量 (\vec{v}_i)更新时[ \vec{v}i^{(new)} \frac{1}{Z_i} \sum{j} a_{ij} W_v \vec{v}_j ]因为 (a_{ij}) 是旋转不变的标量(W_v) 是作用在每个向量上的线性变换可以是 (w_0 \vec{v}_j) 这种标量缩放或者是矩阵乘法但矩阵必须满足和旋转操作对易的条件所以整个式子在分子坐标旋转 (R) 下满足 ( \vec{v}_i^{(new)} \to R\vec{v}_i^{(new)} )。这个推导非常简单把每个 (\vec{v}_j) 换成 (R\vec{v}_j)线性组合系数不变结果自然整体乘上 (R)。第二种是高阶张量路线特征里除了向量还有高阶球谐系数张量比如l2的5维分量。这个时候值的组合就不是简单的数乘了因为你希望不同阶数的张量特征交互最终生成更高阶的等变特征。这就必须用到张量积分解也就是下一小节要讲的Clebsch-Gordan乘积。等变性的保证仍然建立在“标量权重 等变张量线性组合”这个大框架上只不过组合算子从“乘以标量”换成了“按表示矩阵做不可约分解”。3.3 球谐函数与Clebsch-Gordan乘积等变张量特征的必要工具聊到高阶等变张量绕不开两个工具球谐函数Spherical Harmonics和Clebsch-GordanCG系数。球谐函数 (Y_l^m(\hat{r})) 是定义在单位球面上的函数族(l) 是阶数(m) 是磁量子数取值范围 (-l,\dots,l)。它有一个极好的性质在三维旋转下(2l1) 个 (Y_l^m) 之间按一个确定矩阵Wigner D矩阵线性变换。这意味着我们可以把一个函数在球面上的取值按l分层展开成一个个“旋转起来非常规矩”的向量块。l0就是标量球对称l1就是三维向量三个分量对应x/y/zl2是一个5维分量块对应d轨道的形状。把注意力机制里的键向量方向 (\hat{r}{ij}) 用球谐函数展开就得到了一个随方向旋转而等变的边特征块。查询特征和键特征也各有各的球谐展开块。接下来要让查询特征 (Y{l_1}) 和键向量展开 (Y_{l_2}) 交互得到输出特征 (Y_l)数学上就要做两个表示空间的张量积再分解成若干个不可约表示。CG系数就是这个分解的“坐标转换矩阵”它把两个角动量耦合之后的各项分量重新组合成 (l |l_1-l_2|,\dots,l_1l_2) 的若干个块。这在物理上很自然两个轨道的重叠积分、角动量耦合本来就是CG系数在管的事。在代码里你不需要手动去算这些系数成熟的库比如某个等变网络常用工具包已经帮你把各种阶数的张量积操作封装好了。但对使用者来说理解一点很关键不同阶数特征交互后会产生更高阶或更低阶的输出块你得在模型里明确保留哪些阶数。这个选择基本等于炼丹中的超参数。4. 一个可复现的最小等变注意力层4.1 输入特征设计理论聊太多容易飘这节给一个可以直接跑通的最小实现。我们做一个单层等变图注意力输入每节点特征分成两块标量块 (s_i \in \mathbb{R}^{F_s})比如元素类型embedding、电荷数向量块 (\vec{v}_i \in \mathbb{R}^{3 \times F_v})初始可以设成全零也可以从标量特征通过一个线性层“提升”出来。边特征用原子对相对向量 (\vec{d}{ij} \vec{x}j - \vec{x}i)。归一化得到方向单位向量 (\hat{r}{ij})距离 (d{ij}) 用高斯径向基函数编码成 (F{rad}) 维向量。这个设计虽然是简化版但已经具备等变注意力的完整逻辑标量是旋转不变的向量是等变的边信息里同时有旋转不变的距离和等变的方向。为了防止初始向量特征全是零导致注意力权重全相等可以先生成一个小的随机线性层把标量特征投影出一部分向量初始化或者干脆第一层计算权重时只依赖标量和距离。4.2 单头等变注意力的前向计算前向计算分五步第一步计算注意力相关性的标量部分[ e_{ij} \frac{(W_Q s_i)^\top (W_K s_j)}{\sqrt{F_s}} \phi_d(\text{RBF}(d_{ij})) \phi_a( (\vec{v}i \cdot \hat{r}{ij}) (\vec{v}j \cdot \hat{r}{ij}) ) ]这里 (W_Q, W_K) 是标量特征的线性层(\phi_d, \phi_a) 是两个小MLP。关键在于最后一项(\vec{v}i) 和 (\vec{v}j) 是等变向量(\hat{r}{ij}) 是方向单位向量它们的点积在旋转前后保持不变。这样整个 (e{ij}) 就是一个旋转不变的标量。第二步softmax归一化权重 (a_{ij} \frac{\exp(e_{ij})}{\sum_k \exp(e_{ik})})此步完全不影响等变性。第三步计算值的向量部分。这里可以做两件事一是把邻居的向量特征 (\vec{v}_j) 按权重求和[ \vec{v}i^{(attn)} \sum_j a{ij} (W_v \vec{v}_j) ]二是把邻居的标量特征通过一个线性层后乘以键方向单位向量得到一个“方向性贡献”[ \vec{u}i \sum_j a{ij} (W_u s_j) \odot \hat{r}_{ij} ]这个第二项才是三维几何里的精髓它让节点特征能“朝邻居所在方位”聚集信息等价于给节点赋予一个由局部邻域几何决定的方向偏好。等变性依然成立因为 (\hat{r}_{ij}) 随旋转而旋转权重是标量整体输出就随旋转而旋转。第四步更新节点的标量特征。标量特征更新相对自由可以用拼接了向量模长 (||\vec{v}_i^{(attn)}||)、聚合向量模长等不变量的MLP来做[ s_i^{(new)} s_i MLP([s_i, ||\vec{v}_i^{(attn)}||, ||\vec{u}i||, \sum_j a{ij} W_s s_j]) ]模长是旋转不变的所以标量通道的更新安全且信息丰富。第五步更新向量特征。直接把注意力聚合向量和方向贡献加权合并[ \vec{v}_i^{(new)} \vec{v}_i \alpha \vec{v}_i^{(attn)} \beta \vec{u}_i ]其中 (\alpha, \beta) 可以是可学习的标量门控参数。这步以后节点向量特征同时包含了“邻居原始方向信息”和“由几何关系诱导的方向信息”。整个层里没有任何一项需要“绝对坐标”做输入所有用到坐标的地方要么是相对位置要么是相对方向旋转后它们会按统一规则变换。因此每一层都是等变的堆叠起来整体等变。4.3 训练与等变性验证搭建完模型第一个要跑的测试不是下游任务而是等变性一致性测试。做法很简单给同一个分子随机生成几百个旋转用绕xyz轴的欧拉角或者四元数随机旋转把原始坐标过一遍模型得到向量特征再把旋转后坐标过一遍模型然后检查后者的向量特征是否等于前者的向量特征经过同样旋转后的结果。标量特征应该在旋转前后完全一致。这个测试极其重要。我在实际项目里遇到过一种情况模型训练出来性质预测没问题单体向量输出也是合理的但一测等变性误差高达 (10^{-2}) 量级。排查下来居然是在某个残差连接里把原始绝对坐标当作向量特征用了——坐标在SO(3)旋转下确实等变但平移之后坐标整体变掉SE(3)下的等变性里坐标这个分支需要特殊处理不能直接当作可学习向量特征来聚合。如果发现等变性测试不通过按照“先查坐标使用再查线性层再查归一化层最后查数值精度”的顺序排查。大多数问题出在三个地方用了BatchNorm处理向量特征它在训练时依赖batch统计量破坏了逐样本的等变性注意力权重计算中不小心用了绝对坐标激活函数作用在向量特征上时没有保持表示的线性分解结构。5. 主流模型中的等变注意力实现细节5.1 某经典等变Transformer的注意力形式先聊一个经典实现节点特征被拆成多个角动量阶数块 (\mathbf{h} {h^{(0)}, h^{(1)}, h^{(2)}, \dots})其中 (h^{(0)}) 是标量块(h^{(1)}) 是向量块更高阶是类似多极展开的张量块。边特征不是我们自己算的RBF而是把相对方向 (\hat{r}_{ij}) 做球谐展开得到一套随旋转等变的边函数基。注意力计算分为两个阶段。第一阶段先做“按注意力权重聚合传输消息”。这时边的消息不是单纯把邻居特征搬过来而是把邻居的查询特征和边方向做张量积查询特征 (Y_l) 和边特征 (Y_{l_{edge}}) 耦合生成消息块。这一步里CG系数把两个不同阶数表示的张量积分解到几个输出阶数上维度会比输入大很多模型再加一个线性的“张量压缩”把输出 (l) 通道数降到合适大小。第二阶段把预聚合的消息和节点自身的特征做内积来计算注意力权重。这里要注意预聚合消息里既包含邻居信息也包含方向信息和节点特征点乘后得出的标量就是旋转不变的相关性分数。整个过程没有显式调用MLP但表达能力不差因为向量空间的结构已经被CG乘积撑得很丰富。代价是参数量和内存消耗都不小尤其是高阶 (l) 一开特征维度膨胀得厉害。做分子体系时一般 (l_{max}) 选1够用少数需要描述很强方向性的任务比如某些氢键、金属配位才考虑选2。5.2 简化路线向量门控更新与等变图注意力不是每个任务都需要完整的张量积。对纯标量性质预测有一个轻量级思路节点特征保持标量但边符号里加入方向信息用多个不同权重矩阵对邻居的标量特征按方向“投票”。这个方案本质上等于在标准图注意力的边特征里加入足够丰富的旋转不变几何不变量距离、角度、二面角、向量点积通过MLP编码权重。另一种路线更像我们第4节那个迷你实现节点特征标量向量双通道等变更新靠向量通道的加权求和和标量门控。这里的边特征常用的是单位方向向量而不是完整球谐展开。它的表达能力比完整张量积弱一些因为标量通道和向量通道的交互方式受限但对于许多中小分子数据集这种轻量等变注意力已经能给任务带来明显增益而且训练速度快、显存占用小。我用这两种路线跑过同一个力场能量预测任务发现一个有趣的差别轻量等变注意力在小数据集上更容易收敛而且不容易过拟合完整张量积模型在大数据集上上限更高但需要精心设计正则和众多训练trick。如果刚接触这个领域我建议先从轻量等变注意力入手把等变的概念、代码实现和验证方法先跑通再往复杂结构进阶。5.3 不同实现之间的取舍与经验观察把几条路线放一起对比距离编码标准图注意力不变内存最省、训练最快但方向信息全丢等变输出基本做不到。轻量向量等变注意力标量向量通道中等参数量能输出一级等变张量方向信息有限适合中小分子性质预测。球谐张量积等变注意力高阶特征丰富方向建模精确但参数量大、开销高适合需要精确方向作用或者力/偶极等向量输出任务。从实际效果看在不同任务上三者差距并不像论文里展示的那么悬殊。很多标量预测任务把距离、角度的旋转不变不变量精心设计好配合轻量等变通道也能做到接近复杂模型的精度。真正拉开差距的是涉及向量输出的任务力场预测、构象生成、偶极矩预测。这类任务没有真正的等变通道效果会断崖式下跌。所以在工程选型上先列举下游任务到底需要哪阶输出。只需要能量、亲和力、分类基本等变强不变量即可。需要力或坐标更新等变张量通道必须保留并且要确保每一层的更新都严格等变。需要高阶性质比如极化率张量那就只能老老实实上完整张量积没有捷径可走。6. 常见问题与避坑指南6.1 数值精度与收敛小技巧等变网络对数值精度比普通网络敏感。球谐函数数值、CG系数矩阵、Wigner D矩阵这些都是浮点数运算如果库的dtype默认float32有时候会遇到旋转一致性测试误差在 (10^{-4}) 到 (10^{-3}) 量级。这个误差在训练过程中一般影响不大但在做严格等变性验证时不好看。建议关键图模块球谐展开、张量积统一用float64跑验证脚本训练再切回float32以省内存。我还遇到过一个小坑某些深度学习框架的实矩阵乘法对特殊结构矩阵比如稀疏的CG系数矩阵优化不佳导致张量积操作比预期慢很多。解决办法是预先计算好CG系数稀疏索引用索引方式实现张量积而不是构造稠密矩阵做乘法。在小分子体系上这个优化能把单层前向时间压缩一半还多。6.2 初始化与正则化等变网络里向量特征通道的初始化非常影响训练稳定性。常见做法是把向量头线性层的权重初始化成很小的值甚至零初始化向量更新分支。原因是如果一开始就有大的向量噪声注意力机制会在“随机方向”上学出偏好后续很难纠正。直接用零初始化向量更新分支相当于刚开始模型就是一个纯标量模型训练过程中再逐步“长”出向量表达能力这会稳定很多。正则化方面对高阶张量特征加权重衰减要小心。普通L2正则对所有参数一视同仁但等变网络里不同表示块的参数尺度相差很大统一施加正则可能导致高阶块被过度压扁模型退化成纯标量模型。我习惯按表示块的阶数分别设置weight_decay高阶块用小一点的惩罚。6.3 调试中容易忽视的坑第一个坑是维度标注。球谐特征的维度是 ((2l1) \times F)(F) 是特征通道数。做多头注意力时很多人习惯把特征维度直接切成多头结果把 ((2l1)) 这个实质的表示空间维度也拆了破坏表示结构。正确的做法是多头切分通道维 (F)每个头都保留完整的 ((2l1)) 表示空间。第二个坑是残差连接。等变特征残差连接不是问题但如果你把注意力权重直接作用在“未被线性变换”的原始特征上方向性可能被过度保留导致数值震荡。我在做某个跨平台分子表征项目时为了提升性能把残差权重调大发现训练曲线时好时坏最终是给残差路径加了一个可学习的标量门控问题立刻缓解。第三个坑是共享权重与位置偏差。分子经常有对称等价原子比如苯环上几个氢原子在空间旋转后可以互换。理论上等变网络输出也应该随原子编号对称但如果你在特征线性层里不小心引入了“原子序数”这种与几何对称无关的输入等价原子之间的输出就出现偏差。这种情况不会破坏等变性但会影响预测的物理合理性。要在预处理环节把对称等价原子做数据增强旋转和置换同时进行让网络学出一致性。第四个坑是BatchNorm。重申一遍标准BatchNorm对逐样本的等变性是破坏性的因为它引入了跨样本统计量。但也不是说不能用如果你在下游头里、标量分支里用BatchNorm那是允许的因为下游头输出本来就是标量不变特征不要求逐样本等变。只有作用在等变特征上的归一化层才必须用等变版本或者干脆不用。写在最后的个人体会等变注意力机制听起来高大上但落地起来核心思想非常朴素把“跟谁交互”“怎么交互”的权重做成旋转不变的标量把“交互的内容”做成跟着坐标系一起转的等变张量。只要这两点守住了网络层数与层之间怎么堆叠都很安全。我自己复现这类模型时最大的感触是不要急着上完整高阶张量积先实现一个小而正确的等变注意力层把旋转一致性测试跑通再逐级加复杂度。很多问题在轻量模型里暴露得更清晰修起来也容易。等变网络的调试有个特点一旦等变性坏了你后面再怎么调超参数也救不回来反之只要等变性稳住了剩下的基本都是常规的过拟合、参数量、学习率这些熟活儿。最后分享一个实用小技巧在模型里专门留一个“诊断头”每次训练前喂一个随机旋转过的batch输出三组统计量——标量输出最大偏差、向量输出最大偏差、坐标更新最大偏差。如果这三个偏差在训练初期就大于 (10^{-3})说明等变结构被什么东西偷偷破坏了不要犹豫立即停住排查。这个习惯帮我省掉了大量无效训练时间也推荐给你。
做分子表示学习的朋友一定绕不开“旋转”这两个字。一个分子在三维空间里不管怎么转、怎么平移它的化学性质不变但它的原子坐标全变了。如果你的神经网络连这种最基本的物理事实都学不会那后面无论是预测结合能、生成构象还是做力场拟合都会非常吃力。这些年搞3D等变几何深度学习本质上就是在解决一个问题怎么让网络既吃三维几何信息又对坐标系的选取无动于衷。而注意力机制作为当下特征交互的主流算子在三维空间里能不能做到“等变”就直接决定了整个模型的地基牢不牢。这篇文章聚焦等变注意力机制也就是1.1.3节这个位置的内容。我会把标准注意力为什么在三维空间里失效、等变注意力怎么设计、里面那些球谐函数和Clebsch-Gordan乘积到底在干什么以及我自己在复现和调试中踩过的坑一次讲清楚。适合刚入坑等变图神经网络的研究生也适合想把分子表征模型做得更稳的工程师。1. 等变性的本质先搞懂“等变”和“不变”的区别1.1 什么叫等变和“不变”差在哪很多新手容易把等变equivariant和不变invariant混为一谈这俩概念差出去十万八千里。用大白话讲不变是说“你输入转个角度输出结果完全不变”比如分子的总能量、偶极矩的模长这些是标量你分子怎么转它数值不变。等变是说“你输入转个角度输出跟着转同样的角度”比如每个原子上的受力方向、各个原子坐标组成的向量场分子整体旋转后每个向量也必须跟着转。用一个带方向性的例子来想假设你要预测一个分子体系里每个原子的偶极矩贡献方向。你把分子在空间里旋转90度那个原子的偶极矩方向必须也旋转90度不能还是原来的方向。这就是等变要求。如果你只要求预测标量性质比如logP、溶解度那不变就够了但如果你要做原子级别的向量输出力、偶极矩片段、构象生成等变特征就是刚需。所以在网络设计里等变和不变并不是非此即彼的关系。实际上特征里既要有不变的标量部分也要有等变的向量/张量部分。前者告诉你“这个原子周围长什么样”后者告诉你“这个方向上的信息往哪里去”。等变注意力机制干的事就是让标量和等变张量在每一层都能正确交互同时保证整体输出遵守坐标变换规则。1.2 为什么三维分子表示里“必须”等变分子表示学习的输入无非是原子种类、坐标可能还有电荷、自旋等属性。坐标这个输入非常特殊它本身就是一个三维欧氏空间里的向量。你把它喂给普通神经网络相当于默认了输入坐标的“绝对数值”是有意义的。但物理世界不这么想问题分子不会因为你把坐标系原点挪了个位置或者把z轴转了30度就变成另一个分子。那能不能不做等变用数据增强来硬扛比如训练的时候随机旋转分子让网络“见过”各种朝向最后近似学会旋转不变性。这种做法在少量数据、简单任务上确实能撑一撑但问题在于数据增强只能让网络在采样过的旋转角度上表现好不可能穷尽SO(3)群里所有连续旋转。而且分子构象本身是有手性的随意旋转数据增强如果处理不好还可能破坏手性信息。更根本的问题在泛化性。实际应用中测试集的分子构象常常是能量最小化之后的构象和训练集的坐标分布有偏差。等变网络因为内部特征本身就携带变换规则遇到没见过的朝向也能立刻给出物理正确的响应而不是靠“见过类似的”。这也是为什么近几年的SOTA分子表征模型清一色地往等变结构上靠。等变注意力就是这其中的核心交互层。2. 标准注意力机制为何在三维空间“失灵”2.1 在旋转坐标系里做点积会发生什么先回忆一个最朴素的注意力公式给定查询 (Q) 和键 (K)相关性打分 (a_{ij} Q_i \cdot K_j)再通过softmax归一化得到权重最后用权重去加权求和对应的值 (V)。这套机制在序列、图像里表现很好因为输入特征本身是定义在固定的坐标系或像素网格里的。换句话说图像里的“向右”就是像素坐标轴的“向右”这个方向全局一致。但把分子放进去就麻烦了。假设每个原子的特征里包含一个向量分量比如局部坐标系下的某个方向两个原子各自有一个局部坐标朝向。当你把整个分子旋转一个角度后每个原子的向量特征都会跟着旋转。此时你去计算两个向量之间的点积会发现这个值在旋转前后是保持不变的因为点积对旋转是保持不变的。听起来不错但问题在于标准注意力完全没有利用这个几何关系。它只计算了标量相关性然后对值向量做普通的线性组合组合系数是标量注意力权重。如果你更进一步把边特征里同时塞进相对位移向量 (\vec{r}_{ij} \vec{x}_j - \vec{x}i)然后直接用 (\vec{r}{ij}) 和某个节点特征做点积来算注意力这个点积在旋转下是不变的所以标量权重旋转不变没问题。但你最终更新节点特征时如果还是像Transformer那样把邻居的原始向量特征直接加权求和那这个新得到的方向就和全局坐标系绑定死了旋转后它不会跟着分子一起转。等变性就没了。2.2 几何信息缺失与SO(3)群的困境标准注意力在三维分子数据上的第二个大问题是它默认“特征顺序”就是“邻居顺序”没有显式建模原子对之间的几何关系。序列模型里位置编码是把token顺序塞进网络的注意力可以靠位置编码学到相邻关系。但分子的原子没有天然顺序你把原子编号打乱分子不变。于是图神经网络用边连接把相对位置向量编码进边特征这方向对了可如果只用标量距离 (||\vec{r}_{ij}||) 当边特征分子旋转后这距离确实不变然而原子对的方向信息全丢了。单纯丢掉方向信息有时候是故意的如果目标性质是标量那确实够用。但一旦你需要在中间层维护等变张量比如用向量特征表示键的方向、局部偶极方向或者输出原子受力那么注意力就必须知道“这个向量特征相对于某个邻居的方向关系”。可方向是定义在SO(3)群作用下的非标量对象网络层在数学上就要具备相应的群表示变换能力。这里点出了本质困境SO(3)群不是一个小离散群它是连续的李群。你没法用“枚举旋转角度”的方式去让网络内化这个对称性必须在算子本身的结构里注入对称性。注意力机制这种“加权求和”的线性操作本身在群作用下有良好的变换性质——只要权重标量是旋转不变的且值特征是等变的加权求和之后就还是等变的。这其实就是设计等变注意力的理论支点。3. 等变注意力机制的设计核心从标量权重到等变更新3.1 构造等变注意力权重的常见做法既然注意力权重的标量性质天然旋转不变核心问题就变成了怎么让权重更充分地利用三维几何信息。最简单的方案是只用距离(a_{ij} \phi(d_{ij}))其中 (\phi) 是一个标量函数常用的是径向基函数展开比如高斯核加一个MLP。距离是旋转不变的所以权重合法。但这样太浪费了原子对之间的方向信息完全被丢弃。稍微进阶一点可以在计算权重时额外引入“等变向量之间的点积”作为几何先验。比如节点特征里有等变向量型特征 (\vec{h}i)键方向是单位向量 (\hat{r}{ij})那 (\vec{h}i \cdot \hat{r}{ij}) 就是一个旋转不变的标量它反映的是“节点局部方向与这对原子连线方向的对齐程度”。把这样的标量拼接上距离编码再过一个MLP就能得到信息量更丰富的注意力权重。这类做法在文献里很常见也和化学直觉吻合拿一个原子的p轨道方向向量与其邻居原子的键方向作比较这个重叠关系就是化学键形成的本质。我在实际复现中发现把这类方向对齐特征加进注意力权重模型对分子构象的判别能力提升明显尤其是对涉及方向性相互作用的性质例如氢键倾向比纯粹用距离编码稳定很多。3.2 值的等变线性组合等变性如何被保证注意力权重建好之后值的更新方式才是决定模型等变性的关键。这里有两种路线。第一种是纯向量路线假设节点特征里有向量分量 (\vec{v}_i)更新时[ \vec{v}i^{(new)} \frac{1}{Z_i} \sum{j} a_{ij} W_v \vec{v}_j ]因为 (a_{ij}) 是旋转不变的标量(W_v) 是作用在每个向量上的线性变换可以是 (w_0 \vec{v}_j) 这种标量缩放或者是矩阵乘法但矩阵必须满足和旋转操作对易的条件所以整个式子在分子坐标旋转 (R) 下满足 ( \vec{v}_i^{(new)} \to R\vec{v}_i^{(new)} )。这个推导非常简单把每个 (\vec{v}_j) 换成 (R\vec{v}_j)线性组合系数不变结果自然整体乘上 (R)。第二种是高阶张量路线特征里除了向量还有高阶球谐系数张量比如l2的5维分量。这个时候值的组合就不是简单的数乘了因为你希望不同阶数的张量特征交互最终生成更高阶的等变特征。这就必须用到张量积分解也就是下一小节要讲的Clebsch-Gordan乘积。等变性的保证仍然建立在“标量权重 等变张量线性组合”这个大框架上只不过组合算子从“乘以标量”换成了“按表示矩阵做不可约分解”。3.3 球谐函数与Clebsch-Gordan乘积等变张量特征的必要工具聊到高阶等变张量绕不开两个工具球谐函数Spherical Harmonics和Clebsch-GordanCG系数。球谐函数 (Y_l^m(\hat{r})) 是定义在单位球面上的函数族(l) 是阶数(m) 是磁量子数取值范围 (-l,\dots,l)。它有一个极好的性质在三维旋转下(2l1) 个 (Y_l^m) 之间按一个确定矩阵Wigner D矩阵线性变换。这意味着我们可以把一个函数在球面上的取值按l分层展开成一个个“旋转起来非常规矩”的向量块。l0就是标量球对称l1就是三维向量三个分量对应x/y/zl2是一个5维分量块对应d轨道的形状。把注意力机制里的键向量方向 (\hat{r}{ij}) 用球谐函数展开就得到了一个随方向旋转而等变的边特征块。查询特征和键特征也各有各的球谐展开块。接下来要让查询特征 (Y{l_1}) 和键向量展开 (Y_{l_2}) 交互得到输出特征 (Y_l)数学上就要做两个表示空间的张量积再分解成若干个不可约表示。CG系数就是这个分解的“坐标转换矩阵”它把两个角动量耦合之后的各项分量重新组合成 (l |l_1-l_2|,\dots,l_1l_2) 的若干个块。这在物理上很自然两个轨道的重叠积分、角动量耦合本来就是CG系数在管的事。在代码里你不需要手动去算这些系数成熟的库比如某个等变网络常用工具包已经帮你把各种阶数的张量积操作封装好了。但对使用者来说理解一点很关键不同阶数特征交互后会产生更高阶或更低阶的输出块你得在模型里明确保留哪些阶数。这个选择基本等于炼丹中的超参数。4. 一个可复现的最小等变注意力层4.1 输入特征设计理论聊太多容易飘这节给一个可以直接跑通的最小实现。我们做一个单层等变图注意力输入每节点特征分成两块标量块 (s_i \in \mathbb{R}^{F_s})比如元素类型embedding、电荷数向量块 (\vec{v}_i \in \mathbb{R}^{3 \times F_v})初始可以设成全零也可以从标量特征通过一个线性层“提升”出来。边特征用原子对相对向量 (\vec{d}{ij} \vec{x}j - \vec{x}i)。归一化得到方向单位向量 (\hat{r}{ij})距离 (d{ij}) 用高斯径向基函数编码成 (F{rad}) 维向量。这个设计虽然是简化版但已经具备等变注意力的完整逻辑标量是旋转不变的向量是等变的边信息里同时有旋转不变的距离和等变的方向。为了防止初始向量特征全是零导致注意力权重全相等可以先生成一个小的随机线性层把标量特征投影出一部分向量初始化或者干脆第一层计算权重时只依赖标量和距离。4.2 单头等变注意力的前向计算前向计算分五步第一步计算注意力相关性的标量部分[ e_{ij} \frac{(W_Q s_i)^\top (W_K s_j)}{\sqrt{F_s}} \phi_d(\text{RBF}(d_{ij})) \phi_a( (\vec{v}i \cdot \hat{r}{ij}) (\vec{v}j \cdot \hat{r}{ij}) ) ]这里 (W_Q, W_K) 是标量特征的线性层(\phi_d, \phi_a) 是两个小MLP。关键在于最后一项(\vec{v}i) 和 (\vec{v}j) 是等变向量(\hat{r}{ij}) 是方向单位向量它们的点积在旋转前后保持不变。这样整个 (e{ij}) 就是一个旋转不变的标量。第二步softmax归一化权重 (a_{ij} \frac{\exp(e_{ij})}{\sum_k \exp(e_{ik})})此步完全不影响等变性。第三步计算值的向量部分。这里可以做两件事一是把邻居的向量特征 (\vec{v}_j) 按权重求和[ \vec{v}i^{(attn)} \sum_j a{ij} (W_v \vec{v}_j) ]二是把邻居的标量特征通过一个线性层后乘以键方向单位向量得到一个“方向性贡献”[ \vec{u}i \sum_j a{ij} (W_u s_j) \odot \hat{r}_{ij} ]这个第二项才是三维几何里的精髓它让节点特征能“朝邻居所在方位”聚集信息等价于给节点赋予一个由局部邻域几何决定的方向偏好。等变性依然成立因为 (\hat{r}_{ij}) 随旋转而旋转权重是标量整体输出就随旋转而旋转。第四步更新节点的标量特征。标量特征更新相对自由可以用拼接了向量模长 (||\vec{v}_i^{(attn)}||)、聚合向量模长等不变量的MLP来做[ s_i^{(new)} s_i MLP([s_i, ||\vec{v}_i^{(attn)}||, ||\vec{u}i||, \sum_j a{ij} W_s s_j]) ]模长是旋转不变的所以标量通道的更新安全且信息丰富。第五步更新向量特征。直接把注意力聚合向量和方向贡献加权合并[ \vec{v}_i^{(new)} \vec{v}_i \alpha \vec{v}_i^{(attn)} \beta \vec{u}_i ]其中 (\alpha, \beta) 可以是可学习的标量门控参数。这步以后节点向量特征同时包含了“邻居原始方向信息”和“由几何关系诱导的方向信息”。整个层里没有任何一项需要“绝对坐标”做输入所有用到坐标的地方要么是相对位置要么是相对方向旋转后它们会按统一规则变换。因此每一层都是等变的堆叠起来整体等变。4.3 训练与等变性验证搭建完模型第一个要跑的测试不是下游任务而是等变性一致性测试。做法很简单给同一个分子随机生成几百个旋转用绕xyz轴的欧拉角或者四元数随机旋转把原始坐标过一遍模型得到向量特征再把旋转后坐标过一遍模型然后检查后者的向量特征是否等于前者的向量特征经过同样旋转后的结果。标量特征应该在旋转前后完全一致。这个测试极其重要。我在实际项目里遇到过一种情况模型训练出来性质预测没问题单体向量输出也是合理的但一测等变性误差高达 (10^{-2}) 量级。排查下来居然是在某个残差连接里把原始绝对坐标当作向量特征用了——坐标在SO(3)旋转下确实等变但平移之后坐标整体变掉SE(3)下的等变性里坐标这个分支需要特殊处理不能直接当作可学习向量特征来聚合。如果发现等变性测试不通过按照“先查坐标使用再查线性层再查归一化层最后查数值精度”的顺序排查。大多数问题出在三个地方用了BatchNorm处理向量特征它在训练时依赖batch统计量破坏了逐样本的等变性注意力权重计算中不小心用了绝对坐标激活函数作用在向量特征上时没有保持表示的线性分解结构。5. 主流模型中的等变注意力实现细节5.1 某经典等变Transformer的注意力形式先聊一个经典实现节点特征被拆成多个角动量阶数块 (\mathbf{h} {h^{(0)}, h^{(1)}, h^{(2)}, \dots})其中 (h^{(0)}) 是标量块(h^{(1)}) 是向量块更高阶是类似多极展开的张量块。边特征不是我们自己算的RBF而是把相对方向 (\hat{r}_{ij}) 做球谐展开得到一套随旋转等变的边函数基。注意力计算分为两个阶段。第一阶段先做“按注意力权重聚合传输消息”。这时边的消息不是单纯把邻居特征搬过来而是把邻居的查询特征和边方向做张量积查询特征 (Y_l) 和边特征 (Y_{l_{edge}}) 耦合生成消息块。这一步里CG系数把两个不同阶数表示的张量积分解到几个输出阶数上维度会比输入大很多模型再加一个线性的“张量压缩”把输出 (l) 通道数降到合适大小。第二阶段把预聚合的消息和节点自身的特征做内积来计算注意力权重。这里要注意预聚合消息里既包含邻居信息也包含方向信息和节点特征点乘后得出的标量就是旋转不变的相关性分数。整个过程没有显式调用MLP但表达能力不差因为向量空间的结构已经被CG乘积撑得很丰富。代价是参数量和内存消耗都不小尤其是高阶 (l) 一开特征维度膨胀得厉害。做分子体系时一般 (l_{max}) 选1够用少数需要描述很强方向性的任务比如某些氢键、金属配位才考虑选2。5.2 简化路线向量门控更新与等变图注意力不是每个任务都需要完整的张量积。对纯标量性质预测有一个轻量级思路节点特征保持标量但边符号里加入方向信息用多个不同权重矩阵对邻居的标量特征按方向“投票”。这个方案本质上等于在标准图注意力的边特征里加入足够丰富的旋转不变几何不变量距离、角度、二面角、向量点积通过MLP编码权重。另一种路线更像我们第4节那个迷你实现节点特征标量向量双通道等变更新靠向量通道的加权求和和标量门控。这里的边特征常用的是单位方向向量而不是完整球谐展开。它的表达能力比完整张量积弱一些因为标量通道和向量通道的交互方式受限但对于许多中小分子数据集这种轻量等变注意力已经能给任务带来明显增益而且训练速度快、显存占用小。我用这两种路线跑过同一个力场能量预测任务发现一个有趣的差别轻量等变注意力在小数据集上更容易收敛而且不容易过拟合完整张量积模型在大数据集上上限更高但需要精心设计正则和众多训练trick。如果刚接触这个领域我建议先从轻量等变注意力入手把等变的概念、代码实现和验证方法先跑通再往复杂结构进阶。5.3 不同实现之间的取舍与经验观察把几条路线放一起对比距离编码标准图注意力不变内存最省、训练最快但方向信息全丢等变输出基本做不到。轻量向量等变注意力标量向量通道中等参数量能输出一级等变张量方向信息有限适合中小分子性质预测。球谐张量积等变注意力高阶特征丰富方向建模精确但参数量大、开销高适合需要精确方向作用或者力/偶极等向量输出任务。从实际效果看在不同任务上三者差距并不像论文里展示的那么悬殊。很多标量预测任务把距离、角度的旋转不变不变量精心设计好配合轻量等变通道也能做到接近复杂模型的精度。真正拉开差距的是涉及向量输出的任务力场预测、构象生成、偶极矩预测。这类任务没有真正的等变通道效果会断崖式下跌。所以在工程选型上先列举下游任务到底需要哪阶输出。只需要能量、亲和力、分类基本等变强不变量即可。需要力或坐标更新等变张量通道必须保留并且要确保每一层的更新都严格等变。需要高阶性质比如极化率张量那就只能老老实实上完整张量积没有捷径可走。6. 常见问题与避坑指南6.1 数值精度与收敛小技巧等变网络对数值精度比普通网络敏感。球谐函数数值、CG系数矩阵、Wigner D矩阵这些都是浮点数运算如果库的dtype默认float32有时候会遇到旋转一致性测试误差在 (10^{-4}) 到 (10^{-3}) 量级。这个误差在训练过程中一般影响不大但在做严格等变性验证时不好看。建议关键图模块球谐展开、张量积统一用float64跑验证脚本训练再切回float32以省内存。我还遇到过一个小坑某些深度学习框架的实矩阵乘法对特殊结构矩阵比如稀疏的CG系数矩阵优化不佳导致张量积操作比预期慢很多。解决办法是预先计算好CG系数稀疏索引用索引方式实现张量积而不是构造稠密矩阵做乘法。在小分子体系上这个优化能把单层前向时间压缩一半还多。6.2 初始化与正则化等变网络里向量特征通道的初始化非常影响训练稳定性。常见做法是把向量头线性层的权重初始化成很小的值甚至零初始化向量更新分支。原因是如果一开始就有大的向量噪声注意力机制会在“随机方向”上学出偏好后续很难纠正。直接用零初始化向量更新分支相当于刚开始模型就是一个纯标量模型训练过程中再逐步“长”出向量表达能力这会稳定很多。正则化方面对高阶张量特征加权重衰减要小心。普通L2正则对所有参数一视同仁但等变网络里不同表示块的参数尺度相差很大统一施加正则可能导致高阶块被过度压扁模型退化成纯标量模型。我习惯按表示块的阶数分别设置weight_decay高阶块用小一点的惩罚。6.3 调试中容易忽视的坑第一个坑是维度标注。球谐特征的维度是 ((2l1) \times F)(F) 是特征通道数。做多头注意力时很多人习惯把特征维度直接切成多头结果把 ((2l1)) 这个实质的表示空间维度也拆了破坏表示结构。正确的做法是多头切分通道维 (F)每个头都保留完整的 ((2l1)) 表示空间。第二个坑是残差连接。等变特征残差连接不是问题但如果你把注意力权重直接作用在“未被线性变换”的原始特征上方向性可能被过度保留导致数值震荡。我在做某个跨平台分子表征项目时为了提升性能把残差权重调大发现训练曲线时好时坏最终是给残差路径加了一个可学习的标量门控问题立刻缓解。第三个坑是共享权重与位置偏差。分子经常有对称等价原子比如苯环上几个氢原子在空间旋转后可以互换。理论上等变网络输出也应该随原子编号对称但如果你在特征线性层里不小心引入了“原子序数”这种与几何对称无关的输入等价原子之间的输出就出现偏差。这种情况不会破坏等变性但会影响预测的物理合理性。要在预处理环节把对称等价原子做数据增强旋转和置换同时进行让网络学出一致性。第四个坑是BatchNorm。重申一遍标准BatchNorm对逐样本的等变性是破坏性的因为它引入了跨样本统计量。但也不是说不能用如果你在下游头里、标量分支里用BatchNorm那是允许的因为下游头输出本来就是标量不变特征不要求逐样本等变。只有作用在等变特征上的归一化层才必须用等变版本或者干脆不用。写在最后的个人体会等变注意力机制听起来高大上但落地起来核心思想非常朴素把“跟谁交互”“怎么交互”的权重做成旋转不变的标量把“交互的内容”做成跟着坐标系一起转的等变张量。只要这两点守住了网络层数与层之间怎么堆叠都很安全。我自己复现这类模型时最大的感触是不要急着上完整高阶张量积先实现一个小而正确的等变注意力层把旋转一致性测试跑通再逐级加复杂度。很多问题在轻量模型里暴露得更清晰修起来也容易。等变网络的调试有个特点一旦等变性坏了你后面再怎么调超参数也救不回来反之只要等变性稳住了剩下的基本都是常规的过拟合、参数量、学习率这些熟活儿。最后分享一个实用小技巧在模型里专门留一个“诊断头”每次训练前喂一个随机旋转过的batch输出三组统计量——标量输出最大偏差、向量输出最大偏差、坐标更新最大偏差。如果这三个偏差在训练初期就大于 (10^{-3})说明等变结构被什么东西偷偷破坏了不要犹豫立即停住排查。这个习惯帮我省掉了大量无效训练时间也推荐给你。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
