Stan codeでみる確率的最適化

#続・確率的変分ベイズ
ついにRStanにも変分ベイズが実装されたようです。
RStanにおける変分ベイズはADVI(Automatic Differentiation Variational Inference)と呼ばれるテクニックで実現されています。
変分ベイズは変分下限を導出するのが非常に面倒な訳ですが、ADVIはこの面倒な操作を自動化してくれます。なのでRStanユーザは導出に頭を悩ませることなく、今まで通りモデルを定義するだけで良いのです。すばらしいと思います。

概要はこちらのブログにまとまっていますし、さらに詳しく知りたい方は元論文を参考にされるのが良いと思います。

さて、この記事では元論文の中でほんの少し取り上げられている確率的最適化のStan実装を扱いたいと思います。これまでのサンプリングと比較したADVIのメリットは速さな訳ですが、元論文ではさらに確率的最適化を適用してスケーリングを達成しています。
確率的最適化は、以前の記事でも紹介したのですが、今回はそれをStan語に落とすとどうなるか？みたいな内容です。

#確率的最適化おさらい
目的関数$f(\lambda)$を逐次更新で最適化するとします。
更新式はこのようになるそうです。

\lambda^{(t)} = \lambda^{(t-1)} - \nu^{(t)} \nabla_{\lambda} f(\lambda^{(t-1)})

確率モデルの最適化の場合、$f(\lambda)$は尤度関数なので尤度の和で表現されることが多いです($f(\lambda) = \sum_{i=1}^n f_i(\lambda)$)。
$n$が大きいと勾配の計算に大きな負荷がかかりますが、$p(i)=1/n$の確率でサンプリングした$f_i$を使って$f(\lambda)$を近似するのが確率的最適化です。

\lambda^{(t)} = \lambda^{(t-1)} - \nu^{(t)} n \nabla_{\lambda} f_i(\lambda^{(t-1)})

#Codeのdiff
確率的最適化無しverと有りverの比較です。元論文のコードそのまんまですが。。。
ちなみに確率モデルは混合正規分布です。
##data block

確率的最適化なし

data {
  int <lower=0> N; // number of data points in entire dataset
  int <lower=0> K; // number of mixture components
  int <lower=0> D; // dimension
  vector[D] y[N]; // observations
  real<lower=0> alpha0; // dirichlet prior
  real<lower=0> mu_sigma0; // means prior
  real<lower=0> sigma_sigma0; // variances prior
}

確率的最適化あり

data {
  real<lower=0> N; // number of data points in entire dataset
  int<lower=0> S_in_minibatch;
  int<lower=0> K; // number of mixture components
  int<lower=0> D; // dimension
  vector[D] y[S_in_minibatch]; // observations
  real<lower=0> alpha0; // dirichlet prior
  real<lower=0> mu_sigma0; // means prior
  real<lower=0> sigma_sigma0; // variances prior
}

データ数Nがrealになっています。これはtransformed dataステップで登場するreal型のSVI_factorという変数の計算に用いる為です。
S_in_minibatchはミニバッチのサイズです。それに伴い渡すデータの宣言もy[N]からy[S_in_minibatch]になっています。実際に渡すデータはN個の点からS_in_minibatch個ランダムにサンプリングして渡します。

##transformed data block

確率的最適化なし

transformed data {
  vector<lower=0>[K] alpha0_vec;
  for(k in 1:K) {
    alpha0_vec[k] <- alpha0;
  }
}

確率的最適化あり

transformed data {
  real SVI_factor;
  vector<lower=0>[K] alpha0_vec;
  for(k in 1:K) {
    alpha0_vec[k] <- alpha0;
  }
  SVI_factor <- N / S_in_minibatch;
}

SVI_factorという変数の計算を行っています。これはS_in_minibatch個のデータを使って計算した勾配をデータN個分の勾配と見せかける為の変数です。

##parameters block
変更点はありません。最適化手法が変わっただけで、モデル自体は変えていないので当然です。

##model block

確率的最適化なし

model {
  // priors
  theta ~ dirichlet(alpha0_vec);
  for(k in 1:K) {
    mu[k] ~ normal(0.0, mu_sigma0);
    sigma[k] ~ lognormal(0.0 , sigma_sigma0);
  }
  // likelihood
  for(n in 1:N) {
    real ps[K];
    for(k in 1:K) {
      ps[k] <- log(theta[k]) + normal_log(y[n], mu[k], sigma[k]);
    }
    increment_log_prob(log_sum_exp(ps));
  }
}

確率的最適化あり

model {
  // priors
  theta ~ dirichlet(alpha0_vec);
  for(k in 1:K) {
    mu[k] ~ normal(0.0, mu_sigma0);
    sigma[k] ~ lognormal(0.0, sigma_sigma0);
  }
  // likelihood
  for(n in 1:S_in_minibatch) {
    real ps[K];
    for(k in 1:K) {
      ps[k] <- log(theta[k]) + normal_log(y[n], mu[k], sigma[k]);
    }
    increment_log_prob(log_sum_exp(ps));
  }
  increment_log_prob(log(SVI_factor));
}

priorに関しても変更点はありません。
尤度の計算は反復回数がNからS_in_minibatchとなって軽くなっていることが確認できます。S_in_minibatch個のデータに対していつも通り対数尤度を計算して更新した後、log(SVI_factor)を足して辻褄を合わせています。

#さいごに
ミニバッチのサイズの決め方が気になります。
精度やパフォーマンスに結構影響する様なのですが、調べた限りでは見つかりませんでした(サーベイ力不足です)。
この論文では3.8M documentsに対して$S \in \{ 10,50,100,500,1000 \}$で実験していました。当然ですが、ミニバッチのサイズは大きい方が精度は大きくなります。