name: equivariant-architecture-designer description: 当您已验证对称性群并需要设计尊重这些对称性的神经网络架构时使用。当用户提到等变层、G-CNN、e3nn、可转向网络、将对称性构建到模型中,或需要针对特定对称性群的架构推荐时调用。提供架构模式和实现指导。
等变架构设计师
这是什么?
这个技能帮助您设计尊重已识别对称性群的神经网络架构。给定已验证的群规范,它推荐架构模式、特定库和实现策略。
好处:等变架构参数更少、训练更快、泛化更好,并且对分布偏移更鲁棒。
工作流程
复制此清单并跟踪您的进度:
架构设计进度:
- [ ] 步骤1:审查群规范和需求
- [ ] 步骤2:选择架构家族
- [ ] 步骤3:选择特定层和组件
- [ ] 步骤4:设计网络拓扑
- [ ] 步骤5:选择实现库
- [ ] 步骤6:创建架构规范
步骤1:审查群规范和需求
收集已验证的群规范。确认:涉及哪些群、需要不变性还是等变性、数据域(图像、点云、图形等)、任务类型(分类、回归、生成)以及任何计算约束。如果群未指定,首先与用户合作识别它。
步骤2:选择架构家族
使用架构选择指南将对称性群匹配到架构家族。关键家族:用于网格上离散群的G-CNN、用于连续2D群的可转向CNN、用于点数据E(3)的e3nn/NequIP、用于图形置换的GNN、用于集合置换的DeepSets。考虑表达能力和效率之间的权衡。
步骤3:选择特定层和组件
基于层模式选择层类型。对于每层决定:卷积类型(常规、群、可转向)、非线性(必须保持等变性 - 使用门控、基于范数或张量积)、归一化(批归一化打破等变性 - 使用层归一化或等变批归一化)、池化(用于不变输出:使用不变池化;用于等变:保持结构)。有关详细设计方法,请参阅方法细节。
步骤4:设计网络拓扑
设计整体网络结构:编码器架构(如何提取特征)、每个阶段的特征表示(李群的不可约表示)、池化/聚合策略、输出头匹配任务需求。使用拓扑模式进行常见设计。平衡深度与宽度以适应群大小。
步骤5:选择实现库
基于库参考选择库。匹配您的群、框架偏好(PyTorch/JAX)和性能需求。流行选择:e3nn(E(3)/O(3),PyTorch)、escnn(离散群,PyTorch)、pytorch_geometric(置换,PyTorch)。确保库支持您的特定群。
步骤6:创建架构规范
使用输出模板记录设计。包括:逐层规范、表示类型、库依赖、预期参数计数以及伪代码或实际代码骨架。此规范指导实现和后续等变性验证。有关即用型实现模板,请参阅代码模板。此输出的质量标准定义在质量评价标准。
架构选择指南
按对称性群
| 群 | 域 | 推荐架构 | 库 |
|---|---|---|---|
| Cₙ, Dₙ | 2D图像 | G-CNN, 群等变CNN | escnn, e2cnn |
| SO(2), O(2) | 2D图像 | 可转向CNN, 谐波网络 | escnn |
| SO(3) | 球形 | 球形CNN | e3nn, s2cnn |
| SE(3), E(3) | 点云 | 等变GNN, 张量场网络 | e3nn, NequIP |
| Sₙ | 集合 | DeepSets | pytorch, jax |
| Sₙ | 图形 | 消息传递GNN | pytorch_geometric |
| E(3) × Sₙ | 分子 | E(3)等变GNN | e3nn, SchNet |
按任务类型
| 任务 | 输出类型 | 关键考虑 |
|---|---|---|
| 分类 | 不变标量 | 使用不变池化 |
| 回归(标量) | 不变标量 | 与分类相同 |
| 分割 | 等变每点 | 保持等变性到输出 |
| 力预测 | 等变向量 | 输出为l=1不可约表示 |
| 姿态估计 | 等变变换 | 输出旋转+平移 |
| 生成 | 等变结构 | 等变解码器 |
层模式
等变卷积模式
标准G-卷积:
(f ⋆ ψ)(g) = ∫_G f(h) ψ(g⁻¹h) dh
- 输入:群G上的特征图
- 核:G上的函数
- 输出:群G上的特征图
可转向卷积:
- 使用可预测转换的可转向核
- 通过不可约表示参数化
- 对于连续群更高效
e3nn张量积层:
# 结合不同角动量的特征
tp = o3.FullyConnectedTensorProduct(
irreps_in1, irreps_in2, irreps_out
)
output = tp(input1, input2)
等变非线性
问题:标准非线性(ReLU等)打破等变性。
解决方案:
| 类型 | 如何工作 | 何时使用 |
|---|---|---|
| 基于范数 | 对 | |
| 门控 | 使用不变门控等变 | 通用目的 |
| 张量积 | 通过Clebsch-Gordan进行非线性 | e3nn、高质量 |
| 不变特征 | 仅应用于l=0组件 | 简单、快速 |
等变归一化
批归一化:打破等变性(不同方向的不同统计) 解决方案:
- 层归一化(每样本归一化)
- 等变批归一化(每不可约表示通道归一化)
- 实例归一化(通常OK)
用于不变性的池化
从等变特征获取不变输出:
| 方法 | 公式 | 何时使用 |
|---|---|---|
| 平均池化 | 群上的平均 | 连续群 |
| 求和池化 | 元素求和 | 集合、图形 |
| 最大池化 | 最大 | |
| 注意力池化 | 加权和 | 当重要性变化时 |
拓扑模式
编码器-解码器(分割、生成)
输入 → [等变编码器] → 潜在(等变) → [等变解码器] → 输出
- 编码器:渐进特征提取
- 潜在:等变表示
- 解码器:用对称性重建
编码器-池化(分类)
输入 → [等变编码器] → 特征(等变) → [不变池化] → [MLP] → 类
- 在末端池化以获取不变特征
- 最终MLP操作于不变表示
消息传递(图形/点云)
节点 → [消息传递层1] → [消息传递层2] → ... → [聚合] → 输出
- 每层:聚合邻居、更新节点
- 聚合:求和/平均用于不变性、每节点用于等变性
库参考
e3nn(PyTorch)
群:E(3), O(3), SO(3) 优势:全不可约表示支持、张量积、球谐函数 用于:分子建模、3D点云、物理
from e3nn import o3
irreps = o3.Irreps("2x0e + 2x1o + 1x2e") # 2标量、2向量、1张量
escnn(PyTorch)
群:离散群(Cₙ, Dₙ)、连续2D(SO(2), O(2)) 优势:图像处理、文档齐全 用于:具有旋转/反射对称性的2D图像
from escnn import gspaces, nn
gspace = gspaces.rot2dOnR2(N=4) # C4旋转群
pytorch_geometric(PyTorch)
群:置换(Sₙ) 优势:图形、批处理、许多GNN层 用于:图形分类/回归、节点预测
from torch_geometric.nn import GCNConv, global_mean_pool
其他库
| 库 | 群 | 框架 | 注释 |
|---|---|---|---|
| NequIP | E(3) | PyTorch | 分子动力学 |
| MACE | E(3) | PyTorch | 分子势能 |
| jraph | Sₙ | JAX | 图形网络 |
| geomstats | 李群 | NumPy/PyTorch | 流形学习 |
输出模板
架构规范
==========================
目标对称性:[群名称和表示法]
对称性类型:[不变/等变]
任务:[分类/回归等]
域:[图像/点云/图形等]
架构家族:[例如,E(3)等变GNN]
库:[例如,e3nn]
层规范:
1. 输入层
- 输入类型:[例如,3D坐标+特征]
- 表示:[例如,位置(l=1)+标量(l=0)]
2. [层名称]
- 类型:[卷积/张量积/消息传递]
- 输入不可约表示:[规范]
- 输出不可约表示:[规范]
- 非线性:[门控/范数/无]
3. [继续每层...]
N. 输出层
- 聚合:[平均/求和/注意力]
- 输出:[不变标量/等变向量等]
估计参数:[计数]
关键依赖:[库版本]
代码骨架:
[提供实现大纲或伪代码]
下一步:
- 使用指定库实现架构
- 实现后通过数值测试验证等变性