TransE [NIPS'13] を実装(と実験再現)した
Graph embedding を調べる上で避けては通れないっぽいTransEを実装して実験再現してみた。モデルがシンプルでカッコイイし実装も簡単だった。データもパラメータも公開されてて実験を再現できたのもポイント高い。
TransE
NIPS'13で提案されたGraph embeddingをする手法。Google scholarで既に100以上引用されていろんな拡張モデルが提案されてる。論文は以下。
TransEはKnowledge graph(Freebaseとか)をベクトル空間上にembeddingする。入力としてKnowledge graphを与えると、全てのsubject, predicate, objectに対してそれぞれk次元のベクトルを出力する。ポイントは出力されたベクトル空間に構造があること。例えば、
v(Kei Nishikori) + v(profession) = v(tennis player)*1
みたいな足し算ができるようになる*2。これは、("Kei Nishikori", "profession", "tennis player")というトリプルがある(ありそう)ということに対応している。
これの何が嬉しいかというと、例えばKnowledge graphのスキーマを意識しながらSPARQLを書いて問合せするっていうめんどくさいことをしなくても、単にベクトルの足し算をするだけで ("Kei Nishikori", "profession", ?) みたいな質問に答えられるようになる。もっと言うと ("Kei Nishikori", "profession", "tennis player") というトリプルがKnowledge graph上になくても、ベクトルの足し算をするだけで"tennis player"という答えを予測できるようになる。
学習
というトリプルの集合(つまりKnowledge graph)が与えられた時、全てのトリプルについてできるだけ以下を満たすようにしたい:
逆に、Dに含まれないトリプル(例えば ("Kei Nishikori", "profession", "politician") とか)については上記の足し算が成り立たないようにしたい。
これをどうやるかというと、まず、トリプルの "それっぽさ" を測る以下のスコア関数を定義する:
右辺はノルムを表す。これは足し算が成り立つ時は0に近い値を取り、そうじゃない時は大きな値を取る。それっぽいトリプルは足し算が成り立つようにしたいので、このスコア関数が小さい値を取るように上手いことベクトルを割り当ててやる。逆に全然それっぽくないトリプルについてはこのスコア関数が大きな値を取るようにしたい。
うまいことベクトルを割り当ててやるには、以下の目的関数を最小化するようなベクトルを学習すれば良い*3:
f(h,r,t) が小さくて f(h',r,t') が大きければ全体として小さくなる(ガンマはマージンみたいな役割)。ここで (h',r,t') というトリプルが突然出てきてるんだけどこれがすごく重要な役割を果たす。Dには正例しか含まれない、つまりそれっぽくないトリプルは含まれないので、じゃあ人工的にそれっぽくないトリプルを作っちゃいましょうということで出てきたのが (h',r,t')。どうやって作るかというと、あるトリプル (h,r,t) に対して、hかtのいずれかを、ランダムに取ってきたentity(h'もしくはt')と交換する。つまり一つのトリプル (h,r,t) にたいして一つの "それっぽくない" トリプル (h',r,t')を作る。こうすることで正例と負例の両方を使って学習することができるようになる*4。
上記の目的関数の最小化には mini-batch SGD を使う。勾配は単に(L1 or L2)ノルムを微分すれば良い。
実装
transe.pyが本体で、それ以外は論文の実験を再現するためのスクリプト。注意するポイントは特に無さそう。論文に書いてあるとおりに実装できる(えらい)。論文にはearly stoppingしたほうが良いよって書いてあったけどそこだけ無視した(あんまり変わらない)。
Reproducing TransE experiments NIPS'13
実験
論文に書いてある実験結果を部分的に再現してみた。データもパラメータもちゃんと公開されててすばらしい。
データはここで公開されてるFB15kをつかった。これはFreebaseからサンプルされた15,000エンティティを含むデータ。割りと小さいけど評価実験するくらいならこれくらいでいいのかな。
まずデータを用意する。
% wget 'https://everest.hds.utc.fr/lib/exe/fetch.php?media=en:fb15k.tgz' % mv fetch.php\?media=en:fb15k.tgz fb15k.tgz % tar zxvf fb15k.tgz
学習する。パラメータは論文に書いてあるとおり。SGD は mini-batch の size を5000で 1000 epochs 回した*5。標準出力に表示されている 10 1528.218 みたいなのは 10 epochs 回した直後の validation set での mean rank の値を表している(Mean rankについては後述)。結果として学習されたモデルが "transe.model" というファイルに書き出される(pickle)。
% python train.py FB15k/freebase_mtr100_mte100-train.txt FB15k/freebase_mtr100_mte100-valid.txt 0 3560.291 10 1528.218 . . . 990 198.681
論文中の実験を再現する前に少し遊んでみる。predict.pyに 予測したいトリプル (h,r,?) を入力するとそれっぽい t を上位 n 件返してくれる。使用したデータセットはエンティティをFreebaseのmidで管理してるのでわかりづらいんだけど、"/m/051q39/"は有名なテニスプレイヤーの Rafael Nadal を表していて、以下の例ではそれのnationalityを知りたいということになる。
% python predict.py transe.model /m/051q39 /people/person/nationality 10 /m/0b1t1 7.87806854062 /m/0d04z6 8.14848562699 /m/06mkj 8.1821732884 /m/03rk0 8.24374933051 /m/027rn 8.27577702421 /m/0hzlz 8.2857627988 /m/0160w 8.31141651263 /m/015fr 8.37179251603 /m/05r4w 8.4572799401 /m/04g5k 8.4820680158
結果としてFreebaseのmidが出てくるので非常にわかりづらいんだけど、top3は"London", "Cuba", "Spain"となった。正しいnationalityはスペインなので、まあそれなりの結果なのかな。右の列の数値はスコア関数の値で、結果はこの値でソートされている(小さいほどそれっぽい)。ちなみに ("Rafael Nadal", "nationality", "Spain") というトリプルは学習データに含まれていないので、上手いこと正しいトリプルを予測できたことになる。データセットが小さいのであまりおもしろい結果は出ないんだけど、色々試してみるといいかも。
さてさて、論文中の実験を再現する。論文中ではMeanRankとHit10という二つの評価尺度が使われている。MeanRankというのは ("Kei Nishikori", "profession", ?) みたいなトリプルを入力した時に正しい答え "tennis player" が予測されたランクを全てのテストトリプルについて平均をとったもの。小さいほどよい。Hit10というのは同じくテストトリプルを入力した時に正しい答えが上位10位以内に予測された割合。大きいほどよい。
% python test.py transe.model FB15k/freebase_mtr100_mte100-test.txt MeanRank: 190.301450796 Hit10: 0.417192869598
論文中ではMeanRankが243、Hit10が34.9になってるのでそれより結構良い結果が出ちゃったけどまぁいいでしょうということで。
まとめ
とにかくシンプル。ムダなところを一切削ぎ落としたようなモデルでカッコイイね。ただシンプル過ぎてこれでうまくいくのかよという感もある。実際TransEではうまくいかない例もいくつもあって、後続の(他の人が提案した)いろんなモデルで修正されてる。ただ後続の手法も枠組みはTransEとほぼおなじで、おおまかに言えばスコア関数 f を変えてるだけ。たくさんの人に使ってもらうには surprisingly simple and moderately effective なものを作るのが重要だといういい例だった。