Training a Neural Network 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. The training step itself is GeometricOptimizers' too: its TrainingOptimizer, which GeometricMachineLearning re-exports, steps a whole parameter set with a gradient the caller computed, and applies the right update to each leaf — a retraction for a weight on a manifold, ordinary arithmetic for a Euclidean one. It is built on the parameters of a network, nn.params.
What GeometricMachineLearning adds is the part that is about neural networks: train!, which draws the minibatches of a data loader over epochs, computes the gradient of the loss on each with automatic differentiation or a symbolic pullback, and takes one step per minibatch.
The gradient comes from one minibatch at a time, so there is no objective function to hand a line search; the step size is a number or a GeometricOptimizers.DecayingStatic schedule:
opt = TrainingOptimizer(nn.params; algorithm = Adam(), linesearch = 1e-3)
opt = TrainingOptimizer(nn.params; AdamOptimizerWithDecay(n_steps)...)
history = train!(nn, opt, dl, Batch(batch_size), n_epochs)Library Functions
ReducedComplexityModeling.train! — Function
train!(nn, opt, dl, batch, n_epochs; loss, show_progress = true)Train the neural network nn with the GeometricOptimizers.TrainingOptimizer opt on the data of the DataLoader or ParametricDataLoader dl for n_epochs epochs, and return the loss history.
Each epoch draws the minibatches of dl with the Batch batch and takes one step per minibatch: loss computes the loss on the minibatch and its gradient, and GeometricOptimizers.optimization_step! writes the new parameters into params(nn), with the whole gradient, a zero leaf included. Entry i of the history is the mean of the minibatch losses of epoch i, each evaluated before its step. The history has the element type of the loss value, so a Float64 network on Float32 data gives a Vector{Float64}.
loss is a NetworkLoss of AbstractNeuralNetworks or an AbstractPullback; a NetworkLoss is trained through a ZygotePullback. Without it, a TransformerIntegrator with a Batch{:Transformer} trains on TransformerLoss(batch), a NeuralNetworkIntegrator with a Batch{:FeedForward} on FeedForwardLoss(), and an autoencoder with a Batch{:FeedForward} on AutoEncoderLoss(), and any network on a ParametricDataLoader on ParametricLoss(); any other combination needs loss, and train! raises an ArgumentError without it.
A network of HamiltonianNeuralNetworks.jl trains on vector-field data with loss = SymbolicPullback(arch) of that package, and on consecutive states with a loss such as its SymplecticEulerLoss(arch, timestep); see its manual.
opt is built for params(nn), with TrainingOptimizer(params(nn); algorithm, linesearch, retraction). It keeps the iteration number, the moments and the global sections, so a second call with the same opt resumes the training where the first one stopped; the only random numbers train! draws are those of the minibatch shuffle. n_epochs < 1 is an ArgumentError, raised before any step.