技術要約：Graph Set Transformer (GST)

問題提起

グラフの集合（順序のないグラフの集まりであり、予測が局所的なグラフ構造とグローバルな集合全体のコンテキストの両方に依存するもの）を扱う機械学習タスクは、重大なアーキテクチャ上の課題を抱えています。DeepSetsやSet Transformerといった既存のアプローチは、通常、逐次的なパイプラインに従います。すなわち、まずGNNを用いて個々のグラフを固定されたベクトル埋め込みにエンコードし、その後に置換不変なセットヘッドを適用してこれらの埋め込みを集計するという手法です。

著者らは、このような逐次的なステージングがボトルネックを生み出すと主張しています。このようなパイプラインでは、集合レベルの情報は、他のグラフに関する認識なしに計算された要約埋め込みから最終段階でのみ再構成されます。この分離により、特定のグラフの局所的な構造が集合全体のグローバルなコンテキストとどのように相互作用するかをモデルが効果的に推論することを妨げます。これは、反応収率の予測や、反応中心の特定（原子の挙動がすべての参加分子に依存する場合など）のようなタスクにおいて不可欠な要素です。

手法

本論文では、グラフの集合 $S = \{G_1, G_2, \dots\}$ を処理するために設計されたニューラルアーキテクチャである Graph Set Transformer (GST) を提案します。このアーキテクタは、局所的なノードレベルの特徴伝播と、クロスグラフのコンテキストモデリングを、各層で逐次的ではなく交互に組み合わせるように設計されています。

アーキテクチャの概要

GSTの核となるのは、ノード特徴 $x_i$ を更新する graphsetconv ブロック です。このブロックは、以下の5つの明確なステップで動作します。

1. 局所的メッセージパッシング (Local Message Passing): 標準的なグラフ畳み込みネットワーク (GCN) 層が、1ホップ隣接ノードから特徴を集約します。著者らは、集合内のグラフのサイズの違いを扱うために、LayerNorm の代わりに GraphNorm を使用しており、これによりグラフレベルの活性化強度の安定化を図っています。
2. アテンショナル・プーリング (Attentional Pooling): 各グラフは、ノード特徴のソフトマックス重み付き和を介して、単一の $d$ 次元トークン ($z_g$) へと集約されます。これが、クロスグラフの情報流におけるボトルネックとなります。
3. セットレベルのコンテキスト化 (Set-Level Contextualization): すべてのグラフのグラフ・トークン ($z_g$) が、マルチヘッド自己注意 (MHSA) ブロックを用いて互いに注意を向け合います。このステップは、集合全体のコンテキスト（例：極値の特定、分布の要約、ランキングなど）を捉え、グラフ・トークン ($\tilde{z}_g$) を更新します。
4. コンテキスト・ルーティング (Context Routing): 更新されたセット・コンテキストが、個々のノードへとルーティングされます。著者らは、このステップに対して2つのバリアントを提案しています。
   • GST-bc (Broadcast): 更新されたグラフ・トークン $\tilde{z}_g$ が、そのグラフ内の全ノードに直接ブロードキャストされます。
   • GST-ca (Cross-Attention): 各ノードが、クエリ機構を介して更新されたグラフ・トークンの集合に対して注意を向けます。
5. ゲート付き融合 (Gated Fusion): 元のノード特徴、局所コンテキスト、およびセットコンテキストが、学習された非対称ゲート機構を用いて融合されます。浅いMLPがゲートを生成し、ノードがセットコンテキスト ($v_i$) と局所コンテキスト ($u_i$) のどちらをどの程度吸収すべきかを決定します。これにより、ノードは、自身のグラフに十分な情報が既に含まれている場合、セットコンテキストを無視することが可能になります。

フルモデルは、これら $L$ 個のブロックのスタックと、タスク固有の出力のためのMLPヘッドで構成されます。

主な貢献

1. アーキテクチャの革新: 本論文は、特徴抽出とセットのコンテキスト化を各層で融合させる新しいアーキテクチャを提案しており、標準的な「エンコードしてから集計する」パイプラインとは対照的です。
2. ゲーティング機構: 非対称ゲート機構の導入により、局所的情報とグローバル情報の適応的な混合が可能になります。これは、「あるノードの親グラフは常にプライマリである」という帰納バイアスを尊重しつつ、セットの影響をコンテキスト依存にするものです。
3. 包括的な評価: 著者らは、GSTを以下の対象で評価しています。
   • 合成ベンチマーク: セット条件付きの構造的推論（例：RING-K, BALL-K, DIST, および外挿タスク）を分離するように設計されたもの。
   • 化学的応用: 反応収率予測（Buchwald-Hartwig）および反応中心の特定（USPTO-15K）。
   • 画像分類: CIFAR-10 セット分類（グラフエンコードされた画像を使用）。
4. アブレーションと分析: パラメータ数を一致させた比較およびアーキテクチャのアブレーションを通じて、性能の向上はモデル容量の増加によるものではなく、交互に組み合わせるアーキテクチャに由来することを証明しています。

結果

パラメータ予算が一致している条件下で、GSTはすべての評価設定において、ベースライン（GCN + DeepSets および GCN + Set Transformer）を一貫して上回りました。

• 合成ベンチマーク: GSTのバリアントは、特にセット条件付き推論の複雑さが増すにつれて（例：RING-K タスクにおける大きな $K$）、ベースラインを大幅に上回ります。性能差はセットサイズのシフトが増加するにつれて拡大しており、これはGSTが単なる訓練セットの統計量ではなく、セット操作を捉えていることを示しています。
• 反応収率予測: GSTは優れた $R^2$ スコアを達成しました。最も顕著な利得は、低データかつノイズの多いレジームで見られました（例：2.5% の訓練データにおいて、GCN-ST の $R^2 = 0.18$ に対し、GST-bc は $R^2 = 0.53$）。
• 反応中心予測: USPTO-15K ベンチマークにおいて、GST-bc は自然なパラメータ予算において、最良のベースラインに対して 5.7 AP ポイントのリードを達成しました。この優位性は、GSTをベースラインのパラメータ数に合わせて縮小した場合や、ベースラインをGSTの容量に合わせて拡張した場合でも維持されました。
• 画像分類: CIFAR-10 セット分類において、GSTのバリアントは、ROC-AUC および PR-AUC メトリックの両方で、GCNベースのベースラインを一貫して上回り、その差はセットサイズが増加しても有意なまま維持されました。

計算コスト: 著者らはトレードオフについても認めています。各層での交互的なセットレベルのアテンションにより、GSTはベースラインと比較して、1トレーニングステップあたりの実時間コストが高くなります（約4〜6倍）。しかし、ブロードキャスト・バリアント (GST-bc) は、精度への影響を最小限に抑えつつ、より効率的な代替案を提供します。

重要性と主張

本論文の主な重要性は、グラフのエンコーディングとセット推論の間のボトルネックを取り除く能力にあると主張されています。局所的情報とセットレベルの情報を交互に組み合わせることで、このアーキテクチャは、最終的なプーリング段階を待つことなく、集合内の要素が特徴抽出フェード中に動的に互いにコンテキスト化することを可能にします。

著者らは、このアーキテクチャの転換は、分子（または分子内の原子）の挙動が反応混合物中の他の分子の存在とその特性に本質的に依存する化学のような領域において、特に重要であると主張しています。結果は、「交互に組み合わせる帰納バイアス」が、セット条件付きの構造的推論を必要とするタスクにおいて、アーキテクチャの限界を補うための生のモデル容量に頼ることなく、強力な代替手段を提供するという根本的な利点を持っていることを示唆しています。