def rmsprop(
learning_rate: ScalarOrSchedule,
decay: float = 0.9,
eps: float = 1e-8,
initial_scale: float = 0.,
centered: bool = False,
momentum: Optional[float] = None,
nesterov: bool = False
) -> base.GradientTransformation:
# pylint: disable=line-too-long
"""A flexible RMSProp optimiser.
RMSProp is an SGD variant with learning rate adaptation. The `learning_rate`
used for each weight is scaled by a suitable estimate of the magnitude of the
gradients on previous steps. Several variants of RMSProp can be found
in the literature. This alias provides an easy to configure RMSProp
optimiser that can be used to switch between several of these variants.
References:
Tieleman and Hinton, 2012: http://www.cs.toronto.edu/~hinton/coursera/lecture6/lec6.pdf
Graves, 2013: https://arxiv.org/abs/1308.0850
Args:
learning_rate: this is a fixed global scaling factor.
decay: the decay used to track the magnitude of previous gradients.
eps: a small numerical constant to avoid dividing by zero when rescaling.
initial_scale: (default `0.`), initialisation of accumulators tracking the
magnitude of previous updates. PyTorch uses `0`, TF1 uses `1`. When
reproducing results from a paper, verify the value used by the authors.
centered: (default `False`), whether the second moment or the variance of
the past gradients is used to rescale the latest gradients.
momentum: (default `None`), the `decay` rate used by the momentum term,
when it is set to `None`, then momentum is not used at all.
nesterov (default `False`): whether nesterov momentum is used.
Returns:
the corresponding `GradientTransformation`.
"""
# pylint: enable=line-too-long
if centered:
return combine.chain(
transform.scale_by_stddev(
decay=decay, eps=eps, initial_scale=initial_scale),
_scale_by_learning_rate(learning_rate),
(transform.trace(decay=momentum, nesterov=nesterov)
if momentum is not None else base.identity())
)
return combine.chain(
transform.scale_by_rms(
decay=decay, eps=eps, initial_scale=initial_scale),
_scale_by_learning_rate(learning_rate),
(transform.trace(decay=momentum, nesterov=nesterov)
if momentum is not None else base.identity())
)
def sgd(
learning_rate: ScalarOrSchedule,
momentum: Optional[float] = None,
nesterov: bool = False
) -> base.GradientTransformation:
"""A canonical Stochastic Gradient Descent optimiser.
This implements stochastic gradient descent. It also includes support for
momentum, and nesterov acceleration, as these are standard practice when
using stochastic gradient descent to train deep neural networks.
References:
Sutskever et al, 2013: http://proceedings.mlr.press/v28/sutskever13.pdf
Args:
learning_rate: this is a fixed global scaling factor.
momentum: (default `None`), the `decay` rate used by the momentum term,
when it is set to `None`, then momentum is not used at all.
nesterov (default `False`): whether nesterov momentum is used.
Returns:
A `GradientTransformation`.
"""
return combine.chain(
(transform.trace(decay=momentum, nesterov=nesterov)
if momentum is not None else base.identity()),
_scale_by_learning_rate(learning_rate)
)
def dpsgd(
learning_rate: ScalarOrSchedule,
l2_norm_clip: float,
noise_multiplier: float,
seed: int,
momentum: Optional[float] = None,
nesterov: bool = False
) -> base.GradientTransformation:
"""The DPSGD optimiser.
Differential privacy is a standard for privacy guarantees of algorithms
learning from aggregate databases including potentially sensitive information.
DPSGD offers protection against a strong adversary with full knowledge of the
training mechanism and access to the model’s parameters.
WARNING: This `GradientTransformation` expects input updates to have a batch
dimension on the 0th axis. That is, this function expects per-example
gradients as input (which are easy to obtain in JAX using `jax.vmap`).
References:
Abadi et al, 2016: https://arxiv.org/abs/1607.00133
Args:
learning_rate: this is a fixed global scaling factor.
l2_norm_clip: maximum L2 norm of the per-example gradients.
noise_multiplier: ratio of standard deviation to the clipping norm.
seed: initial seed used for the jax.random.PRNGKey
momentum: (default `None`), the `decay` rate used by the momentum term,
when it is set to `None`, then momentum is not used at all.
nesterov (default `False`): whether nesterov momentum is used.
Returns:
A `GradientTransformation`.
"""
return combine.chain(
privacy.differentially_private_aggregate(
l2_norm_clip=l2_norm_clip,
noise_multiplier=noise_multiplier,
seed=seed),
(transform.trace(decay=momentum, nesterov=nesterov)
if momentum is not None else base.identity()),
_scale_by_learning_rate(learning_rate)
)