Ch12 §05 3D 图网络 · 讲解稿

章节定位: 教"3D 物理世界"的 GNN — 必须满足旋转/平移等变性。覆盖几何图 / SchNet / DimeNet / EGNN / 等变网络 / 分子动力学应用。 承接: §04 GAT 注意力 → §05 在 3D 上做"等变"网络 走向: 后续 Ch13+ 章节可结合 3D GNN 应用 (蛋白结构、推荐系统) 篇幅: ~4400 字 / 阅读 13 分钟 / 讲解 25-30 分钟

---

一 · 一句话总览

3D 图网络 = 满足 SE(3) / E(n) 等变性的 GNN。SchNet (距离基, 不变)、DimeNet (加角度)、EGNN (位置等变更新)、Equiformer (球谐等变) 是四代 3D GNN 演进, 物理化学/蛋白/分子动力学核心。

---

二 · 心智模型

3D 数据有两个核心特性:

1. 几何对称: 旋转 / 平移 / 反射不应改变物理量
2. 距离/角度/二面角: 三种基本几何量

量	不变性	等变性
距离 \(d_{ij}\)	✓ (旋转/平移都不变)	—
角度 \(\theta_{ijk}\)	✓	—
二面角 \(\phi_{ijkl}\)	✓	—
相对位置 \(\mathbf{r}_{ij} = \mathbf{r}_j - \mathbf{r}_i\)	平移 ✓, 旋转 ✗	旋转等变

---

三 · 章节主线 (5 段)

第 1 段 · 几何图与对称性 (8 分钟)

3.1.1 几何图

定义: 节点 \(i\) 在 \(\mathbb{R}^3\) 中有位置 \(\mathbf{r}_i\), 同时有特征 \(\mathbf{h}_i\)。

数据	节点	边	位置
分子	原子 (类型/电荷)	化学键	原子坐标
蛋白	氨基酸残基	肽键/空间邻近	Cα 原子坐标
点云	LiDAR 点 (颜色/强度)	kNN 邻接	3D 坐标
晶体	原子	晶格	晶胞坐标

3.1.2 关键几何量

量	公式	物理含义
距离	\(d_{ij} = \|\mathbf{r}_i - \mathbf{r}_j\|\)	原子对空间距离
键角	\(\theta_{ijk} = \angle(\mathbf{r}_j - \mathbf{r}_i, \mathbf{r}_k - \mathbf{r}_i)\)	三原子几何
二面角	\(\phi_{ijkl}\)	四原子扭转
相对位置	\(\mathbf{r}_{ij} = \mathbf{r}_j - \mathbf{r}_i\)	旋转等变

3.1.3 SE(3) 与 E(n) 群

群	包含	物理意义
SO(3)	旋转	刚体方向
SE(3)	旋转 + 平移	刚体运动
E(3)	旋转 + 平移 + 反射	物理对称 (含手性翻转)

3.1.4 输出决定架构

输出	要求	架构选择
标量 (能量)	旋转不变	Invariant (距离/角度)
向量 (力)	旋转等变	Equivariant (TFN / EGNN)
张量 (偶极矩)	旋转等变	高阶等变

---

第 2 段 · SchNet: 距离基消息传递 (10 分钟)

SchNet (Schütt 2017) 是 3D GNN 的奠基 — 用距离作为连续滤波。

3.2.1 距离 → 特征 (RBF)

距离是连续值, 网络需要把它"展开"成高维向量。SchNet 用径向基函数 (RBF):

\[\text{RBF}(d_{ij}) = \left[\exp(-\gamma_1 (d_{ij} - \mu_1)^2), \dots, \exp(-\gamma_K (d_{ij} - \mu_K)^2)\right]\]

\(K\) 个高斯, 中心 \(\mu_k\) 在 \([0, r_{\text{cut}}]\) 均匀分布, 宽度 \(\gamma_k\)。

类比: 类似 NLP 的 positional encoding — 把标量距离编码成可学习的"距离 embedding"。

3.2.2 消息公式

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

• \(W_{\text{filter}}\): MLP, 把 RBF 距离 → 滤波向量
• \(\odot\): 元素乘 (Hadamard 积, 类似门控机制)

3.2.3 关键性质

• 只用了距离 — 整个模型自动满足旋转/平移不变
• 不需要专门处理对称性
• 输出: 标量 (如能量 \(E\))

3.2.4 PyTorch 实现

import torch
import torch.nn as nn

class RBF(nn.Module):
    def __init__(self, n_rbf=64, cutoff=10.0):
        super().__init__()
        centers = torch.linspace(0, cutoff, n_rbf)
        widths = torch.full((n_rbf,), 2.0 / cutoff)
        self.register_buffer('centers', centers)
        self.register_buffer('widths', widths)

    def forward(self, d):
        # d: [..., 1]
        diff = d - self.centers
        return torch.exp(-self.widths * diff**2)  # [..., n_rbf]

class SchNetLayer(nn.Module):
    def __init__(self, dim, n_rbf=64):
        super().__init__()
        self.rbf = RBF(n_rbf)
        self.filter_mlp = nn.Sequential(
            nn.Linear(n_rbf, dim),
            nn.SiLU(),
            nn.Linear(dim, dim)
        )
        self.msg_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU())
        self.update_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU())

    def forward(self, h, pos, edge_index):
        src, dst = edge_index
        # 距离
        dist = torch.norm(pos[src] - pos[dst], dim=-1, keepdim=True)
        # 滤波
        rbf_feat = self.rbf(dist)
        filter_w = self.filter_mlp(rbf_feat)
        # 消息
        msg = h[src] * filter_w
        msg = self.msg_mlp(msg)
        # 聚合
        agg = torch.zeros_like(h).index_add_(0, dst, msg)
        # 更新
        return self.update_mlp(h + agg)

---

第 3 段 · DimeNet 与 SphereNet: 角度与二面角 (8 分钟)

3.3.1 距离不够

两分子可距离完全相同但构型不同 — 这叫"距离几何模糊性"。

例: - 顺式 vs 反式: 距离相同, 角度不同 - 蛋白折叠: 主链二面角决定形状

3.3.2 DimeNet 加角度

DimeNet (Gasteiger 2020) 在消息中加键角:

\[\mathbf{m}_{kj \to ji} = f(\mathbf{m}_{kj}, d_{ji}, \theta_{kji})\]

• 沿有向边传播消息
• 角度 \(\theta_{kji}\) 在边上 \(j \to i\) 处 (即 \(j\) 的入射角)
• 用球 Bessel 函数展开 (球面的自然基)

3.3.3 SphereNet 加二面角

SphereNet (Liu 2022) 进一步加二面角:

量	复杂度
距离	\(O(\|E\|)\)
角度	\(O(\|E\| \cdot k)\)
二面角	\(O(\|E\| \cdot k^2)\)

其中 \(k\) 是平均度。层次越深, 信息越丰富, 计算越贵。

3.3.4 三种几何量的物理意义

• 距离: 原子间远近 (范德华力)
• 角度: 局部几何 (弯曲 vs 直线)
• 二面角: 3D 扭转 (蛋白骨架, 药物结合姿态)

---

第 4 段 · 等变 GNN: EGNN 与球谐 (10 分钟)

3.4.1 限制: 不变架构

SchNet / DimeNet 只输出标量。要预测力向量 (分子动力学) 或偶极矩 (向量), 必须用等变架构。

3.4.2 EGNN: 位置等变更新

EGNN (Satorras 2021) 同时更新特征和位置:

\[\mathbf{m}_{ij} = \phi_e(\mathbf{h}_i, \mathbf{h}_j, d_{ij}^2, a_{ij})\]

\[\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(\mathbf{h}_i, \sum_j \mathbf{m}_{ij})\]

关键: 位置更新用相对位置 \((\mathbf{r}_i - \mathbf{r}_j)\) 加权求和。如果输入旋转 \(R\), 输出也旋转 \(R\) — 严格等变。

3.4.3 Tensor Field Network (TFN)

TFN (Thomas 2018) 用球谐函数当基, 把特征编码为不可约表示:

\[(\mathbf{h}_i)_l = \text{l阶表示, 2l+1 个分量}\]

• type-0 (\(l=0\)): 标量, 旋转不变
• type-1 (\(l=1\)): 向量, 旋转等变
• type-2 (\(l=2\)): 二阶张量

每层按 Wigner-Eckart 定理作用。

3.4.4 Equiformer: Transformer + 等变

Equiformer (Liao, ICLR 2023) = E(3) 等变 + Transformer:

class EquiformerLayer(nn.Module):
    def __init__(self, dim, lmax=2):
        super().__init__()
        self.lmax = lmax
        # 每个 l 阶有自己的 MLP / 注意力
        self.type0_mlp = nn.Linear(dim, dim)  # 标量
        self.type1_mlp = nn.Linear(dim, dim)  # 向量
        # 球谐基
        self.sphere_harm = SphereHarmonics(lmax)

    def forward(self, x_type0, x_type1, pos):
        # x_type0: 标量特征
        # x_type1: 向量特征
        # pos: 原子坐标
        # 1. 计算球谐表示
        sh = self.sphere_harm(pos)  # [n_atom, n_sh]
        # 2. 标量 + 向量做 attention (类型感知)
        msg0 = self.type0_mlp(x_type0)
        # 3. 等变更新位置
        return msg0

3.4.5 OGB-mol SOTA 对比

模型	OGB-mol ROC-AUC
GIN	76.0%
SchNet	73.0%
DimeNet++	76.4%
EGNN	78.6%
Equiformer	80.4%
GemNet	81.5%

等变 + Transformer 是当前 SOTA。

---

第 5 段 · 实践: 分子动力学应用 (10 分钟)

3.5.1 用 EGNN 模拟 MD

分子动力学 (MD) 需要快速预测能量 + 力场。EGNN 替代 DFT:

import torch
from torch_geometric.nn import EGNNConv

class EGNN(torch.nn.Module):
    def __init__(self, in_dim, hid_dim, out_dim):
        super().__init__()
        self.conv1 = EGNNConv(in_dim, hid_dim)
        self.conv2 = EGNNConv(hid_dim, hid_dim)
        self.head = torch.nn.Linear(hid_dim, out_dim)

    def forward(self, h, pos, edge_index):
        # h: [n, d]
        # pos: [n, 3]
        h, pos = self.conv1(h, pos, edge_index)
        h = torch.relu(h)
        h, pos = self.conv2(h, pos, edge_index)
        # 能量 (READOUT)
        energy = self.head(h).sum()
        return energy, pos

3.5.2 力的预测 (等变)

# 自动微分得到力
pos.requires_grad_(True)
energy, _ = model(h, pos, edge_index)
force = -torch.autograd.grad(energy, pos)[0]
# force 旋转等变: 如果 pos 旋转 R, force 也旋转 R

3.5.3 应用: 药物发现

分子 SMILES → 3D 构象生成 (RDKit) → 几何图 → EGNN → 性质预测

任务	模型	指标
OGB-molhiv	GIN/EGNN	ROC-AUC ~80%
QM9 能量	SchNet	MAE 0.5 meV
MD17 力场	EGNN	MAE 0.1 kcal/mol/Å
OCP 催化	Equiformer	E MAE 0.1 eV

3.5.4 应用: AlphaFold 3

AlphaFold 3 (DeepMind, 2024) 处理蛋白 + DNA + RNA + 小分子复合结构, 在三角注意力 (AlphaFold 2) 基础上引入 GNN 模块:

• 节点: 残基 + 核苷酸 + 原子
• 边: 化学键 + 距离 (≤ 30 Å)
• 消息: 距离 + 角度 + 边特征
• 输出: 3D 坐标 (等变更新)

3.5.5 应用: 点云 (Ch08/Ch11)

# PointNet++ 处理 LiDAR 点云
from torch_geometric.nn import PointNetConv

class PointNetPlusPlus(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = PointNetConv(local_nn=torch.nn.Linear(3+3, 64), add_self_loops=False)
        self.conv2 = PointNetConv(local_nn=torch.nn.Linear(64+3, 128), add_self_loops=False)
        # ...

    def forward(self, pos, batch):
        # kNN 邻接
        edge_index = knn_graph(pos, k=16, batch=batch)
        h = self.conv1(pos, pos, edge_index)
        h = self.conv2(h, pos, edge_index)
        # READOUT (max pool)
        return global_max_pool(h, batch)

---

四 · 四代 3D GNN 对比

模型	几何量	不变/等变	输出类型	复杂度	OGB-mol
SchNet	距离	不变	标量	\(O(E)\)	73%
DimeNet++	距离 + 角度	不变	标量	\(O(E k)\)	76%
EGNN	距离 + 位置	等变	标量+向量	\(O(E)\)	78%
Equiformer	球谐 (l≤2)	等变	多阶	\(O(E k^2)\)	80%
GemNet	二面角	等变	多阶	\(O(E k^2)\)	81%

---

五 · 易错点 (5 条)

1. 平移不变 vs 旋转等变: 距离对两者都不变; 相对位置 \(\mathbf{r}_{ij}\) 平移不变但旋转等变。
2. SchNet 只能标量输出: 想预测力向量, 必须用 EGNN/TFN 等变架构。
3. RBF 中心范围: 必须覆盖物理距离 (典型 [0, 10] Å); 超出范围预测不准。
4. 球谐函数的 l: \(l=0\) 标量, \(l=1\) 向量, \(l=2\) 二阶张量; \(l\) 越大计算越贵。
5. EGNN 的位置更新: 只在推理时更新, 训练时通常冻结 (否则优化困难)。

---

六 · 跨章连接

章节	关联
Ch02 §04 矩阵论	旋转矩阵 \(R \in SO(3)\), 等变操作
Ch03 §05 梯度	力 = \(-\nabla E\) 自动微分
Ch08 §04 Diffusion	等变扩散模型 (分子生成)
Ch08 §05 NeRF	3D 表征另一路径 (隐式辐射场)
Ch11 §05 VLA	机器人空间推理可用 EGNN
Ch05 §04 贝叶斯	贝叶斯优化分子性质

---

七 · 自测 5 题

1. SchNet 的 RBF 有什么作用? 为什么必须?
2. 距离不变 vs 相对位置等变的区别?
3. DimeNet 相对 SchNet 多了什么? 复杂度?
4. EGNN 如何保证位置等变更新?
5. 球谐函数在 Equiformer 中的角色?

---

八 · 一句话总结

3D 图网络 = 在物理空间满足 SE(3) / E(n) 等变性的 GNN。SchNet (距离基) → DimeNet (加角度) → EGNN (位置等变) → Equiformer (球谐等变), 表达力递增。分子动力学、蛋白结构 (AlphaFold 3)、药物发现是其核心应用。Ch12 至此讲完: GDL 范式 (§01) + 图论 (§02) + GNN 三件套 (§03) + 图注意力 (§04) + 3D 图 (§05)。

---

九 · 板书建议

[板书 1: 几何量]
   d_ij = ||r_i - r_j||       旋转+平移 不变
   θ_ijk = angle at i          旋转+平移 不变
   r_ij = r_j - r_i            平移 不变, 旋转 等变

[板书 2: SchNet]
   RBF(d) = [exp(-γ₁(d-μ₁)²), ...]
   m_{j→i} = h_j ⊙ W_filter(RBF(d_ij))

[板书 3: EGNN 位置更新]
   r_i' = r_i + C Σ (r_i - r_j) · φ_r(m_ij)
   ← 等变 (旋转 R 输入 → 旋转 R 输出)

[板书 4: 球谐]
   l=0: 标量
   l=1: 向量
   l=2: 二阶张量
   SO(3) 作用: Y_l^m → D^l(R) Y_l^m

---

十 · 参考资料

1. Schütt et al., "SchNet: A Continuous-filter Convolutional Neural Network for Modeling Quantum Interactions", NeurIPS 2017
2. Gasteiger et al., "Directional Message Passing for Molecular Graphs", ICLR 2020 (DimeNet)
3. Satorras et al., "E(n) Equivariant Graph Neural Networks", ICML 2021 (EGNN)
4. Thomas et al., "Tensor Field Networks", ICML 2018 (TFN)
5. Liao et al., "Equiformer: A Generic Graph Transformer for Molecular Property Prediction", ICLR 2023
6. Jumper et al., "AlphaFold 3", Nature 2024
7. Ch03 §05 梯度下降 (本套书)
8. Ch08 §04 Diffusion (本套书)