線形回帰を１つ１つ改造して変分オートエンコーダ（VAE）を作る
こんばんは．
今日は統計や機械学習において最も基本となる手法である線形回帰から出発し，１つ１つモデルや学習方法に変更を加えていき，最終的に深層学習の分野で非常に良く使われている生成モデルである変分オートエンコーダ（variational auto-encoder，VAE）*1*2を導いていきたいと思います．
2014年に発表されたVAEは，勾配近似を得るためのreparametrization trickや，効率的に潜在変数を近似推論する認識モデル（recognition model, inference model）の導入，確率的勾配法（stochastic gradient descent，SGD）の2重適用など，様々なアイデアが散りばめられている確率的生成モデルです．背景としては，当時ニューラルネットワークを用いて画像を生成するといったアプローチが（CNNを使った識別などと比べて）それほどメジャーではなかったこと，またニューラルネットワークに対して変分推論（variational inference）を適用する事例もそれほど多くなかったことがブームのきっかけになったのではないかと推察しています．
生成モデルや変分推論に詳しくない方は，VAEの論文を読んで「てんこ盛り」感があるように思えるかもしれないですが，１つ１つの要素を分解して考えてみればどのような必然性を元にVAEが構成されているのかを知ることができます．具体的には，線形回帰モデルに対して次のような近似推論の導入とモデルの拡張を施せばVAEが得られます．
個人的には，最先端の「全部載せ」アルゴリズムをただ漫然と受け入れるのではなく，個々の要素の必然性を理解することが機械学習エンジニアの技術力につながってくると思っています．逆に言えば，この記事によって「俺の課題，VAEじゃなくてPCA（主成分分析）で十分じゃん！」みたいになってくれれば大成功です．
線形回帰
Zを入力データセット，Xを出力データセット，Wをこのモデルのパラメータであるとします．これらの同時分布を次のようにモデル化します．
ただし，p(Z)はZの事前分布で，各潜在変数z_nに対して次のような等方ガウスを使うことにします．
観測モデルp_w(X|Z)は次のような線形性を仮定します．
ここでは，計算の簡便性からガウス分布Nを観測分布として採用しており，λはその精度パラメータ（分散の逆数，今回は学習せず固定）です．
式（３）のパラメータWを尤度最大化（最尤推定）で学習することにすると，
と書けます．せっかく設定した事前分布p(Z)ですが，結果的にWの学習に何の影響も与えないことがわかりますね．
尤度が最大となるWは
を解くことによって次のように解析的に求められます．
あとは新規の入力データz_testに対する未知の出力値x_testを次のように点として予測するのが一般的です．
線形次元削減（確率的主成分分析）
さて，先ほどとまったく同じモデルを使い，今度はZが観測できないようなシチュエーションを考えてみます．つまり，Zを潜在変数として取り扱います．この場合，最尤推定を実施するためには次のようにZを周辺化した上でWに関して尤度最大化を行わなければなりません．
式（９）は残念ながらWに関して解析的に最大値が求められないことが知られています．ここでは次のようなEMアルゴリズムによる繰り返し最適化を用いて，尤度の局所的な最大値を求めることにします*3．
＜E-step＞
＜M-step＞
iは最適化の繰り返しのインデックスです．EMアルゴリズムの詳しい意味に関してはPRMLなどをご参照ください．
このモデルでは，各E-stepおよびM-stepは解析的に実行することができます．E-stepでは，パラメータWを固定した上でのZの事後分布は
とガウス分布で表すことができます．一方，M-stepはパラメータWで微分すれば線形回帰とほとんど同じ計算により
と解析解が求められます．
ところで，このような主成分分析（PCA）でもVAEのように画像を生成したり，次のように欠損値を補完したりすることができます．
また，モデルに対してちょこっとだけ変更を加えれば次のようなモザイク除去も作ることができます．
このように，単純な確率的主成分分析を使いこなすだけでも本質的にVAEなどの深層生成モデルでやられていることは実現できます．もちろん，データが大量にある場合や，データが高度に非線形な構造を持っていると思われる場合は，VAEをはじめとしたフレキシブルなモデルの方が性能が良くなります．
ミニバッチ化
先ほどのEMアルゴリズムをミニバッチで学習できるように変更しましょう．これにより，メモリ効率はもちろん，学習の効率性も上がることが知られています．いくつかやり方がありますが，ここではstepwise EM*4というシンプルでメモリ効率の良い手法を使います．
簡単のため，データセットからランダムに１つのデータx_nを取ってくることにします．E-stepでは，パラメータWを固定した上で，１つのデータx_nに対する事後分布のみを推論してあげればよく，
＜E-step＞
となります．ここで計算した事後分布による期待値を次のようにM-stepのパラメータ更新で使います．
＜M-step＞
ここでη_iはステップサイズです．イメージとしては，ひとつ前のステップiまでに学習されたパラメータの値W(i)を保ちつつ，ステップi+1で新しい更新量を重み付きで加えてあげる感じです．
また，同様の考え方は完全なベイズ版であるstochastic variational inference*5にも利用されています．
非線形関数の導入
次にモデルに非線形関数を導入してみましょう．ちょっと深層学習っぽくなってきますね．次のように，p(X|Z)に対して，パラメータθに関して微分可能な非線形関数f_gを導入します．
非線形関数としては，例えば次のように単純にtanhを導入し，
としても良いですし，非線形関数を再帰的に適用してもっと深い層を考えてあげても良いです（ここではθ={W, b}としています）． さて，ここで生じる問題点としては，非線形関数f_gを導入したことにより，E-stepおよびM-stepがそれぞれ解析的に計算できなくなってしまうことです．したがって，それぞれのstepで勾配を利用した最適化を導入することにします． まず，E-stepですが，解析的に簡単な分布（前の例ではガウス分布）が導出できないので，事後分布がガウス分布であると無理やり仮定し（=近似する事後分布のクラスを制限する），その中でKLダイバージェンスがなるべく小さくなるように調整します．つまり，まずz_nに対する近似分布を
と置きます．ここでφ_n＝{μ_n，σ_n}は変分パラメータと呼ばれます．さらに，次のように確率的勾配法によりKLダイバージェンスを小さくするように変分パラメータを更新します．
＜E-step＞
α_Eはステップサイズです．実はこの勾配計算もまだ解析的には実行することが出来ません．というのも，式（２１）のKLダイバージェンスの項を書き下してみると，
となり，式（２２）の第１項に非線形関数f_gを含んだ複雑な期待値計算が現れてしまいます．
ここでは解析的に積分を実行し期待値を計算することは諦め，z_nを１つサンプルすることにより近似的に期待値を求めることにします．z_nをそのままサンプルしてしまうと，式（２２）の第１項目が変分パラメータφ_nの関数でなくなってしまうので，後で変分パラメータの勾配が計算できません．これを解決する方法はいくつかありますが，VAEの元論文で使われているのはreparametrization trickと呼ばれている手法で，次のように変数変換を行ったうえでz_nをサンプリングすることにします．
こうすることにより，微分対象となる変分パラメータのないq(ε)の分布に対する期待値を取っていることになり，サンプリングで積分を近似した後に変分パラメータφ_nに関して微分ができるようになります．
M-stepはもう少しシンプルで，単純にz_nをサンプルしたあとにθで微分すれば勾配が求まります．
＜M-step＞
ここで式（２６）における負の対数尤度は，ニューラルネットワークにおける誤差関数と等価であり，通常の合成関数の微分で勾配が計算できます（誤差逆伝播法，backprop）．
認識モデルの導入
VAEの元論文では，上記のミニバッチを使ったE-stepの更新をより効率化するために，次のようなニューラルネットワークを使った認識モデル（recognition model）*6を提案しています．
ここではμやσは一時的に置かれている変数に過ぎず，新しく導入した非線形関数f_rの出力になっています．先ほどの例だと，１つのデータ点x_nに対して潜在変数z_nの近似事後分布を１つ１つ最適化しなければならなかったのですが，変分パラメータφ={W,b}を使った非線形関数を使うことにより，異なる複数のデータ点に対しても対応するzの近似事後分布が更新されるようになります．これは，似たような入力xに対しては似たような潜在変数zの事後分布になるはず，というヒューリスティクスから由来しています（amortized inference*7）．これにより収束の高速化を狙おうというのがここでの認識モデル導入の狙いです．「変分パラメータが複数のデータ点にわたってシェアされている」と元論文では説明されています．
認識モデルを導入することにより，変分パラメータの最適化にも誤差逆伝播法を適用する必要が出てきます．E-stepにおける式（２１）のKLダイバージェンスの項は，今回の認識モデルを用いた場合は次のようになります．
ここでz_nは先ほどと同様，reparametrization trickを使います．式（２９）の最後の項はニューラルネットワーク（非線形関数）ですので，通常の誤差逆伝搬によりφに関する勾配が計算できます．
今回は，統計学や機械学習における最も基本的な予測モデルである線形回帰に対して１つ１つ変更を施していき，深層生成モデルとして有名な変分オートエンコーダの成り立ちを確認してみました．１つ１つの技術はそれなりに必然性を持って導入されているものであり，逆に言えばデータ数の多くない場合やデータの複雑な非線形構造を抽出する必要がない場合は最も単純な潜在変数モデルである線形次元削減でも十分であることも説明しました．
ちなみに今回の記事はすべてパラメータを確率変数として扱っていない非ベイズのアルゴリズムになっています*8．今回紹介したアルゴリズムはあくまで解説目的であり，実際はパラメータに対する正則化項を付け加えたMAP推定版にしたり，パラメータの事後分布も考慮するような完全な近似ベイズ推論の枠組みで実装するのが理想です（議論の大筋は今回と変わりません）．
モデルの拡張は実はこれで終わりではありません．前回の記事（深層学習はガウス過程）でも解説しましたが，ニューラルネットワークにおける重みパラメータの数を無限大にしてあげるとガウス過程に一致します．今回のVAEの場合でも，非線形関数f_gの重みパラメータを無限大に増やすことにより，ガウス過程潜在変数モデル（Gaussian process latent variable model, GPLVM）*9と呼ばれるより洗練されたモデルを得ることができます．
今回のようなベイズ推論と深層学習の融合領域に関しては次の書籍が詳しいです． www.kspub.co.jp
線形回帰，潜在変数モデル，変分推論，ベイズニューラルネットの推論等は下記書籍に基本的な解説があります． www.kspub.co.jp
*1:https://arxiv.org/abs/1312.6114
*2:https://arxiv.org/abs/1401.4082
*3:https://rss.onlinelibrary.wiley.com/doi/abs/10.1111/1467-9868.00196
*4:https://cs.stanford.edu/~pliang/papers/online-naacl2009.pdf
*5:http://www.columbia.edu/~jwp2128/Papers/HoffmanBleiWangPaisley2013.pdf
*6:変分推論における事後分布に対する近似分布の表現手段を表しており，推論モデル（inference model），変分モデル（variational model）とも呼ばれます．要は「凝った近似分布」です．
*7:https://web.stanford.edu/~ngoodman/papers/amortized_inference.pdf
*8:したがって元論文で言われているauto-encoding variationa Bayesはベイズではありません．ただし，full Bayes版はappendixにちゃんと載っています．
*9:http://proceedings.mlr.press/v9/titsias10a/titsias10a.pdf
