Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

Studi Kasus Probabilitas TensorFlow: Estimasi Kovarian

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

Saya menulis notebook ini sebagai studi kasus untuk mempelajari Probabilitas TensorFlow. Masalah yang saya pilih untuk dipecahkan adalah mengestimasi matriks kovarians untuk sampel variabel acak 2-D mean 0 Gaussian. Masalahnya memiliki beberapa fitur bagus:

  • Jika kita menggunakan Wishart terbalik sebelum kovarian (pendekatan umum), masalahnya memiliki solusi analitik, sehingga kita dapat memeriksa hasil kita.
  • Masalahnya melibatkan pengambilan sampel parameter terbatas, yang menambahkan beberapa kompleksitas yang menarik.
  • Solusi yang paling mudah bukanlah yang tercepat, jadi ada beberapa pekerjaan pengoptimalan yang harus dilakukan.

Saya memutuskan untuk menulis pengalaman saya saat saya pergi. Saya butuh beberapa saat untuk memahami poin-poin penting TFP, jadi notebook ini memulai dengan cukup sederhana dan kemudian secara bertahap bekerja hingga fitur TFP yang lebih rumit. Saya mengalami banyak masalah di sepanjang jalan, dan saya telah mencoba menangkap kedua proses yang membantu saya mengidentifikasi mereka dan solusi yang akhirnya saya temukan. Saya telah mencoba memasukkan banyak detail (termasuk banyak tes untuk memastikan langkah individu sudah benar).

Mengapa mempelajari Probabilitas TensorFlow?

Saya menemukan Probabilitas TensorFlow menarik untuk proyek saya karena beberapa alasan:

  • Probabilitas TensorFlow memungkinkan Anda membuat prototipe mengembangkan model kompleks secara interaktif di notebook. Anda dapat memecah kode Anda menjadi potongan-potongan kecil yang dapat Anda uji secara interaktif dan dengan pengujian unit.
  • Setelah Anda siap untuk meningkatkan, Anda dapat memanfaatkan semua infrastruktur yang kami miliki untuk membuat TensorFlow berjalan pada beberapa prosesor yang dioptimalkan pada banyak mesin.
  • Akhirnya, meskipun saya sangat menyukai Stan, saya merasa cukup sulit untuk melakukan debug. Anda harus menulis semua kode pemodelan Anda dalam bahasa mandiri yang memiliki sangat sedikit alat untuk memungkinkan Anda menyodok kode Anda, memeriksa status perantara, dan seterusnya.

Sisi negatifnya adalah Probabilitas TensorFlow jauh lebih baru daripada Stan dan PyMC3, jadi dokumentasinya sedang dalam proses, dan ada banyak fungsi yang belum dibuat. Untungnya, saya menemukan fondasi TFP kokoh, dan dirancang dengan cara modular yang memungkinkan seseorang untuk memperluas fungsinya dengan cukup mudah. Dalam buku catatan ini, selain menyelesaikan studi kasus, saya akan menunjukkan beberapa cara untuk memperluas TFP.

Untuk siapa ini

Saya berasumsi bahwa pembaca datang ke buku catatan ini dengan beberapa prasyarat penting. Anda harus:

Percobaan pertama

Inilah upaya pertama saya untuk mengatasi masalah tersebut. Spoiler: solusi saya tidak berfungsi, dan itu akan membutuhkan beberapa upaya untuk memperbaikinya! Meskipun prosesnya memakan waktu cukup lama, setiap upaya di bawah ini berguna untuk mempelajari bagian baru TFP.

Satu catatan: TFP saat ini tidak menerapkan distribusi Wishart terbalik (kita akan melihat di bagian akhir cara menggulung Wishart terbalik kita sendiri), jadi sebagai gantinya saya akan mengubah masalah menjadi masalah memperkirakan matriks presisi menggunakan Wishart sebelumnya.

import collections
import math
import os
import time

import numpy as np
import pandas as pd
import scipy
import scipy.stats
import matplotlib.pyplot as plt

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

Langkah 1: kumpulkan observasi

Data saya di sini semuanya sintetis, jadi ini akan tampak sedikit lebih rapi daripada contoh dunia nyata. Namun, tidak ada alasan Anda tidak dapat membuat beberapa data sintetis Anda sendiri.

Kiat : Setelah Anda memutuskan bentuk model, Anda dapat memilih beberapa nilai parameter dan menggunakan model yang Anda pilih untuk menghasilkan beberapa data sintetis. Sebagai pemeriksaan kesehatan penerapan, Anda kemudian dapat memverifikasi bahwa taksiran Anda menyertakan nilai sebenarnya dari parameter yang Anda pilih. Untuk membuat siklus debugging / pengujian Anda lebih cepat, Anda dapat mempertimbangkan versi model yang disederhanakan (misalnya, gunakan lebih sedikit dimensi atau lebih sedikit sampel).

Tip: Paling mudah bekerja dengan pengamatan Anda sebagai array NumPy. Satu hal penting yang perlu diperhatikan adalah NumPy secara default menggunakan float64, sedangkan TensorFlow secara default menggunakan float32.

Secara umum, operasi TensorFlow ingin semua argumen memiliki jenis yang sama, dan Anda harus melakukan transmisi data eksplisit untuk mengubah jenis. Jika Anda menggunakan observasi float64, Anda harus menambahkan banyak operasi cast. Sebaliknya, NumPy akan menangani transmisi secara otomatis. Oleh karena itu, jauh lebih mudah untuk mengonversi data Numpy Anda menjadi float32 daripada memaksa TensorFlow menggunakan float64.

Pilih beberapa nilai parameter

# We're assuming 2-D data with a known true mean of (0, 0)
true_mean = np.zeros([2], dtype=np.float32)
# We'll make the 2 coordinates correlated
true_cor = np.array([[1.0, 0.9], [0.9, 1.0]], dtype=np.float32)
# And we'll give the 2 coordinates different variances
true_var = np.array([4.0, 1.0], dtype=np.float32)
# Combine the variances and correlations into a covariance matrix
true_cov = np.expand_dims(np.sqrt(true_var), axis=1).dot(
    np.expand_dims(np.sqrt(true_var), axis=1).T) * true_cor
# We'll be working with precision matrices, so we'll go ahead and compute the
# true precision matrix here
true_precision = np.linalg.inv(true_cov)
# Here's our resulting covariance matrix
print(true_cov)
# Verify that it's positive definite, since np.random.multivariate_normal
# complains about it not being positive definite for some reason.
# (Note that I'll be including a lot of sanity checking code in this notebook -
# it's a *huge* help for debugging)
print('eigenvalues: ', np.linalg.eigvals(true_cov))
[[4.  1.8]
 [1.8 1. ]]
eigenvalues:  [4.843075   0.15692513]

Hasilkan beberapa pengamatan sintetis

Perhatikan bahwa Probabilitas TensorFlow menggunakan konvensi bahwa dimensi awal data Anda mewakili indeks sampel, dan dimensi akhir data Anda mewakili dimensi sampel Anda.

Di sini kita ingin 100 sampel, yang masing-masing adalah vektor dengan panjang 2. Kita akan menghasilkan array my_data dengan bentuk (100, 2). my_data[i, :] adalah sampel $ i $ th, dan merupakan vektor dengan panjang 2.

(Ingatlah untuk membuat my_data memiliki tipe float32!)

# Set the seed so the results are reproducible.
np.random.seed(123)

# Now generate some observations of our random variable.
# (Note that I'm suppressing a bunch of spurious about the covariance matrix
# not being positive semidefinite via check_valid='ignore' because it really is
# positive definite!)
my_data = np.random.multivariate_normal(
    mean=true_mean, cov=true_cov, size=100,
    check_valid='ignore').astype(np.float32)
my_data.shape
(100, 2)

Sanity memeriksa pengamatan

Salah satu sumber bug potensial adalah mengacaukan data sintetis Anda! Mari lakukan beberapa pemeriksaan sederhana.

# Do a scatter plot of the observations to make sure they look like what we
# expect (higher variance on the x-axis, y values strongly correlated with x)
plt.scatter(my_data[:, 0], my_data[:, 1], alpha=0.75)
plt.show()

png

print('mean of observations:', np.mean(my_data, axis=0))
print('true mean:', true_mean)
mean of observations: [-0.24009615 -0.16638893]
true mean: [0. 0.]

print('covariance of observations:\n', np.cov(my_data, rowvar=False))
print('true covariance:\n', true_cov)
covariance of observations:
 [[3.95307734 1.68718486]
 [1.68718486 0.94910269]]
true covariance:
 [[4.  1.8]
 [1.8 1. ]]

Oke, sampel kami terlihat masuk akal. Langkah berikutnya.

Langkah 2: Terapkan fungsi kemungkinan di NumPy

Hal utama yang perlu kita tulis untuk melakukan pengambilan sampel MCMC di TF Probability adalah fungsi kemungkinan log. Secara umum, menulis TF sedikit lebih sulit daripada NumPy, jadi saya merasa terbantu untuk melakukan implementasi awal di NumPy. Saya akan membagi fungsi kemungkinan menjadi 2 bagian, fungsi kemungkinan data yang sesuai dengan $ P (data | parameter) $ dan fungsi kemungkinan sebelumnya yang sesuai dengan $ P (parameter) $.

Perhatikan bahwa fungsi NumPy ini tidak harus dioptimasi / di-vektorisasi secara super karena tujuannya hanya untuk menghasilkan beberapa nilai untuk pengujian. Ketepatan adalah pertimbangan utama!

Pertama kita akan menerapkan potongan kemungkinan log data. Itu sangat mudah. Satu hal yang perlu diingat adalah kita akan bekerja dengan matriks presisi, jadi kita akan membuat parameter yang sesuai.

def log_lik_data_numpy(precision, data):
  # np.linalg.inv is a really inefficient way to get the covariance matrix, but
  # remember we don't care about speed here
  cov = np.linalg.inv(precision)
  rv = scipy.stats.multivariate_normal(true_mean, cov)
  return np.sum(rv.logpdf(data))

# test case: compute the log likelihood of the data given the true parameters
log_lik_data_numpy(true_precision, my_data)
-280.81822950593767

Kita akan menggunakan Wishart sebelum matriks presisi karena ada solusi analitis untuk posterior (lihat tabel praktis Wikipedia tentang prior konjugasi ).

Distribusi Wishart memiliki 2 parameter:

  • jumlah derajat kebebasan (berlabel $ \ nu $ di Wikipedia)
  • matriks skala (berlabel $ V $ di Wikipedia)

Rata-rata untuk distribusi Wishart dengan parameter $ \ nu, V $ adalah $ E [W] = \ nu V $, dan variansnya adalah $ \ text {Var} (W_ {ij}) = \ nu (v_ {ij} ^ 2 + v_ {ii} v_ {jj}) $

Beberapa intuisi yang berguna: Anda dapat membuat sampel Wishart dengan menghasilkan $ \ nu $ penarikan independen $ x_1 \ ldots x _ {\ nu} $ dari variabel acak normal multivarian dengan mean 0 dan kovarian $ V $ lalu membentuk jumlah $ W = \ sum_ {i = 1} ^ {\ nu} x_i x_i ^ T $.

Jika Anda mengubah skala sampel Wishart dengan membaginya dengan $ \ nu $, Anda mendapatkan matriks kovarian sampel dari $ x_i $. Contoh matriks kovarians ini seharusnya mengarah ke $ V $ karena $ \ nu $ meningkat. Jika $ \ nu $ kecil, ada banyak variasi dalam contoh matriks kovarian, jadi nilai kecil $ \ nu $ sesuai dengan prior yang lebih lemah dan nilai besar $ \ nu $ sesuai dengan prior yang lebih kuat. Perhatikan bahwa $ \ nu $ setidaknya harus sebesar dimensi ruang yang Anda ambil sampelnya atau Anda akan menghasilkan matriks tunggal.

Kita akan menggunakan $ \ nu = 3 $ sehingga kita memiliki nilai sebelumnya yang lemah, dan kita akan mengambil $ V = \ frac {1} {\ nu} I $ yang akan menarik perkiraan kovarians kita menuju identitas (ingat bahwa mean adalah $ \ nu V $).

PRIOR_DF = 3
PRIOR_SCALE = np.eye(2, dtype=np.float32) / PRIOR_DF

def log_lik_prior_numpy(precision):
  rv = scipy.stats.wishart(df=PRIOR_DF, scale=PRIOR_SCALE)
  return rv.logpdf(precision)

# test case: compute the prior for the true parameters
log_lik_prior_numpy(true_precision)
-9.103606346649766

Distribusi Wishart adalah konjugasi sebelum untuk memperkirakan matriks presisi dari normal multivariat dengan mean $ \ mu $ yang diketahui.

Misalkan parameter Wishart sebelumnya adalah $ \ nu, V $ dan bahwa kita memiliki $ n $ observasi normal multivariasi kita, $ x_1, \ ldots, x_n $. Parameter posterior adalah $ n + \ nu, \ kiri (V ^ {- 1} + \ sum_ {i = 1} ^ n (x_i- \ mu) (x_i- \ mu) ^ T \ kanan) ^ {- 1 } $.

n = my_data.shape[0]
nu_prior = PRIOR_DF
v_prior = PRIOR_SCALE
nu_posterior = nu_prior + n
v_posterior = np.linalg.inv(np.linalg.inv(v_prior) + my_data.T.dot(my_data))
posterior_mean = nu_posterior * v_posterior
v_post_diag = np.expand_dims(np.diag(v_posterior), axis=1)
posterior_sd = np.sqrt(nu_posterior *
                       (v_posterior ** 2.0 + v_post_diag.dot(v_post_diag.T)))

Plot singkat dari posterior dan nilai-nilai sebenarnya. Perhatikan bahwa posterior dekat dengan posterior sampel tetapi sedikit menyusut ke arah identitas. Perhatikan juga bahwa nilai sebenarnya cukup jauh dari mode posterior - mungkin ini karena prior tidak cocok untuk data kami. Dalam masalah nyata kami mungkin akan melakukannya lebih baik dengan sesuatu seperti Wishart terbalik berskala sebelum kovarians (lihat, misalnya, komentar Andrew Gelman tentang subjek), tetapi kemudian kami tidak akan memiliki posterior analitik yang bagus.

sample_precision = np.linalg.inv(np.cov(my_data, rowvar=False, bias=False))
fig, axes = plt.subplots(2, 2)
fig.set_size_inches(10, 10)
for i in range(2):
  for j in range(2):
    ax = axes[i, j]
    loc = posterior_mean[i, j]
    scale = posterior_sd[i, j]
    xmin = loc - 3.0 * scale
    xmax = loc + 3.0 * scale
    x = np.linspace(xmin, xmax, 1000)
    y = scipy.stats.norm.pdf(x, loc=loc, scale=scale)
    ax.plot(x, y)
    ax.axvline(true_precision[i, j], color='red', label='True precision')
    ax.axvline(sample_precision[i, j], color='red', linestyle=':', label='Sample precision')
    ax.set_title('precision[%d, %d]' % (i, j))
plt.legend()
plt.show()

png

Langkah 3: Terapkan fungsi kemungkinan di TensorFlow

Spoiler: Upaya pertama kami tidak akan berhasil; kami akan membahas alasannya di bawah.

Tips : gunakan mode bersemangat TensorFlow saat mengembangkan fungsi kemungkinan Anda. Mode Eager membuat TF berperilaku lebih seperti NumPy - semuanya dijalankan dengan segera, jadi Anda bisa men-debug secara interaktif daripada harus menggunakan Session.run() . Lihat catatannya di sini .

Pendahuluan: Kelas distribusi

TFP memiliki kumpulan kelas distribusi yang akan kita gunakan untuk menghasilkan probabilitas log kita. Satu hal yang perlu diperhatikan adalah bahwa kelas-kelas ini bekerja dengan tensor sampel bukan hanya sampel tunggal - ini memungkinkan vektorisasi dan percepatan terkait.

Distribusi dapat bekerja dengan tensor sampel dalam 2 cara berbeda. Paling sederhana untuk mengilustrasikan 2 cara ini dengan contoh konkret yang melibatkan distribusi dengan parameter skalar tunggal. Saya akan menggunakan distribusi Poisson , yang memiliki parameter rate .

  • Jika kita membuat Poisson dengan nilai tunggal untuk parameter rate , panggilan ke metode sample() mengembalikan satu nilai. Nilai ini disebut event , dan dalam hal ini semua peristiwa adalah skalar.
  • Jika kita membuat Poisson dengan tensor nilai untuk parameter rate , panggilan ke metode sample() sekarang mengembalikan beberapa nilai, satu untuk setiap nilai dalam tensor laju. Objek bertindak sebagai kumpulan Poissons independen, masing-masing dengan rate-nya sendiri, dan masing-masing nilai yang dikembalikan oleh panggilan ke sample() sesuai dengan salah satu Poissons ini. Kumpulan peristiwa independen tetapi tidak terdistribusi secara identik ini disebut batch .
  • Metode sample() mengambil parameter sample_shape yang defaultnya adalah tupel kosong. Meneruskan nilai yang tidak kosong untuk sample_shape menghasilkan sampel yang mengembalikan beberapa batch. Kumpulan kumpulan ini disebut sample .

Metode log_prob() distribusi menggunakan data dengan cara yang paralel dengan cara sample() menghasilkannya. log_prob() mengembalikan probabilitas untuk sampel, yaitu untuk beberapa, kumpulan peristiwa independen.

  • Jika kita memiliki objek Poisson yang dibuat dengan rate skalar, setiap kelompok adalah skalar, dan jika kita memasukkan tensor sampel, kita akan mendapatkan tensor dengan ukuran probabilitas log yang sama.
  • Jika kita memiliki objek Poisson yang dibuat dengan tensor berbentuk T nilai rate , setiap batch adalah tensor bentuk T Jika kita memasukkan tensor sampel bentuk D, T, kita akan mendapatkan tensor probabilitas log bentuk D, T.

Di bawah ini adalah beberapa contoh yang menggambarkan kasus tersebut. Lihat buku catatan ini untuk tutorial yang lebih mendetail tentang acara, kelompok, dan bentuk.

# case 1: get log probabilities for a vector of iid draws from a single
# normal distribution
norm1 = tfd.Normal(loc=0., scale=1.)
probs1 = norm1.log_prob(tf.constant([1., 0.5, 0.]))

# case 2: get log probabilities for a vector of independent draws from
# multiple normal distributions with different parameters.  Note the vector
# values for loc and scale in the Normal constructor.
norm2 = tfd.Normal(loc=[0., 2., 4.], scale=[1., 1., 1.])
probs2 = norm2.log_prob(tf.constant([1., 0.5, 0.]))

print('iid draws from a single normal:', probs1.numpy())
print('draws from a batch of normals:', probs2.numpy())
iid draws from a single normal: [-1.4189385 -1.0439385 -0.9189385]
draws from a batch of normals: [-1.4189385 -2.0439386 -8.918939 ]

Kemungkinan log data

Pertama kita akan menerapkan fungsi kemungkinan log data.

VALIDATE_ARGS = True
ALLOW_NAN_STATS = False

Satu perbedaan utama dari kasus NumPy adalah bahwa fungsi kemungkinan TensorFlow kami perlu menangani vektor matriks presisi, bukan hanya matriks tunggal. Vektor parameter akan digunakan saat kami mengambil sampel dari beberapa rantai.

Kita akan membuat objek distribusi yang bekerja dengan sekumpulan matriks presisi (yaitu satu matriks per rantai).

Saat menghitung probabilitas log dari data kami, kami memerlukan data kami untuk direplikasi dengan cara yang sama seperti parameter kami sehingga ada satu salinan per variabel batch. Bentuk data yang direplikasi kami harus seperti berikut:

[sample shape, batch shape, event shape]

Dalam kasus kami, bentuk acara adalah 2 (karena kami bekerja dengan Gaussians 2-D). Bentuk sampelnya 100, karena kita punya 100 sampel. Bentuk batch hanya akan menjadi jumlah matriks presisi yang sedang kita kerjakan. Membuang-buang waktu untuk mereplikasi data setiap kali kami memanggil fungsi kemungkinan, jadi kami akan mereplikasi data terlebih dahulu dan meneruskan versi yang direplikasi.

Perhatikan bahwa ini adalah implementasi yang tidak efisien: MultivariateNormalFullCovariance mahal dibandingkan dengan beberapa alternatif yang akan kita bicarakan di bagian pengoptimalan di akhir.

def log_lik_data(precisions, replicated_data):
  n = tf.shape(precisions)[0]  # number of precision matrices
  # We're estimating a precision matrix; we have to invert to get log
  # probabilities.  Cholesky inversion should be relatively efficient,
  # but as we'll see later, it's even better if we can avoid doing the Cholesky
  # decomposition altogether.
  precisions_cholesky = tf.linalg.cholesky(precisions)
  covariances = tf.linalg.cholesky_solve(
      precisions_cholesky, tf.linalg.eye(2, batch_shape=[n]))
  rv_data = tfd.MultivariateNormalFullCovariance(
      loc=tf.zeros([n, 2]),
      covariance_matrix=covariances,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)

  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# For our test, we'll use a tensor of 2 precision matrices.
# We'll need to replicate our data for the likelihood function.
# Remember, TFP wants the data to be structured so that the sample dimensions
# are first (100 here), then the batch dimensions (2 here because we have 2
# precision matrices), then the event dimensions (2 because we have 2-D
# Gaussian data).  We'll need to add a middle dimension for the batch using
# expand_dims, and then we'll need to create 2 replicates in this new dimension
# using tile.
n = 2
replicated_data = np.tile(np.expand_dims(my_data, axis=1), reps=[1, 2, 1])
print(replicated_data.shape)
(100, 2, 2)

Tip: Satu hal yang menurut saya sangat membantu adalah menulis sedikit pemeriksaan kewarasan pada fungsi TensorFlow saya. Sangat mudah untuk mengacaukan vektorisasi di TF, jadi memiliki fungsi NumPy yang lebih sederhana adalah cara yang bagus untuk memverifikasi keluaran TF. Anggap ini sebagai tes unit kecil.

# check against the numpy implementation
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
n = precisions.shape[0]
lik_tf = log_lik_data(precisions, replicated_data=replicated_data).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.8182

Kemungkinan log sebelumnya

Prioritas lebih mudah karena kita tidak perlu khawatir tentang replikasi data.

@tf.function(autograph=False)
def log_lik_prior(precisions):
  rv_precision = tfd.WishartTriL(
      df=PRIOR_DF,
      scale_tril=tf.linalg.cholesky(PRIOR_SCALE),
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)
  return rv_precision.log_prob(precisions)
# check against the numpy implementation
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
n = precisions.shape[0]
lik_tf = log_lik_prior(precisions).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]))
  print('tensorflow:', lik_tf[i])
0
numpy: -2.2351873809649625
tensorflow: -2.2351875
1
numpy: -9.103606346649766
tensorflow: -9.103608

Bangun fungsi kemungkinan log bersama

Fungsi kemungkinan log data di atas bergantung pada pengamatan kita, tetapi sampler tidak akan memilikinya. Kita bisa menghilangkan ketergantungan tanpa menggunakan variabel global dengan menggunakan [closure] (https://en.wikipedia.org/wiki/Closure_ (program_komputer). Penutupan melibatkan fungsi luar yang membangun lingkungan yang berisi variabel yang dibutuhkan oleh fungsi batin.

def get_log_lik(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik(precision):
    return log_lik_data(precision, replicated_data) + log_lik_prior(precision)

  return _log_lik

Langkah 4: Contoh

Oke, waktunya mencicipi! Untuk mempermudah, kita hanya akan menggunakan 1 rantai dan kita akan menggunakan matriks identitas sebagai titik awal. Kami akan melakukannya dengan lebih hati-hati nanti.

Sekali lagi, ini tidak akan berhasil - kami akan mendapatkan pengecualian.

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  init_precision = tf.expand_dims(tf.eye(2), axis=0)

  # Use expand_dims because we want to pass in a tensor of starting values
  log_lik_fn = get_log_lik(my_data, n_chains=1)

  # we'll just do a few steps here
  num_results = 10
  num_burnin_steps = 10
  states = tfp.mcmc.sample_chain(
     num_results=num_results,
     num_burnin_steps=num_burnin_steps,
     current_state=[
         init_precision,
     ],
     kernel=tfp.mcmc.HamiltonianMonteCarlo(
         target_log_prob_fn=log_lik_fn,
         step_size=0.1,
         num_leapfrog_steps=3,
         seed=123),
     trace_fn=None,
     parallel_iterations=1)
  return states

try:
  states = sample()
except Exception as e:
  # shorten the giant stack trace
  lines = str(e).split('\n')
  print('\n'.join(lines[:5]+['...']+lines[-3:]))
 Cholesky decomposition was not successful. The input might not be valid.
     [[{ {node mcmc_sample_chain/trace_scan/while/body/_79/smart_for_loop/while/body/_371/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/body/_537/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/StatefulPartitionedCall/Cholesky} }]] [Op:__inference_sample_2849]

Function call stack:
sample
...
Function call stack:
sample


Mengidentifikasi masalah

InvalidArgumentError (see above for traceback): Cholesky decomposition was not successful. The input might not be valid. Itu tidak terlalu membantu. Mari kita lihat apakah kita bisa mencari tahu lebih banyak tentang apa yang terjadi.

  • Kami akan mencetak parameter untuk setiap langkah sehingga kami dapat melihat nilai hal-hal yang gagal
  • Kami akan menambahkan beberapa pernyataan untuk mencegah masalah tertentu.

Pernyataan itu rumit karena merupakan operasi TensorFlow, dan kami harus berhati-hati agar dieksekusi dan tidak dioptimalkan dari grafik. Sebaiknya Anda membaca ringkasan debugging TensorFlow ini jika Anda tidak terbiasa dengan pernyataan TF. Anda dapat secara eksplisit memaksa pernyataan untuk dieksekusi menggunakan tf.control_dependencies (lihat komentar pada kode di bawah).

Fungsi Print asli TensorFlow memiliki perilaku yang sama dengan assertion - ini adalah operasi, dan Anda perlu berhati-hati untuk memastikannya dijalankan. Print menyebabkan sakit kepala tambahan saat kita bekerja di notebook: outputnya dikirim ke stderr , dan stderr tidak ditampilkan di sel. Kami akan menggunakan trik di sini: alih-alih menggunakan tf.Print , kami akan membuat operasi cetak TensorFlow kami sendiri melalui tf.pyfunc . Seperti halnya pernyataan, kita harus memastikan metode kita dijalankan.

def get_log_lik_verbose(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  def _log_lik(precisions):
    # An internal method we'll make into a TensorFlow operation via tf.py_func
    def _print_precisions(precisions):
      print('precisions:\n', precisions)
      return False  # operations must return something!
    # Turn our method into a TensorFlow operation
    print_op = tf.compat.v1.py_func(_print_precisions, [precisions], tf.bool)

    # Assertions are also operations, and some care needs to be taken to ensure
    # that they're executed
    assert_op = tf.assert_equal(
        precisions, tf.linalg.matrix_transpose(precisions),
        message='not symmetrical', summarize=4, name='symmetry_check')

    # The control_dependencies statement forces its arguments to be executed
    # before subsequent operations
    with tf.control_dependencies([print_op, assert_op]):
      return (log_lik_data(precisions, replicated_data) +
              log_lik_prior(precisions))

  return _log_lik
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  init_precision = tf.eye(2)[tf.newaxis, ...]
  log_lik_fn = get_log_lik_verbose(my_data)
  # we'll just do a few steps here
  num_results = 10
  num_burnin_steps = 10
  states = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          init_precision,
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=log_lik_fn,
          step_size=0.1,
          num_leapfrog_steps=3,
          seed=123),
      trace_fn=None,
      parallel_iterations=1)

try:
  states = sample()
except Exception as e:
  # shorten the giant stack trace
  lines = str(e).split('\n')
  print('\n'.join(lines[:5]+['...']+lines[-3:]))
precisions:
 [[[1. 0.]
  [0. 1.]]]
precisions:
 [[[1. 0.]
  [0. 1.]]]
precisions:
 [[[ 0.24315196 -0.2761638 ]
  [-0.33882257  0.8622    ]]]
 assertion failed: [not symmetrical] [Condition x == y did not hold element-wise:] [x (leapfrog_integrate_one_step/add:0) = ] [[[0.243151963 -0.276163787][-0.338822573 0.8622]]] [y (leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/matrix_transpose/transpose:0) = ] [[[0.243151963 -0.338822573][-0.276163787 0.8622]]]
     [[{ {node mcmc_sample_chain/trace_scan/while/body/_96/smart_for_loop/while/body/_381/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/body/_503/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/symmetry_check_1/Assert/AssertGuard/else/_577/Assert} }]] [Op:__inference_sample_4837]

Function call stack:
sample
...
Function call stack:
sample


Mengapa ini gagal

Nilai parameter baru pertama yang dicoba oleh sampler adalah matriks asimetris. Hal ini menyebabkan dekomposisi Cholesky gagal, karena dekomposisi ini hanya ditentukan untuk matriks simetris (dan pasti positif).

Masalahnya di sini adalah bahwa parameter yang kita minati adalah matriks presisi, dan matriks presisi harus nyata, simetris, dan pasti positif. Sampler tidak tahu apa-apa tentang batasan ini (kecuali mungkin melalui gradien), jadi sangat mungkin bahwa sampler akan mengusulkan nilai yang tidak valid, yang mengarah ke pengecualian, terutama jika ukuran langkahnya besar.

Dengan sampel Hamiltonian Monte Carlo, kita mungkin dapat mengatasi masalah tersebut dengan menggunakan ukuran langkah yang sangat kecil, karena gradien harus menjauhkan parameter dari wilayah yang tidak valid, tetapi ukuran langkah yang kecil berarti konvergensi yang lambat. Dengan sampler Metropolis-Hastings, yang tidak tahu apa-apa tentang gradien, kami dikutuk.

Versi 2: melakukan parameter ulang ke parameter yang tidak dibatasi

Ada solusi langsung untuk masalah di atas: kita dapat mengatur ulang model kita sedemikian rupa sehingga parameter baru tidak lagi memiliki batasan ini. TFP menyediakan seperangkat alat yang berguna - bijector - untuk melakukan hal itu.

Reparameterisasi dengan bijector

Matriks presisi kita harus nyata dan simetris; kami menginginkan parameterisasi alternatif yang tidak memiliki batasan ini. Titik awal adalah faktorisasi Cholesky dari matriks presisi. Faktor Cholesky masih dibatasi - mereka adalah segitiga bawah, dan elemen diagonalnya harus positif. Namun, jika kita mengambil log dari diagonal faktor Cholesky, log tidak lagi dibatasi menjadi positif, dan kemudian jika kita meratakan bagian segitiga bawah menjadi vektor 1-D, kita tidak lagi memiliki batasan segitiga bawah. . Hasil dalam kasus kami akan menjadi vektor panjang 3 tanpa batasan.

( Manual Stan memiliki bab yang bagus tentang penggunaan transformasi untuk menghilangkan berbagai jenis batasan pada parameter.)

Reparameterisasi ini memiliki sedikit pengaruh pada fungsi kemungkinan log data kita - kita hanya perlu membalik transformasi kita sehingga kita mendapatkan kembali matriks presisi - tetapi efek pada prior lebih rumit. Kami telah menetapkan bahwa probabilitas matriks presisi yang diberikan diberikan oleh distribusi Wishart; berapa probabilitas matriks yang kita ubah?

Ingatlah bahwa jika kita menerapkan fungsi monotonik $ g $ ke variabel acak 1-D $ X $, $ Y = g (X) $, kepadatan untuk $ Y $ diberikan oleh

$$ f_Y(y) = | \frac{d}{dy}(g^{-1}(y)) | f_X(g^{-1}(y)) $$

Turunan dari $ g ^ {- 1} $ term menjelaskan cara $ g $ mengubah volume lokal. Untuk variabel acak berdimensi lebih tinggi, faktor korektif adalah nilai absolut dari determinan Jacobian $ g ^ {- 1} $ (lihat di sini ).

Kita harus menambahkan Jacobian dari transformasi terbalik ke dalam fungsi kemungkinan log prior kita. Untungnya, kelas Bijector TFP bisa menangani ini untuk kita.

Kelas Bijector digunakan untuk merepresentasikan fungsi halus yang dapat dibalik yang digunakan untuk mengubah variabel dalam fungsi kepadatan probabilitas. Semua Bijector memiliki metode forward() yang melakukan transformasi, metode inverse() yang membalikkannya, dan metode forward_log_det_jacobian() dan inverse_log_det_jacobian() yang menyediakan koreksi Jacobian yang kami butuhkan saat kami melakukan reparaterisasi pdf.

TFP menyediakan kumpulan bijector berguna yang dapat kita gabungkan melalui komposisi melalui operator Chain untuk membentuk transformasi yang cukup rumit. Dalam kasus kami, kami akan membuat 3 bijector berikut (operasi dalam rantai dilakukan dari kanan ke kiri):

  1. Langkah pertama dari transformasi kita adalah melakukan faktorisasi Cholesky pada matriks presisi. Tidak ada kelas Bijector untuk itu; namun, bijector CholeskyOuterProduct mengambil produk dari 2 faktor Cholesky. Kita dapat menggunakan kebalikan dari operasi itu menggunakan operator Invert .
  2. Langkah selanjutnya adalah mengambil log dari elemen diagonal dari faktor Cholesky. Kami melakukannya melalui bijector TransformDiagonal dan kebalikan dari bijector Exp .
  3. Akhirnya kita meratakan bagian segitiga bawah dari matriks menjadi vektor menggunakan kebalikan dari FillTriangular FillTriangular.
# Our transform has 3 stages that we chain together via composition:
precision_to_unconstrained = tfb.Chain([
    # step 3: flatten the lower triangular portion of the matrix
    tfb.Invert(tfb.FillTriangular(validate_args=VALIDATE_ARGS)),
    # step 2: take the log of the diagonals    
    tfb.TransformDiagonal(tfb.Invert(tfb.Exp(validate_args=VALIDATE_ARGS))),
    # step 1: decompose the precision matrix into its Cholesky factors
    tfb.Invert(tfb.CholeskyOuterProduct(validate_args=VALIDATE_ARGS)),
])
# sanity checks
m = tf.constant([[1., 2.], [2., 8.]])
m_fwd = precision_to_unconstrained.forward(m)
m_inv = precision_to_unconstrained.inverse(m_fwd)

# bijectors handle tensors of values, too!
m2 = tf.stack([m, tf.eye(2)])
m2_fwd = precision_to_unconstrained.forward(m2)
m2_inv = precision_to_unconstrained.inverse(m2_fwd)

print('single input:')
print('m:\n', m.numpy())
print('precision_to_unconstrained(m):\n', m_fwd.numpy())
print('inverse(precision_to_unconstrained(m)):\n', m_inv.numpy())
print()

print('tensor of inputs:')
print('m2:\n', m2.numpy())
print('precision_to_unconstrained(m2):\n', m2_fwd.numpy())
print('inverse(precision_to_unconstrained(m2)):\n', m2_inv.numpy())
single input:
m:
 [[1. 2.]
 [2. 8.]]
precision_to_unconstrained(m):
 [0.6931472 2.        0.       ]
inverse(precision_to_unconstrained(m)):
 [[1. 2.]
 [2. 8.]]

tensor of inputs:
m2:
 [[[1. 2.]
  [2. 8.]]

 [[1. 0.]
  [0. 1.]]]
precision_to_unconstrained(m2):
 [[0.6931472 2.        0.       ]
 [0.        0.        0.       ]]
inverse(precision_to_unconstrained(m2)):
 [[[1. 2.]
  [2. 8.]]

 [[1. 0.]
  [0. 1.]]]

Kelas TransformedDistribution mengotomatiskan proses penerapan bijector ke distribusi dan membuat koreksi Jacobian yang diperlukan ke log_prob() . Sebelumnya baru kami menjadi:

def log_lik_prior_transformed(transformed_precisions):
  rv_precision = tfd.TransformedDistribution(
      tfd.WishartTriL(
          df=PRIOR_DF,
          scale_tril=tf.linalg.cholesky(PRIOR_SCALE),
          validate_args=VALIDATE_ARGS,
          allow_nan_stats=ALLOW_NAN_STATS),
      bijector=precision_to_unconstrained,
      validate_args=VALIDATE_ARGS)
  return rv_precision.log_prob(transformed_precisions)
# Check against the numpy implementation.  Note that when comparing, we need
# to add in the Jacobian correction.
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
transformed_precisions = precision_to_unconstrained.forward(precisions)
lik_tf = log_lik_prior_transformed(transformed_precisions).numpy()
corrections = precision_to_unconstrained.inverse_log_det_jacobian(
    transformed_precisions, event_ndims=1).numpy()
n = precisions.shape[0]

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]) + corrections[i])
  print('tensorflow:', lik_tf[i])
0
numpy: -0.8488930160357633
tensorflow: -0.84889317
1
numpy: -7.305657151741624
tensorflow: -7.305659

Kita hanya perlu membalik transformasi untuk kemungkinan log data kita:

precision = precision_to_unconstrained.inverse(transformed_precision)

Karena kita sebenarnya menginginkan faktorisasi Cholesky dari matriks presisi, akan lebih efisien jika melakukan inversi parsial saja di sini. Namun, kita akan meninggalkan pengoptimalan untuk nanti dan membiarkan invers parsial sebagai latihan untuk pembaca.

def log_lik_data_transformed(transformed_precisions, replicated_data):
  # We recover the precision matrix by inverting our bijector.  This is
  # inefficient since we really want the Cholesky decomposition of the
  # precision matrix, and the bijector has that in hand during the inversion,
  # but we'll worry about efficiency later.
  n = tf.shape(transformed_precisions)[0]
  precisions = precision_to_unconstrained.inverse(transformed_precisions)
  precisions_cholesky = tf.linalg.cholesky(precisions)
  covariances = tf.linalg.cholesky_solve(
      precisions_cholesky, tf.linalg.eye(2, batch_shape=[n]))
  rv_data = tfd.MultivariateNormalFullCovariance(
      loc=tf.zeros([n, 2]),
      covariance_matrix=covariances,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)

  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# sanity check
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
transformed_precisions = precision_to_unconstrained.forward(precisions)
lik_tf = log_lik_data_transformed(
    transformed_precisions, replicated_data).numpy()

for i in range(precisions.shape[0]):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.8182

Sekali lagi kami membungkus fungsi baru kami dalam sebuah closure.

def get_log_lik_transformed(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik_transformed(transformed_precisions):
    return (log_lik_data_transformed(transformed_precisions, replicated_data) +
            log_lik_prior_transformed(transformed_precisions))

  return _log_lik_transformed
# make sure everything runs
log_lik_fn = get_log_lik_transformed(my_data)
m = tf.eye(2)[tf.newaxis, ...]
lik = log_lik_fn(precision_to_unconstrained.forward(m)).numpy()
print(lik)
[-431.5611]

Contoh

Sekarang kita tidak perlu khawatir tentang sampler kita meledak karena nilai parameter yang tidak valid, mari buat beberapa sampel nyata.

Sampler bekerja dengan versi parameter yang tidak dibatasi, jadi kita perlu mengubah nilai awal menjadi versi yang tidak dibatasi. Sampel yang kami hasilkan juga semuanya akan dalam bentuk tidak dibatasi, jadi kami perlu mengubahnya kembali. Bijector dibuat vektor, jadi mudah untuk melakukannya.

# We'll choose a proper random initial value this time
np.random.seed(123)
initial_value_cholesky = np.array(
    [[0.5 + np.random.uniform(), 0.0],
     [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
    dtype=np.float32)
initial_value =  initial_value_cholesky.dot(
  initial_value_cholesky.T)[np.newaxis, ...]

# The sampler works with unconstrained values, so we'll transform our initial
# value
initial_value_transformed = precision_to_unconstrained.forward(
  initial_value).numpy()
# Sample!
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_transformed(my_data, n_chains=1)

  num_results = 1000
  num_burnin_steps = 1000

  states, is_accepted = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          initial_value_transformed,
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=log_lik_fn,
          step_size=0.1,
          num_leapfrog_steps=3,
          seed=123),
      trace_fn=lambda _, pkr: pkr.is_accepted,
      parallel_iterations=1)
  # transform samples back to their constrained form
  precision_samples = [precision_to_unconstrained.inverse(s) for s in states]
  return states, precision_samples, is_accepted

states, precision_samples, is_accepted = sample()

Mari kita bandingkan mean dari keluaran sampler kita dengan mean posterior analitik!

print('True posterior mean:\n', posterior_mean)
print('Sample mean:\n', np.mean(np.reshape(precision_samples, [-1, 2, 2]), axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Sample mean:
 [[ 1.4315274  -0.25587553]
 [-0.25587553  0.5740424 ]]

Kami jauh! Mari kita cari tahu alasannya. Pertama mari kita lihat sampel kita.

np.reshape(precision_samples, [-1, 2, 2])
array([[[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       ...,

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]]], dtype=float32)

Aduh - sepertinya semuanya memiliki nilai yang sama. Mari kita cari tahu alasannya.

Variabel kernel_results_ adalah tupel bernama yang memberikan informasi tentang sampler di setiap status. Bidang is_accepted adalah kuncinya di sini.

# Look at the acceptance for the last 100 samples
print(np.squeeze(is_accepted)[-100:])
print('Fraction of samples accepted:', np.mean(np.squeeze(is_accepted)))
[False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False]
Fraction of samples accepted: 0.0

Semua sampel kami ditolak! Agaknya ukuran langkah kami terlalu besar. Saya memilih stepsize=0.1 murni secara sembarangan.

Versi 3: pengambilan sampel dengan ukuran langkah adaptif

Sejak pengambilan sampel dengan pilihan ukuran langkah sewenang-wenang saya gagal, kami memiliki beberapa item agenda:

  1. menerapkan ukuran langkah adaptif, dan
  2. melakukan beberapa pemeriksaan konvergensi.

Ada beberapa contoh kode yang bagus di tensorflow_probability/python/mcmc/hmc.py untuk mengimplementasikan ukuran langkah adaptif. Saya telah menyesuaikannya di bawah ini.

Perhatikan bahwa ada pernyataan sess.run() terpisah untuk setiap langkah. Ini sangat membantu untuk debugging, karena ini memungkinkan kita dengan mudah menambahkan beberapa diagnostik per langkah jika perlu. Misalnya, kami dapat menunjukkan kemajuan tambahan, waktu setiap langkah, dll.

Tip: Salah satu cara yang tampaknya umum untuk mengacaukan pengambilan sampel Anda adalah dengan membuat grafik Anda tumbuh dalam lingkaran. (Alasan menyelesaikan grafik sebelum sesi dijalankan adalah untuk mencegah masalah seperti itu.) Namun, jika Anda belum menggunakan finalize (), pemeriksaan debugging yang berguna untuk memeriksa apakah kode Anda melambat hingga perayapan adalah untuk mencetak grafik ukuran di setiap langkah melalui len(mygraph.get_operations()) - jika panjangnya bertambah, Anda mungkin melakukan sesuatu yang buruk.

Kami akan menjalankan 3 rantai independen di sini. Melakukan beberapa perbandingan antara rantai akan membantu kami memeriksa konvergensi.

# The number of chains is determined by the shape of the initial values.
# Here we'll generate 3 chains, so we'll need a tensor of 3 initial values.
N_CHAINS = 3

np.random.seed(123)

initial_values = []
for i in range(N_CHAINS):
  initial_value_cholesky = np.array(
      [[0.5 + np.random.uniform(), 0.0],
       [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
      dtype=np.float32)
  initial_values.append(initial_value_cholesky.dot(initial_value_cholesky.T))
initial_values = np.stack(initial_values)

initial_values_transformed = precision_to_unconstrained.forward(
  initial_values).numpy()
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_transformed(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=hmc,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states, is_accepted = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values_transformed,
      kernel=adapted_kernel,
      trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
      parallel_iterations=1)
  # transform samples back to their constrained form
  precision_samples = precision_to_unconstrained.inverse(states)
  return states, precision_samples, is_accepted

states, precision_samples, is_accepted = sample() 

Pemeriksaan cepat: tingkat penerimaan kami selama pengambilan sampel kami mendekati target kami di 0,651.

print(np.mean(is_accepted))
0.6190666666666667

Bahkan lebih baik lagi, mean sampel dan deviasi standar kita mendekati apa yang kita harapkan dari solusi analitik.

precision_samples_reshaped = np.reshape(precision_samples, [-1, 2, 2])
print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.96426415 -1.6519215 ]
 [-1.6519215   3.8614824 ]]

print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13622096 0.25235635]
 [0.25235635 0.5394968 ]]

Memeriksa konvergensi

Secara umum, kami tidak akan memiliki solusi analitik untuk diperiksa, jadi kami perlu memastikan sampler telah berkumpul. Satu pemeriksaan standar adalah statistik Gelman-Rubin $ \ hat {R} $, yang memerlukan banyak rantai pengambilan sampel. $ \ hat {R} $ mengukur sejauh mana varians (sarana) antara rantai melebihi apa yang diharapkan jika rantai didistribusikan secara identik. Nilai $ \ hat {R} $ mendekati 1 digunakan untuk menunjukkan perkiraan konvergensi. Lihat sumber untuk detailnya.

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0038308 1.0005717]
 [1.0005717 1.0006068]]

Kritik model

Jika kami tidak memiliki solusi analitik, ini akan menjadi waktu untuk melakukan kritik model nyata.

Berikut adalah beberapa histogram cepat dari komponen sampel yang berhubungan dengan kebenaran dasar kami (berwarna merah). Perhatikan bahwa sampel telah dikecilkan dari nilai matriks presisi sampel menjadi matriks identitas sebelumnya.

fig, axes = plt.subplots(2, 2, sharey=True)
fig.set_size_inches(8, 8)
for i in range(2):
  for j in range(2):
    ax = axes[i, j]
    ax.hist(precision_samples_reshaped[:, i, j])
    ax.axvline(true_precision[i, j], color='red',
               label='True precision')
    ax.axvline(sample_precision[i, j], color='red', linestyle=':',
               label='Sample precision')
    ax.set_title('precision[%d, %d]' % (i, j))
plt.tight_layout()
plt.legend()
plt.show()

png

Beberapa diagram sebar dari pasangan komponen presisi menunjukkan bahwa karena struktur korelasi posterior, nilai posterior sebenarnya tidak sekecil yang terlihat dari margin di atas.

fig, axes = plt.subplots(4, 4)
fig.set_size_inches(12, 12)
for i1 in range(2):
  for j1 in range(2):
    index1 = 2 * i1 + j1
    for i2 in range(2):
      for j2 in range(2):
        index2 = 2 * i2 + j2
        ax = axes[index1, index2]
        ax.scatter(precision_samples_reshaped[:, i1, j1],
                   precision_samples_reshaped[:, i2, j2], alpha=0.1)
        ax.axvline(true_precision[i1, j1], color='red')
        ax.axhline(true_precision[i2, j2], color='red')
        ax.axvline(sample_precision[i1, j1], color='red', linestyle=':')
        ax.axhline(sample_precision[i2, j2], color='red', linestyle=':')
        ax.set_title('(%d, %d) vs (%d, %d)' % (i1, j1, i2, j2))
plt.tight_layout()
plt.show()

png

Versi 4: pengambilan sampel parameter terbatas yang lebih sederhana

Bijector membuat pengambilan sampel matriks presisi secara langsung, tetapi ada cukup banyak konversi manual ke dan dari representasi yang tidak dibatasi. Ada cara yang lebih mudah!

The TransformedTransitionKernel

TransformedTransitionKernel menyederhanakan proses ini. Ini membungkus sampler Anda dan menangani semua konversi. Sebagai argumen, diperlukan daftar bijector yang memetakan nilai parameter tidak dibatasi ke yang dibatasi. Jadi di sini kita membutuhkan kebalikan dari bijector precision_to_unconstrained kita gunakan di atas. Kita bisa saja menggunakan tfb.Invert(precision_to_unconstrained) , tapi itu akan melibatkan pengambilan invers dari invers (TensorFlow tidak cukup pintar untuk menyederhanakan tf.Invert(tf.Invert()) ke tf.Identity()) , jadi kita Saya baru saja menulis bijector baru.

Bijector pembatas

# The bijector we need for the TransformedTransitionKernel is the inverse of
# the one we used above
unconstrained_to_precision = tfb.Chain([
    # step 3: take the product of Cholesky factors
    tfb.CholeskyOuterProduct(validate_args=VALIDATE_ARGS),
    # step 2: exponentiate the diagonals    
    tfb.TransformDiagonal(tfb.Exp(validate_args=VALIDATE_ARGS)),
    # step 1: map a vector to a lower triangular matrix
    tfb.FillTriangular(validate_args=VALIDATE_ARGS),
])
# quick sanity check
m = [[1., 2.], [2., 8.]]
m_inv = unconstrained_to_precision.inverse(m).numpy()
m_fwd = unconstrained_to_precision.forward(m_inv).numpy()

print('m:\n', m)
print('unconstrained_to_precision.inverse(m):\n', m_inv)
print('forward(unconstrained_to_precision.inverse(m)):\n', m_fwd)
m:
 [[1.0, 2.0], [2.0, 8.0]]
unconstrained_to_precision.inverse(m):
 [0.6931472 2.        0.       ]
forward(unconstrained_to_precision.inverse(m)):
 [[1. 2.]
 [2. 8.]]

Pengambilan sampel dengan TransformedTransitionKernel

Dengan TransformedTransitionKernel , kita tidak lagi harus melakukan transformasi manual pada parameter kita. Nilai awal dan sampel kami semuanya adalah matriks presisi; kita hanya perlu meneruskan bijector kita yang tidak membatasi ke kernel dan itu menangani semua transformasi.

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  ttk = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=hmc, bijector=unconstrained_to_precision)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=ttk,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values,
      kernel=adapted_kernel,
      trace_fn=None,
      parallel_iterations=1)
  # transform samples back to their constrained form
  return states

precision_samples  = sample()

Memeriksa konvergensi

Pemeriksaan konvergensi $ \ hat {R} $ kelihatannya bagus!

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0013582 1.0019467]
 [1.0019467 1.0011805]]

Perbandingan terhadap analitik posterior

Sekali lagi mari kita periksa dengan analitik posterior.

# The output samples have shape [n_steps, n_chains, 2, 2]
# Flatten them to [n_steps * n_chains, 2, 2] via reshape:
precision_samples_reshaped = np.reshape(precision_samples, [-1, 2, 2])
print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.96687526 -1.6552585 ]
 [-1.6552585   3.867676  ]]

print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13329624 0.24913791]
 [0.24913791 0.53983927]]

Optimasi

Sekarang kita sudah menjalankan semuanya, mari kita lakukan versi yang lebih dioptimalkan. Kecepatan tidak terlalu menjadi masalah untuk contoh ini, tetapi begitu matriks menjadi lebih besar, beberapa pengoptimalan akan membuat perbedaan besar.

Satu peningkatan kecepatan besar yang dapat kami lakukan adalah melakukan parameter ulang dalam hal dekomposisi Cholesky. Alasannya adalah fungsi kemungkinan data kami membutuhkan kovarian dan matriks presisi. Inversi matriks mahal ($ O (n ^ 3) $ untuk $ n \ times n $ matriks), dan jika kita membuat parameter dalam hal kovarian atau matriks presisi, kita perlu melakukan inversi untuk mendapatkan yang lain.

Sebagai pengingat, matriks simetris nyata, positif-pasti, dan simetris $ M $ dapat diuraikan menjadi hasil perkalian dari bentuk $ M = LL ^ T $ di mana matriks $ L $ adalah segitiga bawah dan memiliki diagonal positif. Mengingat dekomposisi Cholesky sebesar $ M $, kita dapat lebih efisien memperoleh $ M $ (hasil kali matriks segitiga bawah dan atas) dan $ M ^ {- 1} $ (melalui substitusi balik). Faktorisasi Cholesky sendiri tidak murah untuk dihitung, tetapi jika kita membuat parameter dalam hal faktor Cholesky, kita hanya perlu menghitung faktorisasi Choleksy dari nilai parameter awal.

Menggunakan dekomposisi Cholesky dari matriks kovarians

TFP memiliki versi distribusi normal multivariate, MultivariateNormalTriL , yang diparameterisasi dalam hal faktor Cholesky dari matriks kovarian. Jadi jika kita melakukan parameterisasi dalam hal faktor Cholesky dari matriks kovarians, kita dapat menghitung kemungkinan log data secara efisien. Tantangannya adalah menghitung kemungkinan log sebelumnya dengan efisiensi yang sama.

Jika kita memiliki versi distribusi Wishart terbalik yang bekerja dengan faktor sampel Cholesky, kita sudah siap. Sayangnya, kami tidak. (Namun, tim akan menyambut pengiriman kode!) Sebagai alternatif, kita dapat menggunakan versi distribusi Wishart yang bekerja dengan faktor Cholesky sampel bersama dengan rantai bijector.

Saat ini, kami kehilangan beberapa bijector stok untuk membuat semuanya benar-benar efisien, tetapi saya ingin menunjukkan prosesnya sebagai latihan dan ilustrasi yang berguna tentang kekuatan bijector TFP.

Distribusi Wishart yang beroperasi pada faktor Cholesky

Distribusi Wishart memiliki tanda yang berguna, input_output_cholesky , yang menetapkan bahwa matriks masukan dan keluaran harus berupa faktor Cholesky. Lebih efisien dan menguntungkan secara numerik untuk bekerja dengan faktor Cholesky daripada matriks lengkap, itulah mengapa ini diinginkan. Poin penting tentang semantik bendera: ini hanya indikasi bahwa representasi input dan output ke distribusi harus berubah - ini tidak menunjukkan reparameterisasi penuh dari distribusi, yang akan melibatkan koreksi Jacobian ke log_prob() fungsi. Kami sebenarnya ingin melakukan reparameterisasi penuh ini, jadi kami akan membangun distribusi kami sendiri.

# An optimized Wishart distribution that has been transformed to operate on
# Cholesky factors instead of full matrices.  Note that we gain a modest
# additional speedup by specifying the Cholesky factor of the scale matrix
# (i.e. by passing in the scale_tril parameter instead of scale).

class CholeskyWishart(tfd.TransformedDistribution):
  """Wishart distribution reparameterized to use Cholesky factors."""
  def __init__(self,
      df,
      scale_tril,
      validate_args=False,
      allow_nan_stats=True,
      name='CholeskyWishart'):
    # Wishart has a bunch of methods that we want to support but not
    # implement.  We'll subclass TransformedDistribution here to take care of
    # those.  We'll override the few for which speed is critical and implement
    # them with a separate Wishart for which input_output_cholesky=True
    super(CholeskyWishart, self).__init__(
        distribution=tfd.WishartTriL(
            df=df,
            scale_tril=scale_tril,
            input_output_cholesky=False,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats),
        bijector=tfb.Invert(tfb.CholeskyOuterProduct()),
        validate_args=validate_args,
        name=name
    )
    # Here's the Cholesky distribution we'll use for log_prob() and sample()
    self.cholesky = tfd.WishartTriL(
        df=df,
        scale_tril=scale_tril,
        input_output_cholesky=True,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats)
    
  def _log_prob(self, x):
    return (self.cholesky.log_prob(x) +
            self.bijector.inverse_log_det_jacobian(x, event_ndims=2))

  def _sample_n(self, n, seed=None):
    return self.cholesky._sample_n(n, seed)
# some checks
PRIOR_SCALE_CHOLESKY = np.linalg.cholesky(PRIOR_SCALE)

@tf.function(autograph=False)
def compute_log_prob(m):
  w_transformed = tfd.TransformedDistribution(
      tfd.WishartTriL(df=PRIOR_DF, scale_tril=PRIOR_SCALE_CHOLESKY),
      bijector=tfb.Invert(tfb.CholeskyOuterProduct()))
  w_optimized = CholeskyWishart(
      df=PRIOR_DF, scale_tril=PRIOR_SCALE_CHOLESKY)
  log_prob_transformed = w_transformed.log_prob(m)
  log_prob_optimized = w_optimized.log_prob(m)
  return log_prob_transformed, log_prob_optimized

for matrix in [np.eye(2, dtype=np.float32),
               np.array([[1., 0.], [2., 8.]], dtype=np.float32)]:
  log_prob_transformed, log_prob_optimized = [
      t.numpy() for t in compute_log_prob(matrix)]
  print('Transformed Wishart:', log_prob_transformed)
  print('Optimized Wishart', log_prob_optimized)
Transformed Wishart: -0.84889317
Optimized Wishart -0.84889317
Transformed Wishart: -99.269455
Optimized Wishart -99.269455

Membangun distribusi Wishart terbalik

Kami memiliki matriks kovarian kami $ C $ didekomposisi menjadi $ C = LL ^ T $ di mana $ L $ adalah segitiga bawah dan memiliki diagonal positif. Kami ingin mengetahui probabilitas $ L $ mengingat $ C \ sim W ^ {- 1} (\ nu, V) $ di mana $ W ^ {- 1} $ adalah distribusi Wishart terbalik.

Distribusi Wishart terbalik memiliki properti bahwa jika $ C \ sim W ^ {- 1} (\ nu, V) $, maka matriks presisi $ C ^ {- 1} \ sim W (\ nu, V ^ {- 1 }) $. Jadi kita bisa mendapatkan probabilitas $ L $ melalui TransformedDistribution yang mengambil parameter distribusi Wishart dan bijector yang memetakan faktor Cholesky matriks presisi ke faktor Cholesky dari kebalikannya.

Cara mudah (tetapi tidak super efisien) untuk beralih dari faktor Cholesky dari $ C ^ {- 1} $ ke $ L $ adalah dengan membalikkan faktor Cholesky dengan pemecahan-balik, kemudian membentuk matriks kovarians dari faktor-faktor yang dibalik ini, dan kemudian melakukan faktorisasi Cholesky.

Biarkan dekomposisi Cholesky dari $ C ^ {- 1} = MM ^ T $. $ M $ adalah segitiga bawah, jadi kita bisa membalikkannya menggunakan MatrixInverseTriL MatrixInverseTriL.

Membentuk $ C $ dari $ M ^ {- 1} $ agak rumit: $ C = (MM ^ T) ^ {- 1} = M ^ {- T} M ^ {- 1} = M ^ {- T } (M ^ {- T}) ^ T $. $ M $ adalah segitiga bawah, jadi $ M ^ {- 1} $ juga akan menjadi segitiga bawah, dan $ M ^ {- T} $ akan menjadi segitiga atas. Bijector CholeskyOuterProduct() hanya bekerja dengan matriks segitiga bawah, jadi kita tidak dapat menggunakannya untuk membentuk $ C $ dari $ M ^ {- T} $. Solusi kami adalah rantai bijector yang mengubah baris dan kolom dari sebuah matriks.

Untungnya logika ini dikemas dalam bijector CholeskyToInvCholesky !

Menggabungkan semua bagian

# verify that the bijector works
m = np.array([[1., 0.], [2., 8.]], dtype=np.float32)
c_inv = m.dot(m.T)
c = np.linalg.inv(c_inv)
c_chol = np.linalg.cholesky(c)
wishart_cholesky_to_iw_cholesky = tfb.CholeskyToInvCholesky()
w_fwd = wishart_cholesky_to_iw_cholesky.forward(m).numpy()

print('numpy =\n', c_chol)
print('bijector =\n', w_fwd)
numpy =
 [[ 1.0307764   0.        ]
 [-0.03031695  0.12126781]]
bijector =
 [[ 1.0307764   0.        ]
 [-0.03031695  0.12126781]]

Distribusi terakhir kami

Wishart terbalik kami yang beroperasi pada faktor Cholesky adalah sebagai berikut:

inverse_wishart_cholesky = tfd.TransformedDistribution(
    distribution=CholeskyWishart(
        df=PRIOR_DF,
        scale_tril=np.linalg.cholesky(np.linalg.inv(PRIOR_SCALE))),
    bijector=tfb.CholeskyToInvCholesky())

Kami memiliki Wishart terbalik, tetapi agak lambat karena kami harus melakukan dekomposisi Cholesky di bijector. Mari kembali ke parameterisasi matriks presisi dan lihat apa yang dapat kita lakukan di sana untuk pengoptimalan.

Versi Akhir (!): Menggunakan dekomposisi Cholesky dari matriks presisi

Pendekatan alternatif adalah bekerja dengan faktor Cholesky dari matriks presisi. Di sini fungsi kemungkinan sebelumnya mudah dihitung, tetapi fungsi kemungkinan log data membutuhkan lebih banyak pekerjaan karena TFP tidak memiliki versi multivariate normal yang diparameterisasi dengan presisi.

Kemungkinan log sebelumnya yang dioptimalkan

We use the CholeskyWishart distribution we built above to construct the prior.

# Our new prior.
PRIOR_SCALE_CHOLESKY = np.linalg.cholesky(PRIOR_SCALE)

def log_lik_prior_cholesky(precisions_cholesky):
  rv_precision = CholeskyWishart(
      df=PRIOR_DF,
      scale_tril=PRIOR_SCALE_CHOLESKY,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)
  return rv_precision.log_prob(precisions_cholesky)
# Check against the slower TF implementation and the NumPy implementation.
# Note that when comparing to NumPy, we need to add in the Jacobian correction.
precisions = [np.eye(2, dtype=np.float32),
              true_precision]
precisions_cholesky = np.stack([np.linalg.cholesky(m) for m in precisions])
precisions = np.stack(precisions)
lik_tf = log_lik_prior_cholesky(precisions_cholesky).numpy()
lik_tf_slow = tfd.TransformedDistribution(
    distribution=tfd.WishartTriL(
        df=PRIOR_DF, scale_tril=tf.linalg.cholesky(PRIOR_SCALE)),
    bijector=tfb.Invert(tfb.CholeskyOuterProduct())).log_prob(
    precisions_cholesky).numpy()
corrections = tfb.Invert(tfb.CholeskyOuterProduct()).inverse_log_det_jacobian(
    precisions_cholesky, event_ndims=2).numpy()
n = precisions.shape[0]

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]) + corrections[i])
  print('tensorflow slow:', lik_tf_slow[i])
  print('tensorflow fast:', lik_tf[i])
0
numpy: -0.8488930160357633
tensorflow slow: -0.84889317
tensorflow fast: -0.84889317
1
numpy: -7.442875031036973
tensorflow slow: -7.442877
tensorflow fast: -7.442876

Optimized data log likelihood

We can use TFP's bijectors to build our own version of the multivariate normal. Here is the key idea:

Suppose I have a column vector $X$ whose elements are iid samples of $N(0, 1)$. We have $\text{mean}(X) = 0$ and $\text{cov}(X) = I$

Now let $Y = AX + b$. We have $\text{mean}(Y) = b$ and $\text{cov}(Y) = AA^T$

Hence we can make vectors with mean $b$ and covariance $C$ using the affine transform $Ax+b$ to vectors of iid standard Normal samples provided $AA^T = C$. The Cholesky decomposition of $C$ has the desired property. However, there are other solutions.

Let $P = C^{-1}$ and let the Cholesky decomposition of $P$ be $B$, ie $BB^T = P$. Now

$P^{-1} = (BB^T)^{-1} = B^{-T} B^{-1} = B^{-T} (B^{-T})^T$

So another way to get our desired mean and covariance is to use the affine transform $Y=B^{-T}X + b$.

Our approach (courtesy of this notebook ):

  1. Use tfd.Independent() to combine a batch of 1-D Normal random variables into a single multi-dimensional random variable. The reinterpreted_batch_ndims parameter for Independent() specifies the number of batch dimensions that should be reinterpreted as event dimensions. In our case we create a 1-D batch of length 2 that we transform into a 1-D event of length 2, so reinterpreted_batch_ndims=1 .
  2. Apply a bijector to add the desired covariance: tfb.Invert(tfb.Affine(scale_tril=precision_cholesky, adjoint=True)) . Note that above we're multiplying our iid normal random variables by the transpose of the inverse of the Cholesky factor of the precision matrix $(B^{-T}X)$. The tfb.Invert takes care of inverting $B$, and the adjoint=True flag performs the transpose.
  3. Apply a bijector to add the desired offset: tfb.Affine(shift=shift) Note that we have to do the shift as a separate step from the initial inverted affine transform because otherwise the inverted scale is applied to the shift (since the inverse of $y=Ax+b$ is $x=A^{-1}y - A^{-1}b$).
class MVNPrecisionCholesky(tfd.TransformedDistribution):
  """Multivariate normal parameterized by loc and Cholesky precision matrix."""

  def __init__(self, loc, precision_cholesky, name=None):
    super(MVNPrecisionCholesky, self).__init__(
        distribution=tfd.Independent(
            tfd.Normal(loc=tf.zeros_like(loc),
                       scale=tf.ones_like(loc)),
            reinterpreted_batch_ndims=1),
        bijector=tfb.Chain([
            tfb.Affine(shift=loc),
            tfb.Invert(tfb.Affine(scale_tril=precision_cholesky,
                                  adjoint=True)),
        ]),
        name=name)
@tf.function(autograph=False)
def log_lik_data_cholesky(precisions_cholesky, replicated_data):
  n = tf.shape(precisions_cholesky)[0]  # number of precision matrices
  rv_data = MVNPrecisionCholesky(
      loc=tf.zeros([n, 2]),
      precision_cholesky=precisions_cholesky)
  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# check against the numpy implementation
true_precision_cholesky = np.linalg.cholesky(true_precision)
precisions = [np.eye(2, dtype=np.float32), true_precision]
precisions_cholesky = np.stack([np.linalg.cholesky(m) for m in precisions])
precisions = np.stack(precisions)
n = precisions_cholesky.shape[0]
replicated_data = np.tile(np.expand_dims(my_data, axis=1), reps=[1, 2, 1])
lik_tf = log_lik_data_cholesky(precisions_cholesky, replicated_data).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.81824

Combined log likelihood function

Now we combine our prior and data log likelihood functions in a closure.

def get_log_lik_cholesky(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik_cholesky(precisions_cholesky):
    return (log_lik_data_cholesky(precisions_cholesky, replicated_data) +
            log_lik_prior_cholesky(precisions_cholesky))

  return _log_lik_cholesky

Constraining bijector

Our samples are constrained to be valid Cholesky factors, which means they must be lower triangular matrices with positive diagonals. The TransformedTransitionKernel needs a bijector that maps unconstrained tensors to/from tensors with our desired constraints. We've removed the Cholesky decomposition from the bijector's inverse, which speeds things up.

unconstrained_to_precision_cholesky = tfb.Chain([
    # step 2: exponentiate the diagonals    
    tfb.TransformDiagonal(tfb.Exp(validate_args=VALIDATE_ARGS)),
    # step 1: expand the vector to a lower triangular matrix
    tfb.FillTriangular(validate_args=VALIDATE_ARGS),
])
# some checks
inv = unconstrained_to_precision_cholesky.inverse(precisions_cholesky).numpy()
fwd = unconstrained_to_precision_cholesky.forward(inv).numpy()
print('precisions_cholesky:\n', precisions_cholesky)
print('\ninv:\n', inv)
print('\nfwd(inv):\n', fwd)
precisions_cholesky:
 [[[ 1.         0.       ]
  [ 0.         1.       ]]

 [[ 1.1470785  0.       ]
  [-2.0647411  1.0000004]]]

inv:
 [[ 0.0000000e+00  0.0000000e+00  0.0000000e+00]
 [ 3.5762781e-07 -2.0647411e+00  1.3721828e-01]]

fwd(inv):
 [[[ 1.         0.       ]
  [ 0.         1.       ]]

 [[ 1.1470785  0.       ]
  [-2.0647411  1.0000004]]]

Initial values

We generate a tensor of initial values. We're working with Cholesky factors, so we generate some Cholesky factor initial values.

# The number of chains is determined by the shape of the initial values.
# Here we'll generate 3 chains, so we'll need a tensor of 3 initial values.
N_CHAINS = 3

np.random.seed(123)

initial_values_cholesky = []
for i in range(N_CHAINS):
  initial_values_cholesky.append(np.array(
      [[0.5 + np.random.uniform(), 0.0],
       [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
      dtype=np.float32))
initial_values_cholesky = np.stack(initial_values_cholesky)

Sampling

We sample N_CHAINS chains using the TransformedTransitionKernel .

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_cholesky(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  ttk = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=hmc, bijector=unconstrained_to_precision_cholesky)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=ttk,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values,
      kernel=adapted_kernel,
      trace_fn=None,
      parallel_iterations=1)
  # transform samples back to their constrained form
  samples = tf.linalg.matmul(states, states, transpose_b=True)
  return samples

precision_samples = sample()

Convergence check

A quick convergence check looks good:

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0013583 1.0019467]
 [1.0019467 1.0011804]]

Comparing results to the analytic posterior

# The output samples have shape [n_steps, n_chains, 2, 2]
# Flatten them to [n_steps * n_chains, 2, 2] via reshape:
precision_samples_reshaped = np.reshape(precision_samples, newshape=[-1, 2, 2])

And again, the sample means and standard deviations match those of the analytic posterior.

print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.9668749 -1.6552604]
 [-1.6552604  3.8676758]]

print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13329637 0.24913797]
 [0.24913797 0.53983945]]

Ok, all done! We've got our optimized sampler working.