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.

Oito escolas

Veja em TensorFlow.org

Executar no Google Colab

Ver fonte em GitHub

Baixar notebook

O problema das oito escolas (Rubin 1981) considera a eficácia dos programas de treinamento do equivalente ao vestibular dos Estados Unidos (SAT) em paralelo em oito escolas. É um problema clássico (Bayesian Data Analysis, Stan) que ilustra a utilidade de modelagem hierárquica para compartilhamento de informações entre grupos intercambiáveis.

A implementação abaixo é uma adaptação de um tutorial em Edward 1.0.

Importações

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

Os dados

Da seção 5.5 de Bayesian Data Analysis (Análise bayesiana de dados, Gelman et al. 2013, em tradução livre):

Um estudo foi realizado pelo Educational Testing Service (Serviço de Testes Educacionais) para analisar os efeitos dos programas de treinamento especiais para SAT-V (Scholastic Aptitude Test-Verbal – Teste de Aptidão Escolar - Verbal) em cada uma das oito escolas. A variável resultante em cada estudo foi a nota em uma aplicação especial do SAT-V, um teste de múltipla escolha padronizado aplicado pelo Educational Testing Service e usado para ajudar as faculdades a tomar decisões de admissão. As notas variam de 200 e 800, com média de 500 e desvio padrão de 100. As provas do SAT são resistentes aos efeitos de curto prazo para aumentar o desempenho no teste e foram projetadas para refletir os conhecimentos adquiridos e as habilidades desenvolvidas ao longo de muitos anos de educação. Todavia, cada uma das oito escolas neste estudo consideraram que seu programa de treinamento de curto prazo teve grande sucesso no aumento das notas do SAT. Além disso, não havia motivo anterior para acreditar que qualquer um dos oito programas era mais eficaz que o outro ou que alguns tinham efeito mais similar entre si do que com qualquer outro.

Para cada uma das oito escolas ($J = 8$), temos um efeito de tratamento estimado $y_j$ e um erro padrão da estimativa de efeito $\sigma_j$. Os efeitos de tratamento no estudo foram obtidos por meio de uma regressão linear do grupo de tratamento usando as notas PSAT-M e PSAT-V como variáveis de controle. Como não havia motivo anterior para acreditar que qualquer uma das escolas fosse mais ou menos similar ou que qualquer um dos programas de treinamento fosse mais eficaz, podemos considerar que os efeitos de tratamento são intercambiáveis.

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 os dados, usamos um modelo normal hierárquico, que segue o processo 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*}$$

em que $\mu$ representa o efeito de tratamento médio anterior, e $\tau$ controla quanta variância existe entre as escolas. $y_i$ e $\sigma_i$ são observados. À medida que $\tau \rightarrow \infty$, o modelo se aproxima do modelo sem pooling, ou seja, as estimativas de efeito de tratamento de cada escola podem ser mais independentes. À medida que $\tau \rightarrow 0$, o modelo se aproxima do modelo com pooling completo, ou seja, os efeitos de tratamento de todas as escolas estão mais próximos da média do grupo $\mu$. Para restringir que o desvio padrão seja positivo, obtemos $\tau$ de uma distribuição log-normal (que é equivalente a obter $log(\tau)$ de uma distribuição normal).

Seguindo o artigo Diagnosing Biased Inference with Divergences (Diagnóstico de inferência com bias e divergências), transformamos o modelo acima em um modelo não centralizado 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*}$$

Nós materializamos esse modelo como uma instância de 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))

Inferência bayesiana

Com os dados colhidos, fazemos o Monte Carlo Hamiltoniano (HMC) para calcular a distribuição posterior para os parâmetros do 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 o encolhimento em direção ao efeito médio avg_effect do grupo acima.

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 obter a distribuição preditiva anterior, isto é, um modelo de novos dados $y^*$ usando os dados observados $y$:

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

sobrescrevemos os valores das variáveis aleatórias no modelo, definindo-as como a média da distribuição posterior, e fazemos a amostragem desse modelo para gerar os novos dados $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 conferir os restos entre os dados de efeito de tratamento e as previsões do modelo posterior, que correspondem ao gráfico acima, mostrando o encolhimento dos efeitos estimados rumo à média da população.

treatment_effects - prediction

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

Como temos uma distribuição de previsões para cada escola, também podemos considerar a distribuição dos restos.

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

Agradecimentos

Este tutorial foi escrito originalmente em Edward 1.0 (fonte). Agradecemos a todos os contribuidores por escreverem e revisarem essa versão.

Referências

1. Donald B. Rubin. Estimation in parallel randomized experiments (Estimativa em experimentos randomizados paralelos). Journal of Educational Statistics, 6(4):377-401, 1981.

2. Andrew Gelman, John Carlin, Hal Stern, David Dunson, Aki Vehtari, and Donald Rubin. Bayesian Data Analysis (Análise bayesiana de dados), Terceira edição. Chapman e Hall/CRC, 2013.