GitHub

The Optimizer in GeometricMachineLearning

The general framework for optimization on homogeneous spaces — the Riemannian gradient, the lift to the global tangent space $\mathfrak{g}^\mathrm{hor}$, the optimizer cache, the retraction and the global section — belongs to GeometricOptimizers and is described in its documentation, under Optimization on Homogeneous Spaces and Retractions.

What GeometricMachineLearning adds is the part that is about neural networks: walking the parameter tree of a NeuralNetwork, applying the right update to each leaf — a retraction for a weight on a manifold, ordinary arithmetic for a Euclidean one — and driving that from a data loader over epochs and batches.

The gradient comes from automatic differentiation on one batch at a time, so there is no objective function to hand a line search; the step size is a property of the Optimizer instead. It is either a number, or a GeometricOptimizers.DecayingStatic schedule:

opt = Optimizer(Adam(Float32), nn; step_size = 1e-3)
opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...)

Library Functions

GeometricMachineLearning.Optimizer — Type
Optimizer(method, nn; retraction, step_size)
Optimizer(nn; algorithm, linesearch, retraction)

Optimizer state combining a GeometricOptimizers method with the parameters of a neural network.

step_size is either a number — a fixed learning rate — or a GeometricOptimizers.DecayingStatic, a learning rate that decays geometrically with the iteration number. The second form is the one GeometricOptimizers uses itself, so a method paired with a schedule splats straight in:

opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...)

Extended help

The step size is a property of the optimizer and not of the method: the same Adam() trains at any learning rate. That is the split GeometricOptimizers makes — the method supplies a direction, a SimpleSolvers.LinesearchMethod supplies how far to go along it — and step_size is GML's half of it for a training loop, which has no objective function for a real line search to evaluate.

source
GeometricMachineLearning.optimize_for_one_epoch! — Function
optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY)

Sample the data contained in dl according to batch and optimize for these batches.

This step also performs automatic differentiation on loss.

The output of optimize_for_one_epoch! is the average loss over all batches of the epoch:

\[output = \frac{1}{\mathtt{steps\_per\_epoch}}\sum_{t=1}^\mathtt{steps\_per\_epoch}\mathtt{loss}(\theta^{(t-1)}).\]

This is done because any reverse differentiation routine always has two outputs; for Zygote:

loss_value, pullback = Zygote.pullback(ps -> loss(model, ps, input, output), ps)

So we get the value for the loss for free whenever we compute the pullback with AD.

Arguments

All the arguments are mandatory (there are no defaults):

  1. an instance of Optimizer.
  2. the neural network model.
  3. the neural network parameters ps.
  4. the data (i.e. an instance of DataLoader).
  5. batch::Batch: stores batch_size (and optionally seq_length and prediction_window).
  6. loss::NetworkLoss.
  7. the section λY of the parameters ps.

Implementation

Internally this calls optimization_step! for each minibatch.

The number of minibatches can be determined with number_of_batches:

using GeometricMachineLearning
using GeometricMachineLearning: number_of_batches

data = [1, 2, 3, 4, 5]
batch = Batch(2)
dl = DataLoader(data; suppress_info = true)

number_of_batches(dl, batch)

# output

3
source
GeometricMachineLearning.optimization_step! — Function
optimization_step!(opt, λY, ps, dp)

Apply one optimization step to the parameters ps and their gradient dp, with the method and step size the Optimizer opt carries.

λY is a GlobalSection of ps (or a NamedTuple of them). Note that it is an output here: the section the optimizer carries from step to step lives in opt.state, and λY is written so that callers who inspect it see the updated section. It therefore has to be allocated once and reused, not rebuilt per step – rebuilding it costs a QR decomposition per manifold weight.

The step counter is incremented before the step size is read, so the first step of a run is step 1. This matters for a decaying step_size and is how GeometricOptimizers counts too.

source