Étude de cas de probabilité TensorFlow: estimation de la covariance

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le cahier

J'ai écrit ce cahier comme étude de cas pour apprendre la probabilité TensorFlow. Le problème que j'ai choisi de résoudre est l'estimation d'une matrice de covariance pour des échantillons d'une variable aléatoire gaussienne moyenne 2-D 0. Le problème a quelques fonctionnalités intéressantes:

  • Si nous utilisons un a priori de Wishart inverse pour la covariance (une approche courante), le problème a une solution analytique, nous pouvons donc vérifier nos résultats.
  • Le problème consiste à échantillonner un paramètre contraint, ce qui ajoute une complexité intéressante.
  • La solution la plus simple n'est pas la plus rapide, il y a donc un travail d'optimisation à faire.

J'ai décidé d'écrire mes expériences au fur et à mesure. Il m'a fallu un certain temps pour comprendre les subtilités de la TFP, donc ce cahier démarre assez simplement, puis évolue progressivement vers des fonctionnalités TFP plus compliquées. J'ai rencontré de nombreux problèmes en cours de route et j'ai essayé de capturer à la fois les processus qui m'ont aidé à les identifier et les solutions de contournement que j'ai finalement trouvées. J'ai essayé d'inclure beaucoup de détails (y compris de nombreux tests pour m'assurer que les étapes individuelles sont correctes).

Pourquoi apprendre la probabilité TensorFlow?

J'ai trouvé TensorFlow Probability attrayant pour mon projet pour plusieurs raisons:

  • La probabilité TensorFlow vous permet de développer des prototypes de modèles complexes de manière interactive dans un bloc-notes. Vous pouvez diviser votre code en petits morceaux que vous pouvez tester de manière interactive et avec des tests unitaires.
  • Une fois que vous êtes prêt à évoluer, vous pouvez profiter de toute l'infrastructure que nous avons en place pour faire fonctionner TensorFlow sur plusieurs processeurs optimisés sur plusieurs machines.
  • Enfin, même si j'aime vraiment Stan, je trouve que c'est assez difficile à déboguer. Vous devez écrire tout votre code de modélisation dans un langage autonome qui dispose de très peu d'outils pour vous permettre de fouiller votre code, d'inspecter les états intermédiaires, etc.

L'inconvénient est que TensorFlow Probability est beaucoup plus récent que Stan et PyMC3, donc la documentation est un travail en cours, et il y a beaucoup de fonctionnalités qui doivent encore être construites. Heureusement, j'ai trouvé que la base de TFP était solide, et elle est conçue de manière modulaire qui permet d'étendre ses fonctionnalités assez simplement. Dans ce cahier, en plus de résoudre l'étude de cas, je vais montrer quelques façons d'étendre la TFP.

Pour qui c'est

Je suppose que les lecteurs arrivent à ce cahier avec des conditions préalables importantes. Tu devrais:

  • Connaître les bases de l'inférence bayésienne. (Si vous ne le faites pas, un très bon premier livre est Statistical Rethinking )
  • Avoir une certaine familiarité avec une bibliothèque d'échantillonnage MCMC, par exemple Stan / PyMC3 / BUGS
  • Avoir une solide compréhension de NumPy (une bonne introduction est Python pour l'analyse de données )
  • Avoir au moins une familiarité passagère avec TensorFlow , mais pas nécessairement une expertise. (L' apprentissage de TensorFlow est une bonne chose , mais l'évolution rapide de TensorFlow signifie que la plupart des livres seront un peu datés. Le cours CS20 de Stanford est également bon.)

Premier essai

Voici ma première tentative sur le problème. Spoiler: ma solution ne fonctionne pas, et il faudra plusieurs tentatives pour bien faire les choses! Bien que le processus prenne un certain temps, chaque tentative ci-dessous a été utile pour apprendre une nouvelle partie de la PTF.

Une note: TFP n'implémente actuellement pas la distribution inverse de Wishart (nous verrons à la fin comment rouler notre propre Wishart inverse), donc je vais plutôt changer le problème en celui d'estimer une matrice de précision à l'aide d'un Wishart antérieur.

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

Étape 1: rassembler les observations

Mes données ici sont toutes synthétiques, donc cela va sembler un peu plus ordonné qu'un exemple réel. Cependant, il n'y a aucune raison pour que vous ne puissiez pas générer vos propres données synthétiques.

Conseil : une fois que vous avez choisi la forme de votre modèle, vous pouvez choisir des valeurs de paramètres et utiliser le modèle que vous avez choisi pour générer des données synthétiques. Pour vérifier la cohérence de votre implémentation, vous pouvez ensuite vérifier que vos estimations incluent les vraies valeurs des paramètres que vous avez choisis. Pour accélérer votre cycle de débogage / test, vous pouvez envisager une version simplifiée de votre modèle (par exemple, utilisez moins de dimensions ou moins d'échantillons).

Conseil: il est plus simple de travailler avec vos observations sous forme de tableaux NumPy. Une chose importante à noter est que NumPy utilise par défaut float64, tandis que TensorFlow utilise par défaut float32.

En général, les opérations TensorFlow souhaitent que tous les arguments aient le même type et vous devez effectuer un cast de données explicite pour changer de type. Si vous utilisez des observations float64, vous devrez ajouter de nombreuses opérations de conversion. NumPy, en revanche, se chargera de la diffusion automatiquement. Par conséquent, il est beaucoup plus facile de convertir vos données Numpy en float32 que de forcer TensorFlow à utiliser float64.

Choisissez des valeurs de paramètres

# 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]

Générer des observations synthétiques

Notez que TensorFlow Probability utilise la convention selon laquelle la ou les dimension (s) initiale (s) de vos données représentent des exemples d'indices, et la ou les dimensions finales de vos données représentent la dimensionnalité de vos échantillons.

Ici, nous voulons 100 échantillons, dont chacun est un vecteur de longueur 2. Nous allons générer un tableau my_data avec la forme (100, 2). my_data[i, :] est le $ i $ ème échantillon, et c'est un vecteur de longueur 2.

(N'oubliez pas de faire en my_data que my_data ait le type 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)

Vérifiez la santé mentale des observations

Une source potentielle de bogues est de gâcher vos données synthétiques! Faisons quelques vérifications simples.

# 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. ]]

Ok, nos échantillons semblent raisonnables. L'étape suivante.

Étape 2: implémenter la fonction de vraisemblance dans NumPy

La principale chose que nous devrons écrire pour effectuer notre échantillonnage MCMC dans TF Probability est une fonction de vraisemblance logarithmique. En général, il est un peu plus difficile d'écrire TF que NumPy, donc je trouve utile de faire une implémentation initiale dans NumPy. Je vais diviser la fonction de vraisemblance en 2 parties, une fonction de vraisemblance de données qui correspond à $ P (données | paramètres) $ et une fonction de vraisemblance antérieure qui correspond à $ P (paramètres) $.

Notez que ces fonctions NumPy n'ont pas besoin d'être super optimisées / vectorisées puisque le but est simplement de générer des valeurs pour les tests. L'exactitude est la considération clé!

Nous allons d'abord implémenter l'élément de probabilité du journal de données. C'est assez simple. La seule chose à retenir est que nous allons travailler avec des matrices de précision, nous allons donc paramétrer en conséquence.

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

Nous allons utiliser un a priori de Wishart pour la matrice de précision car il existe une solution analytique pour le postérieur (voir le tableau pratique des priors conjugués de Wikipedia ).

La distribution Wishart a 2 paramètres:

  • le nombre de degrés de liberté (étiqueté $ \ nu $ dans Wikipedia)
  • une matrice d'échelle (étiquetée $ V $ dans Wikipedia)

La moyenne pour une distribution Wishart avec les paramètres $ \ nu, V $ est $ E [W] = \ nu V $, et la variance est $ \ text {Var} (W_ {ij}) = \ nu (v_ {ij} ^ 2 + v_ {ii} v_ {jj}) $

Une intuition utile: vous pouvez générer un échantillon Wishart en générant $ \ nu $ tirages indépendants $ x_1 \ ldots x _ {\ nu} $ à partir d'une variable aléatoire normale multivariée de moyenne 0 et de covariance $ V $, puis en formant la somme $ W = \ sum_ {i = 1} ^ {\ nu} x_i x_i ^ T $.

Si vous redimensionnez les échantillons Wishart en les divisant par $ \ nu $, vous obtenez la matrice de covariance de l'échantillon de $ x_i $. Cet exemple de matrice de covariance devrait tendre vers $ V $ lorsque $ \ nu $ augmente. Lorsque $ \ nu $ est petit, il y a beaucoup de variation dans la matrice de covariance de l'échantillon, donc de petites valeurs de $ \ nu $ correspondent à des a priori plus faibles et de grandes valeurs de $ \ nu $ correspondent à des a priori plus forts. Notez que $ \ nu $ doit être au moins aussi grand que la dimension de l'espace que vous échantillonnez ou vous allez générer des matrices singulières.

Nous utiliserons $ \ nu = 3 $ donc nous avons un a priori faible, et nous prendrons $ V = \ frac {1} {\ nu} I $ qui tirera notre estimation de covariance vers l'identité (rappelons que la moyenne est $ \ 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

La distribution de Wishart est l'a priori conjugué pour estimer la matrice de précision d'une normale multivariée de moyenne connue $ \ mu $.

Supposons que les paramètres Wishart antérieurs soient $ \ nu, V $ et que nous ayons $ n $ observations de notre normale multivariée, $ x_1, \ ldots, x_n $. Les paramètres postérieurs sont $ n + \ nu, \ left (V ^ {- 1} + \ sum_ {i = 1} ^ n (x_i- \ mu) (x_i- \ mu) ^ T \ right) ^ {- 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)))

Un rapide tracé des postérieurs et des vraies valeurs. Notez que les postérieurs sont proches des postérieurs de l'échantillon mais sont un peu rétrécis vers l'identité. Notez également que les vraies valeurs sont assez éloignées du mode du postérieur - probablement parce que prior ne correspond pas très bien à nos données. Dans un problème réel, nous ferions probablement mieux avec quelque chose comme un préalable de Wishart inverse mis à l'échelle pour la covariance (voir, par exemple, le commentaire d' Andrew Gelman sur le sujet), mais alors nous n'aurions pas un joli postérieur analytique.

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

Étape 3: implémentation de la fonction de vraisemblance dans TensorFlow

Spoiler: Notre première tentative ne fonctionnera pas; nous parlerons de pourquoi ci-dessous.

Astuce : utilisez le mode impatient TensorFlow lors du développement de vos fonctions de vraisemblance. Le mode Eager fait que TF se comporte plus comme NumPy - tout s'exécute immédiatement, vous pouvez donc déboguer de manière interactive au lieu d'avoir à utiliser Session.run() . Voir les notes ici .

Préliminaire: classes de distribution

TFP a une collection de classes de distribution que nous utiliserons pour générer nos probabilités de journal. Une chose à noter est que ces classes fonctionnent avec des tenseurs d'échantillons plutôt qu'avec des échantillons uniques - cela permet la vectorisation et les accélérations associées.

Une distribution peut fonctionner avec un tenseur d'échantillons de 2 manières différentes. Il est plus simple d'illustrer ces 2 façons avec un exemple concret impliquant une distribution avec un seul paramètre scalaire. J'utiliserai la distribution de Poisson , qui a un paramètre de rate .

  • Si nous créons un Poisson avec une valeur unique pour le paramètre rate , un appel à sa méthode sample() renvoie une seule valeur. Cette valeur est appelée un event , et dans ce cas, les événements sont tous scalaires.
  • Si nous créons un Poisson avec un tenseur de valeurs pour le paramètre rate , un appel à sa méthode sample() renvoie maintenant plusieurs valeurs, une pour chaque valeur du tenseur de taux. L'objet agit comme une collection de Poissons indépendants, chacun avec son propre taux, et chacune des valeurs renvoyées par un appel à sample() correspond à l'un de ces Poissons. Cette collection d'événements indépendants mais non répartis de manière identique est appelée un batch .
  • La méthode sample() prend un paramètre sample_shape qui par défaut est un tuple vide. Si vous transmettez une valeur non vide pour sample_shape échantillon renvoie plusieurs lots. Cette collection de lots est appelée un sample .

La méthode log_prob() une distribution consomme des données d'une manière parallèle à la façon dont sample() génère. log_prob() renvoie des probabilités pour des échantillons, c'est-à-dire pour plusieurs lots d'événements indépendants.

  • Si nous avons notre objet Poisson qui a été créé avec un rate scalaire, chaque lot est un scalaire, et si nous passons dans un tenseur d'échantillons, nous sortirons un tenseur de la même taille de probabilités log.
  • Si nous avons notre objet Poisson qui a été créé avec un tenseur de forme T de valeurs de rate , chaque lot est un tenseur de forme T Si nous passons dans un tenseur d'échantillons de forme D, T, nous sortirons un tenseur de probabilités logarithmiques de forme D, T.

Voici quelques exemples qui illustrent ces cas. Consultez ce bloc-notes pour un didacticiel plus détaillé sur les événements, les lots et les formes.

# 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 ]

Probabilité du journal de données

Nous allons d'abord implémenter la fonction de vraisemblance du journal de données.

VALIDATE_ARGS = True
ALLOW_NAN_STATS = False

Une différence clé avec le cas NumPy est que notre fonction de vraisemblance TensorFlow devra gérer des vecteurs de matrices de précision plutôt que de simples matrices uniques. Des vecteurs de paramètres seront utilisés lorsque nous échantillonnerons à partir de plusieurs chaînes.

Nous allons créer un objet de distribution qui fonctionne avec un lot de matrices de précision (c'est-à-dire une matrice par chaîne).

Lors du calcul des probabilités de journal de nos données, nous aurons besoin que nos données soient répliquées de la même manière que nos paramètres afin qu'il y ait une copie par variable de lot. La forme de nos données répliquées devra être la suivante:

[sample shape, batch shape, event shape]

Dans notre cas, la forme de l'événement est 2 (puisque nous travaillons avec des gaussiens 2D). La forme de l'échantillon est de 100, puisque nous avons 100 échantillons. La forme du lot sera juste le nombre de matrices de précision avec lesquelles nous travaillons. Il est inutile de répliquer les données chaque fois que nous appelons la fonction de vraisemblance, nous allons donc répliquer les données à l'avance et transmettre la version répliquée.

Notez qu'il s'agit d'une implémentation inefficace: MultivariateNormalFullCovariance est cher par rapport à certaines alternatives dont nous parlerons dans la section d'optimisation à la fin.

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)

Conseil: une chose que j'ai trouvée extrêmement utile est l'écriture de petites vérifications de cohérence de mes fonctions TensorFlow. Il est vraiment facile de gâcher la vectorisation dans TF, donc disposer des fonctions NumPy plus simples est un excellent moyen de vérifier la sortie TF. Considérez-les comme de petits tests unitaires.

# 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

Probabilité du journal antérieur

Le préalable est plus simple car nous n'avons pas à nous soucier de la réplication des données.

@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

Construire la fonction de vraisemblance log conjointe

La fonction de vraisemblance du journal de données ci-dessus dépend de nos observations, mais l'échantillonneur ne les aura pas. Nous pouvons nous débarrasser de la dépendance sans utiliser de variable globale en utilisant une [fermeture] (https://en.wikipedia.org/wiki/Closure_ (computer_programming). Les fermetures impliquent une fonction externe qui construit un environnement contenant les variables nécessaires à un fonction intérieure.

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

Étape 4: Échantillon

Ok, il est temps de goûter! Pour simplifier les choses, nous n'utiliserons qu'une chaîne et nous utiliserons la matrice d'identité comme point de départ. Nous ferons les choses plus attentivement plus tard.

Encore une fois, cela ne fonctionnera pas - nous aurons une exception.

@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),
     trace_fn=None,
     seed=123)
  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

Identifier le problème

InvalidArgumentError (see above for traceback): Cholesky decomposition was not successful. The input might not be valid. Ce n'est pas très utile. Voyons si nous pouvons en savoir plus sur ce qui s'est passé.

  • Nous imprimerons les paramètres pour chaque étape afin que nous puissions voir la valeur pour laquelle les choses échouent
  • Nous ajouterons quelques affirmations pour se prémunir contre des problèmes spécifiques.

Les assertions sont délicates car ce sont des opérations TensorFlow, et nous devons veiller à ce qu'elles soient exécutées et ne soient pas optimisées hors du graphique. Cela vaut la peine de lire cet aperçu du débogage TensorFlow si vous n'êtes pas familier avec les assertions TF. Vous pouvez explicitement forcer les assertions à s'exécuter à l'aide de tf.control_dependencies (voir les commentaires dans le code ci-dessous).

La fonction d' Print native de TensorFlow a le même comportement que les assertions - c'est une opération, et vous devez faire attention pour vous assurer qu'elle s'exécute. Print provoque des maux de tête supplémentaires lorsque nous travaillons dans un notebook: sa sortie est envoyée à stderr , et stderr n'est pas affiché dans la cellule. Nous allons utiliser une astuce ici: au lieu d'utiliser tf.Print , nous allons créer notre propre opération d'impression TensorFlow via tf.pyfunc . Comme pour les assertions, nous devons nous assurer que notre méthode s'exécute.

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),
      trace_fn=None,
      seed=123)

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

Pourquoi cela échoue

La toute première nouvelle valeur de paramètre que l'échantillonneur essaie est une matrice asymétrique. Cela provoque l'échec de la décomposition de Cholesky, car elle n'est définie que pour les matrices symétriques (et définies positives).

Le problème ici est que notre paramètre d'intérêt est une matrice de précision, et les matrices de précision doivent être réelles, symétriques et définies positivement. L'échantillonneur ne sait rien de cette contrainte (sauf éventuellement via des gradients), il est donc tout à fait possible que l'échantillonneur propose une valeur invalide, conduisant à une exception, en particulier si la taille du pas est grande.

Avec l'échantillonneur Hamiltonian Monte Carlo, nous pourrons peut-être contourner le problème en utilisant une très petite taille de pas, car le gradient doit éloigner les paramètres des régions invalides, mais de petites tailles de pas signifient une convergence lente. Avec un échantillonneur Metropolis-Hastings, qui ne sait rien des dégradés, nous sommes condamnés.

Version 2: reparamétrisation en paramètres sans contrainte

Il existe une solution simple au problème ci-dessus: nous pouvons reparamétrer notre modèle de telle sorte que les nouveaux paramètres n'aient plus ces contraintes. TFP fournit un ensemble d'outils utiles - des bijecteurs - pour faire exactement cela.

Reparamétrage avec des bijecteurs

Notre matrice de précision doit être réelle et symétrique; nous voulons un paramétrage alternatif qui n'a pas ces contraintes. Un point de départ est une factorisation de Cholesky de la matrice de précision. Les facteurs de Cholesky sont encore contraints - ils sont triangulaires inférieurs et leurs éléments diagonaux doivent être positifs. Cependant, si nous prenons le log des diagonales du facteur Cholesky, les logs ne sont plus contraints d'être positifs, puis si nous aplatissons la partie triangulaire inférieure en un vecteur 1-D, nous n'avons plus la contrainte triangulaire inférieure . Le résultat dans notre cas sera un vecteur de longueur 3 sans contrainte.

(Le manuel Stan contient un grand chapitre sur l'utilisation des transformations pour supprimer divers types de contraintes sur les paramètres.)

Cette reparamétrie a peu d'effet sur notre fonction de vraisemblance du journal de données - il suffit d'inverser notre transformation pour récupérer la matrice de précision - mais l'effet sur le prior est plus compliqué. Nous avons précisé que la probabilité d'une matrice de précision donnée est donnée par la distribution de Wishart; quelle est la probabilité de notre matrice transformée?

Rappelons que si on applique une fonction monotone $ g $ à une variable aléatoire 1-D $ X $, $ Y = g (X) $, la densité pour $ Y $ est donnée par

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

Le dérivé de $ g ^ {- 1} $ terme explique la façon dont $ g $ change les volumes locaux. Pour les variables aléatoires de dimension supérieure, le facteur correctif est la valeur absolue du déterminant du jacobien de $ g ^ {- 1} $ (voir ici ).

Nous devrons ajouter un Jacobien de la transformée inverse dans notre fonction de vraisemblance log a priori. Heureusement, la classe Bijector de TFP peut s'en charger pour nous.

La classe Bijector est utilisée pour représenter les fonctions inversibles et lisses utilisées pour changer les variables dans les fonctions de densité de probabilité. Bijectors ont tous un forward() méthode qui effectue une transformation, une inverse() méthode qui radiers, et forward_log_det_jacobian() et inverse_log_det_jacobian() méthodes qui fournissent les corrections jacobiens dont nous avons besoin lorsque nous reparaterize un pdf.

TFP fournit une collection de bijecteurs utiles que nous pouvons combiner via la composition via l'opérateur Chain pour former des transformations assez compliquées. Dans notre cas, nous allons composer les 3 bijecteurs suivants (les opérations dans la chaîne sont effectuées de droite à gauche):

  1. La première étape de notre transformation consiste à effectuer une factorisation de Cholesky sur la matrice de précision. Il n'y a pas de classe Bijector pour cela; cependant, le bijecteurCholeskyOuterProduct prend le produit de 2 facteurs Cholesky. Nous pouvons utiliser l'inverse de cette opération en utilisant l'opérateur Invert .
  2. L'étape suivante consiste à prendre le log des éléments diagonaux du facteur Cholesky. Nous accomplissons cela via le bijecteur TransformDiagonal et l'inverse du bijecteur Exp .
  3. Enfin, nous aplatissons la partie triangulaire inférieure de la matrice en un vecteur en utilisant l'inverse du 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.]]]

La classe TransformedDistribution automatise le processus d'application d'un bijecteur à une distribution et apporte la correction jacobienne nécessaire à log_prob() . Notre nouveau prieur devient:

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

Nous avons juste besoin d'inverser la transformation pour la probabilité de notre journal de données:

precision = precision_to_unconstrained.inverse(transformed_precision)

Puisque nous voulons en fait la factorisation de Cholesky de la matrice de précision, il serait plus efficace de ne faire ici qu'un inverse partiel. Cependant, nous laisserons l'optimisation pour plus tard et laisserons l'inverse partiel comme exercice pour le lecteur.

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

Encore une fois, nous enveloppons nos nouvelles fonctions dans une fermeture.

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]

Échantillonnage

Maintenant que nous n'avons plus à nous soucier de l'explosion de notre échantillonneur à cause de valeurs de paramètres invalides, générons de vrais échantillons.

L'échantillonneur fonctionne avec la version sans contrainte de nos paramètres, nous devons donc transformer notre valeur initiale en sa version sans contrainte. Les échantillons que nous générons seront également tous dans leur forme sans contrainte, nous devons donc les reconvertir. Les bijecteurs sont vectorisés, il est donc facile de le faire.

# 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),
      trace_fn=lambda _, pkr: pkr.is_accepted,
      seed=123)
  # 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()

Comparons la moyenne de la sortie de notre échantillonneur à la moyenne analytique postérieure!

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 ]]

Nous sommes loin! Voyons pourquoi. Regardons d'abord nos échantillons.

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)

Oh oh - on dirait qu'ils ont tous la même valeur. Voyons pourquoi.

La variable kernel_results_ est un tuple nommé qui donne des informations sur l'échantillonneur à chaque état. Le champ is_accepted est la clé ici.

# 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

Tous nos échantillons ont été rejetés! Vraisemblablement, la taille de notre pas était trop grande. J'ai choisi pas à stepsize=0.1 purement arbitraire.

Version 3: échantillonnage avec une taille de pas adaptative

Étant donné que l'échantillonnage avec mon choix arbitraire de taille d'étape a échoué, nous avons quelques points à l'ordre du jour:

  1. mettre en œuvre une taille de pas adaptative, et
  2. effectuez quelques vérifications de convergence.

Il existe un bon exemple de code dans tensorflow_probability/python/mcmc/hmc.py pour implémenter des tailles de pas adaptatives. Je l'ai adapté ci-dessous.

Notez qu'il existe une sess.run() distincte pour chaque étape. Ceci est vraiment utile pour le débogage, car cela nous permet d'ajouter facilement des diagnostics par étape si nécessaire. Par exemple, nous pouvons afficher la progression incrémentielle, le temps de chaque étape, etc.

Astuce: une façon apparemment courante de gâcher votre échantillonnage est de faire croître votre graphique dans la boucle. (La raison de la finalisation du graphique avant l'exécution de la session est d'éviter de tels problèmes.) Si vous n'avez pas utilisé finalize (), cependant, une vérification de débogage utile si votre code ralentit à une analyse est d'imprimer le graphique size à chaque étape via len(mygraph.get_operations()) - si la longueur augmente, vous faites probablement quelque chose de mal.

Nous allons gérer 3 chaînes indépendantes ici. Faire des comparaisons entre les chaînes nous aidera à vérifier la convergence.

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

Un petit contrôle: notre taux d'acceptation lors de notre prélèvement est proche de notre objectif de 0,651.

print(np.mean(is_accepted))
0.6190666666666667

Mieux encore, la moyenne et l'écart type de notre échantillon sont proches de ce que nous attendons de la solution analytique.

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 ]]

Vérification de la convergence

En général, nous n'aurons pas de solution analytique à vérifier, nous devrons donc nous assurer que l'échantillonneur a convergé. Une vérification standard est la statistique Gelman-Rubin $ \ hat {R} $, qui nécessite plusieurs chaînes d'échantillonnage. $ \ hat {R} $ mesure le degré auquel la variance (des moyennes) entre les chaînes dépasse ce à quoi on pourrait s'attendre si les chaînes étaient distribuées de manière identique. Des valeurs de $ \ hat {R} $ proches de 1 sont utilisées pour indiquer une convergence approximative. Voir la source pour plus de détails.

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

Critique de modèle

Si nous n'avions pas de solution analytique, ce serait le moment de faire une véritable critique de modèle.

Voici quelques histogrammes rapides des composants de l'échantillon par rapport à notre vérité terrain (en rouge). Notez que les échantillons ont été réduits des valeurs de la matrice de précision de l'échantillon vers la matrice d'identité avant.

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

Certains diagrammes de dispersion de paires de composantes de précision montrent qu'en raison de la structure de corrélation du postérieur, les vraies valeurs postérieures ne sont pas aussi improbables qu'elles apparaissent à partir des marginaux ci-dessus.

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

Version 4: échantillonnage plus simple des paramètres contraints

Les bijecteurs ont simplifié l'échantillonnage de la matrice de précision, mais il y avait une bonne quantité de conversion manuelle vers et à partir de la représentation sans contrainte. Il existe un moyen plus simple!

Le noyau de transition transformé

Le TransformedTransitionKernel simplifie ce processus. Il enveloppe votre échantillonneur et gère toutes les conversions. Il prend comme argument une liste de bijecteurs qui mappent des valeurs de paramètres non contraintes à des valeurs contraintes. Nous avons donc ici besoin de l'inverse du bijecteur precision_to_unconstrained que nous avons utilisé ci-dessus. Nous pourrions simplement utiliser tfb.Invert(precision_to_unconstrained) , mais cela impliquerait de prendre des inverses d'inverses (TensorFlow n'est pas assez intelligent pour simplifier tf.Invert(tf.Invert()) en tf.Identity()) , donc à la place nous J'écrirai juste un nouveau bijecteur.

Bijecteur contraignant

# 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.]]

Échantillonnage avec le TransformedTransitionKernel

Avec le TransformedTransitionKernel , nous n'avons plus à faire de transformations manuelles de nos paramètres. Nos valeurs initiales et nos échantillons sont toutes des matrices de précision; il suffit de passer notre (nos) bijecteur (s) sans contrainte au noyau et il s'occupe de toutes les transformations.

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

Vérification de la convergence

Le test de convergence $ \ hat {R} $ semble bon!

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

Comparaison avec le postérieur analytique

Encore une fois, vérifions par rapport au postérieur analytique.

# 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]]

Optimisations

Maintenant que les choses fonctionnent de bout en bout, faisons une version plus optimisée. La vitesse n'a pas trop d'importance pour cet exemple, mais une fois que les matrices deviennent plus grandes, quelques optimisations feront une grande différence.

Une grande amélioration de la vitesse que nous pouvons faire est de reparamétrer en termes de décomposition de Cholesky. La raison en est que notre fonction de vraisemblance des données nécessite à la fois la covariance et les matrices de précision. L'inversion de matrice est coûteuse ($ O (n ^ 3) $ pour une matrice $ n \ times n $), et si nous paramétrons en termes de covariance ou de matrice de précision, nous devons faire une inversion pour obtenir l'autre.

Pour rappel, une matrice symétrique réelle définie positive $ M $ peut être décomposée en un produit de la forme $ M = LL ^ T $ où la matrice $ L $ est triangulaire inférieure et a des diagonales positives. Compte tenu de la décomposition de Cholesky de $ M $, nous pouvons obtenir plus efficacement à la fois $ M $ (le produit d'une matrice triangulaire inférieure et supérieure) et $ M ^ {- 1} $ (par substitution arrière). La factorisation de Cholesky elle-même n'est pas bon marché à calculer, mais si nous paramétrons en termes de facteurs de Cholesky, nous n'avons besoin que de calculer la factorisation de Choleksy des valeurs initiales des paramètres.

Utilisation de la décomposition de Cholesky de la matrice de covariance

TFP a une version de la distribution normale multivariée , MultivariateNormalTriL , qui est paramétrée en termes du facteur de Cholesky de la matrice de covariance. Donc, si nous devions paramétrer en termes du facteur de Cholesky de la matrice de covariance, nous pourrions calculer efficacement la vraisemblance du journal de données. Le défi consiste à calculer la vraisemblance logarithmique antérieure avec une efficacité similaire.

Si nous avions une version de la distribution inverse de Wishart qui fonctionnait avec des facteurs d'échantillons de Cholesky, nous serions tous prêts. Hélas, nous ne le faisons pas. (L'équipe accepterait cependant les soumissions de code!) Comme alternative, nous pouvons utiliser une version de la distribution Wishart qui fonctionne avec des facteurs d'échantillons de Cholesky avec une chaîne de bijecteurs.

Pour le moment, il nous manque quelques bijecteurs de stock pour rendre les choses vraiment efficaces, mais je veux montrer le processus comme un exercice et une illustration utile de la puissance des bijecteurs de TFP.

Une distribution Wishart qui fonctionne sur les facteurs Cholesky

La distribution Wishart a un indicateur utile, input_output_cholesky , qui spécifie que les matrices d'entrée et de sortie doivent être des facteurs de Cholesky. Il est plus efficace et numériquement avantageux de travailler avec les facteurs Cholesky que des matrices complètes, c'est pourquoi cela est souhaitable. Un point important concernant la sémantique du drapeau: c'est seulement une indication que la représentation de l'entrée et de la sortie de la distribution devrait changer - cela n'indique pas une reparamétrisation complète de la distribution, ce qui impliquerait une correction jacobienne du log_prob() une fonction. Nous voulons en fait effectuer cette reparamétrie complète, nous allons donc créer notre propre distribution.

# 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

Construire une distribution de Wishart inverse

Nous avons notre matrice de covariance $ C $ décomposée en $ C = LL ^ T $ où $ L $ est triangulaire inférieur et a une diagonale positive. Nous voulons connaître la probabilité de $ L $ étant donné que $ C \ sim W ^ {- 1} (\ nu, V) $ où $ W ^ {- 1} $ est la distribution inverse de Wishart.

La distribution inverse de Wishart a la propriété que si $ C \ sim W ^ {- 1} (\ nu, V) $, alors la matrice de précision $ C ^ {- 1} \ sim W (\ nu, V ^ {- 1 }) $. Nous pouvons donc obtenir la probabilité de $ L $ via une TransformedDistribution qui prend comme paramètres la distribution de Wishart et un bijecteur qui mappe le facteur de Cholesky de la matrice de précision à un facteur de Cholesky de son inverse.

Un moyen simple (mais pas très efficace) de passer du facteur Cholesky de $ C ^ {- 1} $ à $ L $ est d'inverser le facteur Cholesky en rétrocédant, puis en formant la matrice de covariance à partir de ces facteurs inversés, et puis faire une factorisation de Cholesky.

Soit la décomposition de Cholesky de $ C ^ {- 1} = MM ^ T $. $ M $ est triangulaire inférieur, nous pouvons donc l'inverser en utilisant le MatrixInverseTriL MatrixInverseTriL.

Former $ C $ à partir de $ M ^ {- 1} $ est un peu délicat: $ C = (MM ^ T) ^ {- 1} = M ^ {- T} M ^ {- 1} = M ^ {- T } (M ^ {- T}) ^ T $. $ M $ est triangulaire inférieur, donc $ M ^ {- 1} $ sera également triangulaire inférieur, et $ M ^ {- T} $ sera triangulaire supérieur. Le bijecteur CholeskyOuterProduct() ne fonctionne qu'avec des matrices triangulaires inférieures, nous ne pouvons donc pas l'utiliser pour former $ C $ à partir de $ M ^ {- T} $. Notre solution de contournement est une chaîne de bijecteurs qui permutent les lignes et les colonnes d'une matrice.

Heureusement, cette logique est encapsulée dans le bijecteur CholeskyToInvCholesky !

Combiner toutes les pièces

# 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]]

Notre distribution finale

Notre Wishart inverse fonctionnant sur les facteurs de Cholesky est le suivant:

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

Nous avons notre Wishart inverse, mais c'est un peu lent car nous devons faire une décomposition de Cholesky dans le bijecteur. Revenons au paramétrage de la matrice de précision et voyons ce que nous pouvons y faire pour l'optimisation.

Version finale (!): En utilisant la décomposition de Cholesky de la matrice de précision

Une approche alternative consiste à travailler avec les facteurs de Cholesky de la matrice de précision. Ici, la fonction de vraisemblance a priori est facile à calculer, mais la fonction de vraisemblance du journal de données demande plus de travail puisque TFP n'a pas de version de la normale multivariée qui est paramétrée par précision.

Probabilité de log précédent optimisée

Nous utilisons la distribution CholeskyWishart nous avons construite ci-dessus pour construire le 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

Probabilité optimisée du journal de données

Nous pouvons utiliser les bijecteurs de TFP pour construire notre propre version de la normale multivariée. Voici l'idée clé:

Supposons que j'ai un vecteur colonne $ X $ dont les éléments sont des échantillons iid de $ N (0, 1) $. Nous avons $ \ text {mean} (X) = 0 $ et $ \ text {cov} (X) = I $

Soit maintenant $ Y = AX + b $. Nous avons $ \ text {mean} (Y) = b $ et $ \ text {cov} (Y) = AA ^ T $

Par conséquent, nous pouvons créer des vecteurs avec la moyenne $ b $ et la covariance $ C $ en utilisant la transformation affine $ Ax + b $ en vecteurs du standard iid Échantillons normaux fournis $ AA ^ T = C $. La décomposition de Cholesky de $ C $ a la propriété souhaitée. Cependant, il existe d'autres solutions.

Soit $ P = C ^ {- 1} $ et que la décomposition de Cholesky de $ P $ soit $ B $, c'est-à-dire $ BB ^ T = P $. À présent

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

Donc, une autre façon d'obtenir la moyenne et la covariance désirées est d'utiliser la transformée affine $ Y = B ^ {- T} X + b $.

Notre approche (gracieuseté de ce cahier ):

  1. Utilisez tfd.Independent() pour combiner un lot de variables aléatoires 1-D Normal en une seule variable aléatoire multidimensionnelle. Le paramètre reinterpreted_batch_ndims pour Independent() spécifie le nombre de dimensions de lot à réinterpréter en tant que dimensions d'événement. 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.ScaleMatvecTriL(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.Shift(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.Shift(shift=loc),
            tfb.Invert(tfb.ScaleMatvecTriL(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.