Source code for gmmvi.optimization.gmmvi_modules.component_adaptation

from math import pi
from typing import Union

import tensorflow as tf
import tensorflow_probability as tfp

from gmmvi.models.gmm_wrapper import GmmWrapper
from gmmvi.optimization.sample_db import SampleDB
from gmmvi.models.diagonal_gmm import DiagonalGMM
from gmmvi.experiments.target_distributions.lnpdf import LNPDF

[docs]class ComponentAdaptation: """ This class provides a common interface for adapting the number of components. There are currently only two options: 1. The :py:class:`FixedComponentAdaptation<gmmvi.optimization.gmmvi_modules.component_adaptation.FixedComponentAdaptation>` a dummy-class, that does not do anything. 2. The :py:class:`VipsComponentAdaptation<gmmvi.optimization.gmmvi_modules.component_adaptation.VipsComponentAdaptation>` uses the procedure of VIPS :cite:p:`Arenz2020` to add and delete components. Parameters: gmm_wrapper: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model where we want to adapt the number of components. sample_db: :py:class:`SampleDB<gmmvi.optimization.sample_db.SampleDB>` The sample database can be used to select candidate locations for adding a new component, without having to perform additional queries to the target distribution. target_distribution: :py:class:`LNPDF<gmmvi.experiments.target_distributions.lnpdf.LNPDF>` The target distribution can be used to evaluate candidate locations for adding a new component. prior_mean: tf.Tensor A one dimensional tensor of size num_dimensions, specifying the mean of the Gaussian that we can use to sample candidate locations for adding a new component. initial_cov: tf.Tensor A two-dimensional tensor of size num_dimensions x num_dimensions, specifying the covariance of the Gaussian that we can use to sample candidate locations for adding a new component. """ def __init__(self): pass
[docs] @staticmethod def build_from_config(config, gmm_wrapper, sample_db, target_distribution, prior_mean, initial_cov): """This static method provides a convenient way to create a :py:class:`FixedComponentAdaptation<gmmvi.optimization.gmmvi_modules.component_adaptation.FixedComponentAdaptation>` or :py:class:`VipsComponentAdaptation<gmmvi.optimization.gmmvi_modules.component_adaptation.VipsComponentAdaptation>` depending on the provided config. Parameters: config: dict The dictionary is typically read from YAML a file, and holds all hyperparameters. gmm_wrapper: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model where we want to adapt the number of components. sample_db: :py:class:`SampleDB<gmmvi.optimization.sample_db.SampleDB>` The sample database can be used to select candidate locations for adding a new component, without having to perform additional queries to the target distribution. target_distribution: :py:class:`LNPDF<gmmvi.experiments.target_distributions.lnpdf.LNPDF>` The target distribution can be used to evaluate candidate locations for adding a new component. prior_mean: tf.Tensor A one dimensional tensor of size num_dimensions, specifying the mean of the Gaussian that we can use to sample candidate locations for adding a new component. initial_cov: tf.Tensor A two-dimensional tensor of size num_dimensions x num_dimensions, specifying the covariance of the Gaussian that we can use to sample candidate locations for adding a new component. """ if config["num_component_adapter_type"] == "adaptive": return VipsComponentAdaptation(gmm_wrapper, sample_db, target_distribution, prior_mean, initial_cov, **config["num_component_adapter_config"]) elif config["num_component_adapter_type"] == "fixed": return FixedComponentAdaptation(**config["num_component_adapter_config"]) else: raise ValueError( f"config['num_component_adapter_type'] is '{config['num_componenter_adapter_type']}' " f"which is an unknown type")
def adapt_number_of_components(self, iteration): raise NotImplementedError
[docs]class FixedComponentAdaptation(ComponentAdaptation): """ This is a dummy class, used when we do not want to adapt the number of components during learning. """ def __init__(self): super(FixedComponentAdaptation, self).__init__()
[docs] def adapt_number_of_components(self, iteration): """ As we do not want to change the number of components, this method does not do anything. Parameters: iteration: int The current iteration (ignored). """ pass
[docs]class VipsComponentAdaptation(ComponentAdaptation): """ This class implements the component adaptation procedure used by VIPS. See :cite:p:`Arenz2020`. Parameters: gmm_wrapper: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model where we want to adapt the number of components. sample_db: :py:class:`SampleDB<gmmvi.optimization.sample_db.SampleDB>` The sample database can be used to select candidate locations for adding a new component, without having to perform additional queries to the target distribution. target_distribution: :py:class:`LNPDF<gmmvi.experiments.target_distributions.lnpdf.LNPDF>` The target distribution can be used to evaluate candidate locations for adding a new component. prior_mean: tf.Tensor A one dimensional tensor of size num_dimensions, specifying the mean of the Gaussian that we can use to sample candidate locations for adding a new component. initial_cov: tf.Tensor A two-dimensional tensor of size num_dimensions x num_dimensions, specifying the covariance of the Gaussian that we can use to sample candidate locations for adding a new component. del_iters: int minimum number of updates a component needs to have received, before it is considered as candidate for deletion. add_iters: int a new component will be added every *add_iters* iterations max_components: int do not add components, if the model has at least *max_components* components num_database_samples: int number of samples from the :py:class:`SampleDB<gmmvi.optimization.sample_db.SampleDB>` that are used for selecting a good initial mean when adding a new component. num_prior_samples: int number of samples from the prior distribution that are used for selecting a good initial mean when adding a new component. """ def __init__(self, model: GmmWrapper, sample_db: SampleDB, target_lnpdf: LNPDF, prior_mean: Union[float, tf.Tensor], initial_cov: Union[float, tf.Tensor], del_iters: int, add_iters: int, max_components: int, thresholds_for_add_heuristic: float, min_weight_for_del_heuristic: float, num_database_samples: int, num_prior_samples: int): super(VipsComponentAdaptation, self).__init__() self.model = model if (prior_mean is not None) and (initial_cov is not None): if tf.rank(initial_cov) == 0: initial_cov = initial_cov * tf.ones(model.num_dimensions) if tf.rank(prior_mean) == 0: prior_mean = prior_mean * tf.ones(model.num_dimensions) self.prior = DiagonalGMM(tf.ones(1), tf.expand_dims(prior_mean, 0), tf.expand_dims(initial_cov, 0)) else: self.prior = None self.num_prior_samples = num_prior_samples self.target_lnpdf = target_lnpdf self.sample_db = sample_db self.del_iters = del_iters self.add_iters = add_iters self.max_components = max_components self.num_db_samples = num_database_samples self.num_calls_to_add_heuristic = tf.Variable(0) self.thresholds_for_addHeuristic = tf.convert_to_tensor(thresholds_for_add_heuristic, dtype=tf.float32) self.min_weight_for_del_heuristic = min_weight_for_del_heuristic self.reward_improvements = tf.Variable(tf.zeros(0), shape=[None], dtype=tf.float32) self.filter_delay = tf.cast(tf.floor(self.del_iters / 3), dtype=tf.int32) gaussian = tfp.distributions.Normal(tf.constant(0, tf.float32), tf.constant(self.del_iters / 8., tf.float32)) kernel = gaussian.prob(tf.range(start=-self.filter_delay, limit=self.filter_delay, dtype=tf.float32)) self.kernel = tf.reshape(kernel / tf.reduce_sum(kernel), [-1, 1, 1])
[docs] def adapt_number_of_components(self, iteration: int): """ This method may change the number of components, either by deleting bad components that have low weights, or by adding new components. Parameters: iteration: int The current iteration, used to decide whether a new component should be added. """ if iteration > self.del_iters: self.delete_bad_components() if iteration > 1 and iteration % self.add_iters == 0: if self.model.num_components < self.max_components: self.add_new_component()
[docs] @tf.function(experimental_relax_shapes=True) def add_at_best_location(self, samples, target_lnpdfs): """ Find the most promising :cite:p:`Arenz2020` location among the provided samples for adding a new component, that is, a new component will be added with mean given by one of the provided samples. Parameters: samples: tf.Tensor candidate locations for initializing the mean of the new component target_lnpdfs: tf.Tensor for each candidate location, this tensor contains the log-density under the (unnormalized) target distribution. """ iter = self.num_calls_to_add_heuristic % len(self.thresholds_for_addHeuristic) model_log_densities = self.model.log_density(samples) init_weight = 1e-29 a = tf.random.uniform([1]) if self.prior is not None: des_entropy = self.model.get_average_entropy() * a + self.prior.get_average_entropy() * (1 - a) else: des_entropy = self.model.get_average_entropy() max_logdensity = tf.reduce_max(model_log_densities) rewards = target_lnpdfs - tf.maximum( max_logdensity - self.thresholds_for_addHeuristic[iter], model_log_densities ) new_mean = samples[tf.argmax(rewards)] H_unscaled = 0.5 * self.model.num_dimensions * (tf.math.log(2.0 * pi) + 1) c = tf.math.exp((2 * (des_entropy - H_unscaled)) / self.model.num_dimensions) if self.model.diagonal_covs: new_cov = c * tf.ones(self.model.num_dimensions) else: new_cov = c * tf.eye(self.model.num_dimensions) self.model.add_component(init_weight, new_mean, new_cov, tf.reshape(self.thresholds_for_addHeuristic[iter], [1]), tf.reshape(des_entropy, [1]))
[docs] @tf.function def select_samples_for_adding_heuristic(self): """ Select a set of samples used as candidates for initializing the mean of the new component. Returns: tuple(tf.Tensor, tf.Tensor, tf.Tensor): **samples** - the selected candidate locations **target_lnpdfs** - log-densities of the *samples* under the unnormalized target distribution **prior_samples** - additional samples drawn from a prior, which have not yet been evaluated on the target distribution. """ self.num_calls_to_add_heuristic.assign_add(1) samples, target_lnpdfs = self.sample_db.get_random_sample(self.num_db_samples) prior_samples = tf.zeros((0, self.model.num_dimensions), tf.float32) if self.num_prior_samples > 0: prior_samples = self.prior.sample(self.num_prior_samples)[0] self.sample_db.num_samples_written.assign_add(self.num_prior_samples) return samples, target_lnpdfs, prior_samples
[docs] def add_new_component(self): """ This method adds a new component by first selecting a set of candidate locations and the choosing the most promising one using the procedure of VIPS :cite:p:`Arenz2020`. """ samples, target_lnpdfs, prior_samples = self.select_samples_for_adding_heuristic() if self.num_prior_samples > 0: samples = tf.concat((samples, prior_samples), 0) target_lnpdfs = tf.concat((target_lnpdfs, self.target_lnpdf.log_density(prior_samples)), 0) self.add_at_best_location(samples, target_lnpdfs)
[docs] def delete_bad_components(self): """ Components are deleted, if all the following criteria are met received: 1. It must have received at least *del_iters* updates 2. It must not have improved significantly during the last iterations. In contrast to VIPS, we use a Gaussian filter to smooth the rewards of the component, to be more robust with respect to noisy target distributions. 3. It must have very low weight, such that the effects on the model are negligible. """ # estimate the relative improvement for every component with respect to # the improvement it would need to catch up (assuming linear improvement) with the best component current_smoothed_reward = tf.reduce_mean( self.model.reward_history[:, -tf.size(self.kernel):] * tf.reshape(self.kernel, [1, -1]), axis=1) old_smoothed_reward = tf.reduce_mean( self.model.reward_history[:, -tf.size(self.kernel)-self.del_iters:-self.del_iters] * tf.reshape(self.kernel, [1, -1]), axis=1) old_smoothed_reward -= tf.reduce_max(current_smoothed_reward) current_smoothed_reward -= tf.reduce_max(current_smoothed_reward) reward_improvements = (current_smoothed_reward - old_smoothed_reward) / tf.abs(old_smoothed_reward) self.reward_improvements.assign(reward_improvements) # compute for each component the maximum weight it had within the last del_iters, # or that it would have gotten when we used greedy updates max_actual_weights = tf.reduce_max(self.model.weight_history[:, -tf.size(self.kernel)-self.del_iters:-1], axis=1) max_greedy_weights = tf.reduce_max(tf.exp( self.model.reward_history[:, -tf.size(self.kernel)-self.del_iters:] - tf.math.reduce_logsumexp(self.model.reward_history[:, -tf.size(self.kernel)-self.del_iters:], axis=0, keepdims=True)), axis=1) max_weights = tf.math.maximum(max_actual_weights, max_greedy_weights) is_stagnating = reward_improvements <= 0.4 is_low_weight = max_weights < self.min_weight_for_del_heuristic is_old_enough = self.model.reward_history[:, -self.del_iters] != -tf.float32.max is_bad = tf.reduce_all((is_stagnating, is_low_weight, is_old_enough), axis=0) bad_component_indices = tf.squeeze(tf.where(is_bad), axis=1) if tf.size(bad_component_indices) > 0: for idx in tf.sort(bad_component_indices, direction='DESCENDING'): self.model.remove_component(idx)