10.8 HMC y Stan
It appears to be quite a general principle that, whenever there is a randomized way of doinf something, then there is a nonrandomized way that delivers better performance but requires more thought. -E.T. Jaynes
Stan es un programa para generar muestras de una distribución posterior de los parámetros de un modelo, el nombre del programa hace referencia a Stanislaw Ulam (1904-1984) que fue pionero en los métodos de Monte Carlo. A diferencia de JAGS y BUGS, los pasos de la cadena de Markov se generan con un método llamado Monte Carlo Hamiltoniano (HMC). HMC es computacionalmente más costoso que Metropolis o Gibbs, sin embargo, sus propuestas suelen ser más eficientes, y por consiguiente no necesita muestras tan grandes. En particular cuando se ajustan modelos grandes y complejos (por ejemplo, con variables con correlación alta) HMC supera a otros.
Muestreo HMC
El uso de HMC en estadística es reciente, sin embargo, gracias a Stan se ha expandido rápidamente tanto en academia como industria. Desafortunadamente, la teoría de HMC está desarrollada en términos de geometría diferencial, lo que hace que su construcción formal requiera de matemáticas avanzadas. En estas notas se presentan las ideas detrás de HMC siguiendo Kruschke (2015), una referencia con mayor detalle es Betancourt (2017) y para el uso de Stan vale la pena tener siempre a la mano el manual.
Stan genera muestras de la posterior usando una variación del algoritmo de Metrópolis. Recordemos como funciona el algoritmo de Metrópolis que vimos en clase:
Tenemos una distribución objetivo \(p(\theta)\) de la cual buscamos generar muestras. Debemos ser capaces de calcular el valor de \(p(\theta)\) para cualquier valor candidato \(\theta\). La distribución objetivo \(p(\theta)\) no tiene que estar normalizada, típicamente \(p(\theta)\) es la distribución posterior de \(\theta\) no normalizada, es decir, es el producto de la verosimilitud y la inicial.
- La muestra de la distribución objetivo se genera mediante una caminata
aleatoria a través del espacio de parámetros.
- La caminata inicia en un lugar arbitrario (definido por el usuario). El punto inicial debe ser tal que \(p(\theta)>0\).
- La caminata avanza en cada tiempo proponiendo un movimiento a una nueva posición y después decidiendo si se acepta o no el valor propuesto. Las distribuciones propuesta pueden tener muchas formas, el objetivo es que la distribución propuesta explore el espacio de parámetros de manera eficiente.
Una vez que tenemos un valor propuesto decidimos si aceptar calculando:
\[p_{mover}=min\bigg( \frac{p(\theta_{propuesta})}{p(\theta_{actual})},1\bigg)\]
Y al final obtenemos valores representativos de la distribución objetivo \(\{\theta_1,...,\theta_n\}\).
Notemos que en la versión de Metrópolis que estudiamos, la forma de la distribución propuesta está centrada de manera simétrica en la posición actual. Es decir, en un espacio paramétrico multidimensional, la distribución propuesta podría ser una Normal multivariada, con la matriz de varianzas y covarianzas seleccionada para mejorar la eficiencia en la aplicación particular. La normal multivariada siempre esta centrada en la posición actual y siempre tiene la misma forma, sin importar en que sección del espacio paramétrico estemos ubicados. Esto puede llevar a ineficiencias, por ejemplo, si nos ubicamos en las colas de la distribución posterior el paso propuesto con la misma probabilidad nos aleja o acerca de la moda de la posterior. Otro ejemplo es si la distriución posterior se curva a lo largo del espacio paramétrico, una distribución propuesta (de forma fija) puede ser eficiente para una parte de la posterior y poco eficiente para otra parte de la misma.
Por su parte HMC, usa una distribución propuesta que cambia dependiendo de la posición actual. HMC utiliza el gradiente de la posterior y envuelve la distribución propuesta hacia el gradiente, como se ve en la siguiente figura.
Figure 10.1: copyright(c) Kruschke, J. K. (2014). Doing Bayesian Data Analysis: A Tutorial with R, JAGS, and Stan. 2nd Edition. Academic Press / Elsevier
HMC genera un movimiento propuesta de manera análoga a rodar una canica en la distribución posterior volteada (también conocida como potencial). El potencial es el negativo del logaritmo de la densidad posterior, en las regiones donde la posterior es alta el potencial es bajo, y en las la regiones donde la posterior es plana el potencial es alto.
La propuesta se genera dando un golpecito la canica en una dirección aleatoria y dejándola rodar cierto tiempo. En el caso del ejemplo de un solo parámetro la dirección del golpecito inicial es hacia la izquierda o derecha, y la magnitud se genera de manera aleatoria muestreando de una distrubución Gaussiana de media cero. El golpecito impone un momento inicial a la canica, y al terminar el tiempo se propone al algoritmo de Metrópilis la posición final de la canica. Es fácil imaginar que la posición propuesta tenderá a estar en regiones de mayor probabilidad posterior.
La última fila de la figura de arriba nos muestra un histograma de los pasos propuestos, notemos que no está centrado en la posición actual sino que hay una inclinación hacia la moda de la posterior.
Para distribuciones posteriores de dimensión alta con valles diagonales o curveados, la dinámica de HMC generará valores propuesta mucho más prometedores que una distribución propuesta simétrica (como la versión de Metrópolis que implementamos) y mejores que un muestreador de Gibbs que puede atorarse en paredes diagonales.
Es así que para pasar del algoritmo de Metrópolis que estudiamos a HMC se modifica la probabilidad de aceptación para tener en cuenta no sola la densidad posterior relativa, sino también el momento (denotado por \(\phi\)) en las posiciones actual y propuesta.
\[p_{aceptar}=min\bigg\{\frac{p(\theta_{propuesta}|x)p(\phi_{propuesta})}{p(\theta_{actual}|x)p(\phi_{actual})}, 1 \bigg\}\]
En un sistema continuo ideal la suma de la energía potencial y cinética (que corresponden a \(-log(p(\theta|x))\) y \(-log(p(\phi))\)) es constante y por tanto aceptaríamos todas las propuestas. Sin embargo, en la práctica las dinámicas continuas se dicretizan a intervalos en el tiempo y los cálculos son solo aproximados conllevando a que se rechacen algunas propuestas.
Si los pasos discretizados son pequeños entonces la aproximación será buena pero se necesitarán más pasos para alejarse de una posición dada, y lo contrario si los pasos son muy grandes. Por tanto se debe ajustar el tamaño del paso (\(\varepsilon\)) y el número de pasos. La duración de la trayectoria es el producto del tamaño del paso y el número de pasos. Es usual buscar una tasa de aceptación alrededor del \(65\%\), moviendo el tamaño de epsilon y compensando con el número de pasos.
Es así que el tamaño del paso controla la suavidad de la trayectoria. También es importante ajustar el número de pasos (es decir la duración del movimiento) pues no queremos alejarnos demasiado o rodar de regreso a la posición actual. La siguiente figura muestra varias trayectorias y notamos que muchas rebotan al lugar inicial.
Figure 10.2: copyright(c) Kruschke, J. K. (2014). Doing Bayesian Data Analysis: A Tutorial with R, JAGS, and Stan. 2nd Edition. Academic Press / Elsevier
Para evitar las ineficiencias de dar vueltas en U, Stan incorpora un algoritmo que generaliza la nación de vueltas en U a espacios de dimensión alta, y así estima cuando parar las trayectorias antes de que reboten hacia la posición inical. El algoritmo se llama No U-turn Sampler (NUTS).
Adicional a ajustar el tamaño del paso y número de pasos debemos ajustar la desviación estándar del momento inicial. Si la desviación estandar del momento es grande también lo será la desviación estándar de las propuestas. Nuevamente, lo más eficiente será una desviación estándar ni muy grande ni muy chica. En Stan la desviación estándar del momento se establece de manera adaptativa para que coincida con la desviación estándar de la posterior.
Figure 10.3: copyright(c) Kruschke, J. K. (2014). Doing Bayesian Data Analysis: A Tutorial with R, JAGS, and Stan. 2nd Edition. Academic Press / Elsevier
Por último, para calcular la trayectoria propuesta debemos ser capaces de calcular el gradiente de la densidad posterior en cualquier valor de los parámetros. Para realizar esto de manera eficiente en espacios de dimensión alta se debe derivar analíticamente, en el caso de modelos complejos las fórmulas se derivan usando algoritmos avanzados.
El paper A Conceptual Introduction to Hamiltonian Monte Carlo de Michael Betancourt explica conceptos e intuición detrás de HMC, y el porqué es apropiado en problemas de alta dimensión.
Stan
Para instalar Stan sigue las instrucciones de aquí.
Nosotros usaremos el paquete rstan, Guo, Gabry, and Goodrich (2019).
library(rstan)
# opcional para correr en paralelo
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())En Stan los programas se organizan mediante secuencias de bloques, cada uno de estos bloques inicia con declaración de variables y después le siguen enunciados.
El siguiente esqueleto ejemplifica los bloques disponibles, sin embargo, veremos que no siempre se utilizan todos los bloques.
functions {
// ... function declarations and definitions ...
}
data {
// ... declarations ...
}
transformed data {
// ... declarations ... statements ...
}
parameters {
// ... declarations ...
}
transformed parameters {
// ... declarations ... statements ...
}
model {
// ... declarations ... statements ...
}
generated quantities {
// ... declarations ... statements ...
}Comenzamos con el ejemplo sencillo de estimar el sesgo de una moneda. El primer paso es especificar el modelo en el lenguaje de Stan. Lo podemos guardar en un archivo de texto separado o simplemente asignarlo a una variable en R.
modelo_bernoulli.stan <-
'
data {
int<lower=0> N;
int y[N];
}
parameters {
real<lower=0,upper=1> theta;
}
model {
theta ~ beta(1,1) ;
y ~ bernoulli(theta) ;
}
'
# notemos que los modelos de Stan deben terminar con una línea en blanco
cat(modelo_bernoulli.stan, file = "src/stan_files/modelo_bernoulli.stan")Notemos que Stan permite operaciones vectorizadas, por lo que podemos poner:
y ~ bernoulli(theta) ;sin necesidad de hacer el ciclo for:
for ( i in 1:N ) {
y[i] ~ dbern(theta)
}Una vez que especificamos el modelo lo siguiente es traducir el modelo a código
de C++ y compilarlo. Para esto usamos la función stan_model() que puede
recibir el archivo de texto con la especificación del modelo.
O el objeto de R,
El paso de compilación puede tardar, pues Stan está calculando los gradientes para las dinámicas Hamiltonianas.
Ahora cargamos los datos y usamos la función sampling para obtener las
simulaciones de la distribución posterior.
N = 50 ; z = 10 ; y = c(rep(1, z), rep(0, N - z))
set.seed(8979)
data_list <- list(y = y, N = N )
stan_fit <- sampling(object = stan_cpp, data = data_list, chains = 3 ,
iter = 1000 , warmup = 200, thin = 1 )La función summary() nos da resúmenes de la distribución posterior de los
parámetros combinando las simulaciones de todas las cadenas y por cadena.
stan_fit
#> Inference for Stan model: modelo_bernoulli.
#> 3 chains, each with iter=1000; warmup=200; thin=1;
#> post-warmup draws per chain=800, total post-warmup draws=2400.
#>
#> mean se_mean sd 2.5% 25% 50% 75% 97.5% n_eff Rhat
#> theta 0.21 0.00 0.06 0.11 0.17 0.21 0.25 0.33 864 1
#> lp__ -27.35 0.02 0.75 -29.49 -27.52 -27.07 -26.88 -26.83 915 1
#>
#> Samples were drawn using NUTS(diag_e) at Tue Dec 10 23:24:41 2019.
#> For each parameter, n_eff is a crude measure of effective sample size,
#> and Rhat is the potential scale reduction factor on split chains (at
#> convergence, Rhat=1).Una alternativa a usar las funciones stan_model() y sampling(), es compilar
y correr las cadenas de manera simultanea con la función stan().
stan_fit_2 <- stan(file = 'src/stan_files/modelo_bernoulli.stan', data = data_list,
chains = 3, iter = 1000 , warmup = 200, thin = 1)Podemos graficar las cadenas con la función traceplot(), esta gráfica nos
permire inspecconar la conducta del muestreador y evaluar si las cadenas se han
mezclado y olvidado el valor inicial.
Ejemplo normal
Recordemos ahora el ejemplo normal con media y varianza desconocidas, para este problema escribimos un muestreador de Gibbs, y ahora veremos como lo haríamos con Stan y compararemos los resultados.
- Escribimos el modelo en Stan:
modelo_normal.stan <-
'
data {
int<lower=0> N;
vector[N] y;
}
parameters {
real mu;
real<lower=0> sigma2;
}
model {
y ~ normal(mu, sqrt(sigma2));
mu ~ normal(1.5, 4);
sigma2 ~ inv_gamma(3, 3);
}
'
cat(modelo_normal.stan, file = "src/stan_files/modelo_normal.stan")Especificamos la verosimilitud normal con media \(\mu\) y varianza \(\sigma^2\), notemos que en Stan los parámetros son la media y desviación estándar.
y ~ normal(mu, sqrt(sigma2));Y al igual que en el ejemplo del muestreador de Gibbs usaremos iniciales Normal para \(\mu\) y Gamma inversa para \(\sigma^2\).
mu ~ normal(1.5, 4);
sigma2 ~ inv_gamma(3, 3);Pasamos al paso de compilación:
El modelo ya esta especificado, pero aun falta indicar los datos observados.
N <- 50 # Observamos 20 realizaciones
set.seed(122)
y <- rnorm(N, 2, 2)
data_list <- list(y = y, N = N )
norm_fit <- sampling(object = stan_norm_cpp, data = data_list,
chains = 3 , iter = 1000 , warmup = 500)Y podemos ver intervalos de la distribución posterior de los parámetros.
norm_fit
#> Inference for Stan model: modelo_normal.
#> 3 chains, each with iter=1000; warmup=500; thin=1;
#> post-warmup draws per chain=500, total post-warmup draws=1500.
#>
#> mean se_mean sd 2.5% 25% 50% 75% 97.5% n_eff Rhat
#> mu 1.91 0.01 0.31 1.30 1.70 1.92 2.12 2.53 1071 1.00
#> sigma2 4.72 0.03 0.92 3.26 4.07 4.62 5.26 6.81 981 1.00
#> lp__ -71.04 0.05 0.99 -73.64 -71.39 -70.76 -70.33 -70.07 455 1.01
#>
#> Samples were drawn using NUTS(diag_e) at Tue Dec 10 23:25:34 2019.
#> For each parameter, n_eff is a crude measure of effective sample size,
#> and Rhat is the potential scale reduction factor on split chains (at
#> convergence, Rhat=1).Podemos realizar graficas de las cadenas con más detalle (y con facilidad),
usando el paquete bayesplot.
library(bayesplot)
#> Error in library(bayesplot): there is no package called 'bayesplot'
norm_posterior_inc_warmup <- extract(norm_fit, inc_warmup = TRUE,
permuted = FALSE)
color_scheme_set("mix-blue-red")
#> Error in color_scheme_set("mix-blue-red"): could not find function "color_scheme_set"
mcmc_trace(norm_posterior_inc_warmup, pars = c("mu", "sigma2"), n_warmup = 300,
facet_args = list(nrow = 2, labeller = label_parsed)) +
facet_text(size = 15)
#> Error in mcmc_trace(norm_posterior_inc_warmup, pars = c("mu", "sigma2"), : could not find function "mcmc_trace"Y podemos graficar la distribución posterior de los parámetros, con intervalos de probabilidad.
norm_posterior <- as.array(norm_fit)
mcmc_areas(norm_posterior, pars = c("mu", "sigma2"), prob = 0.8,
point_est = "median", adjust = 1.4)
#> Error in mcmc_areas(norm_posterior, pars = c("mu", "sigma2"), prob = 0.8, : could not find function "mcmc_areas"
Realiza un histograma de la distribución predictiva
posterior. Construye un intervalo de \(95\%\) de probabilidad para una predicción.
Tip: usa la función extract() para extraer las simulaciones del objeto
stanfit.
Referencias
Betancourt, Michael. 2017. “A Conceptual Introduction to Hamiltonian Monte Carlo.” http://arxiv.org/abs/1701.02434.
Guo, Jiqiang, Jonah Gabry, and Ben Goodrich. 2019. Rstan: R Interface to Stan. https://CRAN.R-project.org/package=rstan.
Kruschke, John. 2015. Doing Bayesian Data Analysis (Second Edition). Boston: Academic Press.