Ocho escuelas

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

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

El IMPLEMENTACIÓN a continuación es una adaptación de un Edward 1,0 tutorial .

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

De Bayesian Data Analysis, sección 5.5 (Gelman et al. 2013):

Se realizó un estudio para el Educational Testing Service para analizar los efectos de los programas especiales de entrenamiento para SAT-V (Scholastic Aptitude Test-Verbal) en cada una de las ocho escuelas secundarias. La variable de resultado en cada estudio fue el puntaje en una administración especial del SAT-V, una prueba estandarizada de opción múltiple administrada por el Educational Testing Service y utilizada para ayudar a las universidades a tomar decisiones de admisión; las puntuaciones pueden variar entre 200 y 800, con una media de 500 y una desviación estándar de 100. Los exámenes SAT están diseñados para resistir los esfuerzos a corto plazo dirigidos específicamente a mejorar el rendimiento en la prueba; en cambio, están diseñados para reflejar los conocimientos adquiridos y las habilidades desarrolladas durante 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 para aumentar los puntajes del SAT. Además, no había ninguna razón previa para creer que alguno de los ocho programas fuera más eficaz que cualquier otro o que algunos fueran más similares en efecto entre sí que cualquier otro.

Para cada una de las ocho escuelas (\(J = 8\)), tenemos un efecto del 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 puntuaciones de PSAT-M y PSAT-V como variables de control. Como no había ninguna creencia previa de que cualquiera de las escuelas eran más o menos similar, o que cualquiera de los programas de entrenamiento sería más eficaz, podemos considerar los efectos del tratamiento como intercambiable .

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

png

Modelo

Para capturar los datos, utilizamos 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 medio del tratamiento y antes \(\tau\) controla la cantidad de variación que hay entre las escuelas. El \(y_i\) y \(\sigma_i\) se observan. Como \(\tau \rightarrow \infty\), el modelo se acerca al modelo sin poner en común, es decir, cada una de las estimaciones de los efectos del tratamiento de la escuela se les permite ser más independiente. Como \(\tau \rightarrow 0\), el modelo se acerca al modelo completo puesta en común, es decir, todos los efectos del tratamiento de la escuela están más cerca de la media del grupo \(\mu\). Para restringir la desviación estándar a ser positivo, dibujamos \(\tau\) de una distribución logarítmica normal (que es equivalente al dibujo \(log(\tau)\) de una distribución normal).

Después de Diagnóstico de Biased inferencia con divergencias , transformamos el modelo anterior en un modelo equivalente no-centrado:

\[ \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*} \]

Nos Reify este modelo como un JointDistributionSequential ejemplo:

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

Dados los datos, realizamos el Hamiltonian 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()

png

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

png

Podemos observar la contracción hacia el grupo avg_effect anteriormente.

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 posterior predictiva, es decir, un modelo de datos nuevos \(y^*\) dado los datos observados \(y\):

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

que anulan los valores de las variables aleatorias en el modelo para ponerlos a la media de la distribución posterior, y la muestra 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()

png

# 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 del modelo posterior. 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)

Debido a que 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()

png

Agradecimientos

Este tutorial fue escrito originalmente en Edward 1,0 ( fuente ). Agradecemos a todos los que contribuyeron a 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.