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.

Regresión lineal de efectos mixtos en {TF Probability, R, Stan}

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

1 Introducción

En este Colab, ajustaremos un modelo de regresión lineal de efectos mixtos a un conjunto de datos de juguetes popular. Haremos que esto se ajuste tres veces, mediante el uso de lme4 de R, el paquete de efectos mixtos de Stan y las primitivas de probabilidad de TensorFlow (TFP). Concluimos con la demostración de que los tres dan aproximadamente los mismos parámetros ajustados y distribuciones posteriores.

Nuestra principal conclusión es que TFP tiene las piezas generales necesarias para adaptarse a modelos similares al HLM y que produce resultados que son consistentes con otros paquetes de software, es decir, lme4, rstanarm. Este Colab no es un reflejo exacto de la eficiencia computacional de ninguno de los paquetes 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 lineal jerárquico

Para nuestra comparación entre R, Stan y TFP, ajustaremos un modelo lineal jerárquico (HLM) al conjunto de datos radon que se popularizó en el análisis de datos bayesiano de Gelman, et. Alabama. (página 559, segunda ed.; página 250, tercera ed.).

Se asume el siguiente 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*}$$

En la "notación de tilde" lme4 de R, este modelo es equivalente a:

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

Encontraremos el MLE para $\omega, \sigma_C, \sigma_N$ mediante el uso de la distribución posterior (condicionada a la evidencia) de $\{\beta_c\}_{c=1}^\text{NumCounties}$.

Si se desea obtener básicamente el mismo modelo, pero con una intersección aleatoria, consulte el Apéndice A.

Para obtener una especificación más general del HLM, consulte el Apéndice B.

3 Manipulación de datos

En esta sección, obtenemos el conjunto de datos radon y ejecutamos un preprocesamiento mínimo para que cumpla con nuestro modelo asumido.

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 Conocimiento de los datos

En esta sección, exploramos el conjunto de datos radon para entender mejor por qué el modelo propuesto podría ser razonable.

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)

Conclusiones:

• Hay una larga cola de 85 condados. (Una ocurrencia común en los GLMM).

• De hecho, $\log(\text{Radon})$ no tiene restricciones. (Por tanto, la regresión lineal podría tener sentido).

• La mayoría de las lecturas se realizan en el piso $0$; no se realizó ninguna lectura por encima del piso $1$. (Por lo tanto, nuestros efectos fijos solo tendrán dos ponderaciones).

4 HLM en R

En esta sección se usa el paquete lme4 de R para ajustar el modelo probabilístico descrito anteriormente.

NOTA: Para ejecutar esta sección, se debe cambiar a un tiempo de ejecución de Colab de 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 en Stan

En esta sección se usa rstanarm para ajustar un modelo Stan con la misma fórmula/sintaxis que el modelo lme4 anterior.

A diferencia de lme4 y el siguiente modelo de TF, rstanarm es un modelo completamente bayesiano, es decir, se supone que todos los parámetros se extraen de una distribución Normal y los propios parámetros se extraen de una distribución.

NOTA: Para ejecutar esta sección, se debe cambiar a un tiempo de ejecución de Colab de 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)

Nota: Los tiempos de ejecución provienen de un único núcleo de CPU. (Este Colab no pretende ser una representación fiel del tiempo de ejecución de Stan ni de 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")

Nota: Vuelva al tiempo de ejecución del núcleo de TF de 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 las estimaciones puntuales y las desviaciones estándar condicionales para los efectos aleatorios del grupo de lme4 para visualizarlas más adelante.

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,)

Extraiga muestras para las ponderaciones de los condados con ayuda de las medias estimadas y las desviaciones estándar de 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)

También recuperamos las muestras posteriores de las ponderaciones del condado del ajuste de 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)

En este ejemplo de Stan se muestra cómo se implementaría LMER en un estilo más cercano a TFP, es decir, al especificar directamente el modelo probabilístico.

6 HLM en TF Probability

En esta sección usaremos primitivas de probabilidad de TensorFlow de bajo nivel (Distributions) para especificar nuestro modelo lineal jerárquico y ajustar los parámetros desconocidos.

# 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 Especificación del modelo

En esta sección especificamos el modelo lineal de efectos mixtos de radón con ayuda de las primitivas de TFP. Para hacer esto, especificamos dos funciones que producen dos distribuciones de TFP:

• make_weights_prior: un valor previo normal multivariado para las ponderaciones aleatorias (que se multiplican por $\log(\text{UraniumPPM}_{c_i})$ para calcular el predictor lineal).

• make_log_radon_likelihood: un lote de distribuciones Normal sobre cada variable dependiente $\log(\text{Radon}_i)$ observada.

Dado que ajustaremos los parámetros de cada una de estas distribuciones, debemos usar variables de TF (es decir, tf.get_variable). Sin embargo, dado que deseamos utilizar optimización sin restricciones, debemos encontrar una manera de restringir los valores reales para lograr la semántica necesaria, por ejemplo, positivos que representen desviaciones estándar.

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

La siguiente función construye nuestra previa, $p(\beta|\sigma_C)$ donde $\beta$ denota los pesos de efectos aleatorios y $\sigma_C$ la desviación estándar.

Usamos tf.make_template para garantizar que la primera llamada a esta función cree una instancia de las variables de TF que usa y que todas las llamadas posteriores reutilicen el valor actual de la variable.

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)

La siguiente función construye nuestra probabilidad, $p(y|x,\omega,\beta,\sigma_N)$ donde $y,x$ denota respuesta y evidencia, $\omega,\beta$ denota ponderaciones de efectos fijos y aleatorios, y $\sigma_N$ la desviación estándar.

Aquí nuevamente usamos tf.make_template para garantizar que las variables de TF se reutilicen en las llamadas.

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)

Finalmente, recurrimos a los generadores de previa y de probabilidad para construir la densidad 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 Entrenamiento (aproximación estocástica de maximización de expectativas)

Para ajustar nuestro modelo de regresión lineal de efectos mixtos, utilizaremos una versión de aproximación estocástica del algoritmo de maximización de expectativas (SAEM). La idea básica consiste en el uso de muestras de la posterior para aproximar la densidad logarítmica conjunta esperada (paso E). Luego, encontramos los parámetros que maximizan este cálculo (paso M). De manera algo más concreta, la iteración de punto fijo viene 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*}$$

donde $x$ denota evidencia, $Z$ alguna variable latente que necesita ser marginada y $\theta,\theta_0$ posibles parametrizaciones.

Para obtener una explicación más detallada, consulte Convergencia de una versión de aproximación estocástica de los algoritmos de maximización de expectativas de Bernard Delyon, Marc Lavielle, Eric, Moulines (Ann. Statist., 1999).

Para calcular el paso E, necesitamos tomar una muestra de la posterior. Dado que no es fácil tomar muestras de nuestra posterior, utilizamos el Hamiltoniano de Monte Carlo (HMC). El HMC es un procedimiento de método de Monte Carlo basado en cadenas de Markov que usa gradientes (estado wrt, no parámetros) de la densidad logarítmica posterior no normalizada para proponer nuevas muestras.

Especificar la densidad logarítmica posterior no normalizada es simple: es solo la densidad logarítmica conjunta "fijada" en cualquier cosa que queramos condicionar.

# 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)

Ahora completamos la configuración del paso E mediante la creación de un núcleo de transición del HMC.

Notas:

• Usamos state_stop_gradient=True para evitar que el paso M retroceda a través de extracciones del MCMC. (Recuerde, no necesitamos retroceder porque nuestro paso E está parametrizado intencionalmente en los estimadores previos más conocidos).

• Usamos tf.placeholder para que cuando finalmente ejecutemos nuestro grafo de TF, podamos alimentar la muestra MCMC aleatoria de la iteración anterior como el valor de la cadena de la siguiente iteración.

• Usamos la heurística adaptativa step_size de 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)

Ahora configuramos el paso M. Esto es esencialmente lo mismo que una optimización que uno podría hacer en 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)

Concluimos con algunas tareas de limpieza. Debemos decirle a TF que todas las variables están inicializadas. También creamos identificadores para nuestras variables de TF para que podamos print sus valores en cada iteración del procedimiento.

# 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 Ejecución

En esta sección ejecutamos nuestro grafo de SAEM de TF. El truco principal aquí es introducir nuestro último extracto del núcleo del HMC en la siguiente iteración. Esto se logra mediante nuestro uso de feed_dict en la llamada 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

Al parecer, después de ~1500 pasos, nuestras estimaciones de los parámetros se han estabilizado.

6.4 Resultados

Ahora que hemos ajustado los parámetros, generemos una gran cantidad de muestras de la posterior y estudiemos los 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

Ahora construimos un diagrama de caja y bigotes del efecto aleatorio $\beta_c \log(\text{UraniumPPM}_c)$. Ordenaremos los efectos aleatorios mediante la disminución de la frecuencia del condado.

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);

En este diagrama de caja y bigotes, observamos que la varianza del efecto aleatorio $\log(\text{UraniumPPM})$ a nivel de condado aumenta a medida que el condado tiene menos representación en el conjunto de datos. Intuitivamente, esto tiene sentido: deberíamos estar menos seguros sobre el impacto de un determinado condado si tenemos menos evidencia de ello.

7 Comparación lado a lado

Ahora comparamos los resultados de los tres procedimientos. Para hacer esto, calcularemos estimaciones no paramétricas de las muestras de la posterior generadas por Stan y TFP. También compararemos con las estimaciones paramétricas (aproximadas) producidas por el paquete lme4 de R.

El siguiente gráfico muestra la distribución posterior de cada ponderación para cada condado de Minnesota. Mostramos resultados para Stan (rojo), TFP (azul) y R's lme4 (naranja). Sombreamos los resultados de Stan y TFP, por lo que esperamos ver violeta cuando los dos coincidan. Para simplificar, no sombreamos los resultados de R. Cada subgráfico representa un solo condado y está ordenado en frecuencia descendente en orden de escaneo ráster (es decir, de izquierda a derecha y luego de arriba a abajo).

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 Conclusión

En este Colab ajustamos un modelo de regresión lineal de efectos mixtos al conjunto de datos radón. Probamos tres paquetes de software diferentes: R, Stan y TensorFlow Probability. Concluimos con el trazado de las 85 distribuciones posteriores calculadas por los tres distintos paquetes de software.

Apéndice A: HLM de radón alternativo (agregar intercepción aleatoria)

En esta sección describimos un HLM alternativo que también tiene una intersección aleatoria que se asocia con 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*}$$

En la "notación de tilde" lme4 de R, este modelo es equivalente a:

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

Apéndice B: Modelos lineales generalizados de efectos mixtos

En esta sección se ofrece una caracterización más general de los modelos lineales jerárquicos que la que se usa en el cuerpo principal. Este modelo más general se conoce como modelo lineal generalizado de efectos mixtos (GLMM).

Los GLMM son generalizaciones de modelos lineales generalizados (GLM). Los GLMM amplían los GLM mediante la incorporación de ruido aleatorio específico de la muestra en la respuesta lineal predicha. Esto es útil en parte porque permite que las características poco vistas compartan información con las características que se ven con más frecuencia.

Como proceso generativo, un Modelo Lineal Generalizado de Efectos Mixtos (GLMM) se caracteriza por lo siguiente:

ParseError: KaTeX parse error: Expected 'EOF', got '#' at position 70: …e{2.45cm}\text{#̲ for each rando…

donde:

$$\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}$$

En palabras, esto significa que cada categoría de cada grupo está asociada con una MVN i. i. d., $\beta_{rc}$. Aunque las extracciones $\beta_{rc}$ son siempre independientes, solo se distribuyen de forma idéntica para un grupo $r$; tenga en cuenta que hay exactamente un $\Sigma_r$ por cada $r\in{1,\ldots,R}$.

Cuando se combina de forma afín con las características del grupo de una muestra, $z_{r,i}$, el resultado es un ruido específico de la muestra en la $i$-ésima respuesta lineal predicha (que de otro modo sería $x_i^\top\omega$).

Cuando estimamos $\{\Sigma_r:r\in\{1,\ldots,R\}\}$ esencialmente estamos estimando la cantidad de ruido que transporta un grupo de efectos aleatorios que, de otro modo, ahogaría la señal presente en $x_i^\top\omega$.

Hay una variedad de opciones para la función $\text{Distribution}$ y la función de enlace inverso, $g^{-1}$. Las opciones comunes son estas:

• $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), and,

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

Si desea conocer más posibilidades, consulte el módulo tfp.glm.