stanで多クラスLogistic Regression

やりたいこと

Stanは素晴らしいベイズ推定のフレームワークです。公式のdocには様々な例が載っていますが、実際にこれを使用して予測をするという内容はほとんど見かけたことがないので試しにやってみました。今回はStanを使って多クラスLogistic Regression(Softmax Regression)を実装してみます。

GitHub repo https://github.com/lucidfrontier45/StanSoftmaxRegression

KクラスのSoftmax Regressionは以下のようなモデルです。

z_k = \mathbf{w}_k^t \mathbf{x} \\
P(C_k|\mathbf{x}, \mathbf{w}) = \frac{e^{z_k}}{\sum_{k'} e^{z_{k'}}}

すなわちKつの重み$\mathbf{w}_k$と入力ベクトル$\mathbf{x}$が与えられたとき、出力がクラスkである確率は$\mathbf{w}_k$と $\mathbf{x}$との内積を取り、それらの指数の値を和で正規化したような形です。

ベイズ推定ではKつの重み$\mathbf{w}_k$に事前分布

P(\mathbf{w}_k)

を仮定し、事後分布

P(\mathbf{w}_k | \mathbf{X})

を何らかの方法で求めます。そしてこの重みの事後分布を用いてクラス確率の予測事後分布を以下のように積分で求めます。

P(C_k|\mathbf{x_{new}}, \mathbf{X}) = \int d\mathbf{w} P(\mathbf{w}_k | \mathbf{X}) P(C_k|\mathbf{x_{new}}, \mathbf{w})

しかし、実際にこの積分を実行することは不可能であり、ラプラス近似、サンプリング、変分ベイズ等の手法を用いる必要があります。
サンプリングを利用した場合は実際に$\mathbf{w}$の事後分布からのサンプル$\mathbf{w}^i$がNつ得られるので

P(C_k|\mathbf{x_{new}}, \mathbf{X}) = \frac{1}{N} \sum_i P(C_k|\mathbf{x_{new}}, \mathbf{w}^i)

のように近似することができます。これをStanとPythonで実際にやってみます。

コード説明

model.stan

data {
  int<lower=1> N;
  int<lower=1> D;
  int<lower=2> K;
  matrix[N, D] X;
  int<lower=1> y[N];
  real<lower=0> s;
}

transformed data {
  row_vector[N] zeros = rep_row_vector(0, N);
}

parameters {
  matrix[K-1, D] w_raw;
}

model {
  matrix[K, N] z = append_row(zeros, w_raw * X');
  for(n in 1:N){
    y[n] ~ categorical_logit(z[,n]);
  }
  to_vector(w_raw) ~ normal(0, s);
}

上記がStanのモデルコードです。公式のdocに書いてあるのとほとんど同じで、softmax関数の不定性に対処するために0番目のクラスの$z_0$の値が常に0になるようにしています。また、事前分布のハイパーパラメータをdataブロックで渡せるようにしています。

これを利用するPythonのコードは次のとおりです。

import pickle
from hashlib import md5
from pathlib import Path

import numpy as np
from pystan import StanModel
from scipy.special import softmax
from sklearn.base import BaseEstimator, ClassifierMixin


def append_const(x, v=0):
    N = len(x)
    return np.column_stack((np.zeros(N)+v, x))


def StanModel_cache(model_code, model_name="anon_model", model_dir="~/.stan"):
    model_dir = Path(model_dir).expanduser()
    if not model_dir.exists():
        model_dir.mkdir()

    code_hash = md5(model_code.encode('ascii')).hexdigest()
    if model_name is None:
        cache_fn = 'cached-model-{}.pkl'.format(code_hash)
    else:
        cache_fn = 'cached-{}-{}.pkl'.format(model_name, code_hash)

    cache_file: Path = model_dir.joinpath(cache_fn)
    if cache_file.exists():
        print("use cached stan model")
        with cache_file.open("rb") as fp:
            sm = pickle.load(fp)
    else:
        print("compile stan model")
        sm = StanModel(model_code=model_code, model_name=model_name)
        with cache_file.open(mode="wb") as fp:
            pickle.dump(sm, fp, pickle.HIGHEST_PROTOCOL)
    return sm


with open("model.stan", "rt") as fp:
    model_code = fp.read()
    stan_model = StanModel_cache(model_code, "softmax_regression")


class StanSoftmaxRegression(BaseEstimator, ClassifierMixin):
    def __init__(self, stan_model: StanModel = stan_model, s=10, **kwds):
        self.stan_model = stan_model
        self.s = s
        self.stan_opts = kwds

    def fit(self, X, y):
        X = append_const(X, 1)
        y = np.asarray(y) + 1
        N, D = X.shape
        K = max(y)
        self.K_ = K
        fit = self.stan_model.sampling(
            data={"N": N, "D": D, "K": K, "X": X, "y": y, "s": self.s},
            **self.stan_opts
        )
        self.w_ = fit["w_raw"]

    def predict_proba(self, X):
        X = append_const(X, 1)
        n_post = len(self.w_)
        z_raw = self.w_.dot(X.T)
        probs = []
        for i in range(n_post):
            z = append_const(z_raw[i].T, 0)
            proba = softmax(z, 1)
            probs.append(proba)

        return np.mean(probs, 0)

    def predict(self, X):
        probs = self.predict_proba(X)
        return probs.argmax(1)

append_const関数はndarrayに前から定数の列を付け足します。線形モデルの切片を扱う際によく用いられる手法です。StanModel_cacheはコンパイル済みのStanのモデルをpickleで保存しておき、モデルコードのハッシュで判別して同じモデルはコンパイル済みのを再利用するようにしています。Stanはモデルのコンパイルにかなり時間がかかる¹ ので開発の際に結構大事です。

StanSoftmaxRegressionがStanを利用したSoftmaxRegressionのsklearn風のクラスで、fitメソッド内でStanのサンプリング関数を呼び、predict_probaでは予測事後分布の近似

P(C_k|\mathbf{x_{new}}, \mathbf{X}) = \frac{1}{N} \sum_i P(C_k|\mathbf{x_{new}}, \mathbf{w}^i)

を実行しています。Pythonの方でもStanのモデルと同様に$z_0$を常に0にするようにappend_constを利用しています。

Irisのデータセットで試してみます。

if __name__ == "__main__":
    from sklearn import datasets, model_selection
    iris = datasets.load_iris()
    X = iris.data
    y = iris.target
    X_train, X_test, y_train, y_test = model_selection.train_test_split(
        X, y, test_size=0.1)
    model = StanSoftmaxRegression()
    model.fit(X_train, y_train)
    probs = model.predict_proba(X_test)
    for y_true, p in zip(y_test, probs):
        print("true class = {}, ".format(y_true),
              "predicted prob = ", p.round(3))

true class = 0,  predicted prob =  [0.992 0.008 0.   ]
true class = 0,  predicted prob =  [0.997 0.003 0.   ]
true class = 1,  predicted prob =  [0.    0.607 0.393]
true class = 1,  predicted prob =  [0.002 0.993 0.004]
true class = 2,  predicted prob =  [0.    0.435 0.565]
true class = 0,  predicted prob =  [0.994 0.006 0.   ]
true class = 0,  predicted prob =  [0.992 0.008 0.   ]
true class = 1,  predicted prob =  [0.001 0.99  0.009]
true class = 2,  predicted prob =  [0.    0.346 0.654]
true class = 1,  predicted prob =  [0.007 0.992 0.001]
true class = 0,  predicted prob =  [0.999 0.001 0.   ]
true class = 0,  predicted prob =  [0.825 0.175 0.   ]
true class = 0,  predicted prob =  [0.992 0.008 0.   ]
true class = 0,  predicted prob =  [0.997 0.003 0.   ]
true class = 0,  predicted prob =  [1. 0. 0.]

Irisのデータセットだとイマイチ効果が実感できないけどもっと小さなデータセットなら通常の最尤推定のSoftmaxRegressionと比べてより汎化能力が高いのかな...?

2019-07-24更新

変分推定対応

Stanはサンプリング以外に変分推定にも対応しています。これによってサンプリングでは対応できない大規模なデータセットに対してもパラメータのベイズ推定を行うことができます。StanSoftmaxRegressionを以下のように修正します。

class StanSoftmaxRegression(BaseEstimator, ClassifierMixin):
    def __init__(self, stan_model: StanModel = stan_model,
                 mode="sampling", s=10, **kwds):
        self.stan_model = stan_model
        self.mode = mode
        self.s = s
        self.stan_opts = kwds

    def fit(self, X, y):
        X = append_const(X, 1)
        y = np.asarray(y) + 1
        N, D = X.shape
        K = max(y)
        self.K_ = K
        data = {"N": N, "D": D, "K": K, "X": X, "y": y, "s": self.s}
        if self.mode == "sampling":
            fit = self.stan_model.sampling(data=data, **self.stan_opts)
            self.w_ = fit["w_raw"]
        elif self.mode == "vb":
            fit = self.stan_model.vb(data=data, **self.stan_opts)
            d = dict(zip(fit["sampler_param_names"], fit["sampler_params"]))
            self.w_ = np.column_stack([
                d["w_raw[{},{}]".format(i, j)]
                for i in range(1, K)
                for j in range(1, D+1)]
            ).reshape(-1, K-1, D)

変更点としてはコンストラクタでmodeという変数も受け取り、fitではこのmodeの応じてsamplingとvbを切り替え、その後の推定した結果を取得する部分も書き分けています。変分推定では本来は通常のパラメトリックな事後分布を関数として得られるはずなんですけど何故かStanではその事後分布からサンプリングした結果を返してきます。(互換性のため？)

sklearnのdigitsのデータセットをsamplingの方で解こうとするといつまで経っても終わりませんでしたが、これなら数十秒で完了しました。

true class = 4,  predicted prob =  [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]
true class = 2,  predicted prob =  [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]
true class = 7,  predicted prob =  [0. 0. 0. 0. 0. 0. 0. 1. 0. 0.]
true class = 1,  predicted prob =  [0.    0.787 0.    0.002 0.038 0.    0.    0.    0.002 0.17 ]
true class = 2,  predicted prob =  [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]
true class = 0,  predicted prob =  [1. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
true class = 4,  predicted prob =  [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]
true class = 5,  predicted prob =  [0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]
....

1. Eigenを利用したC++のバイナリをコンパイルするのでかなり遅い。。。 ↩