def call(self, inputs):
if self.conditional_inputs is None and self.conditional_outputs is None:
covariance_matrix = self.covariance_fn(inputs, inputs)
# Tile locations so output has shape [units, batch_size]. Covariance will
# broadcast to [units, batch_size, batch_size], and we perform
# shape manipulations to get a random variable over [batch_size, units].
loc = self.mean_fn(inputs)
loc = tf.tile(loc[tf.newaxis], [self.units] + [1] * len(loc.shape))
else:
knn = self.covariance_fn(inputs, inputs)
knm = self.covariance_fn(inputs, self.conditional_inputs)
kmm = self.covariance_fn(self.conditional_inputs, self.conditional_inputs)
kmm = tf.matrix_set_diag(
kmm, tf.matrix_diag_part(kmm) + tf.keras.backend.epsilon())
kmm_tril = tf.linalg.cholesky(kmm)
kmm_tril_operator = tf.linalg.LinearOperatorLowerTriangular(kmm_tril)
knm_operator = tf.linalg.LinearOperatorFullMatrix(knm)
# TODO(trandustin): Vectorize linear algebra for multiple outputs. For
# now, we do each separately and stack to obtain a locations Tensor of
# shape [units, batch_size].
loc = []
for conditional_outputs_unit in tf.unstack(self.conditional_outputs,
axis=-1):
center = conditional_outputs_unit - self.mean_fn(
self.conditional_inputs)
loc_unit = knm_operator.matvec(
kmm_tril_operator.solvevec(kmm_tril_operator.solvevec(center),
adjoint=True))
loc.append(loc_unit)
loc = tf.stack(loc) + self.mean_fn(inputs)[tf.newaxis]
covariance_matrix = knn
covariance_matrix -= knm_operator.matmul(
kmm_tril_operator.solve(
kmm_tril_operator.solve(knm, adjoint_arg=True), adjoint=True))
covariance_matrix = tf.matrix_set_diag(
covariance_matrix,
tf.matrix_diag_part(covariance_matrix) + tf.keras.backend.epsilon())
# Form a multivariate normal random variable with batch_shape units and
# event_shape batch_size. Then make it be independent across the units
# dimension. Then transpose its dimensions so it is [batch_size, units].
random_variable = ed.MultivariateNormalFullCovariance(
loc=loc, covariance_matrix=covariance_matrix)
random_variable = ed.Independent(random_variable.distribution,
reinterpreted_batch_ndims=1)
bijector = tfp.bijectors.Inline(
forward_fn=lambda x: tf.transpose(x, [1, 0]),
inverse_fn=lambda y: tf.transpose(y, [1, 0]),
forward_event_shape_fn=lambda input_shape: input_shape[::-1],
forward_event_shape_tensor_fn=lambda input_shape: input_shape[::-1],
inverse_log_det_jacobian_fn=lambda y: tf.cast(0, y.dtype),
forward_min_event_ndims=2)
random_variable = ed.TransformedDistribution(random_variable.distribution,
bijector=bijector)
return random_variable
def variational_sgpr(X,
Z,
ls=1.,
kernel_func=rbf,
ridge_factor=1e-3,
mfvi_mixture=False,
n_mixture=1,
name="",
**kwargs):
"""Defines the mean-field variational family for GPR.
Args:
X: (np.ndarray of float32) input training features, with dimension (Nx, D).
Z: (np.ndarray of float32) inducing points, with dimension (Nz, D).
ls: (float32) length scale parameter.
kernel_func: (function) kernel function.
ridge_factor: (float32) small ridge factor to stabilize Cholesky decomposition
mfvi_mixture: (float32) Whether to output variational family with a
mixture of MFVI.
n_mixture: (int) Number of MFVI mixture component to add.
name: (str) name for the variational parameter/random variables.
kwargs: Dict of other keyword variables.
For compatibility purpose with other variational family.
Returns:
q_f, q_sig: (ed.RandomVariable) variational family.
q_f_mean, q_f_sdev: (tf.Variable) variational parameters for q_f
"""
X = tf.convert_to_tensor(X, dtype=tf.float32)
Z = tf.convert_to_tensor(Z, dtype=tf.float32)
Nx, Nz = X.shape.as_list()[0], Z.shape.as_list()[0]
# 1. Prepare constants
# compute matrix constants
Kxx = kernel_func(X, ls=ls)
Kxz = kernel_func(X, Z, ls=ls)
Kzz = kernel_func(Z, ls=ls, ridge_factor=ridge_factor)
# compute null covariance matrix using Cholesky decomposition
Kzz_chol_inv = tf.matrix_inverse(tf.cholesky(Kzz))
Kzz_inv = tf.matmul(Kzz_chol_inv, Kzz_chol_inv, transpose_a=True)
Kxz_Kzz_chol_inv = tf.matmul(Kxz, Kzz_chol_inv, transpose_b=True)
Kxz_Kzz_inv = tf.matmul(Kxz, Kzz_inv)
Sigma_pre = Kxx - tf.matmul(
Kxz_Kzz_chol_inv, Kxz_Kzz_chol_inv, transpose_b=True)
# 2. Define variational parameters
# define free parameters (i.e. mean and full covariance of f_latent)
m = tf.get_variable(shape=[Nz], name='{}_mean_latent'.format(name))
s = tf.get_variable(shape=[Nz * (Nz + 1) / 2],
name='{}_cov_latent_s'.format(name))
L = fill_triangular(s, name='{}_cov_latent_chol'.format(name))
S = tf.matmul(L, L, transpose_b=True, name='{}_cov_latent'.format(name))
# compute sparse gp variational parameter
# (i.e. mean and covariance of P(f_obs | f_latent))
qf_mean = tf.tensordot(Kxz_Kzz_inv,
m, [[1], [0]],
name='{}_mean'.format(name))
qf_cov = (
Sigma_pre +
tf.matmul(Kxz_Kzz_inv, tf.matmul(S, Kxz_Kzz_inv, transpose_b=True)) +
ridge_factor * tf.eye(Nx, dtype=tf.float32))
# define variational family
mixture_par_list = []
if mfvi_mixture:
gp_dist = tfd.MultivariateNormalFullCovariance(
loc=qf_mean, covariance_matrix=qf_cov)
q_f, mixture_par_list = inference_util.make_mfvi_sgp_mixture_family(
n_mixture=n_mixture, N=Nx, gp_dist=gp_dist, name=name)
else:
q_f = ed.MultivariateNormalFullCovariance(loc=qf_mean,
covariance_matrix=qf_cov,
name=name)
return q_f, qf_mean, qf_cov, mixture_par_list
def call(self, inputs):
self.call_weights()
if (not isinstance(inputs, ed.RandomVariable) and
not isinstance(self.kernel, ed.RandomVariable) and
not isinstance(self.bias, ed.RandomVariable)):
return super(DenseDVI, self).call(inputs)
inputs_mean, inputs_variance, inputs_covariance = get_moments(inputs)
kernel_mean, kernel_variance, _ = get_moments(self.kernel)
if self.use_bias:
bias_mean, _, bias_covariance = get_moments(self.bias)
# E[outputs] = E[inputs] * E[kernel] + E[bias]
mean = tf.tensordot(inputs_mean, kernel_mean, [[-1], [0]])
if self.use_bias:
mean = tf.nn.bias_add(mean, bias_mean)
# Cov = E[inputs**2] Cov(kernel) + E[W]^T Cov(inputs) E[W] + Cov(bias)
# For first term, assume Cov(kernel) = 0 on off-diagonals so we only
# compute diagonal term.
covariance_diag = tf.tensordot(inputs_variance + inputs_mean**2,
kernel_variance, [[-1], [0]])
# Compute quadratic form E[W]^T Cov E[W] from right-to-left. First is
# [..., features, features], [features, units] -> [..., features, units].
cov_w = tf.tensordot(inputs_covariance, kernel_mean, [[-1], [0]])
# Next is [..., features, units], [features, units] -> [..., units, units].
w_cov_w = tf.tensordot(cov_w, kernel_mean, [[-2], [0]])
covariance = w_cov_w
if self.use_bias:
covariance += bias_covariance
covariance = tf.matrix_set_diag(
covariance, tf.matrix_diag_part(covariance) + covariance_diag)
if self.activation in (tf.keras.activations.relu, tf.nn.relu):
# Compute activation's moments with variable names from Wu et al. (2018).
variance = tf.matrix_diag_part(covariance)
scale = tf.sqrt(variance)
mu = mean / (scale + tf.keras.backend.epsilon())
mean = scale * soft_relu(mu)
pairwise_variances = (tf.expand_dims(variance, -1) *
tf.expand_dims(variance, -2)) # [..., units, units]
rho = covariance / tf.sqrt(pairwise_variances +
tf.keras.backend.epsilon())
rho = tf.clip_by_value(rho,
-1. / (1. + tf.keras.backend.epsilon()),
1. / (1. + tf.keras.backend.epsilon()))
s = covariance / (rho + tf.keras.backend.epsilon())
mu1 = tf.expand_dims(mu, -1) # [..., units, 1]
mu2 = tf.matrix_transpose(mu1) # [..., 1, units]
a = (soft_relu(mu1) * soft_relu(mu2) +
rho * tfp.distributions.Normal(0., 1.).cdf(mu1) *
tfp.distributions.Normal(0., 1.).cdf(mu2))
gh = tf.asinh(rho)
bar_rho = tf.sqrt(1. - rho**2)
gr = gh + rho / (1. + bar_rho)
# Include numerically stable versions of gr and rho when multiplying or
# dividing them. The sign of gr*rho and rho/gr is always positive.
safe_gr = tf.abs(gr) + 0.5 * tf.keras.backend.epsilon()
safe_rho = tf.abs(rho) + tf.keras.backend.epsilon()
exp_negative_q = gr / (2. * math.pi) * tf.exp(
-safe_rho / (2. * safe_gr * (1 + bar_rho)) +
(gh - rho) / (safe_gr * safe_rho) * mu1 * mu2)
covariance = s * (a + exp_negative_q)
elif self.activation not in (tf.keras.activations.linear, None):
raise NotImplementedError('Activation is {}. Deterministic variational '
'inference is only available if activation is '
'ReLU or None.'.format(self.activation))
return ed.MultivariateNormalFullCovariance(mean, covariance)
def variational_sgpr(X,
Z,
ls=1.,
kern_func=rbf,
ridge_factor=1e-3,
mfvi_mixture=False,
n_mixture=1):
"""Defines the mean-field variational family for GPR.
Args:
X: (np.ndarray of float32) input training features, with dimension (Nx, D).
Z: (np.ndarray of float32) inducing points, with dimension (Nz, D).
ls: (float32) length scale parameter.
kern_func: (function) kernel function.
ridge_factor: (float32) small ridge factor to stabilize Cholesky decomposition
mfvi_mixture: (float32) Whether to output variational family with a
mixture of MFVI.
n_mixture: (int) Number of MFVI mixture component to add.
Returns:
q_f, q_sig: (ed.RandomVariable) variational family.
q_f_mean, q_f_sdev: (tf.Variable) variational parameters for q_f
mixture_par_list: (list of tf.Variable) variational parameters for
MFVI mixture ('mixture_logits', 'mixture_logits_mfvi_mix',
'mean_mfvi', 'sdev_mfvi') if mfvi_mixture=True, else [].
"""
X = tf.convert_to_tensor(X)
Z = tf.convert_to_tensor(Z)
Nx, Nz = X.shape.as_list()[0], Z.shape.as_list()[0]
# 1. Prepare constants
# compute matrix constants
Kxx = kern_func(X, ls=ls)
Kxz = kern_func(X, Z, ls=ls)
Kzz = kern_func(Z, ls=ls, ridge_factor=ridge_factor)
# compute null covariance matrix using Cholesky decomposition
Kzz_chol_inv = tf.matrix_inverse(tf.cholesky(Kzz))
Kzz_inv = tf.matmul(Kzz_chol_inv, Kzz_chol_inv, transpose_a=True)
Kxz_Kzz_chol_inv = tf.matmul(Kxz, Kzz_chol_inv, transpose_b=True)
Kxz_Kzz_inv = tf.matmul(Kxz, Kzz_inv)
Sigma_pre = Kxx - tf.matmul(
Kxz_Kzz_chol_inv, Kxz_Kzz_chol_inv, transpose_b=True)
# 2. Define variational parameters
# define mean and variance for sigma
q_sig_mean = tf.get_variable(shape=[], name='q_sig_mean')
q_sig_sdev = tf.exp(tf.get_variable(shape=[], name='q_sig_sdev'))
# define free parameters (i.e. mean and full covariance of f_latent)
m = tf.get_variable(shape=[Nz], name='qf_m')
s = tf.get_variable(
shape=[Nz * (Nz + 1) / 2],
# initializer=tf.zeros_initializer(),
name='qf_s')
L = fill_triangular(s, name='qf_chol')
S = tf.matmul(L, L, transpose_b=True)
# compute sparse gp variational parameter (i.e. mean and covariance of P(f_obs | f_latent))
qf_mean = tf.tensordot(Kxz_Kzz_inv, m, [[1], [0]], name='qf_mean')
qf_cov = (
Sigma_pre +
tf.matmul(Kxz_Kzz_inv, tf.matmul(S, Kxz_Kzz_inv, transpose_b=True)) +
ridge_factor * tf.eye(Nx, dtype=tf.float32))
# define variational family
mixture_par_list = []
if mfvi_mixture:
gp_dist = tfd.MultivariateNormalFullCovariance(
loc=qf_mean, covariance_matrix=qf_cov)
q_f, mixture_par_list = inference_util.make_mfvi_sgp_mixture_family(
n_mixture=n_mixture, N=Nx, gp_dist=gp_dist, name='q_f')
else:
q_f = ed.MultivariateNormalFullCovariance(loc=qf_mean,
covariance_matrix=qf_cov,
name='q_f')
q_sig = ed.Normal(loc=q_sig_mean, scale=q_sig_sdev, name='q_sig')
return q_f, q_sig, qf_mean, qf_cov, mixture_par_list