Pyroで遊んでみた【ベイズ多項式回帰】

#概要
手法やライブラリを思いつくままに色々試した内容になっています.
内容と結果は以下の通りです
・Pyroでベイズ推定：Pyroサンプルコード、実行速度評価
　→MCMCはかなり遅い、stanやnumpyroを使用した方が良いかも. 変分ベイズならPyro
・ベイズ多項式回帰
　→特になし。なんだかんだ初めて触った
・モデル選択: 周辺化尤度
　→それっぼく選択できた。その他モデル選択の指標との関連を整理したい.
・モンテカルロ法:ナイーブ・モンテカルロ法、重点サンプリング
　→ナイーブはやっぱりだめだめ、重点サンプリングは一応良さげ
・変分ベイズ:平均場近似、多変量正規分布近似
　→評価中、、、。単純な平均場近似だけではなく、多変量正規分布など近似分布にバリエーションがあるのはPyroの良いとこ.

流れとしては適当に発生させたデータに対してベイズ多項式回帰を行いモデル選択し、データを説明する適切なモデルを構築します.

#ライブラリ
今回はPyroを使ってベイズ推定を行いました.
ちなみに、ライブラリインストールするときに"pyro"だと違うライブラリが指定されます. "pyro-ppl"とする必要があります（ややこしい.
Pytorchをベースとした確率的プログラミング言語で、最近良く聞くので感触確かめたくて使ってます.
Pyro公式HP: https://pyro.ai/

後はおなじみのsklearnやらnumpyやら使ってます.

import logging
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import math
import torch
from torch.distributions import constraints
from sklearn import preprocessing as pc
from sklearn.preprocessing import PolynomialFeatures as PF

# pyro周り
import pyro
import pyro.distributions as dist
import pyro.optim as optim
from pyro.infer.mcmc import MCMC, NUTS, HMC
from pyro.infer import EmpiricalMarginal, TracePredictive, Importance
from pyro.infer.autoguide import AutoMultivariateNormal, init_to_mean

# GPU周りの設定
torch.set_default_tensor_type(torch.cuda.DoubleTensor)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

#関数
以下使用した関数です。
t2n(): pyroのデータをGPUで計算できる形式からnumpyへ変換する関数。pytorch系ではこの変換が地味にめんどくさい
MarginalLikelihood(): ベイズ多項式回帰の周辺化尤度を解析的に求める関数
NaiveBayes(): 周辺化尤度をナイーブベイズで求める関数

def t2n(torchgpu):
    assert torchgpu.is_cuda
    return torchgpu.detach().cpu().numpy()

# marginal likelihood
# P(D|M_i)
def MarginalLikelihood(X, Y):
    X = t2n(X)
    Y = t2n(Y)
    N = X.shape[0]
    m = np.repeat(0, X.shape[1]).reshape(-1, 1)
    lam = 1/(std**2)
    S = np.linalg.inv(np.diag(np.repeat(10, X.shape[1])))
    S_hat = lam*sum([X[i,:].reshape(-1, 1) @ X[i,:].reshape(1, -1) for i in range(N)])+S
    m_hat = lam*np.linalg.inv(S_hat) @ sum([Y[i]*X[i,:] for i in range(N)]).reshape(-1, 1)
    ln_prob = -0.5*(sum([lam*Y[i]**2-np.log(lam)+np.log(2*math.pi) for i in range(N)])+
                    m.T @ S @ m -np.log(np.linalg.det(S)) -m_hat.T @ S_hat @ m_hat +np.log(np.linalg.det(S_hat)))
    return ln_prob[0,0]

def NaiveBayes(x_data, y_data, N = 1000):
    y = 0
    for _ in range(N):
        # prior
        loc = torch.zeros(x_data.shape[1])
        scale = torch.ones(x_data.shape[1])*10
        W = pyro.sample("weight", dist.Normal(loc, scale))

        # likelihood
        with pyro.plate("obs", x.shape[0]):
            mu = x_data @ W
            y += dist.Normal(mu.reshape(-1, 1), std).log_prob(y_data)
    return y.mean()/N

データ

今回使用したデータ.
適当な曲線に正規分布のノイズをのせています.

# Data generation
np.random.seed(100)
def func(x, std = 0):
    return np.sin(x) + np.cos(x/2+1) + np.random.normal(0, std, size = x.size).reshape(-1, 1)

N = 8
std = 0.2
x = (np.random.rand(N)-0.5).reshape(-1, 1) * 3 * math.pi
y = func(x, std=0.2)
lin = np.linspace(-4, 4).reshape(-1, 1)
plt.scatter(x, y)
plt.plot(lin, func(lin))

# 多項式回帰の計画行列作成
Degree = 10
pf = PF(degree=Degree, include_bias=True)
X = pf.fit_transform(x)

x_data, y_data = torch.Tensor(X.reshape(-1, Degree+1)), torch.Tensor(y.reshape(-1, 1))

パラメータ推定

モデル作成

推定に使用するモデルを関数として定義します.
今回は回帰係数のみ推定の対象としています.
細かい文法は公式doc参照. https://pyro.ai/examples/intro_part_i.html

# model
def model(x_data, y_data):
    # prior
    loc = torch.zeros(x_data.shape[1])
    scale = torch.ones(x_data.shape[1])*10
    # 係数に事前分布として正規分布を設定
    W = pyro.sample("weight", dist.Normal(loc, scale))

    # likelihood
    with pyro.plate("obs", x.shape[0]):
        # Pytorchっぽく行列の積できる
        mu = x_data @ W
        # obsでデータを指定
        y = pyro.sample(
            "measure", 
            dist.Normal(mu.reshape(-1, 1), std),
            obs = y_data)
    return y

MCMC

MCMCの実行コード.
次数１から１０まで実行したが、実行時間は合計で 1h15m...
iter=600でこの速度は正直遅い、stanの数十分の１以下な印象... Pyroェ...

と思ったらMCMC用で別のライブラリがある模様.
NumPyro: https://github.com/pyro-ppl/numpyro
こちらはstanと同等か、それ以上早い場合もあるそうなので別途要検証

%%time
mcmc_dic = {}
i = 0
nuts_kernel = NUTS(model, adapt_step_size=True, jit_compile=True)
mcmc_dic["deg" + str(i+1)] = MCMC(
    nuts_kernel, 
    num_samples=300, 
    warmup_steps=300)
mcmc_dic["deg" + str(i+1)].run(x_data[:,:(i+2)], y_data)

for i in range(1, Degree):
    sam = mcmc_dic["deg" + str(i)].get_samples()
    nuts_kernel = NUTS(model, adapt_step_size=True, jit_compile=True)
    mcmc_dic["deg" + str(i+1)] = MCMC(
        nuts_kernel, 
        num_samples=300, 
        warmup_steps=300,
        initial_params={"weight":torch.cat((sam["weight"][299,:], torch.tensor([0.])), dim = 0)})
    mcmc_dic["deg" + str(i+1)].run(x_data[:,:(i+2)], y_data)

予測分布

求めたサンプルを元に標本路を生成.
先程定義したmodel関数がそのまま使えるので、ここはstanと違って嬉しい.
stanだと予測分布作りたかったら最初から指定しないといけないので.

# 標本路生成
pred_sample_y = {}
x_index = torch.linspace(-5, 5).reshape(-1, 1)
X_index = torch.tensor(pf.transform(t2n(x_index)))

for i in range(Degree):
    sample = mcmc_dic["deg" + str(i+1)].get_samples()
    pred_y = pyro.infer.Predictive(model=model,
                            posterior_samples = sample,
                            return_sites=["measure"])
    with torch.no_grad():
        pred_sample_y["deg" + str(i+1)] =pred_y.forward(X_index[:,:(i+2)], None)

# plot
for i in range(Degree):
    plt.figure()
    plt.title("Degree"+str(i+1))
    plt.scatter(x, y, color = "black")
    plt.plot(lin, func(lin))
    plt.ylim(-2, 2)
    for j in range(20):
        plt.plot(t2n(x_index), 
                t2n(pred_sample_y["deg" + str(i+1)]["measure"][j,:,1]),
                alpha = 0.1, color = "blue")

以下、予測分布.
次数1-6で次数が増えるに連れてデータに合っていき、真の関数に近づいていっている様子が確認できます.
次数4か5くらいが丁度良さそうな気がしますね（多分
次数7-10は似たような結果になったので省略して次数10の場合だけ載せているのですが、
これは予想と違って大人しい結果になりました.本当はもっとバタバタ暴れた分布になる予定だったのですが、多分MCMCのiterけちったので局所的な挙動しか見れてないのだと思います.

モデル選択

周辺化尤度

上の結果からなんとなく今回のデータを説明するには次数４，５あたり良さそうな感じがしましたが、
これを定量的に評価するため各次数における周辺化尤度を求めます.
周辺化尤度っていうのはよく出てくるベイズの定理の分母の部分に相当します（というと語弊がありますが、そんな感じです.
普段は無視しがちな部分ですが、ここの値が大きいほど与えられたデータを説明する良いモデルが選択できるらしいです.

周辺化尤度を使ったモデル選択についてはこちらを参照. https://sci-tech.ksc.kwansei.ac.jp/~tohhiro/bioinformatics19/bioinfo-jags8.pdf
個人的にはモデルというパラメータに対するMAP推定みたいなイメージもってます.

ベイズ多項式回帰の場合には周辺化尤度の解析式が得られているため、これを使って厳密解を求めます.
ただ、一般に厳密解が与えらえる場合は多くなくそのような場合はモンテカルロ法による近似解を使用します.
今回は厳密解に加えて最もシンプルな方法であるナイーブ・モンテカルロ法と重点サンプリングによる近似も実施しました.
重点サンプリング: http://larrymei.club/2019/05/13/Importance-sampling/

# 重点サンプリングに使用する提案分布の定義
def is_guide(x_data, y_data):
    with pyro.plate("w_plate", x_data.shape[1]):
        W = pyro.sample("weight", dist.Normal(0, 0.1))

Evi = []
for i in range(Degree):
    Evi.append(MarginalLikelihood(x_data[:,:(i+2)], y_data))

EviNaive = []
for i in range(Degree):
    posterior = Importance(model, num_samples=5000)
    posterior.run(x_data[:,:(i+2)], y_data)
    EviNaive.append(posterior.get_log_normalizer())

EviIS = []
for i in range(Degree):
    posterior = Importance(model, guide = is_guide, num_samples=5000)
    posterior.run(x_data[:,:(i+2)], y_data)
    EviIS.append(posterior.get_log_normalizer())

結果

以下、上から厳密解、近似解（ナイーブ・モンテカルロ）、近似解（重点サンプリング）です.
・厳密解：次数３，４が良いという結果となりまあ良さげです. ただ、これだけ拮抗していたら一つに選択するのではなくフルベイズ的にいくのが良さそうな気がしますが.
あと、次数５以上においては線形に周辺可尤度が減少していてAIC味を感じますね. 何か関係ありそうなので後で確認したいです.
・近似解（ナイーブ）：とにかく値が小さな方へ発散しており、だめだめです.
・近似解（重点）：絶対値は異なりますが、次数３，４あたりが良いという傾向は再現できています.
提案分布工夫したらもうちょっと良くなりそう

実際にモンテカルロ法で周辺化尤度を求めようと思ったら重点サンプリングを発展させたブリッジサンプリングという方法が良いらしいです.
これも要確認ですね.
https://www.slideshare.net/daikihojo/bridgesampling

plt.bar(range(1,1+len(Evi)), Evi)
plt.title("Exact Evidence")
plt.show()

plt.bar(range(1,1+len(Evi)), EviNaive)
plt.ylim(-10000, -100)
plt.title("Naive Evidence")
plt.show()

plt.bar(range(1,1+len(Evi)), EviIS)
plt.ylim(-1000, -100)
plt.title("Importance Sampling Evidence")
plt.show()

変分ベイズ

パラメータ推定からモデル選択までやりましたが、最後に変分ベイズまで試します.
Pyroの変分ベイズは平均場近似だけではなく、分布の近似として多変量正規分布が使えたりとバリエーションが豊富です.
ここは明らかにstanより優れている点だと思います.

以下、平均場近似と多変量正規分布による近似のコードです.
実際に次数１０の場合において計算したところ、実行時間は 約15s と当たり前ですが、MCMCとは比べ物にならないほど早いです.
ここらへんで力尽きたので結果の詳細はまた別途まとめます.....

#独立モデル
def svi_guide(x_data, y_data):
    # prior
    N = x_data.shape[1]
    mu = pyro.param("mu", torch.tensor(0.).expand(N))
    sigma = pyro.param("sigma", torch.tensor(1.).expand(N),
                        constraint=constraints.positive)
    pyro.sample("weight", dist.Normal(mu, sigma))

#多変量正規分布
multi_guide = AutoMultivariateNormal(model, init_loc_fn=init_to_mean)

# set up the optimizer
adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)}
optimizer = pyro.optim.Adam(adam_params)

# setup the inference algorithm
# guideで近似分布を指定
# ここでは多変量正規分布を使用
svi = pyro.infer.SVI(model=model,
                     guide=multi_guide,
                     optim=optimizer,
                     loss=pyro.infer.Trace_ELBO())

%%time
n_steps = 5000
# do gradient steps
pyro.clear_param_store()
losses = np.zeros(n_steps)
for step in range(n_steps):
    losses[step] = svi.step(x_data, y_data)

plt.plot(losses)
plt.title("KL divergence")