Google I / O에서 기조 연설, 제품 세션, 워크샵 등을 시청 하세요. 재생 목록보기

변동 추론을 사용하여 일반화 선형 혼합 효과 모델 피팅하기

TensorFlow.org에서보기 Google Colab에서 실행 GitHub에서 소스보기 노트북 다운로드

설치

설치

요약

이 colab에서는 TensorFlow Probability의 변형 추론을 사용하여 일반화 된 선형 혼합 효과 모델을 맞추는 방법을 보여줍니다.

모델 패밀리

일반화 선형 혼합 효과 모델 (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 $가 있습니다. .

표본 그룹의 특징 ($ z_ {r, i} $)과 유사하게 결합되면 결과는 $ i $ 번째 예측 선형 반응 (그렇지 않으면 $ x_i ^ \ top \ omega $)에 대한 표본 별 노이즈입니다.

$ {\ Sigma_r : r \ in {1, \ ldots, R}} $를 추정 할 때 우리는 본질적으로 $ x_i ^ \ top에있는 신호를 익사시킬 수있는 랜덤 효과 그룹이 전달하는 노이즈의 양을 추정합니다. \ omega $.

$ \ text {Distribution} $ 및 역 링크 함수 , $ g ^ {-1} $에 대한 다양한 옵션이 있습니다. 일반적인 선택은 다음과 같습니다.

  • $ 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 $를 찾습니다.

분포 군은 적절한 차원의 독립적 인 가우시안이 될 것이며 "목표 밀도에 가까움"이란 " Kullback-Leibler 발산 최소화"를 의미합니다. 예를 들어 잘 작성된 유도 및 동기 는 "변이 추론 : 통계학 자를위한 검토"의 섹션 2.2를 참조하십시오. 특히 KL 발산을 최소화하는 것은 부정적 증거 하한 (ELBO)을 최소화하는 것과 동일 함을 보여줍니다.

장난감 문제

Gelman 등 (2007)의 "라돈 데이터 세트" 는 회귀에 대한 접근 방식을 입증하는 데 사용되는 데이터 세트입니다. (예 : 밀접하게 관련된 PyMC3 블로그 게시물 ) 라돈 데이터 세트에는 미국 전역에서 측정 한 라돈의 실내 측정 값이 포함되어 있습니다. 라돈 은 자연적으로 발생하는 방사성 가스로 고농도에서 독성 이 있습니다.

데모를 위해 지하실이있는 가정에서 라돈 수치가 더 높다는 가설을 검증하는 데 관심이 있다고 가정 해 보겠습니다. 우리는 또한 라돈 농도가 토양 유형, 즉 지리 문제와 관련이 있다고 의심합니다.

이것을 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 $ th 판독 값이 적용되는 바닥과 일정한 절편에 의해 (기대 적으로) 관리된다고 가정합니다. 의사 코드에서는 다음과 같이 작성할 수 있습니다.

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은 위의 두 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_scale , county_meanrandom_effect 우리의 관측 데이터를 사용합니다. global 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)

사후 대리 지정

이제 매개 변수 $ \ lambda $를 학습 할 수있는 $ q _ {\ lambda} $ 대리 패밀리를 구성합니다. 이 경우, 우리 가족은 각 매개 변수에 대해 하나씩 독립적 인 다변량 정규 분포이고 $ \ lambda = {(\ mu_j, \ sigma_j)} $입니다. 여기서 $ j $는 4 개의 매개 변수를 색인화합니다.

서로 게이트 패밀리를 맞추기 위해 사용하는 방법은 tf.Variables 사용 tf.Variables . 또한 tfp.util.TransformedVariable 와 함께 Softplus 을 사용하여 (학습 가능한) 스케일 매개 변수를 양수로 제한합니다. 또한 양수 매개 변수 인 전체 scale_prior Softplus 를 적용합니다.

최적화를 돕기 위해 약간의 지터를 사용하여 이러한 학습 가능한 변수를 초기화합니다.

# 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_posteriortfp.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 , tfp.vi.fit_surrogate_posterior 를 사용할 수 있습니다.이 tfp.vi.fit_surrogate_posterior 는 최적화 프로그램과 주어진 단계 수를 받아 음성 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

추정 된 평균 카운티 효과와 해당 평균의 불확실성을 플롯 할 수 있습니다. 우리는 이것을 관측치 수로 정렬했으며 가장 큰 것은 왼쪽에 있습니다. 관측치가 많은 카운티에서는 불확실성이 작지만 관측치가 한두 개만있는 카운티에서는 더 큽니다.

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

R의 lme4 와 비교

%%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 결과가 lme4 의 ~ 10 % 내에 있음을 나타냅니다. 이것은 다음과 같은 이유로 다소 놀랍습니다.

  • lme4 는 VI가 아닌 Laplace의 방법을 기반으로합니다.
  • 이 colab에서 실제로 수렴하려는 노력은 없었습니다.
  • 초 매개 변수를 조정하는 데 최소한의 노력을 기울 였지만
  • 데이터 (예 : 센터 기능 등)를 정규화하거나 전처리하는 데 노력을 기울이지 않았습니다.

결론

이 colab에서는 일반화 된 선형 혼합 효과 모델을 설명하고 TensorFlow Probability를 사용하여 변형 추론을 사용하여 피팅하는 방법을 보여주었습니다. 장난감 문제에는 수백 개의 훈련 샘플 만 있었지만 여기에 사용 된 기술은 규모에서 필요한 것과 동일합니다.