2019年夏季インターンに参加した東京大学の阿部健信です。Chemistryの大規模グラフのタスクへの拡張」というテーマで取り組んだ内容を説明させていただきます。インターン内容のスライドはこちらにアップロードされています。Chemistryで大規模グラフのデータを扱えるようにしました。GPU上で学習できるようにしました。Network(GNN)という分野が近年注目を集めています。[3]といった高機能で最適化されたGNNライブラリの開発が盛んに進められています。Chemistryは、PFNが開発しているGNNのオープンソースのライブラリです。名前からも分かるとおり、もともと分子など化学データへの適用を目的として作られたもので、qm9などの化学データセットが手厚くサポートされています。他にGNNの研究でよく用いられるSNSなどのネットワークデータのサポートはなされていませんでした。Chemistryでネットワークデータのサポートを行うことです。そのために、大きく以下の2つの内容を行いました。classification/regressionといった、グラフ全体の性質を学習するのに用いられます。巨大なネットワークデータに対しては、1つのグラフを入力として各頂点ラベルの分類を行うといった異なるタスクに用いられる事が多いです。classificationへの対応を行いました。具体的なフレームワークの違いはスライドをご参照ください。巨大でsparseなグラフのためのGNNの効率的な実装こちらが今回のインターン内容のメインで、巨大なグラフを動かすためには必要不可欠な内容でした。\)個の辺からなるグラフを考えます。passingにもとづくGNNでは、各頂点に対して近傍の頂点の特徴量のaggregationの操作を行います。invariantな様々な関数が用いられ、例えばよく使われるsumの場合は以下の式になります。aggregateされた特徴量)既存の実装は全てこの行列演算に基づくものでしたが、これは2つ問題点があります。1つめは、グラフが疎な際にメモリ的にも実行時間的にも無駄が生じてしまうことです。2つめは、batch化の際のゼロパディングのオーバーヘッドです。これらの問題を解決するために、辺の情報を密な隣接行列ではなく、疎なデータ形式で持たせるという事が考えられます。patternとして紹介されているデータの持ち方を新たに実装しました。Geometricでも採用されています。\)で行うことができます。batch化の際に、複数のグラフを全体として大きな1つのグラフとしてみなすことによってゼロパディングのオーバーヘッド完全になくすことができます。こちらも、より詳細な手法が知りたい方はスライドをご覧ください。patternによる実装の速度比較は以下のようになりました。3312頂点、4660辺の疎なネットワークグラフに対しては、CPUでは50倍以上、行列演算との相性が良いGPU上でも2倍以上の速度改善が見られました。1つ予想外だったのは、最大でも38頂点という比較的小さなグラフからなる化学データセットに対してもGPU上でも1.5倍程度の速度改善が見られたことです。これには、バッチ化のオーバーヘッドをなくす工夫が効いていると考えられます。1100万辺)を動かすことはできませんでした。(16GB)で動かすことができました。patternが早く、それではメモリが足りない場合はCooMatrixを使うとよい、という結果になりました。この結果を踏まえて、3つのパターンを場合に応じて使い分けられるように実装しています。Chemistryを使うモチベーションの1つになると思います。新しいモデルを自分で実装したいときに、ChainerでサポートされているCooMatrix演算を普通の行列演算と同じようなインターフェースで直感的に使えるのも魅力です。Chemistryにマージされています。新しい実装方針に対応しているモデルはまだ多くはありませんが、これからどんどん対応していく予定です。exampleのコードを動かすことで簡単に巨大グラフ上での学習ができるようになっているので、ぜひ試してみてください。分散処理、専用ハードウェア、バイオインフォマティクス、ケモインフォマティクスといった幅広い分野で研究開発を行っています。