Eight_Schools.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.

Ocho escuelas

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

El problema de las ocho escuelas (Rubin 1981) considera la efectividad de los programas de entrenamiento del SAT realizados en paralelo en ocho escuelas. Se ha convertido en un problema clásico (Análisis de datos bayesianos, Stan) que ilustra la utilidad del modelado jerárquico para compartir información entre grupos intercambiables.

La implementación que se muestra a continuación es una adaptación de un tutorial de Edward 1.0.

Importaciones

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
import warnings

tf.enable_v2_behavior()

plt.style.use("ggplot")
warnings.filterwarnings('ignore')

Los datos

Del análisis de datos bayesiano, sección 5.5 (Gelman et al. 2013):

Se realizó un estudio para el Educational Testing Service con el fin de analizar los efectos de los programas especiales de preparación para el SAT-V (Scholastic Aptitude Test-Verbal) en cada uno de las ocho escuelas secundarias. La variable de resultado en cada estudio fue la puntuación obtenida en una administración especial del SAT-V, del SAT-V, un examen estandarizado de opción múltiple administrado por el Educational Testing Service y uusado para ayudar a las universidades a tomar decisiones de admisión; las puntuaciones pueden variar entre 200 y 800, con una media de aproximadamente 500 y una desviación estándar de aproximadamente 100. Los exámenes SAT están diseñados para que sean resistentes a los esfuerzos a corto plazo dirigidos específicamente a mejorar el rendimiento en la prueba; en su lugar, están diseñados para reflejar los conocimientos adquiridos y las capacidades desarrolladas a lo largo de muchos años de educación. Sin embargo, cada una de las ocho escuelas en este estudio consideró que su programa de entrenamiento a corto plazo fue muy exitoso en el aumento de los puntajes del SAT. Además, no había ninguna razón previa para creer que alguno de los ocho programas fuera más efectivo que cualquier otro o que algunos tuvieran efectos más similares entre sí que con cualquier otro.

Para cada una de las ocho escuelas ($J = 8$), tenemos un efecto de tratamiento estimado $y_j$ y un error estándar de la estimación del efecto $\sigma_j$. Los efectos del tratamiento en el estudio se obtuvieron mediante una regresión lineal en el grupo de tratamiento utilizando las puntuaciones PSAT-M y PSAT-V como variables de control. Como no había ninguna creencia previa de que alguna de las escuelas fuera más o menos similar o que alguno de los programas de entrenamiento fuera más efectivo, podemos considerar los efectos del tratamiento como intercambiables.

num_schools = 8  # number of schools
treatment_effects = np.array(
    [28, 8, -3, 7, -1, 1, 18, 12], dtype=np.float32)  # treatment effects
treatment_stddevs = np.array(
    [15, 10, 16, 11, 9, 11, 10, 18], dtype=np.float32)  # treatment SE

fig, ax = plt.subplots()
plt.bar(range(num_schools), treatment_effects, yerr=treatment_stddevs)
plt.title("8 Schools treatment effects")
plt.xlabel("School")
plt.ylabel("Treatment effect")
fig.set_size_inches(10, 8)
plt.show()

Modelo

Para capturar los datos, usamos un modelo normal jerárquico. Sigue el proceso generativo,

$$\begin{align*} \mu &\sim \text{Normal}(\text{loc}{=}0,\ \text{scale}{=}10) \\ \log\tau &\sim \text{Normal}(\text{loc}{=}5,\ \text{scale}{=}1) \\ \text{for } & i=1\ldots 8:\\ & \theta_i \sim \text{Normal}\left(\text{loc}{=}\mu,\ \text{scale}{=}\tau \right) \\ & y_i \sim \text{Normal}\left(\text{loc}{=}\theta_i,\ \text{scale}{=}\sigma_i \right) \end{align*}$$

donde $\mu$ representa el efecto promedio del tratamiento anterior y $\tau$ controla cuánta variación hay entre escuelas. Se observan $y_i$ y $\sigma_i$. Como $\tau \rightarrow \infty$, el modelo se acerca al modelo sin agrupación, es decir, se permite que cada una de las estimaciones del efecto del tratamiento escolar sea más independiente. Como $\tau \rightarrow 0$, el modelo se aproxima al modelo de agrupación completa, es decir, todos los efectos del tratamiento escolar están más cerca del promedio del grupo $\mu$. Para restringir la desviación estándar para que sea positiva, extraemos $\tau$ de una distribución lognormal (lo que equivale a extraer $log(\tau)$ de una distribución normal).

De acuerdo con el Diagnóstico de inferencia sesgada con divergencias, transformamos el modelo anterior en un modelo no centrado equivalente:

$$\begin{align*} \mu &\sim \text{Normal}(\text{loc}{=}0,\ \text{scale}{=}10) \\ \log\tau &\sim \text{Normal}(\text{loc}{=}5,\ \text{scale}{=}1) \\ \text{for } & i=1\ldots 8:\\ & \theta_i' \sim \text{Normal}\left(\text{loc}{=}0,\ \text{scale}{=}1 \right) \\ & \theta_i = \mu + \tau \theta_i' \\ & y_i \sim \text{Normal}\left(\text{loc}{=}\theta_i,\ \text{scale}{=}\sigma_i \right) \end{align*}$$

Cosificamos este modelo como una instancia JointDistributionSequential:

model = tfd.JointDistributionSequential([
  tfd.Normal(loc=0., scale=10., name="avg_effect"),  # `mu` above
  tfd.Normal(loc=5., scale=1., name="avg_stddev"),  # `log(tau)` above
  tfd.Independent(tfd.Normal(loc=tf.zeros(num_schools),
                             scale=tf.ones(num_schools),
                             name="school_effects_standard"),  # `theta_prime` 
                  reinterpreted_batch_ndims=1),
  lambda school_effects_standard, avg_stddev, avg_effect: (
      tfd.Independent(tfd.Normal(loc=(avg_effect[..., tf.newaxis] +
                                      tf.exp(avg_stddev[..., tf.newaxis]) *
                                      school_effects_standard),  # `theta` above
                                 scale=treatment_stddevs),
                      name="treatment_effects",  # `y` above
                      reinterpreted_batch_ndims=1))
])

def target_log_prob_fn(avg_effect, avg_stddev, school_effects_standard):
  """Unnormalized target density as a function of states."""
  return model.log_prob((
      avg_effect, avg_stddev, school_effects_standard, treatment_effects))

Inferencia bayesiana

A partir de los datos, aplicamos el algoritmo hamiltoniano de Monte Carlo (HMC) para calcular la distribución posterior sobre los parámetros del modelo.

num_results = 5000
num_burnin_steps = 3000

# Improve performance by tracing the sampler using `tf.function`
# and compiling it using XLA.
@tf.function(autograph=False, jit_compile=True)
def do_sampling():
  return tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          tf.zeros([], name='init_avg_effect'),
          tf.zeros([], name='init_avg_stddev'),
          tf.ones([num_schools], name='init_school_effects_standard'),
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          step_size=0.4,
          num_leapfrog_steps=3))

states, kernel_results = do_sampling()

avg_effect, avg_stddev, school_effects_standard = states

school_effects_samples = (
    avg_effect[:, np.newaxis] +
    np.exp(avg_stddev)[:, np.newaxis] * school_effects_standard)

num_accepted = np.sum(kernel_results.is_accepted)
print('Acceptance rate: {}'.format(num_accepted / num_results))

Acceptance rate: 0.5974

fig, axes = plt.subplots(8, 2, sharex='col', sharey='col')
fig.set_size_inches(12, 10)
for i in range(num_schools):
  axes[i][0].plot(school_effects_samples[:,i].numpy())
  axes[i][0].title.set_text("School {} treatment effect chain".format(i))
  sns.kdeplot(school_effects_samples[:,i].numpy(), ax=axes[i][1], shade=True)
  axes[i][1].title.set_text("School {} treatment effect distribution".format(i))
axes[num_schools - 1][0].set_xlabel("Iteration")
axes[num_schools - 1][1].set_xlabel("School effect")
fig.tight_layout()
plt.show()

print("E[avg_effect] = {}".format(np.mean(avg_effect)))
print("E[avg_stddev] = {}".format(np.mean(avg_stddev)))
print("E[school_effects_standard] =")
print(np.mean(school_effects_standard[:, ]))
print("E[school_effects] =")
print(np.mean(school_effects_samples[:, ], axis=0))

E[avg_effect] = 5.57183933258
E[avg_stddev] = 2.47738981247
E[school_effects_standard] =
0.08509017
E[school_effects] =
[15.0051     7.103311   2.4552586  6.2744603  1.3364682  3.1125953
 12.762501   7.743602 ]

# Compute the 95% interval for school_effects
school_effects_low = np.array([
    np.percentile(school_effects_samples[:, i], 2.5) for i in range(num_schools)
])
school_effects_med = np.array([
    np.percentile(school_effects_samples[:, i], 50) for i in range(num_schools)
])
school_effects_hi = np.array([
    np.percentile(school_effects_samples[:, i], 97.5)
    for i in range(num_schools)
])

fig, ax = plt.subplots(nrows=1, ncols=1, sharex=True)
ax.scatter(np.array(range(num_schools)), school_effects_med, color='red', s=60)
ax.scatter(
    np.array(range(num_schools)) + 0.1, treatment_effects, color='blue', s=60)

plt.plot([-0.2, 7.4], [np.mean(avg_effect),
                       np.mean(avg_effect)], 'k', linestyle='--')

ax.errorbar(
    np.array(range(8)),
    school_effects_med,
    yerr=[
        school_effects_med - school_effects_low,
        school_effects_hi - school_effects_med
    ],
    fmt='none')

ax.legend(('avg_effect', 'HMC', 'Observed effect'), fontsize=14)

plt.xlabel('School')
plt.ylabel('Treatment effect')
plt.title('HMC estimated school treatment effects vs. observed data')
fig.set_size_inches(10, 8)
plt.show()

Podemos observar la reducción hacia el grupo avg_effect anterior.

print("Inferred posterior mean: {0:.2f}".format(
    np.mean(school_effects_samples[:,])))
print("Inferred posterior mean se: {0:.2f}".format(
    np.std(school_effects_samples[:,])))

Inferred posterior mean: 6.97
Inferred posterior mean se: 10.41

Crítica

Para obtener la distribución predictiva posterior, es decir, un modelo de nuevos datos $y^*$ a partir de los datos observados $y$:

$$p(y^*|y) \propto \int_\theta p(y^* | \theta)p(\theta |y)d\theta$$

anulamos los valores de las variables aleatorias en el modelo para establecerlos en la media de la distribución posterior y tomamos muestras de ese modelo para generar nuevos datos $y^*$.

sample_shape = [5000]

_, _, _, predictive_treatment_effects = model.sample(
    value=(tf.broadcast_to(np.mean(avg_effect, 0), sample_shape),
           tf.broadcast_to(np.mean(avg_stddev, 0), sample_shape),
           tf.broadcast_to(np.mean(school_effects_standard, 0),
                           sample_shape + [num_schools]),
           None))

fig, axes = plt.subplots(4, 2, sharex=True, sharey=True)
fig.set_size_inches(12, 10)
fig.tight_layout()
for i, ax in enumerate(axes):
  sns.kdeplot(predictive_treatment_effects[:, 2*i].numpy(),
              ax=ax[0], shade=True)
  ax[0].title.set_text(
      "School {} treatment effect posterior predictive".format(2*i))
  sns.kdeplot(predictive_treatment_effects[:, 2*i + 1].numpy(),
              ax=ax[1], shade=True)
  ax[1].title.set_text(
      "School {} treatment effect posterior predictive".format(2*i + 1))
plt.show()

# The mean predicted treatment effects for each of the eight schools.
prediction = np.mean(predictive_treatment_effects, axis=0)

Podemos observar los residuos entre los datos de los efectos del tratamiento y las predicciones de la posterior del modelo. Estos se corresponden con el gráfico anterior que muestra la contracción de los efectos estimados hacia el promedio de la población.

treatment_effects - prediction

array([14.905351 ,  1.2838383, -5.6966295,  0.8327627, -2.3356671,
       -2.0363257,  5.997898 ,  4.3731265], dtype=float32)

Como tenemos una distribución de predicciones para cada escuela, también podemos considerar la distribución de residuos.

residuals = treatment_effects - predictive_treatment_effects

fig, axes = plt.subplots(4, 2, sharex=True, sharey=True)
fig.set_size_inches(12, 10)
fig.tight_layout()
for i, ax in enumerate(axes):
  sns.kdeplot(residuals[:, 2*i].numpy(), ax=ax[0], shade=True)
  ax[0].title.set_text(
      "School {} treatment effect residuals".format(2*i))
  sns.kdeplot(residuals[:, 2*i + 1].numpy(), ax=ax[1], shade=True)
  ax[1].title.set_text(
      "School {} treatment effect residuals".format(2*i + 1))
plt.show()

Agradecimientos

Este tutorial fue escrito originalmente en Edward 1.0 (fuente). Agradecemos a todos los contribuyentes por escribir y revisar esa versión.

Referencias

1. Donald B. Rubin. Estimación en experimentos aleatorios paralelos. Revista de Estadísticas Educativas, 6(4):377-401, 1981.

2. Andrew Gelman, John Carlin, Hal Stern, David Dunson, Aki Vehtari y Donald Rubin. Análisis de datos bayesianos, tercera edición. Chapman y Hall/CRC, 2013.