3D 图网络

一句话总结：把图神经网络 (Graph Neural Network, GNN) 从抽象图扩展到 3D 物理空间，用 SE(3)/E(n) 等变性建模分子、蛋白质、晶体和点云——从 SchNet 到 EGNN 到张量场网络，学会尊重旋转与平移的对称性。

3D 图网络把 GNN 扩展到带空间几何的数据，处理旋转与平移必须做对。 本节涵盖几何图、SE(3)/E(n) 等变性、SchNet、DimeNet、EGNN、张量场网络，以及在分子性质预测、蛋白质结构、材料科学和药物发现里的应用——这些架构学习的是 3D 物理世界.

• 前两节 (第 3、4 节) 里的 GNN 跑在抽象图上：节点有特征，边编码连接，但没有 3D 空间的概念。 一个社交网络图没有几何。 但 GNN 最有影响力的许多应用都涉及生活在 3D 物理空间里的数据：分子、蛋白质、晶体、点云。 对它们来说，节点的空间位置携带关键信息，抽象 GNN 会把这些信息扔掉。

• 挑战在于 3D 数据有几何对称性 (第 1 节)：把一个分子旋转一下，它的性质不变；平移一下也不变。 3D GNN 必须尊重这些对称性。 如果你旋转分子，它预测的能量就变了，那物理上就是错的。

几何图

• 几何图 (geometric graph) 是嵌在 3D 空间里的图。 每个节点 \(i\) 除了特征向量 \(\mathbf{h}_i\)，还有一个位置 \(\mathbf{r}_i \in \mathbb{R}^3\). 边可以由空间邻近定义 (把距离小于 \(r_{\text{cut}}\) 的节点连起来)，而不必依赖显式的化学键。

• 对分子来说，几何图以原子为节点 (特征：元素类型、电荷等)，以化学键为边。 3D 位置 \(\mathbf{r}_i\) 是原子坐标，由量子力学计算或实验测量 (X 射线晶体学、冷冻电镜) 得到。

• 对点云 (来自 LiDAR 或 3D 扫描仪，见第 8 章与第 11 章)，每个点是一个节点，带位置和可选特征 (颜色、强度). 边连接近邻点，形成 k 近邻 (k-Nearest Neighbour, kNN) 图或半径图。

• 消息传递里的关键几何量:

  • 原子间距离: \(d_{ij} = \|\mathbf{r}_i - \mathbf{r}_j\|\). 距离对旋转与平移不变。 两个分子如果原子间距离全部相同，它们就是同一个形状，不管朝向。

  • 键角：在节点 \(i\) 处，向量 \(\mathbf{r}_j - \mathbf{r}_i\) 与 \(\mathbf{r}_k - \mathbf{r}_i\) 之间的夹角 \(\theta_{ijk}\). 角度捕捉的是超越两两距离的局部几何。

  • 二面角 (dihedral, torsion)：平面 \((i, j, k)\) 与 \((j, k, l)\) 之间的夹角 \(\phi_{ijkl}\). 二面角刻画结构在 3D 中如何扭转，对蛋白质骨架几何至关重要。

  • 相对位置向量: \(\mathbf{r}_{ij} = \mathbf{r}_j - \mathbf{r}_i\). 它们对平移不变，但对旋转不不变. 用它们就得上等变 (而不只是不变) 架构。

SE(3) 与 E(n) 等变性

• 3D 物理数据的对称群是欧几里得群 (Euclidean group) \(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) \quad \text{(invariance)}\]

• 对向量预测 \(\mathbf{F}\):

\[\mathbf{F}(R\mathbf{r}_1, R\mathbf{r}_2, \ldots) = R \cdot \mathbf{F}(\mathbf{r}_1, \mathbf{r}_2, \ldots) \quad \text{(equivariance)}\]

• 这些约束直接对应第 1 节的不变性/等变性框架，现在专门用到 3D 旋转与平移群上。

• 有两种设计思路:

  1. 不变架构 (invariant architectures)：只把不变的几何特征 (距离、角度) 送进消息传递。 网络内部表示都是标量 (不变). 简单高效，但除非破坏对称，否则做不出向量输出。
  2. 等变架构 (equivariant architectures)：全网维持向量 (乃至更高阶张量) 表示，保证每一层都等变。 表达力更强，天然能预测向量和张量，但更复杂。

SchNet：基于距离的消息传递

• SchNet (Schütt 等，2017) 是奠基性的不变 3D GNN. 它的关键创新是连续滤波卷积：不再用一组固定的边类型 (像分子 GNN 里那样固定化学键类型), SchNet 直接从原子间距离生成消息滤波器。

• 距离 \(d_{ij}\) 先用径向基函数 (Radial Basis Function, 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\) 的高斯。 这就类似于给距离用的可学习位置编码：连续的距离被映射到一个高维特征空间，网络就能学到依赖距离的相互作用。 中心 \(\mu_k\) 通常从 0 到截断半径均匀取点。

• SchNet 从节点 \(j\) 到节点 \(i\) 的消息是:

\[\mathbf{m}_{j \to i} = \mathbf{h}_j \odot W_{\text{filter}}(\text{RBF}(d_{ij}))\]

• 其中 \(W_{\text{filter}}\) 是一个把 RBF 展开映射到滤波器向量的多层感知机 (Multi-Layer Perceptron, MLP), \(\odot\) 是逐元素乘 (Hadamard 积，见第 2 章). 滤波器依赖距离，所以近处原子与远处原子的相互作用方式不同。 这个逐元素乘就像一种门控机制 (见第 6 章)：距离相关的滤波器控制每个特征维度能通过多少。

• 因为 SchNet 只用距离 (不变量)，整个模型自动对旋转和平移不变。 这一设计选择之外，不需要额外处理对称性。

DimeNet 与 SphereNet：角度与二面角

• 光靠距离没法完全确定 3D 结构。 两种不同的分子构象可以有完全相同的两两距离，却有不同的键角 (这就是所谓"距离几何模糊"问题). DimeNet (Gasteiger 等，2020) 把键角加进了消息传递。

• DimeNet 用的是方向性消息传递：消息沿着有向边流动，边 \((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 等，2022) 再往前一步，加入二面角 \(\phi_{lkji}\)，捕捉完整的 3D 扭转结构。 层次是:

  • 距离 → 捕捉两两邻近关系
  • 角度 → 捕捉局部几何 (弯的还是直的)
  • 二面角 → 捕捉 3D 扭转 (对蛋白质骨架、药物结合都很关键)

• 每加一层几何分辨率就付一次计算代价 (距离是 \(O(|E|)\)，角度是 \(O(|E| \cdot k)\)，二面角是 \(O(|E| \cdot k^2)\)，其中 \(k\) 是平均度).

E(n) 等变 GNN (EGNN)

• EGNN (Satorras 等，2021) 走的是等变路线：不只用不变特征，它在每一层同时更新节点特征和节点位置，全程保持等变。

• EGNN 对节点 \(i\) 的更新:

\[\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\)，而 \(\phi_r\) 只依赖不变量 (特征与距离). 这个构造可证明等变：如果所有输入位置都按 \(R\) 旋转，所有输出位置也按同一个 \(R\) 旋转。

• EGNN 的优雅之处在于，它没有显式用球谐或不可约表示就实现了等变。 相对位置向量携带方向信息，而不变的消息函数控制这些方向信息怎么用。

• 简洁是有代价的：EGNN 只用了向量表示 (1 阶). 不做扩展就没法表示更高阶张量，比如四极矩或应力张量。

张量场网络与高阶表示

• 张量场网络 (Tensor Field Networks) (Thomas 等，2018) 及其后续 (SE(3)-Transformer、MACE、Equiformer) 用上了旋转群不可约表示 (irreducible representations) 的整套机器来搭建等变层。

• 从表示论的角度看 (这就接上了第 2 章的线性代数), 3D 里的旋转可以分解成由整数阶 \(\ell\) 刻画的不可约分量:

  • \(\ell = 0\)：标量 (1 个分量，不变). 能量、电荷。
  • \(\ell = 1\)：向量 (3 个分量，像位置向量那样旋转). 力、偶极矩。
  • \(\ell = 2\)：二阶对称无迹张量 (5 个分量). 四极矩、应力张量。
  • 更高的 \(\ell\)：捕捉越来越复杂的角结构。

• 这些叫球张量 (spherical tensors)，在旋转 \(R\) 下按 Wigner-D 矩阵 \(D^\ell(R)\) 变换：标量不变，向量按 \(R\) 转，二阶张量按更复杂的矩阵转。

• 用球张量做等变消息传递用的是 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 等，2022) 用高阶消息 (多个邻居特征的乘积) 来在更少消息传递层数下达到高精度。 通过构造"体阶有序"的相互作用 (从距离得两体、从角度得三体、从张量积得多体), MACE 高效地捕捉复杂的原子相互作用。

• Equiformer (Liao & Smidt, 2023) 把等变球张量特征与 Transformer 注意力机制 (第 4 节) 结合，造出一个 SE(3) 等变的图 Transformer. 注意力分数从不变特征算，值聚合则在等变张量特征上做。

应用

• 分子性质预测：给一个分子的 3D 结构，预测能量、力、偶极矩、HOMO-LUMO 能隙、毒性、溶解度这些性质。 这是 3D GNN 最成熟的应用。 在量子化学数据集 (QM9、OC20) 上训练的模型在很多性质上达到化学精度，能对上百万候选分子做虚拟筛选。

• 加速分子动力学：用量子力学 (密度泛函理论 (Density Functional Theory, DFT)) 算原子间的力极贵 (\(n\) 个电子要 \(O(n^3)\)). 一个训练好预测力的 3D GNN 可以在分子动力学模拟里替掉 DFT，提速 \(10^3\)–\(10^6\) 倍，精度接近 DFT. 这就能模拟更大的体系、更长的时间尺度，揭示传统方法看不到的现象。

• 蛋白质结构：蛋白质是氨基酸链，折叠成复杂的 3D 结构。 蛋白质骨架就是一个几何图，节点是残基，边连接空间上靠近的残基。 3D GNN 被用来做蛋白质功能预测、结合位点识别和蛋白质设计 (逆折叠：给一个想要的结构，预测氨基酸序列). AlphaFold 就用几何与图推理来从序列预测蛋白质结构。

• 材料科学与催化：晶体材料有周期性的 3D 结构。 GNN 建模重复的晶胞并预测材料性质：带隙、生成能、力学强度。 Open Catalyst 项目 (OC20/OC22) 给 GNN 定了预测催化表面吸附能的基准，加速新催化剂 (用于可再生能源) 的搜索。

• 药物发现: 3D GNN 预测药物分子如何与目标蛋白结合。 结合亲和力取决于药物与蛋白结合口袋之间的 3D 形状互补与化学相互作用。 像 DiffDock 这样的模型把等变 GNN 与扩散模型 (第 8 章) 结合起来预测结合构象 (药物在蛋白口袋里的 3D 朝向).

图生成

• 上面所有架构都是分析已有图。 图生成 (Graph generation) 则要造新图：设计一个具备目标性质的分子、生成用于测试的合成社交网络、或提出一种全新的蛋白质结构。 这是图级预测的生成式对偶。

• 挑战是图是离散、变长、组合式的。 生成一张图要决定造多少个节点、它们的特征是什么、哪些节点两两相连。 可能的图空间随节点数超指数地增长。

• 自回归生成 (autoregressive generation) 一次一个节点 (或一条边) 地构建图。 GraphRNN (You 等，2018) 按序列生成图：一个循环神经网络 (Recurrent Neural Network, RNN) 维护一个状态，每一步生成一个新节点，并决定要跟哪些已有节点连接。 生成顺序给本来无序的图强加了一个人为序列，但按广度优先搜索 (Breadth-First Search, BFS) 排序能保证最近生成的节点仍然相关。

• 基于 VAE 的生成：把图编码到一个连续隐空间 (用 GNN 编码器)，再从采样的隐向量解码出新图。 GraphVAE 一次性生成一个概率邻接矩阵 \(\hat{A} \in [0, 1]^{n \times n}\)，但复杂度是 \(O(n^2)\)，而且输出是稠密的，得阈值化。 隐空间支持平滑插值：在两个分子嵌入之间移动，会生成化学上合法的中间结构。

• 基于扩散的生成：把扩散框架 (第 8 章) 用到图上。 前向过程逐步给节点特征和边结构加噪。 反向过程学着去噪，从噪声里生成合法图。 DiGress (Vignac 等，2023) 对节点类型和边类型都用离散扩散，自然处理图数据的类别性质。

• 对分子生成，关键约束是化学有效性：生成的分子必须遵守化合价规则 (碳形成 4 个键、氧形成 2 个，等等). 像 Junction Tree VAE (JT-VAE) 这样的方法把分子分解成合法的子结构 (环、链、官能团)，用这些积木组装，从构造上保证有效性。

• 目标驱动生成 针对特定性质做优化：生成一个对目标蛋白结合亲和力高、毒性低、溶解度好的分子。 这就是把图生成和性质预测 (用 3D GNN 做性质评估器) 组成一个闭环：生成 → 评估 → 迭代。 强化学习 (第 6 章) 或贝叶斯优化引导在化学空间的搜索。

• DiffDock (Corso 等，2023) 用 SE(3) 等变扩散来预测药物分子如何对接进蛋白结合口袋。 模型从随机放置开始，通过去噪生成 3D 结合构象 (药物相对蛋白的位置与朝向)，把本节的 3D 等变网络与第 8 章的扩散框架结合起来。

编程练习 (用 Colab 或 Notebook)

1. 基于原子间距离搭一个简单的不变 3D 消息传递层。 应用到一个小分子 (水：H-O-H) 上，验证输出对旋转不变。

   import jax
   import jax.numpy as jnp
   
   # 水分子: O 在原点, 两个 H 原子
   positions = jnp.array([[0.0, 0.0, 0.0],     # O
                           [0.96, 0.0, 0.0],    # H1
                           [-0.24, 0.93, 0.0]])  # H2
   
   # 节点特征: [原子序数]
   features = jnp.array([[8.0], [1.0], [1.0]])
   
   # 计算两两距离 (不变量)
   def pairwise_distances(pos):
       diff = pos[:, None, :] - pos[None, :, :]
       return jnp.sqrt(jnp.sum(diff**2, axis=-1) + 1e-8)
   
   # 简单的基于距离的消息传递
   def invariant_message_pass(features, positions):
       dists = pairwise_distances(positions)
       # 用 4 个中心做 RBF 展开
       centres = jnp.array([0.5, 1.0, 1.5, 2.0])
       rbf = jnp.exp(-5.0 * (dists[:, :, None] - centres[None, None, :]) ** 2)
   
       # 消息: 特征被距离相关的滤波器加权
       messages = jnp.einsum("ij,jd->id", rbf.sum(axis=-1), features)
       return messages
   
   output1 = invariant_message_pass(features, positions)
   
   # 绕 z 轴旋转 90 度
   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)}")
   

2. 计算三个原子之间的键角，验证它对旋转不变。

   import jax.numpy as jnp
   
   def bond_angle(r_i, r_j, r_k):
       """节点 j 处, 边 j->i 与 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))
   
   # 三个原子
   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}°")
   
   # 加一个随机旋转
   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)}")
   

3. 演示等变的位置更新 (EGNN 风格). 用距离加权的相对向量更新节点位置，验证等变性。

   import jax
   import jax.numpy as jnp
   
   def egnn_position_update(positions, features):
       """简单的 EGNN 风格等变位置更新."""
       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 = 1.0 / (d_ij + 1.0)
                   # 按特征相似度缩放
                   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 个原子
   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]])
   
   # 更新位置
   pos_new = egnn_position_update(pos, feat)
   
   # 现在旋转输入、更新, 检查输出是否也一致地被旋转
   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)
   
   # 应该与"先原始更新, 再旋转"的结果相同
   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)}")
   

---

📌 本节要点

• 几何图：节点带 3D 位置 \(\mathbf{r}_i\) 的图，边可由空间邻近定义，关键几何量是距离、键角、二面角与相对位置向量。
• SE(3)/E(n) 对称性：标量输出对旋转平移不变，向量输出 (力、偶极矩) 需按同样旋转等变——违反就是物理错。
• 不变 vs 等变架构：前者只用距离/角度等标量输入 (简单，出不了向量)；后者全程保留向量或高阶张量 (表达力更强，更复杂).
• SchNet：用 RBF 把距离展开成特征，逐元素乘作距离相关的滤波，是最经典的不变 3D GNN.
• DimeNet / SphereNet：分别把键角、二面角引入消息传递，越高阶几何分辨率越强但计算量越大。
• EGNN：在每层同步更新特征与位置，用相对位置向量 + 不变消息函数，优雅地做到 SE(3) 等变。
• 张量场网络 / MACE / Equiformer：用旋转群的不可约表示 (\(\ell = 0, 1, 2, \ldots\)) 与 Clebsch-Gordan 张量积做多阶等变消息传递。
• 典型应用：分子性质预测、分子动力学加速 (替代 DFT)、蛋白质结构 (AlphaFold)、材料/催化剂设计、药物结合 (DiffDock).
• 图生成：从自回归 (GraphRNN)、VAE (GraphVAE、JT-VAE) 到扩散 (DiGress、DiffDock)，关键约束是化学有效性与可扩展性。
• 闭环设计：图生成 + 3D GNN 性质评估器 + 强化学习/贝叶斯优化，可在化学空间做目标驱动搜索。