equivariant-architecture-designer

等变架构设计器

工作流

复制此清单并跟踪进度：

Architecture Design Progress:
- [ ] Step 1: Review group specification and requirements
- [ ] Step 2: Select architecture family
- [ ] Step 3: Choose specific layers and components
- [ ] Step 4: Design network topology
- [ ] Step 5: Select implementation library
- [ ] Step 6: Create architecture specification

步骤 1：审查群规范和需求

收集经过验证的群规范。确认：涉及哪些群、需要不变性还是等变性、数据域（图像、点云、图等）、任务类型（分类、回归、生成）以及任何计算限制。如果未指定群，请先与用户一起确定。

步骤 2：选择架构系列

使用架构选择指南将对称群与架构系列匹配。主要系列：适用于网格上离散群的 G-CNN、适用于连续二维群的 Steerable CNN、适用于点数据上 E(3) 的 e3nn/NequIP、适用于图上置换的 GNN、适用于集合上置换的 DeepSets。权衡表达能力与效率。

步骤 3：选择具体层和组件

根据层模式选择层类型。对于每一层，请确定：卷积类型（常规、群、Steerable）、非线性（必须保持等变性——使用门控、基于范数或张量积）、归一化（批归一化会破坏等变性——使用层归一化或等变批归一化）、池化（对于不变输出：使用不变池化；对于等变输出：保留结构）。有关详细设计方法，请参阅方法详解。

步骤 4：设计网络拓扑

设计整体网络结构：编码器架构（如何提取特征）、各阶段的特征表示（李群使用 irreps）、池化/聚合策略，以及与任务需求相匹配的输出头。常见设计请参考拓扑模式。根据群的大小平衡深度与宽度。

步骤 5：选择实现库

根据库参考选择库。结合你的群、框架偏好（PyTorch/JAX）和性能需求进行选择。常见选择：e3nn（E(3)/O(3)，PyTorch）、escnn（离散群，PyTorch）、pytorch_geometric（置换，PyTorch）。确保库支持你的特定群。

步骤 6：创建架构规范

使用输出模板记录设计。包括：逐层规范、表示类型、库依赖项、预期参数量，以及伪代码或实际代码框架。此规范将指导实现及后续的等变性验证。如需可直接使用的实现模板，请参阅代码模板。此输出的质量标准见质量评分标准。

架构选择指南

按对称群分类

群	领域	推荐架构	库
Cₙ, Dₙ	二维图像	G-CNN、群等变 CNN	escnn, e2cnn
SO(2), O(2)	二维图像	Steerable 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 张量积层：

# Combine features with different angular momenta
tp = o3.FullyConnectedTensorProduct(
    irreps_in1, irreps_in2, irreps_out
)
output = tp(input1, input2)

等变非线性

问题：标准非线性（ReLU 等）会破坏等变性。

解决方案：

类型	工作原理	适用场景
基于范数	对
门控	使用不变量对等变特征进行门控	通用
张量积	通过 Clebsch-Gordan 实现非线性	e3nn、高质量
不变特征	仅应用于 l=0 分量	简单、快速

等变归一化

批归一化：会破坏等变性（每种方向的统计量不同） 解决方案：

• 层归一化（按样本归一化）
• 等变批归一化（按不可约表示通道归一化）
• 实例归一化（通常可用）

不变性池化

从等变特征获取不变输出：

方法	公式	适用场景
均值池化	对群取均值	连续群
求和池化	对元素求和	集合、图
最大池化	max
注意力池化	加权求和	重要性不同时

拓扑模式

编码器-解码器（分割、生成）

Input → [Equiv. Encoder] → Latent (equiv.) → [Equiv. Decoder] → Output

• 编码器：逐步提取特征
• 潜在层：等变表示
• 解码器：利用对称性进行重建

编码器-池化（分类）

Input → [Equiv. Encoder] → Features (equiv.) → [Invariant Pool] → [MLP] → Class

• 在末尾进行池化，获取不变特征
• 最终 MLP 处理不变表示

消息传递（图/点云）

Nodes → [MP Layer 1] → [MP Layer 2] → ... → [Aggregation] → Output

• 每层：聚合邻居并更新节点
• 聚合：为实现不变性使用求和/均值，为实现等变性按节点处理

库参考

e3nn（PyTorch）

群：E(3)、O(3)、SO(3) 优势：完整支持不可约表示、张量积、球谐函数 适用于：分子建模、3D 点云、物理

from e3nn import o3
irreps = o3.Irreps("2x0e + 2x1o + 1x2e")  # 2 scalars, 2 vectors, 1 tensor

escnn（PyTorch）

群：离散群（Cₙ、Dₙ）、连续 2D 群（SO(2)、O(2)） 优势：图像处理、文档完善 适用于：具有旋转/反射对称性的 2D 图像

from escnn import gspaces, nn
gspace = gspaces.rot2dOnR2(N=4)  # C4 rotation group

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	流形学习

输出模板

ARCHITECTURE SPECIFICATION
==========================

Target Symmetry: [Group name and notation]
Symmetry Type: [Invariant/Equivariant]
Task: [Classification/Regression/etc.]
Domain: [Images/Point clouds/Graphs/etc.]

Architecture Family: [e.g., E(3) Equivariant GNN]
Library: [e.g., e3nn]

Layer Specification:
1. Input Layer
   - Input type: [e.g., 3D coordinates + features]
   - Representation: [e.g., positions (l=1) + scalars (l=0)]

2. [Layer Name]
   - Type: [Convolution/Tensor Product/Message Passing]
   - Input irreps: [specification]
   - Output irreps: [specification]
   - Nonlinearity: [Gated/Norm/None]

3. [Continue for each layer...]

N. Output Layer
   - Aggregation: [Mean/Sum/Attention]
   - Output: [Invariant scalar / Equivariant vector / etc.]

Estimated Parameters: [count]
Key Dependencies: [library versions]

Code Skeleton:
[Provide implementation outline or pseudo-code]

NEXT STEPS:
- Implement the architecture using the specified library
- Verify equivariance through numerical testing after implementation