等变网络家族大比武：GATr vs SE(3)-Transformer vs SEGNN vs EGNN

# 等变网络家族大比武：GATr vs SE(3)-Transformer vs SEGNN vs EGNN > **参考对象**：AnandTech 的 CPU/GPU 横评风格——用数据说话，但每行数据背后都有「为什么」的解释 --- ## 为什么需要这场比武 3D 几何深度学习领域，2021-2024 年间涌现了十几种等变网络架构。它们都声称「保持旋转/平移不变性」，但实现方式、适用场景、性能特征天差地别。 如果你正在选一个架构做项目，你会遇到这些问题： - "EGNN 代码最简单，但效果是不是不够好？" - "SE(3)-Transformer 名字响亮，但据说很慢？" - "SEGNN 在小数据集上无敌，大数据集呢？" - "GATr 最新最 fancy，但值得学习曲线吗？" 这篇文章用同一套 benchmark 数据，给你一个清晰的决策地图。 --- ## 四剑客简介 | 架构 | 年份 | 核心思想 | 等变群 | 计算复杂度 | |------|------|----------|--------|-----------| | **EGNN** | 2021 | 在消息传递中直接编码距离 | E(n) | O(N²) | | **SE(3)-Transformer** | 2020 | 球谐函数 + 自注意力 | SE(3) | O(N² × L³) | | **SEGNN** | 2022 | 可操控向量 + 消息传递 | SE(3) | O(N² × L²) | | **GATr** | 2023 | 几何代数 + Transformer | E(3) | O(N²) | 注：N = 节点数，L = 特征阶数（球谐阶数）。 --- ## Round 1：N-Body 动力学预测 这是等变网络的「高考」。给定 5 个粒子的初始位置和速度，预测 1000 个时间步后的位置。旋转/平移数据必须不影响预测精度。 ### 成绩表（MSE，越低越好） | 模型 | 1000 ts | 1500 ts | 2000 ts | Forward Time (ms) | |------|---------|---------|---------|-------------------| | Linear | 6.831 | 20.012 | 39.513 | 0.1 | | GNN | 1.077 | 5.059 | 10.591 | 2.1 | | Radial Field | 1.060 | 12.514 | 26.388 | 3.0 | | **EGNN** | **0.716** | **2.201** | **4.049** | **5.5** | | SE(3)-Transformer | 2.483 | 18.891 | 36.730 | 114.6 | | TFN | 1.544 | 11.116 | 23.823 | 27.2 | | **SEGNN** | **0.481** | **1.552** | **3.294** | **24.5** | | **GATr** | **0.320** | **~0.52** | **~1.0** | **~20** | 数据来源：SEGNO (ICLR 2024) 与 GATr (NeurIPS 2023) 论文实验。 ### 解读 **EGNN：性价比之王** - 用 O(N²) 的复杂度做到了接近 SEGNN 的效果 - 核心 trick：在消息传递时直接用距离公式 $m_{ij} = \phi_e(h_i, h_j, ||x_i - x_j||^2)$，把距离作为标量输入 - 缺点：不是严格等变（是 E(n) 等变，但实现方式不如球谐函数优雅），表达能力上限较低 **SE(3)-Transformer：理论满分，实践拉胯** - 用球谐函数（spherical harmonics）做等变注意力，数学上最干净 - 但球谐函数的计算复杂度是 O(L³)，特征阶数一高就爆炸 - 前向传播 114.6ms，是 EGNN 的 20 倍 - 适合：研究新等变机制的原型验证；不适合：生产环境 **SEGNN：小数据集的王者** - 在可操控向量（steerable vectors）上做消息传递，比 EGNN 多了「方向信息」的处理 - 效果比 EGNN 好，但慢 4-5 倍 - 适合：数据量 < 10K 的科学计算任务（分子模拟、材料科学） **GATr：全面发展的优等生** - 效果最好的同时，速度比 SE(3)-Transformer 快 5 倍 - 几何代数让「等变线性层」比球谐函数简单得多（16 维矩阵乘法 vs 球谐卷积） - 适合：需要 Transformer 级别可扩展性的几何任务（大规模点云、机器人规划） --- ## Round 2：分子属性预测（QM9 数据集） QM9 是分子深度学习的标准 benchmark，13K 个小分子，预测 12 种量子化学属性。 ### 成绩表（MAE，越低越好） | 模型 | α (bohr³) | Δε (meV) | μ (D) | C_v (cal/mol·K) | |------|-----------|----------|-------|-----------------| | SchNet | 0.297 | 71 | 0.049 | 0.043 | | DimeNet++ | 0.047 | 45 | 0.028 | 0.027 | | PaiNN | 0.049 | 49 | 0.029 | 0.028 | | SphereNet | 0.056 | 54 | 0.031 | 0.027 | | **EGNN** | **0.071** | **48** | **0.028** | **0.032** | | **SEGNN** | **0.077** | **46** | **0.033** | **0.035** | | SE(3)-Trans. | 0.142 | 53 | 0.051 | 0.054 | 数据来源：DuSEGO (2024) 与 SEGNO (2024) 论文。 ### 解读 **EGNN 和 SEGNN 被打败了？** 看数据，EGNN/S EGNN 在 QM9 上并不是 SOTA。DimeNet++、PaiNN、SphereNet 这些「非等变」的专用架构效果更好。 这说明一个重要事实：**等变性不是万能药**。在 QM9 这种「小分子、大量训练数据」的场景下： - 数据量足够（13K），网络可以从数据中「学到」旋转规律 - 分子结构固定（有机小分子），不需要处理任意 3D 变换 - 专用架构（如 PaiNN 的交互层）比通用等变架构更适配分子物理 **等变网络真正的优势场景**： - 训练数据少（几百个样本） - 需要严格物理一致性（如分子动力学模拟） - 测试时需要泛化到不同旋转/平移的数据（如机器人规划） --- ## Round 3：数据效率对比 等变网络最核心的卖点不是「准确率最高」，而是「用更少的数据达到同样的准确率」。 ### N-Body 数据效率曲线（定性） ``` 准确率 ↑ │ GATr ████████████████████ │ ██（500样本≈其他模型5000样本） │ SEGNN ██████████████ │ ██ │ EGNN ██████████ │ ██ │ 标准 ████ │ Transformer ██ │ ██ └─────────────────────────────→ 训练样本数 100 500 1K 5K ``` ### 为什么等变网络数据效率高？ **标准 Transformer**：要学的东西 = 物理规律 + 旋转不变性 + 平移不变性 + 其他对称性。1000 个样本中，大部分「信息」被浪费在教网络「旋转后东西不变」上。 **GATr**：旋转/平移不变性被写进结构，不需要学。1000 个样本全部用来学物理规律。相当于把考试大纲从 100 页缩减到 10 页，每页都学到精髓。 --- ## 决策树：你该怎么选 ``` 你的任务是什么？ │ ├─→ 3D 点云分类/分割（如 ModelNet40） │ └─→ 数据量 > 10K？ │ ├─→ 是 → 标准 PointNet++ / Point Transformer（不需要等变） │ └─→ 否 → EGNN（最简单）或 GATr（效果最好） │ ├─→ 分子动力学/量子化学 │ └─→ 需要严格物理守恒？ │ ├─→ 是 → SEGNN（最稳定）或 GATr（最通用） │ └─→ 否 → PaiNN / DimeNet++（专用架构更快） │ ├─→ 机器人运动规划 │ └─→ 需要生成模型（如扩散模型）？ │ ├─→ 是 → GATr（已验证的扩散模型骨干） │ └─→ 否 → EGNN（实时性更好） │ ├─→ 高能物理 / 粒子模拟 │ └─→ 需要洛伦兹等变性？ │ ├─→ 是 → L-GATr（2024 扩展） │ └─→ 否 → SEGNN │ └─→ 大规模 3D 场景理解（如自动驾驶） └─→ 点数 > 100K？ ├─→ 是 → LaB-GATr（几何 tokenization）或标准 Transformer └─→ 否 → GATr ``` --- ## 架构哲学对比 | 维度 | EGNN | SE(3)-Transformer | SEGNN | GATr | |------|------|-------------------|-------|------| | **数学基础** | 距离度量 | 球谐函数 | 可操控表示 | 几何代数 | | **设计哲学** | 实用主义 | 数学洁癖 | 物理驱动 | 统一框架 | | **学习曲线** | ★★☆☆☆ | ★★★★★ | ★★★★☆ | ★★★☆☆ | | **代码复杂度** | 低（~200行） | 高（~2000行） | 中（~800行） | 中（~1000行） | | **社区生态** | 成熟 | 研究用 | 较小 | 成长中 | | **扩展性** | 中 | 差 | 中 | 好 | --- ## 隐藏成本 选择架构时，除了准确率，还要算这些账： ### 1. 学习成本 - **EGNN**：看完论文就能写，2 小时上手 - **SE(3)-Transformer**：需要理解球谐函数、Clebsch-Gordan 系数、表示论。2 周上手。 - **SEGNN**：需要理解可操控向量、Wigner D 矩阵、不可约表示。1 周上手。 - **GATr**：需要理解几何代数基础（grade、multivector、geometric product）。3 天上手。 ### 2. 调试成本 - **EGNN**：和普通 GNN 一样调试 - **SE(3)-Transformer**：球谐函数数值不稳定是常见问题 - **SEGNN**：可操控向量的 shape 不匹配是噩梦 - **GATr**：几何积的维度计算容易出错 ### 3. 维护成本 - **EGNN**：社区最大，issue 最多，但也有人修 - **SE(3)-Transformer**：基本停止维护，e3nn 库接手部分功能 - **SEGNN**：作者维护，但更新频率低 - **GATr**：Qualcomm AI Research 官方维护，文档最全 --- ## 未来趋势 | 方向 | EGNN | SE(3)-Trans. | SEGNN | GATr | |------|------|-------------|-------|------| | 扩展到 E(3) | ❌ E(n) | ❌ SE(3) | ❌ SE(3) | ✅ E(3) | | 扩展到共形 | ❌ | ❌ | ❌ | ✅ (CGA) | | 扩展到洛伦兹 | ❌ | ❌ | ❌ | ✅ (L-GATr) | | 扩展到 LLM | ❌ | ❌ | ❌ | 🔶 (Versor 探索中) | | O(N) 复杂度 | ❌ | ❌ | ❌ | ✅ (Versor RRA) | --- ## 一句话总结 - **想最快上手**：EGNN，200 行代码跑起来 - **想最数学正确**：SE(3)-Transformer，但准备好等 20 倍时间 - **想小数据集无敌**：SEGNN，500 样本就能出效果 - **想一次投资长期收益**：GATr，几何代数是通用语言，扩展到哪里都能用 --- > **参考对象**：AnandTech 评测方法论——不只看跑分，还看「为什么这个跑分属于你」> > **信息来源**：SEGNO (ICLR 2024)、DuSEGO (2024)、GATr (NeurIPS 2023)、QM9 benchmark 汇总 #GATr #EGNN #SEGNN #SE3Transformer #等变网络 #横评 #benchmark #深度学习 #小凯