<br />

|---|
| [![](https://www.tensorflow.org/images/GitHub-Mark-32px.png) View source on GitHub](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/optimizer/sgld.py#L32-L294) |

An optimizer module for stochastic gradient Langevin dynamics.

    tfp.optimizer.StochasticGradientLangevinDynamics(
        learning_rate,
        preconditioner_decay_rate=0.95,
        data_size=1,
        burnin=25,
        diagonal_bias=1e-08,
        name=None,
        parallel_iterations=10
    )

### Used in the notebooks

| Used in the tutorials |
|---|
| - [Fitting Dirichlet Process Mixture Model Using Preconditioned Stochastic Gradient Langevin Dynamics](https://www.tensorflow.org/probability/examples/Fitting_DPMM_Using_pSGLD) |

This implements the preconditioned Stochastic Gradient Langevin Dynamics
optimizer \[(Li et al., 2016)\]\[1\]. The optimization variable is regarded as a
sample from the posterior under Stochastic Gradient Langevin Dynamics with
noise rescaled in each dimension according to [RMSProp](http://www.cs.toronto.edu/%7Etijmen/csc321/slides/lecture_slides_lec6.pdf).

> [!NOTE]
> **Note:** If a prior is included in the loss, it should be scaled by `1/data_size`, where `data_size` is the number of points in the data set. I.e., it should be divided by the `data_size` term described below.

#### Examples

##### Optimizing energy of a 3D-Gaussian distribution

This example demonstrates that for a fixed step size SGLD works as an
approximate version of MALA (tfp.mcmc.MetropolisAdjustedLangevinAlgorithm).

    import tensorflow as tf
    import tensorflow_probability as tfp
    import numpy as np

    tfd = tfp.distributions
    dtype = np.float32

    with tf.Session(graph=tf.Graph()) as sess:
      # Set up random seed for the optimizer
      tf.random.set_seed(42)
      true_mean = dtype([0, 0, 0])
      true_cov = dtype([[1, 0.25, 0.25], [0.25, 1, 0.25], [0.25, 0.25, 1]])
      # Loss is defined through the Cholesky decomposition
      chol = tf.linalg.cholesky(true_cov)

      var_1 = tf.Variable(name='var_1', initial_value=[1., 1.])
      var_2 = tf.Variable(name='var_2', initial_value=[1.])

      def loss_fn():
        var = tf.concat([var_1, var_2], axis=-1)
        loss_part = tf.linalg.cholesky_solve(chol, var[..., tf.newaxis])
        return tf.linalg.matvec(loss_part, var, transpose_a=True)

      # Set up the learning rate with a polynomial decay
      step = tf.Variable(0, dtype=tf.int64)
      starter_learning_rate = .3
      end_learning_rate = 1e-4
      decay_steps = 1e4
      learning_rate = tf.compat.v1.train.polynomial_decay(
          starter_learning_rate,
          step,
          decay_steps,
          end_learning_rate,
          power=1.)

      # Set up the optimizer
      optimizer_kernel = tfp.optimizer.StochasticGradientLangevinDynamics(
          learning_rate=learning_rate, preconditioner_decay_rate=0.99)
      optimizer_kernel.iterations = step
      optimizer = optimizer_kernel.minimize(loss_fn, var_list=[var_1, var_2])

      # Number of training steps
      training_steps = 5000
      # Record the steps as and treat them as samples
      samples = [np.zeros([training_steps, 2]), np.zeros([training_steps, 1])]
      sess.run(tf.compat.v1.global_variables_initializer())
      for step in range(training_steps):
        sess.run(optimizer)
        sample = [sess.run(var_1), sess.run(var_2)]
        samples[0][step, :] = sample[0]
        samples[1][step, :] = sample[1]

      samples_ = np.concatenate(samples, axis=-1)
      sample_mean = np.mean(samples_, 0)
      print('sample mean', sample_mean)

Args:
learning_rate: Scalar `float`-like `Tensor`. The base learning rate for the
optimizer. Must be tuned to the specific function being minimized.
preconditioner_decay_rate: Scalar `float`-like `Tensor`. The exponential
decay rate of the rescaling of the preconditioner (RMSprop). (This is
"alpha" in Li et al. (2016)). Should be smaller than but nearly `1` to
approximate sampling from the posterior. (Default: `0.95`)
data_size: Scalar `int`-like `Tensor`. The effective number of
points in the data set. Assumes that the loss is taken as the mean over a
minibatch. Otherwise if the sum was taken, divide this number by the
batch size. If a prior is included in the loss function, it should be
normalized by `data_size`. Default value: `1`.
burnin: Scalar `int`-like `Tensor`. The number of iterations to collect
gradient statistics to update the preconditioner before starting to draw
noisy samples. (Default: `25`)
diagonal_bias: Scalar `float`-like `Tensor`. Term added to the diagonal of
the preconditioner to prevent the preconditioner from degenerating.
(Default: `1e-8`)
name: Python `str` describing ops managed by this function.
(Default: `"StochasticGradientLangevinDynamics"`)
parallel_iterations: the number of coordinates for which the gradients of
the preconditioning matrix can be computed in parallel. Must be a
positive integer.

<br />

<br />

<br />

| ## Raises ||
|---|---|
| `InvalidArgumentError` | If preconditioner_decay_rate is a `Tensor` not in `(0,1]`. |
| `NotImplementedError` | If eager execution is enabled. |

<br />

#### References

\[1\]: Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin.
Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural
Networks. In *Association for the Advancement of Artificial
Intelligence* , 2016. <https://arxiv.org/abs/1512.07666>

<br />

<br />

<br />

| ## Args ||
|---|---|
| `name` | String. The name to use for momentum accumulator weights created by the optimizer. |
| `gradient_aggregator` | The function to use to aggregate gradients across devices (when using [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy)). If `None`, defaults to summing the gradients across devices. The function should accept and return a list of `(gradient, variable)` tuples. |
| `gradient_transformers` | Optional. List of functions to use to transform gradients before applying updates to Variables. The functions are applied after `gradient_aggregator`. The functions should accept and return a list of `(gradient, variable)` tuples. |
| `**kwargs` | keyword arguments. Allowed arguments are `clipvalue`, `clipnorm`, `global_clipnorm`. If `clipvalue` (float) is set, the gradient of each weight is clipped to be no higher than this value. If `clipnorm` (float) is set, the gradient of each weight is individually clipped so that its norm is no higher than this value. If `global_clipnorm` (float) is set the gradient of all weights is clipped so that their global norm is no higher than this value. |

<br />

<br />

<br />

<br />

| ## Attributes ||
|---|---|
| `clipnorm` | `float` or `None`. If set, clips gradients to a maximum norm. |
| `clipvalue` | `float` or `None`. If set, clips gradients to a maximum value. |
| `global_clipnorm` | `float` or `None`. <br /> If set, clips gradients to a maximum norm. Check [`tf.clip_by_global_norm`](https://www.tensorflow.org/api_docs/python/tf/clip_by_global_norm) for more details. |
| `iterations` | Variable. The number of training steps this Optimizer has run. |
| `variable_scope` | Variable scope of all calls to `tf.get_variable`. |
| `weights` | Returns variables of this Optimizer based on the order created. |

<br />

## Methods

### `add_slot`

    add_slot(
        var, slot_name, initializer='zeros', shape=None
    )

Add a new slot variable for `var`.

A slot variable is an additional variable associated with `var` to
train. It is allocated and managed by optimizers, e.g. `Adam`.

<br />

<br />

<br />

| Args ||
|---|---|
| `var` | a `Variable` object. |
| `slot_name` | name of the slot variable. |
| `initializer` | initializer of the slot variable |
| `shape` | (Optional) shape of the slot variable. If not set, it will default to the shape of `var`. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| A slot variable. ||

<br />

### `add_weight`

    add_weight(
        name,
        shape,
        dtype=None,
        initializer='zeros',
        trainable=None,
        synchronization=tf.VariableSynchronization.AUTO,
        aggregation=tf.VariableAggregation.NONE
    )

### `apply_gradients`

    apply_gradients(
        grads_and_vars, name=None, experimental_aggregate_gradients=True
    )

Apply gradients to variables.

This is the second part of `minimize()`. It returns an `Operation` that
applies gradients.

The method sums gradients from all replicas in the presence of
[`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) by default. You can aggregate gradients
yourself by passing `experimental_aggregate_gradients=False`.

#### Example:

    grads = tape.gradient(loss, vars)
    grads = tf.distribute.get_replica_context().all_reduce('sum', grads)
    # Processing aggregated gradients.
    optimizer.apply_gradients(zip(grads, vars),
        experimental_aggregate_gradients=False)

<br />

<br />

<br />

| Args ||
|---|---|
| `grads_and_vars` | List of (gradient, variable) pairs. |
| `name` | Optional name for the returned operation. When `None`, uses the name passed to the `Optimizer` constructor. Defaults to `None`. |
| `experimental_aggregate_gradients` | Whether to sum gradients from different replicas in the presence of [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy). If False, it's user responsibility to aggregate the gradients. Default to `True`. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| An `Operation` that applies the specified gradients. The `iterations` will be automatically increased by 1. ||

<br />

<br />

<br />

<br />

| Raises ||
|---|---|
| `TypeError` | If `grads_and_vars` is malformed. |
| `ValueError` | If none of the variables have gradients. |
| `RuntimeError` | If called in a cross-replica context. |

<br />

### `from_config`

    @classmethod
    from_config(
        config, custom_objects=None
    )

Creates an optimizer from its config.

This method is the reverse of `get_config`,
capable of instantiating the same optimizer from the config
dictionary.

<br />

<br />

<br />

| Args ||
|---|---|
| `config` | A Python dictionary, typically the output of get_config. |
| `custom_objects` | A Python dictionary mapping names to additional Python objects used to create this optimizer, such as a function used for a hyperparameter. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| An optimizer instance. ||

<br />

### `get_config`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/optimizer/sgld.py#L214-L217)

    get_config()

Returns the config of the optimizer.

An optimizer config is a Python dictionary (serializable)
containing the configuration of an optimizer.
The same optimizer can be reinstantiated later
(without any saved state) from this configuration.

<br />

<br />

<br />

| Returns ||
|---|---|
| Python dictionary. ||

<br />

### `get_gradients`

    get_gradients(
        loss, params
    )

Returns gradients of `loss` with respect to `params`.

Should be used only in legacy v1 graph mode.

<br />

<br />

<br />

| Args ||
|---|---|
| `loss` | Loss tensor. |
| `params` | List of variables. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| List of gradient tensors. ||

<br />

<br />

<br />

<br />

| Raises ||
|---|---|
| `ValueError` | In case any gradient cannot be computed (e.g. if gradient function not implemented). |

<br />

### `get_slot`

    get_slot(
        var, slot_name
    )

### `get_slot_names`

    get_slot_names()

A list of names for this optimizer's slots.

### `get_updates`

    get_updates(
        loss, params
    )

### `get_weights`

    get_weights()

Returns the current weights of the optimizer.

The weights of an optimizer are its state (ie, variables).
This function returns the weight values associated with this
optimizer as a list of Numpy arrays. The first value is always the
iterations count of the optimizer, followed by the optimizer's state
variables in the order they were created. The returned list can in turn
be used to load state into similarly parameterized optimizers.

For example, the RMSprop optimizer for this simple model returns a list
of three values-- the iteration count, followed by the root-mean-square
value of the kernel and bias of the single Dense layer:

    opt = tf.keras.optimizers.legacy.RMSprop()
    m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
    m.compile(opt, loss='mse')
    data = np.arange(100).reshape(5, 20)
    labels = np.zeros(5)
    results = m.fit(data, labels)  # Training.
    len(opt.get_weights())
    3

<br />

<br />

<br />

| Returns ||
|---|---|
| Weights values as a list of numpy arrays. ||

<br />

### `minimize`

    minimize(
        loss, var_list, grad_loss=None, name=None, tape=None
    )

Minimize `loss` by updating `var_list`.

This method simply computes gradient using [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/tf/GradientTape) and calls
`apply_gradients()`. If you want to process the gradient before applying
then call [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/tf/GradientTape) and `apply_gradients()` explicitly instead
of using this function.

<br />

<br />

<br />

| Args ||
|---|---|
| `loss` | `Tensor` or callable. If a callable, `loss` should take no arguments and return the value to minimize. If a `Tensor`, the `tape` argument must be passed. |
| `var_list` | list or tuple of `Variable` objects to update to minimize `loss`, or a callable returning the list or tuple of `Variable` objects. Use callable when the variable list would otherwise be incomplete before `minimize` since the variables are created at the first time `loss` is called. |
| `grad_loss` | (Optional). A `Tensor` holding the gradient computed for `loss`. |
| `name` | (Optional) str. Name for the returned operation. |
| `tape` | (Optional) [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/tf/GradientTape). If `loss` is provided as a `Tensor`, the tape that computed the `loss` must be provided. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| An `Operation` that updates the variables in `var_list`. The `iterations` will be automatically increased by 1. ||

<br />

<br />

<br />

<br />

| Raises ||
|---|---|
| `ValueError` | If some of the variables are not `Variable` objects. |

<br />

### `set_weights`

    set_weights(
        weights
    )

Set the weights of the optimizer.

The weights of an optimizer are its state (ie, variables).
This function takes the weight values associated with this
optimizer as a list of Numpy arrays. The first value is always the
iterations count of the optimizer, followed by the optimizer's state
variables in the order they are created. The passed values are used to
set the new state of the optimizer.

For example, the RMSprop optimizer for this simple model takes a list of
three values-- the iteration count, followed by the root-mean-square
value of the kernel and bias of the single Dense layer:

    opt = tf.keras.optimizers.legacy.RMSprop()
    m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
    m.compile(opt, loss='mse')
    data = np.arange(100).reshape(5, 20)
    labels = np.zeros(5)
    results = m.fit(data, labels)  # Training.
    new_weights = [np.array(10), np.ones([20, 10]), np.zeros([10])]
    opt.set_weights(new_weights)
    opt.iterations
    <tf.Variable 'RMSprop/iter:0' shape=() dtype=int64, numpy=10>

<br />

<br />

<br />

| Args ||
|---|---|
| `weights` | weight values as a list of numpy arrays. |

<br />

### `variables`

    variables()

Returns variables of this Optimizer based on the order created.