Embeddingを高速に取り出すMagnitude - Technical Hedgehog

Technical Hedgehog

機械学習と気になる技術を語るはりねずみ

Embeddingを高速に取り出すMagnitude

word2vecなど分散表現が活躍するシーンは多いですが、実行速度やメモリなど利用する上で気がかりになる面もあります。この記事では分散表現をすばやく便利に使うためのパッケージであるMagnitudeの説明と実行速度について実験した結果を紹介します。

What's Magnitude?

Magnitude は簡単にすばやく単語ベクトルを扱うためのライブラリです。

github.com

EMNLP2018で発表された Magnitude: A Fast, Efficient Universal Vector Embedding Utility Package の著者実装となります。論文中で挙げられているように、Gensim の代替を想定しているようです。

Magnitude Object について

Magnitudeではベクトル取得を高速化するために、事前に分散表現を変換したものを用います。それぞれの分散表現ごとに独自のフォーマット".magnitude"のオブジェクトとして用意しています。初期ロード時は遅延ロードするようになっているためオブジェクトの読み込みに時間がかからないようにし、一度呼び出された単語はメモリ上にキャッシュすることで再取得時には高速に取り出すことができます。SQLiteを用いることにより実現されているようです。

こちらに事前学習済みのMagnitude Objectが公開されており、後述する未知語に対する処理が異なるLight, Medium, Heavey の3種類が用意されています。用意されているモデルはword2vecやGlove、ELMoなど様々であり、今後は話題のBERTも追加される予定のようです。

Out-of-Vocabulary について

ここではMagnitudeの特徴的な機能(だと私は思っている)である、未知語(Out-of-Vocabulary)に関する機能を紹介します。

基本的な機能を含めて利用方法自体は簡単であり、githubのREADMEを見ていただければ困らないと思います。

Basic Out-of-Vocabulary Keys

新しい単語やスペルミスなど、分散表現を構築する際の学習データに含まれない未知語に対し、Magnitudeは適当であるランダムなベクトル値を割り振ります。ざっくり言うと、同じ単語であれば2回目以降は同じベクトルを返すようにすること(辞書に新規登録するイメージ)と、n-gramが近しい未知語はベクトルも近しくなるように割り振るようにしているそうです。Lightモデルで利用することが可能です。

Advanced Out-of-Vocabulary Keys

こちらも未知語に対する処理方法ですが、ランダムにベクトルを割り当てるのではなく、辞書に登録されている近しい単語に似るようにベクトルを割り当てます。スペルミスやtypoも近似できるようになっています。MediumやHeavyモデルで利用することができます。

内部でどのような処理を行っているのかの説明と、日本語に対しても処理できるようにした実装が以下の記事で紹介されています。とてもおもしろいです。

yag-ays.github.io

Speed について

論文ではGoogle Newsコーパスで学習したword2vecについてMagnitudeとGensimの実行時間についての比較がされており、初期ロードでは97倍、1単語の初期呼び出しキー(cold key)については等倍なものの、再呼び出しキー(warm key)では110倍高速に処理することができるとされています。その他のベンチマーク結果はこちらから確認することができます。

以下からはMagunitudeを用いることで学習速度を向上させることができるのか、実験で確かめていきます。

速度実験

Magnitudeが他のライブラリに比べて学習に掛かる時間を短縮させることができるのか確認してみます。今回は速度比較を主眼とするため、いずれもライブラリの場合でも全結合層を1層だけ用いただけのモデルであり、違いは分散表現の取得処理だけが違いとなるように実装しています。私のgithubにおいてますのでよければ見てみてください。

  • 比較対象
    Gensimとの比較ではword2vecを、TensorFlow Hubとの比較ではELMoを用いて比較しました。

  • データセット・タスク
    MovieLens の映画ジャンルにおけるマルチラベル分類。 Building a text classification model with TensorFlow Hub and Estimators で紹介されているものを利用。これはMovieLens のうち特定の9種類ジャンルに関するデータだけを抽出した14,993件を用いています。このうち80%を学習データとして、20%をテストデータとして利用します。

  • ハイパーパラメータ

    • 隠れ層ユニット数 : 64
    • バッチサイズ : 32
    • ステップ数 : 375
    • エポック数 : 5
  • マシンスペック

Magnitude vs Gensim

word2vecを利用した場合以下の通りになりました。Magnitudeはword2vecのMediumを利用しています。結果を確認してみましょう。

Magnitude with word2vec

Epoch 1/5
375/375 [==============================] - 5455s 15s/step - loss: 0.4881 - acc: 0.8169 - val_loss: 0.4049 - val_acc: 0.8411
Epoch 2/5
375/375 [==============================] - 4s 12ms/step - loss: 0.4279 - acc: 0.8209 - val_loss: 0.3760 - val_acc: 0.8497
Epoch 3/5
375/375 [==============================] - 2s 6ms/step - loss: 0.4057 - acc: 0.8312 - val_loss: 0.3557 - val_acc: 0.8588
Epoch 4/5
375/375 [==============================] - 1s 4ms/step - loss: 0.3843 - acc: 0.8386 - val_loss: 0.3360 - val_acc: 0.8654
Epoch 5/5
375/375 [==============================] - 1s 4ms/step - loss: 0.3662 - acc: 0.8453 - val_loss: 0.3239 - val_acc: 0.8699

Gensim with word2vec

375/375 [==============================] - 13s 36ms/step - loss: 0.4396 - acc: 0.8246 - val_loss: 0.3887 - val_acc: 0.8565
Epoch 2/5
375/375 [==============================] - 11s 29ms/step - loss: 0.3607 - acc: 0.8502 - val_loss: 0.3194 - val_acc: 0.8759
Epoch 3/5
375/375 [==============================] - 9s 24ms/step - loss: 0.3254 - acc: 0.8630 - val_loss: 0.2967 - val_acc: 0.8825
Epoch 4/5
375/375 [==============================] - 9s 24ms/step - loss: 0.3084 - acc: 0.8702 - val_loss: 0.2813 - val_acc: 0.8867
Epoch 5/5
375/375 [==============================] - 10s 25ms/step - loss: 0.2984 - acc: 0.8743 - val_loss: 0.2764 - val_acc: 0.8877

すべてのエポックをトータルで見るとGensimのほうが速いですね…ただGithubにあげているソースを見ると分かるように、Gensimはword2vecのモデルをロードする際に2,3分時間がかかっていましたが、Magunitudeは一瞬でロードが終わっています。加えて、2エポック目以降はMagnitudeのほうが速かったです。

また、実装上の問題で全く同じ処理にできているかと言うと異なります。Gensimでは未知語をすべて同一の特定単語に置き換えている一方、Magnitudeでは未知語の推測を行っているためです。この当たりが速度に影響を及ぼしているかは今回は確認できていません。

Magnitude vs TensorFlow Hub

ELMoを利用した場合の結果は以下の通りになりました。MagnitudeはELMoのMediumを利用しています。ELMoの概要については以前に以下の記事で紹介しました。

kamujun.hatenablog.com

Magnitude with ELMo

Epoch 1/5
375/375 [==============================] - 8449s 23s/step - loss: 0.2545 - acc: 0.8932 - val_loss: 0.2394 - val_acc: 0.9064
Epoch 2/5
375/375 [==============================] - 147s 391ms/step - loss: 0.2501 - acc: 0.8950 - val_loss: 0.2361 - val_acc: 0.9067
Epoch 3/5
375/375 [==============================] - 169s 450ms/step - loss: 0.2461 - acc: 0.8970 - val_loss: 0.2338 - val_acc: 0.9079
Epoch 4/5
375/375 [==============================] - 185s 494ms/step - loss: 0.2419 - acc: 0.8985 - val_loss: 0.2338 - val_acc: 0.9072
Epoch 5/5
375/375 [==============================] - 163s 436ms/step - loss: 0.2385 - acc: 0.9001 - val_loss: 0.2325 - val_acc: 0.9082

TensoFlow Hub with ELMo

Epoch 1/5
375/375 [==============================] - 13045s 35s/step - loss: 0.3439 - acc: 0.8597 - val_loss: 0.2966 - val_acc: 0.8945
Epoch 2/5
375/375 [==============================] - 10262s 27s/step - loss: 0.2938 - acc: 0.8787 - val_loss: 0.3074 - val_acc: 0.8933
Epoch 3/5
375/375 [==============================] - 10268s 27s/step - loss: 0.2845 - acc: 0.8830 - val_loss: 0.2725 - val_acc: 0.8972
Epoch 4/5
375/375 [==============================] - 10247s 27s/step - loss: 0.2746 - acc: 0.8872 - val_loss: 0.2568 - val_acc: 0.8966
Epoch 5/5
375/375 [==============================] - 10244s 27s/step - loss: 0.2654 - acc: 0.8888 - val_loss: 0.2553 - val_acc: 0.8963

こちらは明確にMagnitudeの方が速くなりました。あまりにも速くなっているので心配になるのですが、学習が進んでいるところを見ると実装が明らかに間違えているわけではなさそうです。(Sequenceを作る際のシードを固定し、TensorFlow Hubと比較してやってもよかったかも)

速度以外に確認した点として、MagnitudeでどのようにELMoを実装しているのかソースを確認しました。ELMoは文脈を考慮した単語ベクトルを得るため文毎に計算が必要ですが、Magnitude内部で本家であるAllen NLPのELMoを用いて計算しているようでした。

まとめ

Magnitudeを用いることでEmbeddingを高速に取り出すことができることを確認しました。word2vecにおいてはGensimも相当早いので、短文では明確なアドバンテージがあるとは言えなかったものの、ELMoではTensorFlow Hubより高速に取り出すことができました。Magnitudeは速さ以外にも未知語の処理やメモリ使用量なども強みがあるので、公開されている事前学習済みモデルを用いる際は利用する価値があるのではないかと思います。自分で作った分散表現もMagnitude Objectに変換できるようなので、自作に挑戦しても良いですね。

フォローやコメントなどで応援して頂けると励みになります。いつもありがとうございます!