企业官网建设流程全解析

3D 图神经网络与 SE(3)/E(n) 等变架构实战指南：从 SchNet 到 Equiformer（maths-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）或半径图。

消息传递依赖的关键几何量有四类，它们的对称性属性各不相同：

最后一行是关键：两个分子若原子间距完全相同，则无论朝向如何都是同一形状；但相对位置向量虽然平移不变，却会随旋转而改变。要使用它们，就必须采用等变（而非仅仅不变）的架构。

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 旋转与平移群上的具体化。为什么必须内建这种对称性？正如该节强调的：把对称性构建进架构，远比让网络从数据中学习它高效——一个能量预测会随旋转而改变的模型，在物理上是错误的；而强行让网络学会所有朝向的等价性，会指数级扩大假设空间。

由此产生两条设计路线：

1. 不变架构（invariant architectures）：只用不变几何特征（距离、角度）作为消息传递输入，内部表示全部是标量。简单高效，但无法在保持对称性的前提下输出向量。
2. 等变架构（equivariant architectures）：在整个网络中维护向量（乃至更高阶张量）表示，保证每一层都等变。表达能力更强，能自然预测向量与张量，但更复杂。

SchNet：基于距离的连续滤波器消息传递

SchNet（Schü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 结构：两个不同的分子构象可能具有完全相同的成对距离，但键角不同（即「距离几何歧义」问题）。DimeNet（Gasteiger 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 之于距离。这一设计让模型在保持不变性的同时获得方向信息。

SphereNet（Liu et al., 2022）更进一步，把二面角$\phi_{lkji}$ 也纳入进来，捕捉完整的 3D 扭转结构。几何分辨率形成清晰的层级：

• 距离→ 成对邻近性
• 角度→ 局部几何（弯曲 vs 直线）
• 二面角→ 3D 扭转（蛋白质主链、药物结合的关键）

每一层分辨率的提升都伴随计算代价：距离是 $O(|E|)$，角度是 $O(|E| \cdot k)$，二面角是 $O(|E| \cdot k^2)$，其中 $k$ 是平均度。实际应用中需要在几何分辨率与算力之间权衡，这也是为什么许多生产级模型止步于角度级而很少全面使用二面角。

EGNN：等变位置更新

EGNN（Satorras 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 Networks（Thomas 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)$ 等变意义下的「矩阵乘法」——正是这一操作把不同阶的张量特征组合起来并输出合法阶数的新张量。

在此框架下涌现的现代模型各有侧重：

• MACE（Batatia et al., 2022）使用高阶消息（多个邻居特征的乘积），用体序交互（body-ordered interactions）——2 体来自距离、3 体来自角度、多体来自张量积——以更少的消息传递层达到高精度，高效捕捉复杂原子相互作用。
• Equiformer（Liao & 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 Project（OC20/OC22）用 GNN 基准预测催化表面的吸附能，加速可再生能源新催化剂的搜索。
• 药物发现：3D GNN 预测药物分子与靶蛋白的结合亲和力——结合亲和力取决于药物与蛋白结合口袋的 3D 形状互补性与化学相互作用。DiffDock这类模型将等变 GNN 与扩散模型（见第 8 章）结合，预测结合姿态（药物在蛋白口袋中的 3D 朝向）。

图生成：从分析到创造

以上所有架构都在分析已有图。图生成（graph generation）则创造新图：设计具有目标性质的分子、为测试生成合成社交网络、提出新型蛋白质结构。这是图级预测的生成式对应物。

挑战在于：图是离散、变长、组合的。生成一个图意味着决定创建多少个节点、它们具有什么特征、连接哪些节点对。可能图空间随节点数超指数增长。主流方法有四类：

• 自回归生成：一次生成一个节点（或一条边）。GraphRNN（You et al., 2018）顺序生成图：RNN 维护状态，每步生成一个新节点并决定与哪些已有节点相连。生成顺序给本质上无序的图强加了一个人为序列，BFS 序通过保持最近生成节点的相关性来缓解这一问题。
• 基于 VAE 的生成：用 GNN 编码器把图编码进连续潜空间，再从采样潜向量解码新图。GraphVAE一次性生成概率邻接矩阵 $\hat{A} \in [0, 1]^{n \times n}$，但开销为 $O(n^2)$ 且产生需要阈值的稠密输出。潜空间支持平滑插值：在两个分子嵌入间移动可生成化学有效的中间结构。
• 基于扩散的生成：把扩散框架（见第 8 章）应用到图上。前向过程逐步向节点特征与边结构加噪，反向过程学习去噪，从噪声生成有效图。DiGress（Vignac et al., 2023）对节点类型与边类型都应用离散扩散，天然处理图数据的类别特性。
• 面向分子的生成：关键约束是化学有效性——生成分子必须遵守价键规则（碳成 4 键、氧成 2 键等）。Junction Tree VAE（JT-VAE）把分子分解为有效子结构（环、链、官能团），通过组装这些积木生成，从构造上保证有效性。

目标导向生成（goal-directed generation）针对特定性质优化：生成对靶蛋白高结合亲和力、低毒性、良溶解度的分子。这以循环方式结合图生成与性质预测（用 3D GNN 作为性质评估器）：生成 → 评估 → 精化。强化学习（见第 6 章）或贝叶斯优化引导化学空间搜索。

值得一提的交叉点是DiffDock（Corso 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]], dtype=float) rotated_positions = (R @ positions.T).T output2 = invariant_message_pass(features, rotated_positions) print(f"Original output:\n{output1}") print(f"\nRotated output:\n{output2}") print(f"\nInvariant: {jnp.allclose(output1, output2, atol=1e-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(f"Original 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(f"Rotated angle: {jnp.degrees(angle_rotated):.1f}°") print(f"Invariant: {jnp.allclose(angle_original, angle_rotated, atol=1e-4)}")

键角由向量点积与范数的比值（即余弦值）给出，余弦值天然旋转不变，因此角度也是旋转不变的。jnp.clip(cos_angle, -1, 1)防止浮点误差导致arccos越界。这正是 DimeNet 中角度信息进入消息传递时依赖的不变性来源。

实验 3：EGNN 式等变位置更新

演示等变位置更新：用距离加权的相对向量更新节点位置，并验证「先旋转再更新」与「先更新再旋转」结果一致。

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(f"Update 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, atol=1e-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 Transformer，Equiformer 的注意力基础；
• 蛋白质设计与药物发现：3D GNN 在应用层的落地场景。

如需将本文档作为知识库供 AI 助手检索使用，本仓库还附带了一个 MCP server（mcp/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