arXiv cs.LG

FuseSampleAgg: 融合した近隣サンプリングと集約を用いたミニバッチGNN

FuseSampleAgg: Fused Neighbor Sampling and Aggregation for Mini-batch GNNs

http://arxiv.org/abs/2511.13645v1


本稿ではFuseSampleAggというCUDAオペレーターを提案しています。このオペレーターは、GraphSAGEの1ホップおよび2ホップに対して近隣サンプリングと平均集約を一度のパスで融合します。これにより、ブロックのマテリアライズや追加のカーネル起動が不要になり、メモリトラフィックとオーバーヘッドが削減される一方で、GraphSAGEの平均セマンティクスを保つことが可能になります。Redditやogbn-arxiv、ogbn-productsのベンチマークで検証した結果、ogbn-productsでは最大51倍のスピードアップ、Redditではファンアウト10-10および15-10で約4倍、ogbn-arxivでは大きなファンアウトにおいて約3.3倍の速さを実現し、ピークGPUメモリの削減も最大100倍に達しました。このオペレーターは決定論的で、標準的なPyTorchオプティマイザと統合可能であり、全てのテーブルや図をCSVログから再現するスクリプトも提供されています。