HELLO CYBERNETICS

深層学習、機械学習、強化学習、信号処理、制御工学、量子計算などをテーマに扱っていきます

【凄く楽しいぞ!Stan!】給料モデリング【Python】

 

 

follow us in feedly

f:id:s0sem0y:20181113234135p:plain

はじめに

とあるように感動したので、怒られない程度にアヒル本の中に出てくるモデルを使ってPyStanで実行してみました。 (TF-PはEdwardだったときは割と分かりやすいAPIだった気がするのですが…PyMC4で持ち直すのだろうか)

データを見てみましょう

まずはインポート

import numpy as np
import matplotlib.pyplot as plt
import seaborn
import pystan
import pandas as pd
from scipy.stats import mstats
%matplotlib inline

下記の関数でデータを作ります

N = 200
K = 4

def get_data():
    
    a0 = 350.
    b0 = 20.
    s_a = 40.
    s_b = 5.
    s_Y = 30.

    
    a = np.random.normal(loc=a0, scale=s_a, size=(K,))
    b = np.random.normal(loc=b0, scale=s_b, size=(K,))
    
    KID = []
    X = []
    Y = []
    
    for n in range(N):
        kid = np.random.randint(0, K, 1)
        KID.append(int(kid))
        
        x = np.random.randint(22, 45, 1)
        X.append(x)
        
        Y.append(np.random.normal(loc=a[kid]+b[kid]*(x-22), scale=s_Y))

    return (np.array(X).reshape(N, 1).astype(np.float32), 
            np.array(Y).reshape(N, 1), 
            np.array(KID).reshape(N, 1).astype(np.int32))

これは下記の数式に基づいています。

$$ \begin{align} KID &= {1, ..., 4} \\ n & = {1, ..., 200} \\ a[k] &\sim \mathcal N(10, 15) \\ b[k] &\sim \mathcal N(20, 3) \\ y[n] & \sim \mathcal N (a[KID[n]] + b[KID[n]] * x) \end{align} $$

$KID$ は会社を識別するID

$n$ は労働者を識別するID

( $KID[n]$ は 労働者 $n$ の所属する会社IDを表す)

$a[k]$ は会社 $k$ の基本的な給与を表す確率変数

$b[k]$ は会社 $k$ の昇給幅を表す確率変数

$ x$ は年齢

データを可視化

要するに、異なる4つの会社から、合計200 人の年収を1年毎に調べ上げたデータというわけです。兎にも角にも、データをサンプリングしてみると下記のようになりました。

X_data, Y_data, KID_data = get_data()
plt.scatter(X_data, Y_data, c=KID_data, cmap='tab10')

f:id:s0sem0y:20181113232205p:plain

df = pd.DataFrame(np.hstack([X_data, Y_data, KID_data]), 
                  columns=['age', 'salary', 'KID'])

直線フィッティング

データの素性を知らない場合

このデータが4つの会社から抽出されていることを知らない場合には、下記のようにデータが見えます。

f:id:s0sem0y:20181113232618p:plain

このデータに対して無情報事前分布を引いたガウス分布を用いたモデルを使いましょう。

$$ \begin{align} n & = 1, 2, ..., N \\ z &= a + b(x - 22) \\ y &\sim \mathcal N(z, \sigma_y^2) \end{align} $$

複雑に見えるかもしれないが、概ね書きように読めます。

・新卒時年収は大体$a$には無情報事前分布を使う。

・1年毎の昇給は大体$b$にも無情報事前分布を用いる。

・合計 $N$ 人を選出している。

・$x$ 歳の基本年収 $z$は新卒年収 $a$ から毎年 $b$ だけ上がる

・個人、年ごとの出来栄えで実際の年収 $y[n]$ は $\sigma_y^2$だけバラつく。

Stanでモデルを書く

上記のモデルは以下のように書くことで実現できます。 Stanそのものの解説はココではしません。

stan_model = """
  data {
    int N;
    real X[N];
    real Y[N];
    int N_s;
    real X_s[N_s];
  }
  parameters {
    real a;
    real b;
    real<lower=0> sigma; 
  }
  model{
    for (n in 1:N){
      Y[n] ~ normal(a + b * (X[n] - 22), sigma);
    }
  }
  generated quantities {
    real Y_s[N_s];
    for (n in 1:N_s){
      Y_s[n] = normal_rng(a + b * (X_s[n] - 22), sigma);
    }
  }
"""

下記のコードでコンパイルができます(StanはC++実装です)。

sm = pystan.StanModel(model_code=stan_model)

データを辞書で渡す

X_s = np.arange(22, 60, 1)
N_s = X_s.shape[0]
stan_data = {"N":df.shape[0], "X":df["age"], "Y":df["salary"], "N_s":N_s, "X_s":X_s}

$X_s$ は 推論済モデルを利用して予測を行う歳の年齢の範囲です。 Stanモデルを書いたときの data で宣言したものを辞書でPythonからわたしてやります。TensorFlowのtf.placeholdertf.Session(feed_dict)の関係に近いですね。

推論(学習)開始

恐るべしStan下記の一行で学習を開始できます。利用するアルゴリズムはNUTS(オート変分ベイズも選べるらしい)です。

fit = sm.sampling(data=stan_data, iter=2000, warmup=500, chains=3, seed=1992)

事後分布

例えば新卒初任給に相当するパラメータ $a$ の事後分布は下記で取り出せます。

ms_a = fit.extract("a")["a"]
plt.hist(ms_a)

f:id:s0sem0y:20181113233530p:plain

うーん、割と日本の現実を表していそうだ…(まあ今回は人工データを作る関数の調整次第なのだが)。

ベイズ予測分布

下記でベイズ予測分布の95%ベイズ信頼区間(95%ベイズ予測区間とも言う)を表示します。

Y_p = fit.extract("Y_s")["Y_s"]
low_y, high_y = mstats.mquantiles(Y_p, [0.025, 0.975], axis=0)

plt.scatter(df["age"], df["salary"])
plt.fill_between(X_s, low_y, high_y, alpha=0.3, color="gray")
a_ = 319.61
b_ = 15.56
x_ = np.arange(22, 60, 1) 
y_ = a_ + b_ * (x_ - 22)
plt.plot(x_, y_, c='r')

a_ = 319.61は事後分布でのパラメータ $a$ の平均です。(b も同様。 fit オブジェクトに入っています)

f:id:s0sem0y:20181113234135p:plain

会社毎のバラツキを考慮したモデリング

階層ベイズモデリング

よく調べると、データは4つの会社からランダムに従業員が選ばれていたことが分かった(すなわち、本来のデータの生成過程に気づいた)。 このような場合には下記のようにモデリングをしてみましょう。

・ある会社$k$の新卒時年収は大体$a_{mean}$であり、会社によって $\sigma_a^2$ くらいバラける。

・ある会社$k$の1年毎の昇給は大体$b_{mean}$であり、会社によって $\sigma_a^2$ くらいバラける。

・会社の種類 $KID$ は 4種類であり、4つの会社から合計 $N$ 人を選出している。

・ある会社員$n$の$x$ 歳の基本年収 $z[n]$は所属する会社の新卒年収$a[KID[n]]$から毎年 $b[KID[n]]$だけ上がる

・個人、年ごとの出来栄えで実際の年収 $y[n]$ は $\sigma_y^2$だけバラつく。

stan_model2 = """
  data {
    int N;
    int K;
    real X[N];
    real Y[N];
    int<lower=1, upper=K> KID[N];
    int N_s;
    real X_s[N_s];
  }
  
  parameters {
    real a0;
    real b0;
    real a[K];
    real b[K];
    real<lower=0> s_a;
    real<lower=0> s_b;
    real<lower=0> s_Y;
  }
  
  model{
    for (k in 1:K){
      a[k] ~ normal(a0, s_a);
      b[k] ~ normal(b0, s_b);
    }
  
    for (n in 1:N){
      Y[n] ~ normal(a[KID[n]] + b[KID[n]] * (X[n] - 22) , s_Y);
    }
  }
  
  generated quantities {
    real Y_s[N_s];
    for (n in 1:N_s){
      Y_s[n] = normal_rng(a[KID[n]] + b[KID[n]] * (X_s[n] - 22), s_Y);
    }
  }
"""

辞書でデータを渡し学習

ここで、$KID$の方をPythonで作ったのだが、StanとPythonではIndexの開始が異なるので注意。Stanに渡す前にIndexを1からに移動しておきましょう。

df['KID'] = df['KID'].astype(np.int64) + 1

X_s = np.arange(22, 60, 1)
N_s = X_s.shape[0]
stan_data = {"N":df.shape[0], "K":4, "X":df["age"], "Y":df["salary"], 
             "KID": df["KID"], "N_s":N_s, "X_s":X_s}

fit2 = sm2.sampling(data=stan_data, iter=10000, warmup=2000, chains=3, seed=1992)

結果の可視化

fit2オブジェクトにはパラメータの推論に関する情報が入っているので、これらを使って可視化します。

Y_p = fit.extract("Y_s")["Y_s"]
low_y, high_y = mstats.mquantiles(Y_p, [0.025, 0.975], axis=0)

plt.scatter(df["age"], df["salary"])
plt.fill_between(X_s, low_y, high_y, alpha=0.3, color="gray")

a_ = 258.16
b_ = 11.27
x_ = np.arange(22, 60, 1) 
y_ = a_ + b_ * (x_ - 22)
plt.plot(x_, y_, c='r')

a_ = 336.26
b_ = 20.24
x_ = np.arange(22, 60, 1) 
y_ = a_ + b_ * (x_ - 22)
plt.plot(x_, y_, c='g')

a_ = 379.52
b_ = 12.3
x_ = np.arange(22, 60, 1) 
y_ = a_ + b_ * (x_ - 22)
plt.plot(x_, y_, c='b')

a_ = 336.18
b_ = 16.25
x_ = np.arange(22, 60, 1) 
y_ = a_ + b_ * (x_ - 22)
plt.plot(x_, y_, c='y')


plt.legend(["KID[1]", "KID[2]", "KID[3]", "KID[4]"])

f:id:s0sem0y:20181113235203p:plain

完全に会社で給与が決まっていますね…。現実的。 全体データの平均など意味がなかったのじゃ。

まとめ

データの生成過程を想像してモデリング

まず、データに対して安易に直線フィッティングしても意味が無いかもしれません。 ​今回は人工的なデータで検証したため「4つの会社からの抽出」ということをキッチリとモデリングすることで、はっきりと違いが出ることを見ました。

このように階層ベイズでは容易に個体差やバラツキを表現することが可能です。 ​

なぜに個々にフィッティングしないか

さて、最後の図を見て感じたことがあるのではないでしょうか。 ​ それは、4つの会社からデータを抽出していることがわかった時点で、それぞれの会社のデータ毎にフィッティング(モデルを4つ考えてしまえばいい)という方向性でやれば良いのではないか?ということです。 ​ 今回は、社会的な背景によって会社によらずに平均的な新卒年収があり、そこを中心として会社ごとにずれるというモデルにすることで、「全てのデータを使って」一つの正規分布をフィットすれば良くなったのです(昇給額も同様)。これはデータが多く集まっていない場合に有効です。また、個々の会社でデータをフィッティングするより、現実的なモデルになっているともいえます。実際に同じ国、あるいは同じ業界に属する企業群は似たような給与のなります。

このことを考えると、更に業界差を考えたモデルを作ることができそうですね。複数の業界から複数の会社を選出し、その中から更に従業員の年収データを集めると言った具合です。

このようなモデリングに興味を持った方なら、間違いなくベイズ、面白いです。

www.hellocybernetics.tech

なぜかgithubで数式死んでるのでcolabリンクで見てください。 github.com