GitHub

Optimisation Rules

Any optimization rule from Optimisers.jl can be used with train! and other training functions.

For full details of how the interface works, see the Optimisers.jl documentation.

Optimisers Reference

All optimisers return an object that, when passed to train!, will update the parameters passed to it.

Optimisers.Descent — Type
Descent(η = 1f-1)
Descent(; [eta])

Classic gradient descent optimiser with learning rate η. For each parameter p and its gradient dp, this runs p -= η*dp.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
source
Optimisers.Momentum — Type
Momentum(η = 0.01, ρ = 0.9)
Momentum(; [eta, rho])

Gradient descent optimizer with learning rate η and momentum ρ.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Momentum (ρ == rho): Controls the acceleration of gradient descent in the prominent direction, in effect dampening oscillations.
source
Optimisers.Nesterov — Type
Nesterov(η = 0.001, ρ = 0.9)
Nesterov(; [eta, rho])

Gradient descent optimizer with learning rate η and Nesterov momentum ρ.

Parameters

  • Learning rate (η): Amount by which gradients are discounted before updating the weights.
  • Nesterov momentum (ρ): Controls the acceleration of gradient descent in the prominent direction, in effect dampening oscillations.
source
Optimisers.RMSProp — Type
RMSProp(η = 0.001, ρ = 0.9, ϵ = 1e-8; centred = false)
RMSProp(; [eta, rho, epsilon, centred])

Optimizer using the RMSProp algorithm. Often a good choice for recurrent networks. Parameters other than learning rate generally don't need tuning.

Centred RMSProp is a variant which normalises gradients by an estimate their variance, instead of their second moment.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Momentum (ρ == rho): Controls the acceleration of gradient descent in the prominent direction, in effect dampening oscillations.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
  • Keyword centred (or centered): Indicates whether to use centred variant of the algorithm.
source
Optimisers.Adam — Type
Adam(η = 0.001, β = (0.9, 0.999), ϵ = 1e-8)
Adam(; [eta, beta, epsilon])

Adam optimiser.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.RAdam — Type
RAdam(η = 0.001, β = (0.9, 0.999), ϵ = 1e-8)
RAdam(; [eta, beta, epsilon])

Rectified Adam optimizer.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.AdaMax — Type
AdaMax(η = 0.001, β = (0.9, 0.999), ϵ = 1e-8)
AdaMax(; [eta, beta, epsilon])

AdaMax is a variant of Adam based on the ∞-norm.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.AdaGrad — Type
AdaGrad(η = 0.1, ϵ = 1e-8)
AdaGrad(; [eta, epsilon])

AdaGrad optimizer. It has parameter specific learning rates based on how frequently it is updated. Parameters don't need tuning.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.AdaDelta — Type
AdaDelta(ρ = 0.9, ϵ = 1e-8)
AdaDelta(; [rho, epsilon])

AdaDelta is a version of AdaGrad adapting its learning rate based on a window of past gradient updates. Parameters don't need tuning.

Parameters

  • Rho (ρ == rho): Factor by which the gradient is decayed at each time step.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.AMSGrad — Type
AMSGrad(η = 0.001, β = (0.9, 0.999), ϵ = 1e-8)
AMSGrad(; [eta, beta, epsilon])

The AMSGrad version of the Adam optimiser. Parameters don't need tuning.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.NAdam — Type
NAdam(η = 0.001, β = (0.9, 0.999), ϵ = 1e-8)
NAdam(; [eta, beta, epsilon])

NAdam is a Nesterov variant of Adam. Parameters don't need tuning.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.AdamW — Type
AdamW(η = 0.001, β = (0.9, 0.999), λ = 0, ϵ = 1e-8; couple = true)
AdamW(; [eta, beta, lambda, epsilon, couple])

AdamW is a variant of Adam fixing (as in repairing) its weight decay regularization. Implemented as an OptimiserChain of Adam and WeightDecay`.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Weight decay (λ == lambda): Controls the strength of $L_2$ regularisation.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
  • Keyword couple: If true, the weight decay is coupled with the learning rate, as in pytorch's AdamW. This corresponds to an update of the form x = x - η * (dx + λ * x), where dx is the update from Adam with learning rate 1. If false, the weight decay is decoupled from the learning rate, in the spirit of the original paper. This corresponds to an update of the form x = x - η * dx - λ * x. Default is true.
Breaking change in v0.4

With version 0.4 the default update rule for AdamW has changed to match the pytorch implementation. The previous rule, which is closer to the original paper, can be obtained by setting AdamW(..., couple=false). See this issue for more details.

source
Optimisers.OAdam — Type
OAdam(η = 0.001, β = (0.5, 0.9), ϵ = 1e-8)
OAdam(; [eta, beta, epsilon])

OAdam (Optimistic Adam) is a variant of Adam adding an "optimistic" term suitable for adversarial training.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.AdaBelief — Type
AdaBelief(η = 0.001, β = (0.9, 0.999), ϵ = 1e-16)
AdaBelief(; [eta, beta, epsilon])

The AdaBelief optimiser is a variant of the well-known Adam optimiser.

Parameters

  • Learning rate (η == eta): Amount by which gradients are discounted before updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
  • Machine epsilon (ϵ == epsilon): Constant to prevent division by zero (no need to change default)
source
Optimisers.Lion — Type
Lion(η = 0.001, β = (0.9, 0.999))
Lion(; [eta, beta])

Lion optimiser.

Parameters

  • Learning rate (η == eta): Magnitude by which gradients are updating the weights.
  • Decay of momentums (β::Tuple == beta): Exponential decay for the first (β1) and the second (β2) momentum estimate.
source

Composing Optimisers

Flux (through Optimisers.jl) defines a special kind of optimiser called OptimiserChain which takes in arbitrary optimisers as input. Its behaviour is similar to the usual optimisers, but differs in that it acts by calling the optimisers listed in it sequentially. Each optimiser produces a modified gradient that will be fed into the next, and the resultant update will be applied to the parameter as usual. A classic use case is where adding decays is desirable. Optimisers.jl defines the basic decay corresponding to an $L_2$ regularization in the loss as WeightDecay.

opt = OptimiserChain(WeightDecay(1e-4), Descent())

Here we apply the weight decay to the Descent optimiser. The resulting optimiser opt can be used as any optimiser.

w = [randn(10, 10), randn(10, 10)]
opt_state = Flux.setup(opt, w)

loss(w, x) = Flux.mse(w[1] * x, w[2] * x)

loss(w, rand(10)) # around 0.9

for t = 1:10^5
  g = gradient(w -> loss(w[1], w[2], rand(10)), w)
  Flux.update!(opt_state, w, g)
end

loss(w, rand(10)) # around 0.9

It is possible to compose optimisers for some added flexibility.

Optimisers.OptimiserChain — Type
OptimiserChain(opts...)

Compose a sequence of optimisers so that each opt in opts updates the gradient, in the order specified.

With an empty sequence, OptimiserChain() is the identity, so update! will subtract the full gradient from the parameters. This is equivalent to Descent(1).

Example

julia> o = OptimiserChain(ClipGrad(1.0), Descent(0.1));

julia> m = (zeros(3),);

julia> s = Optimisers.setup(o, m)
(Leaf(OptimiserChain(ClipGrad(1.0), Descent(0.1)), (nothing, nothing)),)

julia> Optimisers.update(s, m, ([0.3, 1, 7],))[2]  # clips before discounting
([-0.03, -0.1, -0.1],)
source

Decays

Similar to optimisers, Flux also defines some simple decays that can be used in conjunction with other optimisers, or standalone.

Optimisers.SignDecay — Type
SignDecay(λ = 1e-3)
SignDecay(; [lambda])

Implements $L_1$ regularisation, also known as LASSO regression, when composed with other rules as the first transformation in an OptimiserChain.

It does this by adding λ .* sign(x) to the gradient. This is equivalent to adding λ * sum(abs, x) == λ * norm(x, 1) to the loss.

See also [WeightDecay] for $L_2$ normalisation. They can be used together: OptimiserChain(SignDecay(0.012), WeightDecay(0.034), Adam()) is equivalent to adding 0.012 * norm(x, 1) + 0.017 * norm(x, 2)^2 to the loss function.

Parameters

  • Penalty (λ ≥ 0): Controls the strength of the regularisation.
source
Optimisers.WeightDecay — Type
WeightDecay(λ = 5e-4)
WeightDecay(; [lambda])

Implements $L_2$ regularisation, also known as ridge regression, when composed with other rules as the first transformation in an OptimiserChain.

It does this by adding λ .* x to the gradient. This is equivalent to adding λ/2 * sum(abs2, x) == λ/2 * norm(x)^2 to the loss.

See also [SignDecay] for $L_1$ normalisation.

Parameters

  • Penalty (λ ≥ 0): Controls the strength of the regularisation.
source

Gradient Clipping

Gradient clipping is useful for training recurrent neural networks, which have a tendency to suffer from the exploding gradient problem. An example usage is

opt = OptimiserChain(ClipGrad(1e-3), Adam(1e-3))
Optimisers.ClipNorm — Type
ClipNorm(ω = 10, p = 2; throw = true)

Scales any gradient array for which norm(dx, p) > ω to stay at this threshold (unless p==0).

Throws an error if the norm is infinite or NaN, which you can turn off with throw = false.

Typically composed with other rules using OptimiserChain.

See also ClipGrad.

source