Inferensi Variasi pada Model Grafis Probabilistik dengan Distribusi Gabungan

Lihat di TensorFlow.org Jalankan di Google Colab Lihat sumber di GitHub Unduh buku catatan

Inferensi Variasi (VI) menampilkan perkiraan inferensi Bayesian sebagai masalah optimasi dan mencari distribusi posterior 'pengganti' yang meminimalkan divergensi KL dengan posterior sebenarnya. VI berbasis gradien seringkali lebih cepat daripada metode MCMC, menyusun secara alami dengan optimalisasi parameter model, dan memberikan batas bawah pada bukti model yang dapat digunakan secara langsung untuk perbandingan model, diagnosis konvergensi, dan inferensi yang dapat disusun.

TensorFlow Probability menawarkan alat untuk VI yang cepat, fleksibel, dan skalabel yang cocok secara alami dengan tumpukan TFP. Alat-alat ini memungkinkan pembangunan posterior pengganti dengan struktur kovarians yang disebabkan oleh transformasi linier atau aliran normalisasi.

VI dapat digunakan untuk memperkirakan Bayesian interval kredibel untuk parameter dari model regresi untuk memperkirakan efek dari berbagai pengobatan atau fitur yang diamati pada hasil yang menarik. Interval kredibel mengikat nilai parameter yang tidak teramati dengan probabilitas tertentu, menurut distribusi posterior dari parameter yang dikondisikan pada data yang diamati dan diberikan asumsi pada distribusi parameter sebelumnya.

Dalam CoLab ini, kami menunjukkan cara menggunakan VI untuk mendapatkan interval kredibel untuk parameter dari Bayesian linier model regresi untuk tingkat radon diukur dalam rumah (menggunakan Gelman et al (2007) Radon dataset. ; Lihat contoh-contoh serupa di Stan). Kami menunjukkan bagaimana TFP JointDistribution s menggabungkan dengan bijectors untuk membangun dan cocok dua jenis posteriors pengganti ekspresif:

  • distribusi Normal standar yang ditransformasikan oleh matriks blok. Matriks dapat mencerminkan kemandirian di antara beberapa komponen posterior dan ketergantungan antara lain, merelaksasi asumsi bidang rata-rata atau posterior kovarians penuh.
  • lebih kompleks, lebih tinggi-kapasitas terbalik aliran autoregressive .

Posterior pengganti dilatih dan dibandingkan dengan hasil dari baseline posterior pengganti lapangan rata-rata, serta sampel ground-truth dari Hamiltonian Monte Carlo.

Ikhtisar Inferensi Variasi Bayesian

Misalkan kita memiliki proses generatif berikut, di mana \(\theta\) mewakili parameter acak, \(\omega\) mewakili parameter deterministik, dan \(x_i\) adalah fitur dan \(y_i\) adalah nilai-nilai target untuk \(i=1,\ldots,n\) diamati titik data: \ begin {menyelaraskan } &\theta \sim r(\Theta) && \text{(Sebelumnya)}\ &\text{untuk } i = 1 \ldots n: \nonumber \ &\quad y_i \sim p(Y_i|x_i, \theta , \ omega) && \ text {(Kemungkinan)} \ end {menyelaraskan}

VI kemudian dicirikan oleh: $\newcommand{\E}{\operatorname{\mathbb{E} } } \newcommand{\K}{\operatorname{\mathbb{K} } } \newcommand{\defeq}{\overset {\tiny\text{def} }{=} } \DeclareMathOperator*{\argmin}{arg\,min}$

\ begin {menyelaraskan} - \ log p ({y_i} _i ^ n | {x_i} _i ^ n, \ omega) & \ defeq - \ log \ int \ textrm {d} \ theta \, r (\ theta) \ prod_i^np(y_i|x_i,\theta, \omega) && \text{(Integral yang sangat sulit)} \ &= -\log \int \textrm{d}\theta\, q(\theta) \frac{1 }{q(\theta)} r(\theta) \prod_i^np(y_i|x_i,\theta, \omega) && \text{(Kalikan dengan 1)}\ &\le - \int \textrm{d} \ theta \, q (\ theta) \ log \ frac {r (\ theta) \ prod_i ^ np (y_i | x i, \ theta, \ omega)} {q (\ theta)} && \ text {(ketimpangan Jensen )} \ & \ defeq \ E {q (\ Theta)} [- \ log p (y_i | x_i, \ Theta, \ omega)] + \ K [q (\ Theta), r (\ Theta)] \ & \ defeq \text{expected negative log likelihood"} + \ text {kl regularizer"} \ end {menyelaraskan}

(Secara teknis kami menganggap \(q\) adalah benar-benar terus menerus terhadap \(r\). Lihat juga, ketimpangan Jensen .)

Karena terikat berlaku untuk semua q, itu jelas paling ketat untuk:

\[q^*,w^* = \argmin_{q \in \mathcal{Q},\omega\in\mathbb{R}^d} \left\{ \sum_i^n\E_{q(\Theta)}\left[ -\log p(y_i|x_i,\Theta, \omega) \right] + \K[q(\Theta), r(\Theta)] \right\}\]

Mengenai terminologi, kami menyebutnya

  • \(q^*\) yang "posterior pengganti," dan,
  • \(\mathcal{Q}\) "keluarga pengganti."

\(\omega^*\) mewakili nilai-nilai maksimum-kemungkinan parameter deterministik pada VI kerugian. Lihat survei ini untuk informasi lebih lanjut tentang inferensi variational.

Contoh: Regresi linier hierarkis Bayesian pada pengukuran Radon

Radon adalah gas radioaktif yang memasuki rumah melalui titik kontak dengan tanah. Ini adalah karsinogen yang merupakan penyebab utama kanker paru-paru pada non-perokok. Tingkat radon sangat bervariasi dari rumah tangga ke rumah tangga.

EPA melakukan studi tingkat radon di 80.000 rumah. Dua prediktor penting adalah:

  • Lantai tempat pengukuran dilakukan (radon lebih tinggi di ruang bawah tanah)
  • Tingkat uranium kabupaten (korelasi positif dengan tingkat radon)

Memprediksi tingkat radon di rumah-rumah dikelompokkan berdasarkan county adalah masalah klasik dalam Bayesian pemodelan hirarki, diperkenalkan oleh Gelman dan Hill (2006) . Kami akan membangun model linier hierarkis untuk memprediksi pengukuran radon di rumah, di mana hierarki adalah pengelompokan rumah berdasarkan wilayah. Kami tertarik pada interval yang kredibel untuk pengaruh lokasi (county) pada tingkat radon rumah di Minnesota. Untuk mengisolasi efek ini, efek lantai dan tingkat uranium juga disertakan dalam model. Selain itu, kami akan memasukkan efek kontekstual yang sesuai dengan rata-rata lantai tempat pengukuran dilakukan, menurut wilayah, sehingga jika ada variasi di antara wilayah lantai tempat pengukuran dilakukan, hal ini tidak dikaitkan dengan efek wilayah.

pip3 install -q tf-nightly tfp-nightly

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp
import warnings

tfd = tfp.distributions
tfb = tfp.bijectors

plt.rcParams['figure.facecolor'] = '1.'

# Load the Radon dataset from `tensorflow_datasets` and filter to data from
# Minnesota.
dataset = tfds.as_numpy(
    tfds.load('radon', split='train').filter(
        lambda x: x['features']['state'] == 'MN').batch(10**9))

# Dependent variable: Radon measurements by house.
dataset = next(iter(dataset))
radon_measurement = dataset['activity'].astype(np.float32)
radon_measurement[radon_measurement <= 0.] = 0.1
log_radon = np.log(radon_measurement)

# Measured uranium concentrations in surrounding soil.
uranium_measurement = dataset['features']['Uppm'].astype(np.float32)
log_uranium = np.log(uranium_measurement)

# County indicator.
county_strings = dataset['features']['county'].astype('U13')
unique_counties, county = np.unique(county_strings, return_inverse=True)
county = county.astype(np.int32)
num_counties = unique_counties.size

# Floor on which the measurement was taken.
floor_of_house = dataset['features']['floor'].astype(np.int32)

# Average floor by county (contextual effect).
county_mean_floor = []
for i in range(num_counties):
  county_mean_floor.append(floor_of_house[county == i].mean())
county_mean_floor = np.array(county_mean_floor, dtype=log_radon.dtype)
floor_by_county = county_mean_floor[county]

Model regresi ditentukan sebagai berikut:

\(\newcommand{\Normal}{\operatorname{\sf Normal} }\)\ begin {menyelaraskan} & \ text {uranium_weight} \ sim \ Normal (0, 1) \ & \ text {county_floor_weight} \ sim \ Normal (0, 1) \ & \ text {untuk} j = 1 \ ldots \text{num_counties}:\ &\quad \text{county_effect}_j \sim \Normal (0, \sigma_c)\ &\text{untuk } i = 1\ldots n:\ &\quad \mu_i = ( \ & \ quad \ quad \ text {Bias} \ & \ quad \ quad + \ text {efek county} {\ text {county} _i} \ & \ quad \ quad + \ text {log_uranium} _i \ kali \ text {uranium_weight } \ & \ quad \ quad + \ text {floor_of_house} _i \ kali \ text {floor_weight} \ & \ quad \ quad + \ text {floor_by county} {\ text {county} _i} \ kali \ text {county_floor_weight}) \ & \ quad \ text {log_radon} _i \ sim \ normal (\ mu_i, \ sigma_y) \ end {menyelaraskan} dimana \(i\) indeks pengamatan dan \(\text{county}_i\) adalah county di mana \(i\)pengamatan th adalah diambil.

Kami menggunakan efek acak tingkat kabupaten untuk menangkap variasi geografis. Parameter uranium_weight dan county_floor_weight dimodelkan secara probabilistik, dan floor_weight dan konstanta bias deterministik. Pilihan pemodelan ini sebagian besar bersifat arbitrer, dan dibuat untuk tujuan mendemonstrasikan VI pada model probabilistik dengan kompleksitas yang masuk akal. Untuk diskusi yang lebih menyeluruh dari pemodelan multilevel dengan efek tetap dan acak dalam TFP, menggunakan dataset radon, lihat Multilevel Modeling Primer dan Fitting Generalized Linear Mixed-efek Model Menggunakan Variasional Inference .

# Create variables for fixed effects.
floor_weight = tf.Variable(0.)
bias = tf.Variable(0.)

# Variables for scale parameters.
log_radon_scale = tfp.util.TransformedVariable(1., tfb.Exp())
county_effect_scale = tfp.util.TransformedVariable(1., tfb.Exp())

# Define the probabilistic graphical model as a JointDistribution.
@tfd.JointDistributionCoroutineAutoBatched
def model():
  uranium_weight = yield tfd.Normal(0., scale=1., name='uranium_weight')
  county_floor_weight = yield tfd.Normal(
      0., scale=1., name='county_floor_weight')
  county_effect = yield tfd.Sample(
      tfd.Normal(0., scale=county_effect_scale),
      sample_shape=[num_counties], name='county_effect')
  yield tfd.Normal(
      loc=(log_uranium * uranium_weight + floor_of_house* floor_weight
           + floor_by_county * county_floor_weight
           + tf.gather(county_effect, county, axis=-1)
           + bias),
      scale=log_radon_scale[..., tf.newaxis],
      name='log_radon') 

# Pin the observed `log_radon` values to model the un-normalized posterior.
target_model = model.experimental_pin(log_radon=log_radon)

Posterior pengganti ekspresif

Selanjutnya kami memperkirakan distribusi posterior dari efek acak menggunakan VI dengan dua jenis posterior pengganti yang berbeda:

  • Distribusi normal multivariat terbatas, dengan struktur kovarians yang diinduksi oleh transformasi matriks blok.
  • Sebuah distribusi Normal Standar multivariat diubah oleh Inverse Autoregressive Arus , yang kemudian dibagi dan direstrukturisasi untuk mencocokkan dukungan posterior.

Multivariat Pengganti normal posterior

Untuk membangun posterior pengganti ini, operator linier yang dapat dilatih digunakan untuk menginduksi korelasi di antara komponen-komponen posterior.

# Determine the `event_shape` of the posterior, and calculate the size of each
# `event_shape` component. These determine the sizes of the components of the
# underlying standard Normal distribution, and the dimensions of the blocks in
# the blockwise matrix transformation.
event_shape = target_model.event_shape_tensor()
flat_event_shape = tf.nest.flatten(event_shape)
flat_event_size = tf.nest.map_structure(tf.reduce_prod, flat_event_shape)

# The `event_space_bijector` maps unconstrained values (in R^n) to the support
# of the prior -- we'll need this at the end to constrain Multivariate Normal
# samples to the prior's support.
event_space_bijector = target_model.experimental_default_event_space_bijector()

Buatlah sebuah JointDistribution dengan vektor-dihargai komponen normal standar, dengan ukuran ditentukan oleh komponen sebelumnya yang sesuai. Komponen harus bernilai vektor sehingga dapat ditransformasikan oleh operator linier.

base_standard_dist = tfd.JointDistributionSequential(
      [tfd.Sample(tfd.Normal(0., 1.), s) for s in flat_event_size])

Bangun operator linier segitiga bawah blok yang dapat dilatih. Kami akan menerapkannya pada distribusi Normal standar untuk menerapkan transformasi matriks blok (yang dapat dilatih) dan menginduksi struktur korelasi posterior.

Dalam operator blockwise linear, sebuah dilatih blok penuh-matriks merupakan kovarians penuh antara dua komponen dari posterior, sementara blok nol (atau None ) mengungkapkan kemerdekaan. Balok pada diagonal adalah matriks segitiga bawah atau matriks diagonal, sehingga seluruh struktur blok mewakili matriks segitiga bawah.

Menerapkan bijektor ini ke distribusi dasar menghasilkan distribusi Normal multivariat dengan rata-rata 0 dan kovarians (difaktorkan Cholesky) sama dengan matriks blok segitiga bawah.

operators = (
    (tf.linalg.LinearOperatorDiag,),  # Variance of uranium weight (scalar).
    (tf.linalg.LinearOperatorFullMatrix,  # Covariance between uranium and floor-by-county weights.
     tf.linalg.LinearOperatorDiag),  # Variance of floor-by-county weight (scalar).
    (None,  # Independence between uranium weight and county effects.
     None,  #  Independence between floor-by-county and county effects.
     tf.linalg.LinearOperatorDiag)  # Independence among the 85 county effects.
    )

block_tril_linop = (
    tfp.experimental.vi.util.build_trainable_linear_operator_block(
        operators, flat_event_size))
scale_bijector = tfb.ScaleMatvecLinearOperatorBlock(block_tril_linop)

Setelah menerapkan operator linear dengan distribusi normal standar, menerapkan multi Shift bijector untuk memungkinkan rata-rata untuk mengambil nilai-nilai nol.

loc_bijector = tfb.JointMap(
    tf.nest.map_structure(
        lambda s: tfb.Shift(
            tf.Variable(tf.random.uniform(
                (s,), minval=-2., maxval=2., dtype=tf.float32))),
        flat_event_size))

Distribusi Normal multivariat yang dihasilkan, diperoleh dengan mentransformasikan distribusi Normal standar dengan bijektor skala dan lokasi, harus dibentuk ulang dan direstrukturisasi agar sesuai dengan prior, dan akhirnya dibatasi pada dukungan prior.

# Reshape each component to match the prior, using a nested structure of
# `Reshape` bijectors wrapped in `JointMap` to form a multipart bijector.
reshape_bijector = tfb.JointMap(
    tf.nest.map_structure(tfb.Reshape, flat_event_shape))

# Restructure the flat list of components to match the prior's structure
unflatten_bijector = tfb.Restructure(
        tf.nest.pack_sequence_as(
            event_shape, range(len(flat_event_shape))))

Sekarang, gabungkan semuanya -- ikat bijektor yang dapat dilatih bersama-sama dan terapkan ke distribusi Normal standar dasar untuk membangun posterior pengganti.

surrogate_posterior = tfd.TransformedDistribution(
    base_standard_dist,
    bijector = tfb.Chain(  # Note that the chained bijectors are applied in reverse order
        [
         event_space_bijector,  # constrain the surrogate to the support of the prior
         unflatten_bijector,  # pack the reshaped components into the `event_shape` structure of the posterior
         reshape_bijector,  # reshape the vector-valued components to match the shapes of the posterior components
         loc_bijector,  # allow for nonzero mean
         scale_bijector  # apply the block matrix transformation to the standard Normal distribution
         ]))

Melatih multivariat Normal pengganti posterior.

optimizer = tf.optimizers.Adam(learning_rate=1e-2)
mvn_loss = tfp.vi.fit_surrogate_posterior(
    target_model.unnormalized_log_prob,
    surrogate_posterior,
    optimizer=optimizer,
    num_steps=10**4,
    sample_size=16,
    jit_compile=True)

mvn_samples = surrogate_posterior.sample(1000)
mvn_final_elbo = tf.reduce_mean(
    target_model.unnormalized_log_prob(*mvn_samples)
    - surrogate_posterior.log_prob(mvn_samples))

print('Multivariate Normal surrogate posterior ELBO: {}'.format(mvn_final_elbo))

plt.plot(mvn_loss)
plt.xlabel('Training step')
_ = plt.ylabel('Loss value')

Multivariate Normal surrogate posterior ELBO: -1065.705322265625

png

Karena posterior pengganti terlatih adalah distribusi TFP, kita dapat mengambil sampel darinya dan memprosesnya untuk menghasilkan interval kredibel posterior untuk parameter.

Kotak-dan-kumis plot di bawah ini menunjukkan 50% dan 95% interval kredibel untuk efek county dari dua kabupaten terbesar dan bobot regresi pada pengukuran uranium tanah dan rata lantai oleh county. Interval kredibel posterior untuk efek daerah menunjukkan bahwa lokasi di daerah St. Louis dikaitkan dengan tingkat radon yang lebih rendah, setelah memperhitungkan variabel lain, dan bahwa pengaruh lokasi di daerah Hennepin mendekati netral.

Interval kredibel posterior pada bobot regresi menunjukkan bahwa tingkat uranium tanah yang lebih tinggi dikaitkan dengan tingkat radon yang lebih tinggi, dan wilayah di mana pengukuran dilakukan di lantai yang lebih tinggi (kemungkinan karena rumah tersebut tidak memiliki ruang bawah tanah) cenderung memiliki tingkat radon yang lebih tinggi, yang dapat berhubungan dengan sifat-sifat tanah dan pengaruhnya terhadap jenis struktur yang dibangun.

Koefisien (deterministik) lantai negatif, menunjukkan bahwa lantai bawah memiliki tingkat radon yang lebih tinggi, seperti yang diharapkan.

st_louis_co = 69  # Index of St. Louis, the county with the most observations.
hennepin_co = 25  # Index of Hennepin, with the second-most observations.

def pack_samples(samples):
  return {'County effect (St. Louis)': samples.county_effect[..., st_louis_co],
          'County effect (Hennepin)': samples.county_effect[..., hennepin_co],
          'Uranium weight': samples.uranium_weight,
          'Floor-by-county weight': samples.county_floor_weight}

def plot_boxplot(posterior_samples):
  fig, axes = plt.subplots(1, 4, figsize=(16, 4))

  # Invert the results dict for easier plotting.
  k = list(posterior_samples.values())[0].keys()
  plot_results = {
      v: {p: posterior_samples[p][v] for p in posterior_samples} for v in k}
  for i, (var, var_results) in enumerate(plot_results.items()):
    sns.boxplot(data=list(var_results.values()), ax=axes[i],
                width=0.18*len(var_results), whis=(2.5, 97.5))
    # axes[i].boxplot(list(var_results.values()), whis=(2.5, 97.5))
    axes[i].title.set_text(var)
    fs = 10 if len(var_results) < 4 else 8
    axes[i].set_xticklabels(list(var_results.keys()), fontsize=fs)

results = {'Multivariate Normal': pack_samples(mvn_samples)}

print('Bias is: {:.2f}'.format(bias.numpy()))
print('Floor fixed effect is: {:.2f}'.format(floor_weight.numpy()))
plot_boxplot(results)

Bias is: 1.40
Floor fixed effect is: -0.72

png

Inverse Autoregressive Flow pengganti posterior

Inverse Autoregressive Flows (IAFs) adalah aliran normalisasi yang menggunakan jaringan saraf untuk menangkap dependensi nonlinier yang kompleks di antara komponen distribusi. Selanjutnya kita membangun posterior pengganti IAF untuk melihat apakah model berkapasitas lebih tinggi dan lebih fleksibel ini mengungguli Multivariat yang dibatasi Normal.

# Build a standard Normal with a vector `event_shape`, with length equal to the
# total number of degrees of freedom in the posterior.
base_distribution = tfd.Sample(
    tfd.Normal(0., 1.), sample_shape=[tf.reduce_sum(flat_event_size)])

# Apply an IAF to the base distribution.
num_iafs = 2
iaf_bijectors = [
    tfb.Invert(tfb.MaskedAutoregressiveFlow(
        shift_and_log_scale_fn=tfb.AutoregressiveNetwork(
            params=2, hidden_units=[256, 256], activation='relu')))
    for _ in range(num_iafs)
]

# Split the base distribution's `event_shape` into components that are equal
# in size to the prior's components.
split = tfb.Split(flat_event_size)

# Chain these bijectors and apply them to the standard Normal base distribution
# to build the surrogate posterior. `event_space_bijector`,
# `unflatten_bijector`, and `reshape_bijector` are the same as in the
# multivariate Normal surrogate posterior.
iaf_surrogate_posterior = tfd.TransformedDistribution(
    base_distribution,
    bijector=tfb.Chain([
         event_space_bijector,  # constrain the surrogate to the support of the prior
         unflatten_bijector,  # pack the reshaped components into the `event_shape` structure of the prior
         reshape_bijector,  # reshape the vector-valued components to match the shapes of the prior components
         split] +  # Split the samples into components of the same size as the prior components
         iaf_bijectors  # Apply a flow model to the Tensor-valued standard Normal distribution
         ))

Latih posterior pengganti IAF.

optimizer=tf.optimizers.Adam(learning_rate=1e-2)
iaf_loss = tfp.vi.fit_surrogate_posterior(
  target_model.unnormalized_log_prob,
  iaf_surrogate_posterior,
  optimizer=optimizer,
  num_steps=10**4,
  sample_size=4,
  jit_compile=True)

iaf_samples = iaf_surrogate_posterior.sample(1000)
iaf_final_elbo = tf.reduce_mean(
    target_model.unnormalized_log_prob(*iaf_samples)
    - iaf_surrogate_posterior.log_prob(iaf_samples))
print('IAF surrogate posterior ELBO: {}'.format(iaf_final_elbo))

plt.plot(iaf_loss)
plt.xlabel('Training step')
_ = plt.ylabel('Loss value')

IAF surrogate posterior ELBO: -1065.3663330078125

png

Interval kredibel untuk posterior pengganti IAF tampak serupa dengan Normal multivariat terbatas.

results['IAF'] = pack_samples(iaf_samples)
plot_boxplot(results)

png

Baseline: Posterior pengganti bidang rata-rata

Posterior pengganti VI sering diasumsikan sebagai bidang rata-rata (independen) Distribusi normal, dengan rata-rata dan varians yang dapat dilatih, yang dibatasi untuk mendukung prior dengan transformasi bijektif. Kami mendefinisikan posterior pengganti medan rata-rata di samping dua posterior pengganti yang lebih ekspresif, menggunakan rumus umum yang sama dengan posterior pengganti Normal multivariat.

# A block-diagonal linear operator, in which each block is a diagonal operator,
# transforms the standard Normal base distribution to produce a mean-field
# surrogate posterior.
operators = (tf.linalg.LinearOperatorDiag,
             tf.linalg.LinearOperatorDiag,
             tf.linalg.LinearOperatorDiag)
block_diag_linop = (
    tfp.experimental.vi.util.build_trainable_linear_operator_block(
        operators, flat_event_size))
mean_field_scale = tfb.ScaleMatvecLinearOperatorBlock(block_diag_linop)

mean_field_loc = tfb.JointMap(
    tf.nest.map_structure(
        lambda s: tfb.Shift(
            tf.Variable(tf.random.uniform(
                (s,), minval=-2., maxval=2., dtype=tf.float32))),
        flat_event_size))

mean_field_surrogate_posterior = tfd.TransformedDistribution(
    base_standard_dist,
    bijector = tfb.Chain(  # Note that the chained bijectors are applied in reverse order
        [
         event_space_bijector,  # constrain the surrogate to the support of the prior
         unflatten_bijector,  # pack the reshaped components into the `event_shape` structure of the posterior
         reshape_bijector, # reshape the vector-valued components to match the shapes of the posterior components
         mean_field_loc,   # allow for nonzero mean
         mean_field_scale  # apply the block matrix transformation to the standard Normal distribution
         ]))

optimizer=tf.optimizers.Adam(learning_rate=1e-2)
mean_field_loss = tfp.vi.fit_surrogate_posterior(
    target_model.unnormalized_log_prob,
    mean_field_surrogate_posterior,
    optimizer=optimizer,
    num_steps=10**4,
    sample_size=16,
    jit_compile=True)

mean_field_samples = mean_field_surrogate_posterior.sample(1000)
mean_field_final_elbo = tf.reduce_mean(
    target_model.unnormalized_log_prob(*mean_field_samples)
    - mean_field_surrogate_posterior.log_prob(mean_field_samples))
print('Mean-field surrogate posterior ELBO: {}'.format(mean_field_final_elbo))

plt.plot(mean_field_loss)
plt.xlabel('Training step')
_ = plt.ylabel('Loss value')

Mean-field surrogate posterior ELBO: -1065.7652587890625

png

Dalam hal ini, rata-rata bidang pengganti posterior memberikan hasil yang serupa dengan posterior pengganti yang lebih ekspresif, menunjukkan bahwa model yang lebih sederhana ini mungkin memadai untuk tugas inferensi.

results['Mean Field'] = pack_samples(mean_field_samples)
plot_boxplot(results)

png

Kebenaran dasar: Hamiltonian Monte Carlo (HMC)

Kami menggunakan HMC untuk menghasilkan sampel "kebenaran dasar" dari posterior yang sebenarnya, untuk dibandingkan dengan hasil dari posterior pengganti.

num_chains = 8
num_leapfrog_steps = 3
step_size = 0.4
num_steps=20000

flat_event_shape = tf.nest.flatten(target_model.event_shape)
enum_components = list(range(len(flat_event_shape)))
bijector = tfb.Restructure(
    enum_components,
    tf.nest.pack_sequence_as(target_model.event_shape, enum_components))(
        target_model.experimental_default_event_space_bijector())

current_state = bijector(
    tf.nest.map_structure(
        lambda e: tf.zeros([num_chains] + list(e), dtype=tf.float32),
    target_model.event_shape))

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=target_model.unnormalized_log_prob,
    num_leapfrog_steps=num_leapfrog_steps,
    step_size=[tf.fill(s.shape, step_size) for s in current_state])

hmc = tfp.mcmc.TransformedTransitionKernel(
    hmc, bijector)
hmc = tfp.mcmc.DualAveragingStepSizeAdaptation(
    hmc,
    num_adaptation_steps=int(num_steps // 2 * 0.8),
    target_accept_prob=0.9)

chain, is_accepted = tf.function(
    lambda current_state: tfp.mcmc.sample_chain(
        current_state=current_state,
        kernel=hmc,
        num_results=num_steps // 2,
        num_burnin_steps=num_steps // 2,
        trace_fn=lambda _, pkr:
        (pkr.inner_results.inner_results.is_accepted),
        ),
    autograph=False,
    jit_compile=True)(current_state)

accept_rate = tf.reduce_mean(tf.cast(is_accepted, tf.float32))
ess = tf.nest.map_structure(
    lambda c: tfp.mcmc.effective_sample_size(
        c,
        cross_chain_dims=1,
        filter_beyond_positive_pairs=True),
    chain)

r_hat = tf.nest.map_structure(tfp.mcmc.potential_scale_reduction, chain)
hmc_samples = pack_samples(
    tf.nest.pack_sequence_as(target_model.event_shape, chain))
print('Acceptance rate is {}'.format(accept_rate))

Acceptance rate is 0.9008625149726868

Plot jejak sampel untuk memeriksa kewarasan hasil HMC.

def plot_traces(var_name, samples):
  fig, axes = plt.subplots(1, 2, figsize=(14, 1.5), sharex='col', sharey='col')
  for chain in range(num_chains):
    s = samples.numpy()[:, chain]
    axes[0].plot(s, alpha=0.7)
    sns.kdeplot(s, ax=axes[1], shade=False)
    axes[0].title.set_text("'{}' trace".format(var_name))
    axes[1].title.set_text("'{}' distribution".format(var_name))
    axes[0].set_xlabel('Iteration')

warnings.filterwarnings('ignore')
for var, var_samples in hmc_samples.items():
  plot_traces(var, var_samples)

png

png

png

png

Ketiga posterior pengganti menghasilkan interval yang kredibel yang secara visual mirip dengan sampel HMC, meskipun kadang-kadang kurang tersebar karena efek hilangnya ELBO, seperti yang biasa terjadi pada VI.

results['HMC'] = hmc_samples
plot_boxplot(results)

png

Hasil tambahan

Fungsi plot

Bukti Batas Bawah (ELBO)

IAF, sejauh ini merupakan posterior pengganti terbesar dan paling fleksibel, menyatu dengan Bukti Batas Bawah (ELBO) tertinggi.

plot_loss_and_elbo()

png

Sampel posterior

Sampel dari masing-masing posterior pengganti, dibandingkan dengan sampel kebenaran dasar HMC (visualisasi sampel yang berbeda ditunjukkan dalam plot kotak).

plot_kdes()

png

Kesimpulan

Dalam Colab ini, kami membangun posterior pengganti VI menggunakan distribusi gabungan dan bijektor multi-bagian, dan menyesuaikannya untuk memperkirakan interval yang kredibel untuk bobot dalam model regresi pada dataset radon. Untuk model sederhana ini, posterior pengganti yang lebih ekspresif dilakukan serupa dengan posterior pengganti medan rata-rata. Alat yang kami tunjukkan, bagaimanapun, dapat digunakan untuk membangun berbagai posterior pengganti fleksibel yang cocok untuk model yang lebih kompleks.