グラフは、エッジで接続されたノードの潜在的に広範なウェブであり、社会的つながり、金融取引、交通、エネルギー グリッド、分子相互作用などのデータ間の関係を表現し、調査するために使用できます。 研究者がより多くのデータを収集し、これらのグラフィック画像を構築するにつれて、研究者は、グラフ ニューラル ネットワーク (GNN) の方法でディープ ラーニングを実行するために、より高速で効率的な方法と、より多くの計算能力を必要とします。
現在、MIT と IBM Research の研究者によって開発された SALIENT (サンプリング、スライシング、データ移動) と呼ばれる新しい方法は、計算の 100 つの主要なボトルネックに対処することで、トレーニングと推論のパフォーマンスを向上させます。 これにより、大規模なデータセットでの GNN の実行時間が大幅に短縮されます。たとえば、1 億ノードと 16 億エッジの規模で含まれます。 さらに、チームは、計算能力が XNUMX から XNUMX のグラフィカル プロセッシング ユニット (GPU) に追加されると、この手法が適切にスケーリングされることを発見しました。 この作品は、機械学習とシステムに関する第 XNUMX 回会議で発表されました。
「グラフの最先端の機械学習技術を非常に大きなデータセットにスケーリングする際に、現在のシステムが経験する課題に注目し始めました。 既存のシステムの多くは、主に GPU メモリに収まる小さなデータセットで優れたパフォーマンスを達成していたため、やるべきことがたくさんあることが判明しました」および人工知能研究所（CSAIL）。
膨大なデータセットとは、ビットコイン ネットワーク全体のようなスケールを意味し、特定のパターンとデータの関係がトレンドや不正行為を明らかにする可能性があります。 「ブロックチェーンには約 XNUMX 億のビットコイン取引があり、そのような共同ネットワーク内での違法行為を特定したい場合、そのような規模のグラフに直面することになります」と共著者のシニア リサーチ サイエンティスト兼マネージャーである Jie Chen 氏は述べています。 IBM Research および MIT-IBM Watson AI Lab. 「毎日生成される新しいデータのペースに追いつきたいので、そのようなグラフを処理でき、処理をできるだけ効率的にできるシステムを構築したいと考えています。」
Kaler と Chen の共著者には、卒業研究の一環として SALIENT を開発した Jump Trading の Nicolas Stathas MEng '21 が含まれます。 元 MIT-IBM ワトソン AI ラボのインターンで MIT 大学院生の Anne Ouyang。 MIT CSAIL ポスドク Alexandros-Stavros Iliopoulos; MIT CSAIL 研究科学者 Tao B. Schardl; Charles E. Leiserson は、MIT の Edwin Sibley Webster 電気工学教授で、MIT-IBM Watson AI Lab の研究者です。
この問題に対して、チームは方法を開発する際にシステム指向のアプローチを採用しました。SALIENT です。 これを行うために、研究者は、機械学習モデルを構築するためのインターフェイスである PyTorch ジオメトリックやディープ グラフ ライブラリ (DGL) などの既存の機械学習フレームワークに適合するコンポーネントの重要な基本的な最適化を実装しました。 このプロセスは、より速い車を作るためにエンジンを交換するようなものだとスタタスは言います。 彼らの方法は既存の GNN アーキテクチャに適合するように設計されているため、ドメインの専門家はこの作業を指定された分野に簡単に適用して、モデルのトレーニングを促進し、推論中に洞察をより迅速に引き出すことができます。 チームが決定した秘訣は、すべてのハードウェア (CPU、データ リンク、および GPU) を常にビジー状態に保つことでした。その間、CPU はグラフをサンプリングし、データ リンクを介して転送されるデータのミニバッチを準備します。 、より重要な GPU が機械学習モデルのトレーニングまたは推論の実行に取り組んでいます。
研究者はまず、一般的に使用されている GNN 用の機械学習ライブラリ (PyTorch ジオメトリック) のパフォーマンスを分析し、利用可能な GPU リソースの使用率が驚くほど低いことを示しました。 単純な最適化を適用することで、研究者は GPU 使用率を 10% から 30% に改善し、公開ベンチマーク コードと比較して 1.4 倍から 50.4 倍のパフォーマンス向上を実現しました。 この高速なベースライン コードは、XNUMX 秒でアルゴリズム (エポック) を介して大規模なトレーニング データセットに対して XNUMX つの完全なパスを実行できます。
さらなるパフォーマンスの向上を求めて、研究者はデータ パイプラインの最初に発生するボトルネック、つまりグラフ サンプリングとミニバッチ準備のアルゴリズムの調査に着手しました。 他のニューラル ネットワークとは異なり、GNN は近傍集約操作を実行します。これは、グラフ内の他の近くのノードに存在する情報を使用して、ノードに関する情報を計算します。たとえば、ソーシャル ネットワーク グラフでは、ユーザーの友人の友人からの情報です。 GNN のレイヤー数が増えると、ネットワークが情報を得るために到達しなければならないノードの数が爆発的に増加し、コンピューターの限界を超える可能性があります。 近隣サンプリング アルゴリズムは、収集するノードのより小さなランダムなサブセットを選択することで役立ちます。 ただし、研究者は、これの現在の実装は遅すぎて、最新の GPU の処理速度についていけないことを発見しました。 それに応じて、彼らはサンプリング速度を向上させるデータ構造、アルゴリズムの最適化などの組み合わせを特定し、最終的にサンプリング操作だけで約 50.4 倍改善し、エポックごとの実行時間を 34.6 秒から XNUMX 秒に短縮しました。 また、適切なレートでのサンプリングを推論中に行うことができ、全体的なエネルギー効率とパフォーマンスが向上することも発見しました。これは、文献では見過ごされていた点です。
以前のシステムでは、このサンプリング ステップはマルチプロセス アプローチであり、余分なデータとプロセス間の不要なデータ移動が作成されていました。 研究者は、CPU 上のデータを共有メモリに保持する軽量スレッドで単一のプロセスを作成することにより、SALIENT メソッドをより機敏にしました。 さらに、SALIENT は最新のプロセッサのキャッシュを利用し、CPU コア キャッシュの共有メモリ内で、関心のあるノードとその周囲の隣接ノードおよびエッジから関連情報を抽出する機能スライシングを並列化します。 これにより、エポックごとの全体的なランタイムが 34.6 秒から 27.8 秒に短縮されました。
研究者が対処した最後のボトルネックは、必要になる直前にデータを準備するプリフェッチ ステップを使用して、CPU と GPU 間のミニバッチ データ転送をパイプライン化することでした。 チームは、これによりデータ リンクでの帯域幅の使用が最大化され、メソッドが完全に使用されるようになると計算しました。 ただし、約 90% しか見られませんでした。 彼らは、人気のある PyTorch ライブラリで、CPU と GPU 間の不要な往復通信を引き起こすパフォーマンス バグを特定して修正しました。 このバグを修正したことで、チームは SALIENT でエポックあたり 16.5 秒のランタイムを達成しました。
「私たちの研究は、悪魔は細部に宿るということを示したと思います」と Kaler は言います。 「グラフ ニューラル ネットワークのトレーニング時にパフォーマンスに影響を与える詳細に細心の注意を払うと、膨大な数のパフォーマンスの問題を解決できます。 私たちのソリューションでは、このようなシステムの理想的な目標である GPU 計算が完全にボトルネックになってしまいました。」
SALIENT の速度は、ogbn-arxiv、ogbn-products、および ogbn-papers100M の 100 つの標準データセットと、ファンアウト (CPU が GPU 用に準備するデータの量) のレベルが異なるマルチマシン設定で評価されました。最新の最先端の GraphSAGE-RI を含むいくつかのアーキテクチャにまたがっています。 各設定で、SALIENT は PyTorch ジオメトリックを上回りました。特に、100 億のノードと 16 億を超えるエッジを含む大規模な ogbn-papersXNUMXM データセットで、SALIENT は XNUMX つの GPU で実行され、最初に作成された最適化されたベースラインよりも XNUMX 倍高速でした。この作品; XNUMX 個の GPU を使用すると、SALIENT はさらに XNUMX 倍速くなりました。
他のシステムはハードウェアと実験のセットアップがわずかに異なるため、常に直接比較できるわけではありませんが、SALIENT は依然としてそれらを上回っていました。 同様の精度を達成したシステムの中で、代表的なパフォーマンス数値には、99 つの GPU と 32 の CPU を使用した場合に 13 秒、1,536 の CPU を使用した場合に 20 秒が含まれます。 対照的に、16.5 つの GPU と 16 の CPU を使用した SALIENT の実行時間は 320 秒で、16 の GPU と XNUMX の CPU ではわずか XNUMX 秒でした。 「以前の研究で報告された最終的な数値を見ると、XNUMX GPU ランタイム (XNUMX 秒) は、このデータセットで以前に報告された他の数値よりも桁違いに高速です」と Kaler 氏は言います。 研究者は、分散設定に移行する前に単一のマシン用にコードを最適化するというアプローチがパフォーマンスの改善の一部であると考えています。 Stathas 氏によると、ここでの教訓は、「複数のコンピューターにスケールアップする前に、所有しているハードウェアを効率的に、そして極端に使用する方が理にかなっています」ということです。これにより、コストと二酸化炭素排出量を大幅に節約できます。これはモデル トレーニングに付属しています。
この新しい容量により、研究者はますます大きなグラフに取り組み、より深く掘り下げることができるようになります。 たとえば、前述のビットコイン ネットワークには 100,000 のノードが含まれていました。 SALIENT システムは、1,000 倍 (または XNUMX 桁) 大きいグラフを処理できます。
「将来的には、各ノードのプロパティを分類または予測するために実装した既存のアルゴリズムでこのグラフ ニューラル ネットワーク トレーニング システムを実行するだけでなく、次のようなより詳細なタスクも実行したいと考えています。グラフ内の一般的なパターン (サブグラフ パターン) を特定することで、[これは] 金融犯罪を示すのに実際に興味深い可能性があります」とチェンは言います。 「また、金融犯罪の同じ悪役に対応している可能性があるという意味で類似しているグラフ内のノードを特定したいと考えています。 これらのタスクには、追加のアルゴリズムを開発する必要があり、場合によってはニューラル ネットワーク アーキテクチャも開発する必要があります。」
この研究は、MIT-IBM ワトソン AI ラボの支援を受け、一部は米国空軍研究所と米国空軍人工知能アクセラレーターの支援を受けました。
