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、連続 2D 群には Steerable 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 画像	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