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分のパラメータを複製するなどのアイデアがありえます。
R　package https://cran.r-project.org/web/packages/KSD/
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
https://github.com/xiangze/SVGD

Reference

Stein’s Method for Practical Machine Learning ここに多くの論文がまとめられている
Probabilistic Learning and Inference Using Stein Discrepancy　概要スライド
Stein's Method
機械学習論文読みメモ_14
http://chunyuan.li/
論文
Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm arxiv
LEARNING TO DRAW SAMPLES: WITH APPLICATION TO AMORTIZED MLE FOR GENERATIVE ADVERSARIAL LEARNING pdf
A Kernelized Stein Discrepancy for Goodness-of-fit Tests and Model Evaluationarxiv