def adam(learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8,
eps_root: float = 0.0) -> base.GradientTransformation:
"""The classic Adam optimiser.
Adam is an SGD variant with learning rate adaptation. The `learning_rate`
used for each weight is computed from estimates of first- and second-order
moments of the gradients (using suitable exponential moving averages).
References:
[Kingma et al, 2014](https://arxiv.org/abs/1412.6980)
Args:
learning_rate: this is a fixed global scaling factor.
b1: the exponential decay rate to track the first moment of past gradients.
b2: the exponential decay rate to track the second moment of past gradients.
eps: a small constant applied to denominator outside of the square root
(as in the Adam paper) to avoid dividing by zero when rescaling.
eps_root: (default `0`), a small constant applied to denominator inside the
square root (as in RMSProp), to avoid dividing by zero when rescaling.
This is needed for example when computing (meta-)gradients through Adam.
Returns:
the corresponding `GradientTransformation`.
"""
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
_scale_by_learning_rate(learning_rate),
)
def test_chain(self):
transformations = [
transform.scale_by_adam(),
transform.trace(decay=0, nesterov=False),
transform.scale(-LR)]
# Apply updates with chain.
chain_params = self.init_params
chained_transforms = combine.chain(*transformations)
state = chained_transforms.init(chain_params)
@self.variant
def update_fn(updates, state):
return chained_transforms.update(updates, state)
for _ in range(STEPS):
updates, state = update_fn(self.per_step_updates, state)
chain_params = update.apply_updates(chain_params, updates)
# Manually apply sequence of transformations.
manual_params = self.init_params
states = [t.init(manual_params) for t in transformations]
for _ in range(STEPS):
updates = self.per_step_updates
new_states = []
for t, s in zip(transformations, states):
updates, state = t.update(updates, s)
new_states.append(state)
manual_params = update.apply_updates(manual_params, updates)
states = new_states
# Check equivalence.
chex.assert_tree_all_close(manual_params, chain_params, rtol=1e-4)
def adam(learning_rate: float,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8) -> GradientTransformation:
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps),
transform.scale(-learning_rate),
)
def adam(learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8,
eps_root: float = 0.0) -> GradientTransformation:
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
_scale_by_learning_rate(learning_rate),
)
def adamw(learning_rate: float,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8,
weight_decay: float = 1e-4) -> GradientTransformation:
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps),
transform.additive_weight_decay(weight_decay),
transform.scale(-learning_rate),
)
def adamw(learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8,
eps_root: float = 0.0,
weight_decay: float = 1e-4) -> GradientTransformation:
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
transform.additive_weight_decay(weight_decay),
_scale_by_learning_rate(learning_rate),
)
def lamb(learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-6,
eps_root: float = 0.0,
weight_decay: float = 0.) -> base.GradientTransformation:
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
transform.add_decayed_weights(weight_decay),
transform.scale_by_trust_ratio(),
_scale_by_learning_rate(learning_rate),
)
def lamb(learning_rate: float,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-6,
eps_root: float = 0.0,
weight_decay: float = 0.) -> GradientTransformation:
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
transform.additive_weight_decay(weight_decay),
transform.scale_by_trust_ratio(),
transform.scale(-learning_rate),
)
def adamw(
learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8,
eps_root: float = 0.0,
mu_dtype: Optional[Any] = None,
weight_decay: float = 1e-4,
mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None,
) -> base.GradientTransformation:
"""Adam with weight decay regularization.
AdamW uses weight decay to regularise learning towards small weights, as
this leads to better generalisation. In SGD you can also use L2 regularisation
to implement this as an additive loss term, however L2 regularization
does not behave as intended for adaptive gradient algorithms such as Adam.
WARNING: Sometimes you may want to skip weight decay for BatchNorm scale or
for the bias parameters. You can use `optax.masked` to make your own AdamW
variant where `additive_weight_decay` is applied only to a subset of `params`.
References:
Loshchilov et al, 2019: https://arxiv.org/abs/1711.05101
Args:
learning_rate: this is a fixed global scaling factor.
b1: the exponential decay rate to track the first moment of past gradients.
b2: the exponential decay rate to track the second moment of past gradients.
eps: a small constant applied to denominator outside of the square root
(as in the Adam paper) to avoid dividing by zero when rescaling.
eps_root: (default `0`), a small constant applied to denominator inside the
square root (as in RMSProp), to avoid dividing by zero when rescaling.
This is needed for instance when computing (meta-)gradients through Adam.
mu_dtype: optional `dtype` to be used for the first order accumulator; if
`None` then the `dtype` is inferred from `params` and `updates`.
weight_decay: strength of the weight decay regularization.
mask: a tree with same structure as (or a prefix of) the params PyTree,
or a Callable that returns such a pytree given the params/updates.
The leaves should be booleans, `True` for leaves/subtrees you want to
apply the weight decay to, and `False` for those you want to skip. Note
that the Adam gradient transformations are applied to all parameters.
Returns:
the corresponding `GradientTransformation`.
"""
return combine.chain(
transform.scale_by_adam(
b1=b1, b2=b2, eps=eps, eps_root=eps_root, mu_dtype=mu_dtype),
transform.add_decayed_weights(weight_decay, mask),
_scale_by_learning_rate(learning_rate),
)
def lamb(
learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-6,
eps_root: float = 0.0,
weight_decay: float = 0.,
mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None,
) -> base.GradientTransformation:
"""The LAMB optimiser.
LAMB is a general purpose layer-wise adaptive large batch optimiser designed
to provide consistent training performance across a wide range of tasks,
including those that use attention-based models (such as Transformers) and
ResNet-50. The optimiser is able to work with small and large batch sizes.
LAMB was inspired by the LARS learning algorithm.
References:
You et al, 2019: https://arxiv.org/abs/1904.00962
Args:
learning_rate: this is a fixed global scaling factor.
b1: the exponential decay rate to track the first moment of past gradients.
b2: the exponential decay rate to track the second moment of past gradients.
eps: a small constant applied to denominator outside of the square root
(as in the Adam paper) to avoid dividing by zero when rescaling.
eps_root: (default `0.0`), a small constant applied to denominator inside
the square root (as in RMSProp), to avoid dividing by zero when rescaling.
This is needed for instance when computing (meta-)gradients through Adam.
weight_decay (default `0.`): strength of the weight decay regularization.
mask: a tree with same structure as (or a prefix of) the params PyTree,
or a Callable that returns such a pytree given the params/updates.
The leaves should be booleans, `True` for leaves/subtrees you want to
apply the transformation to, and `False` for those you want to skip.
Returns:
the corresponding `GradientTransformation`.
"""
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
transform.add_decayed_weights(weight_decay=weight_decay, mask=mask),
transform.scale_by_trust_ratio(),
_scale_by_learning_rate(learning_rate),
)
def adamw(
learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-8,
eps_root: float = 0.0,
weight_decay: float = 1e-4
) -> base.GradientTransformation:
"""Adam with weight decay regularization.
AdamW uses weight decay to regularise learning towards small weights, as
this leads to better generalisation. In SGD you can also use L2 regularisation
to implement this as an additive loss term, however L2 regularization
does not behave as intended for adaptive gradient algorithms such as Adam.
WARNING: Sometimes you may want to skip weight decay for BatchNorm scale or
for the bias parameters. You can use `optax.masked` to make your own AdamW
variant where `additive_weight_decay` is applied only to a subset of `params`.
References:
Loshchilov et al, 2019: https://arxiv.org/abs/1711.05101
Args:
learning_rate: this is a fixed global scaling factor.
b1: the exponential decay rate to track the first moment of past gradients.
b2: the exponential decay rate to track the second moment of past gradients.
eps: a small constant applied to denominator outside of the square root
(as in the Adam paper) to avoid dividing by zero when rescaling.
eps_root: (default `0`), a small constant applied to denominator inside the
square root (as in RMSProp), to avoid dividing by zero when rescaling.
This is needed for instance when computing (meta-)gradients through Adam.
weight_decay: strength of the weight decay regularization.
Returns:
the corresponding `GradientTransformation`.
"""
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
transform.additive_weight_decay(weight_decay),
_scale_by_learning_rate(learning_rate),
)
def lamb(learning_rate: ScalarOrSchedule,
b1: float = 0.9,
b2: float = 0.999,
eps: float = 1e-6,
eps_root: float = 0.0,
weight_decay: float = 0.) -> base.GradientTransformation:
"""The LAMB optimiser.
LAMB is a general purpose layer-wise adaptive large batch optimiser designed
to provide consistent training performance across a wide range of tasks,
including those that use attention-based models (such as Transformers) and
ResNet-50. The optimiser is able to work with small and large batch sizes.
LAMB was inspired by the LARS learning algorithm.
References:
[You et al, 2019](https://arxiv.org/abs/1904.00962)
Args:
learning_rate: this is a fixed global scaling factor.
b1: the exponential decay rate to track the first moment of past gradients.
b2: the exponential decay rate to track the second moment of past gradients.
eps: a small constant applied to denominator outside of the square root
(as in the Adam paper) to avoid dividing by zero when rescaling.
eps_root: (default `0.0`), a small constant applied to denominator inside
the square root (as in RMSProp), to avoid dividing by zero when rescaling.
This is needed for instance when computing (meta-)gradients through Adam.
weight_decay (default `0.`): strength of the weight decay regularization.
Returns:
the corresponding `GradientTransformation`.
"""
return combine.chain(
transform.scale_by_adam(b1=b1, b2=b2, eps=eps, eps_root=eps_root),
transform.add_decayed_weights(weight_decay),
transform.scale_by_trust_ratio(),
_scale_by_learning_rate(learning_rate),
)
