Monthly Hacker's Blog

プログラミングや機械学習の記事を中心に書きます。

chainerでVQ-VAEを実装した

できること

この記事では、次のことができるようになります。

  • CVTK-Corpusを使ってVQ-VAE(speaker-unconditional)
  • chainer v3でfunctionおよびlinkを自作する(WIP)

はじめに


こんなツイートをしたら、WaveNet著書の1人の全さんをはじめ、多くの方に反応していただき引くに引けなくなったので、記事を書くことにしました。

VQ-VAEとは

VQ-VAEはVariational AutoEncoder(VAE)の一種です。1番の特徴は、潜在表現(encoderの出力)がVector Quantization(VQ)されたことで、意味のある潜在表現を獲得したことです。とはいえ、これだけで本質を理解するのは難しいので、噛み砕いて考えていきます。

VQ-VAEを一旦忘れて、conv-deconvで構成されたAutoEncoderを考えます。潜在表現がチャンネル数d、高さh、幅wの画像だと解釈します。この画像の1ピクセルは、d個の連続値のベクトルだとみなすことができます。とすると、このピクセルの取りうる値は \mathbb{R}^dという集合の1要素です。つまりありとあらゆる実数値を取りえます。

ではVQ-VAEはどうかというと、潜在表現の1ピクセル(d個の連続値のベクトル)は「要素数kのd次元ベクトル集合」の1要素です。これがVQたる所以です。潜在表現が取りうる値の集合がぐんと小さくなるため、1つ1つのベクトルの持つ意味が相対的に大きくなることが期待されます。実際に音声で実験をすると、要素数kのd次元ベクトル集合の各要素が音素に相当する表現を学習していることからも、この期待が的外れではないことが分かります。

ネットワーク構造

ぼんやりとVQ-VAEの利点がつかめたでしょうか。正直これだけだとまだ理解しにくいと思うので、ネットワークの構造を見て理解を深めていきます。今回は音声の場合の図を用意しました。
f:id:d-higurashi:20171117011623p:plain
簡単のため、ミニバッチのうち1サンプルのみに注目しています。直方体は多次元配列を表していて、shapeはチャンネルxたてx横です。quantizeはVQのことではなくmu-law変換のことを表しています。

AEの名を冠していながら、実はAEではなく自己回帰モデルのWaveNetが根幹にあるのが気になりますが、細かいことは置いておきましょう。DNNに慣れた方ならこれでおおよその全容がつかめるかと思いますが、キモはstraight through(ST)というネットワークです。これは入力の各ピクセル(d次元ベクトル)を量子化するネットワークです。具体例を挙げて簡単化してみます。

あるピクセルの値が \begin{eqnarray} \begin{bmatrix} 1 & 1 & 2 \end{bmatrix} \end{eqnarray} だとします(d=3)。このとき、STはd次元ベクトルをk個(このkは初期化の際に決める数値)もったネットワークです。今回は例としてk=2で、d次元ベクトルはそれぞれ a = \begin{eqnarray} \begin{bmatrix} 1 & 2 & 3 \end{bmatrix} \end{eqnarray} と b = \begin{eqnarray} \begin{bmatrix} 3 & 2 & 1 \end{bmatrix} \end{eqnarray} だとします。STはベクトルを入力するとk個のベクトルのうち、入力との距離が最も小さいベクトルを返します。計算を簡単にするためマンハッタン距離で計算すると、aと入力の距離は 0 + 1 + 1 = 2でbと入力の距離は 2 + 0 + 1 = 3です。そのため、 \begin{eqnarray} \begin{bmatrix} 1 & 2 & 3 \end{bmatrix} \end{eqnarray} が出力です。

以上がフォワード時の振る舞いです。もう少し平たい言葉で言うならば、k個のベクトルの中から最も似ているベクトルを取り出すと言えます。これが量子化に当たるわけです。では、この関数は微分できるかというと、できません。つまり、誤差逆伝播ができないということです。一体どうやってこの問題を解決しているかというと、バックワード時はSTが恒等関数だとして計算しています。入力と出力の値が近いのであれば、恒等関数と近似しても問題ないということでしょうか。たしかに学習時にはSTの入力と出力の差の二乗和を最小化するように学習しています。

chainer v3でfunctionおよびlinkを自作する(WIP)

実装時の注意ですが、chainerではST層がありません。そのため、VQ-VAEのためにST層を実装しました。この辺りも得られたtipsとかまとめたいのですが、うまくまとめられるほど自分の中で整理できていないのでWIPとします。しばらくお待ちください。なお、コード自体はgithubにありますので、そちらを参考にしてください。

CVTK-Corpusを使ってVQ-VAE(speaker-unconditional)

実際に学習した結果を聞いてみます。どうやらはてなブログでは音声の埋め込みができないらしいので、githubへのリンクです。

入力
https://raw.githubusercontent.com/dhgrs/chainer-VQ-VAE/audio/input.wav

出力
https://raw.githubusercontent.com/dhgrs/chainer-VQ-VAE/audio/output.wav

話者変換などない単純な再構築です。著者サイトの音源に比べてかなり音質が落ちています。これは論文に明記されていないパラメータの差が原因かと考えています。ただ、私の実装に不備があるかもしれないので、見つけた方はぜひPRしてください。

パラメータ調整

論文中では画像向けVQ-VAEのネットワークはある程度詳しく記載されていますが、音声向けは結構省略されています。例によってWaveNetのチャンネル数は詳細が記されていませんし、encoderもチャンネル数が分かりません。また、バッチサイズも画像のときとは違うのではないかと予想しています。他にも、入力サンプル長は何秒分かなど、不明な点が多いです。計算資源がある方には色々試してみるといいかと思います。調整できるパラメータはopt.pyにまとめてあります。

最後に

TF実装やpytorch実装はあったものの、音声に適用していなかったので、実装してみました。CapsNetに話題を持って行かれましたが、VQ-VAEは音素に頼らない音声合成の新たな一歩になる重要な技術だと考えています。この実装を元に多くの方が興味を持ってさらに音声合成が発展するのを楽しみにしています。