変分推論を使用した一般化線形混合効果モデルのフィッティング

TensorFlow.orgで表示 GoogleColabで実行 GitHubでソースを表示 ノートブックをダウンロード

インストール

インストール

概要

このコラボでは、TensorFlow確率の変分推論を使用して、一般化線形混合効果モデルを適合させる方法を示します。

モデルファミリ

混合効果モデル線形一般(GLMM)に類似しているモデル線形一般それらが予測線形応答にノイズサンプル特定を組み込むことを除いて(GLM)。これは、めったに見られない機能がより一般的に見られる機能と情報を共有できるため、部分的には便利です。

生成プロセスとして、一般化線形混合効果モデル(GLMM)は次の特徴があります。

$$ \begin{align} \text{for } & r = 1\ldots R: \hspace{2.45cm}\text{# for each random-effect group}\\ &\begin{aligned} \text{for } &c = 1\ldots |C_r|: \hspace{1.3cm}\text{# for each category ("level") of group $r$}\\ &\begin{aligned} \beta_{rc} &\sim \text{MultivariateNormal}(\text{loc}=0_{D_r}, \text{scale}=\Sigma_r^{1/2}) \end{aligned} \end{aligned}\\\\ \text{for } & i = 1 \ldots N: \hspace{2.45cm}\text{# for each sample}\\ &\begin{aligned} &\eta_i = \underbrace{\vphantom{\sum_{r=1}^R}x_i^\top\omega}_\text{fixed-effects} + \underbrace{\sum_{r=1}^R z_{r,i}^\top \beta_{r,C_r(i) } }_\text{random-effects} \\ &Y_i|x_i,\omega,\{z_{r,i} , \beta_r\}_{r=1}^R \sim \text{Distribution}(\text{mean}= g^{-1}(\eta_i)) \end{aligned} \end{align} $$

どこ:

$$ \begin{align} R &= \text{number of random-effect groups}\\ |C_r| &= \text{number of categories for group $r$}\\ N &= \text{number of training samples}\\ x_i,\omega &\in \mathbb{R}^{D_0}\\ D_0 &= \text{number of fixed-effects}\\ C_r(i) &= \text{category (under group $r$) of the $i$th sample}\\ z_{r,i} &\in \mathbb{R}^{D_r}\\ D_r &= \text{number of random-effects associated with group $r$}\\ \Sigma_{r} &\in \{S\in\mathbb{R}^{D_r \times D_r} : S \succ 0 \}\\ \eta_i\mapsto g^{-1}(\eta_i) &= \mu_i, \text{inverse link function}\\ \text{Distribution} &=\text{some distribution parameterizable solely by its mean} \end{align} $$

言い換えると、これは、各グループのすべてのカテゴリが、多変量正規分布からのサンプル$ \ beta_ {rc} $に関連付けられていることを意味します。 $ \ beta_ {rc} $の描画は常に独立していますが、グループ$ r $に対してのみ同じように分散されます。$ r \ in {1、\ ldots、R} $ごとに$ \ Sigma_r $が1つだけあることに注意してください。 。

サンプルのグループの機能($ z_ {r、i} $)と密接に組み合わせると、結果は$ i $番目の予測線形応答(それ以外の場合は$ x_i ^ \ top \ omega $)のサンプル固有のノイズになります。

$ {\ Sigma_r:r \ in {1、\ ldots、R}} $を見積もるとき、基本的にランダム効果グループが運ぶノイズの量を見積もっています。そうしないと、$ x_i ^ \ topに存在する信号をかき消してしまいます。 \ omega $。

$ \ {テキスト配布} $とのためのさまざまなオプションがあります逆リンク関数{ - 1} $、$ G ^が。一般的な選択肢は次のとおりです。

  • $ Y_i \ sim \ text {Normal}(\ text {mean} = \ eta_i、\ text {scale} = \ sigma)$、
  • $ Y_i \ sim \ text {Binomial}(\ text {mean} = n_i \ cdot \ text {sigmoid}(\ eta_i)、\ text {total_count} = n_i)$、および、
  • $ Y_i \ sim \ text {Poisson}(\ text {mean} = \ exp(\ eta_i))$。

より多くの可能性については、 tfp.glmモジュールを。

変分推論

残念ながら、パラメーター$ \ beta、{\ Sigma_r} _r ^ R $の最尤推定値を見つけるには、非分析積分が必要です。この問題を回避するために、代わりに

  1. 付録で$ q _ {\ lambda} $と示されている、パラメーター化された分布のファミリー(「代理密度」)を定義します。
  2. $ q _ {\ lambda} $が実際のターゲット密度に近くなるように、パラメーター$ \ lambda $を見つけます。

ディストリビューションの家族は、適切な寸法の独立したガウス分布になり、「私たちの目標濃度に近い」によって、我々は「最小化を意味しますカルバック・ライブラー情報量を」。例のために、参照してください:「統計学者のためのAレビュー変分推論」のセクション2.2よく書かれ導出し、モチベーションのために。特に、KL発散を最小化することは、負の証拠の下限(ELBO)を最小化することと同等であることを示しています。

トイプロブレム

ゲルマンらの(2007)「ラドンデータセットは、」時々回帰のアプローチを実証するために使用されるデータセットです。 (例えば、この密接に関連PyMC3のブログ記事が。)ラドンデータセットは、米国全体で撮影したラドンの屋内の測定が含まれています。ラドンは当然である放射性ガスocurringれる毒性高濃度です。

私たちのデモンストレーションでは、地下室のある世帯ではラドンレベルが高いという仮説を検証することに関心があるとしましょう。また、ラドン濃度は土壌の種類、つまり地理的な問題に関連していると思われます。

これをML問題としてフレーム化するために、読み取りが行われたフロアの線形関数に基づいて対数ラドンレベルを予測しようとします。また、郡を変量効果として使用し、その際に地理的条件による差異を考慮します。言い換えれば、我々は使います一般化線形混合効果モデルを

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from six.moves import urllib

import matplotlib.pyplot as plt; plt.style.use('ggplot')
import numpy as np
import pandas as pd
import seaborn as sns; sns.set_context('notebook')
import tensorflow_datasets as tfds

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()

import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

また、GPUの可用性を簡単に確認します。

if tf.test.gpu_device_name() != '/device:GPU:0':
  print("We'll just use the CPU for this run.")
else:
  print('Huzzah! Found GPU: {}'.format(tf.test.gpu_device_name()))
We'll just use the CPU for this run.

データセットの取得:

TensorFlowデータセットからデータセットを読み込み、軽い前処理を行います。

def load_and_preprocess_radon_dataset(state='MN'):
  """Load the Radon dataset from TensorFlow Datasets and preprocess it.

  Following the examples in "Bayesian Data Analysis" (Gelman, 2007), we filter
  to Minnesota data and preprocess to obtain the following features:

  - `county`: Name of county in which the measurement was taken.
  - `floor`: Floor of house (0 for basement, 1 for first floor) on which the
    measurement was taken.

  The target variable is `log_radon`, the log of the Radon measurement in the
  house.
  """
  ds = tfds.load('radon', split='train')
  radon_data = tfds.as_dataframe(ds)
  radon_data.rename(lambda s: s[9:] if s.startswith('feat') else s, axis=1, inplace=True)
  df = radon_data[radon_data.state==state.encode()].copy()

  df['radon'] = df.activity.apply(lambda x: x if x > 0. else 0.1)
  # Make county names look nice. 
  df['county'] = df.county.apply(lambda s: s.decode()).str.strip().str.title()
  # Remap categories to start from 0 and end at max(category).
  df['county'] = df.county.astype(pd.api.types.CategoricalDtype())
  df['county_code'] = df.county.cat.codes
  # Radon levels are all positive, but log levels are unconstrained
  df['log_radon'] = df['radon'].apply(np.log)

  # Drop columns we won't use and tidy the index 
  columns_to_keep = ['log_radon', 'floor', 'county', 'county_code']
  df = df[columns_to_keep].reset_index(drop=True)

  return df

df = load_and_preprocess_radon_dataset()
df.head()

GLMMファミリーの専門

このセクションでは、GLMMファミリーをラドンレベルの予測タスクに特化します。これを行うために、最初にGLMMの固定効果の特殊なケースを検討します。

$$ \mathbb{E}[\log(\text{radon}_j)] = c + \text{floor_effect}_j $$

このモデルは、観測値$ j $のログラドンが(予想どおり)$ j $番目の読み取りが行われるフロアと一定の切片によって支配されることを前提としています。擬似コードでは、

def estimate_log_radon(floor):
    return intercept + floor_effect[floor]

すべての床と普遍のために学んで重量がありますintercept用語は。 0階と1階からのラドン測定値を見ると、これは良いスタートかもしれないように見えます。

fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12, 4))
df.groupby('floor')['log_radon'].plot(kind='density', ax=ax1);
ax1.set_xlabel('Measured log(radon)')
ax1.legend(title='Floor')

df['floor'].value_counts().plot(kind='bar', ax=ax2)
ax2.set_xlabel('Floor where radon was measured')
ax2.set_ylabel('Count')
fig.suptitle("Distribution of log radon and floors in the dataset");

png

モデルをもう少し洗練されたものにするためには、地理に関するものを含めて、おそらくさらに良いでしょう。ラドンは、地面に存在する可能性のあるウランの崩壊系列の一部であるため、地理が説明の鍵となるようです。

$$ \mathbb{E}[\log(\text{radon}_j)] = c + \text{floor_effect}_j + \text{county_effect}_j $$

繰り返しますが、擬似コードでは、

def estimate_log_radon(floor, county):
    return intercept + floor_effect[floor] + county_effect[county]

郡固有の重みを除いて、以前と同じです。

十分に大きなトレーニングセットを考えると、これは妥当なモデルです。ただし、ミネソタからのデータを考えると、測定数が少ない郡が多数あることがわかります。たとえば、85の郡のうち39の観測値は5つ未満です。

これにより、郡ごとの観測数が増えるにつれて上記のモデルに収束するように、すべての観測間で統計的強度を共有するようになります。

fig, ax = plt.subplots(figsize=(22, 5));
county_freq = df['county'].value_counts()
county_freq.plot(kind='bar', ax=ax)
ax.set_xlabel('County')
ax.set_ylabel('Number of readings');

png

我々はこのモデルを当てはめるならば、 county_effectベクトルは、おそらく、おそらく過剰適合と貧しい一般化につながる、わずか数のトレーニングサンプルを持っていた郡の結果を記憶することになります。

GLMMは、上記の2つのGLMの中間に位置します。フィッティングを検討するかもしれません

$$ \log(\text{radon}_j) \sim c + \text{floor_effect}_j + \mathcal{N}(\text{county_effect}_j, \text{county_scale}) $$

このモデルは、最初と同じですが、私たちは、正規分布であることを私たちの可能性を修正した、と(シングル)変数を介してすべての郡全体で分散を共有することになりますcounty_scale 。擬似コードでは、

def estimate_log_radon(floor, county):
    county_mean = county_effect[county]
    random_effect = np.random.normal() * county_scale + county_mean
    return intercept + floor_effect[floor] + random_effect

私たちは、以上の同時分布を推測しますcounty_scalecounty_mean 、およびrandom_effect私たちの観測データを使用しています。グローバルcounty_scale 、多くの観測とそれらのいくつかの観察と郡の分散でヒットを提供します。私たちは、郡全体で統計的な強さを共有することができます。さらに、より多くのデータを収集すると、このモデルは、プールされたスケール変数のないモデルに収束します。このデータセットを使用しても、どちらのモデルでも最も観察された郡について同様の結論に達します。

実験

ここで、TensorFlowの変分推論を使用して、上記のGLMMを適合させようとします。まず、データをフィーチャとラベルに分割します。

features = df[['county_code', 'floor']].astype(int)
labels = df[['log_radon']].astype(np.float32).values.flatten()

モデルの指定

def make_joint_distribution_coroutine(floor, county, n_counties, n_floors):

  def model():
    county_scale = yield tfd.HalfNormal(scale=1., name='scale_prior')
    intercept = yield tfd.Normal(loc=0., scale=1., name='intercept')
    floor_weight = yield tfd.Normal(loc=0., scale=1., name='floor_weight')
    county_prior = yield tfd.Normal(loc=tf.zeros(n_counties),
                                    scale=county_scale,
                                    name='county_prior')
    random_effect = tf.gather(county_prior, county, axis=-1)

    fixed_effect = intercept + floor_weight * floor
    linear_response = fixed_effect + random_effect
    yield tfd.Normal(loc=linear_response, scale=1., name='likelihood')
  return tfd.JointDistributionCoroutineAutoBatched(model)

joint = make_joint_distribution_coroutine(
    features.floor.values, features.county_code.values, df.county.nunique(),
    df.floor.nunique())

# Define a closure over the joint distribution 
# to condition on the observed labels.
def target_log_prob_fn(*args):
  return joint.log_prob(*args, likelihood=labels)

代理後部を指定する

ここで、サロゲートファミリ$ q _ {\ lambda} $をまとめました。ここで、パラメータ$ \ lambda $はトレーニング可能です。この場合、私たちの家族は、パラメーターごとに1つずつ、独立した多変量正規分布であり、$ \ lambda = {(\ mu_j、\ sigma_j)} $です。ここで、$ j $は4つのパラメーターにインデックスを付けます。

我々は、サロゲート家族の用途を合わせて使う方法tf.Variables 。また、使用tfp.util.TransformedVariable一緒にSoftplus (トレーニング可能)、スケールパラメータが正であることを制約します。さらに、当社は、適用Softplus全体にscale_prior正のパラメータであり、。

最適化を支援するために、これらのトレーニング可能な変数を少しのジッターで初期化します。

# Initialize locations and scales randomly with `tf.Variable`s and 
# `tfp.util.TransformedVariable`s.
_init_loc = lambda shape=(): tf.Variable(
    tf.random.uniform(shape, minval=-2., maxval=2.))
_init_scale = lambda shape=(): tfp.util.TransformedVariable(
    initial_value=tf.random.uniform(shape, minval=0.01, maxval=1.),
    bijector=tfb.Softplus())
n_counties = df.county.nunique()

surrogate_posterior = tfd.JointDistributionSequentialAutoBatched([
  tfb.Softplus()(tfd.Normal(_init_loc(), _init_scale())),           # scale_prior
  tfd.Normal(_init_loc(), _init_scale()),                           # intercept
  tfd.Normal(_init_loc(), _init_scale()),                           # floor_weight
  tfd.Normal(_init_loc([n_counties]), _init_scale([n_counties]))])  # county_prior

このセルを置き換えることができることに注意してくださいtfp.experimental.vi.build_factored_surrogate_posterior :のように、

surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
  event_shape=joint.event_shape_tensor()[:-1],
  constraining_bijectors=[tfb.Softplus(), None, None, None])

結果

私たちの目標は、扱いやすいパラメーター化された分布のファミリーを定義し、次にパラメーターを選択して、ターゲット分布に近い扱いやすい分布を作成することであることを思い出してください。

我々は、上記のサロゲート分布を構築し、そして使用することができtfp.vi.fit_surrogate_posterior間のKullback-Liebler発散を最小限にするためにどのcorresonds(負ELBOを最小代替モデルのパラメータを見つけるためにオプティマイザとステップの所定の数を受け入れます、サロゲートとターゲット分布)。

戻り値は、各ステップで否定ELBOあり、かつにおける分布surrogate_posterior 、オプティマイザによって発見パラメータで更新されています。

optimizer = tf.optimizers.Adam(learning_rate=1e-2)

losses = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn, 
    surrogate_posterior,
    optimizer=optimizer,
    num_steps=3000, 
    seed=42,
    sample_size=2)

(scale_prior_, 
 intercept_, 
 floor_weight_, 
 county_weights_), _ = surrogate_posterior.sample_distributions()
print('        intercept (mean): ', intercept_.mean())
print('     floor_weight (mean): ', floor_weight_.mean())
print(' scale_prior (approx. mean): ', tf.reduce_mean(scale_prior_.sample(10000)))
intercept (mean):  tf.Tensor(1.4352839, shape=(), dtype=float32)
     floor_weight (mean):  tf.Tensor(-0.6701997, shape=(), dtype=float32)
 scale_prior (approx. mean):  tf.Tensor(0.28682157, shape=(), dtype=float32)
fig, ax = plt.subplots(figsize=(10, 3))
ax.plot(losses, 'k-')
ax.set(xlabel="Iteration",
       ylabel="Loss (ELBO)",
       title="Loss during training",
       ylim=0);

png

その平均の不確実性とともに、推定された平均郡効果をプロットできます。これを観測数で並べ替えました。左側が最大です。観測値が多い郡では不確実性は小さいが、観測値が1つか2つしかない郡では不確実性が大きいことに注意してください。

county_counts = (df.groupby(by=['county', 'county_code'], observed=True)
                   .agg('size')
                   .sort_values(ascending=False)
                   .reset_index(name='count'))

means = county_weights_.mean()
stds = county_weights_.stddev()

fig, ax = plt.subplots(figsize=(20, 5))

for idx, row in county_counts.iterrows():
  mid = means[row.county_code]
  std = stds[row.county_code]
  ax.vlines(idx, mid - std, mid + std, linewidth=3)
  ax.plot(idx, means[row.county_code], 'ko', mfc='w', mew=2, ms=7)

ax.set(
    xticks=np.arange(len(county_counts)),
    xlim=(-1, len(county_counts)),
    ylabel="County effect",
    title=r"Estimates of county effects on log radon levels. (mean $\pm$ 1 std. dev.)",
)
ax.set_xticklabels(county_counts.county, rotation=90);

png

実際、推定された標準偏差に対して観測値の対数をプロットすることで、これをより直接的に確認でき、関係がほぼ線形であることがわかります。

fig, ax = plt.subplots(figsize=(10, 7))
ax.plot(np.log1p(county_counts['count']), stds.numpy()[county_counts.county_code], 'o')
ax.set(
    ylabel='Posterior std. deviation',
    xlabel='County log-count',
    title='Having more observations generally\nlowers estimation uncertainty'
);

png

比較lme4 Rに

%%shell
exit  # Trick to make this block not execute.

radon = read.csv('srrs2.dat', header = TRUE)
radon = radon[radon$state=='MN',]
radon$radon = ifelse(radon$activity==0., 0.1, radon$activity)
radon$log_radon = log(radon$radon)

# install.packages('lme4')
library(lme4)
fit <- lmer(log_radon ~ 1 + floor + (1 | county), data=radon)
fit

# Linear mixed model fit by REML ['lmerMod']
# Formula: log_radon ~ 1 + floor + (1 | county)
#    Data: radon
# REML criterion at convergence: 2171.305
# Random effects:
#  Groups   Name        Std.Dev.
#  county   (Intercept) 0.3282
#  Residual             0.7556
# Number of obs: 919, groups:  county, 85
# Fixed Effects:
# (Intercept)        floor
#       1.462       -0.693
<IPython.core.display.Javascript at 0x7f90b888e9b0>
<IPython.core.display.Javascript at 0x7f90b888e780>
<IPython.core.display.Javascript at 0x7f90b888e780>
<IPython.core.display.Javascript at 0x7f90bce1dfd0>
<IPython.core.display.Javascript at 0x7f90b888e780>
<IPython.core.display.Javascript at 0x7f90b888e780>
<IPython.core.display.Javascript at 0x7f90b888e780>
<IPython.core.display.Javascript at 0x7f90b888e780>

次の表は、結果をまとめたものです。

print(pd.DataFrame(data=dict(intercept=[1.462, tf.reduce_mean(intercept_.mean()).numpy()],
                             floor=[-0.693, tf.reduce_mean(floor_weight_.mean()).numpy()],
                             scale=[0.3282, tf.reduce_mean(scale_prior_.sample(10000)).numpy()]),
                   index=['lme4', 'vi']))
intercept   floor     scale
lme4   1.462000 -0.6930  0.328200
vi     1.435284 -0.6702  0.287251

このテーブルには、VIの結果は、の〜10%の範囲内であることを示しlme4さん。これは、次の理由からやや驚くべきことです。

  • lme4基づいているラプラスの方法(ないVI)、
  • このコラボでは、実際に収束するための努力は行われませんでした。
  • ハイパーパラメータを調整するために最小限の労力が費やされました。
  • データを正規化または前処理するための努力は必要ありませんでした(たとえば、中心の特徴など)。

結論

このコラボでは、一般化線形混合効果モデルについて説明し、TensorFlow確率を使用して変分推論を使用してそれらを適合させる方法を示しました。トイプロブレムには数百のトレーニングサンプルしかありませんでしたが、ここで使用される手法は、大規模に必要とされるものと同じです。