HLM_TFP_R_Stan.ipynb

Kernel: Python 3

Copyright 2018 The TensorFlow Probability Authors.

Licensed under the Apache License, Version 2.0 (the "License");

#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Regressão linear de efeitos mistos no {TF Probability, R, Stan}

Veja em TensorFlow.org

Executar no Google Colab

Ver fonte em GitHub

Baixar notebook

1 – Introdução

Neste Colab, vamos ajustar um modelo de regressão linear de efeitos mistos para um dataset de exemplo popular. Vamos fazer o ajuste três vezes utilizando o lme4 do R, o pacote mixed-effects do Stan e os primitivos do TensorFlow Probability (TFP). Concluiremos mostrando que todos os três fornecem aproximadamente os mesmos parâmetros ajustados e distribuições posteriores.

Nossa principal conclusão é de que o TFP tem as partes gerais necessárias para fazer o ajuste de modelos tipo HLM e gera resultados consistentes com outros pacotes de software, como lme4 e rstanarm. Este Colab não mostra com exatidão a eficiência computacional de nenhum dos pacotes comparados.

%matplotlib inline

import os
from six.moves import urllib
import numpy as np
import pandas as pd
import warnings

from matplotlib import pyplot as plt
import seaborn as sns

from IPython.core.pylabtools import figsize
figsize(11, 9)

import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp

2 – Modelo linear hierárquico

Para a comparação entre o R, Stan e TFP, vamos ajustar um modelo linear hierárquico (HLM, na sigla em inglês) para o dataset Radon, popularizado em Bayesian Data Analysis de Gelman, et. al. (Análise bayesiana de dados, página 559, segunda edição; página 250, terceira edição).

Vamos supor o seguinte modelo generativo:

$$\begin{align*} \text{for } & c=1\ldots \text{NumCounties}:\\ & \beta_c \sim \text{Normal}\left(\text{loc}=0, \text{scale}=\sigma_C \right) \\ \text{for } & i=1\ldots \text{NumSamples}:\\ &\eta_i = \underbrace{\omega_0 + \omega_1 \text{Floor}_i}_\text{fixed effects} + \underbrace{\beta_{ \text{County}_i} \log( \text{UraniumPPM}_{\text{County}_i}))}_\text{random effects} \\ &\log(\text{Radon}_i) \sim \text{Normal}(\text{loc}=\eta_i , \text{scale}=\sigma_N) \end{align*}$$

Na "notação de til" do lme4 do R, esse modelo é equivalente a:

log_radon ~ 1 + floor + (0 + log_uranium_ppm | county)

Vamos encontrar a estimativa de máxima verossimilhança para $\omega, \sigma_C, \sigma_N$ usando a distribuição posterior (condicionada às evidências) de $\{\beta_c\}_{c=1}^\text{NumCounties}$.

Para ver essencialmente o mesmo modelo, porém com um intercepto aleatório, confira o Apêndice A.

Para ver uma especificação mais geral de HLMs, confira o Apêndice B.

3 – Transformação dos dados

Nesta seção, vamos obter o dataset radon e fazer um pré-processamento mínimo para deixá-lo adequado ao modelo presumido.

def load_and_preprocess_radon_dataset(state='MN'):
  """Preprocess Radon dataset as done in "Bayesian Data Analysis" book.
  
  We filter to Minnesota data (919 examples) and preprocess to obtain the
  following features:
  - `log_uranium_ppm`: Log of soil uranium measurements.
  - `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()

  # For any missing or invalid activity readings, we'll use a value of `0.1`.
  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).
  county_name = sorted(df.county.unique())
  df['county'] = df.county.astype(
      pd.api.types.CategoricalDtype(categories=county_name)).cat.codes
  county_name = list(map(str.strip, county_name))

  df['log_radon'] = df['radon'].apply(np.log)
  df['log_uranium_ppm'] = df['Uppm'].apply(np.log)
  df = df[['idnum', 'log_radon', 'floor', 'county', 'log_uranium_ppm']]

  return df, county_name

radon, county_name = load_and_preprocess_radon_dataset()

# We'll use the following directory to store our preprocessed dataset.
CACHE_DIR = os.path.join(os.sep, 'tmp', 'radon')

# Save processed data. (So we can later read it in R.)
if not tf.gfile.Exists(CACHE_DIR):
  tf.gfile.MakeDirs(CACHE_DIR)
with tf.gfile.Open(os.path.join(CACHE_DIR, 'radon.csv'), 'w') as f:
  radon.to_csv(f, index=False)

3.1 – Conheça seus dados

Nesta seção, vamos explorar o dataset radon para entender melhor por que o modelo proposto poderá ser razoável.

radon.head()

fig, ax = plt.subplots(figsize=(22, 5));
county_freq = radon['county'].value_counts()
county_freq.plot(kind='bar', color='#436bad');
plt.xlabel('County index')
plt.ylabel('Number of radon readings')
plt.title('Number of radon readings per county', fontsize=16)
county_freq = np.array(zip(county_freq.index, county_freq.values))  # We'll use this later.

fig, ax = plt.subplots(ncols=2, figsize=[10, 4]);

radon['log_radon'].plot(kind='density', ax=ax[0]);
ax[0].set_xlabel('log(radon)')

radon['floor'].value_counts().plot(kind='bar', ax=ax[1]);
ax[1].set_xlabel('Floor');
ax[1].set_ylabel('Count');

fig.subplots_adjust(wspace=0.25)

Conclusões:

• Há uma cauda longa de 85 condados (uma ocorrência comum em modelos lineares generalizados mistos).

• $\log(\text{Radon})$ não é restringido (portanto, a regressão linear pode fazer sentido).

• As medições são feitas principalmente no $0$º andar, nenhuma medição foi feita acima do $1$º andar (portanto, os efeitos fixos terão somente dois pesos).

4 – HLM no R

Nesta seção, usaremos o pacote lme4 do R para ajustar o modelo probabilístico descrito acima.

OBSERVAÇÃO: para executar o código nesta seção, você precisa mudar para um runtime de Colab do R.

suppressMessages({
  library('bayesplot')
  library('data.table')
  library('dplyr')
  library('gfile')
  library('ggplot2')
  library('lattice')
  library('lme4')
  library('plyr')
  library('rstanarm')
  library('tidyverse')
  RequireInitGoogle()
})

data = read_csv(gfile::GFile('/tmp/radon/radon.csv'))

Parsed with column specification:
cols(
  log_radon = col_double(),
  floor = col_integer(),
  county = col_integer(),
  log_uranium_ppm = col_double()
)

head(data)

# A tibble: 6 x 4
  log_radon floor county log_uranium_ppm
      <dbl> <int>  <int>           <dbl>
1     0.788     1      0          -0.689
2     0.788     0      0          -0.689
3     1.06      0      0          -0.689
4     0         0      0          -0.689
5     1.13      0      1          -0.847
6     0.916     0      1          -0.847

# https://github.com/stan-dev/example-models/wiki/ARM-Models-Sorted-by-Chapter
radon.model <- lmer(log_radon ~ 1 + floor  + (0 + log_uranium_ppm | county), data = data)

summary(radon.model)

Linear mixed model fit by REML ['lmerMod']
Formula: log_radon ~ 1 + floor + (0 + log_uranium_ppm | county)
   Data: data

REML criterion at convergence: 2166.3

Scaled residuals: 
    Min      1Q  Median      3Q     Max 
-4.5202 -0.6064  0.0107  0.6334  3.4111 

Random effects:
 Groups   Name            Variance Std.Dev.
 county   log_uranium_ppm 0.7545   0.8686  
 Residual                 0.5776   0.7600  
Number of obs: 919, groups:  county, 85

Fixed effects:
            Estimate Std. Error t value
(Intercept)  1.47585    0.03899   37.85
floor       -0.67974    0.06963   -9.76

Correlation of Fixed Effects:
      (Intr)
floor -0.330

qqmath(ranef(radon.model, condVar=TRUE))

$county

write.csv(as.data.frame(ranef(radon.model, condVar = TRUE)), '/tmp/radon/lme4_fit.csv')

5 – HLM no Stan

Nesta seção, usaremos o rstanarm para ajustar um modelo do Stan usando a mesma fórmula/sintaxe do modelo do lme4 acima.

Diferentemente do modelo do lme4 e do TF abaixo, o rstanarm é um modelo totalmente bayesiano, isto é, todos os parâmetros são presumidamente obtidos a partir de uma distribuição normal, com os parâmetros em si obtidos a partir de uma distribuição.

OBSERVAÇÃO: para executar o código nesta seção, você precisa mudar para um runtime de Colab do R.

fit <- stan_lmer(log_radon ~ 1 + floor  + (0 + log_uranium_ppm | county), data = data)

SAMPLING FOR MODEL 'continuous' NOW (CHAIN 1).

Chain 1, Iteration:    1 / 2000 [  0%]  (Warmup)
Chain 1, Iteration:  200 / 2000 [ 10%]  (Warmup)
Chain 1, Iteration:  400 / 2000 [ 20%]  (Warmup)
Chain 1, Iteration:  600 / 2000 [ 30%]  (Warmup)
Chain 1, Iteration:  800 / 2000 [ 40%]  (Warmup)
Chain 1, Iteration: 1000 / 2000 [ 50%]  (Warmup)
Chain 1, Iteration: 1001 / 2000 [ 50%]  (Sampling)
Chain 1, Iteration: 1200 / 2000 [ 60%]  (Sampling)
Chain 1, Iteration: 1400 / 2000 [ 70%]  (Sampling)
Chain 1, Iteration: 1600 / 2000 [ 80%]  (Sampling)
Chain 1, Iteration: 1800 / 2000 [ 90%]  (Sampling)
Chain 1, Iteration: 2000 / 2000 [100%]  (Sampling)
 Elapsed Time: 7.73495 seconds (Warm-up)
               2.98852 seconds (Sampling)
               10.7235 seconds (Total)


SAMPLING FOR MODEL 'continuous' NOW (CHAIN 2).

Chain 2, Iteration:    1 / 2000 [  0%]  (Warmup)
Chain 2, Iteration:  200 / 2000 [ 10%]  (Warmup)
Chain 2, Iteration:  400 / 2000 [ 20%]  (Warmup)
Chain 2, Iteration:  600 / 2000 [ 30%]  (Warmup)
Chain 2, Iteration:  800 / 2000 [ 40%]  (Warmup)
Chain 2, Iteration: 1000 / 2000 [ 50%]  (Warmup)
Chain 2, Iteration: 1001 / 2000 [ 50%]  (Sampling)
Chain 2, Iteration: 1200 / 2000 [ 60%]  (Sampling)
Chain 2, Iteration: 1400 / 2000 [ 70%]  (Sampling)
Chain 2, Iteration: 1600 / 2000 [ 80%]  (Sampling)
Chain 2, Iteration: 1800 / 2000 [ 90%]  (Sampling)
Chain 2, Iteration: 2000 / 2000 [100%]  (Sampling)
 Elapsed Time: 7.51252 seconds (Warm-up)
               3.08653 seconds (Sampling)
               10.5991 seconds (Total)


SAMPLING FOR MODEL 'continuous' NOW (CHAIN 3).

Chain 3, Iteration:    1 / 2000 [  0%]  (Warmup)
Chain 3, Iteration:  200 / 2000 [ 10%]  (Warmup)
Chain 3, Iteration:  400 / 2000 [ 20%]  (Warmup)
Chain 3, Iteration:  600 / 2000 [ 30%]  (Warmup)
Chain 3, Iteration:  800 / 2000 [ 40%]  (Warmup)
Chain 3, Iteration: 1000 / 2000 [ 50%]  (Warmup)
Chain 3, Iteration: 1001 / 2000 [ 50%]  (Sampling)
Chain 3, Iteration: 1200 / 2000 [ 60%]  (Sampling)
Chain 3, Iteration: 1400 / 2000 [ 70%]  (Sampling)
Chain 3, Iteration: 1600 / 2000 [ 80%]  (Sampling)
Chain 3, Iteration: 1800 / 2000 [ 90%]  (Sampling)
Chain 3, Iteration: 2000 / 2000 [100%]  (Sampling)
 Elapsed Time: 8.14628 seconds (Warm-up)
               3.01001 seconds (Sampling)
               11.1563 seconds (Total)


SAMPLING FOR MODEL 'continuous' NOW (CHAIN 4).

Chain 4, Iteration:    1 / 2000 [  0%]  (Warmup)
Chain 4, Iteration:  200 / 2000 [ 10%]  (Warmup)
Chain 4, Iteration:  400 / 2000 [ 20%]  (Warmup)
Chain 4, Iteration:  600 / 2000 [ 30%]  (Warmup)
Chain 4, Iteration:  800 / 2000 [ 40%]  (Warmup)
Chain 4, Iteration: 1000 / 2000 [ 50%]  (Warmup)
Chain 4, Iteration: 1001 / 2000 [ 50%]  (Sampling)
Chain 4, Iteration: 1200 / 2000 [ 60%]  (Sampling)
Chain 4, Iteration: 1400 / 2000 [ 70%]  (Sampling)
Chain 4, Iteration: 1600 / 2000 [ 80%]  (Sampling)
Chain 4, Iteration: 1800 / 2000 [ 90%]  (Sampling)
Chain 4, Iteration: 2000 / 2000 [100%]  (Sampling)
 Elapsed Time: 7.6801 seconds (Warm-up)
               3.23663 seconds (Sampling)
               10.9167 seconds (Total)

Observação: os runtimes têm um único núcleo de CPU (este Colab não tem a pretensão de ser uma representação fiel do runtime do Stan ou do TFP).

fit

stan_lmer(formula = log_radon ~ 1 + floor + (0 + log_uranium_ppm | 
    county), data = data)

Estimates:
            Median MAD_SD
(Intercept)  1.5    0.0  
floor       -0.7    0.1  
sigma        0.8    0.0  

Error terms:
 Groups   Name            Std.Dev.
 county   log_uranium_ppm 0.87    
 Residual                 0.76    
Num. levels: county 85 

Sample avg. posterior predictive 
distribution of y (X = xbar):
         Median MAD_SD
mean_PPD 1.2    0.0   

Observations: 919  Number of unconstrained parameters: 90 

color_scheme_set("red")
ppc_dens_overlay(y = fit$y,
                 yrep = posterior_predict(fit, draws = 50))

color_scheme_set("brightblue")
ppc_intervals(
  y = data$log_radon,
  yrep = posterior_predict(fit),
  x = data$county,
  prob = 0.8
) +
  labs(
    x = "County",
    y = "log radon",
    title = "80% posterior predictive intervals \nvs observed log radon",
    subtitle = "by county"
  ) +
  panel_bg(fill = "gray95", color = NA) +
  grid_lines(color = "white")

# Write the posterior samples (4000 for each variable) to a CSV.
write.csv(tidy(as.matrix(fit)), "/tmp/radon/stan_fit.csv")

Observação: volte para o runtime do kernel do TF no Python.

with tf.gfile.Open('/tmp/radon/lme4_fit.csv', 'r') as f:
  lme4_fit = pd.read_csv(f, index_col=0)

lme4_fit.head()

Recupere as estimativas de ponto e os desvios padrão condicionais para os efeitos aleatórios do grupo a partir do lme4 para visualização posterior.

posterior_random_weights_lme4 = np.array(lme4_fit.condval, dtype=np.float32)
lme4_prior_scale = np.array(lme4_fit.condsd, dtype=np.float32)
print(posterior_random_weights_lme4.shape, lme4_prior_scale.shape)

(85,) (85,)

Obtenha amostras para os pesos de condado usando as médias e desvios padrão estimados pelo lme4.

with tf.Session() as sess:
  lme4_dist = tfp.distributions.Independent(
      tfp.distributions.Normal(
          loc=posterior_random_weights_lme4,
          scale=lme4_prior_scale),
      reinterpreted_batch_ndims=1)
  posterior_random_weights_lme4_final_ = sess.run(lme4_dist.sample(4000))

posterior_random_weights_lme4_final_.shape

(4000, 85)

Também recuperamos as amostras posteriores dos pesos de condado a partir do ajuste do Stan.

with tf.gfile.Open('/tmp/radon/stan_fit.csv', 'r') as f:
  samples = pd.read_csv(f, index_col=0)

samples.head()

posterior_random_weights_cols = [
    col for col in samples.columns if 'b.log_uranium_ppm.county' in col
]
posterior_random_weights_final_stan = samples[
    posterior_random_weights_cols].values
print(posterior_random_weights_final_stan.shape)

(4000, 85)

Este exemplo do Stan mostra como implementar a regressão linear de efeitos mistos em um estilo próximo ao do TFP, isto é, especificando o modelo probabilístico diretamente.

6 – HLM no TF Probability

Nesta seção, usaremos os primitivos de baixo nível do TensorFlow Probability (Distributions) para especificar o modelo linear hierárquico, bem como para ajustar os parâmetros desconhecidos.

# Handy snippet to reset the global graph and global session.
with warnings.catch_warnings():
  warnings.simplefilter('ignore')
  tf.reset_default_graph()
  try:
    sess.close()
  except:
    pass
  sess = tf.InteractiveSession()

6.1 – Especifique o modelo

Nesta seção, vamos especificar o modelo linear de efeitos mistos de radônio usando primitivos do TFP. Para fazer isso, especificamos duas funções que geram duas distribuições do TFP:

• make_weights_prior: uma distribuição anterior normal multivariada para os pesos aleatórios (que são multiplicados por $\log(\text{UraniumPPM}_{c_i})$ para computar o preditor linear).

• make_log_radon_likelihood: um lote de distribuições Normal para cada variável dependente $\log(\text{Radon}_i)$ observada.

Como estamos ajustando os parâmetros de cada uma dessas distribuições, precisamos usar variáveis do TF (isto é, tf.get_variable). Porém, como queremos usar otimização não restringida, precisamos encontrar uma forma de restringir valores reais para conseguir a semântica necessária, ou seja, positivos que representam desvios padrão.

inv_scale_transform = lambda y: np.log(y)  # Not using TF here.
fwd_scale_transform = tf.exp

A seguinte função constrói a distribuição anterior: $p(\beta|\sigma_C)$, em que $\beta$ denota os pesos de efeitos aleatórios, e $\sigma_C$ denota o desvio padrão.

Usamos tf.make_template para garantir que a primeira chamada a essa função instancie as variáveis do TF usadas e que todas as chamadas subsequentes reutilizem o valor atual das variáveis.

def _make_weights_prior(num_counties, dtype):
  """Returns a `len(log_uranium_ppm)` batch of univariate Normal."""
  raw_prior_scale = tf.get_variable(
      name='raw_prior_scale',
      initializer=np.array(inv_scale_transform(1.), dtype=dtype))
  return tfp.distributions.Independent(
      tfp.distributions.Normal(
          loc=tf.zeros(num_counties, dtype=dtype),
          scale=fwd_scale_transform(raw_prior_scale)),
      reinterpreted_batch_ndims=1)


make_weights_prior = tf.make_template(
    name_='make_weights_prior', func_=_make_weights_prior)

A seguinte função constrói a verossimilhança: $p(y|x,\omega,\beta,\sigma_N)$, em que $y,x$ denota a resposta e a evidência, $\omega,\beta$ denota pesos de efeitos aleatórios e fixos, e $\sigma_N$ denota o desvio padrão.

Novamente, usamos tf.make_template para garantir que as variáveis do TF sejam reutilizadas por chamadas subsequentes.

def _make_log_radon_likelihood(random_effect_weights, floor, county,
                               log_county_uranium_ppm, init_log_radon_stddev):
  raw_likelihood_scale = tf.get_variable(
      name='raw_likelihood_scale',
      initializer=np.array(
          inv_scale_transform(init_log_radon_stddev), dtype=dtype))
  fixed_effect_weights = tf.get_variable(
      name='fixed_effect_weights', initializer=np.array([0., 1.], dtype=dtype))
  fixed_effects = fixed_effect_weights[0] + fixed_effect_weights[1] * floor
  random_effects = tf.gather(
      random_effect_weights * log_county_uranium_ppm,
      indices=tf.to_int32(county),
      axis=-1)
  linear_predictor = fixed_effects + random_effects
  return tfp.distributions.Normal(
      loc=linear_predictor, scale=fwd_scale_transform(raw_likelihood_scale))


make_log_radon_likelihood = tf.make_template(
    name_='make_log_radon_likelihood', func_=_make_log_radon_likelihood)

Por fim, usamos os geradores de distribuição anterior e verossimilhança para construir a densidade logarítmica conjunta.

def joint_log_prob(random_effect_weights, log_radon, floor, county,
                   log_county_uranium_ppm, dtype):
  num_counties = len(log_county_uranium_ppm)
  rv_weights = make_weights_prior(num_counties, dtype)
  rv_radon = make_log_radon_likelihood(
      random_effect_weights,
      floor,
      county,
      log_county_uranium_ppm,
      init_log_radon_stddev=radon.log_radon.values.std())
  return (rv_weights.log_prob(random_effect_weights)
          + tf.reduce_sum(rv_radon.log_prob(log_radon), axis=-1))

6.2 – Treinamento (maximização de expectativa com aproximação estocástica)

Para ajustar o modelo de regressão linear de efeitos mistos, usaremos uma versão de aproximação estocástica do algoritmo de maximização de expectativa (SAEM, na sigla em inglês). A ideia básica é usar amostras da distribuição posterior para fazer a aproximação da densidade logarítmica conjunta esperada (passo E). Em seguida, encontramos os parâmetros que maximizam esse cálculo (passo M). De uma forma um pouco mais concreta, a iteração de ponto fixo é dada por:

$$\begin{align*} \text{E}[ \log p(x, Z | \theta) | \theta_0] &\approx \frac{1}{M} \sum_{m=1}^M \log p(x, z_m | \theta), \quad Z_m\sim p(Z | x, \theta_0) && \text{E-step}\\ &=: Q_M(\theta, \theta_0) \\ \theta_0 &= \theta_0 - \eta \left.\nabla_\theta Q_M(\theta, \theta_0)\right|_{\theta=\theta_0} && \text{M-step} \end{align*}$$

em que $x$ denota a evidência, $Z$ denota uma variável latente que precisa ser marginalizada, e $\theta,\theta_0$ denotam possíveis parametrizações.

Confira uma explicação mais detalhada em Convergence of a stochastic approximation version of the EM algorithms (Convergência de uma versão de aproximação estocástica dos algoritmos de EM) de Bernard Delyon, Marc Lavielle, Eric, Moulines (Ann. Statist., 1999).

Para computar o passo E, precisamos fazer a amostragem da distribuição posterior. Como não é fácil fazer isso, usamos o Monte Carlo Hamiltoniano (HMC). O HMC é um procedimento de Monte Carlo via Cadeias de Markov que usa gradientes (estado wrt, não parâmetros) da densidade logarítmica posterior não normalizada para propor novas amostras.

Especificar a densidade logarítmica posterior não normalizada é simples: ela é simplesmente a densidade logarítmica conjunta "fixada" para a condição que desejamos.

# Specify unnormalized posterior.

dtype = np.float32

log_county_uranium_ppm = radon[
    ['county', 'log_uranium_ppm']].drop_duplicates().values[:, 1]
log_county_uranium_ppm = log_county_uranium_ppm.astype(dtype)

def unnormalized_posterior_log_prob(random_effect_weights):
  return joint_log_prob(
      random_effect_weights=random_effect_weights,
      log_radon=dtype(radon.log_radon.values),
      floor=dtype(radon.floor.values),
      county=np.int32(radon.county.values),
      log_county_uranium_ppm=log_county_uranium_ppm,
      dtype=dtype)

Agora, computamos o passo E criando um kernel de transição do HMC.

Observações:

• Usamos state_stop_gradient=True para evitar que o passo M faça a retropropagação por meio de obtenções de dados do MCMC (lembre-se de que não precisamos da retropropagação porque nosso passo E é intencionalmente parametrizado nos estimadores anteriores mais bem conhecidos.)

• Usamos tf.placeholder para que, quando executarmos o grafo do TF, possamos usar a amostra do MCMC aleatória da iteração anterior como o valor da cadeia da próxima iteração.

• Usamos a heurística step_size adaptativa do TFP, tfp.mcmc.hmc_step_size_update_fn.

# Set-up E-step.

step_size = tf.get_variable(
    'step_size',
    initializer=np.array(0.2, dtype=dtype),
    trainable=False)

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=unnormalized_posterior_log_prob,
    num_leapfrog_steps=2,
    step_size=step_size,
    step_size_update_fn=tfp.mcmc.make_simple_step_size_update_policy(
      num_adaptation_steps=None),
    state_gradients_are_stopped=True)

init_random_weights = tf.placeholder(dtype, shape=[len(log_county_uranium_ppm)])

posterior_random_weights, kernel_results = tfp.mcmc.sample_chain(
    num_results=3,
    num_burnin_steps=0,
    num_steps_between_results=0,
    current_state=init_random_weights,
    kernel=hmc)

Agora, configuramos o passo M. É basicamente igual ao que um passo de otimização faz no TF.

# Set-up M-step.

loss = -tf.reduce_mean(kernel_results.accepted_results.target_log_prob)

global_step = tf.train.get_or_create_global_step()

learning_rate = tf.train.exponential_decay(
    learning_rate=0.1,
    global_step=global_step,
    decay_steps=2,
    decay_rate=0.99)

optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss, global_step=global_step)

Concluímos com algumas tarefas de organização. Precisamos indicar ao TF que todas as variáveis estão inicializadas. Também criamos identificadores para as variáveis do TF para que possamos exibir seus valores via print em cada iteração do procedimento.

# Initialize all variables.

init_op = tf.initialize_all_variables()

# Grab variable handles for diagnostic purposes.

with tf.variable_scope('make_weights_prior', reuse=True):
  prior_scale = fwd_scale_transform(tf.get_variable(
      name='raw_prior_scale', dtype=dtype))
  
with tf.variable_scope('make_log_radon_likelihood', reuse=True):
  likelihood_scale = fwd_scale_transform(tf.get_variable(
      name='raw_likelihood_scale', dtype=dtype))
  fixed_effect_weights = tf.get_variable(
      name='fixed_effect_weights', dtype=dtype)

6.3 – Execute

Nesta seção, vamos executar nosso grafo do TF relativo ao SAEM. O principal truque é alimentar a última obtenção de dados do kernel do HMC na próxima iteração, o que é feito utilizando feed_dict na chamada sess.run.

init_op.run()
w_ = np.zeros([len(log_county_uranium_ppm)], dtype=dtype)

%%time
maxiter = int(1500)
num_accepted = 0
num_drawn = 0
for i in range(maxiter):
  [
      _,
      global_step_,
      loss_,
      posterior_random_weights_,
      kernel_results_,
      step_size_,
      prior_scale_,
      likelihood_scale_,
      fixed_effect_weights_,
  ] = sess.run([
      train_op,
      global_step,
      loss,
      posterior_random_weights,
      kernel_results,
      step_size,
      prior_scale,
      likelihood_scale,
      fixed_effect_weights,
  ], feed_dict={init_random_weights: w_})
  w_ = posterior_random_weights_[-1, :]
  num_accepted += kernel_results_.is_accepted.sum()
  num_drawn += kernel_results_.is_accepted.size
  acceptance_rate = num_accepted / num_drawn
  if i % 100 == 0 or i == maxiter - 1:
    print('global_step:{:>4}  loss:{: 9.3f}  acceptance:{:.4f}  '
          'step_size:{:.4f}  prior_scale:{:.4f}  likelihood_scale:{:.4f}  '
          'fixed_effect_weights:{}'.format(
              global_step_, loss_.mean(), acceptance_rate, step_size_,
              prior_scale_, likelihood_scale_, fixed_effect_weights_))              

global_step:   0  loss: 1966.948  acceptance:1.0000  step_size:0.2000  prior_scale:1.0000  likelihood_scale:0.8529  fixed_effect_weights:[ 0.  1.]
global_step: 100  loss: 1165.385  acceptance:0.6205  step_size:0.2040  prior_scale:0.9568  likelihood_scale:0.7611  fixed_effect_weights:[ 1.47523439 -0.66043079]
global_step: 200  loss: 1149.851  acceptance:0.6766  step_size:0.2081  prior_scale:0.7465  likelihood_scale:0.7665  fixed_effect_weights:[ 1.48918796 -0.67058587]
global_step: 300  loss: 1163.464  acceptance:0.6811  step_size:0.2040  prior_scale:0.8445  likelihood_scale:0.7594  fixed_effect_weights:[ 1.46291411 -0.67586178]
global_step: 400  loss: 1158.846  acceptance:0.6808  step_size:0.2081  prior_scale:0.8377  likelihood_scale:0.7574  fixed_effect_weights:[ 1.47349834 -0.68823022]
global_step: 500  loss: 1154.193  acceptance:0.6766  step_size:0.1961  prior_scale:0.8546  likelihood_scale:0.7564  fixed_effect_weights:[ 1.47703862 -0.67521363]
global_step: 600  loss: 1163.903  acceptance:0.6783  step_size:0.2040  prior_scale:0.9491  likelihood_scale:0.7587  fixed_effect_weights:[ 1.48268366 -0.69667786]
global_step: 700  loss: 1163.894  acceptance:0.6767  step_size:0.1961  prior_scale:0.8644  likelihood_scale:0.7617  fixed_effect_weights:[ 1.4719094  -0.66897118]
global_step: 800  loss: 1153.689  acceptance:0.6742  step_size:0.2123  prior_scale:0.8366  likelihood_scale:0.7609  fixed_effect_weights:[ 1.47345769 -0.68343043]
global_step: 900  loss: 1155.312  acceptance:0.6718  step_size:0.2040  prior_scale:0.8633  likelihood_scale:0.7581  fixed_effect_weights:[ 1.47426116 -0.6748783 ]
global_step:1000  loss: 1151.278  acceptance:0.6690  step_size:0.2081  prior_scale:0.8737  likelihood_scale:0.7581  fixed_effect_weights:[ 1.46990883 -0.68891817]
global_step:1100  loss: 1156.858  acceptance:0.6676  step_size:0.2040  prior_scale:0.8716  likelihood_scale:0.7584  fixed_effect_weights:[ 1.47386014 -0.6796245 ]
global_step:1200  loss: 1166.247  acceptance:0.6653  step_size:0.2000  prior_scale:0.8748  likelihood_scale:0.7588  fixed_effect_weights:[ 1.47389269 -0.67626756]
global_step:1300  loss: 1165.263  acceptance:0.6636  step_size:0.2040  prior_scale:0.8771  likelihood_scale:0.7581  fixed_effect_weights:[ 1.47612262 -0.67752427]
global_step:1400  loss: 1158.108  acceptance:0.6640  step_size:0.2040  prior_scale:0.8748  likelihood_scale:0.7587  fixed_effect_weights:[ 1.47534692 -0.6789524 ]
global_step:1499  loss: 1161.030  acceptance:0.6638  step_size:0.1941  prior_scale:0.8738  likelihood_scale:0.7589  fixed_effect_weights:[ 1.47624075 -0.67875224]
CPU times: user 1min 16s, sys: 17.6 s, total: 1min 33s
Wall time: 27.9 s

Parece que, após cerca de 1.500 passos, as estimativas dos parâmetros estabilizaram.

6.4 – Resultados

Agora que ajustamos os parâmetros, vamos gerar um grande número de amostras posteriores e estudar os resultados.

%%time
posterior_random_weights_final, kernel_results_final = tfp.mcmc.sample_chain(
    num_results=int(15e3),
    num_burnin_steps=int(1e3),
    current_state=init_random_weights,
    kernel=tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=unnormalized_posterior_log_prob,
      num_leapfrog_steps=2,
      step_size=step_size))

[
    posterior_random_weights_final_,
    kernel_results_final_,
] = sess.run([
    posterior_random_weights_final,
    kernel_results_final,
], feed_dict={init_random_weights: w_})

CPU times: user 1min 42s, sys: 26.6 s, total: 2min 8s
Wall time: 35.1 s

print('prior_scale: ', prior_scale_)
print('likelihood_scale: ', likelihood_scale_)
print('fixed_effect_weights: ', fixed_effect_weights_)
print('acceptance rate final: ', kernel_results_final_.is_accepted.mean())

prior_scale:  0.873799
likelihood_scale:  0.758913
fixed_effect_weights:  [ 1.47624075 -0.67875224]
acceptance rate final:  0.7448

Agora vamos construir um diagrama de caixa estreita do efeito aleatório $\beta_c \log(\text{UraniumPPM}_c)$. Vamos ordenar os efeitos aleatórios pela frequência decrescente de condados.

x = posterior_random_weights_final_ * log_county_uranium_ppm
I = county_freq[:, 0]
x = x[:, I]
cols = np.array(county_name)[I]
pw = pd.DataFrame(x)
pw.columns = cols

fig, ax = plt.subplots(figsize=(25, 4))
ax = pw.boxplot(rot=80, vert=True);

No diagrama de caixa estreita, observamos que a variância do efeito aleatório $\log(\text{UraniumPPM})$ no nível de condado aumenta à medida que o condado é menos representado no dataset, o que faz sentido intuitivamente – devemos ter menos certeza sobre o impacto de um determinado condado se tivermos menos evidências para ele.

7 – Comparativo triplo lado a lado

Agora, vamos comparar os resultados dos três procedimentos. Para isso, vamos computar as estimativas não paramétricas das amostras posteriores geradas pelo Stan e TFP. Também vamos comparar com as estimativas paramétricas (aproximadas) geradas pelo pacote lme4 do R.

O gráfico abaixo mostra a distribuição posterior de cada peso para cada condado de Minnesota. Mostramos os resultados para o Stan (vermelho), TFP (azul) e lme4 do R (laranja). Sombreamos os resultados do Stan e TFP, portanto, você verá em roxo quando os dois coincidirem. Por questões de simplicidade, não sombreamos os resultados do R. Cada subgráfico representa um único condado, e os subgráficos são ordenados em frequência decrescente, da esquerda para a direita e de cima para baixo.

nrows = 17
ncols = 5
fig, ax = plt.subplots(nrows, ncols, figsize=(18, 21), sharey=True, sharex=True)
with warnings.catch_warnings():
  warnings.simplefilter('ignore')
  ii = -1
  for r in range(nrows):
    for c in range(ncols):
      ii += 1
      idx = county_freq[ii, 0]
      sns.kdeplot(
          posterior_random_weights_final_[:, idx] * log_county_uranium_ppm[idx],
          color='blue',
          alpha=.3,
          shade=True,
          label='TFP',
          ax=ax[r][c])
      sns.kdeplot(
          posterior_random_weights_final_stan[:, idx] *
          log_county_uranium_ppm[idx],
          color='red',
          alpha=.3,
          shade=True,
          label='Stan/rstanarm',
          ax=ax[r][c])
      sns.kdeplot(
          posterior_random_weights_lme4_final_[:, idx] *
          log_county_uranium_ppm[idx],
          color='#F4B400',
          alpha=.7,
          shade=False,
          label='R/lme4',
          ax=ax[r][c])
      ax[r][c].vlines(
          posterior_random_weights_lme4[idx] * log_county_uranium_ppm[idx],
          0,
          5,
          color='#F4B400',
          linestyle='--')
      ax[r][c].set_title(county_name[idx] + ' ({})'.format(idx), y=.7)
      ax[r][c].set_ylim(0, 5)
      ax[r][c].set_xlim(-1., 1.)
      ax[r][c].get_yaxis().set_visible(False)
      if ii == 2:
        ax[r][c].legend(bbox_to_anchor=(1.4, 1.7), fontsize=20, ncol=3)
      else:
        ax[r][c].legend_.remove()
  fig.subplots_adjust(wspace=0.03, hspace=0.1)

8 – Conclusão

Neste Colab, ajustamos um modelo de regressão linear de efeitos mistos para o dataset radon. Utilizamos três pacotes de software diferentes: R, Stan e TensorFlow Probability. Concluímos plotando as 85 distribuições posteriores conforme computadas pelos três pacotes.

Apêndice A: HLM do Radon alternativo (com intercepto aleatório)

Nesta seção, vamos descrever um HLM alternativo que também tem um intercepto aleatório associado a cada condado.

$$\begin{align*} \text{for } & c=1\ldots \text{NumCounties}:\\ & \beta_c \sim \text{MultivariateNormal}\left(\text{loc}=\left[ \begin{array}{c} 0 \\ 0 \end{array}\right] , \text{scale}=\left[\begin{array}{cc} \sigma_{11} & 0 \\ \sigma_{12} & \sigma_{22} \end{array}\right] \right) \\ \text{for } & i=1\ldots \text{NumSamples}:\\ & c_i := \text{County}_i \\ &\eta_i = \underbrace{\omega_0 + \omega_1\text{Floor}_i \vphantom{\log( \text{CountyUraniumPPM}_{c_i}))}}_{\text{fixed effects}} + \underbrace{\beta_{c_i,0} + \beta_{c_i,1}\log( \text{CountyUraniumPPM}_{c_i}))}_{\text{random effects}} \\ &\log(\text{Radon}_i) \sim \text{Normal}(\text{loc}=\eta_i , \text{scale}=\sigma) \end{align*}$$

Na "notação de til" do lme4 do R, esse modelo é equivalente a:

log_radon ~ 1 + floor + (1 + log_county_uranium_ppm | county)

Apêndice B: modelos lineares generalizados de efeitos mistos

Nestas seção, mostraremos uma caracterização mais geral dos modelos lineares hierárquicos do que a usada no corpo do documento. Esse modelo mais geral é conhecido como modelo linear generalizado de efeitos mistos (GLMM, na sigla em inglês).

Os GLMMs são generalizações dos modelos lineares generalizados (GLMs). Os GLMMs estendem os GLMs, incorporando ruído aleatório específico das amostras à resposta linear prevista, o que é útil em parte porque permite que características raramente vistas compartilhem informações com características vistas com maior frequência.

Como um processo generativo, um modelo linear generalizado de efeitos mistos (GLMM) é caracterizado por:

ParseError: KaTeX parse error: Expected 'EOF', got '#' at position 70: …e{2.45cm}\text{#̲ para cada grup…

onde:

$$\begin{align} R &= \text{número de grupos de efeitos aleatórios}\ |C_r| &= \text{número de categorias para o grupo $r$}\ N &= \text{número de amostras de treinamento}\ x_i,\omega &\in \mathbb{R}^{D_0}\ D_0 &= \text{número de efeitos fixos}\ C_r(i) &= \text{categoria (sob o grupo $r$) da $i$ésima amostra}\ z_{r,i} &\in \mathbb{R}^{D_r}\ D_r &= \text{número de efeitos aleatórios associados ao grupo $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{função de ligação inversa}\ \text{Distribution} &=\text{uma distribuição parametrizável somente por sua média} \end{align}$$

Basicamente, toda categoria de cada grupo está associada a um iid MVN, $\beta_{rc}$. Embora as obtenções de $\beta_{rc}$ sejam sempre independentes, são distribuídas identicamente somente para um grupo $r$; observe que existe exatamente um $\Sigma_r$ para cada $r\in{1,\ldots,R}$.

Ao combinar com as características de grupo de uma amostra, $z_{r,i}$, o resultado é um ruído específico para a amostra na $i$ésima resposta linear prevista (que também é $x_i^\top\omega$).

Quando estimamos $\{\Sigma_r:r\in\{1,\ldots,R\}\}$ , estamos basicamente estimando a quantidade de ruído que um grupo de efeitos aleatórios transporta, que, caso contrário, abafaria o sinal presente em $x_i^\top\omega$.

Há diversas opções para $\text{Distribution}$ e a função de ligação inversa, $g^{-1}$. Confira algumas escolhas comuns:

• $Y_i\sim\text{Normal}(\text{mean}=\eta_i, \text{scale}=\sigma)$

• ParseError: KaTeX parse error: Expected 'EOF', got '_' at position 81: …i), \text{total_̲count}=n_i)

• $Y_i\sim\text{Poisson}(\text{mean}=\exp(\eta_i))$.

Confira mais possibilidades no módulo tfp.glm.