3D 图神经网络与 SE(3)/E(n) 等变架构实战指南从 SchNet 到 Equiformermaths-cs-ai-compendium 第 12 章第 5 节【免费下载链接】maths-cs-ai-compendiumBecome a cracked AI/ML researcher/engineer with this unconventional textbook covering maths, computing, and ML with intuition.项目地址: https://gitcode.com/GitHub_Trending/mat/maths-cs-ai-compendium本文是开源教科书 maths-cs-ai-compendium一个以直觉优先、从零讲解数学、计算机与 AI 的开源教材仓库中第 12 章「Graph Neural Networks」的第五篇技术专题。普通 GNN 只在抽象图上操作节点有特征、边编码连接关系但分子、蛋白质、晶体与点云等最有影响力的图数据都存在于物理 3D 空间中节点坐标承载着抽象图所忽略的关键信息。本文将从几何图的基本概念出发系统讲解 SE(3)/E(n) 等变性约束、SchNet、DimeNet、SphereNet、EGNN、Tensor Field Networks 等经典与前沿架构覆盖分子性质预测、蛋白质结构、材料科学、药物发现等核心应用并给出可直接在 JAX 中运行的等变性验证实验。读完本文你将掌握「不变架构」与「等变架构」两条 3D GNN 设计路线并能够亲手实现并验证旋转不变性与旋转等变性。在进入 3D 世界之前先回顾本仓库中已经建立的两块基石几何深度学习基础第 12 章第 1 节建立了对称性、群、群作用、不变性与等变性的统一框架CNN 利用图像的平移对称性Transformer 利用序列的置换对称性GNN 利用图的置换对称性。3D 图网络正是把这一框架应用到 3D 旋转与平移群上。图神经网络第 12 章第 3 节与图注意力网络第 12 章第 4 节讲解的消息传递框架、GCN、GraphSAGE、GIN、GAT 与 Graph Transformer是本文所有 3D 架构的基础骨架。几何图把图嵌入物理 3D 空间几何图geometric graph是嵌入 3D 空间的图。每个节点 $i$ 除了特征向量 $\mathbf{h}_i$ 之外还有一个位置 $\mathbf{r}i \in \mathbb{R}^3$。与显式化学键驱动的分子图不同几何图的边也可以由空间邻近性定义连接距离在截断半径 $r{\text{cut}}$ 以内的节点这使它可以统一处理分子、点云与晶体等不同数据形态分子节点是原子特征为元素类型、电荷等边是化学键3D 坐标 $\mathbf{r}_i$ 由量子力学计算或实验测定X 射线晶体学、冷冻电镜 cryo-EM获得。点云来自 LiDAR 或 3D 扫描仪见第 8 章与第 11 章每个点是一个节点带位置与可选特征颜色、强度边连接邻近点构成k 近邻图kNN graph或半径图。消息传递依赖的关键几何量有四类它们的对称性属性各不相同几何量定义平移不变旋转不变用途原子间距$d_{ij} |\mathbf{r}_i - \mathbf{r}_j|$✅✅分子形状的完整刻画与朝向无关键角节点 $i$ 处向量 $\mathbf{r}_j - \mathbf{r}_i$ 与 $\mathbf{r}_k - \mathbf{r}i$ 的夹角 $\theta{ijk}$✅✅捕捉超越成对距离的局部几何二面角扭转角平面 $(i,j,k)$ 与 $(j,k,l)$ 的夹角 $\phi_{ijkl}$✅✅结构在 3D 中的扭转蛋白质主链几何的关键相对位置向量$\mathbf{r}_{ij} \mathbf{r}_j - \mathbf{r}_i$✅❌携带方向信息需要等变架构才能使用最后一行是关键两个分子若原子间距完全相同则无论朝向如何都是同一形状但相对位置向量虽然平移不变却会随旋转而改变。要使用它们就必须采用等变而非仅仅不变的架构。SE(3) 与 E(n) 等变性3D 数据的基本对称性3D 物理数据的对称群是欧氏群 $E(3)$——包含所有旋转、反射与平移。其子群$SE(3)$特殊欧氏群只含旋转与平移排除反射。一个正确的 3D GNN 必须满足标量输出能量、结合亲和力平移不变所有原子整体平移相同向量预测不变标量输出旋转不变旋转分子能量不变向量/张量输出力、偶极矩旋转等变旋转分子预测的力向量应随之旋转同样的角度。形式化地对标量预测 $f$ 与旋转 $R \in SO(3)$有不变性$$f(R\mathbf{r}_1, R\mathbf{r}_2, \ldots) f(\mathbf{r}_1, \mathbf{r}_2, \ldots)$$而对向量预测 $\mathbf{F}$要求等变性$$\mathbf{F}(R\mathbf{r}_1, R\mathbf{r}_2, \ldots) R \cdot \mathbf{F}(\mathbf{r}_1, \mathbf{r}_2, \ldots)$$这两条约束直接对应第 12 章第 1 节中 invariance/equivariance 框架在 3D 旋转与平移群上的具体化。为什么必须内建这种对称性正如该节强调的把对称性构建进架构远比让网络从数据中学习它高效——一个能量预测会随旋转而改变的模型在物理上是错误的而强行让网络学会所有朝向的等价性会指数级扩大假设空间。由此产生两条设计路线不变架构invariant architectures只用不变几何特征距离、角度作为消息传递输入内部表示全部是标量。简单高效但无法在保持对称性的前提下输出向量。等变架构equivariant architectures在整个网络中维护向量乃至更高阶张量表示保证每一层都等变。表达能力更强能自然预测向量与张量但更复杂。SchNet基于距离的连续滤波器消息传递SchNetSchütt et al., 2017是最具奠基性的不变 3D GNN。其核心创新是连续滤波器卷积continuous filter convolution不再使用固定的边类型如分子 GNN 中的键型而是直接从原子间距生成消息滤波器。距离 $d_{ij}$ 首先通过**径向基函数RBF**展开为特征向量$$\text{RBF}(d_{ij}) \left[\exp\left(-\gamma_1 (d_{ij} - \mu_1)^2\right), \ldots, \exp\left(-\gamma_K (d_{ij} - \mu_K)^2\right)\right]$$每个基函数是以 $\mu_k$ 为中心、宽度为 $\gamma_k$ 的高斯函数。这相当于一个可学习的距离位置编码连续的 1D 距离被映射到高维特征空间网络可以在此空间中学习依赖距离的相互作用。中心 $\mu_k$ 通常从 0 到截断半径均匀分布$\gamma_k$ 控制每个基函数的感受宽度——高斯越窄对距离的分辨率越高但需要的基函数 $K$ 也越多。SchNet 从节点 $j$ 到节点 $i$ 的消息为$$\mathbf{m}{j \to i} \mathbf{h}j \odot W{\text{filter}}(\text{RBF}(d{ij}))$$其中 $W_{\text{filter}}$ 是将 RBF 展开映射为滤波器向量的 MLP$\odot$ 是逐元素相乘Hadamard 积见第 2 章矩阵。滤波器依赖距离因此近邻原子与远距原子产生不同交互逐元素乘法类似门控机制见第 6 章机器学习距离依赖的滤波器控制每个特征维度的信息通过量。由于 SchNet 只使用距离不变量整个模型自动对旋转和平移不变——无需任何特殊对称性处理这正是不变架构优雅之处。DimeNet 与 SphereNet引入角度与二面角仅有距离无法完整确定 3D 结构两个不同的分子构象可能具有完全相同的成对距离但键角不同即「距离几何歧义」问题。DimeNetGasteiger et al., 2020把键角引入消息传递。DimeNet 采用定向消息传递directional message passing消息沿有向边流动边 $(j \to i)$ 上的消息受到边 $(k \to j)$ 与 $(j \to i)$ 夹角的影响$$\mathbf{m}{kj \to ji} f\left(\mathbf{m}{kj}, d_{ji}, \theta_{kji}\right)$$角度 $\theta_{kji}$ 用球贝塞尔函数与球谐函数展开——这是球面上角信息最自然的基底正如 RBF 之于距离。这一设计让模型在保持不变性的同时获得方向信息。SphereNetLiu et al., 2022更进一步把二面角$\phi_{lkji}$ 也纳入进来捕捉完整的 3D 扭转结构。几何分辨率形成清晰的层级距离→ 成对邻近性角度→ 局部几何弯曲 vs 直线二面角→ 3D 扭转蛋白质主链、药物结合的关键每一层分辨率的提升都伴随计算代价距离是 $O(|E|)$角度是 $O(|E| \cdot k)$二面角是 $O(|E| \cdot k^2)$其中 $k$ 是平均度。实际应用中需要在几何分辨率与算力之间权衡这也是为什么许多生产级模型止步于角度级而很少全面使用二面角。EGNN等变位置更新EGNNSatorras et al., 2021采取等变路线每层不仅更新节点特征还更新节点位置并在整个过程中维持等变性。其更新为$$\mathbf{m}_{ij} \phi_e\left(\mathbf{h}i, \mathbf{h}j, d{ij}^2, a{ij}\right)$$$$\mathbf{r}_i \mathbf{r}i C \sum{j \neq i} (\mathbf{r}_i - \mathbf{r}j) \cdot \phi_r(\mathbf{m}{ij})$$$$\mathbf{h}_i \phi_h\left(\mathbf{h}i, \sum_j \mathbf{m}{ij}\right)$$关键在于位置更新节点位置被调整为相对位置向量 $(\mathbf{r}_i - \mathbf{r}_j)$ 的加权和。权重来自消息函数 $\phi_r$而该函数只依赖不变量特征与距离。这一构造可证明等变若所有输入位置被 $R$ 旋转则所有输出位置被同一个 $R$ 旋转——这正是 SE(3)-等变性要求。EGNN 的优雅之处在于它无需显式使用球谐函数或不可约表示就实现了等变性。相对位置向量承载方向信息不变的消息函数控制方向信息的使用方式。但这种简洁也有代价EGNN 只使用向量表示1 阶无法在没有扩展的情况下表示四极矩、应力张量等高阶张量。Tensor Field Networks 与高阶表示Tensor Field NetworksThomas et al., 2018及其后继者SE(3)-Transformers、MACE、Equiformer使用旋转群不可约表示irreducible representations的完整机制构建等变层。联系第 2 章的线性代数3D 中的旋转可以分解为由整数阶 $\ell$ 刻画的不可约分量$\ell 0$标量1 个分量不变。能量、电荷。$\ell 1$向量3 个分量像位置向量一样旋转。力、偶极矩。$\ell 2$秩 2 对称无迹张量5 个分量。四极矩、应力张量。更高阶 $\ell$捕捉日益复杂的角结构。这些称为球张量spherical tensors在旋转 $R$ 下通过Wigner-D 矩阵$D^\ell(R)$ 变换标量不变向量被 $R$ 旋转秩 2 张量被更复杂的矩阵旋转。球张量上的等变消息传递使用Clebsch-Gordan 张量积组合不同阶的特征$$(\mathbf{f}^{\ell_1} \otimes \mathbf{f}^{\ell_2})^{\ell_{\text{out}}} \sum_{m_1, m_2} C^{\ell_{\text{out}}, m_{\text{out}}}{\ell_1, m_1, \ell_2, m_2} \cdot f^{\ell_1}{m_1} \cdot f^{\ell_2}_{m_2}$$Clebsch-Gordan 系数 $C$ 是固定的数学常数确保张量积保持等变。这是 $SO(3)$ 等变意义下的「矩阵乘法」——正是这一操作把不同阶的张量特征组合起来并输出合法阶数的新张量。在此框架下涌现的现代模型各有侧重MACEBatatia et al., 2022使用高阶消息多个邻居特征的乘积用体序交互body-ordered interactions——2 体来自距离、3 体来自角度、多体来自张量积——以更少的消息传递层达到高精度高效捕捉复杂原子相互作用。EquiformerLiao Smidt, 2023把等变球张量特征与 Transformer 注意力机制见第 12 章第 4 节结合打造 SE(3) 等变的 Graph Transformer注意力分数由不变特征计算而值聚合在等变张量特征上进行——不变性负责「该注意谁」等变性负责「注意到的信息如何旋转」。应用从分子性质到药物发现3D GNN 最成熟的应用场景包括分子性质预测给定分子 3D 结构预测能量、力、偶极矩、HOMO-LUMO 能隙、毒性、溶解度等。基于量子化学数据集QM9、OC20训练的模型在多项性质上达到化学精度支撑数百万候选分子的虚拟筛选。分子动力学加速用量子力学密度泛函理论 DFT计算原子间力的代价极高$n$ 个电子约 $O(n^3)$。训练好的 3D GNN 在分子动力学模拟中可替代 DFT实现 $10^3$–$10^6$ 倍的加速且保持接近 DFT 的精度从而支持更大体系、更长时间尺度的模拟。蛋白质结构蛋白质主链可建模为几何图节点为残基边连接空间邻近残基。3D GNN 用于蛋白质功能预测、结合位点识别与蛋白质设计逆向折叠给定目标结构预测氨基酸序列。AlphaFold结合几何与图推理从序列预测蛋白质结构。材料科学与催化晶体材料具有周期性 3D 结构GNN 对重复单元胞建模并预测带隙、形成能、机械强度等性质。Open Catalyst ProjectOC20/OC22用 GNN 基准预测催化表面的吸附能加速可再生能源新催化剂的搜索。药物发现3D GNN 预测药物分子与靶蛋白的结合亲和力——结合亲和力取决于药物与蛋白结合口袋的 3D 形状互补性与化学相互作用。DiffDock这类模型将等变 GNN 与扩散模型见第 8 章结合预测结合姿态药物在蛋白口袋中的 3D 朝向。图生成从分析到创造以上所有架构都在分析已有图。图生成graph generation则创造新图设计具有目标性质的分子、为测试生成合成社交网络、提出新型蛋白质结构。这是图级预测的生成式对应物。挑战在于图是离散、变长、组合的。生成一个图意味着决定创建多少个节点、它们具有什么特征、连接哪些节点对。可能图空间随节点数超指数增长。主流方法有四类自回归生成一次生成一个节点或一条边。GraphRNNYou et al., 2018顺序生成图RNN 维护状态每步生成一个新节点并决定与哪些已有节点相连。生成顺序给本质上无序的图强加了一个人为序列BFS 序通过保持最近生成节点的相关性来缓解这一问题。基于 VAE 的生成用 GNN 编码器把图编码进连续潜空间再从采样潜向量解码新图。GraphVAE一次性生成概率邻接矩阵 $\hat{A} \in [0, 1]^{n \times n}$但开销为 $O(n^2)$ 且产生需要阈值的稠密输出。潜空间支持平滑插值在两个分子嵌入间移动可生成化学有效的中间结构。基于扩散的生成把扩散框架见第 8 章应用到图上。前向过程逐步向节点特征与边结构加噪反向过程学习去噪从噪声生成有效图。DiGressVignac et al., 2023对节点类型与边类型都应用离散扩散天然处理图数据的类别特性。面向分子的生成关键约束是化学有效性——生成分子必须遵守价键规则碳成 4 键、氧成 2 键等。Junction Tree VAEJT-VAE把分子分解为有效子结构环、链、官能团通过组装这些积木生成从构造上保证有效性。目标导向生成goal-directed generation针对特定性质优化生成对靶蛋白高结合亲和力、低毒性、良溶解度的分子。这以循环方式结合图生成与性质预测用 3D GNN 作为性质评估器生成 → 评估 → 精化。强化学习见第 6 章或贝叶斯优化引导化学空间搜索。值得一提的交叉点是DiffDockCorso et al., 2023它用 SE(3) 等变扩散预测药物分子如何对接进蛋白结合口袋。模型从随机放置去噪生成 3D 结合姿态药物相对蛋白的位置与朝向把本文件中的 3D 等变网络与第 8 章的扩散框架合二为一。动手实验用 JAX 验证不变性与等变性下面三个实验可在 Colab 或 notebook 中运行把本文的数学结论变成可运行的验证。仓库中该文件与其他章节一样使用 MathJax 渲染公式、Pyodide runner 提供网页内运行代码的环境见 mkdocs.yml因此这些代码可直接在教材在线版本中交互运行。实验 1不变消息传递层 旋转验证构建一个基于原子间距的简单不变 3D 消息传递层应用在小分子水H-O-H上并验证输出对旋转不变。import jax import jax.numpy as jnp # Water molecule: O at origin, two H atoms positions jnp.array([[0.0, 0.0, 0.0], # O [0.96, 0.0, 0.0], # H1 [-0.24, 0.93, 0.0]]) # H2 # Node features: [atomic number] features jnp.array([[8.0], [1.0], [1.0]]) # Compute pairwise distances (invariant) def pairwise_distances(pos): diff pos[:, None, :] - pos[None, :, :] return jnp.sqrt(jnp.sum(diff**2, axis-1) 1e-8) # Simple distance-based message passing def invariant_message_pass(features, positions): dists pairwise_distances(positions) # RBF expansion with 4 centres centres jnp.array([0.5, 1.0, 1.5, 2.0]) rbf jnp.exp(-5.0 * (dists[:, :, None] - centres[None, None, :]) ** 2) # Message: features weighted by distance-dependent filter messages jnp.einsum(ij,jd-id, rbf.sum(axis-1), features) return messages output1 invariant_message_pass(features, positions) # Rotate the molecule by 90 degrees around z-axis R jnp.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtypefloat) rotated_positions (R positions.T).T output2 invariant_message_pass(features, rotated_positions) print(fOriginal output:\n{output1}) print(f\nRotated output:\n{output2}) print(f\nInvariant: {jnp.allclose(output1, output2, atol1e-5)})代码中jnp.einsum(ij,jd-id, ...)完成「按距离权重聚合邻居特征」的求和rbf.sum(axis-1)把 4 个高斯基函数的输出折叠成单个标量权重模拟了 SchNet 滤波器的雏形。1e-8的常数用于避免零距离除零。实验 2键角计算 旋转不变性验证计算三个原子之间的键角并验证其旋转不变。import jax.numpy as jnp def bond_angle(r_i, r_j, r_k): Angle at node j between edges j-i and j-k. v1 r_i - r_j v2 r_k - r_j cos_angle jnp.dot(v1, v2) / (jnp.linalg.norm(v1) * jnp.linalg.norm(v2)) return jnp.arccos(jnp.clip(cos_angle, -1, 1)) # Three atoms r1 jnp.array([1.0, 0.0, 0.0]) r2 jnp.array([0.0, 0.0, 0.0]) r3 jnp.array([0.0, 1.0, 0.0]) angle_original bond_angle(r1, r2, r3) print(fOriginal angle: {jnp.degrees(angle_original):.1f}°) # Apply random rotation R jnp.array([[0.36, 0.48, -0.80], [-0.80, 0.60, 0.00], [0.48, 0.64, 0.60]]) r1_rot, r2_rot, r3_rot R r1, R r2, R r3 angle_rotated bond_angle(r1_rot, r2_rot, r3_rot) print(fRotated angle: {jnp.degrees(angle_rotated):.1f}°) print(fInvariant: {jnp.allclose(angle_original, angle_rotated, atol1e-4)})键角由向量点积与范数的比值即余弦值给出余弦值天然旋转不变因此角度也是旋转不变的。jnp.clip(cos_angle, -1, 1)防止浮点误差导致arccos越界。这正是 DimeNet 中角度信息进入消息传递时依赖的不变性来源。实验 3EGNN 式等变位置更新演示等变位置更新用距离加权的相对向量更新节点位置并验证「先旋转再更新」与「先更新再旋转」结果一致。import jax import jax.numpy as jnp def egnn_position_update(positions, features): Simple EGNN-style equivariant position update. n positions.shape[0] new_positions jnp.zeros_like(positions) for i in range(n): shift jnp.zeros(3) for j in range(n): if i ! j: r_ij positions[i] - positions[j] d_ij jnp.linalg.norm(r_ij) # Weight based on distance (simple: inverse distance) weight 1.0 / (d_ij 1.0) # Scale by feature similarity feat_sim jnp.dot(features[i], features[j]) shift shift weight * feat_sim * r_ij new_positions new_positions.at[i].set(positions[i] 0.1 * shift) return new_positions # 3 atoms pos jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) feat jnp.array([[1.0, 0.5], [0.5, 1.0], [0.8, 0.3]]) # Update positions pos_new egnn_position_update(pos, feat) # Now rotate input, update, and check if output is rotated consistently R jnp.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) pos_rot (R pos.T).T pos_new_from_rot egnn_position_update(pos_rot, feat) # Should be the same as rotating the original output pos_new_then_rot (R pos_new.T).T print(fUpdate then rotate:\n{jnp.round(pos_new_then_rot, 4)}) print(f\nRotate then update:\n{jnp.round(pos_new_from_rot, 4)}) print(f\nEquivariant: {jnp.allclose(pos_new_then_rot, pos_new_from_rot, atol1e-4)})这个实验验证了 EGNN 位置更新的核心机制位移项 $\sum_{j} \text{weight}_{ij} \cdot (\mathbf{r}_i - \mathbf{r}_j)$ 中权重只依赖距离与特征不变量而 $(\mathbf{r}_i - \mathbf{r}_j)$ 是相对向量——当输入整体旋转时每个相对向量和它们的加权和都旋转相同的 $R$因此「更新后再旋转」与「旋转后再更新」数学上完全等价。注意真实 EGNN 用学习到的 $\phi_e$、$\phi_r$ 函数替代这里的固定权重但等变性论证完全相同。小结与延伸阅读3D 图网络把 GNN 从「抽象关系图」扩展到「物理空间中的几何图」其设计哲学完全由对称性驱动标量输出必须对 SE(3) 不变向量/张量输出必须对 SE(3) 等变。不变架构SchNet、DimeNet、SphereNet通过只消费距离、角度、二面角等不变特征来简化问题等变架构EGNN、Tensor Field Networks、MACE、Equiformer通过维护球张量表示与等变运算来释放向量级预测能力。这两条路线分别对应分子性质预测与力场学习、以及需要高阶张量输出的场景是当前 AI for Science 的核心引擎。继续深入本仓库的相关内容几何深度学习基础对称群、不变性/等变性的统一理论框架图论邻接矩阵、图拉普拉斯、谱图理论图神经网络消息传递框架、GCN/GraphSAGE/GIN、过平滑图注意力网络GAT、Graph TransformerEquiformer 的注意力基础蛋白质设计与药物发现3D GNN 在应用层的落地场景。如需将本文档作为知识库供 AI 助手检索使用本仓库还附带了一个 MCP servermcp/src/index.ts它会按「chapter N - name/MM. section.md」的命名约定扫描全部章节并暴露为可检索的知识工具——本文即属于其中第 12 章第 5 节「3D Graph Networks」。【免费下载链接】maths-cs-ai-compendiumBecome a cracked AI/ML researcher/engineer with this unconventional textbook covering maths, computing, and ML with intuition.项目地址: https://gitcode.com/GitHub_Trending/mat/maths-cs-ai-compendium创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考