deeplearn.js で Autoencoder

MNIST のデータを使って、 deeplearn.js で Autoencoder を実装しました。
はじめに断っておくと、僕は Web フロントエンジニアで、機械学習はずぶの素人です。

作ったもの

動かすと、 cost の変化や、その時点での Autoencoder の出力結果がブラウザに表示されます。
よかったら遊んでみてください。

Autoencoder とは

詳細な説明は省きますが、自己符号化器と訳されるもので、 encoder, decoder で変換した出力が、入力と同じになるように学習を行います。
DNN の事前学習や、データ圧縮、異常検知などの目的で利用されます。

作ったもの解説

主にこちらの TensorFlow の実装を参考にして作りました。
学習部分だけ抜き出すと、以下のような感じです。

const graph = new Graph();
const rni = new RandomNormalInitializer(0, 0.1);
const ci = new ConstantInitializer(0.001);
const x = graph.placeholder('x', [INPUT_VECTOR_LENGTH]);
const math = new NDArrayMathGPU();

// 1st encoding phase
const WEncode1 = graph.variable('WEncode1', rni.initialize([INPUT_VECTOR_LENGTH, 256], 0, 0));
const bEncode1 = graph.variable('bEncode1', ci.initialize([256], 0, 0));
const hEncode1 = graph.sigmoid(graph.add(graph.matmul(x, WEncode1), bEncode1));

// 2nd encoding phase
const WEncode2 = graph.variable('WEncode2', rni.initialize([256, 128], 0, 0));
const bEncode2 = graph.variable('bEncode2', ci.initialize([128], 0, 0));
const hEncode2 = graph.sigmoid(graph.add(graph.matmul(hEncode1, WEncode2), bEncode2));

// 1st decoding phase
const WDecode1 = graph.variable('WDecode1', rni.initialize([128, 256], 0, 0));
const bDecode1 = graph.variable('bDecode1', ci.initialize([256], 0, 0));
const hDecode1 = graph.sigmoid(graph.add(graph.matmul(hEncode2, WDecode1), bDecode1));

// 2nd decoding phase
const WDecode2 = graph.variable('WDecode2', rni.initialize([256, INPUT_VECTOR_LENGTH], 0, 0));
const bDecode2 = graph.variable('bDecode2', ci.initialize([INPUT_VECTOR_LENGTH], 0, 0));
const hDecode2 = graph.sigmoid(graph.add(graph.matmul(hDecode1, WDecode2), bDecode2));
const costTensor = graph.meanSquaredCost(x, hDecode2);

INPUT_VECTOR_LENGTH は 28 * 28 = 784 で、ここでは encoder で 784 -> 256 -> 128 と圧縮し、　decoder で 128 -> 256 -> 784 と次元を元に戻しています。
簡単ですね。

細かい話をすると、 initializer (RandomNormalInitializer, ConstantInitializer) の initialize メソッドの第二、第三引数に 0 を指定していますが、この値はなんでもいいです。
なぜなら使われていないから。
でも TypeScript の定義上必須の引数になっているので、渡さないとコンパイルが通りません。悲しい。

上のコードにはない部分ですが、元の Tensorflow の実装と違うところとして、以下のような点があります。

• batch size を 256 -> 32 に変更
  • 256 だとすごく重かったので。。
• optimizer を RMSPropOptimizer -> AdadeltaOptimizer に変更
  • adadelta は Keras の Autoencoder のチュートリアルで使われていたもので、試しに使ってみたら RMSPropOptimizer より結果がだいぶ良かった

結果

30,000 step ほど学習を回すと、以下のような出力が取れるようになりました。 (左が元画像、右が出力)

この間の cost の遷移は以下のような感じ。 (x 座標が step 数、 y 座標が cost)

ちょっと分かりづらいですが、一番最初は 0.13 くらいから始まり、そこから 0.3~0.4 くらいまではすぐに (100 step くらい) 下がって一旦そこで落ち着きます。
なので最初、その落ち着いたところで大体学習が終わったのかな、と思ってしまい、その時の結果がこんな感じだったので、

実装失敗したかな？と思いました。
5,000 step くらいまで行くと、ようやく数字ごとに特徴的な形状が現れるようになり、安心しました。

どうも学習は二段階で進んでいるような印象があって、一段階目では、全ての数字に共通の、真ん中付近に白いものがある、というような特徴を学習し、その後二段階目で、数字ごとの特徴を学習するようです。
前者の学習はすぐに終わりますが、後者の方はかなり時間がかかリます。
こういう学習パターンもあるんですね。勉強になりました。

Convolutional Autoencoder

本当は、 Convolutional Autoencoder を実装したかったのですが、現時点では deeplearn.js による Convolutional Autoencoder の実装は難しそうです。
Tensorflow で Convolutional Autoencoder を実装する場合、 decode 時に tf.nn.conv2d_transpose() というメソッドを使うのですが、これに相当するものが deeplearn.js に用意されていないためです。
ひとまず今回は、関連する deeplearn.js の issue に、「こういうメソッドいりますよね？」みたいなコメントをつけるにとどめておきました。

終わりに

ブラウザ上での機械学習は、学習済みモデルをブラウザ上で動かすようなものが多いですが、学習からブラウザ上で動かせると、色々と可能性が広がってくると思います。
Autoencoder はモデルが単純なので、これで何か面白いことができないか、探っていきたいです。