Stein discrepancyと最適化(SVGD)
手法
目的
KL divergenceの最小化
$\phi^*=\arg\min [ \partial_\epsilon KL(q_\epsilon // p) ]$
を行いたい。
Stein Operator $A_p$（後述)を使って勾配を
$\nabla_\epsilon KL(q_\epsilon//p)|_{\epsilon=0} $
$= -E_{x \sim q}[A_p \phi(x)]$
$ (x' = x + \phi(x) )$
としてこれを用いて計算する。
Stein Operator
Stein operatorを
$A_p \equiv f(x)\nabla_x \log p(x)+ \nabla_x f(x)$
と定義する以下のように導出される。
分布p(x)に対して
$\lim_{|x| <- \infty} f(x)p(x)=0$
となるようなf(x)に対して恒等式(Stein identity)
$E_{x \sim p}[ f(x)\nabla_x \log p(x)+ \nabla_x f(x) ]=0$
が成り立つ(部分積分から)。分布p(x)と近いq(x)に対してこれを
$E_{x\sim q}[\nabla_x \log p(x)+ \nabla_x f(x)] \equiv E_{x\sim q}[A_p f(x)]$
$A_p f(x) \equiv \nabla_x \log p(x)+ \nabla_x f(x)$
と定義する(記号もあいまって共変微分のように見える)。また
$E_{x \sim q}[A_p f(x)]=E_{x \sim q}[A_p f(x)]-E_{x \sim q}[A_q f(x)] =E_{x \sim q}[ f(x)(\nabla \log p(x) -\nabla \log q(x))]$
という関係式も成り立つ。
Stein disrepancy
Stein Operator$A _p$を用いて
$\sqrt{S(p,q)} \equiv \max_{f \in F} E_{x\sim q} [A_p f(x)]$
とStein disrepancyを定義する。
これをどのように計算するのか
関数f(x)を
$f(x) = \sum_i \omega_i f_i(x)$
と展開すると
$E_q[A_p f(x) ]= E_q[A_p \sum_i \omega_i f_i(x)]= \sum \omega_i \beta_i$
$\beta_i \equiv E_{x\sim q}[A_p f_i(x)]$
さらにカーネルを使って
$f(x)=\sum_i \omega_i k(x,x_i)$
とかくと$ H =\sum{i,j} w_i v_j k(x_i,x_j)$という形に書けることから
$E_{x \sim q}[A_p f(x)]=$
添字iに関する和を有限個とすることができて(基底定理？)
$A_p f(x)=\frac{1}{n}\sum_{i=0}^n [ f(x_i) k(x_i,x) \nabla_x \log p(x_i)+ \nabla_{x_i}k(x_i,x) ]$
となる。
結局$\phi^*=\arg\min [ \partial_\epsilon KL(q_\epsilon // p) ]$に対しては
$\phi$がxの関数であるとしてparticle iに対し
$x_i^{l+1} \leftarrow x_i^l + \epsilon_l \phi^*(x_i^l)$
$\phi^*(x)=\frac{1}{n}\sum_{j=0}^n [ k(x_j^l,x) \nabla_x \log p(x_j^l)+ \nabla_{x_j}k(x_j^l,x) ]$
という更新式を使う。$k(x_i,x_j)$はカーネルで行列を使って表すことができる(未知の値$k(x_i,x)$に対しては関数として実装する)。
サンプリングの方法として(GANへの応用)
論文 https://arxiv.org/abs/1611.01722
Amortized SVGD
乱数$\xi_i$に対し
$x_i=f_n(\xi_i) $
であるとする更新式は
$ \Delta x_i =E_{ x\sim {x_i} } [\nabla p(x) k(x,x_i) + \nabla_x k(x,x_i)]$
$x_i'=x_i+\epsilon \Delta x_i$
関数$f(\eta;\xi_i)$のパラメータ(として)$\eta$があるとすると
$\eta=\eta+\epsilon \sum_i \partial_\eta f_\eta (\xi_i)\Delta x_i$
としてfを更新していくのがSVGD
特にKL divergence$KL(q_\eta//p)$に関して書き直すと
$\nabla_\eta KL(q_\eta//p)=-E_{\xi \sim q_0}[\partial_\eta f(\eta ;\xi)(\nabla_x \log p(x) - \nabla_x \log q_\eta(x) ] $
となる。
さらに$\bar{\Delta}x_i= \nabla_x \log p(x_i) - \nabla_x \log q_\eta(x_i) $
と書くと
$\Delta x_i \simeq E_{x \sim q}[k(x,x_i) \nabla_x \log p(x)+ \nabla_{x}k(x,x_i) ] $
$= E_{x \sim q}[k(x,x_i)( \nabla_x \log p(x)+ \nabla_{x} \log q(x)] $
$= E_{x \sim q}[\bar{\Delta}x_i k(x,x_i]$
という近似的関係が成り立つ($\Delta x_i$は$\bar{\Delta} x_i$をカーネル重み付けしたもの
Amortized KSD
$\kappa_p(x,x')=\nabla_x \log p(x) \nabla_x \log p(x') k(x,x') +\nabla_x \log p(x) \nabla_x k(x,x') + \nabla_x \log p(x') \nabla_x k(x,x')+
\nabla_x \nabla_x' k(x,x')$
これによって
$D^2(q//p)=E_{x,x'\sim q}[\kappa_p(x,x')]$
の最小化を行う。
Amortized MLE
$p(x|\theta)=\exp(-\phi(x,\theta)-\Phi(\theta)$
という形の関数の最適化
$\xi_i \sim q_{\xi}, x_i=f_\eta(\xi_i)$
として
$\eta \leftarrow \epsilon \sum_i \partial_\eta f_\eta (\xi_i) \Delta x_i$
${x_{obs}}$ を新たに引く
$\theta \leftarrow \theta -E_{obs}[\nabla_\theta \phi(x,\theta)]+E_\eta[\nabla_\theta \phi(x,\theta)]$
を交互に繰り返す。
GANの場合は
$p(x|\theta) \propto \exp(-||x-D(E(x;\theta);\theta)||)$
実装
オリジナル実装はTheanoで書かれている。
- https://github.com/DartML/Stein-Variational-Gradient-Descent/tree/master/python
- https://github.com/DartML/Stein-Variational-Gradient-Descent/tree/master/python
SteinGAN もベタコード
- https://github.com/DartML/Stein-Variational-Gradient-Descent/blob/master/python/bayesian_logistic_regression.py
- https://github.com/DartML/SteinGAN/tree/master/lib
- https://github.com/ChunyuanLI/SVGD/blob/master/demo_svgd.ipynb
Edward, tensorflowのoptimizerとして組み込めるかもしれません。確率分布pとその微分を使うのでEdwardやtf.distributionsを使う必要がありそうです。particle分のパラメータを複製するなどのアイデアがありえます。
-
Julia 実装
- https://github.com/krisrs1128/readings/tree/master/svgd
- https://github.com/jgorham/stein_discrepancy/tree/master/src/experiments kernelは決め打ちです。自動微分を使えば任意のkernelのgradientが計算できるかと思ったのですが、Juliaのautodiffは偏微分をやるのが難しそうだったので断念しました
-
Tensorflow