# NgBasedComponentUpdater

class gmmvi.optimization.gmmvi_modules.ng_based_component_updater.NgBasedComponentUpdater(model: GmmWrapper, temperature: float)

  This class provides a common interface for updating the Gaussian components along the natural gradient.

  The Gaussian components of the mixture model, are updated by updating their parameters (their mean and covariance) based on previously computed estimates of the natural gradient (see `NgEstimator`) and stepsizes (see `ComponentStepsizeAdaptation`).

  There are currently three options for updating the components:

  1. The `DirectNgBasedComponentUpdater` straightforwardly performs the natural gradient with the given stepsize.

  2. The `NgBasedComponentUpdaterIblr` uses the improved Bayesian learning rate to ensure positive definite covariance matrices.

  3. The `KLConstrainedNgBasedComponentUpdater` treats the stepsize as a trust-region constraint.

  Parameters:

    - gmm_wrapper – `GmmWrapper` The wrapped model where we want to update the components.

    - temperature – float Usually temperature=1., can be used to scale the importance of maximizing the model entropy.

  apply_NG_update(expected_hessians_neg: Tensor, expected_gradients_neg: Tensor, stepsizes: Tensor)

    Update the components based on the estimates of the natural gradients (provided in terms of the negated expected Hessian and expected gradients) and the given component-specific stepsizes.

    Parameters:

      - expected_hessians_neg – tf.Tensor A three-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - expected_gradients_neg – tf.Tensor A two-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - stepsizes – tf.Tensor A one-dimensional tensor of shape num_components, containing the stepsize for each component.

  static build_from_config(config, gmm_wrapper)

    This static method provides a convenient way to create a `DirectNgBasedComponentUpdater`, `NgBasedComponentUpdaterIblr` or `KLConstrainedNgBasedComponentUpdater` depending on the provided config.

    Parameters:

      - config – dict The dictionary is typically read from YAML a file, and holds all hyperparameters.

      - gmm_wrapper – `GmmWrapper` The wrapped model for which we want to update the components

## DirectNgBasedComponentUpdater

class gmmvi.optimization.gmmvi_modules.ng_based_component_updater.DirectNgBasedComponentUpdater(model: GmmWrapper, temperature: float)

  This class straightforwardly performs the natural gradient update with the given stepsize.

  Parameters:

    - gmm_wrapper – `GmmWrapper` The wrapped model where we want to update the components.

    - temperature – float Usually temperature=1., can be used to scale the importance of maximizing the model entropy.

  apply_NG_update(expected_hessians_neg: Tensor, expected_gradients_neg: Tensor, stepsizes: Tensor)

    Update the components based on the estimates of the natural gradients (provided in terms of the negated expected Hessian and expected gradients) and the given component-specific stepsizes.

    Parameters:

      - expected_hessians_neg – tf.Tensor A three-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - expected_gradients_neg – tf.Tensor A two-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - stepsizes – tf.Tensor A one-dimensional tensor of shape num_components, containing the stepsize for each component.

## KLConstrainedNgBasedComponentUpdater

class gmmvi.optimization.gmmvi_modules.ng_based_component_updater.KLConstrainedNgBasedComponentUpdater(model: GmmWrapper, temperature: float)

  Updates the component by treating the stepsize as a constraint on the KL-divergence to the current component.

  This class updates the component by performing the largest update along the natural gradient direction, that confines with a trust-region constraint on the Kullback-Leibler divergence with respect to the current component [AZN20].

  Parameters:

    - gmm_wrapper – `GmmWrapper` The wrapped model where we want to update the components.

    - temperature – float Usually temperature=1., can be used to scale the importance of maximizing the model entropy.

  apply_NG_update(expected_hessians_neg: Tensor, expected_gradients_neg: Tensor, stepsizes: Tensor)

    Update the components based on the estimates of the natural gradients (provided in terms of the negated expected Hessian and expected gradients) and the given component-specific stepsizes.

    Parameters:

      - expected_hessians_neg – tf.Tensor A three-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - expected_gradients_neg – tf.Tensor A two-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - stepsizes – tf.Tensor A one-dimensional tensor of shape num_components, containing the stepsize for each component.

  bracketing_search(kl_bound: tf.float32, lower_bound: tf.float32, upper_bound: tf.float32, old_lin_term: Tensor, old_precision: Tensor, old_inv_chol: Tensor, reward_lin_term: Tensor, reward_quad_term: Tensor, kl_const_part: tf.float32, old_mean: Tensor, eta_in_logspace: tf.float32) → [tf.float32, tf.float32]

    This method finds the largest stepsize eta, such that the updated component stays within a KL-constrained trust-region around the current component. Whereas, [AZN20] used L-BFGS-B to solve this convex optimization problem, here we use a simple bracketing search, which seems to be more robust (by not relying on gradients, and can be efficiently performed within a Tensorflow graph. The procedure simple keeps track of a lower bound and an upper bound on the optimal stepsize and recursively evaluates the arithmetic mean of both bounds. If this mean-stepsize results in a two large KL divergence, or in a non-positive-definite covariance matrix, it becomes the new lower bound; otherwise the new upper bound.

    Parameters:

      - kl_bound – tf.float32 The trust region constraint

      - lower_bound – tf.float32 The initial lower bound on the stepsize

      - upper_bound – The initial upper bound on the stepsize

      - old_lin_term – tf.Tensor The linear term of the canonical Gaussian form of the current component. A one-dimensional tensor of shape num_dimensions.

      - old_precision – The precision matrix of the current component. A two-dimensional tensor of shape num_dimensions x num_dimensions.

      - old_inv_chol – The inverse of the Cholesky matrix of the current component. A two-dimensional tensor of shape num_dimensions x num_dimensions.

      - reward_lin – When using `MORE` to estimate the natural gradient, this tensor correspond to the linear coefficient of the quadratic reward model. When using Stein’s Lemma, this term can be computed from the expected gradient and expected Hessian.

      - reward_quad – When using `MORE` to estimate the natural gradient, this tensor correspond to the quadratic coefficient of the quadratic reward model. When using Stein’s Lemma, this term can be computed from the expected Hessian.

      - kl_const_part – A term of the KL divergence that can be precomputed as it does not depend on the parameters of the updated component.

      - old_mean – The mean of the current component.

      - eta_in_logspace – if true, the bracketing search should be performed in log-space (requires fewer iterations)

    Returns:

      new_lower_bound - The lower bound after a stopping criterion was reached.

      new_upper_bound - The upper bound after a stopping criterion was reached.

    Return type:

      tuple(tf.float32, tf.float32)

  kl(eta: tf.float32, old_lin_term: Tensor, old_precision: Tensor, old_inv_chol: Tensor, reward_lin: Tensor, reward_quad: Tensor, kl_const_part: tf.float32, old_mean: Tensor, eta_in_logspace: bool) → [tf.float32, <class 'tensorflow.python.framework.ops.Tensor'>, <class 'tensorflow.python.framework.ops.Tensor'>, <class 'tensorflow.python.framework.ops.Tensor'>]

    Computes the Kullback-Leibler divergence between the updated component and current component, when updating with stepsize eta along the natural gradient.

    Parameters:

      - eta – tf.float32 The stepsize for which the KL divergence should be computed.

      - old_lin_term – tf.Tensor The linear term of the canonical Gaussian form of the current component. A one-dimensional tensor of shape num_dimensions.

      - old_precision – tf.Tensor The precision matrix of the current component. A two-dimensional tensor of shape num_dimensions x num_dimensions.

      - old_inv_chol – tf.Tensor The inverse of the Cholesky matrix of the current component. A two-dimensional tensor of shape num_dimensions x num_dimensions.

      - reward_lin – tf.Tensor When using `MORE` to estimate the natural gradient, this tensor correspond to the linear coefficient of the quadratic reward model. When using Stein’s Lemma, this term can be computed from the expected gradient and expected Hessian.

      - reward_quad – tf.Tensor When using `MORE` to estimate the natural gradient, this tensor correspond to the quadratic coefficient of the quadratic reward model. When using Stein’s Lemma, this term can be computed from the expected Hessian.

      - kl_const_part – tf.float32 A term of the KL divergence that can be precomputed as it does not depend on the parameters of the updated component.

      - old_mean – tf.Tensor The mean of the current component.

      - eta_in_logspace – bool if true, the provided eta is given in log-space.

    Returns:

      kl - a float corresponding to the KL between the updated component and the old component.

      new_mean - a tensor specifying the mean of the updated component.

      new_precision - a tensor specifying the precision of the updated component.

      inv_chol_inv - a tensor specifying the inverse of the Cholesky of the precision matrix of the updated component.

    Return type:

      tuple(float, tf.Tensor, tf.Tensor, tf.Tensor)

## NgBasedComponentUpdaterIblr

class gmmvi.optimization.gmmvi_modules.ng_based_component_updater.NgBasedComponentUpdaterIblr(model: GmmWrapper, temperature: float)

  This class updates the component using the improved Bayesian learning rule.

  > The iBLR performs Riemannian gradient descent and ensures positive definite covariance matrices [LSK20].

  Parameters:

    - gmm_wrapper – `GmmWrapper` The wrapped model where we want to update the components.

    - temperature – float Usually temperature=1., can be used to scale the importance of maximizing the model entropy.

  apply_NG_update(expected_hessians_neg: Tensor, expected_gradients_neg: Tensor, stepsizes: Tensor)

    Update the components based on the estimates of the natural gradients (provided in terms of the negated expected Hessian and expected gradients) and the given component-specific stepsizes.

    Parameters:

      - expected_hessians_neg – tf.Tensor A three-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - expected_gradients_neg – tf.Tensor A two-dimensional tensor of shape num_components x num_dimensions x num_dimensions, containing an estimate of for each component.

      - stepsizes – tf.Tensor A one-dimensional tensor of shape num_components, containing the stepsize for each component.