GLM:階層回帰モデルのミニバッチADVI

ガウス混合モデルと異なり、(階層)回帰モデルは、独立した変数です。これらの変数は尤度関数に影響しますが、ランダム変数ではありません。ミニバッチを使う場合、それに注意すべきです。

%env PYTENSOR_FLAGS=device=cpu, floatX=float32, warn_float64=ignore

import os

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt
import seaborn as sns

from scipy import stats

print(f"Running on PyMC v{pm.__version__}")

env: PYTENSOR_FLAGS=device=cpu, floatX=float32, warn_float64=ignore
Running on PyMC v5.0.1

%config InlineBackend.figure_format = 'retina'
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

try:
    data = pd.read_csv(os.path.join("..", "data", "radon.csv"))
except FileNotFoundError:
    data = pd.read_csv(pm.get_data("radon.csv"))

data

919 rows × 30 columns

county_idx = data["county_code"].values
floor_idx = data["floor"].values
log_radon_idx = data["log_radon"].values

coords = {"counties": data.county.unique()}

　ここで、log_randon_idx_tは独立した変数です。floor_idx_tとcounty_idx_tは独立変数を決定します。

log_radon_idx_t = pm.Minibatch(log_radon_idx, batch_size=100)
floor_idx_t = pm.Minibatch(floor_idx, batch_size=100)
county_idx_t = pm.Minibatch(county_idx, batch_size=100)

with pm.Model(coords=coords) as hierarchical_model:
    # Hyperpriors for group nodes
    mu_a = pm.Normal("mu_alpha", mu=0.0, sigma=100**2)
    sigma_a = pm.Uniform("sigma_alpha", lower=0, upper=100)
    mu_b = pm.Normal("mu_beta", mu=0.0, sigma=100**2)
    sigma_b = pm.Uniform("sigma_beta", lower=0, upper=100)

　各州の区間、グループ平均mu_aの周辺に分布します。上記で私たちがmuとsdを固定値に設定します。すべてのaとbのために、共通のグループ分布をプラグインします。（それは私たちの例では、ユニークな州の数として同じ長さのベクターになります）

with hierarchical_model:

    a = pm.Normal("alpha", mu=mu_a, sigma=sigma_a, dims="counties")
    # Intercept for each county, distributed around group mean mu_a
    b = pm.Normal("beta", mu=mu_b, sigma=sigma_b, dims="counties")

　radonレベルa[county_index]のモデル予測は、a[0,0,0,1,1,..]に変換されます。私たちはこのように、多重の州の住宅保有尺度を係数にリンクします。

with hierarchical_model:

    radon_est = a[county_idx_t] + b[county_idx_t] * floor_idx_t

　最後に、私たちは、尤度を規定します。

with hierarchical_model:

    # Model error
    eps = pm.Uniform("eps", lower=0, upper=100)

    # Data likelihood
    radon_like = pm.Normal(
        "radon_like", mu=radon_est, sigma=eps, observed=log_radon_idx_t, total_size=len(data)
    )

　Random変数、randon_like、log_randon_idx_tに関連した、尤度項の観測値として示すために、ADVIへの関数で与えられます。

　一方で、minibatchesは上の三つの変数を含む必要があります。

　その後、ADVIをミニバッチで実行します。

with hierarchical_model:
    approx = pm.fit(100_000, callbacks=[pm.callbacks.CheckParametersConvergence(tolerance=1e-4)])

idata_advi = approx.sample(500)

　ELBOのトレースをチェックし、MCMCの結果と比較してください。

plt.plot(approx.hist);

　私たちは、平均領域の近似から共分散行列を抽出することができます。そして、NUTSアルゴリズムのために、スケーリング行列として使います。

scaling = approx.cov.eval()

　そしてまた、私たちは、サンプラーの開始ポイントとして使うためにサンプル(各チェインに一つ)を生成します。

n_chains = 4
sample = approx.sample(return_inferencedata=False, size=n_chains)
start_dict = list(sample[i] for i in range(n_chains))

# Inference button (TM)!
with pm.Model(coords=coords):

    mu_a = pm.Normal("mu_alpha", mu=0.0, sigma=100**2)
    sigma_a = pm.Uniform("sigma_alpha", lower=0, upper=100)
    mu_b = pm.Normal("mu_beta", mu=0.0, sigma=100**2)
    sigma_b = pm.Uniform("sigma_beta", lower=0, upper=100)

    a = pm.Normal("alpha", mu=mu_a, sigma=sigma_a, dims="counties")
    b = pm.Normal("beta", mu=mu_b, sigma=sigma_b, dims="counties")

    # Model error
    eps = pm.Uniform("eps", lower=0, upper=100)

    radon_est = a[county_idx] + b[county_idx] * floor_idx

    radon_like = pm.Normal("radon_like", mu=radon_est, sigma=eps, observed=log_radon_idx)

    # essentially, this is what init='advi' does
    step = pm.NUTS(scaling=scaling, is_cov=True)
    hierarchical_trace = pm.sample(draws=2000, step=step, chains=n_chains, initvals=start_dict)

az.plot_density(
    data=[idata_advi, hierarchical_trace],
    var_names=["~alpha", "~beta"],
    data_labels=["ADVI", "NUTS"],
);