Mô hình hỗn hợp Bayesian Gaussian và MCMC Hamilton

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

Trong chuyên mục này, chúng ta sẽ khám phá việc lấy mẫu từ phần sau của Mô hình hỗn hợp Bayesian Gaussian (BGMM) chỉ sử dụng các giá trị ban đầu của Xác suất TensorFlow.

Mô hình

Đối với \(k\in\{1,\ldots, K\}\) thành phần hỗn hợp mỗi chiều \(D\), chúng tôi muốn mô hình \(i\in\{1,\ldots,N\}\) mẫu bằng cách sử dụng sau đây Bayesian Gaussian Mixture Mẫu 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*}\]

Lưu ý, các scale đối số đều có cholesky ngữ nghĩa. Chúng tôi sử dụng quy ước này vì nó là của Phân phối TF (bản thân quy ước này sử dụng quy ước này một phần vì nó có lợi về mặt tính toán).

Mục tiêu của chúng tôi là tạo ra các mẫu từ phía sau:

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

Chú ý rằng \(\{Z_i\}_{i=1}^N\) là không có mặt - -chúng tôi quan tâm đến việc chỉ những biến ngẫu nhiên mà không mở rộng quy mô với \(N\). (Và may mắn là có một phân phối TF mà xử lý marginalizing ra \(Z_i\).)

Không thể lấy mẫu trực tiếp từ phân phối này do thuật ngữ chuẩn hóa khó tính toán được.

Thuật toán Metropolis-Hastings là kỹ thuật để cho lấy mẫu từ các bản phân phối khó-to-bình thường hóa.

TensorFlow Probability cung cấp một số tùy chọn MCMC, bao gồm một số tùy chọn dựa trên Metropolis-Hastings. Trong máy tính xách tay này, chúng tôi sẽ sử dụng Hamiltonian Monte Carlo ( tfp.mcmc.HamiltonianMonteCarlo ). HMC thường là một lựa chọn tốt vì nó có thể hội tụ nhanh chóng, lấy mẫu không gian trạng thái cùng nhau (trái ngược với theo chiều tọa độ), và tận dụng một trong những ưu điểm của TF: phân biệt tự động. Điều đó nói rằng, lấy mẫu từ một sau BGMM thực sự có thể được thực hiện tốt hơn bằng cách tiếp cận khác, ví dụ như, lấy mẫu Gibb của .

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

Trước khi thực sự xây dựng mô hình, chúng ta sẽ cần xác định một kiểu phân phối mới. Từ đặc tả mô hình ở trên, rõ ràng là chúng tôi đang tham số hóa MVN bằng ma trận hiệp phương sai ngược, tức là [ma trận chính xác] (https://en.wikipedia.org/wiki/Pre precision_ (Statistics% 29). Để thực hiện điều này trong TF, chúng tôi sẽ cần phải tung ra của chúng tôi Bijector này. Bijector sẽ sử dụng sự chuyển đổi về phía trước:

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

Và log_prob tính chỉ là sự nghịch đảo, ví dụ:

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

Vì tất cả các nhu cầu chúng tôi cho HMC là log_prob , phương tiện này, chúng ta tránh bao giờ gọi tf.linalg.triangular_solve (như sẽ là trường hợp cho tfd.MultivariateNormalTriL ). Đây là thuận lợi kể từ tf.linalg.matmul thường nhanh hơn do địa phương Cache tốt hơn.

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)

Các tfd.Independent lượt phân phối độc lập thu hút của một phân phối, vào một phân phối đa biến với tọa độ độc lập về mặt thống kê. Xét về tính log_prob , điều này "meta-phân phối" biểu hiện như một khoản tiền đơn giản trên khía cạnh sự kiện (s).

Cũng cần lưu ý rằng chúng tôi đã adjoint ( "chuyển vị") của ma trận quy mô. Điều này là do nếu chính xác là hiệp phương sai nghịch, tức là \(P=C^{-1}\) và nếu \(C=AA^\top\), sau đó \(P=BB^{\top}\) nơi \(B=A^{-\top}\).

Kể từ khi phân phối này là loại khó khăn, chúng ta hãy nhanh chóng xác minh rằng chúng tôi MVNCholPrecisionTriL hoạt động như chúng tôi nghĩ rằng nó nên.

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

Vì trung bình và hiệp phương sai của mẫu gần với giá trị trung bình và hiệp phương sai thực sự, nên có vẻ như phân phối được thực hiện chính xác. Bây giờ, chúng ta sẽ sử dụng MVNCholPrecisionTriL tfp.distributions.JointDistributionNamed để xác định mô hình BGMM. Đối với mô hình quan sát, chúng tôi sẽ sử dụng tfd.MixtureSameFamily để tự động tích hợp ra \(\{Z_i\}_{i=1}^N\) rút.

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)

Tạo dữ liệu "đào tạo"

Đối với bản demo này, chúng tôi sẽ lấy mẫu một số dữ liệu ngẫu nhiên.

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

Suy luận Bayes bằng HMC

Bây giờ chúng tôi đã sử dụng TFD để chỉ định mô hình của mình và thu được một số dữ liệu quan sát được, chúng tôi có tất cả các phần cần thiết để chạy HMC.

Để làm điều này, chúng tôi sẽ sử dụng một ứng dụng phần để "pin xuống" những điều chúng ta không muốn mẫu. Trong trường hợp này có nghĩa là chúng ta chỉ cần ghim xuống observations . (Các siêu thông số đã được nướng vào các bản phân phối trước và không nằm trong joint_log_prob chức năng chữ ký.)

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

Đại diện không bị giới hạn

Hamiltonian Monte Carlo (HMC) yêu cầu hàm xác suất log mục tiêu phải có thể phân biệt được đối với các đối số của nó. Hơn nữa, HMC có thể thể hiện hiệu quả thống kê cao hơn đáng kể nếu không gian trạng thái không bị giới hạn.

Điều này có nghĩa là chúng tôi sẽ phải giải quyết hai vấn đề chính khi lấy mẫu từ phần sau BGMM:

  1. \(\theta\) đại diện cho một vector khả năng rời rạc, tức là phải được như vậy mà \(\sum_{k=1}^K \theta_k = 1\) và \(\theta_k>0\).
  2. \(T_k\) đại diện cho một nghịch đảo ma trận hiệp phương sai, tức là phải được như vậy mà \(T_k \succ 0\), tức là tích cực nhất định .

Để giải quyết yêu cầu này, chúng tôi cần:

  1. chuyển đổi các biến bị ràng buộc thành một không gian không bị giới hạn
  2. chạy MCMC trong không gian không bị giới hạn
  3. biến đổi các biến không bị giới hạn trở lại không gian bị ràng buộc.

Như với MVNCholPrecisionTriL , chúng tôi sẽ sử dụng Bijector s chuyển biến ngẫu nhiên để không gian không bị giới hạn.

  • Các Dirichlet được chuyển thành không gian không bị giới hạn thông qua hàm softmax .

  • Biến ngẫu nhiên chính xác của chúng tôi là một phân phối trên ma trận bán xác định vị trí. Để huỷ cố định những chúng ta sẽ sử dụng FillTriangular và TransformDiagonal bijectors. Chúng chuyển đổi vectơ thành ma trận tam giác thấp hơn và đảm bảo đường chéo là dương. Điều thứ nhất là hữu ích vì nó cho phép lấy mẫu chỉ \(d(d+1)/2\) nổi hơn \(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()

Bây giờ chúng ta sẽ thực hiện chuỗi và in các phương tiện sau.

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

png

Sự kết luận

Chuyên mục đơn giản này đã trình bày cách sử dụng các giá trị nguyên thủy của Xác suất TensorFlow để xây dựng các mô hình hỗn hợp Bayes có thứ bậc.