WaveGlowの要点整理とメモ

昔からWaveGlowの論文は目を通していたが，最近になりようやく理解が追いついてきたので要点整理とメモを行う．

参考リンク

• 論文
• [DL論読会]Glow: Generative Flow with Invertible 1×1 Convolutions
• [DL輪読会]Flow-based Deep Generative Models
• NVIDIA公式の実装

#WaveGlowのここが凄い

• 自己回帰しないので推論が速い（最近だとデフォルトだが）
• 直接尤度の最大化で最適化できる

#Flow-based generative modelの復習
まずはFlow-basedの生成モデルの簡単な復習から行う．こちらの記事が分かりやすいので，詳しく知りたい方はこちらを参照
##簡単に言うと
サンプリングや尤度を求めるのが簡単な潜在変数から，簡単かつ逆変換可能な関数を複数回適用することでデータの分布を求める生成モデル．学習時はデータに逆変換した関数を適用して潜在変数を求めて，その尤度を最大化する．多くは，オーバーフローを起こさないように負の対数尤度を最小化する．
##WaveGlowの場合
球面ガウスに従う潜在変数$\boldsymbol{z}\sim\mathcal{N}(\boldsymbol{z};0,\boldsymbol{I})$から関数$\boldsymbol{f}$を用いて，音声波形$\boldsymbol{x}$を以下のように求める．

\boldsymbol{x}=\boldsymbol{f}_0 \circ \boldsymbol{f}_1\circ \dots \boldsymbol{f}_k(\boldsymbol{z})

この関数$\boldsymbol{f}$の学習は負の対数尤度を直接最小化することで学習する．尤度の計算は以下のようになる．
$$\log p_{\theta}(\boldsymbol{x})=\log p_{\theta}(\boldsymbol{z})+\sum^k_{i=1}\log|\det(\boldsymbol{J}(\boldsymbol{f}_i^{-1}(\boldsymbol{x})))|$$
ただし，

\boldsymbol{z}=\boldsymbol{f}_k^{-1} \circ \boldsymbol{f}_{k-1}^{-1} \circ \dots \boldsymbol{f}_0^{-1}(x)

である．また，$\boldsymbol{J}$はヤコビアンである．ヤコビアンの計算に関して，上記の記事から引用する．

ただし, サイズが$D×D$のJacobianの行列式の計算には$O(D^3)$の計算量が必要なので, $D$が大きくなるとlog-detが計算時間のボトルネックになります. 実際, Glowで生成しているCelebAの画像は$D$=256×256×3≃200,000次元です. したがって, 現実的にはJacobian行列式を高速に計算するための工夫が必要になります. 詳しくは後で述べますが, Jacobianが三角行列になるような工夫をしたり, log-det自体をtraceによって近似したりといったアプローチがあります.

ゆえに，ネットワーク構造の箇所で，ヤコビアンの計算が簡単になる点についても触れていきたい．
#ネットワーク構造
ネットワーク構造は以下のようになる．

R.Prenger et al.2018

基本的には，Glowのアーキテクチャを，音声波形の生成ができるように置き換えたと考えて良い．以下にネットワークの詳細を記す．

##squeeze to vector
入力の音声を8サンプルごとにベクトルとしてグループ化する．実験では，音声データを16000サンプルに区切って学習を行っているため，データサイズを8×2000に変換している．
###なぜ行うか
Real NVPとは違い，所謂step of flowのたびに行っているわけではない．これは，音声波形の時間方向をそのまま扱うのは冗長で，8サンプルでグルーピングしたほうが学習の効率が良いからなのでは，と考えている．正直なところあまり言及されていないのでここはもう少し理解を深めたいところである．

##invertible 1×1 convolution
Glowで用いられている手法である．後述のaffine coupling layerで，$\boldsymbol{x}_a$が常に同じにならないように，1×1 convolutionを用いてaffine coupling layerへの入力の入れ替えを行っている．この変換のヤコビアンは

\begin{align}
\boldsymbol{f}_{conv}^{-1}&=\boldsymbol{Wx} \\
\log|\det (\boldsymbol{J}(\boldsymbol{f}_{conv}^{-1}(\boldsymbol{x})|&=\log|\det\boldsymbol{W}|
\end{align}

となる．
###Glowでのinvertible 1×1 convolutionとの違い
Glowでは，このconvolutionの重みをLU分解を行うことで対角行列にして，行列式の計算を簡単にしている．一方，WaveGlowではこのconvolutionの重みを正規直交で初期化することで逆行列が存在することを担保している．直交行列の性質より，
$$\det\boldsymbol{W}=\pm1$$
であることから
$$\log|\det\boldsymbol{W}|=0$$
と対数尤度が0になる．このため，WaveGlowではLU分解を行わずに直接$\boldsymbol{W}$の行列式の計算を行っていると思われる．数学に弱いので間違えていたらコメントが欲しい

##affine coupling layer
基本的な構造はGlowと同じであるが，WNはWaveNetに似たDilated Convolutionsとskip connectionsで構成されている．違いとしては，causal convolutionではなく，未来方向の情報も畳み込んでいることである．
###計算
まず，入力をチャンネル方向に半分にする．

\boldsymbol{x}_a,\boldsymbol{x}_b=split(\boldsymbol{x})

そして，$\boldsymbol{x}_a$と時間方向にアップサンプリングしたmel-spectrogramを条件付きベクトルとして，WNの計算を行う．

(\log \boldsymbol{s},\boldsymbol{t})=WN(\boldsymbol{x}_a,mel-spectrogram)

affine xformで，$\boldsymbol{x'}_b$を計算する．

\boldsymbol{x'}_b=\boldsymbol{s}\odot \boldsymbol{x}_b + \boldsymbol{t}

ただし，$\odot$は，行列の要素積である．最後に，$\boldsymbol{x}_a,\boldsymbol{x'}_b$を結合して，次のステップへと渡す．

\boldsymbol{f}_{coupling}^{-1}(\boldsymbol{x})=concat(\boldsymbol{x}_a,\boldsymbol{x'}_b)

###affine coupling layerの対数尤度
論文では省略されているが，affine coupling layerのヤコビアンは以下のようになる．

\begin{align}
\det(\boldsymbol{J}(\boldsymbol{f}_{coupling}^{-1}(\boldsymbol{x}))) &= 
\begin{vmatrix}
\frac{\partial \boldsymbol{x}_a}{\partial \boldsymbol{x}_a} & \frac{\partial \boldsymbol{x}_a}{\partial \boldsymbol{x}_b} \\
\frac{\partial \boldsymbol{x'}_b}{\partial \boldsymbol{x}_a} & \frac{\partial \boldsymbol{x'}_b}{\partial \boldsymbol{x}_b}
\end{vmatrix} \\
&=
\begin{vmatrix}
\boldsymbol{I}_d & 0 \\
\frac{\partial \boldsymbol{x'}_b}{\partial \boldsymbol{x}_a} & \boldsymbol{s}
\end{vmatrix}\\
&= \boldsymbol{s}
\end{align}

affine xformのおかげで，ヤコビアンが三角行列になる．ゆえに，
$$\log|\det(\boldsymbol{J}(\boldsymbol{f}_{coupling}^{-1}(\boldsymbol{x})))| = \log |\boldsymbol{s}|$$
とaffine coupling layerのヤコビアンは簡単に求めることができる．

##WaveGlowの対数尤度
以上より，WaveGlowの対数尤度は以下の式で表せる．

\begin{align}
\log p_{\theta}(\boldsymbol{x})=&-\frac{\boldsymbol{z}(\boldsymbol{x})^T\boldsymbol{z}(\boldsymbol{x})}{2\sigma^2} \\
&+\sum_{j=0}^{\#coupling}\log\boldsymbol{s}_j(\boldsymbol{x},mel-spectrogram) \\
&+\sum_{k=0}^{\#conv}\log\det|\boldsymbol{W}_k|
\end{align}

最初の項は，球面ガウスの対数尤度である．$\sigma^2$はガウス分布の分散である．
##Early outputs
ネットワークの図には記されていないが，学習時は4つのcoupling layerごとに出力から2チャンネルのデータを損失関数に加えている．実験ではinvertible 1×1 convolution，affine coupling layerそれぞれ12層で構成されているため，４層目と８層目のときにEarly outputsを行っている．こうすることにより，浅い層の学習の効率が良くなるとのこと．発想としてはReal NVPのmulti-scale architectureに似ている．

##推論
推論時は，ガウス分布からサンプリングした潜在変数$\boldsymbol{z}$を入力として，音声波形を生成する．学習時は$\sigma=\sqrt{0.5}$で学習を行うが，推論時は$\sigma=0.6$と値を変更してサンプリングを行う方が良いとのこと．
推論時，invertible 1×1 convolutionは重みの逆行列を求める．一方，affine coupling layerでは，

\boldsymbol{x}_a=\frac{\boldsymbol{x'}_a-\boldsymbol{t}}{\boldsymbol{s}}

という計算によって逆変換を行う．これは学習時の$\boldsymbol{x'}_b=\boldsymbol{s}\odot \boldsymbol{x}_b + \boldsymbol{t}$という計算の式変形を行うことで得ることができ，これにより逆行列の計算コストを削減している．
#評価
サンプル音声．聞いた印象だと

• WaveNetよりも，「スーッ」といった少しノイジーな音声の再現度がより自然に聞こえる．（サンプル音声の上から3番目は分かりやすいかと）
• ただし，Tacotron2+WaveGlowだと，少し機械音のようになってる箇所がある．これがTacotron2のメルスペクトログラムの精度によるものなのかは判断が難しい

#感じたこと

• 他のFlow-basedのモデルにも通ずることだが，教師データとの差分を取って誤差を最小化するのではなく，尤度を直接最大化する考え方は面白いなと
• ただ，1つ1つの重みの表現力が制限されることで，結果としてかなり深いネットワークになることが考えられる．GPUメモリの効率としてはどうなのだろうと疑問
• VAEやGANと比べて，最適化の難易度が気になる．

#まとめ
Flow-basedな生成モデルを用いたニューラルボコーダであるWaveGlowの要点整理を行った．かなり数学的なモデルになっていて，尤度を直接最大化するというのは非常にかっこいい．しかし，いかんせん数学の知識がまだまだ及ばないため，これら周辺技術の調査をしつつ，わからない知識は適宜勉強し直して行きたいところである．また，過去に実装を動かしたことがあるがソースコードの理解ができていなかったため，ソースコードの理解にも努めていきたい．