 Visualizza su TensorFlow.org |  Esegui in Google Colab |  Visualizza la fonte su GitHub |  Scarica taccuino |
In questa collaborazione esploreremo il campionamento dal posteriore di un modello di miscela gaussiana bayesiana (BGMM) utilizzando solo le primitive di probabilità TensorFlow.
Modello
Per \(k\in\{1,\ldots, K\}\) componenti della miscela ciascuno di dimensione \(D\), vorremmo modello \(i\in\{1,\ldots,N\}\) campioni utilizzando la seguente miscela modello bayesiano gaussiana IID:
\[\begin{align*} \theta &\sim \text{Dirichlet}(\text{concentration}=\alpha_0)\\ \mu_k &\sim \text{Normal}(\text{loc}=\mu_{0k}, \text{scale}=I_D)\\ T_k &\sim \text{Wishart}(\text{df}=5, \text{scale}=I_D)\\ Z_i &\sim \text{Categorical}(\text{probs}=\theta)\\ Y_i &\sim \text{Normal}(\text{loc}=\mu_{z_i}, \text{scale}=T_{z_i}^{-1/2})\\ \end{align*}\]
Nota, la scale argomenti tutti hanno cholesky semantica. Usiamo questa convenzione perché è quella di TF Distributions (che a sua volta usa questa convenzione in parte perché è computazionalmente vantaggiosa).
Il nostro obiettivo è generare campioni dal posteriore:
\[p\left(\theta, \{\mu_k, T_k\}_{k=1}^K \Big| \{y_i\}_{i=1}^N, \alpha_0, \{\mu_{ok}\}_{k=1}^K\right)\]
Si noti che \(\{Z_i\}_{i=1}^N\) non è presente -: siamo interessava solo in quelle variabili casuali che non scala con \(N\). (E per fortuna c'è una distribuzione di TF che gestisce emarginando fuori \(Z_i\).)
Non è possibile campionare direttamente da questa distribuzione a causa di un termine di normalizzazione computazionalmente intrattabile.
Algoritmi metropoli-Hastings sono tecnica per per campionare da distribuzioni intrattabili a normalizzare.
TensorFlow Probability offre una serie di opzioni MCMC, incluse molte basate su Metropolis-Hastings. In questo notebook, useremo hamiltoniano Monte Carlo ( tfp.mcmc.HamiltonianMonteCarlo ). L'HMC è spesso una buona scelta perché può convergere rapidamente, campionare lo spazio statale congiuntamente (anziché in modo coordinato) e sfrutta una delle virtù di TF: la differenziazione automatica. Detto questo, il campionamento da un posteriore BGMM potrebbe effettivamente essere meglio fatto da altri approcci, ad esempio, il campionamento di Gibb .
%matplotlib inline
import functools
import matplotlib.pyplot as plt; plt.style.use('ggplot')
import numpy as np
import seaborn as sns; sns.set_context('notebook')
import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
Prima di costruire effettivamente il modello, dovremo definire un nuovo tipo di distribuzione. Dalla specifica del modello di cui sopra, è chiaro che stiamo parametrizzando l'MVN con una matrice di covarianza inversa, ovvero [matrice di precisione](https://en.wikipedia.org/wiki/Precision_(statistics%29). Per ottenere ciò in TF, avremo bisogno di implementare la nostra Bijector questo. Bijector utilizzerà il trasformazione in avanti:
-
Y = tf.linalg.triangular_solve((tf.linalg.matrix_transpose(chol_precision_tril), X, adjoint=True) + loc.
E la log_prob calcolo è proprio l'inverso, vale a dire:
-
X = tf.linalg.matmul(chol_precision_tril, X - loc, adjoint_a=True).
Dal momento che tutti abbiamo bisogno di HMC è log_prob , questo significa che evitano sempre chiamare tf.linalg.triangular_solve (come sarebbe il caso per tfd.MultivariateNormalTriL ). Questo è vantaggioso poiché tf.linalg.matmul è solitamente più veloce dovuta al meglio frazione cache.
class MVNCholPrecisionTriL(tfd.TransformedDistribution):
"""MVN from loc and (Cholesky) precision matrix."""
def __init__(self, loc, chol_precision_tril, name=None):
super(MVNCholPrecisionTriL, self).__init__(
distribution=tfd.Independent(tfd.Normal(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=chol_precision_tril,
adjoint=True)),
]),
name=name)
Il tfd.Independent spire distribuzione indipendente attira di una distribuzione, in una distribuzione multivariata con coordinate statisticamente indipendenti. In termini di calcolo log_prob , questo "meta-distribuzione" si manifesta come una semplice somma sulla dimensione evento (s).
Si noti inoltre che abbiamo preso adjoint ( "trasposizione") della matrice scala. Questo perché se la precisione è covarianza inversa, cioè, \(P=C^{-1}\) e se \(C=AA^\top\), allora \(P=BB^{\top}\) dove \(B=A^{-\top}\).
Dal momento che questa distribuzione è una specie di difficile, cerchiamo di verificare rapidamente che il nostro MVNCholPrecisionTriL funziona come pensiamo che dovrebbe.
def compute_sample_stats(d, seed=42, n=int(1e6)):
x = d.sample(n, seed=seed)
sample_mean = tf.reduce_mean(x, axis=0, keepdims=True)
s = x - sample_mean
sample_cov = tf.linalg.matmul(s, s, adjoint_a=True) / tf.cast(n, s.dtype)
sample_scale = tf.linalg.cholesky(sample_cov)
sample_mean = sample_mean[0]
return [
sample_mean,
sample_cov,
sample_scale,
]
dtype = np.float32
true_loc = np.array([1., -1.], dtype=dtype)
true_chol_precision = np.array([[1., 0.],
[2., 8.]],
dtype=dtype)
true_precision = np.matmul(true_chol_precision, true_chol_precision.T)
true_cov = np.linalg.inv(true_precision)
d = MVNCholPrecisionTriL(
loc=true_loc,
chol_precision_tril=true_chol_precision)
[sample_mean, sample_cov, sample_scale] = [
t.numpy() for t in compute_sample_stats(d)]
print('true mean:', true_loc)
print('sample mean:', sample_mean)
print('true cov:\n', true_cov)
print('sample cov:\n', sample_cov)
true mean: [ 1. -1.] sample mean: [ 1.0002806 -1.000105 ] true cov: [[ 1.0625 -0.03125 ] [-0.03125 0.015625]] sample cov: [[ 1.0641273 -0.03126175] [-0.03126175 0.01559312]]
Poiché la media e la covarianza del campione sono vicine alla media e alla covarianza vere, sembra che la distribuzione sia implementata correttamente. Ora, useremo MVNCholPrecisionTriL tfp.distributions.JointDistributionNamed specificare il modello BGMM. Per il modello di osservazione, useremo tfd.MixtureSameFamily di integrare automaticamente il \(\{Z_i\}_{i=1}^N\) disegna.
dtype = np.float64
dims = 2
components = 3
num_samples = 1000
bgmm = tfd.JointDistributionNamed(dict(
mix_probs=tfd.Dirichlet(
concentration=np.ones(components, dtype) / 10.),
loc=tfd.Independent(
tfd.Normal(
loc=np.stack([
-np.ones(dims, dtype),
np.zeros(dims, dtype),
np.ones(dims, dtype),
]),
scale=tf.ones([components, dims], dtype)),
reinterpreted_batch_ndims=2),
precision=tfd.Independent(
tfd.WishartTriL(
df=5,
scale_tril=np.stack([np.eye(dims, dtype=dtype)]*components),
input_output_cholesky=True),
reinterpreted_batch_ndims=1),
s=lambda mix_probs, loc, precision: tfd.Sample(tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=mix_probs),
components_distribution=MVNCholPrecisionTriL(
loc=loc,
chol_precision_tril=precision)),
sample_shape=num_samples)
))
def joint_log_prob(observations, mix_probs, loc, chol_precision):
"""BGMM with priors: loc=Normal, precision=Inverse-Wishart, mix=Dirichlet.
Args:
observations: `[n, d]`-shaped `Tensor` representing Bayesian Gaussian
Mixture model draws. Each sample is a length-`d` vector.
mix_probs: `[K]`-shaped `Tensor` representing random draw from
`Dirichlet` prior.
loc: `[K, d]`-shaped `Tensor` representing the location parameter of the
`K` components.
chol_precision: `[K, d, d]`-shaped `Tensor` representing `K` lower
triangular `cholesky(Precision)` matrices, each being sampled from
a Wishart distribution.
Returns:
log_prob: `Tensor` representing joint log-density over all inputs.
"""
return bgmm.log_prob(
mix_probs=mix_probs, loc=loc, precision=chol_precision, s=observations)
Genera dati di "allenamento".
Per questa demo, analizzeremo alcuni dati casuali.
true_loc = np.array([[-2., -2],
[0, 0],
[2, 2]], dtype)
random = np.random.RandomState(seed=43)
true_hidden_component = random.randint(0, components, num_samples)
observations = (true_loc[true_hidden_component] +
random.randn(num_samples, dims).astype(dtype))
Inferenza bayesiana utilizzando HMC
Ora che abbiamo utilizzato TFD per specificare il nostro modello e ottenuto alcuni dati osservati, abbiamo tutti i pezzi necessari per eseguire HMC.
Per fare questo, useremo un'applicazione parziale a "pin down" le cose che non vogliamo campione. In questo caso questo significa che dobbiamo solo PIN giù observations . (Iper-parametri sono già cotti in alle distribuzioni a priori e non parte del joint_log_prob firma funzione.)
unnormalized_posterior_log_prob = functools.partial(joint_log_prob, observations)
initial_state = [
tf.fill([components],
value=np.array(1. / components, dtype),
name='mix_probs'),
tf.constant(np.array([[-2., -2],
[0, 0],
[2, 2]], dtype),
name='loc'),
tf.linalg.eye(dims, batch_shape=[components], dtype=dtype, name='chol_precision'),
]
Rappresentazione illimitata
Hamiltonian Monte Carlo (HMC) richiede che la funzione di probabilità logaritmica target sia differenziabile rispetto ai suoi argomenti. Inoltre, HMC può mostrare un'efficienza statistica notevolmente superiore se lo spazio degli stati non è vincolato.
Ciò significa che dovremo risolvere due problemi principali durante il campionamento dal BGMM posteriore:
- \(\theta\) rappresenta un vettore probabilità discreta, cioè, deve essere tale che \(\sum_{k=1}^K \theta_k = 1\) e \(\theta_k>0\).
- \(T_k\) rappresenta una matrice di covarianza inversa, cioè, deve essere tale che \(T_k \succ 0\), cioè, è definita positiva .
Per soddisfare questo requisito dovremo:
- trasformare le variabili vincolate in uno spazio non vincolato
- eseguire l'MCMC in uno spazio illimitato
- trasformare le variabili non vincolate nello spazio vincolato.
Come con MVNCholPrecisionTriL , useremo Bijector s di trasformare variabili casuali allo spazio senza vincoli.
Il
Dirichletsi trasforma in uno spazio senza vincoli tramite la funzione di SoftMax .La nostra variabile casuale di precisione è una distribuzione su matrici semidefinite positive. Per Sblocca questi useremo le
FillTriangulareTransformDiagonalbijectors. Questi convertono i vettori in matrici triangolari inferiori e assicurano che la diagonale sia positiva. Il primo è utile perché consente il campionamento solo \(d(d+1)/2\) galleggia anziché \(d^2\).
unconstraining_bijectors = [
tfb.SoftmaxCentered(),
tfb.Identity(),
tfb.Chain([
tfb.TransformDiagonal(tfb.Softplus()),
tfb.FillTriangular(),
])]
@tf.function(autograph=False)
def sample():
return tfp.mcmc.sample_chain(
num_results=2000,
num_burnin_steps=500,
current_state=initial_state,
kernel=tfp.mcmc.SimpleStepSizeAdaptation(
tfp.mcmc.TransformedTransitionKernel(
inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=unnormalized_posterior_log_prob,
step_size=0.065,
num_leapfrog_steps=5),
bijector=unconstraining_bijectors),
num_adaptation_steps=400),
trace_fn=lambda _, pkr: pkr.inner_results.inner_results.is_accepted)
[mix_probs, loc, chol_precision], is_accepted = sample()
Ora eseguiremo la catena e stamperemo i mezzi posteriori.
acceptance_rate = tf.reduce_mean(tf.cast(is_accepted, dtype=tf.float32)).numpy()
mean_mix_probs = tf.reduce_mean(mix_probs, axis=0).numpy()
mean_loc = tf.reduce_mean(loc, axis=0).numpy()
mean_chol_precision = tf.reduce_mean(chol_precision, axis=0).numpy()
precision = tf.linalg.matmul(chol_precision, chol_precision, transpose_b=True)
print('acceptance_rate:', acceptance_rate)
print('avg mix probs:', mean_mix_probs)
print('avg loc:\n', mean_loc)
print('avg chol(precision):\n', mean_chol_precision)
acceptance_rate: 0.5305 avg mix probs: [0.25248723 0.60729516 0.1402176 ] avg loc: [[-1.96466753 -2.12047249] [ 0.27628865 0.22944732] [ 2.06461244 2.54216122]] avg chol(precision): [[[ 1.05105032 0. ] [ 0.12699955 1.06553113]] [[ 0.76058015 0. ] [-0.50332767 0.77947431]] [[ 1.22770457 0. ] [ 0.70670027 1.50914164]]]
loc_ = loc.numpy()
ax = sns.kdeplot(loc_[:,0,0], loc_[:,0,1], shade=True, shade_lowest=False)
ax = sns.kdeplot(loc_[:,1,0], loc_[:,1,1], shade=True, shade_lowest=False)
ax = sns.kdeplot(loc_[:,2,0], loc_[:,2,1], shade=True, shade_lowest=False)
plt.title('KDE of loc draws');
Conclusione
Questa semplice collaborazione ha dimostrato come le primitive TensorFlow Probability possono essere utilizzate per costruire modelli di miscele bayesiane gerarchiche.