用pytorch實現stan的概率推斷

stan，pymc與edward都是基於一套自動求導機制（stan是自己實現的，pymc基於theano，edward利用tensorflow）在上面提供一些封裝來簡化概率推斷的描述。但它們同時也不那麼靈活，pymc和edward經常要和它們背後的theano,tensorflow的機制交互，而這兩個框架都不好調試。本文利用pytorch實現stan風格的幾個推斷方法來揭示它們內部做的類似的事情。

stan的模型簡介

stan的模型可以說最精簡的，如pystan例子的這段代碼

data { int<lower=0> J; // number of schools real y[J]; // estimated treatment effects real<lower=0> sigma[J]; // s.e. of effect estimates } parameters { real mu; real<lower=0> tau; real eta[J]; } transformed parameters { real theta[J]; for (j in 1:J) theta[j] = mu + tau * eta[j]; } model { eta ~ normal(0, 1); y ~ normal(theta, sigma); }

可以說是十分直觀的。

stan的data段表示外部程序向其導入的數據，parameters和transformed parameters里聲明和計算的變數會被記錄下來，model則在用前面定義的參數和數據描述一個聯合概率函數。

其中類似

X ~ normal(mu,sigma)

之類的代碼都等價於向聯合對數概率target累加概率

target += normal_lp(X | mu,sigma)

由於這個target可以任意編寫，比pymc和edward的寫法要靈活得多。

X ~ normal(mu,sigma)這種所謂的採樣語句在任何時候都不對應著真正的隨機採樣，在MAP——極大後驗點估計中我們只是在極大化target，在ADVI自動微分變分推斷中，使用蒙特卡洛積分計算一個期望時，我們是在變分分布里隨機採樣而不是這裡的原分布，在HMC中，也不是直接用這個分布進行採樣。

事實上，為了在stan里進行真正的採樣，要調用normal_rng之類的函數，上述的sampling statement為了避免歧義被刻意限制為只能在model塊中使用，而XXX_rng函數則必須不能在model塊里使用。

stan里的模型只是花式寫一個聯合概率函數的本質使得這類框架變得十分清晰。pytorch也提供了一套自動求導機制，所以我們可以很容易平移它們的實現。

MAP點估計

下面用pystan實現一個簡單的正態模型：

import pystan import numpy as np import torch from torch.autograd import Variable

ocode = """
data {
int<lower=1> N;
real y[N];
}
parameters {
real mu;
}
model {
y ~ normal(mu, 1);
}
"""

sm = pystan.StanModel(model_code=ocode)
np.random.seed(13)
y2 = np.random.normal(size=20)
np.mean(y2)
# 0.36640264498852165

op = sm.optimizing(data=dict(y=y2, N=len(y2)))
op
#OrderedDict([(mu, array(0.36640264))])

下面用pytorch實現基本等效的代碼（當然stan並不是直接用SGD優化的）

def normal_lp(x,mu,sigma): #雖然在stan里這個函數是包含常數的，但這裡把它棄掉 return -((x - mu)**2/(2*sigma**2)).sum()

y = Variable(torch.from_numpy(y2).float())
mu = Variable(torch.zeros(1),requires_grad=True)

def model(y):
target = Variable(torch.zeros(1))
target += normal_lp(y,mu,1)
return target

optimizer = torch.optim.SGD([mu],lr=0.01)
for epoch in range(100):
optimizer.zero_grad()
y = Variable(torch.from_numpy(y2).float())
target = model(y)
loss = -target
loss.backward()
optimizer.step()

Variable containing:
0.3664
[torch.FloatTensor of size 1]

自動微分變分推斷（AVDI）

這裡的代碼參考的Fully Automatic Variational Inference of Differentiable Probability Models和Automatic Differentiation Variational Inference的描述，stan的對應源碼參考github.

基本上AVDI就是把有約束變數變成無約束變數，用正態分布表示這個無約束變數，然後優化後者的參數。