Optimisers
Consider a simple linear regression. We create some dummy data, calculate a loss, and backpropagate to calculate gradients for the parameters W
and b
.
using Flux
W = rand(2, 5)
b = rand(2)
predict(x) = (W * x) .+ b
loss(x, y) = sum((predict(x) .- y).^2)
x, y = rand(5), rand(2) # Dummy data
l = loss(x, y) # ~ 3
θ = Params([W, b])
grads = gradient(() -> loss(x, y), θ)
We want to update each parameter, using the gradient, in order to improve (reduce) the loss. Here's one way to do that:
using Flux.Optimise: update!
η = 0.1 # Learning Rate
for p in (W, b)
update!(p, -η * grads[p])
end
Running this will alter the parameters W
and b
and our loss should go down. Flux provides a more general way to do optimiser updates like this.
opt = Descent(0.1) # Gradient descent with learning rate 0.1
for p in (W, b)
update!(opt, p, grads[p])
end
An optimiser update!
accepts a parameter and a gradient, and updates the parameter according to the chosen rule. We can also pass opt
to our training loop, which will update all parameters of the model in a loop. However, we can now easily replace Descent
with a more advanced optimiser such as ADAM
.
Optimiser Reference
All optimisers return an object that, when passed to train!
, will update the parameters passed to it.
Flux.Optimise.update!
— Functionupdate!(opt, p, g)
update!(opt, ps::Params, gs)
Perform an update step of the parameters ps
(or the single parameter p
) according to optimizer opt
and the gradients gs
(the gradient g
).
As a result, the parameters are mutated and the optimizer's internal state may change.
update!(x, x̄)
Update the array x
according to x .-= x̄
.
Flux.Optimise.Descent
— TypeDescent(η)
Classic gradient descent optimiser with learning rate η
. For each parameter p
and its gradient δp
, this runs p -= η*δp
Parameters
- Learning Rate (η): The amount by which the gradients are discounted before updating the weights. Defaults to
0.1
.
Example
opt = Descent() # uses default η (0.1)
opt = Descent(0.3) # use provided η
ps = params(model)
gs = gradient(ps) do
loss(x, y)
end
Flux.Optimise.update!(opt, ps, gs)
Flux.Optimise.Momentum
— TypeMomentum(η, ρ)
Gradient descent with learning rate η
and momentum ρ
.
Parameters
- Learning Rate (
η
): Amount by which gradients are discounted before updating the weights. Defaults to0.01
. - Momentum (
ρ
): Parameter that accelerates descent in the relevant direction and dampens oscillations. Defaults to0.9
.
Examples
opt = Momentum() # uses defaults of η = 0.01 and ρ = 0.9
opt = Momentum(0.01, 0.99)
Flux.Optimise.Nesterov
— TypeNesterov(η, ρ)
Gradient descent with learning rate η
and Nesterov momentum ρ
.
Parameters
- Learning Rate (η): Amount by which the gradients are dicsounted berfore updating the weights. Defaults to
0.001
. - Nesterov Momentum (ρ): Parameters controlling the amount of nesterov momentum to be applied. Defaults to
0.9
.
Examples
opt = Nesterov() # uses defaults η = 0.001 and ρ = 0.9
opt = Nesterov(0.003, 0.95)
Flux.Optimise.RMSProp
— TypeRMSProp(η, ρ)
Implements the RMSProp algortihm. Often a good choice for recurrent networks. Parameters other than learning rate generally don't need tuning.
Parameters
- Learning Rate (η): Defaults to
0.001
. - Rho (ρ): Defaults to
0.9
.
Examples
opt = RMSProp() # uses default η = 0.001 and ρ = 0.9
opt = RMSProp(0.002, 0.95)
References
Flux.Optimise.ADAM
— TypeADAM(η, β::Tuple)
Implements the ADAM optimiser.
Paramters
- Learning Rate (
η
): Defaults to0.001
. - Beta (
β::Tuple
): The first element refers to β1 and the second to β2. Defaults to(0.9, 0.999)
.
Examples
opt = ADAM() # uses the default η = 0.001 and β = (0.9, 0.999)
opt = ADAM(0.001, (0.9, 0.8))
References
ADAM optimiser.
Flux.Optimise.AdaMax
— TypeAdaMax(η, β::Tuple)
Variant of ADAM based on ∞-norm.
Parameters
- Learning Rate (η): Defaults to
0.001
- Beta (β::Tuple): The first element refers to β1 and the second to β2. Defaults to
(0.9, 0.999)
.
Examples
opt = AdaMax() # uses default η and β
opt = AdaMax(0.001, (0.9, 0.995))
References
AdaMax optimiser.
Flux.Optimise.ADAGrad
— TypeADAGrad(η)
Implements AdaGrad. It has parameter specific learning rates based on how frequently it is updated.
Parameters
- Learning Rate (η): Defaults to
0.1
Examples
opt = ADAGrad() # uses default η = 0.1
opt = ADAGrad(0.001)
References
ADAGrad optimiser. Parameters don't need tuning.
Flux.Optimise.ADADelta
— TypeADADelta(ρ)
Version of ADAGrad that adapts learning rate based on a window of past gradient updates. Parameters don't need tuning.
Parameters
- Rho (ρ): Factor by which gradient is decayed at each time step. Defaults to
0.9
.
Examples
opt = ADADelta() # uses default ρ = 0.9
opt = ADADelta(0.89)
References
ADADelta optimiser.
Flux.Optimise.AMSGrad
— TypeAMSGrad(η, β::Tuple)
Implements AMSGrad version of the ADAM optimiser. Parameters don't need tuning.
Parameters
- Learning Rate (η): Defaults to
0.001
. - Beta (β::Tuple): The first element refers to β1 and the second to β2. Defaults to
(0.9, 0.999)
.
Examples
opt = AMSGrad() # uses default η and β
opt = AMSGrad(0.001, (0.89, 0.995))
References
AMSGrad optimiser.
Flux.Optimise.NADAM
— TypeNADAM(η, β::Tuple)
Nesterov variant of ADAM. Parameters don't need tuning.
Parameters
- Learning Rate (η): Defaults to
0.001
. - Beta (β::Tuple): The first element refers to β1 and the second to β2. Defaults to
(0.9, 0.999)
.
Examples
opt = NADAM() # uses default η and β
opt = NADAM(0.002, (0.89, 0.995))
References
NADAM optimiser.
Flux.Optimise.ADAMW
— FunctionADAMW(η, β::Tuple, decay)
Variant of ADAM defined by fixing weight decay regularization.
Parameters
- Learning Rate (η): Defaults to
0.001
. - Beta (β::Tuple): The first element refers to β1 and the second to β2. Defaults to (0.9, 0.999).
- decay: Decay applied to weights during optimisation. Defaults to 0.
Examples
opt = ADAMW() # uses default η, β and decay
opt = ADAMW(0.001, (0.89, 0.995), 0.1)
References
Optimiser Interface
Flux's optimisers are built around a struct
that holds all the optimiser parameters along with a definition of how to apply the update rule associated with it. We do this via the apply!
function which takes the optimiser as the first argument followed by the parameter and its corresponding gradient.
In this manner Flux also allows one to create custom optimisers to be used seamlessly. Let's work this with a simple example.
mutable struct Momentum
eta
rho
velocity
end
Momentum(eta::Real, rho::Real) = Momentum(eta, rho, IdDict())
The Momentum
type will act as our optimiser in this case. Notice that we have added all the parameters as fields, along with the velocity which we will use as our state dictionary. Each parameter in our models will get an entry in there. We can now define the rule applied when this optimiser is invoked.
function apply!(o::Momentum, x, Δ)
η, ρ = o.eta, o.rho
v = get!(o.velocity, x, zero(x))::typeof(x)
@. v = ρ * v - η * Δ
@. Δ = -v
end
This is the basic definition of a Momentum update rule given by:
The apply!
defines the update rules for an optimiser opt
, given the parameters and gradients. It returns the updated gradients. Here, every parameter x
is retrieved from the running state v
and subsequently updates the state of the optimiser.
Flux internally calls on this function via the update!
function. It shares the API with apply!
but ensures that multiple parameters are handled gracefully.
Composing Optimisers
Flux defines a special kind of optimiser simply called Optimiser
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. Flux defines some basic decays including ExpDecay
, InvDecay
etc.
opt = Optimiser(ExpDecay(0.001, 0.1, 1000, 1e-4), Descent())
Here we apply exponential decay to the Descent
optimiser. The defaults of ExpDecay
say that its learning rate will be decayed every 1000 steps. It is then applied like any optimiser.
w = randn(10, 10)
w1 = randn(10,10)
ps = Params([w, w1])
loss(x) = Flux.mse(w * x, w1 * x)
loss(rand(10)) # around 9
for t = 1:10^5
θ = Params([w, w1])
θ̄ = gradient(() -> loss(rand(10)), θ)
Flux.Optimise.update!(opt, θ, θ̄)
end
loss(rand(10)) # around 0.9
In this manner it is possible to compose optimisers for some added flexibility.
Decays
Similar to optimisers, Flux also defines some simple decays that can be used in conjunction with other optimisers, or standalone.
Flux.Optimise.ExpDecay
— TypeExpDecay(eta, decay, decay_step, clip)
Discount the learning rate eta
by a multiplicative factor decay
every decay_step
till a minimum of clip
.
Parameters
- Learning Rate (eta): Defaults to
0.001
. - decay: Factor by which the learning rate is discounted. Defaults to
0.1
. - decay_step: Schedules decay operations by setting number of steps between two decay operations. Defaults to
1000
. - clip: Minimum value of learning rate. Defaults to
1e-4
.
Example
To apply exponential decay to an optimiser:
Optimiser(ExpDecay(..), Opt(..))
opt = Optimiser(ExpDecay(), ADAM())
Flux.Optimise.InvDecay
— TypeInvDecay(γ)
Applies inverse time decay to an optimiser, i.e., the effective step size at iteration n
is eta / (1 + γ * n)
where eta
is the initial step size. The wrapped optimiser's step size is not modified.
Parameters
- gamma (γ): Defaults to
0.001
Example
Optimiser(InvDecay(..), Opt(..))
Flux.Optimise.WeightDecay
— TypeWeightDecay(wd)
Decays the weight by wd
Parameters
- weight decay (wd): 0