Bayesian_Gaussian_Mixture_Model.ipynb

Kernel: Python 3

Copyright 2018 The TensorFlow Probability Authors.

Licensed under the Apache License, Version 2.0 (the "License");

#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Modelo de mistura gaussiana bayesiana e MCMC Hamiltoniano

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Neste Colab, veremos a amostragem da distribuição posterior de um modelo de mistura gaussiana bayesiana (BGMM, na sigla em inglês) usando somente os primitivos do TensorFlow Probability.

Modelo

Para $k\in{1,\ldots, K}$ componentes de mistura, cada um com dimensão $D$, gostaríamos de modelar $i\in{1,\ldots,N}$ iid amostras usando o seguinte modelo de mistura gaussiana bayesiana:

$$\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*}$$

Observação: todos os argumentos scale têm semântica cholesky. Usamos essa convenção porque é a mesma do TF Distributions (que usa essa convenção em parte porque tem vantagens computacionais).

Nosso objetivo é gerar amostras da distribuição posterior:

$$p\left(\theta, \{\mu_k, T_k\}_{k=1}^K \Big| \{y_i\}_{i=1}^N, \alpha_0, \{\mu_{ok}\}_{k=1}^K\right)$$

Observe que ${Z_i}_{i=1}^N$ não está presente, pois temos interesse somente nas variáveis aleatórias que não têm a escala ajustada com $N$ (e felizmente existe uma distribuição do TF que trata a marginalização de $Z_i$.)

Não é possível fazer a amostragem diretamente dessa distribuição devido a um termo de normalização intratável computacionalmente.

Os algoritmos de Metropolis-Hastings são técnicas para amostrar distribuições com normalização intratável.

O TensorFlow Probability oferece diversas opções de MCMC, incluindo várias baseadas nos algoritmos de Metropolis-Hastings. Neste notebook, usaremos o Monte Carlo Hamiltoniano (tfp.mcmc.HamiltonianMonteCarlo). O HMC costuma ser uma boa escolha porque pode convergir rapidamente, faz a amostragem do espaço de estados conjuntamente (ao contrário de por coordenadas) e aproveita uma das virtudes do TF: a diferenciação automática. Dito isso, a amostragem de uma distribuição posterior BGMM pode ser feita de uma forma melhor por outras técnicas, como, por exemplo, a amostragem de Gibbs.

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

Antes de criar o modelo, precisamos definir um novo tipo de distribuição. Pela especificação do modelo acima, está claro que estamos parametrizando a MVN com uma matriz de covariância inversa, isto é, [matriz de precisão](https://en.wikipedia.org/wiki/Precision_(statistics)). Para fazer isso no TF, precisamos utilizar o Bijector. Esse Bijector usará a transformação forward:

• Y = tf.linalg.triangular_solve((tf.linalg.matrix_transpose(chol_precision_tril), X, adjoint=True) + loc.

E o cálculo de log_prob é simplesmente o inverso:

• X = tf.linalg.matmul(chol_precision_tril, X - loc, adjoint_a=True).

Como só o que precisamos para o HMC é log_prob, podemos evitar chamar tf.linalg.triangular_solve (que seria o caso para tfd.MultivariateNormalTriL). Isso traz vantagens, pois tf.linalg.matmul costuma ser mais rápido devido a uma melhor localidade do 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)

A distribuição tfd.Independent transforma obtenções independentes de uma distribuição em uma distribuição multivariada com coordenadas estatisticamente independentes. Quanto à computação de log_prob, essa "metadistribuição" se manifesta como uma simples soma das dimensões do evento.

Observe também que pegamos o adjoint ("transposição") da matriz de escala, pois, se a precisão tiver covariância inversa, ou seja, $P=C^{-1}$ e se $C=AA^\top$, então $P=BB^{\top}$, em que $B=A^{-\top}$.

Como essa distribuição é complicada, vamos verificar rapidamente se o MVNCholPrecisionTriL funciona como achamos que deveria funcionar.

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

Como a média e a covariância da amostra estão próximas da média e covariância reais, parece que a distribuição foi implementada corretamente. Agora, usaremos MVNCholPrecisionTriL tfp.distributions.JointDistributionNamed para especificar o modelo de BGMM. Para o modelo observacional, usaremos tfd.MixtureSameFamily para fazer a integral das obtenções $\{Z_i\}_{i=1}^N$ automaticamente.

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)

Gere dados de "treinamento"

Para esta demonstração, faremos a amostragem de dados aleatórios.

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

Inerência bayesiana usando HMC

Agora que usamos o TFD para especificar nosso modelo e obtivemos alguns dados observados, temos todos os componentes necessários para executar o HMC.

Usaremos uma aplicação parcial para "imobilizar" o que não queremos amostrar. Neste caso, precisamos imobilizar somente observations (os hiperparâmetros já estão incorporados às distribuições a priori e não fazem parte da assinatura da função joint_log_prob).

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

Representação não restrita

O Monte Carlo Hamiltoniano (HMC, na sigla em inglês) requer que a função de probabilidade logarítmica alvo seja diferenciável quanto aos seus argumentos. Além disso, o HMC pode ter uma eficiência estatística drasticamente superior se o espaço de estados for não restrito.

Portamos, teremos que tratar dois problemas principais ao fazer a amostragem da distribuição BGMM posterior:

1. $\theta$ representa um vetor de probabilidades discretas, isto é, precisa ser de tal forma que $\sum_{k=1}^K \theta_k = 1$ e $\theta_k>0$.

2. $T_k$ representa uma matriz de covariância inversa, isto é, precisa ser de tal forma que $T_k \succ 0$, ou seja, é definida positiva.

Para tratar esse requisito, precisaremos:

1. transformar as variáveis restritas em um espaço não restrito

2. executar o MCMC no espaço não restrito

3. transformar as variáveis não restritas de volta no espaço restrito

Da mesma forma que em MVNCholPrecisionTriL, usaremos Bijectors para transformar variáveis aleatórias no espaço não restrito.

• Dirichlet é transformado no espaço não restrito por meio da função softmax.

• Nossa variável aleatória de precisão é uma distribuição de matrizes semidefinidas positivas. Para remover a restrição delas, usaremos os bijetores FillTriangular e TransformDiagonal, que convertem os vetores em matrizes triangulares inferiores e garantem que a diagonal seja positiva. O primeiro é útil porque permite a amostragem somente de $d(d+1)/2$ pontos flutuantes, em vez de $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()

Agora vamos executar a cadeia e exibir as médias posteriores via print.

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

Conclusão

Este Colab simples demonstrou como os primitivos do TensorFlow Probability podem ser usados para criar modelos de mistura bayesiana hierárquicos.