SUREL+: スケーラブルな部分グラフベースのグラフ表現学習のためのウォークからセットへの移行

部分グラフに基づくグラフ表現学習（Subgraph-based graph representation learning, SGRL）は、モデルの表現力および一般化能力の優位性から、近年、グラフ上の多くの予測タスクにおいて強力なツールとして注目されている。これまでの多数のSGRLモデルは、各訓練またはテストクエリに対して部分グラフを抽出する際の計算コストが高いため、計算上の課題に直面していた。最近、SUREL（Subgraph Representation via Efficient Random Walks）が提案され、オフラインでランダムウォークをサンプリングし、オンラインでこれらのウォークを部分グラフの代替として結合することで、SGRLの高速化を実現した。サンプリングされたウォークが異なるクエリ間で再利用可能であるため、SURELはスケーラビリティおよび予測精度の面で最先端の性能を達成している。しかし、SURELはサンプリングされたウォークにおけるノードの重複によって引き起こされる高い計算負荷という課題を依然として抱えている。本研究では、ウォークではなくノード集合を用いて部分グラフを表現する新たなフレームワーク、SUREL+を提案する。この集合ベースの表現は定義上ノードの重複を排除するが、サイズが不規則になる可能性がある。この問題に対処するため、ノード集合を効率的に格納・アクセスできるカスタムなスパースデータ構造を設計し、並列バッチでの結合を可能にする専用演算子を提供している。SUREL+はモジュール化されており、複数種類の集合サンプラー、構造的特徴、ニューラルエンコーダーをサポートすることで、ウォークから集合への変換に伴う構造情報の損失を補完する。リンク予測、関係タイプ予測、高次パターン予測といった複数の予測タスクにおいて、広範な実験を実施した結果、SUREL+はSURELと比較して3～11倍の高速化を達成しつつ、同等またはより優れた予測性能を維持した。また、他のSGRLベースラインと比較して、約20倍の高速化を実現し、予測精度も顕著に向上した。