Gated Graph Sequence Neural Networks (GGNN) 再実装

Qiitaからの移植

概要

目的：プログラミングタスクにDNNを使うために、DNNでGraph構造を扱いたい
- プログラムのASTを入力とするDNN
- コード中に登場するデータ構造を入力とするDNN
- etc…
再実装対象：Gated Graph Neural Network (GGNN)
- ICLR 2016で発表された技術
- プログラミングタスク（loop invariantの発見）のための技術で、目的に近そう
やったこと
- 再実装した
- bAbIタスクでは確かに高いaccuracyを達成できた

目的

去年夏くらいから、DNN/機械学習を使ったProgram Synthesisに興味を持って、論文の再実装をしたり、その他実験をしてみたりしている。その際、Program SynthesisにDNNを使おうとすると、以下のあたりで困ることが多かった。

扱えるデータ構造が少なすぎる（現在の研究だと、整数、配列/リスト、文字列くらい）
- 合成対象のDSLの構成に強い制約がかかる
- 何らかのアプリケーションとしてまとめようと思っても、結局できることがほとんどない
合成までに時間がかかりすぎる
- DeepCoderの場合、5文程度のシンプルなプログラムに1時間くらいはかかる
- （ただし、RobustFillは実行時間が書かれていないが、アルゴリズムを見る限りではこの問題は解決できそう）

2はRobustFillで解決されていそう（DeepCoderでは全探索していたのをRobustFillはビームサーチに変えている）なので、1を解決できる手法を勉強してみようと考えた。整数・配列/リスト・文字列以外だと、bool・実数・木/グラフ・集合・連想配列あたりが多くの言語で用いられるデータ構造として挙げられる。このうち、boolと実数は整数の応用でなんとかなりそう。また、集合・連想配列は内部では木（二分木）か配列（ハッシュ）を使っていることが多い。従って、木/グラフをDNNで扱えるようになれば、1の問題はだいぶ緩和されるのではと考え、DNNでグラフを扱う研究の再実装を目論みた。

Gated Graph Neural Networks (GGNN)

グラフをDNNの入力として食わせる手法は化学への応用を主目的として複数存在する。例えば、chainer-chemistryは2018/5現在、

NFP
GGNN
WeaveNet
SchNet
Renormalized Spectral Graph Convolutional Network

の5つをサポートしている。今回はこれらの中で、上から2番目のGGNN（Gated Graph Neural Networks）を選んで再実装した。理由は元論文で「loop invariantの発見のために研究した」と書かれていて、プログラミングタスクに応用が効きやすそうな雰囲気がしたため。

技術概要

Graph Neural Network (GNN)という手法を元に、以下の変更を加えたもの

GRUを用いる
勾配計算にAlmeida-Pinedaアルゴリズムではなく、backprop through timeを用いる
sequentialな出力も学習できるようにする

おそらく目的は2と3で、2によってモデルの表現力を強化し、3によって適用できる問題の種類を増やすことが手法の狙い。しかし、2によって勾配消失・爆発のような問題が生じたので、1の変更を追加してそれを抑えている。

再実装

実装方法

肝となるのは、グラフの接続関係に従った結合の実現方法（GGNN独自というよりはGNNから引き継いだ内容に見える）。一度ばらばらにして、辺のlabelごとにaffineをして、もう一度もとに戻す。簡単に書くと以下のようなコードとなる。

# input: 入力側のVariable
# edges: 辺集合

xs = F.split(input)                                  # 頂点ごとに分割する
h_next = []                                          # 出力側のための一次変数
for _ in range(len(xs)):                             # 初期化
    h_next.append([])                                # 初期化続き

for (i, j), label in edges:
    with nn.parameter_scope(label):                  # ラベルごとにパラメータを変える
        h_next[j].append(PF.affine(xs[i], 1))        # 隣接した頂点同士を結合 (順方向)
    with nn.parameter_scope("{}_rev".format(label)): # ラベルごとにパラメータを変える
        h_next[i].append(PF.affine(xs[j], 1))        # 隣接した頂点同士を結合 (逆方向)

h_next = list(map(lambda xs: F.sum(*xs), h_next))    # 頂点ごとに総和を取る
return F.stack(*h_next)                              # 一つのVariableにまとめ直す

ただし、実際にはレイヤー数を減らして高速化を図るため、ラベルごとにまとめてaffineしている。コードはこのあたり。

実験と結果

論文ではbAbI 4, 15, 16, 18, 19について評価しているが、そこまでやるのは面倒だったので、Single Outputから1つ(15)、Sequential Outputから1つ（19）のみを選んで評価。

タスク名	反復回数 [epoch]	学習データ個数	精度 [%]	論文中の精度 [%]
bAbI 15	1	50	100	100
bAbI 19	216	250	95	99

論文と比べるとbAbI 19の精度が低いが、これは精度95%に到達した時点で実験を打ち切るよう訓練スクリプトを書いたためかもしれない。learning curveを見るとまだ下がっていきそうな感じがしたので、もう少し実験を続けると99%まで到達した可能性はある（が飽きた）。

感想

実装自体は面倒な箇所はいっぱいあったけどそこまで難しくはなかった。また、少ない学習データセットで十分実験できるし、学習にそんなに時間がかからないので、そのへんにあったGPUでも十分試せる。DNNの手法再実装の勉強・訓練には（厳密にやるのでなければ）ちょうどよい難易度だった気がする。

論文中に書かれていなくて気になる点としては、問題が複雑になってきた時にdata augmentationどうすればよいのかは気になる。画像であればコントラストを変える等のaugmentationをするが、グラフの場合はどうするのがよいのだろうか。ソースコードを扱う場合は変数名・関数名を変えるなど、まずはタスクごとにaugmentationの方法は考えていく必要がありそう。

その他 (chainer-chemistry)

前述の通り、GGNNはすでにchainer-chemistryで実装されている。自分が今回再実装したのはその存在を知らなかったからというのが大体の理由。ただし、chainer-chemistryのGGNNは、以下の3点が自分の再実装と違いそう。

chainer-chemistryはミニバッチ学習をサポートしている（ただし、頂点数は揃うことを仮定していそう）
- 自分の実装では、頂点数がバッチ間で可変でも効率よく処理する方法が思い浮かばずミニバッチ学習の実装を諦めたが、chainerは頂点数がバッチ間で揃う前提で実装しているように見える
- 化学計算目的だと頂点数固定で事足りるのかも？あるいは最大の頂点数のグラフにそろえて適宜マスクをかければよいということなのかもしれない
- プログラミングタスクのためには頂点数はバラバラであっても効率よく扱いたい場面が多そう。
  例えば、GitHubからコードを落としてきて訓練データにする時、一番巨大なコードに合わせてミニバッチ化するとメモリ消費量が馬鹿にならなさそう。
chainer-chemistryの方はsequential-outputに対応していない
- ただし、つなぐ部分を適切に書けばcahinerの実装でもsequential outputを学習できる気はする
chainerの実装だと、元論文に記載のあるtanh全部なくなっているようにみえる。自分の見落としだろうか？