Graph Neural Network を用いたグラフの木幅予測

Masahiro Sakai

2019-10-10 12:10:39

本記事は、2019年夏のインターンシップに参加された中野裕太さんによる寄稿です。


皆様はじめまして。2019 年 PFN 夏季インターンシップに参加していた北海道大学の中野裕太です。本ブログでは、私が夏季インターンで取り組んだテーマである、「Graph Neural Network を用いたグラフの木幅予測」について説明します。

要旨

与えられた無向グラフがどれくらい木に近いかを表す値である木幅は、グラフ上の組み合わせ最適化問題に対するアルゴリズムの効率性や解そのものと深く関係しています。しかし、木幅を計算することは NP 困難なため、木幅を計算するには頂点数に対し指数時間かかってしまいます。そこで、今回 Graph Neural Network を用いた 2 つの方法でこの問題にアプローチしました。1 つ目は、よく知られた既存のアルゴリズムと組み合わせ探索木の枝刈りを行い高速化を図り計算する方法です。結果は、多くの場合で枝刈りがうまく働き、関数呼び出し回数が 90% 以上減少しました。2 つ目はグラフ回帰問題として直接木幅を計算する方法で、平均絶対誤差が 0.75 となりました。また、真値との誤差が 5 を超えるような、大きく間違えた例はありませんでした。

はじめに

木幅とは?

グラフ理論において木幅とは、ある無向グラフに対し定義されるグラフ不変量の一つであり、大雑把にいうとグラフがどれくらい木に近いかを表す値です。この値を厳密に定義するためにまず、グラフに対する木分解*1というものを定義します。

ある無向グラフ \(G=(V,E)\) に対する木分解とは、下記の条件を満たす \((X = \{X_i\ |\ i \in I \},\ T = (I,\ F))\) のことをいいます。

  1. \(i \in I\) に対し、\(X_i \subseteq V\) である。
  2. \(V = \bigcup_{i \in I}\ X_i\)
  3. 全ての \(\{v, w\} \in E\) に対し、\(v, w \in X_i\) となる \(i \in I\) が存在する。
  4. 全ての \(v \in V\) に対し、\(\{i \in I\ |\ v \in X_i\}\) で誘導される、\(T\) の部分グラフが木となる。

下図に、グラフと木分解の一例を示します。

また、木分解 \((T = (I,\ F),\ X)\) に対する幅を \(\max_{i \in I} |X_i| – 1\) と定義します。例えば、先程の図における木分解の幅は 2 です。ここで、ある無向グラフ \(G=(V,E)\) に対する木幅 \(tw(G)\) は、\(G\) に対し可能な全ての木分解における幅の値の最小値として定義されます。

例えば、木に対する木幅を考えてみると、\(X_i\) を各辺の端点 2 点とすることで、定義を満たす \(T\) が構成できるため、頂点数 2 以上の木は木幅 1 を持ちます。このように、木幅が 1 に近いほどそのグラフが木に近いこと、つまり木へと分解しやすいということを意味します。また、任意の 3 頂点がサイクルを形成し、最も木から離れている完全グラフは、\(X_1\) に全ての頂点を含める以外に条件を満たす \(T\) を構成できないため、\(|V| – 1\)の木幅を持ちます。

*1 機械学習の分野などでは Junction Tree とも呼ばれます。

木幅は何の役に立つ?

グラフ上における組み合わせ最適化問題は、大きく 2 種類に分けられます。1 つは与えられたグラフのサイズ(辺数や頂点数)に対し多項式時間で解ける問題、もう 1 つは多項式時間で解けるアルゴリズムが \(P \neq NP\) 予想の下で存在が知られていない問題です。前者の例として、最短経路問題や最大流問題、後者の例として、最大独立点集合問題や巡回セールスマン問題があります。

しかし後者の問題のほとんどにおいて、入力として与えられるグラフが木である場合に、グラフのサイズに対する多項式時間で解けるということが知られています。では、木に似ているグラフ、つまり木幅が小さいグラフではどうでしょうか?これは、先程述べた木分解を用いることで、 木幅を固定した際、つまり木幅 \(k\) が定数の場合、グラフのサイズに対し多項式時間で解けることが知られています。これは、アルゴリズム分野における Fixed Parameter Tractability (FPT) と深く関係しています。詳しくは [1] をご参照ください。

また、先程も述べたように木幅はグラフ不変量であるため、他のグラフ不変量との間で、上界や下界の関係が存在します。例えばグラフ彩色数 [2] に関して、木幅 \(k\) を持つグラフは高々 \(k + 1\) の彩色数を持つといった関係があることが知られています。詳しくは、[3] をご参照ください。

このようなことから、木幅を計算することや予測することは、グラフ上の組み合わせ最適化問題に取り組む上で重要なことといえます。

Graph Neural Network とは?

Graph Neural Network (GNN) とは、グラフ構造を扱う新しいニューラルネットワークです。一般的なニューラルネットワークでは、入力として実数ベクトルや行列、テンソルなどを扱っていますが、GNN では下記に示すフレームワークを用いるなどで、グラフ構造を扱うことを可能にしています。GNN が主に扱う問題としては、教師ありのグラフ分類・回帰問題やノード分類・回帰問題があります。GNN に関するサーベイ論文として [4] などがあります。詳細はこちらを参照してください。

GNN にも様々な種類がありますが、今回用いる大きなフレームワークとして [6] で提案されているものを用います。これは、各ノードに特徴ベクトルが与えられており、それらを一定回数更新していくことで、ノードに対する特徴表現を獲得し、グラフ全体における全てのノード特徴表現を集約することで、グラフに対する特徴表現を獲得するという仕組みになっています。

入力として、無向グラフ \(G=(V, E)\) が与えられ、各頂点 \(v \in G\) に対し特徴ベクトル \(h_v\) が割り当てられていると仮定し、具体的なノード特徴ベクトルの更新方法を下図に示します。

 

Aggregate で更新ノードの近傍(\(\mathcal{N}(v)\) で表される部分です)の特徴ベクトルを集約し、Combine で、自身の特徴ベクトルと集約した特徴ベクトルを用い、新しい特徴ベクトルを生成します。グラフ分類・回帰タスクを行う場合、最終的に Readout でグラフ全体の特徴ベクトルを生成します。

実際のモデルにおいて、Aggregation や Readout の部分にあたる特徴ベクトルの多重集合に対する関数は、Max, Average, Sum などが用いられることが多いです。例えば、Kipf らにより提案されたノード分類における GNN モデルである、 Graph Convolutional Networks (GCN) [5] では、Aggregation の部分に Element-wise mean を、Combine の部分に1層のパーセプトロンを用いています。

では、GNN の性能を最大限に引き出すためにはどのような操作が良いでしょうか?これに関し Xu らは、ノード近傍の特徴ベクトルの多重集合における区別可能性を基に考察しています [6]。これを簡単に説明します。

上図におけるノードの色は、特徴ベクトルを表しており、同じ色は同じ特徴ベクトルです。今、黄色のノードの特徴ベクトルを更新する際の Aggregation について、具体的な例として上げた Max と Average と Sum の 3 つについて考えます。Max や Average を取った場合は、左右どちらとも緑色が 1 つになってしまい、2 つのグラフを区別できません。しかし、Sum をとった場合、左は緑が 3 つ、右は緑が 2 つとなり、区別できるようになります。

このようないくつかの考察を通し、Xu らは Graph Isomorphism Network (GIN) と呼ばれる、Aggregation や Readout を Sum ベースに行う、シンプルなネットワーク構造を提案し、多くのグラフ分類タスクで State-of-The-Art を達成しています。今回、木幅を予測するために試したいくつかの GNN 構造は、この GIN における、Aggregation や Readout の部分を変えて得られるものを用いました。詳しくは後述するコードをご覧ください。

どのように木幅を予測するか

今回、下記に示す 2 つの手法を用い木幅の予測を行いました。

  1. 既存のアルゴリズムと GNN を組み合わせての予測
  2. GNN を用いた直接予測

それぞれの手法について説明する前に、今回用いたグラフデータセットについて簡単に説明します。

今回用いたグラフデータセット

GNN を学習させるためのデータセットが必要になりますが、木幅に関するデータセットはほとんどありません。見つけたデータセットとして、Parameterized Algorithms and Computational Experiments Challenge (PACE) と呼ばれる様々な組合せ最適化問題に対するアルゴリズムの設計を競うコンテストにおいて、2016 年、2017 年に木幅に関するコンテスト [7, 8] が開かれた際に用いられたデータセット [9] がありましたが、200 グラフほどしか含まれておらず、GNN の学習に用いるためには少し心許ない数量でした。そこで、ランダムにグラフを生成し、木幅を既存の木幅計算プログラムで計算することで、GNN の学習に用いるためのデータセットを作成しました。

今回、グラフを生成する際に用いたランダムグラフ生成モデルは下記の 3 つです。

各モデルからそれぞれ 10000 グラフを生成し混ぜ合わせたデータセットを用いました。後述する各分類タスクにおいては、ラベルが均等になるように選択し、回帰タスクにおいては、30000 グラフの中から、10000 グラフをランダムに選択し、それぞれのタスクにおいてデータセットとして 1 割をテストデータ、残りを訓練データとしています。

また、ラベル付けのために木幅を計算するための既存のプログラムとして、先程も取り上げた 2017年に開催された PACE の木幅の計算コンテストにおいて提出された、明治大学の玉木先生の実装 [10] を用いました。

アプローチ 1

既存の木幅を計算するアルゴリズムに対し、GNN による木幅に関する分類タスクの予測結果を探索状態の枝刈りに用いることで、既存のアルゴリズムの高速化を図りました。

既存アルゴリズムについて

木幅を計算する既存アルゴリズムは複数あります。詳しくは [11] をご参照ください。今回は、[11] にて取り上げられていた 空間計算量、時間計算量ともに \(O^{\ast}(2^n)\)*2 である木幅計算アルゴリズムを参考にした、下記に示す Algorithm 1 を用います。アルゴリズム中に登場する、\(\mathrm{Q}(G, S, v)\) の詳細については [11] を御覧ください。これは、あるグラフ \(G=(V, E)\) に対し \(\mathbf{TWCheck}(G, V, opt)\) を呼び出すことで、\(G\) が木幅 \(opt\) 以下であるかどうかを判定します。つまり、\(opt = 1\) から値を 1 ずつ増やしながら順に代入していき、初めて true を返した値、または \(opt = |V| – 1\) から値を 1 ずつ減らしながら順に代入していき、初めて false を返した値に 1 を加えた値として、\(G\) の木幅を計算できます。以降、前者を Lower-Calculation、後者を Upper-Calculation と呼びます。

*2 Big-O Star 記法と呼ばれ、ある評価関数の指数部オーダのみを記述するためのものです。Parameterized Algorithm 分野でよく用いられます。詳しくは文献 [14] をご参照ください。

今回導入する枝刈り

上記の擬似コードを見ていただくとわかるように、一度の木幅判定において最大で \(\mathbf{TWCheck}\) が \(O(|V|!)\) 回呼び出されてしまいます。そこで、 \(\mathbf{TWCheck}(G, S, opt)\) が、\(G\) において、\(V\) の部分集合 \(S\) で誘導される部分グラフ \(G[S]\) の木幅が \(opt\) 以下であるかどうかを返す関数であることに着目し、この誘導部分グラフの木幅を GNN により分類し、その結果次第で枝刈りを行う部分を追加します。これにより、結果が不正確なものになる可能性と引き換えに、再帰回数が減るため関数呼び出し回数の減少、実行時間の削減が見込まれます。この枝刈りについて具体的に説明していきます。

まず、木幅がしきい値 \(thresh\) 以下かどうかでラベル付けを行った二値分類タスクの学習を行った GNN を考えます。枝刈りを行うために、 \(G[S]\) に対しこの GNN を用い推論を行います。

ここで、現在の \(opt\) の値が \(thresh\) より小さく、GNN による予測結果が \(thresh\) より大きい場合、部分グラフの木幅が十分に大きいと判断し、枝刈りを行って false を返します。これを Upper-Prune と呼びます。一方、現在の \(opt\) の値が \(thresh\) より大きく、GNN による予測結果が \(thresh\) より小さい場合には、部分グラフの木幅が十分小さいと判断し、枝刈りを行って true を返します。これを Lower-Prune と呼びます。これら 2 つの枝刈りをまとめたものを下図に示します。

この枝刈り方法の利点として、下記の 3 つがあります。

  • 先述の Lower-Calculation や Upper-Calculation と適切に組み合わせることで、正確に上界や下界の計算が可能な点
  • Softmax Cross Entropy 関数を損失関数に用いモデルを学習させることで、モデルの予測確率を考慮した枝刈りが可能になる点
  • 今回の二値分類モデルだけではなく、多値分類モデルを用いる際にも容易に拡張できる点

正確な上界・下界の計算が可能

これら 2 つの枝刈り方法は、先程述べた Lower-Calculation や Upper-Calculation と適切に組み合わせることにより、与えられたグラフに対する木幅の上界や下界を計算できます。前者は Upper-Prune と、後者は Lower-Prune と組み合わせることにより、それぞれ正確に上界と下界を計算できます。

前者について少し説明します。Lower-Calculation のアルゴリズムを、\(opt = |V| – 1\) が入力された場合必ず \(true\) を返すように書き換えます。ここで Upper-Prune において、本来は True であるべき場所で False を返してしまった場合を考えます。これにより、本来はある木幅 \(k\) 以下であると判定されるべき値で、木幅は \(k\) より大きいと判定することとなり、本来よりも大きい値として算出してしまいます。後者についても同様に示すことができます。

モデルの予測確率を用いた枝刈りが可能

Cross Entropy 関数を損失関数としてニューラルネットワークを学習するということは、モデルの出力ベクトルを活性化関数(Softmax 関数)に通した際の各要素が、各クラスへの割り当て確率となるように学習することとみなせます。そのため、出力ベクトルの要素の最大値 \(x_{max}\) に対し、\(x_{max} \geq \mathrm{prob\_bound}\) となる場合にのみ、枝刈りを行うことで、モデルの予測確率が \(\mathrm{prob\_bound}\) 以上である予測のみを考慮でき、最終的に出力する木幅の精度の向上が見込まれます。

多値分類モデルへの拡張が可能

二値分類を行う GNN だけでなく、多値分類タスクに対して訓練済みの GNN を用いる場合にも容易に拡張可能です。また、前述の予測確率を組み合わせることで、さらなる精度の向上が見込まれます。

GNN は木幅に関する分類タスクを適切に学習できるか?

今回導入した 2 つの枝刈りにおいて、GNN による予測性能が重要になるため、分類タスクを適切に学習した GNN が要求されます。しかし、木幅は [12] でも述べられているように、大域的な連結度を表す値であり、近傍の集約を繰り返し特徴ベクトルを更新していく今回の GNN では、木幅に関する特徴を適切に学習できない可能性があります。そのため、まずは木幅に関する分類タスクの実験による検証を行いました。

今回は、しきい値 \(thresh\) に関して \(\frac{|V|}{2}\) を用いる二値分類タスクと、\(\lfloor \frac{5tw(G)}{|V|} \rfloor\) の値でラベル付けした 5 クラス分類タスクを考えました。詳しいハイパパラメータの設定などは後述するコードをご覧ください。

GNN における木幅に関する分類タスク実験の結果

それぞれのタスクにおける実験結果として、テストデータに対する精度の推移を下図に示します。

Operation Type における、A は Aggregation、R は Readout の種類を表しています。どちらの場合も、Aggregation に max を使用したもの以外は 9 割前後、8 割前後といった高い精度が得られています。この結果より、GNN は Aggregation, Readout の方法をうまく選択することで、木幅に関する分類タスクを適切に学習できると判断しました。今回用いる GNN モデルにおける Aggregation, Readout の方法として、総合的に精度が良かった、どちらも sum を用いる組み合わせを使用しました。

アプローチ 2

先程は分類タスクに対する GNN を利用しましたが、こちらは直接グラフ回帰問題として GNN により End-to-End で木幅を予測しようというアプローチです。先程の実験結果より、GNN を用いた木幅に対する分類タスクの精度が良く、GNN が、木幅に関する特徴を適切に捉えていると考えました。そのため、End-to-End で予測を行うことで、より精度良い予測が可能になると考えました。しかし先程とは異なり、この方法では正確な上界や下界などは得られません。

今回、アプローチ 1 で用いた GNN 構造により出力されるベクトルの次元数を 32 に変え、続いてそれを 2 層の MLP に入力し、最終的に 1 つの実数値を得るといったネットワーク構造を用いました。詳しいネットワーク構造は後述するコードをご覧ください。

実験結果

アプローチ 1

今回、先程紹介した既存のアルゴリズムと GNN による枝刈りを導入したアルゴリズムの両方を python 3.7 にて実装し、先程のランダムグラフ生成モデルより生成された 50 個のグラフに対し、アプローチ 1 の有効性を検証するための実験を行いました。

今回、GNN による枝刈りを導入したアルゴリズムのうち下記の 3 つを実装しました。

  • 上界を求める Upper-Prune と Lower-Calculation を組み合わせたもの(Upper-Alg)
  • 下界を求める、 Lower-Prune と Upper-Calculation を組み合わせたもの(Lower-Alg)
  • 双方の枝刈りを採用した上で、Lower-Calculation を行うもの(Combine-Alg)

これらと Lower-Calculation ベースの既存のアルゴリズム(Existing-Alg)との比較を行いました。また、グラフの大きさによっては途方も無い時間がかかってしまうことが予期されるため、タイムアウトを 10 分に設定しました。

結果を下図に示します。

10 分以内に実行が完了したインスタンス数は、既存のアルゴリズムでは 9 となったのに対し、 Combine-Alg では 18 と、既存アルゴリズムより多くのインスタンスに対し計算できていることがわかりました。また Lower-Alg や Upper-Alg ではそれぞれ 8, 10 となり、既存アルゴリズムとほぼ変わりませんでした。

続いて、この枝刈りによりどの程度既存のアルゴリズムが効率よく動作するようになったかを見るために、関数呼び出し回数を比較しました。今回、Upper-Alg で枝刈りが行われたインスタンスは、Combine-Alg や Lower-Alg ではタイムアウトしていたため、Combine-Alg と Lower-Alg に関して、平均の減少率を計算しました。これは、どちらも 9 割前後となり、GNN による枝刈りが効率よく動作していることがわかりました。

また、精度を比較するために、近似解と真の値との絶対値の平均を計算したところ、Combine-Alg では 1.6 、Lower-Alg では 2.6 となりました。より積極的な枝刈りを行う Combine-Alg のほうが精度が良く、直感と反する結果となっているため、今後より多くのインスタンスで実験を行い検証する必要があります。

しかし、表には載せていませんが総合的な実行時間を比較すると、Combine-Alg などで平均よりも多く枝刈りを行っているいくつかのインスタンスを除き、既存のアルゴリズムのほうが実行時間が速いです。これは、GNN の推論に時間がかかっているためと考えられます。多くの場合、総実行時間の 9 割ほどを占めていました。

下記に、どの程度関数呼び出し回数が減少しているかの例として 、2 つのインスタンスに対し、実際のグラフと関数呼び出し回数をプロットした棒グラフを示します。

アプローチ 2

テストデータに対する Mean Absolute Error (MAE) の値の推移を下図に示します。

このグラフを見ていただくとわかるように、分類タスクと同様に、Aggregation, Readout ともに sum を用いるモデルが最も精度がよく、最終的な値は 0.75 ほどとなりました。これは、予想したように手法 1 よりも良い結果が得られています。また、テストデータに対し、予測した木幅と実際の木幅をプロットした散布図を下記に示します。左側が GNN (Aggregation, Readout ともに sum を用いるモデル)による結果で、右側が同じデータを用い、頂点数、辺数、平均次数で線形回帰を行った際の結果です。

散布図を見るとわかるように、大きく外している例は無いことがわかります。これら結果を比較してもわかるように、比較的精度良く予測ができていると考えられます。

まとめ

まず、アプローチ 1 についてですが、いくつかのインスタンスで枝刈りルールがうまく働き、関数呼び出し回数の観点において 90% ほどの減少が見られました。しかし、推論に時間がかかっており、総実行時間の面では改善が見られないインスタンスも多くありました。今回のインターンシップでは推論の高速化については特別には取り組みませんでしたが、推論フレームワーク(例えば Menohchainer-trt, chainer-compiler)の利用や各種の最適化によって、ある程度は高速化できる可能性があります。

続いて、アプローチ 2 についてですが、平均絶対誤差が 0.75 ほどのモデルが学習でき、いくつか計算が容易なグラフ不変量を用い線形回帰を行った結果と比較し、精度よく予測ができていました。

今後の課題として、より素早く正確に推論が可能なモデルや推論フレームワークを使用すること、素早く計算が終わる状態から探索木を展開していくために、探索順序を学習させ高速化を図るなどがあります。また、ハイパパラメータのチューニングなどは一切行っていないため、ハイパパラメータのチューニングなどで更に精度が向上すると考えています。

今回、実験に使ったコードはこちらにまとめられています。もし、興味がありましたらご覧ください。

インターンシップ全体を通しての感想ですが、今回取り組むテーマとして選んだこの課題は、自分で興味を持ったテーマを選んだため、結果が出るかどうかがわからなく不安なことも多くありましたが、メンターの酒井さんや劉さん、並びに他の社員の皆さんやインターンの皆さんに助けられながら、なんとかここまで進められました。お世話になった皆様、本当にありがとうございました。

参考文献

  1. Hans L. Bodlaender, Arie M. C. A. Koster, Combinatorial Optimization on Graphs of Bounded Treewidth, The Computer Journal, Volume 51, Issue 3, May 2008, pp. 255–269.
  2. Paul Erdős and András Hajnal, On chromatic number of graphs and set-systems,  Acta Mathematica Hungarica, 17.1-2, 1966.
  3. Manuel Sorge,  The graph parameter hierarchy, 2013.
  4. Ziwei Zhang, Peng Cui and Wenwu Zhu, Deep learning on graphs: A survey, arXiv preprint, 2018.
  5. Thomas N. Kipf and Max Welling, Semi-supervised classification with graph convolutional networks, ICLR 2017.
  6. Keyulu Xu, Weihua Hu, Jure Leskovec and Stefanie Jegelka, How powerful are graph neural networks?, ICLR 2019.
  7. Holger Dell, Thore Husfeldt, Bart M. P. Jansen, Petteri Kaski, Christian Komusiewicz and Frances A. Rosamond, The First Parameterized Algorithms and Computational Experiments Challenge, IPEC 2016.
  8. Holger Dell, Christian Komusiewicz, Nimrod Talmon and Mathias Weller, The PACE 2017 Parameterized Algorithms and Computational Experiments Challenge: The Second Iteration, IPEC 2017.
  9. PACE-challenge/Treewidth: List of Treewidth solvers, instances, and tools, https://github.com/PACE-challenge/Treewidth
  10. Tamaki, Hisao. Positive-instance driven dynamic programming for treewidth, ESA 2017.
    著者実装 https://github.com/TCS-Meiji/PACE2017-TrackA
  11. Hans L. Bodlaender, Fedor V. Fomin, Arie M. C. A. Koster, Dieter Kratsch, Dimitrios M. Thiliko. On exact algorithms for treewidth, ACM Transactions on Algorithms, Vol 9, 1, Article 12, 2012.
  12. 河原林 健一, グラフの木分解と木幅に関する研究とグラフマイナー, 数学, 2015, 67 巻, 4 号, p. 337-355, 2017.
  13. 木分解とグラフ・アルゴリズム (最適化と数え上げ) – IBISML 2008,  岡本 吉央
  14. Woeginger, Gerhard J. “Exact algorithms for NP-hard problems: A survey.” Combinatorial optimization—eureka, you shrink!. Springer, Berlin, Heidelberg, 2003.

メンターからのコメント

メンターを担当したPFNの酒井と劉です。

当初こちら側では、組合せ最適化問題に機械学習・深層学習を応用するようなテーマを考えていたのですが、中野さんから「直接計算するのが難しいグラフの不変量をGNNを使って予測できないか」という提案を受け、劉のグラフアルゴリズムに関する知識や興味もあって、トントン拍子に決まったのが今回のテーマです。

最初は、「木幅は大域的な性質なので、GNNによる予測はあまりうまく行かないのではないか」と予想していて、そういう結果が出たときに、それをどう改善したり問題を単純化すれば解けるのかを考えていくことになると想定していたのですが、中野さんに一通りの実装と最初の実験をしてもらってみると、想像以上の精度で予測することができ、メンター二人もとても驚きました。

インターンシップの限られた期間では、より本格的なアルゴリズムへの組み込みや、現実的な応用方法については踏み込むことは出来ず、それらは今後の課題になりますが、今回の結果だけでも非常に面白い結果だと考えています。

微分可能なシミュレータ上での方策最適化

Kengo Nakasuka

2019-10-08 14:40:58

本記事は、2019年インターンシップに参加された金川さんによる寄稿です。


こんにちは。PFNの2019夏季インターンシップに参加した、東京大学の金川です。大学では、強化学習を使って賢くロバストなゲームAIを作る研究をしています。

今回のインターンシップでは、「微分可能なシミュレータ上での方策最適化」というテーマで研究を行いました。この記事では、その内容と結果について、簡単に紹介させていただきたいと思います。

なぜ微分可能なシミュレータか

現実世界では、時間が途切れることなく流れているので、次々と行動を決定していかなくてはいけません。そのため現実世界でロボットを動かしたい時にとれる選択肢の一つは、状態Sを観測した時、あらかじめ学習しておいた方策\(\pi(\cdot|S)\)から行動aをサンプリングし、それにしたがって行動させることです。方策が例えばニューラルネットで表現されていれば、各タイムステップで行動計画を最適化する方法と比べ高速に行動を計算できるため、リアルタイム性が求められるドメインでは大きなメリットがあります。

方策を訓練する方法として、強化学習が注目され、さかんに研究されています。強化学習とは、エージェントがある行動をとると次の状態と報酬が返ってくる、というブラックボックスな環境で、報酬和の期待値を最大化するように方策を訓練する枠組みです。強化学習の課題として、しばしばそのサンプル効率の悪さが挙げられます。例えば、昨年のリサーチブログ で紹介されたEmergence of Locomotion Behaviours in Rich Environmentsでは、シミュレータ上のロボットに歩行動作を獲得させるため、1千万回程度ネットワークのパラメタを更新する必要がありました。

しかし、もしブラックボックスな環境を仮定しないならば、シミュレータの内部情報を上手に使って、学習を効率化できるのではないでしょうか。方策としてニューラルネットを使う場合、報酬をネットワークのパラメタについて微分した値が得られれば、それを利用して直接パラメタを更新できそうです。

手法の説明

そこで今回は微分可能なシミュレータを作成し、勾配という形でシミュレータの内部情報をエージェントに渡すことで、効率よく方策を訓練する手法を試みました。具体的に、シミュレータ内部での遷移関数・報酬関数の計算を、微分可能な演算のみに限定します。この演算過程をchainerの提供するdefine by runスタイルの自動微分を使って記述することで、状態遷移・報酬関数をパラメタで微分できるようになります。

すると、方策と未来の状態・報酬が微分可能な計算グラフでつながり、強化学習における最大化ターゲットである報酬和を、直接方策のパラメタについて微分できます。こうして得られた勾配\(\frac{\partial\sum_{t}\gamma^t R_t}{\partial \theta}\)を使って、方策を勾配上昇法により直接最適化します。

実際のアルゴリズムでは、適当な定数Nをとって、Nステップ分の割引報酬和\(\sum_{t=t_0}^{t_0 + N} \gamma^{t-t_o} R_t\) を最大化します。また学習高速化のため、A2C スタイルの環境並列化およびバッチ学習を行いました。これはシミュレータ内部の物理パラメタをベクトルとして持つだけで、簡単かつ効率的に実装できます。行動についてはすべて連続値をとるものとし、決定的に行動を出力する方策と、行動の正規分布を出力する確率的な方策の2種類を試しました。正規分布からのサンプリングは、再パラメータ化トリック を使うことで逆伝播可能な演算として表現できます。

実験結果

今回のインターンシップではOpen AI gymのClassical Controlタスクを中心に、いくつかの環境をchainerを用いて微分可能なシミュレーションとして実装しました。その上でシミュレータの勾配を使った方策の学習を行って、以下の3点を検証します。

  1. この手法がうまくいくような問題があるか
  2. 決定的な方策を使うものと確率的な方策を使うものでは、どちらの性能が良いか
  3. この手法は、どのような問題に対してうまくいかないか

まず、特に簡単な問題であるCartPole での結果を示します。

これは、台車に力を加え、上に乗っている棒を安定させる問題です。棒が一定以上倒れるか台車が画面の外に出ると、ゲームオーバーになります。

オリジナルの問題での行動は、 「-1or1の力を加える」という離散的でなものですが、今回の設定では[-1, 1]の連続値をとるように変更しました。報酬についてはオリジナルの問題ではゲームオーバーなら0、それ以外なら1が貰えますが、今回は微分可能になるように改造した \(1.0 -\frac{2|\theta |}{\pi} \) を与えました。実装した他の環境に対しても同様に、なるべく元の問題と変わらないよう行動を連続値に、報酬を微分可能にする変更を行っています。

上図に結果を示します(シミュレータのステップ数で性能を比較していますが、同じステップ数の場合ネットワークの更新回数は提案手法の方が20倍程度少ないです。)

どちらの方策を使う場合も、ベースラインとして使用した深層強化学習の手法(TD3) と比べ、20倍以上早く収束していることがわかります。

次にFlappyBirdでの実験結果を示します。これは元々iOS向けに公開されていたゲームですが、今回はOSSとして公開されているPythonクローン をもとに、chainerにより再実装しました。

ゲームの内容としては、鳥に上方向の力を加え土管にぶつからないように飛行させる、という単純なものです。オリジナルの問題では画面をタップするとかなり強い力が加わりますが、今回は連続値の力を加えるため、やや簡単になっています。しかし内部状態(=位置、速さ、角度など)だけが与えられた先程のCartPoleと比べると、この問題では土管の位置など多様な観測情報をもとに行動決定を行う必要があるため、学習は難しくなりそうです。下図に、実験の結果を示します。

実際に行動を観察(上図左)してみると、スタート直後から急激に上昇し、画面外で土管に激突していることがわかります。

この原因について考えるため、この手法での最大化ターゲット\(\sum_{t=t_0}^{t_0 + N} \gamma^{t-t_o} R_t\) を再訪してみましょう。方策を変えれば将来貰える報酬は変化しますから、この式は現在の方策のもとでの報酬の期待値\(\sum_{t=t_0}^{t_0 + N} \mathbb{E}_\pi  \left[ \gamma^{t-t_o} R_t\right]\) (…式1) を最大化しています。

そのため方策がランダム性を持っていないとき、局所最適解に陥りやすくなってしまうのです。強化学習では、局所最適に陥らないよう広範な行動データを収集する必要があり、そのための行動を探査(exploration)といいます。

一方で確率的な方策を使うものは、やや不安定ではあるもののベースラインを大きく上回る性能を残しています。この問題では探査に成功しているようです。(上図右)に示すように、飛行動作も安定しています。

以上の2つの実験から、この手法が単純な状態遷移のダイナミクスを持つ問題に対しては非常にうまくいくこと、探査が必要な問題では確率的な方策の方が性能がよいことがわかりました。

しかしもっと複雑な、カオス的な挙動をする系では、少しの方策の違いが将来の報酬を大きく変動させます。そのため式(1)の分散は非常に大きくなり、学習は難しくなりそうです。

カオス的な問題として、CartPoleSwingUp(下図)での実験を行いました。これは先程のCartPoleとほとんど同じ問題ですが、通常のCartPoleで初期位置がθ=0の近傍から始まるのに対し、θの値が[-π, π]からランダムに与えられます。

この変更はダイナミクスに強い非線形性をもたらすだけでなく、探査も難しくします。

報酬については原論文に似せた \(\frac{\cos (\theta) + 1)}{2} \)を使用しました。下図に結果を示します。

どちらの方策を使う場合も、ベースラインより大きく劣るという結果になりました。

次に探査に失敗する例として、MountainCar(下図)での結果を紹介します。この問題では右上の旗にたどり着かないと正の報酬が貰えませんが、そのためにはいったん左の坂を登って勢いをつけないといけません。

下図に結果を示します。

これも、ベースラインよりはるかに劣るという結果になりました。

この問題では車を動かすため使用した燃料の量に応じて負の報酬が入りますが、それに過剰に反応し、なるべく力をかけないような方策を学習してしまいます。

以上の2つの実験から、この手法がカオス的なダイナミクスを持つ問題や、局所最適に陥りやすい問題に対してはあまり性能が出ないことがわかりました。

まとめ

微分可能な物理シミュレータは新しい研究分野です。決定的な方策を最適化するもの物理パラメタの推定に用いるもの などが提案されていますが、いまだにロボット制御など複雑なシミュレータで検証を行った研究は、我々の知識ではありません。そのため今回のインターンでトイモデルではあるものの、様々な問題で微分可能シミュレータ上での方策最適化の性質を調べられたことは、この手法をより複雑なドメインへ適用するために大きくプラスになったと考えています。

今回の実験では様々な課題が見つかり、単にシミュレータを微分可能にすれば何でもうまくいくという訳ではない、ということがわかりました。しかし簡単な問題では性能が出ることを確認できましたし、またシミュレータ自身が訓練可能になるなど、今回は検証できなかった多くの興味深い性質を持っています。今後も様々な側面から、微分可能シミュレータの可能性を探っていければと思います。

謝辞

慣れないテーマで困惑することや大変に感じることもありましたが、メンターのお二方、インターン同期の方々、同じチームの方々と相談しながら楽しく作業を進めることができました。ありがとうございました。

おまけ:メンターより

金川さんのメンターを担当した、PFNの中須賀と今城です。

PFNでは微分可能シミュレーションの研究開発を行っています。微分可能シミュレーションによって初期値同定など様々な問題が解けることがわかりつつありますが、ただオプティマイザーに投げるだけでは十分な探索が行われないがゆえの問題があり、これを強化学習のアイディアをうまく取り入れて解決できないかということで始まったプロジェクトでした。金川さんには短い期間で様々なアプローチで強化学習のアイディアを取り入れ検証し微分可能シミュレータの可能性を見出していただきました。

深層学習をただ適用するだけではなく問題の性質を考えた学習方法を提案できるというのがPFNの強みの一つであり、問題の性質を理解したり新たな学習方法を理解し現実で解かれていない難しい問題を解くために、このように基礎研究に近いようなアイディアから取り組むということも行っています。今回のプロジェクトはそのような取り組みの一環でもあります。興味のある学生の皆さんは、ぜひ来年のPFNインターンシップへの応募をご検討ください。 また、もちろん中途・新卒の人材募集も通年で行っていますので、こちらも興味のある方はぜひご検討ください!PFNの人材募集のページはこちらです。

 

Disentangled な表現の教師なし学習手法の検証

Keisuke Nakata

2019-10-08 11:20:57

本記事は、2019年インターンシップに参加された蕭喬仁さんによる寄稿です。


はじめまして。PFN の2019夏季インターンシップに参加した東京大学の蕭喬仁です。 大学では自然言語処理について研究しており、SNS からのマイニングに興味があります。
今回のインターンでは「Disentangled な表現の教師なし学習手法の検証 (Unsupervised Disentangled Representation Learning)」というテーマで研究を行いましたので、その紹介をいたします。

実験に使用したコードはこちら https://github.com/pfnet-research/chainer-disentanglement-lib で公開しています。

Disentangledな表現

映画 Star Wars がお好きな方は ”imperial entanglements” という表現でおなじみかもしれませんが、entangle とはもつれるという意味の英単語です。したがって Disentangled Representation とは直訳するともつれを解いた表現ということを指します。
では、もつれを解いた表現とは何なのでしょうか?実は disentangled な表現の定義は研究者の間でもまだ定まったものが無いのですが、多くの研究では潜在空間中の各次元が観測データ中の因子や性状ごとに分かれているような状態を disentangled な表現としています。たとえば、画像認識における disentangled な表現の各次元は被写体の「色」「形」「大きさ」などをそれぞれ表すことが期待されます。このような性質を持つ表現はある1つの次元を変動させても観測データ中の複数の要素が同時に変わる事が無いため、観測データの情報が解釈可能で低次元な潜在空間に圧縮された表現とも言えます。そのため教師有り・半教師有り学習などの機械学習タスクや few-shot learning, domain adaptation などに有用であるとされており、様々な手法が近年提案されてきました。

先行研究

教師なしで disentangled な表現を得る手法として state-of-the-art を主張している論文の多くは変分オートエンコーダ (Variational AutoEncoder; VAE) [1] をベースにした手法を採用しています。

図1. VAE の概念図

VAE では潜在変数の確率分布を標準正規分布に近づけながら学習を行いますが、disentangled な表現を得るためには潜在変数の事後分布がより標準正規分布に近い必要があります。データ \(X\) を生成している真の因子は各々独立であることを仮定しているからです。このような仮定のもと近年提案されてきた手法は VAE にどのような正則化項を加えるかで以下のように分類することができます。

  • 近似事後分布 \(q(z \mid x) \) をより事前分布である標準正規分布に近づける正則化項を加えた β-VAE [2]
  • aggregated variational posterior \(q(z) \) が各成分独立になるような正則化項を加えた FactorVAE [3] や β-TCVAE [4]
  • aggregated variational posterior \(q(z) \) が事前分布と近くなるような正則化項を加えた DIPVAE-1/DIPVAE-2 [5]

また、潜在変数に離散変数を使えるようにした、JointVAE [6] や CascadeVAE [7] などの手法もあります。

様々なモデルの提案と同時に disentangled な表現を定量的に評価する指標も数多く提案されています。評価指標には以下のような2系統が存在しますが、いずれの指標でも値の計算のためにデータの真の生成因子が必要となります。こちらでは紹介していませんが、インターンに参加していた8月中にも新しい評価指標 [8] が提案されており、どの指標を用いるべきかは研究者の間でも定まっていません。

  • データ \(X \) のある生成因子を固定した状態でその他の因子を変化させた時の潜在変数の変化をみる BetaVAE metric [2], FactorVAE metric [3], IRS [9]
  • エンコードした潜在変数から真の生成因子をどれだけ予測可能かをみる MIG [4], SAP-SCORE [5], DCI [10]

一方で、disentangled な表現を完全に教師無しで学習することは困難なのではないかという主張が最近なされるようになりました。Locattello らの研究 [11] では近年 state-of-the-art を主張してきた手法に対して大規模な検証実験を行い以下のようなことを示しています。

  1. モデルや正則化項等のハイパーパラメータよりも乱数による影響が大きいということ。
  2. モデルが学習した disentangled な表現が人間の直感に合う保証が無いこと。
  3. 再構成誤差や ELBO などの教師無しで得られるモデルの評価指標と教師有りの disentangled な表現の評価指標には相関が見られないこと。
  4. データ \(X \) を生成する確率 \(P(X) \) しか持っていない状況では、disentangled な潜在変数 \(p(z) \) と entangled な潜在変数 \(p(z^\prime) \) の識別が数学的に不可能であること。

ただし、論文では適切な帰納バイアスを取り入れることで教師ラベルを用いる事なく disentangled な表現を得ることは可能かもしれないとされており、別の論文 [12] でも損失関数に潜在変数の事前分布を適切に調整することでより良い disentangled な表現を得ることが可能であったという報告もされています。disentangled な表現を得る試みはまだまだ始まったばかりと言えるでしょう。

実験設定

今回のインターンでは、先行研究では検証されていなかった潜在変数の次元数や種類がパフォーマンスにどのような影響を与えるかを検証するために state-of-the-art を主張してきた BetaVAE, FactorVAE, DIPVAE-1/-2 に加えて離散変数を潜在変数として扱える JointVAE の Chainer 実装を行い、2つのデータセットを用いて実験をしました。モデル間で公平な比較を行うために、encoder/decoder の構造や optimizer・バッチサイズ・iteration の数などは共通にした上で、モデルごとに30個のシードで学習を実施しました。

本記事では代表的な実験結果として、データセットのうちの dSprites データに関する結果を紹介したいと思います。このデータセットにはハートや楕円などの物体がサイズや位置を変えて配置された2次元画像が納められており、データを生成する真の因子には物体の種類・物体の大きさ・物体の回転・x軸座標・y軸座標の5種類があります。

図2. データセットの例

実験設定1

上で紹介したモデルを使用する時にまず気になるのが潜在因子の次元数をどのように設定するべきかだと思います。そこで、今回は連続変数を潜在変数とする BetaVAE, FactorVAE, DIPVAE-1/-2 に関して、潜在変数の次元数を真の生成因子の数の半分・同じ・倍の3パターンのモデルを用意して学習結果を比較しました。

実験設定2

実務で使用する際には、dSprites データのように因子の全組み合わせに対応する網羅的なデータを用意することは現実的ではないことが考えられます。そこで、x座標と物体の大きさに相関ができるように因子の組み合わせを半分削除したデータと、比較対象として因子の組み合わせをランダムに半分削除したデータを dSprites データセットから人工的に作成し、パフォーマンスがどの程度落ちるのかを検証しました。

結果

実験1

以下に、潜在変数のある次元だけを動かして、その他の次元を固定した時にデータがどのように変化するかを見た latent traversal という図を掲載します。一番左の列は VAE に与えた元のデータを示しており、その右隣の画像は VAE による再構成画像です。その他の列は潜在変数中の各次元を \(-1.5 \) から \(1.5 \) まで動かした時の画像の変化を順番に描画しています。

 

潜在変数を3次元とした場合
BetaVAE
FactorVAE
DIPVAE-1
DIPVAE-2

 

潜在変数を5次元とした場合
BetaVAE
FactorVAE
DIPVAE-1
DIPVAE-2

 

潜在変数を10次元とした場合
BetaVAE
FactorVAE
DIPVAE-1
DIPVAE-2

 

潜在変数の次元数が小さい場合は、各画像の左から2番目、つまり再構成画像を見ると、どのモデルを使っても画像の再構成が上手くいかず、物体の形がぼやけてしまっているのが見て取れます。どうやら次元数を少なくしすぎるのは問題のようです。しかし、BetaVAE は潜在変数を大きくしても再構成があまり上手くいっておらず物体の形がぼやけてしまっています。実はこの問題は FactorVAE を提唱した論文でも述べられており、BetaVAE に観測データ \(X \) と潜在変数 \(Z \) の相互情報量を小さくするような効果があるから起きるとされています。次元数を真の生成因子と同じ5次元にした場合は、FactorVAE はx軸座標とy軸座標、サイズの変化を disentangle できていますが、DIPVAE-1/-2 では entangled された特徴が学習されてしまっています。次元数が10次元の場合も FactorVAE が5つの因子を disentangle できているのに対して、DIPVAE-1/-2 は直感的ではない特徴を学んでしまっているように見えます。次元数を真の生成因子の数よりも多くした場合は全てのモデルで latent traversal 上で動かない次元が出てきていますが、FactorVAE では真の生成因子の学習が上手くいっていることから、実務で使用する際には潜在変数の次元数は出来るだけ余分にとっておいたほうがいいことが推察されます。

実験2

以下では実験1で定性的に性能がよかった FactorVAE による結果を掲載しています。x座標と物体の大きさに相関があるデータで学習したものは、latent traversal 上でもx座標の移動とともに物体の大きさが変化していることがわかります。また、y軸方向の移動とともに物体の形状が変化しているのをみると、その他の因子の学習にも悪影響を及ぼしていてそうです。ランダムな欠損を与えたデータでは、完全なデータと同じレベルとまではいかないものの、ある程度は5つの因子を獲得できています。今回は全データの半分を削除したので、データを削除しすぎた可能性もあります。しかし、実務で使用する際に生成因子の全ての組み合わせのデータの準備が難しい場合は、せめて因子に相関がないようなデータに整えた上で学習を行うべきである事がこの結果から推察できます。

完全なデータ

x座標と物体の大きさに相関があるデータ
ランダムな欠損を与えたデータ

 

メンターからのコメント

メンターを担当した PFN の仲田と吉川です。

本文中でも触れられている通り、disentangled な表現の教師なし学習は few-shot learning や domain adaptation といった分野への応用や、機械学習モデルの解釈性の向上が期待できます。これらは現実世界への機械学習の適用をする際にいつも直面する課題です。
深層学習などの先端技術による現実世界の課題解決をミッションとする PFN としても以前から取り組んではいるのですが、タスク自体の難しさはもちろん、問題設定をおこなうことも難しい分野です。インターン期間中は手法の再現性や各種文献の実装の読解などにお互い苦労しましたが、最終的には代表的な手法を一通り Chainer で実装し、様々なデータやパラメータ、評価指標を広く調査して頂いた蕭さんの今回の成果は、今後の PFN における研究開発にあたっての足がかりとなることと思います。

参考文献

微分可能なMPCのChainer実装および追加実験

Yo Iida

2019-10-04 14:30:30

本記事は、2019年インターンシップとして勤務した新幡 駿さんによる寄稿です。


こんにちは、2019年夏にPFNでインターンをしていた新幡です。大学では連続最適化の研究室に所属しています。PFNではモデル予測制御 (MPC; Model Predictive Control) を「微分可能」にするAmosら (2018) の論文“Differentiable MPC for End-to-end Planning and Control”[1] について、著者によるPyTorch実装のChainerへの移植と追加実験を行いました。MPCはモデルを用いて将来の状態を予測しながら制御を行う手法であり、化学プラントの制御などに用いられています。

【2019/10/7追記】Chainerへの移植実装はGitHubで公開しています。
https://github.com/pfnet-research/chainer-differentiable-mpc

1. OptNet

Amosらは、2017年に二次計画法を微分可能にする層として“OptNet: differentiable optimization as a layer in neural networks”[2] を提案しました。二次計画法を微分可能にすることで、勾配降下法を用いて二次計画法の目的関数や制約条件を、ネットワークの損失関数に対して最適化することができます。OptNetの順伝播は、二次計画問題を主双対内点法という収束するまで反復を繰り返す方法を用いて解いていますOptNetの誤差逆伝播は、主双対内点法を全て連鎖律で逆にたどるのではなく、最適解の満たすべき条件のうち等式で書けるもの (KKT行列) から効率的に偏微分を求めています。大量のデータを用いて学習する為に、二次計画問題を大量に解く必要があるので、Amosらはバッチで計算できるような主双対内点法を実装しており、著者実装は A fast and differentiable QP solver for PyTorch として公開されています。また、副メンターの酒井さんによるChainerへの移植の試みとして chainer-optnet があります

図1. OptNetの概要図

2. Differentiable MPC

OptNetで二次計画法を微分可能にする手法が提案されたので、有限離散時間線形二次レギュレータ (Linear-Quadratic Regulator; LQR) などの二次計画法を用いて解ける最適制御問題はOptNetを用いて微分可能にできます。しかし、LQR二次計画問題として解くのではなく、動的計画法を用いた方が効率的に解くことができることが知られています。Amosらは、LQRの順伝播の計算に動的計画法を用い、誤差逆伝播の計算もKKT行列の時間的構造を用いてOptNetより効率的に微分を求められるLQRを提案しました。

さらに有限離散時間MPCのようなコスト関数やダイナミクスが非線形な問題に対しても、微分可能にした論文がAmosらによるDifferentiable MPCです。著者による実装はアルゴリズムの主な部分が mpc.pytorch としてライブラリになっており、著者の実験コードは differentiable-mpc といレポジトリにあります。

3. 前提知識

3.1. LQR

LQRは一般的には、コスト関数が二次でダイナミクスの制約が線形、操作量が無制約な次のような形で表される最適制御問題です。

$$
\begin{align*}
&\underset{\tau_{1:T}}{\text{minimize}} & & \sum_t \frac{1}{2} \tau^T_t C_t \tau_t+c^T_t \tau_t \\
&\text{subject to} & & x_1 = x_\text{init}\\
& & & x_{t+1} = F_t \tau_t + f_t
\end{align*}
$$

ここで \(x\) は状態、 \(\tau\) は状態と操作量を連結したもの、 \(C\) はコスト関数の二次の項を表す行列、 \(c\) はコスト関数の1次の項を表すベクトル、\(F\) はダイナミクスを表す行列になります。Amosらは、LQRを動的計画法を用いて、ベルマン方程式を再帰的に解くことで解いています。詳しくは[3]を参照してください。

3.2. MPC

MPCはLQRと異なり、コスト関数は2次に限らず、ダイナミクスも線形に限りません。さらに、操作量に制約がつく場合もあります。元論文で扱ったMPCは以下の式です。

$$
\begin{align*}
&\underset{\tau_{1:T}}{\text{minimize}} && \sum_t C_{\theta, t}(\tau_t) \\
& \text{subject to } && x_1= x_{\mathrm{init}}\\
&&& x_{t+1} = f_{\theta}(\tau_t)\\
&&& \underline{u} \leq u \leq \overline{u}
\end{align*}
$$

Amosらは、コスト関数を二次近似、ダイナミクスを一次近似して、収束するまでLQRソルバを繰り返し用いるiLQRという手法を用いて解いています。また、制約については、Tassaら (2012) による“Control-Limited Differential Dynamic Programming”[4] と同じように扱っています。

4. 応用

Differentiable MPCの応用として、Amosらが行っているのがエキスパートの操作列を模倣する模倣学習 (imitation learning) です。Amosらによる模倣学習は、ダイナミクスやコスト関数をパラメータ化した形で表して、エキスパートの操作列との二乗誤差をネットワークの損失関数としてパラメータを最適化させています。最適化によって得られたダイナミクスやコスト関数をMPCに用いることで、エキスパートの操作列を模倣する方策を得られると期待できます。Amosらは、実験設定として、倒立振子 (inverted pendulum) を題材に、学習するものがダイナミクスのパラメータのみ・コスト関数のパラメータのみ・両方の三通りの学習を試しています。

図2. 模倣学習の設定

例として、ダイナミクスを既知として、コスト関数をパラメータ化した形で学習を行うと次のようになります。エキスパートの操作列から、正しく倒立振子を安定化させることができていることがわかります。

図3. 学習したコスト関数を用いた制御

5. 追加実験

追加実験として、コスト関数をAmosらのように特定の形にパラメータ化した形ではなく一般の二次関数の場合もうまく学習できるのかを試しました。ここでAmosらのパラメータ化を詳しく説明すると、以下の数式のようになっていて、\(q_{\text{logit}}\)と\(p\)という二つのベクトルをパラメータとしています。

$$
\begin{align*}
q &= \mathrm{sigmoid}(q_{\rm{logit}})\\
C(\tau) &= \frac{1}{2}\tau^T q \tau + (\sqrt{q} \circ p)^T \tau
\end{align*}
$$

追加実験として試したパラメータ化は、半正定値行列\(Q\)とベクトル\(p\)をパラメータとして用いてコスト関数をより多くのパラメータで表現しています。

$$
\begin{align*}
C(\tau) = \frac{1}{2}\tau^T Q \tau + p\tau
\end{align*}
$$

5.1. 真のコスト関数がAmosらのパラメータ化で表現できる場合

図4. 真のコスト関数がAmosらのパラメータ化で表現できる時

この場合は、真のコスト関数の構造を取りいれているAmosらによる手法の方が損失関数の値が下がりました。

5.2. 真のコスト関数がAmosらのパラメータ化で表現できない場合

真のコスト関数がAmosらのパラメータ化で表せないように問題設定を変えました。

この場合、一般的な二次関数のパラメータを用いた手法とAmosらの手法を比較すると、損失関数がどの程度まで下がるかは同等でしたAmosらの手法では、真のコスト関数を表現できないので、一般的な二次関数を用いて学習させる方が損失関数の値が下がることを期待していたのですが、これは意外な結果でした。Amosらは、LQRの場合に、ダイナミクスを一般的な行列の形にして模倣学習を行い、その際は半分程度は損失関数がうまく下がらないという現象を確認していますが、これを模倣学習の損失関数が、ダイナミクスのパラメータに対して、勾配降下法では最適化することが難しい非凸な関数であったからと結論づけています。今回、一般的な形のコスト関数でうまく学習できなかったのも、コスト関数を一般的な形にすると、損失関数がコスト関数のパラメータに対して非凸な関数になったからではないかと考えられます。

図5. 真のコスト関数がAmosのパラメータ化で表現できない時

6. 感想

今回のインターンでは、Differentiable MPCのChainerへの移植を行い、色々と実験設定を変えてみました。コスト関数を凸なニューラルネットワークで置き換えられるとおもしろそうだなという思いつきから、まずはコスト関数のパラメータを増やして実験してみました。元論文では模倣学習の損失関数が制御器のパラメータに関して非凸な関数になることがあると指摘されており、追加実験でもパラメータを増やしたところあまり損失関数の値が下がらないという結果になりました。

Differentiable MPCは理論、実装共にかなり複雑で、Chainerにライブラリを移植するのにかなり苦労し、時間がかかってしまいましたが、メンターさん方の助けもあってなんとかできました。ありがとうございました。

参考文献

[1] B. Amos, I. Jimenez, J. Sacks, B. Boots, and Z. Kolter. “Differentiable MPC for End-to-end Planning and Control” In NIPS’18 Proceedings of Advances in Neural Information Processing Systems. Available: https://arxiv.org/abs/1810.13400

[2] Brandon Amos, and J. Zico Kolter, “OptNet: differentiable optimization as a layer in neural networks” In ICML’17 Proceedings of the 34th International Conference on Machine Learning. Available: https://arxiv.org/abs/1703.00443

[3] Sergey Levine.  Optimal control and planning. berkeley cs 294-112:  Deep reinforcement learning. Available: http://rail.eecs.berkeley.edu/deeprlcourse-fa17/f17docs/lecture_8_model_based_planning.pdf

[4] Yuval Tassa, Nicolas Mansard and Emo Todorov. “Control-Limited Differential Dynamic Programming” ICRA’14 IEEE International Conference on Robotics and Automation. Available: https://homes.cs.washington.edu/~todorov/papers/TassaICRA14.pdf


メンターからのコメント

メンターを担当したPFNの飯田と酒井です。

PFNは深層学習に強みがありますが、現実の問題を解決するには深層学習だけでは不十分なこともあり、最適化や制御の技術を組み合わせて用いることも多いです。しかし、昨年までのインターンシップではそういったバックグラウンドの学生の方にはあまりリーチ出来ていませんでした。

そこで、今年のインターンシップ募集でのテーマ「機械学習アルゴリズムの応用 (数理最適化、シミュレーション、時系列予測など) に関する研究開発」では「数理最適化」もキーワードに含め、結果として最適化のバックグラウンドを持つ学生さんとの研究開発を幾つか実施することができました。微分可能MPCに関するこのテーマはその一つです。

今回のテーマは、再現実装の負荷がメンター二人の想定していたよりも重く、新幡さんには随分苦労させてしまいましたが、再現実装と追加実験までを行っていただくことができました。社内の他のプロジェクトでも深層学習とMPCを組み合わせて用いており、今回の新幡さんの成果も今後それらにうまく活用していきたいと考えています。

スポーツ映像に対するシーンのアノテーション効率化

Tatsuya Takamura

2019-10-02 12:03:57

本記事は、2019年インターンシップとして勤務した佐々木 克仁さんによる寄稿です。


はじめまして。PFNの2019年夏季インターンシップに参加させていただいた東京大学修士1年の佐々木克仁です。大学ではHCIの研究をしています。WEB開発が好きです。

テーマとその背景

今回のインターンシップで私が取り組んだ研究テーマは「スポーツ映像に対するシーンのアノテーション効率化」です。

PFNでは、スポーツ映像の中でチームが取っている戦術を推定し、スポーツの戦術解析に応用するシステムを開発しています。このような推定を実現する機械学習モデルを学習するためには、チームが取っている戦術とその時間範囲(以降シーンと呼びます)がスポーツ映像にアノテーションされた大量のデータセットが要求されます。しかし、スポーツ映像におけるシーンの戦術レベルでの詳細な区別を一般の人々が行うのは困難で、そのスポーツに精通した専門家しかアノテーションできず、彼らの限られたリソースを効率的に利用するためにアノテーションの効率化が必要となります。

今回私が着目したのは「シーンの境界線を決めるのに時間がかかっている」という課題でした。アノテーターは先頭からスポーツ映像を見ていって各シーン間の境界線の時間を記録することでアノテーションしていきます。境界線はルールに基づいて決定するのですが、この作業に時間がかかっており、またアノテーターのストレスに繋がっています。インターンシップではこの課題を解決するための取り組みを行いました。

手法

私が取り組んだ手法では、各シーンに対してアノテーターは「シーンの境界線を決める」代わりに「シーンの中の適当な一時点を決める」作業を行います。これによってシーンの境界線を決める労力が削減されるためアノテーションが効率化されると考えられます。以降では「シーンの境界線を決める」アノテーションを「Boundary型」、「シーンの中の適当な一時点を決める」アノテーションを「Timestamp型」と呼びます。Timestamp型アノテーションの考え方は[1]から取り入れました。

annotation-image

図1: アノテーターはシーンの中の適当な一時点を決める

しかし、Timestamp型のアノテーションはそのままデータセットとして利用することはできず、一つのTimestampを元にそのTimestampが示すシーンの境界線を推定する必要があります。そこで今回は、学習させた機械学習モデルの反応を見て、与えられたTimestampを元に暫定的に決めた境界線から徐々に境界線を更新させるという方法を採用しました(図2)。境界線の更新と機械学習モデルの学習を並行して行うことで、相互に作用しながら解に近づいていきます。

iteration

図2: Timestampから境界線を推定する手法の概要

境界線は、機械学習モデルが出力した各クラスの確率分布に基づいて更新されます(図3)。図2は、クラスAとクラスBのシーンの境界線を示しています。現在の境界線(点線)において機械学習モデルが出力したAの確率とBの確率の差Δが、閾値よりも大きい場合に境界線をB方向に更新します。更新の大きさは、Aの確率を示す曲線とBの確率を示す曲線の交点を基準に学習率をかけたものとなります:

eq1

ここで、bは境界線の時刻、cは曲線の交点の時刻、rは学習率です。

boundary-update

図3: 境界線の更新

この境界線の更新手法は、[2]を参考にして改良を加えたものになっています。[1]の方法は、Backgroundクラスが大部分を占めるAction Recognitionのタスクでは有効ですが、スポーツ映像のようなほぼ全てのフレームがBackgroundでないクラスを持つタスクではうまく機能しないと考えられるため採用しませんでした。

実験・評価

境界線の推定精度

まず、境界線を推定する手法が適切に機能するかどうかをシミュレーションによって評価しました。アノテーターによるTimestamp型アノテーションのシミュレーションを、Ground-Truthのデータにおけるそれぞれのシーンの時間範囲からガウス分布に基づいてサンプリングすることで行います。推定された境界線の精度は、Ground-Truthの境界線と推定された境界線のIoU(Intersection over Union)によって評価しました。IoUは2つの領域が重なっている割合を示す値であり、IoUが高いほど境界線の推定精度が高いことを示しています。

以下に結果を示します(図4)。「学習→境界線の更新」の繰り返しを行うたびにIoUの平均値が向上することが示されています。最終的に0.76まで到達しています。今回推定された境界線に基づいたデータセットの有用性の検証は今後の課題となります。

estimated boundary

図4: IoUの平均値(横軸は境界線の更新回数)

Timestamp型アノテーションの効率

次に、Timestamp型アノテーションによって実際にアノテーションの速度が向上することをユーザースタディによって検証しました。ユーザースタディに向けてBoundary型とTimestamp型のWEBインターフェースを実装しました(図5)。被験者としてスポーツ解析の専門家2人にご協力いただきました。被験者は制限時間内に可能な限りのアノテーションを行い、それぞれの方式でアノテーションできる数を比較します。今回は被験者数が十分ではありませんが、初期検討の材料としての実験とし、追加実験は今後の取り組みとします。

compare mode

図5: アノテーションツールの実装(左はBoundary型、右はTimestamp型)
※スポーツ映像の部分はイメージ図です

ユーザースタディの結果、Timestamp型ではBoundary型の約2倍の速度でアノテーションできることが分かりました。アノテーションが速くなった原因としては、当初想定していた境界線を決める時間が削減されたこと以外にも、Timestamp型のほうがツールとして操作が容易であったこともあるようでした。被験者となっていただいた専門家の方からは「直感的」「映像を見ている感覚でできるのでストレスが少ない」という今回の手法に対する肯定的なフィードバックをいただくことができました。一方で、アノテーターによってシーンの中のどのタイミングでアノテーションするかの傾向が異なるという課題も発見されました。シーンの中のどのタイミングでアノテーションするるかは境界線の推定における初期値を決める重要な因子です。ユーザスタディを実際に行ったからこそ判明した重要な発見でした。

研究のまとめ

今回の研究では、スポーツ映像に対するシーンのアノテーションを効率化するための手法として、Timestamp型のアノテーションに取り組みました。Timestampからシーンの境界線を推定する手法では、Ground-Truthと比較してIoUで0.76の精度を実現しました。また、スポーツ解析の専門家2名によるユーザースタディによってTimestamp型のアノテーションでは従来の方式より2倍速くアノテーションできることが分かりました。今後の課題としては、シーンの境界線推定の精度が得られた結果で十分であるかどうか検証する必要があり、またユーザによるアノテーションのタイミングのばらつきに対処する必要があります。

インターンシップの感想・謝辞

今回2ヶ月間のインターンシップに参加させていただき、非常に充実した時間を過ごすことができました。2ヶ月もの間1つの研究について考え続けるのはとても楽しく、かつてない経験をさせていただきました。PFN社員の方々のレベルが非常に高く、せっかくこの環境にいるのだから難しい問題に取り組んでみようと前向きな気持ちでチャレンジすることができました。また、インターンの同期も同世代と思えないほど実力のある人ばかりで自分の至らない部分を知ることができ、多大な刺激を受けました。本当に参加してよかったと思えるインターンシップでした。

最後に、メンターをはじめとしたチームの皆さま、研究のサポートをして頂き誠にありがとうございました。また、ポスターセッションでフィードバックをくださった皆さま、適切な助言をありがとうございました。

参考文献

[1] Moltisanti, Davide, Sanja Fidler, and Dima Damen. “Action Recognition from Single Timestamp Supervision in Untrimmed Videos.” Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019.

[2] Ding, Li, and Chenliang Xu. “Weakly-supervised action segmentation with iterative soft boundary assignment.” Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018.

Chainer Chemistryの大規模グラフのタスクへの拡張

Kosuke Nakago

2019-10-01 13:09:08

本記事は、2019年インターンシップで勤務した 阿部健信 さんによる寄稿です。

こんにちは。2019年夏季インターンに参加した東京大学の阿部健信です。「Chainer Chemistryの大規模グラフのタスクへの拡張」というテーマで取り組んだ内容を説明させていただきます。インターン内容のスライドはこちらにアップロードされています。

PFN Summer Internship 2019 / Kenshin Abe: Extension of Chainer-Chemistry for Large and Sparse Graph from Preferred Networks

 

TLDR;

  • Chainer Chemistryで大規模グラフのデータを扱えるようにしました。
  • convolution演算を\( O(V^2) \)から\( O(E) \)にしました。
  • メモリ使用量も抑えて、PyTorch Geometricでは動かないRedditデータセット(23万頂点, 1100万辺)を16GBのsingle GPU上で学習できるようにしました。

 

はじめに

 

graph convolution [1]

画像に対する2D ConvolutionとGraph Convolutionの比較 [1]

入力としてグラフを受け取ることのできる、Graph Neural Network(GNN)という分野が近年注目を集めています。
その注目の高まりから、PyTorch Geometric [2]やDeep Graph Library (DGL) [3]といった高機能で最適化されたGNNライブラリの開発が盛んに進められています。

Chainer Chemistryは、PFNが開発しているGNNのオープンソースのライブラリです。
名前からも分かるとおり、もともと分子など化学データへの適用を目的として作られたもので、qm9などの化学データセットが手厚くサポートされています。一方、他にGNNの研究でよく用いられるSNSなどのネットワークデータのサポートはなされていませんでした。

 

課題内容

今回のインターンのタスクは、Chainer Chemistryでネットワークデータのサポートを行うことです。そのために、大きく以下の2つの内容を行いました。

1. node classificationのサポート
化学分子データなどたくさんの小さなグラフのデータセットに対してGNNは、graph classification/regressionといった、グラフ全体の性質を学習するのに用いられます。
一方、巨大なネットワークデータに対しては、1つのグラフを入力として各頂点ラベルの分類を行うといった異なるタスクに用いられる事が多いです。
[4]で提案されているようなsemi-supervised node classificationへの対応を行いました。
具体的なフレームワークの違いはスライドをご参照ください。

2. 巨大でsparseなグラフのためのGNNの効率的な実装
こちらが今回のインターン内容のメインで、巨大なグラフを動かすためには必要不可欠な内容でした。
以下、\( V \) 個の頂点、\( E \)個の辺からなるグラフを考えます。
Message passingにもとづくGNNでは、各頂点に対して近傍の頂点の特徴量のaggregationの操作を行います。このaggregationの関数はpermutation invariantな様々な関数が用いられ、例えばよく使われるsumの場合は以下の式になります。
\( H’ = AH \)
(\( H \): 頂点の特徴量行列, \( A \): 隣接行列, \( H’ \): aggregateされた特徴量)

既存の実装は全てこの行列演算に基づくものでしたが、これは2つ問題点があります。
1つめは、グラフが疎な際にメモリ的にも実行時間的にも無駄が生じてしまうことです。
2つめは、batch化の際のゼロパディングのオーバーヘッドです。

これらの問題を解決するために、辺の情報を密な隣接行列ではなく、疎なデータ形式で持たせるという事が考えられます。今回のインターンでは、こちらのレポジトリでsparse patternとして紹介されているデータの持ち方を新たに実装しました。
これは辺の情報を\( [2, E] \)のサイズの行列で持つ手法で、PyTorch Geometricでも採用されています。

Sparse patternでは、scatter演算と呼ばれる命令を用いることでaggregation部分の計算量を\( O(E) \)で行うことができます。
またbatch化の際に、複数のグラフを全体として大きな1つのグラフとしてみなすことによってゼロパディングのオーバーヘッド完全になくすことができます。
こちらも、より詳細な手法が知りたい方はスライドをご覧ください。

 

結果

行列演算による既存実装と、sparse patternによる実装の速度比較は以下のようになりました。
まず、3312頂点、4660辺の疎なネットワークグラフに対しては、CPUでは50倍以上、行列演算との相性が良いGPU上でも2倍以上の速度改善が見られました。
また、1つ予想外だったのは、最大でも38頂点という比較的小さなグラフからなる化学データセットに対してもGPU上でも1.5倍程度の速度改善が見られたことです。
これには、バッチ化のオーバーヘッドをなくす工夫が効いていると考えられます。

sparse patternはグラフのconvolution演算に特化して実装されているため速いもののメモリ使用量にまだ無駄があり、Redditデータセット(23万頂点, 1100万辺)を動かすことはできませんでした。
これについては、ChainerのサポートしているCooMatrix演算によるモデルを用いたところsingle GPU (16GB)で動かすことができました。

これまで触れた、既存の隣接行列・sparse pattern・CooMatrixの3パターンについてまとめると、グラフが疎であったりバッチ化のオーバーヘッドが大きかったりすれば基本的にsparse patternが早く、それではメモリが足りない場合はCooMatrixを使うとよい、という結果になりました。
この結果を踏まえて、3つのパターンを場合に応じて使い分けられるように実装しています。
特に、現状のPyTorch Geometricでは動かすことができないredditなどの超巨大なグラフを動かせるという点は、Chainer Chemistryを使うモチベーションの1つになると思います。
新しいモデルを自分で実装したいときに、ChainerでサポートされているCooMatrix演算を普通の行列演算と同じようなインターフェースで直感的に使えるのも魅力です。

 

 

まとめ

今回の成果はChainer Chemistryにマージされています。新しい実装方針に対応しているモデルはまだ多くはありませんが、これからどんどん対応していく予定です。
exampleのコードを動かすことで簡単に巨大グラフ上での学習ができるようになっているので、ぜひ試してみてください。

 

参考資料

[1] https://arxiv.org/pdf/1901.00596.pdf
[2] https://rlgm.github.io/papers/2.pdf
[3] https://rlgm.github.io/papers/49.pdf
[4] https://arxiv.org/pdf/1609.02907.pdf

 

CuPy カーネル融合の拡張

Daisuke Nishino

2019-09-27 14:44:40

本記事は、2019年インターンシップとして勤務した徐 子健さんによる寄稿です。


2019年度夏季インターンのJoeです。この度インターンプロジェクトとしてCuPyのカーネル融合の拡張に取り組み、既存のカーネル融合の適用範囲を大幅に拡張しました。さらにその応用として、ResNet50のバッチ正規化においてCPU実行時間を30%ほど、GPU実行時間を(入力サイズに大きく依存しますがおおよそ)70%ほど削減することに成功しましたので、その取り組みをご紹介します。

背景

CuPyはNumPyと同じAPIを提供するPythonのライブラリで、CUDAを用いて演算を高速に行います。具体的には、行列・ベクトルの要素ごとの演算や、リダクションと呼ばれる、演算によって配列の次元が落ちる演算(たとえばcupy.sum)など、GPUが得意とする計算を高速に行うことができます。

さて、CuPyのGPU演算は強力ですが、CPUからGPUカーネルを呼び出す際には、GPU内で実際に行われる処理によらずおおよそ10 μsほどのオーバーヘッドがCPU側にかかってしまいます。cupy.addcupy.mulのような演算が呼ばれるたびにGPUカーネルを呼び出していては、無視できない量のオーバーヘッドが発生してしまいます。

そこで、いくつかの演算を融合したCUDAコードを実行時に動的に生成することにより、GPUカーネルの呼び出しを減らすことによる高速化が重要となります。この技術をカーネル融合といいます。詳しくは後述しますが、要素ごとの演算を組み合わせた関数やリダクション演算を一度まで含む関数に対するカーネル融合は現在のCuPyにもすでに実装されています。

https://docs-cupy.chainer.org/en/v6.4.0/reference/generated/cupy.fuse.html

カーネル融合で行っている処理を具体的に説明するために、以下のPythonで書かれた関数に配列(cupy.ndarray)を引数に与えて実行したときに、カーネル融合をしない場合とする場合に生成されるCUDAコードを比較してみます。


def f(x):
    return x + x * x

カーネル融合をしない場合

このとき掛け算と足し算の両方の呼び出しでCUDAコードが生成され、2度のカーネル呼び出しが発生します。

具体的には以下のようなコードが生成され、関数single_op_mulsingle_op_addがそれぞれ実行されます。


__device__ void mul_element(int &src1, int &src2, int &dst) {
    dst = src1 * src2;
}

__global__ void single_op_mul(Array& in, Array& out) {
    int a = in[tid], b;
    mul_element(a, a, b);
    out[tid] = b;
}

__device__ void add_element(int &src1, int &src2, int &dst) {
    dst = src1 + src2;
}

__global__ void single_op_add(Array& in1, Array& in2, Array& out) {
    int a = in1[tid], b = in2[tid], c;
    add_element(a, b, c);
    out[tid] = c;
}

カーネル融合をする場合

この場合、2つの演算は1つの関数にまとめられて、以下のように生成される関数fusedが呼ばれ、カーネル呼び出しは1回で済みます。


__device__ void add_element(int &src1, int &src2, int &dst);

__device __ void mul_element(int &src1, int &src2, int &dst);

__global__ void fused(Array& in, Array& out) {
    int a = x[tid], b, c;
    mul_element(a, a, b);
    add_element(a, b, c);
    out[tid] = c;
}

このとき、カーネル呼び出しの回数が1回に減るのでCPU時間が短縮されます。さらに、x[tid]部分はグローバルな配列にアクセスしていて、GPU側のボトルネックになりがちですが、処理をまとめることで、こうしたグローバル領域のアクセスの回数を減らすことができ、GPU側の高速化にも繋がります。

要素ごとの演算とリダクション演算のカーネル融合

既存のカーネル融合は要素ごとの演算だけではなく、一度までリダクションを含む関数ならば融合することができます。上述した例に示したような要素ごとの演算を組み合わせた関数に対するカーネル融合は比較的イメージしやすく、生成するCUDAコードもcupy.addcupy.mulに相当する演算をつなげるだけでいいので簡単ですが、要素ごとの演算とリダクション演算を融合しようとすると話は少しややこしくなります。要素ごとの演算のカーネル融合ではGPUの各スレッドが特定の位置の要素を独立に扱っていましたが、リダクション演算では、各スレッドは適切に同期されながらシェアードメモリに読み書きを行い、演算を処理します。

上の図は、既存のカーネル融合で、要素ごとの演算とリダクション演算からなる関数を融合した場合にGPUで行われる処理を表したものです。ここで、リダクション前の演算はすべてリダクションの入力と同じサイズ(正しくは、cupy.ndarrayshapeのこと。以下同様)をもち、 リダクション後の演算はすべてリダクションの出力と同じサイズをもちます。逆にこのような制約のもとで表せない関数には既存のカーネル融合を適用できません。

このような、リダクションを一度行う関数を融合することは、先程説明したような要素ごとの演算のみからなる関数を融合する場合よりは難しいですが、それでも元のリダクション演算の前後に、要素ごとの演算を織り込むだけで実現することができます。

課題

今回のインターンの課題は、カーネル融合で扱える関数の範囲を広げ、要素ごとの演算とリダクションからなる関数ならなんでも融合できるようにする、というものでした。実際既存のカーネル融合において、リダクションが一度までというのは大きな制約であり、例えば平均や分散の計算を伴うバッチ正規化の関数を融合することはできませんでした。バッチ正規化はResNet50などで何度も実行される関数であり、特にカーネル融合の効果が期待される処理でした。

取り組んだこと

リダクションを二度以上含む関数の融合は、一見リダクションを一度だけ含む関数の融合方法を簡単に拡張するだけで実現できるように思えますが、このニ者の間には大きな違いがあります。というのも、複数回リダクションが行われるとき、一度前のリダクションの結果を一時的にGPUのグローバル領域に格納する必要が生まれます。同じサイズの要素ごとの演算のみだと、各スレッドは配列の特定の要素を見続けるので、演算結果の変数をスレッドのローカル変数に保持しておくことができます(リダクション演算が一度だけの場合、出力用配列にそのまま保存するので問題ありません)。しかし、配列のサイズが演算間で変化するとき、スレッドが担当する要素がずれるため、一度結果をグローバル(スレッド間で共通してアクセスできる位置)に保存してやる必要があるのです。具体的には新たな空の配列をグローバル領域に確保しておき、それを使います。

カーネル融合の処理の流れ

具体的なアルゴリズムを説明する前に、カーネル融合の処理の流れについて整理しておきます。

各々の部分の詳細な説明は後述しますが、カーネル融合の処理は大まかに、解析パートと実行時パートに別れます。解析パートでは、CUDAコードを生成するために必要な情報を解析し、CUDAコードの最適化を行います。解析パートは同様な引数の組み合わせに対しては一度だけ実行されることを想定しています。一方、実行時パートは、毎回の関数呼び出しで実行されるパートのことを指します。ここでは、配列の次元圧縮とよばれる実行時最適化を行い、解析パートで得られた情報から実際にCUDAコードを生成し、コンパイルしてカーネルを呼び出します。何度も呼び出されることが想定される関数の場合、関数呼び出しあたりの平均的なCPUの処理時間は、この実行時パートの処理時間にかかっています。それでは、解析パートおよび実行時パートで行っていることを順に説明していきます。

解析パート

サイズの抽象化と制約生成

多段階のリダクションを実現する上で重要なことは、サイズが異なる演算が発生するたびに一時変数をグローバル領域に保存してやることでした。そうした一時変数のメモリを割り当ててやるためには、カーネル呼び出しの直前にサイズを決定してやる必要があります。なぜならば、同一の関数が異なるサイズの入力で繰り返し呼ばれる可能性があるからです。サイズが異なる入力が来るたびにPython関数を解析し直すのは非常に高コストなので、入力の配列のサイズが変わった場合にも同じコードを使えるように変数のサイズ情報を抽象化してを保存したいです。

サイズを抽象化して保存するためには、まず関数の引数で渡される配列のサイズを抽象化して保持します。たとえば、三次元配列xのサイズを(x0, x1, x2)のようにして保持します。これらの抽象化したサイズをもとに、演算を順に追っていき、一時変数の抽象的なサイズを決定します。こうした演算を追う過程で、演算がただしく成立するために必要な、サイズ間の制約が発生します。たとえば、(x0, x1, x2)のサイズの配列と、(y0, y1)のサイズの配列の要素ごとの演算を行うには、後者をブロードキャストする必要があり、このとき、x1, x2, y0およびy1がいずれも1でないとすると、y0 == x1かつy1 == x2がサイズ間の制約となります。これらの制約は関数の解析時に順次保存され、2度目以降の関数呼び出しの際に、与えられた引数のサイズが保存された制約を満たすならば、以前に解析した結果およびコードを使い回すことができます。

CUDAコードの生成

CUDAコードを生成する上での基本的なアイデアは、本稿の初めで説明した、要素ごとの演算のみからなる関数のカーネル融合とほとんど同じです。すなわち、関数内に登場するそれぞれの演算を呼び出し順に記録し、それをつなげて一つのCUDAコードにします。リダクション演算が続いたり、サイズが異なる要素ごとの演算が連続したりするときは、スレッド間で同期をとる必要が生まれますが、基本的にはそれぞれの演算を順につなげたコードを生成すればよいです。

CUDAコードの最適化

さて、ここまでの説明だけだと、カーネル融合が簡単なものに思えるかもしれません。実際、GPUの実行時間を考慮しないならば、演算ごとにサブルーチンを生成し、そのサブルーチンをすべて順につなげればひとつのCUDAコードが完成します。しかし、このままだと本当にGPUを呼び出すための10 μs程度のCPU側の高速化しかなされません。カーネル融合ではCPUのオーバーヘッドをへらすだけではなく、GPUコードの無駄な処理を省くことでGPU側の高速化をも期待することができます。

さて、各演算を順に結合しただけのCUDAコードは無駄な処理を多く含みます。たとえば直前の演算で結果をグローバルな配列に格納し、次の演算でまた同じ要素をグローバルな配列から読み出すのは明らかに無駄です。こうしたケースでは演算結果をスレッドごとのローカル変数に記録してやると大きな高速化に繋がります。

たとえば、愚直に生成したCUDAコード内で以下のような部分があったとします。


tmp[tid] = a;

// ここまで一つの演算に対応。変数aを出力として一時変数に格納
// ここから新しい演算が開始。変数aを新しい入力として一時変数から読み出す

int b = tmp[tid];

このとき、配列tmpが一時変数で、しかもここでしかアクセスされない変数であるならば、


int b = a;

と処理を短縮してやることができます。GPU演算で各スレッドが担当する演算は非常に軽量であることが多いため、このようにグローバル配列へのメモリアクセスを削減してやるだけでも大きな高速化に繋がります。

今回、要素ごとの演算間のみならず、リダクション演算が混じる場合にもこうしたグローバルなメモリアクセスを減らす最適化を行いました。他にもスレッド同期の削減など非常に多様な最適化を試行錯誤し、結果として、従来のカーネル融合を適用できる演算については、ほとんど劣ることがないレベルまでに最適化がなされています。また、高速化とは直接関係ないですが、一時変数の生存解析をすることにより、メモリを共有できる変数はメモリを共有することで、実行時の消費メモリを抑える最適化も実装されています。

実行時パート

配列の次元圧縮

詳細な説明は省きますが、多次元配列の要素にアクセスするとき、一次元配列の要素アクセスよりも時間がかかります。そこで、配列のデータ領域がメモリ上で連続していて低次元配列だとみなせるときに、要素アクセスをより単純に行うための実行時最適化を行います。完全なCUDAコードを生成するためにはこの最適化後の次元を知る必要があるため、完全なコードの生成およびコンパイルは実行時に行われます。CUDAコードのコンパイルはCPU時間に含まれますが、キャッシュされるので一般に高速であり、次元圧縮によるGPU実行時間の削減のほうが効果が大きいため、この方式を採用しています。

実験結果

ResNet50のバッチ正規化部分のコードを融合し、CPU時間とGPU時間を比較しました。実験を行った環境は以下のとおりです。CUDA: 10.1, GCC: 5.4.0 GPU: Tesla P100. また、ResNet50でバッチ正規化される配列のサイズは様々ですが、今回はその中の一つである(32,256,56,56)を採用しました。これは縦横56×56の画像を256チャネル分、32枚を正規化する処理に対応し、軸(0,2,3)について正規化をすることで、出力のサイズは(256,)になります。

比較対象は、カーネル融合をせずにそのまま呼び出した場合(naive)と、実際に採用されている実装である、要素ごとの演算部分だけを融合した実装(custom kernel)です。今回実装したカーネル融合(new fusion)はCPU時間、GPU時間の両方の削減に成功しました。

CPU時間は、実行時の配列のサイズに依存せずほぼ一定の割合短縮されました。また、GPU時間についても、従来のボトルネックであった、分散計算時の余分なメモリアクセスを潰すことができたため大幅に短縮されました。

まとめ

今回のインターンでは、既存のカーネル融合を拡張し、複数回のリダクション演算が含まれる場合にも対応し、無事高速化を達成できました。複数回のリダクション演算にも対応させたカーネル融合を実現する上での理論的な大枠はかなりシンプルで、一時変数に多くの情報(もっとも重要なものは抽象的なサイズ情報)を記録させながら演算をたどれば実現できるのですが、そうして生成されたコードは既存のカーネルで生成されるコードに比べ極めて冗長で、最適化を十分に施さない限り高速化が見込めない、という点が大変でした。アルゴリズムの大枠は序盤から見えていましたが、実際にベンチマークを更新したのはインターン終了直前という形になってしまいました。

アプリケーションに新しい機能を追加をする開発と違い、高速化を目指すための機能追加は、高速化を実現できて初めて世に広まるものなので、プロジェクト自体が成功するかどうかはインターン後半までかなり心配していましたが、無事に高速化を実現できてよかったです。

今回のインターンを通して、頼りになるメンター・副メンターとともにこのような挑戦的なプロジェクトに携われてよかったです。

ベイズ最適化を用いた高次元ブラックボックス最適化手法の検証

Yuhei Otomo

2019-09-27 14:12:18

本記事は、2019年インターンシップとして勤務した森 雄人さんによる寄稿です。


概要

2019年PFN夏季インターンに参加していた森 雄人です. 本インターンでは高次元空間中のベイズ最適化問題について, 実装・研究していました. 特にREMBO [1] [2]  やLINEBO [3] を用いた最適化を行うことで, どのようにアルゴリズムの振る舞いが異なるかを数値的に検証していました.

ブラックボックス最適化

本稿ではブラックボックス最適化問題, 特に最小化問題について考えます. 変数を\(\lambda\), 目的関数を\(L(\lambda)\)と書き, 探索空間を\(\Lambda \subset \mathbb{R}^{D}\)と書くことにします. ここで「ブラックボックス」とは, 最小化したい目的関数\(L(\lambda)\)にまつわる情報が手に入らないという想定を指します. すなわち, 自分である入力\(\lambda\)を選び, その出力\(L(\lambda)\)を得ることはできるけれども, \(L(\lambda)\)の勾配\(\nabla L(\lambda)\)などの情報は手に入れることができない, という最小化問題を扱います. このような状況では勾配法に基づく最適化アルゴリズムを用いることができないため, 異なる手法を適用する必要があります.

ブラックボックス最適化が必要となる代表的なシーンとしては「ハイパーパラメータ最適化」があります. 例えばSupport Vector Machine (SVM)の訓練時において, validationデータに対する損失を, SVMの最適化問題に現れる正則化項係数\(C\)や, カーネルのスケールパラメータ\(\gamma\)の関数として捉えてみましょう. するとこれはvalidationに対する損失を目的関数とした\(C, \gamma\)についてのブラックボックス最適化問題になります. 通常ハイパーパラメータ探索はランダムサーチやグリッドサーチが用いられますが, 近年次に示すベイズ最適化などのアプローチが注目されてきています.

ベイズ最適化 (Bayesian Optimization)

ベイズ最適化とは不確かさを用いながら次に探索すべき点を決定するブラックボックス最適化アルゴリズムです. 具体的なアルゴリズムとしては次のようになります.

キーとなるステップは非常に単純です.

  1. 今まで調べた点を用いて真の関数\(L(\lambda)\)を近似する関数\(a(\lambda)\)を作る
  2. \(a(\lambda)\)を最小化する点\(\hat{\lambda}\)を求める
  3. \(L(\hat{\lambda})\)を求め, 調べた点のリストに追加する

真の関数\(L(\lambda)\)を近似した関数\(a(\lambda)\)を獲得関数 (acquisition function) と言います. この獲得関数の構成に確率過程, 特にガウス過程を用いる手法がベイズ最適化と呼ばれます.

次に示す図は\(a(\lambda)\)としてLower Confidence Bound (LCB) を用いた図です. 今までチェックした点から, 探索できていない領域の不確かさをも含めて関数を近似し, 次に探索すべき点を探す流れを表しています.

ベイズ最適化は一般に, 20次元以上ほどの高次元空間では局所解に陥ってしまったり, 獲得関数の最小化がうまくいかないなどの理由で, 度々うまく動かないことがあります. そこで高次元空間中のベイズ最適化問題が研究されてきています.

REMBO

Random EMbedding Bayesian Optimization (REMBO) [1] [2] は, ランダム行列を用いて低次元空間を高次元空間に埋め込むことで, 高次元空間を陽に扱わずにベイズ最適化を行う手法です. この手法は真の関数\(L(\lambda)\)に
$$ L(\lambda_{\top} \oplus \lambda_{\perp}) = L(\lambda_{\top} \oplus 0) $$
という「縮退した線形空間があって, ある方向に動かしても\(L(\lambda)\)の値が変わらない」性質がある場合に効果を発揮します. これは\(\lambda\)の中に冗長なパラメータがある場合や, \(L(\lambda)\)に寄与しないノイズのような次元があるときに対応しています.

アルゴリズムとしては直接, 高次元空間\(\Lambda\)を扱うのではなく, 低次元空間\(Z \subset \mathbb{R}^{d}\)上でガウス過程を構成し, この空間内で最適化を行います. このとき\(z \in Z\)に対し, 各成分が独立な正規分布\(\mathcal{N}(0, 1)\)に従うランダム行列\(A \in \mathbb{R}^{D \times d}\)を用いて\(L(Az)\)を求めることで真の関数の評価を得ます. なお, 埋め込んだ\(\lambda = Az\)が現在考えている探索空間\(\Lambda\)の外に出てしまっている場合は探索空間上に射影を行います.

LINEBO

Line Bayesian Optimization (LINEBO) [3] は獲得関数の最小化を1次元空間に制限した上で行う手法です. 探索空間を1次元空間に制限する方法はいくつか考えられますが, \(\lambda\)を座標降下させる方法, または\(– \nabla a(\lambda)\)の方向に下る方法が比較的優位であることが数値的に示されています.

 

実験・観察

本稿ではテスト関数\(L(\lambda)\)としてSTYBLINSKI-TANG関数を用います. これはブラックボックス最適化のベンチマークとして使われる標準的な関数の一つで, 入力の次元数\(D\)を柔軟に設定することができます. 今回はベイズ最適化にとって比較的高次元とされる次元数である25に設定し, 探索空間\(\Lambda\)は通常用いられる\([-5.0, 5.0]^{25}\)としています.

実装としてはGPyTorch [4] をベースにしたベイズ最適化ライブラリBoTorchを用いています. GPyTorch・BoTorchを用いることの利点として, 従来ベイズ最適化の計算量のボトルネックであった「ガウス過程の予測」の段階における高速化が挙げられます.

さて, 次に示すのはアルゴリズムごとに各ステップにおける評価値\(L(\lambda)\)をプロットした図です.

 

ここで実験設定の詳細についてさらに述べておきます. Basic-BOとは獲得関数の最大化の際に特に何もせず, そのまま25次元空間上を探索するアルゴリズムを指します. またREMBOにおいて, 低次元空間の次元数は10次元, 探索空間\(Z\)は\([-1.0, 1.0]^{10}\)に設定しています. LINEBOでは降下方向として各次元の軸に沿った方向, すなわち座標降下法を用いています. 降下する軸の選び方は25次元の中から一様ランダムに選んでいます.  global optとは25次元STYBLINSKI-TANG関数の大域的最適解\(\lambda^{*}\)における\(L(\lambda^{*})\)の値を示しています.

また, アルゴリズムに依らずカーネル関数はRBFカーネル, 獲得関数はLCB, 初期点の数は200で統一しています.  REMBOのみ探索空間が低次元空間であるため, 初期点は低次元探索空間\(Z\)上から, Basic-BOとLINEBOは\(\Lambda\)から一様にサンプリングしてきた共通のものを使用しています.

実験結果から得られる観察として, まずどのアルゴリズムも大域的な最適解に収束できていない様子が見られます. 各アルゴリズムごとに見ると, Basic-BOはステップ数が増え, 探索点が増えていった場合にほとんど同じ評価値を取り続けており, 探索が十分に行われていない様子が見られます. 一方でREMBO, LINEBOはステップ数が増えた場合にも探索が行われている様子が確認できますが, 評価値の大きな改善には至っていないことがわかります. しかし, REMBO, LINEBOはBasic-BOに比べ少ないステップ数でより小さい評価値を得ていることがわかります.

結論

本稿では高次元空間中のベイズ最適化手法の収束性について数値的に検証しました. 結果として, どのアルゴリズムもテスト関数に対して大域的な解を得ることはできていなかったものの, REMBO, LINEBOは高次元空間中でも探索する様子が見られ, 標準的なベイズ最適化よりも速い収束性が観察されました.

一点注意すべき点としてベイズ最適化はフレームワークとして柔軟であるが故に, カスタマイズできる点が複数あることが挙げられます. 例えば, 「どのような獲得関数を使うか」, 「カーネル関数はどのようなものを用いるか」, 「初期点はどれぐらいにするか」などといった点です. そのため, 本稿での実験結果がそのままベイズ最適化の限界を表している訳ではないことに留意すべきです.

実際の業務ではテスト関数に留まらず, より複雑な関数に対してベイズ最適化を適用し, どのような結果が得られるか, どのような発展が考えられるかを試行していました. 今回, このような大変刺激的なテーマで研究する機会を頂いたことに感謝致します.

参考文献

[1] Wang, Ziyu, et al. “Bayesian optimization in high dimensions via random embeddings.” Twenty-Third International Joint Conference on Artificial Intelligence. 2013.

[2] Wang, Ziyu, et al. “Bayesian optimization in a billion dimensions via random embeddings.” Journal of Artificial Intelligence Research 55 (2016): 361-387.

[3] Kirschner, Johannes, et al. “Adaptive and Safe Bayesian Optimization in High Dimensions via One-Dimensional Subspaces.” arXiv preprint arXiv:1902.03229 (2019).

[4] Gardner, Jacob, et al. “GPyTorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration.” Advances in Neural Information Processing Systems. 2018.

GraphNVP

Kosuke Nakago

2019-07-16 12:00:06

本記事は、2018年インターンシップ・PEとして勤務した Kaushalya Madhawa さんによる寄稿です。
本記事はこちらの英語ブログの日本語訳です。

本記事では、私たちが最近投稿した論文 “GraphNVP: An Invertible Flow Model for Generating Molecular Graphs” を紹介します。
コードも公開しています → Github repo

分子の生成モデル

望ましい薬理学的な物性値をもつ新しい分子の発見は創薬の分野において重要な問題です。これまで、候補となる化合物を実際に合成して実験を行うといったことがされてきました。
しかし、化合物のとりえる形の可能性は膨大でそれらを合成し、実験するということはとても時間がかかります。
このように望む物性値をもつ分子を探す代わりに、”de novo drug design” では、新しい化合物を興味対象とする物性値とともに設計するということを行います。

近年の深層学習の技術発展、特に深層生成モデルの分野は “de novo drug design” にとって重要な技術となってきています。

分子の表現方法

深層学習で分子の生成モデルを考える際、初めの重要なステップはどのように化合物を表現するかということです。
初期の研究では SMILES という、文字列ベースの表現方法を採用していました。
これらの研究は、 RNN ベースの言語モデルや、Variational AutoEncoder (VAE) を用いることで、SMILES の文字列を生成しそれを分子に変換します。
このアプローチの問題点はSMILES 文字列での小さな変化が分子全体でみると大きく変化するような場合があるということです。
そこで、分子をグラフを用いて表現しようというアプローチが出てきました。この問題は “molecular graph generation” として扱われます。

分子は無向グラフとして表現されます。原子をノード、結合 (bond) を辺に対応させます。この表記を用いると分子の構造・結合情報は Adjacency tensor (隣接テンソル) \( A \) として、また各ノードの原子(酸素、フッ素など)を node feature matrix \( X \) で表すことができます。この定式化の下で、分子の生成問題はグラフの生成問題として扱うことができ、これまでに技術発展してきた GANやVAEといった深層学習のモデルを用いることができます。さらにグラフの生成問題は2つのタイプに分類することができます。1つめは分子のグラフを”順番に” 生成する方法で ノード (原子) や 辺 (結合) を1つずつ生成していきます。もう1つの方法が直接グラフを生成するタイプで、画像の生成と似たように一度に全体を生成します。

可逆性

上記で紹介したVAEやGANと比較して、可逆な Flow を用いたモデルは尤度の最大化を直接行えるという長所があります。

さらに、Flow モデルは可逆な変換のみで構成されているので、逆変換が可能で100% 分子→潜在変数→分子 の再構成を行うことができます。

GANのみの生成モデルなどEncoder がない場合は、特定の分子を操作して生成・最適化を行うことは困難です。例えば、特定のベースとなる分子 (query) に似ている分子を生成して最適化を行っていくということが (創薬の分野で Lead optimization と呼ばれる) 困難です。一方Flow モデルは 分子の潜在変数を直接計算できるので、query 分子に対する潜在変数を簡単に求めることができそこから最適化を行うことも可能です。

提案モデル

model_reverse

提案モデルであるGraphNVPは上図のような構成をとります。GraphNVPは私たちの知る限りではOne-shot Graph生成モデルとして初めてFlow を用いたモデルです。
このモデルは辺に対応する変数とノードに対応する変数の2つの潜在変数を使用し、グラフ構造や原子種類の分布を表現できるようにします。
これらのをFlowの定式下で計算できるよう、 Adjacency Coupling および Node Feature Coupling と呼ばれる2つのCoupling layerを提案しました。
グラフの生成を行う際は、まずAdjacency tensorを生成し、次に Node feature tensor を graph convolutional network を用いて生成します。

結果

訓練データから分子をとり、その潜在変数 \( z_0 \) をGraphNVPで計算します。そして潜在空間上でランダムに選んだ2つの方向の軸に対して \( z_0\) が中心となるような2次元グリッドを構成し、その各点で潜在変数をデコードして分子を生成したものが下図となります。
下図の可視化結果では、似たような分子が並んでおり、隣どおしでは少しの原子が変わっているだけというようなものとなっており、提案モデルがスムーズな潜在空間の学習をできていることが見て取れます。

zinc_neighborhood_2d

 

 

さいごに:メンターより

Kaushalyaさんのメンターを担当した、中郷・石黒です。
今回の研究は2018年夏のインターンシップから始めたものです。このころはGraphの生成モデルの研究が盛んになり様々なアプローチでの生成モデルが提案され始めたころでしたが、まだFlowを用いたグラフの生成モデルはないので取り組んでみたいというKaushalyaさんの持ち込みからはじまったものでした。

グラフ生成モデルとしては初めての取り組みであり、またFlowを用いるモデルはかなり深いネットワークを使用する必要があり計算量も大きくなるということで、研究には時間がかかりましたが、最終的には論文として形にすることができ、コードも公開することができました。

最後に宣伝となりますが、PFNでは今回の “Drug Discovery / Material Discovery” の分野だけでなく、多種多様な業界で機械学習適用の横断プロジェクトを実施しており、中途・新卒の人材募集も通年で行っております!
興味のある方はぜひこちらのページをご覧ください。

2019年 PFN夏季インターンシップのコーディング課題公開

楠本充
エンジニア

2019-06-25 15:27:28

Preferred Networks 2019 夏季インターンシップの選考で用いたコーディング課題を github 上で公開しました。

https://github.com/pfnet/intern-coding-tasks

PFN 楠本です。PFN では毎年8月9月の夏季休暇に約2ヶ月間の長期インターンシップを行っています。コーディング課題はその選考で応募者のプログラミング能力や問題解決能力を見るために出題させて頂いているものです。今年は「機械学習・数理」「バックエンド」「フロントエンド」「チップ」「性能最適化」「コンピュータービジョン(Chainer)」の6種類のコーディング課題を用意し、応募者の希望テーマに応じてこのうちのいずれかを解いていただく形にしていました。

今年のインターンシップでは去年をさらに上回る数の応募を頂きました。PFN では来年以降もインターンシップを開催する予定ですので、これらの課題を見て PFN に興味を持っていただけた方はぜひ応募を検討ください。また、フルタイムの採用も通年で行っております。こちらもご検討いただければ幸いです。