From 4d9af899b270f23d2379b034f5020031020b42e1 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Thu, 7 May 2026 16:32:18 +0200 Subject: [PATCH 1/7] Removed optimizers from GeometricMachineLearning and started importing them from GeometricOptimizers. --- Project.toml | 2 + docs/src/assets/custom.sty | 2 +- src/GeometricMachineLearning.jl | 717 +++++++++--------- src/loss/hnn_loss.jl | 18 +- src/loss/losses.jl | 86 ++- src/optimizers/adam_optimizer.jl | 56 -- ...adam_optimizer_with_learning_rate_decay.jl | 41 - src/optimizers/bfgs_cache.jl | 56 -- src/optimizers/bfgs_optimizer.jl | 67 -- src/optimizers/default_optimizer.jl | 3 - src/optimizers/gradient_optimizer.jl | 26 - src/optimizers/init_optimizer_cache.jl | 14 - .../manifold_related/global_sections.jl | 259 ------- .../manifold_related/modified_exponential.jl | 79 -- .../manifold_related/retraction_types.jl | 6 - .../manifold_related/retractions.jl | 188 ----- src/optimizers/momentum_optimizer.jl | 32 - src/optimizers/optimizer.jl | 177 ----- src/optimizers/optimizer_caches.jl | 120 --- src/optimizers/optimizer_method.jl | 15 - src/pullbacks/zygote_pullback.jl | 8 +- src/utils.jl | 74 +- test/runtests.jl | 273 +++++-- test/symplectic_autoencoder_tests.jl | 15 +- 24 files changed, 681 insertions(+), 1653 deletions(-) delete mode 100644 src/optimizers/adam_optimizer.jl delete mode 100644 src/optimizers/adam_optimizer_with_learning_rate_decay.jl delete mode 100644 src/optimizers/bfgs_cache.jl delete mode 100644 src/optimizers/bfgs_optimizer.jl delete mode 100644 src/optimizers/default_optimizer.jl delete mode 100644 src/optimizers/gradient_optimizer.jl delete mode 100644 src/optimizers/init_optimizer_cache.jl delete mode 100644 src/optimizers/manifold_related/global_sections.jl delete mode 100644 src/optimizers/manifold_related/modified_exponential.jl delete mode 100644 src/optimizers/manifold_related/retraction_types.jl delete mode 100644 src/optimizers/manifold_related/retractions.jl delete mode 100644 src/optimizers/momentum_optimizer.jl delete mode 100644 src/optimizers/optimizer.jl delete mode 100644 src/optimizers/optimizer_caches.jl delete mode 100644 src/optimizers/optimizer_method.jl diff --git a/Project.toml b/Project.toml index c226c75c5..9f882a6bd 100644 --- a/Project.toml +++ b/Project.toml @@ -13,6 +13,7 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GeometricBase = "9a0b12b7-583b-4f04-aa1f-d8551b6addc9" GeometricEquations = "c85262ba-a08a-430a-b926-d29770767bf2" +GeometricOptimizers = "fc236c15-5557-4942-aa65-b650f329279e" GeometricSolutions = "7843afe4-64f4-4df4-9231-049495c56661" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" @@ -42,6 +43,7 @@ Distances = "0.10" ForwardDiff = "0.10, 1" GeometricBase = "0.14" GeometricEquations = "0.21" +GeometricOptimizers = "1.0.0" GeometricSolutions = "0.6" HDF5 = "0.16, 0.17" KernelAbstractions = "0.9" diff --git a/docs/src/assets/custom.sty b/docs/src/assets/custom.sty index ba1673bd1..c2df8a857 100644 --- a/docs/src/assets/custom.sty +++ b/docs/src/assets/custom.sty @@ -125,7 +125,7 @@ % custom hyphenation \usepackage[htt]{hyphenat} -% renew texttt for better line breaks; taken from https://tex.stackexchange.com/questions/579789/combining-wrapping-texttt-with-sections-and-toc-improper-alphabetic-constan +% renew texttt for better line breaks; taken from https://tex.stackexchange.com/questions/579789/combining-wrapping-texttt-with-sections-and-toc-improper-alphabetic-constant \DeclareRobustCommand{\texttt}[1]{% \begingroup \ttfamily diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 7c7e6564f..64065803a 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -1,412 +1,399 @@ module GeometricMachineLearning - using AbstractNeuralNetworks - using BandedMatrices - using ChainRulesCore - using Distances - using GeometricBase - using GeometricSolutions: GeometricSolution, EnsembleSolution, DataSeries, StateVariable, TimeSeries - using GeometricEquations: EnsembleProblem, ODEProblem, HODEProblem, ODEEnsemble, HODEEnsemble - using KernelAbstractions - using LinearAlgebra - using NNlib - using ProgressMeter - using Random - using Zygote - using ForwardDiff - using InteractiveUtils - using TimerOutputs - import LazyArrays - import SymbolicNeuralNetworks - import SymbolicNeuralNetworks: input_dimension, output_dimension, SymbolicPullback - using SymbolicNeuralNetworks: derivative, _get_contents, _get_params, SymbolicNeuralNetwork - using Symbolics: @variables, substitute - - import AbstractNeuralNetworks: Architecture, Model, AbstractExplicitLayer, AbstractExplicitCell, AbstractNeuralNetwork , NeuralNetwork, UnknownArchitecture, FeedForwardLoss - import AbstractNeuralNetworks: Chain, GridCell - import AbstractNeuralNetworks: Dense, Linear, Recurrent - import AbstractNeuralNetworks: IdentityActivation, ZeroVector - import AbstractNeuralNetworks: add!, update! - import AbstractNeuralNetworks: layer - import AbstractNeuralNetworks: initialparameters - import AbstractNeuralNetworks: parameterlength - import AbstractNeuralNetworks: GlorotUniform - import AbstractNeuralNetworks: params, architecture, model, dim - import AbstractNeuralNetworks: AbstractPullback, NetworkLoss, _compute_loss - import AbstractNeuralNetworks: networkbackend - # export params, architetcure, model - export dim - import NNlib: σ, sigmoid, softmax - import Base: iterate, eltype - #import LogExpFunctions: softmax - - export CPU, GPU - export Chain, NeuralNetwork - export Dense, Linear - export initialparameters - export parameterlength - export NeuralNetworkParameters - - export σ, sigmoid, softmax - - # from GeometricBase to print docs - export description - - include("utils.jl") - - include("data_loader/data_loader.jl") - - # INCLUDE ARRAYS - include("arrays/skew_symmetric.jl") - include("arrays/symmetric.jl") - include("arrays/poisson_tensor.jl") - include("arrays/abstract_lie_algebra_horizontal.jl") - include("arrays/stiefel_lie_algebra_horizontal.jl") - include("arrays/grassmann_lie_algebra_horizontal.jl") - include("arrays/triangular.jl") - include("arrays/lower_triangular.jl") - include("arrays/upper_triangular.jl") - - export SymmetricMatrix, PoissonTensor, SkewSymMatrix - export StiefelLieAlgHorMatrix - export SymplecticLieAlgMatrix, SymplecticLieAlgHorMatrix - export GrassmannLieAlgHorMatrix - export StiefelProjection, SymplecticProjection - export LowerTriangular, UpperTriangular - - include("kernels/assign_q_and_p.jl") - include("kernels/tensor_mat_mul.jl") - include("kernels/tensor_tensor_mul.jl") - include("kernels/tensor_transpose_tensor_mul.jl") - include("kernels/tensor_tensor_transpose_mul.jl") - include("kernels/tensor_transpose_mat_mul.jl") - include("kernels/tensor_transpose_tensor_transpose_mul.jl") - include("kernels/mat_tensor_mul.jl") - include("kernels/tensor_transpose.jl") - include("kernels/exponentials/tensor_exponential.jl") - include("kernels/inverses/cpu_inverse.jl") - include("kernels/inverses/inverse_2x2.jl") - include("kernels/inverses/inverse_3x3.jl") - include("kernels/inverses/inverse_4x4.jl") - include("kernels/inverses/inverse_5x5.jl") - include("kernels/inverses/tensor_cayley.jl") - include("kernels/inverses/tensor_mat_skew_sym_assign.jl") - include("kernels/vec_tensor_mul.jl") - - include("kernels/kernel_ad_routines/assign_q_and_p.jl") - include("kernels/kernel_ad_routines/tensor_mat_mul.jl") - include("kernels/kernel_ad_routines/mat_tensor_mul.jl") - include("kernels/kernel_ad_routines/tensor_tensor_mul.jl") - include("kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl") - include("kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl") - include("kernels/kernel_ad_routines/tensor_transpose.jl") - include("kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl") - include("kernels/kernel_ad_routines/vec_tensor_mul.jl") - # export tensor_mat_mul - - # this defines empty retraction type structs (doesn't rely on anything) - include("optimizers/manifold_related/retraction_types.jl") - - export MatrixSoftmax, VectorSoftmax - include("activations/softmax.jl") - - # are these needed? - export UnknownProblem, NothingFunction - - # + operation has been overloaded to work with NamedTuples! - export _add, apply_toNT, split_and_flatten, add! - - # GPU specific operations - export convert_to_dev, Device, CPUDevice - - export GradientLayerQ, GradientLayerP, ActivationLayerQ, ActivationLayerP, LinearLayerQ, LinearLayerP - export Linear - export ResidualLayer - export LinearSymplecticLayerP, LinearSymplecticLayerQ - export SymplecticStiefelLayer - - include("manifolds/abstract_manifold.jl") - include("manifolds/stiefel_manifold.jl") - # include("manifolds/symplectic_stiefel_manifold.jl") - include("manifolds/grassmann_manifold.jl") - - include("arrays/stiefel_projection.jl") - - export StiefelManifold, SymplecticStiefelManifold, GrassmannManifold, Manifold - export rgrad, metric - - include("optimizers/optimizer_method.jl") - include("optimizers/optimizer_caches.jl") - include("optimizers/optimizer.jl") - include("optimizers/gradient_optimizer.jl") - include("optimizers/momentum_optimizer.jl") - include("optimizers/adam_optimizer.jl") - include("optimizers/adam_optimizer_with_learning_rate_decay.jl") - include("optimizers/bfgs_cache.jl") - include("optimizers/bfgs_optimizer.jl") - include("optimizers/init_optimizer_cache.jl") - - include("optimizers/manifold_related/global_sections.jl") - include("optimizers/manifold_related/modified_exponential.jl") - include("optimizers/manifold_related/retractions.jl") - - include("layers/sympnets.jl") - include("layers/bias_layer.jl") - include("layers/resnet.jl") - include("layers/manifold_layer.jl") - include("layers/stiefel_layer.jl") - include("layers/grassmann_layer.jl") - include("layers/multi_head_attention.jl") - include("layers/volume_preserving_attention.jl") - include("layers/volume_preserving_feedforward.jl") - include("layers/transformer.jl") - include("layers/psd_like_layer.jl") - include("layers/classification.jl") - - # include("layers/symplectic_stiefel_layer.jl") - export StiefelLayer, GrassmannLayer, ManifoldLayer - export PSDLayer - export MultiHeadAttention - export VolumePreservingAttention - export VolumePreservingFeedForwardLayer - export VolumePreservingLowerLayer - export VolumePreservingUpperLayer - export VolumePreservingTransformer - export NeuralNetworkIntegrator - export ResNet - export Transformer - export TransformerIntegrator, StandardTransformerIntegrator - - # INCLUDE OPTIMIZERS - export OptimizerMethod, AbstractCache - export GradientOptimizer, GradientCache - export MomentumOptimizer, MomentumCache - export AdamOptimizerWithDecay - export AdamOptimizer, AdamCache - export AdamOptimizerWithDecay - export BFGSOptimizer, BFGSCache - - export Optimizer - export optimization_step! - - export GlobalSection, apply_section, apply_section! - export global_rep - export Geodesic, Cayley - export geodesic, cayley - export retraction - # export ⊙², √ᵉˡᵉ, /ᵉˡᵉ, scalar_add - export update! - export check - - #INCLUDE ABSTRACT TRAINING integrator - export AbstractTrainingMethod - - export loss_single #, loss - - export HnnTrainingMethod - export LnnTrainingMethod - export SympNetTrainingMethod - - include("training_method/abstract_training_method.jl") - - # INCLUDE DATA TRAINING STRUCTURE - export AbstractDataShape, TrajectoryData, SampledData - export get_length_trajectory, get_Δt, get_nb_point, get_nb_trajectory, get_data - - include("data/data_shape.jl") - - export AbstractDataSymbol - export PositionSymbol, PhaseSpaceSymbol, DerivativePhaseSpaceSymbol, PosVeloAccSymbol, PosVeloSymbol - export DataSymbol - export can_reduce, symbols, symboldiff - - include("data/data_symbol.jl") - - # INCLUDE TRAINING INTEGRATOR - - export TrainingMethod - export symbol, shape - export min_length_batch - - - include("training_method/training_method.jl") - - # INCLUDE DATA TRAINING STRUCTURE - export AbstractTrainingData - export TrainingData - export shape, symbols, dim, noisemaker, data_symbols # , problem - export reduce_symbols, reshape_intoSampledData - export aresame - - include("data/data_training.jl") - - export get_batch, complete_batch_size, check_batch_size - - include("data/batch.jl") +using AbstractNeuralNetworks +using BandedMatrices +using ChainRulesCore +using Distances +using GeometricBase +using GeometricSolutions: GeometricSolution, EnsembleSolution, DataSeries, StateVariable, + TimeSeries +using GeometricEquations: EnsembleProblem, ODEProblem, HODEProblem, ODEEnsemble, + HODEEnsemble +using KernelAbstractions +using LinearAlgebra +using NNlib +using ProgressMeter +using Random +using Zygote +using ForwardDiff +using InteractiveUtils +using TimerOutputs +import LazyArrays +import SymbolicNeuralNetworks +import SymbolicNeuralNetworks: input_dimension, output_dimension, SymbolicPullback +using SymbolicNeuralNetworks: derivative, _get_contents, _get_params, SymbolicNeuralNetwork +using Symbolics: @variables, substitute + +using GeometricOptimizers +using GeometricOptimizers: OptimizerSolution, Geodesic, OptimizerMethod +export StiefelManifold, GrassmannManifold, rgrad + +using GeometricOptimizers: Adam +export Adam + +import AbstractNeuralNetworks: Architecture, Model, AbstractExplicitLayer, + AbstractExplicitCell, AbstractNeuralNetwork, NeuralNetwork, + UnknownArchitecture, FeedForwardLoss +import AbstractNeuralNetworks: Chain, GridCell +import AbstractNeuralNetworks: Dense, Linear, Recurrent +import AbstractNeuralNetworks: IdentityActivation, ZeroVector +import AbstractNeuralNetworks: add!, update! +import AbstractNeuralNetworks: layer +import AbstractNeuralNetworks: initialparameters +import AbstractNeuralNetworks: parameterlength +import AbstractNeuralNetworks: GlorotUniform +import AbstractNeuralNetworks: params, architecture, model, dim +import AbstractNeuralNetworks: AbstractPullback, NetworkLoss, _compute_loss +import AbstractNeuralNetworks: networkbackend +# export params, architetcure, model +export dim +import NNlib: σ, sigmoid, softmax +import Base: iterate, eltype +#import LogExpFunctions: softmax + +export CPU, GPU +export Chain, NeuralNetwork +export Dense, Linear +export initialparameters +export parameterlength +export NeuralNetworkParameters + +export σ, sigmoid, softmax + +# from GeometricBase to print docs +export description + +include("utils.jl") + +include("data_loader/data_loader.jl") + +# INCLUDE ARRAYS +include("arrays/skew_symmetric.jl") +include("arrays/symmetric.jl") +include("arrays/poisson_tensor.jl") +include("arrays/abstract_lie_algebra_horizontal.jl") +include("arrays/stiefel_lie_algebra_horizontal.jl") +include("arrays/grassmann_lie_algebra_horizontal.jl") +include("arrays/triangular.jl") +include("arrays/lower_triangular.jl") +include("arrays/upper_triangular.jl") + +export SymmetricMatrix, PoissonTensor, SkewSymMatrix +export StiefelLieAlgHorMatrix +export SymplecticLieAlgMatrix, SymplecticLieAlgHorMatrix +export GrassmannLieAlgHorMatrix +export StiefelProjection, SymplecticProjection +export LowerTriangular, UpperTriangular + +include("kernels/assign_q_and_p.jl") +include("kernels/tensor_mat_mul.jl") +include("kernels/tensor_tensor_mul.jl") +include("kernels/tensor_transpose_tensor_mul.jl") +include("kernels/tensor_tensor_transpose_mul.jl") +include("kernels/tensor_transpose_mat_mul.jl") +include("kernels/tensor_transpose_tensor_transpose_mul.jl") +include("kernels/mat_tensor_mul.jl") +include("kernels/tensor_transpose.jl") +include("kernels/exponentials/tensor_exponential.jl") +include("kernels/inverses/cpu_inverse.jl") +include("kernels/inverses/inverse_2x2.jl") +include("kernels/inverses/inverse_3x3.jl") +include("kernels/inverses/inverse_4x4.jl") +include("kernels/inverses/inverse_5x5.jl") +include("kernels/inverses/tensor_cayley.jl") +include("kernels/inverses/tensor_mat_skew_sym_assign.jl") +include("kernels/vec_tensor_mul.jl") + +include("kernels/kernel_ad_routines/assign_q_and_p.jl") +include("kernels/kernel_ad_routines/tensor_mat_mul.jl") +include("kernels/kernel_ad_routines/mat_tensor_mul.jl") +include("kernels/kernel_ad_routines/tensor_tensor_mul.jl") +include("kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl") +include("kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl") +include("kernels/kernel_ad_routines/tensor_transpose.jl") +include("kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl") +include("kernels/kernel_ad_routines/vec_tensor_mul.jl") +# export tensor_mat_mul + +export MatrixSoftmax, VectorSoftmax +include("activations/softmax.jl") + +# are these needed? +export UnknownProblem, NothingFunction + +# + operation has been overloaded to work with NamedTuples! +export _add, apply_toNT, split_and_flatten, add! + +# GPU specific operations +export convert_to_dev, Device, CPUDevice + +export GradientLayerQ, GradientLayerP, ActivationLayerQ, ActivationLayerP, LinearLayerQ, + LinearLayerP +export Linear +export ResidualLayer +export LinearSymplecticLayerP, LinearSymplecticLayerQ +export SymplecticStiefelLayer + +include("manifolds/abstract_manifold.jl") +include("manifolds/stiefel_manifold.jl") +# include("manifolds/symplectic_stiefel_manifold.jl") +include("manifolds/grassmann_manifold.jl") + +include("arrays/stiefel_projection.jl") + +include("layers/sympnets.jl") +include("layers/bias_layer.jl") +include("layers/resnet.jl") +include("layers/manifold_layer.jl") +include("layers/stiefel_layer.jl") +include("layers/grassmann_layer.jl") +include("layers/multi_head_attention.jl") +include("layers/volume_preserving_attention.jl") +include("layers/volume_preserving_feedforward.jl") +include("layers/transformer.jl") +include("layers/psd_like_layer.jl") +include("layers/classification.jl") + +# include("layers/symplectic_stiefel_layer.jl") +export StiefelLayer, GrassmannLayer, ManifoldLayer +export PSDLayer +export MultiHeadAttention +export VolumePreservingAttention +export VolumePreservingFeedForwardLayer +export VolumePreservingLowerLayer +export VolumePreservingUpperLayer +export VolumePreservingTransformer +export NeuralNetworkIntegrator +export ResNet +export Transformer +export TransformerIntegrator, StandardTransformerIntegrator + +# INCLUDE OPTIMIZERS +export GradientOptimizer, GradientCache +export MomentumOptimizer, MomentumCache +export AdamOptimizerWithDecay +export AdamOptimizer, AdamCache +export AdamOptimizerWithDecay +export BFGSOptimizer, BFGSCache + +export Optimizer +export optimization_step! + +export GlobalSection, apply_section, apply_section! +export global_rep +export Geodesic, Cayley +export geodesic, cayley +export retraction +# export ⊙², √ᵉˡᵉ, /ᵉˡᵉ, scalar_add +export update! +export check + +#INCLUDE ABSTRACT TRAINING integrator +export AbstractTrainingMethod + +export loss_single #, loss + +export HnnTrainingMethod +export LnnTrainingMethod +export SympNetTrainingMethod + +include("training_method/abstract_training_method.jl") + +# INCLUDE DATA TRAINING STRUCTURE +export AbstractDataShape, TrajectoryData, SampledData +export get_length_trajectory, get_Δt, get_nb_point, get_nb_trajectory, get_data + +include("data/data_shape.jl") + +export AbstractDataSymbol +export PositionSymbol, PhaseSpaceSymbol, DerivativePhaseSpaceSymbol, PosVeloAccSymbol, + PosVeloSymbol +export DataSymbol +export can_reduce, symbols, symboldiff + +include("data/data_symbol.jl") + +# INCLUDE TRAINING INTEGRATOR + +export TrainingMethod +export symbol, shape +export min_length_batch + +include("training_method/training_method.jl") + +# INCLUDE DATA TRAINING STRUCTURE +export AbstractTrainingData +export TrainingData +export shape, symbols, dim, noisemaker, data_symbols # , problem +export reduce_symbols, reshape_intoSampledData +export aresame + +include("data/data_training.jl") + +export get_batch, complete_batch_size, check_batch_size + +include("data/batch.jl") + +# INCLUDE BACKENDS +export LuxBackend +export NeuralNetwork +export arch + +include("backends/backends.jl") +include("backends/lux.jl") + +export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss + +#INCLUDE ARCHITECTURES +include("architectures/neural_network_integrator.jl") +include("architectures/resnet.jl") +include("architectures/transformer_integrator.jl") +include("architectures/standard_transformer_integrator.jl") +include("architectures/sympnet.jl") +include("architectures/autoencoder.jl") +include("architectures/symplectic_autoencoder.jl") +include("architectures/psd.jl") +include("architectures/fixed_width_network.jl") +include("architectures/hamiltonian_neural_network.jl") +include("architectures/lagrangian_neural_network.jl") +include("architectures/variable_width_network.jl") +include("architectures/recurrent_neural_network.jl") +include("architectures/LSTM_neural_network.jl") +include("architectures/transformer_neural_network.jl") +include("architectures/volume_preserving_feedforward.jl") +include("architectures/volume_preserving_transformer.jl") - # INCLUDE BACKENDS - export LuxBackend - export NeuralNetwork - export arch +export HamiltonianArchitecture +export LagrangianNeuralNetwork +export SympNet, LASympNet, GSympNet +export RecurrentNeuralNetwork +export LSTMNeuralNetwork +export ClassificationTransformer, ClassificationLayer +export VolumePreservingFeedForward +export SymplecticAutoencoder, PSDArch +export HamiltonianArchitecture, StandardHamiltonianArchitecture, + GeneralizedHamiltonianArchitecture - include("backends/backends.jl") - include("backends/lux.jl") +export solve!, encoder, decoder - export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss +export train!, apply!, jacobian! +export iterate - #INCLUDE ARCHITECTURES - include("architectures/neural_network_integrator.jl") - include("architectures/resnet.jl") - include("architectures/transformer_integrator.jl") - include("architectures/standard_transformer_integrator.jl") - include("architectures/sympnet.jl") - include("architectures/autoencoder.jl") - include("architectures/symplectic_autoencoder.jl") - include("architectures/psd.jl") - include("architectures/fixed_width_network.jl") - include("architectures/hamiltonian_neural_network.jl") - include("architectures/lagrangian_neural_network.jl") - include("architectures/variable_width_network.jl") - include("architectures/recurrent_neural_network.jl") - include("architectures/LSTM_neural_network.jl") - include("architectures/transformer_neural_network.jl") - include("architectures/volume_preserving_feedforward.jl") - include("architectures/volume_preserving_transformer.jl") +export default_arch - export HamiltonianArchitecture - export LagrangianNeuralNetwork - export SympNet, LASympNet, GSympNet - export RecurrentNeuralNetwork - export LSTMNeuralNetwork - export ClassificationTransformer, ClassificationLayer - export VolumePreservingFeedForward - export SymplecticAutoencoder, PSDArch - export HamiltonianArchitecture, StandardHamiltonianArchitecture, GeneralizedHamiltonianArchitecture +include("architectures/default_architecture.jl") - export solve!, encoder, decoder +include("loss/losses.jl") +include("loss/hnn_loss.jl") - export train!, apply!, jacobian! - export iterate +export AbstractPullback, ZygotePullback, SymbolicPullback +include("pullbacks/zygote_pullback.jl") +include("pullbacks/symbolic_hnn_pullback.jl") - export default_arch +export DataLoader, onehotbatch +export Batch, optimize_for_one_epoch! +include("data_loader/tensor_assign.jl") +include("data_loader/matrix_assign.jl") +include("data_loader/mnist_utils.jl") +include("data_loader/batch.jl") +include("data_loader/optimize.jl") - include("architectures/default_architecture.jl") +# INCLUDE TRAINING parameters - include("loss/losses.jl") - include("loss/hnn_loss.jl") +export TrainingParameters - export AbstractPullback, ZygotePullback, SymbolicPullback - include("pullbacks/zygote_pullback.jl") - include("pullbacks/symbolic_hnn_pullback.jl") +include("training/training_parameters.jl") - export DataLoader, onehotbatch - export Batch, optimize_for_one_epoch! - include("data_loader/tensor_assign.jl") - include("data_loader/matrix_assign.jl") - include("data_loader/mnist_utils.jl") - include("data_loader/batch.jl") - include("data_loader/optimize.jl") +# INCLUDE NEURALNET SOLUTION - export default_optimizer +export SingleHistory +export parameters, datashape +export History +export last, sizemax, nbtraining, show - include("optimizers/default_optimizer.jl") +include("nnsolution/history.jl") - # INCLUDE TRAINING parameters +export NeuralNetSolution +export problem, timestep, history, size_history +export set_sizemax_history - export TrainingParameters +include("nnsolution/neural_net_solution.jl") - include("training/training_parameters.jl") +export EnsembleNeuralNetSolution +export push!, merge! - # INCLUDE NEURALNET SOLUTION +include("nnsolution/neural_net_solution_ensemble.jl") - export SingleHistory - export parameters, datashape - export History - export last, sizemax, nbtraining, show +# INCLUDE TRAINING integrator - include("nnsolution/history.jl") +export TrainingSet +export parameters # , data - export NeuralNetSolution - export problem, timestep, history, size_history - export set_sizemax_history +include("training/training_set.jl") - include("nnsolution/neural_net_solution.jl") +export EnsembleTraining +export isnnShared, isParametersShared, isDataShared +export parameters, data +export push!, merge!, size - export EnsembleNeuralNetSolution - export push!, merge! +include("training/ensemble_training.jl") - include("nnsolution/neural_net_solution_ensemble.jl") +include("training/nn_parameters_transformation.jl") - # INCLUDE TRAINING integrator +export loss_gradient +export train! - export TrainingSet - export parameters # , data +include("training/train.jl") - include("training/training_set.jl") +export SymplecticEuler +export SymplecticEulerA, SymplecticEulerB +export SEuler, SEulerA, SEulerB - export EnsembleTraining - export isnnShared, isParametersShared, isDataShared - export parameters, data - export push!, merge!, size +include("training_method/symplectic_euler.jl") - include("training/ensemble_training.jl") +export HnnExactMethod +export ExactHnn - include("training/nn_parameters_transformation.jl") +include("training_method/hnn_exact_method.jl") - export loss_gradient - export train! +export VariationalMethod +export VariationalMidPointMethod +export VariaMidPoint - include("training/train.jl") +include("training_method/variational_method.jl") - export SymplecticEuler - export SymplecticEulerA, SymplecticEulerB - export SEuler, SEulerA, SEulerB +export LnnExactMethod +export ExactLnn - include("training_method/symplectic_euler.jl") +include("training_method/lnn_exact_method.jl") - export HnnExactMethod - export ExactHnn +export BasicSympNetMethod +export BasicSympNet - include("training_method/hnn_exact_method.jl") +include("training_method/sympnet_basic_method.jl") - export VariationalMethod - export VariationalMidPointMethod - export VariaMidPoint +export default_method - include("training_method/variational_method.jl") +include("training/default_method.jl") - export LnnExactMethod - export ExactLnn +# INCLUDE ASSERTION Function +export matching +include("training/matching.jl") - include("training_method/lnn_exact_method.jl") +include("reduced_system/reduced_system.jl") - export BasicSympNetMethod - export BasicSympNet +export HRedSys, reduction_error, projection_error, integrate_reduced_system, + integrate_full_system - include("training_method/sympnet_basic_method.jl") +include("layers/linear_symplectic_attention.jl") +include("layers/symplectic_attention.jl") +include("architectures/linear_symplectic_transformer.jl") +include("architectures/symplectic_transformer.jl") - export default_method +export LinearSymplecticAttention, LinearSymplecticAttentionQ, LinearSymplecticAttentionP +export LinearSymplecticTransformer +export SymplecticAttention, SymplecticAttentionQ, SymplecticAttentionP +export SymplecticTransformer - include("training/default_method.jl") - - - # INCLUDE ASSERTION Function - export matching - include("training/matching.jl") - - include("reduced_system/reduced_system.jl") - - export HRedSys, reduction_error, projection_error, integrate_reduced_system, integrate_full_system - - include("layers/linear_symplectic_attention.jl") - include("layers/symplectic_attention.jl") - include("architectures/linear_symplectic_transformer.jl") - include("architectures/symplectic_transformer.jl") - - export LinearSymplecticAttention, LinearSymplecticAttentionQ, LinearSymplecticAttentionP - export LinearSymplecticTransformer - export SymplecticAttention, SymplecticAttentionQ, SymplecticAttentionP - export SymplecticTransformer - - include("map_to_cpu.jl") +include("map_to_cpu.jl") end diff --git a/src/loss/hnn_loss.jl b/src/loss/hnn_loss.jl index 3ffd800e5..397e0602d 100644 --- a/src/loss/hnn_loss.jl +++ b/src/loss/hnn_loss.jl @@ -26,15 +26,17 @@ function HNNLoss(arch::HamiltonianArchitecture) HNNLoss(hamiltonian_vector_field(arch)) end -function (loss::HNNLoss)( ::Union{Chain, AbstractExplicitLayer}, - ps::Union{NeuralNetworkParameters, NamedTuple}, - input::QPTOAT, - output::QPTOAT) +AbstractNeuralNetworks.NetworkLoss(arch::HamiltonianArchitecture) = HNNLoss(arch) + +function (loss::HNNLoss)(::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, + input::QPTOAT, + output::QPTOAT) loss(ps, input, output) end -function (loss::HNNLoss)( ps::Union{NeuralNetworkParameters, NamedTuple}, - input::QPTOAT, - output::QPTOAT) +function (loss::HNNLoss)(ps::Union{NeuralNetworkParameters, NamedTuple}, + input::QPTOAT, + output::QPTOAT) norm(loss.hvf(input, ps) - output) / norm(output) -end \ No newline at end of file +end diff --git a/src/loss/losses.jl b/src/loss/losses.jl index 93764f99a..375787e75 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -1,7 +1,10 @@ +# type piracy! This should go into `AbstractNeuralNetworks`! ... +AbstractNeuralNetworks.NetworkLoss(nn::NeuralNetwork) = NetworkLoss(architecture(nn)) + @doc raw""" TransformerLoss(seq_length, prediction_window) -Make an instance of the transformer loss. +Make an instance of the transformer loss. This should be used together with a neural network of type [`TransformerIntegrator`](@ref). @@ -9,7 +12,7 @@ This should be used together with a neural network of type [`TransformerIntegrat `TransformerLoss` applies a neural network to an input and compares it to the `output` via an ``L_2`` norm: -```jldoctest +```jldoctest using GeometricMachineLearning using LinearAlgebra: norm import Random @@ -40,7 +43,7 @@ So `TransformerLoss` simply does: ```math \mathtt{loss}(\mathcal{NN}, \mathtt{input}, \mathtt{output}) = || \mathcal{NN}(\mathtt{input})[(\mathtt{sl} - \mathtt{pw} + 1):\mathtt{end}] - \mathtt{output} || / || \mathtt{output} ||, ``` -where ``||\cdot||`` is the ``L_2`` norm. +where ``||\cdot||`` is the ``L_2`` norm. # Parameters @@ -52,14 +55,24 @@ struct TransformerLoss <: NetworkLoss prediction_window::Int end +function AbstractNeuralNetworks.NetworkLoss( + arch::TransformerIntegrator, prediction_window::Integer = 1) + TransformerLoss(arch.seq_length, prediction_window) +end + TransformerLoss(seq_length::Int) = TransformerLoss(seq_length, seq_length) -# This crops the output array of the neural network so that it conforms with the output it should be compared to. This is needed for the transformer loss. -function crop_array_for_transformer_loss(nn_output::AT, output::BT) where {T, T2, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T2, 3}} - @view nn_output[axes(output, 1), axes(output, 2) .+ size(nn_output, 2) .- size(output, 2), axes(output, 3)] +# This crops the output array of the neural network so that it conforms with the output it should be compared to. This is needed for the transformer loss. +function crop_array_for_transformer_loss(nn_output::AT, + output::BT) where {T, T2, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T2, 3}} + @view nn_output[ + axes(output, 1), axes(output, 2) .+ size(nn_output, 2) .- size(output, 2), + axes(output, 3)] end -function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, output::AT) where {T, AT <: AbstractArray{T, 3}} +function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, + output::AT) where {T, AT <: AbstractArray{T, 3}} input_dim, input_seq_length = size(input) output_dim, output_prediction_window = size(output) @assert input_dim == output_dim @@ -67,39 +80,46 @@ function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps: @assert output_prediction_window == loss.prediction_window predicted_output_uncropped = model(input, ps) - predicted_output_cropped = crop_array_for_transformer_loss(predicted_output_uncropped, output) + predicted_output_cropped = crop_array_for_transformer_loss( + predicted_output_uncropped, output) _compute_loss(predicted_output_cropped, output) end -function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, output::AT) where {T, AT <: AbstractArray{T, 2}} +function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, + output::AT) where {T, AT <: AbstractArray{T, 2}} loss(model, ps, reshape(input, size(input)..., 1), reshape(output, size(output)..., 1)) end -function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::T, output::T) where {T <: QPT} +function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, + input::T, output::T) where {T <: QPT} loss(model, ps, vcat(input.q, input.p), vcat(output.q, output.p)) end # @doc raw""" # ClassificationTransformerLoss() -# +# # Make an instance of `ClassificationTransformerLoss`. -# +# # This is to be used together with a [`ClassificationTransformer`](@ref). -# +# # It takes an input, parses it to the transformer and then crops it to conform with the desired output size. -# +# # Suppose the input is of dimension ``\mathtt{td}\times\mathtt{sl}``, where `td` is *transformer dimension* and `sl` is *sequence length*. -# The output of the transformer will again be of the same dimension: -# +# The output of the transformer will again be of the same dimension: +# # ```math # \mathrm{output}\in\mathbb{R}^{\mathtt{td}\times\mathtt{sl}}. # ``` -# +# # if the output dimension `cl` of the [`ClassificationLayer`](@ref) is differnt form `td`. # """ struct ClassificationTransformerLoss <: NetworkLoss end -function (loss::ClassificationTransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::AbstractArray, output::AbstractArray) +function (loss::ClassificationTransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, + input::AbstractArray, output::AbstractArray) predicted_output_uncropped = model(input, ps) # predicted_output_cropped = crop_array_for_transformer_loss(predicted_output_uncropped, output) norm(predicted_output_uncropped - output) / norm(output) @@ -110,13 +130,13 @@ end Make an instance of `AutoEncoderLoss`. -This loss should always be used together with a neural network of type [`AutoEncoder`](@ref) (and it is also the default for training such a network). +This loss should always be used together with a neural network of type [`AutoEncoder`](@ref) (and it is also the default for training such a network). # Example `AutoEncoderLoss` applies a neural network to an input and compares it to the `output` via an ``L_2`` norm: -```jldoctest +```jldoctest using GeometricMachineLearning using LinearAlgebra: norm import Random @@ -142,23 +162,29 @@ So `AutoEncoderLoss` simply does: ```math \mathtt{loss}(\mathcal{NN}, \mathtt{input}) = || \mathcal{NN}(\mathtt{input}) - \mathtt{input} || / || \mathtt{input} ||, ``` -where ``||\cdot||`` is the ``L_2`` norm. +where ``||\cdot||`` is the ``L_2`` norm. # Parameters This loss does not have any parameters. """ -struct AutoEncoderLoss <: NetworkLoss end +struct AutoEncoderLoss <: NetworkLoss end + +AbstractNeuralNetworks.NetworkLoss(::AutoEncoder) = AutoEncoderLoss() function (loss::AutoEncoderLoss)(nn::NeuralNetwork, input::QPTOAT) loss(nn.model, params(nn), input, input) end -function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT) +function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT) loss(model, ps, input, input) end -(loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT, output::QPTOAT) = FeedForwardLoss()(model, ps, input, output) +function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT, output::QPTOAT) + FeedForwardLoss()(model, ps, input, output) +end @doc raw""" ReducedLoss(encoder, decoder) @@ -171,7 +197,7 @@ This loss should be used together with a [`NeuralNetworkIntegrator`](@ref) or [` `ReducedLoss` applies the *encoder*, *integrator* and *decoder* neural networks in this order to an input and compares it to the `output` via an ``L_2`` norm: -```jldoctest +```jldoctest using GeometricMachineLearning using LinearAlgebra: norm import Random @@ -196,7 +222,7 @@ loss(transformer, input_mat, output_mat) ≈ norm(output_mat - output_prediction true ``` -So the loss computes: +So the loss computes: ```math \mathrm{loss}_{\mathcal{E}, \mathcal{D}}(\mathcal{NN}, \mathrm{input}, \mathrm{output}) = ||\mathcal{D}(\mathcal{NN}(\mathcal{E}(\mathrm{input}))) - \mathrm{output}||, @@ -204,7 +230,8 @@ So the loss computes: where ``\mathcal{E}`` is the [`Encoder`](@ref), ``\mathcal{D}`` is the [`Decoder`](@ref). ``\mathcal{NN}`` is the neural network we compute the loss of. """ -struct ReducedLoss{ET <: NeuralNetwork{<:Encoder}, DT <: NeuralNetwork{<:Decoder}} <: NetworkLoss +struct ReducedLoss{ET <: NeuralNetwork{<:Encoder}, DT <: NeuralNetwork{<:Decoder}} <: + NetworkLoss encoder::ET decoder::DT end @@ -224,6 +251,7 @@ function ReducedLoss(autoencoder::NeuralNetwork{<:AutoEncoder}) ReducedLoss(encoder(autoencoder), decoder(autoencoder)) end -function (loss::ReducedLoss)(model::Chain, params::NeuralNetworkParameters, input::CT, output::CT) where {CT <: QPTOAT} +function (loss::ReducedLoss)(model::Chain, params::NeuralNetworkParameters, + input::CT, output::CT) where {CT <: QPTOAT} _compute_loss(loss.decoder(model(loss.encoder(input), params)), output) -end \ No newline at end of file +end diff --git a/src/optimizers/adam_optimizer.jl b/src/optimizers/adam_optimizer.jl deleted file mode 100644 index 27781769f..000000000 --- a/src/optimizers/adam_optimizer.jl +++ /dev/null @@ -1,56 +0,0 @@ -@doc raw""" - AdamOptimizer(η, ρ₁, ρ₂, δ) - -Make an instance of the Adam Optimizer. - -Here the cache consists of first and second moments that are updated as - -```math -B_1 \gets ((\rho_1 - \rho_1^t)/(1 - \rho_1^t))\cdot{}B_1 + (1 - \rho_1)/(1 - \rho_1^t)\cdot{}\nabla{}L, -``` -and - -```math -B_2 \gets ((\rho_2 - \rho_1^t)/(1 - \rho_2^t))\cdot{}B_2 + (1 - \rho_2)/(1 - \rho_2^t)\cdot\nabla{}L\odot\nabla{}L. -``` -The final velocity is computed as: - -```math -\mathrm{velocity} \gets -\eta{}B_1/\sqrt{B_2 + \delta}. -``` - -# Implementation - -The *velocity* is stored in the input to save memory: - -```julia -mul!(B, -o.method.η, /ᵉˡᵉ(C.B₁, scalar_add(racᵉˡᵉ(C.B₂), o.method.δ))) -``` -where `B` is the input to the [`update!`] function. - -The algorithm and suggested defaults are taken from [goodfellow2016deep; page 301](@cite). -""" -struct AdamOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - ρ₁::T - ρ₂::T - δ::T - - AdamOptimizer(η = 1f-3, ρ₁ = 9f-1, ρ₂ = 9.9f-1, δ = 3f-7; T=typeof(η)) = new{T}(T(η), T(ρ₁), T(ρ₂), T(δ)) -end - -function AdamOptimizer(T::Type) - AdamOptimizer(T(1f-3)) -end - -function update!(o::Optimizer{<:AdamOptimizer{T}}, C::AdamCache, B::AbstractArray) where T - add!(C.B₁, ((o.method.ρ₁ - o.method.ρ₁^o.step)/(T(1.) - o.method.ρ₁^o.step))*C.B₁, ((T(1.) - o.method.ρ₁)/(T(1.) - o.method.ρ₁^o.step))*B) - add!(C.B₂, ((o.method.ρ₂ - o.method.ρ₂^o.step)/(T(1.) - o.method.ρ₂^o.step))*C.B₂, ((T(1.) - o.method.ρ₂)/(T(1.) - o.method.ρ₂^o.step))*⊙²(B)) - mul!(B, -o.method.η, /ᵉˡᵉ(C.B₁, scalar_add(racᵉˡᵉ(C.B₂), o.method.δ))) -end - -# defaults: -⊙²(A::AbstractVecOrMat) = A.^2 -racᵉˡᵉ(A::AbstractVecOrMat) = sqrt.(A) -/ᵉˡᵉ(A::AbstractVecOrMat, B::AbstractVecOrMat) = A./B -scalar_add(A::AbstractVecOrMat, δ::Real) = A .+ δ \ No newline at end of file diff --git a/src/optimizers/adam_optimizer_with_learning_rate_decay.jl b/src/optimizers/adam_optimizer_with_learning_rate_decay.jl deleted file mode 100644 index f2faefea6..000000000 --- a/src/optimizers/adam_optimizer_with_learning_rate_decay.jl +++ /dev/null @@ -1,41 +0,0 @@ -@doc raw""" - AdamOptimizerWithDecay(n_epochs, η₁=1f-2, η₂=1f-6, ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8) - -Make an instance of the Adam Optimizer with weight decay. - -All except the first argument (the number of epochs) have defaults. - -The difference to the standard [`AdamOptimizer`](@ref) is that we change the learning reate ``\eta`` in each step. -Apart from the *time dependency* of ``\eta`` the two algorithms are however equivalent. -``\eta(0)`` starts with a high value ``\eta_1`` and then exponentially decrease until it reaches ``\eta_2`` with - -```math - \eta(t) = \gamma^t\eta_1, -``` -where ``\gamma = \exp(\log(\eta_1 / \eta_2) / \mathtt{n\_epochs}).`` -""" -struct AdamOptimizerWithDecay{T<:Real} <: OptimizerMethod{T} - η₁::T - η₂::T - ρ₁::T - ρ₂::T - δ::T - γ::T - n_epochs::Int - - function AdamOptimizerWithDecay(n_epochs::Int, η₁=1f-2, η₂=1f-6, ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8; T=typeof(η₁)) - γ = exp(log(η₂ / η₁) / n_epochs) - new{T}(T(η₁), T(η₂), T(ρ₁), T(ρ₂), T(δ), T(γ), n_epochs) - end -end - -function AdamOptimizerWithDecay(n_epochs::Int, T::Type; η₁=1f-2, η₂=1f-6, ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8) - AdamOptimizerWithDecay(n_epochs, T(η₁), T(η₂), T(ρ₁), T(ρ₂), T(δ)) -end - -function update!(o::Optimizer{<:AdamOptimizerWithDecay{T}}, C::AdamCache, B::AbstractArray) where T - η = o.method.γ ^ o.step * o.method.η₁ - add!(C.B₁, ((o.method.ρ₁ - o.method.ρ₁^o.step) / (T(1.) - o.method.ρ₁^o.step)) * C.B₁, ((T(1.) - o.method.ρ₁) / (T(1.) - o.method.ρ₁ ^ o.step)) * B) - add!(C.B₂, ((o.method.ρ₂ - o.method.ρ₂^o.step) / (T(1.) - o.method.ρ₂^o.step)) * C.B₂, ((T(1.) - o.method.ρ₂) / (T(1.) - o.method.ρ₂ ^ o.step)) * ⊙²(B)) - mul!(B, -η, /ᵉˡᵉ(C.B₁, scalar_add(racᵉˡᵉ(C.B₂), o.method.δ))) -end \ No newline at end of file diff --git a/src/optimizers/bfgs_cache.jl b/src/optimizers/bfgs_cache.jl deleted file mode 100644 index 1bbcf26b4..000000000 --- a/src/optimizers/bfgs_cache.jl +++ /dev/null @@ -1,56 +0,0 @@ -@doc raw""" - BFGSCache(B) - -Make the cache for the BFGS optimizer based on the array `B`. - -It stores an array for the gradient of the previous time step `B` and the inverse of the Hessian matrix `H`. - -The cache for the inverse of the Hessian is initialized with the idendity. -The cache for the previous gradient information is initialized with the zero vector. - -Note that the cache for `H` is changed iteratively, whereas the cache for `B` is newly assigned at every time step. -""" -struct BFGSCache{T, BT<:AbstractArray{T}, HT<:AbstractMatrix{T}} <: AbstractCache{T} - B::BT - H::HT - function BFGSCache(B::AbstractArray) - zeroB = zero(B) - H_init = initialize_hessian_inverse(zeroB) - new{eltype(B), typeof(zeroB), typeof(H_init)}(zero(B), H_init) - end -end - -@kernel function assign_diagonal_ones_kernel!(B::AbstractMatrix{T}) where T - i = @index(Global) - B[i, i] = one(T) -end - -function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::BFGSCache) - show(io, raw"`BFGSCache` that currently stores `B`as ...") - show(io, "text/plain", C.B) - println(io, "") - println(io, "... and `H` as") - show(io, "text/plain", C.H) -end - -# @doc raw""" -# initialize_hessian_inverse(B) -# -# Initialize the inverse of the Hessian for various arrays. -# -# # Implementation -# This requires an implementation of a *vectorization operation* `vec`. This is important for custom arrays. -# """ -function initialize_hessian_inverse(B::AbstractArray{T}) where T - length_of_array = length(vec(B)) - backend = networkbackend(B) - H = KernelAbstractions.zeros(backend, T, length_of_array, length_of_array) - assign_diagonal_ones! = assign_diagonal_ones_kernel!(backend) - assign_diagonal_ones!(H, ndrange=length_of_array) - H -end - -setup_bfgs_cache(ps::NeuralNetworkParameters) = setup_bfgs_cache(params(ps)) -setup_bfgs_cache(ps::NamedTuple) = apply_toNT(setup_bfgs_cache, ps) -setup_bfgs_cache(ps::Tuple) = Tuple([setup_bfgs_cache(x) for x in ps]) -setup_bfgs_cache(B::AbstractArray) = BFGSCache(B) \ No newline at end of file diff --git a/src/optimizers/bfgs_optimizer.jl b/src/optimizers/bfgs_optimizer.jl deleted file mode 100644 index d90f7d740..000000000 --- a/src/optimizers/bfgs_optimizer.jl +++ /dev/null @@ -1,67 +0,0 @@ -@doc raw""" - BFGSOptimizer(η, δ) - -Make an instance of the Broyden-Fletcher-Goldfarb-Shanno (BFGS) optimizer. - -`η` is the *learning rate*. -`δ` is a stabilization parameter. -""" -struct BFGSOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - δ::T - - function BFGSOptimizer(η::T = 1f-2, δ=1f-8) where T - new{T}(η, T(δ)) - end -end - -@doc raw""" - update!(o::Optimizer{<:BFGSOptimizer}, C, B) - -Peform an update with the BFGS optimizer. - -`C` is the cache, `B` contains the gradient information (the output of [`global_rep`](@ref) in general). - -First we compute the *final velocity* with -```julia -vecS = -o.method.η * C.H * vec(B) -``` -and then we update `H` -```julia -C.H .= (𝕀 - ρ * SY) * C.H * (𝕀 - ρ * SY') + ρ * vecS * vecS' -``` -where `SY` is `vecS * Y'` and `𝕀` is the idendity. - -# Implementation - -For stability we use `δ` for computing `ρ`: -```julia -ρ = 1. / (vecS' * Y + o.method.δ) -``` - -This is similar to the [`AdamOptimizer`](@ref) - -# Extended help - -If we have weights on a [`Manifold`](@ref) than the updates are slightly more difficult. -In this case the [`vec`](@ref) operation has to be generalized to the corresponding *global tangent space*. -""" -function update!(o::Optimizer{<:BFGSOptimizer}, C::BFGSCache, B::AbstractArray) - T = eltype(o) - # in the first step we compute the difference between the current and the previous mapped gradients: - Y = vec(B - C.B) - # compute the descent direction - P = -C.H * vec(B) - # compute S - vecS = o.method.η * P - # store gradient - assign!(C.B, copy(B)) - # output final velocity - assign!(vec(B), copy(vecS)) - # compute SY and HY - ρ = one(T) / (vecS' * Y + o.method.δ) - SY = vecS * Y' - 𝕀 = one(SY) - # compute H - C.H .= (𝕀 - ρ * SY) * C.H * (𝕀 - ρ * SY') + ρ * vecS * vecS' -end \ No newline at end of file diff --git a/src/optimizers/default_optimizer.jl b/src/optimizers/default_optimizer.jl deleted file mode 100644 index 622e05fbe..000000000 --- a/src/optimizers/default_optimizer.jl +++ /dev/null @@ -1,3 +0,0 @@ -default_optimizer() = GradientOptimizer() - -default_optimizer(arch::Architecture) = GradientOptimizer() \ No newline at end of file diff --git a/src/optimizers/gradient_optimizer.jl b/src/optimizers/gradient_optimizer.jl deleted file mode 100644 index 46c68365b..000000000 --- a/src/optimizers/gradient_optimizer.jl +++ /dev/null @@ -1,26 +0,0 @@ -@doc raw""" - GradientOptimizer(η) - -Make an instance of a gradient optimizer. - -This is the simplest neural network optimizer. It has no cache and computes the final velocity as: -```math - \mathrm{velocity} \gets - \eta\nabla_\mathrm{weight}L. -``` - -# Implementation - -The operations are done as memory efficiently as possible. -This means the provided ``\nabla_WL`` is mutated via: -```julia -rmul!(∇L, -method.η) -``` -""" -struct GradientOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - GradientOptimizer(η = 1e-2) = new{typeof(η)}(η) -end - -function update!(o::Optimizer{<:GradientOptimizer}, ::GradientCache, B::AbstractVecOrMat) - rmul!(B, -o.method.η) -end \ No newline at end of file diff --git a/src/optimizers/init_optimizer_cache.jl b/src/optimizers/init_optimizer_cache.jl deleted file mode 100644 index c12a6a3f9..000000000 --- a/src/optimizers/init_optimizer_cache.jl +++ /dev/null @@ -1,14 +0,0 @@ -# @doc raw""" -# init_optimizer_cache(method, x) -# -# Initialize the cache corresponding to the weights `x` for a specific method. -# -# # Implementation -# -# Wrapper for the functions `setup_adam_cache`, `setup_momentum_cache`, `setup_gradient_cache`, `setup_bfgs_cache`. -# These appear outside of `optimizer_caches.jl` because the `OptimizerMethods` first have to be defined. -# """ -init_optimizer_cache(::GradientOptimizer, x) = setup_gradient_cache(x) -init_optimizer_cache(::MomentumOptimizer, x) = setup_momentum_cache(x) -init_optimizer_cache(::Union{AdamOptimizer, AdamOptimizerWithDecay}, x) = setup_adam_cache(x) -init_optimizer_cache(::BFGSOptimizer, x) = setup_bfgs_cache(x) \ No newline at end of file diff --git a/src/optimizers/manifold_related/global_sections.jl b/src/optimizers/manifold_related/global_sections.jl deleted file mode 100644 index a3ed00d7d..000000000 --- a/src/optimizers/manifold_related/global_sections.jl +++ /dev/null @@ -1,259 +0,0 @@ -@doc raw""" - GlobalSection(Y) - -Construct a global section for `Y`. - -A global section ``\lambda`` is a mapping from a homogeneous space ``\mathcal{M}`` to the corresponding Lie group ``G`` such that - -```math -\lambda(Y)E = Y, -``` - -Also see [`apply_section`](@ref) and [`global_rep`](@ref). - -# Implementation - -For an implementation of `GlobalSection` for a custom array (especially manifolds), the function [`global_section`](@ref) has to be generalized. -""" -struct GlobalSection{T, AT, λT} - Y::AT - # for now the only lift that is implemented is the Stiefel one - these types will have to be expanded! - λ::λT - - function GlobalSection(Y::AbstractVecOrMat) - λ = global_section(Y) - new{eltype(Y), typeof(Y), typeof(λ)}(Y, λ) - end -end - -function GlobalSection(ps::NamedTuple) - apply_toNT(GlobalSection, ps) -end - -function GlobalSection(ps::NeuralNetworkParameters) - setup_cache(GlobalSection, ps) -end - -@doc raw""" - Matrix(λY::GlobalSection) - -Put `λY` into matrix form. - -This is not recommended if speed is important! - -Use [`apply_section`](@ref) and [`global_rep`](@ref) instead! -""" -function Base.Matrix(λY::GlobalSection) - hcat(Matrix(λY.Y), Matrix(λY.λ)) -end - -@doc raw""" - λY * Y - -Apply the element `λY` onto `Y`. - -Here `λY` is an element of a Lie group and `Y` is an element of a homogeneous space. -""" -Base.:*(λY::GlobalSection, Y::Manifold) = apply_section(λY, Y) - -@doc raw""" - apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT <: StiefelManifold{T}} - -Apply `λY` to `Y₂`. - -Mathematically this is the group action of the element ``\lambda{}Y\in{}G`` on the element ``Y_2`` of the homogeneous space ``\mathcal{M}``. - -Internally it calls [`apply_section!`](@ref). -""" -function apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}} - Y = StiefelManifold(zero(Y₂.A)) - apply_section!(Y, λY, Y₂) - - Y -end - -@doc raw""" - apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}} - -Apply `λY` to `Y₂` and store the result in `Y`. - -This is the inplace version of [`apply_section`](@ref). -""" -function apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}} - N, n = size(λY.Y) - - @views Y.A .= λY.Y * Y₂.A[1:n, :] + λY.λ * Y₂.A[(n+1):N, :] -end - -function apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:GrassmannManifold{T}} - Y = GrassmannManifold(zero(Y₂.A)) - apply_section!(Y, λY, Y₂) - - Y -end - -function apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:GrassmannManifold{T}} - N, n = size(λY.Y) - - @views Y.A = λY.Y * Y₂.A[1:n, :] + λY.λ * Y₂.A[(n + 1):N, :] -end - -function apply_section(λY::GlobalSection{T}, Y₂::AbstractVecOrMat{T}) where {T} - Y = copy(Y₂) - apply_section!(Y, λY, Y₂) -end - -function apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AbstractVecOrMat{T}) where {T, AT<:AbstractVecOrMat{T}} - Y .= Y₂ + λY.Y -end - -function apply_section(λY::NamedTuple, Y₂::NamedTuple) - apply_toNT(apply_section, λY, Y₂) -end - -function apply_section!(Y::NamedTuple, λY::NamedTuple, Y₂::NamedTuple) - apply_toNT(apply_section!, Y, λY, Y₂) -end - -function global_rep(λY::NamedTuple, gx::NamedTuple) - apply_toNT(global_rep, λY, gx) -end - -##auxiliary function -function global_rep(::GlobalSection{T}, gx::AbstractVecOrMat{T}) where {T} - gx -end - -@doc raw""" - global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:StiefelManifold{T}} - -Express `Δ` (an the tangent space of `Y`) as an instance of `StiefelLieAlgHorMatrix`. - -This maps an element from ``T_Y\mathcal{M}`` to an element of ``\mathfrak{g}^\mathrm{hor}``. - -These two spaces are isomorphic where the isomorphism where the isomorphism is established through ``\lambda(Y)\in{}G`` via: - -```math -T_Y\mathcal{M} \to \mathfrak{g}^{\mathrm{hor}}, \Delta \mapsto \lambda(Y)^{-1}\Omega(Y, \Delta)\lambda(Y). -``` - -Also see [`GeometricMachineLearning.Ω`](@ref). - -# Examples - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: _round -import Random - -Random.seed!(123) - -Y = rand(StiefelManifold, 6, 3) -Δ = rgrad(Y, randn(6, 3)) -λY = GlobalSection(Y) - -_round(global_rep(λY, Δ); digits = 3) - -# output - -6×6 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}: - 0.0 0.679 1.925 0.981 -2.058 0.4 - -0.679 0.0 0.298 -0.424 0.733 -0.919 - -1.925 -0.298 0.0 -1.815 1.409 1.085 - -0.981 0.424 1.815 0.0 0.0 0.0 - 2.058 -0.733 -1.409 0.0 0.0 0.0 - -0.4 0.919 -1.085 0.0 0.0 0.0 -``` - -# Implementation - -The function `global_rep` does in fact not perform the entire map ``\lambda(Y)^{-1}\Omega(Y, \Delta)\lambda(Y)`` but only - -```math -\Delta \mapsto \mathrm{skew}(Y^T\Delta), -``` - -to get the small skew-symmetric matrix ``A\in\mathcal{S}_\mathrm{skew}(n)`` and - -```math -\Delta \mapsto (\lambda(Y)_{[1:N, n:N]}^T \Delta)_{[1:(N-n), 1:n]}, -``` - -to get the arbitrary matrix ``B\in\mathbb{R}^{(N-n)\times{}n}``. -""" -function global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:StiefelManifold{T}} - N, n = size(λY.Y) - StiefelLieAlgHorMatrix( - SkewSymMatrix(λY.Y.A' * Δ), - λY.λ' * Δ, - N, - n - ) -end - -@doc raw""" - global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:GrassmannManifold{T}} - -Express `Δ` (an element of the tangent space of `Y`) as an instance of [`GrassmannLieAlgHorMatrix`](@ref). - -The method `global_rep` for [`GrassmannManifold`](@ref) is similar to that for [`StiefelManifold`](@ref). - -# Examples - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: _round -import Random - -Random.seed!(123) - -Y = rand(GrassmannManifold, 6, 3) -Δ = rgrad(Y, randn(6, 3)) -λY = GlobalSection(Y) - -_round(global_rep(λY, Δ); digits = 3) - -# output - -6×6 GrassmannLieAlgHorMatrix{Float64, Matrix{Float64}}: - 0.0 0.0 0.0 0.981 -2.058 0.4 - 0.0 0.0 0.0 -0.424 0.733 -0.919 - 0.0 0.0 0.0 -1.815 1.409 1.085 - -0.981 0.424 1.815 0.0 0.0 0.0 - 2.058 -0.733 -1.409 0.0 0.0 0.0 - -0.4 0.919 -1.085 0.0 0.0 0.0 -``` -""" -function global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:GrassmannManifold{T}} - N, n = size(λY.Y) - GrassmannLieAlgHorMatrix( - λY.λ' * Δ, - N, - n - ) -end - -function update_section!(Λ⁽ᵗ⁻¹⁾::GlobalSection{T, MT}, B⁽ᵗ⁻¹⁾::AbstractLieAlgHorMatrix{T}, retraction) where {T, MT <: Manifold} - N, n = B⁽ᵗ⁻¹⁾.N, B⁽ᵗ⁻¹⁾.n - expB = retraction(B⁽ᵗ⁻¹⁾) - apply_section!(expB, Λ⁽ᵗ⁻¹⁾, expB) - Λ⁽ᵗ⁻¹⁾.Y.A .= @view expB.A[:, 1:n] - Λ⁽ᵗ⁻¹⁾.λ .= @view expB.A[:, (n+1):N] - - nothing -end - -function update_section!(Λ⁽ᵗ⁻¹⁾::GlobalSection{T, AT}, B⁽ᵗ⁻¹⁾::AT, retraction) where {T, AT <: AbstractVecOrMat{T}} - expB = retraction(B⁽ᵗ⁻¹⁾) - apply_section!(expB, Λ⁽ᵗ⁻¹⁾, expB) - Λ⁽ᵗ⁻¹⁾.Y .= expB - - nothing -end - -function update_section!(Λ⁽ᵗ⁻¹⁾::NamedTuple, B⁽ᵗ⁻¹⁾::NamedTuple, retraction) - update_section_closure!(Λ⁽ᵗ⁻¹⁾, B⁽ᵗ⁻¹⁾) = update_section!(Λ⁽ᵗ⁻¹⁾, B⁽ᵗ⁻¹⁾, retraction) - apply_toNT(update_section_closure!, Λ⁽ᵗ⁻¹⁾, B⁽ᵗ⁻¹⁾) - - nothing -end \ No newline at end of file diff --git a/src/optimizers/manifold_related/modified_exponential.jl b/src/optimizers/manifold_related/modified_exponential.jl deleted file mode 100644 index 13167391f..000000000 --- a/src/optimizers/manifold_related/modified_exponential.jl +++ /dev/null @@ -1,79 +0,0 @@ -update_algorithm = "while norm(Aⁿ) > ε -mul!(A_temp, Aⁿ, A) -Aⁿ .= A_temp -rmul!(Aⁿ, T(inv(n))) - -𝔄A += Aⁿ -n += 1 -end" - -@doc (raw""" - 𝔄(A) - -Compute ``\mathfrak{A}(A) := \sum_{n=1}^\infty \frac{1}{n!} (A)^{n-1}.`` - -# Implementation - -This uses a Taylor expansion that iteratively adds terms with - -```julia -""" * update_algorithm * raw""" - -``` - -until the norm of `Aⁿ` becomes smaller than machine precision. -The counter `n` in the above algorithm is initialized as `2` -The matrices `Aⁿ` and `𝔄` are initialized as the identity matrix. -""") -function 𝔄(A::AbstractMatrix) - T = eltype(A) - Aⁿ = one(A) - 𝔄A = one(A) - A_temp = zero(A) - n = 2 - ε = eps(T) - while norm(Aⁿ) > ε - mul!(A_temp, Aⁿ, A) - Aⁿ .= A_temp - rmul!(Aⁿ, T(inv(n))) - - 𝔄A += Aⁿ - n += 1 - end - # println("Number of iterations is: ", i) - 𝔄A -end - -@doc raw""" - 𝔄(B̂, B̄) - -Compute ``\mathfrak{A}(B', B'') := \sum_{n=1}^\infty \frac{1}{n!} ((B'')^TB')^{n-1}.`` - -This expression has the property ``\mathbb{I} + B'\mathfrak{A}(B', B'')(B'')^T = \exp(B'(B'')^T).`` - -# Examples - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: 𝔄 -import Random -Random.seed!(123) - -B = rand(StiefelLieAlgHorMatrix, 10, 2) -B̂ = hcat(vcat(.5 * B.A, B.B), vcat(one(B.A), zero(B.B))) -B̄ = hcat(vcat(one(B.A), zero(B.B)), vcat(-.5 * B.A, -B.B)) - -one(B̂ * B̄') + B̂ * 𝔄(B̂, B̄) * B̄' ≈ exp(Matrix(B)) - -# output - -true -``` -""" -function 𝔄(B̂::AbstractMatrix, B̄::AbstractMatrix) - 𝔄(B̄' * B̂) -end - -function 𝔄exp(X::AbstractMatrix{T}, Y::AbstractMatrix{T}) where T - I + X * 𝔄(X, Y) * Y' -end \ No newline at end of file diff --git a/src/optimizers/manifold_related/retraction_types.jl b/src/optimizers/manifold_related/retraction_types.jl deleted file mode 100644 index e38c10bd5..000000000 --- a/src/optimizers/manifold_related/retraction_types.jl +++ /dev/null @@ -1,6 +0,0 @@ -# AbstractRetraction is a type that comprises all retraction methods for manifolds. For every manifold layer one has to specify a retraction method that takes the layer and elements of the (global) tangent space. -abstract type AbstractRetraction end - -struct Cayley <: AbstractRetraction end - -struct Geodesic <: AbstractRetraction end \ No newline at end of file diff --git a/src/optimizers/manifold_related/retractions.jl b/src/optimizers/manifold_related/retractions.jl deleted file mode 100644 index 3dab15aad..000000000 --- a/src/optimizers/manifold_related/retractions.jl +++ /dev/null @@ -1,188 +0,0 @@ -geodesic(A::AbstractVecOrMat) = A -cayley(A::AbstractVecOrMat) = A - -geodesic(B::NamedTuple) = apply_toNT(geodesic, B) - -@doc raw""" - geodesic(Y::Manifold, Δ) - -Take as input an element of a manifold `Y` and a tangent vector in `Δ` in the corresponding tangent space and compute the geodesic (exponential map). - -In different notation: take as input an element ``x`` of ``\mathcal{M}`` and an element of ``T_x\mathcal{M}`` and return ``\mathtt{geodesic}(x, v_x) = \exp(v_x).`` - - -# Examples - -```jldoctest -using GeometricMachineLearning - -Y = StiefelManifold([1. 0. 0.;]' |> Matrix) -Δ = [0. .5 0.;]' |> Matrix -Y₂ = geodesic(Y, Δ) - -Y₂' * Y₂ ≈ [1.;] - -# output - -true -``` - -# Implementation - -Internally this `geodesic` method calls [`geodesic(::StiefelLieAlgHorMatrix)`](@ref). -""" -function geodesic(Y::Manifold{T}, Δ::AbstractMatrix{T}) where T - λY = GlobalSection(Y) - B = global_rep(λY, Δ) - E = StiefelProjection(B) - expB = geodesic(B) - λY * typeof(Y)(expB * E) -end - -@doc raw""" - geodesic(B̄::StiefelLieAlgHorMatrix) - -Compute the geodesic of an element in [`StiefelLieAlgHorMatrix`](@ref). - -# Implementation - -Internally this is using: - -```math -\mathbb{I} + B'\mathfrak{A}(B', B'')B'', -``` - -with - -```math -\bar{B} = \begin{bmatrix} - A & -B^T \\ - B & \mathbb{O} -\end{bmatrix} = \begin{bmatrix} \frac{1}{2}A & \mathbb{I} \\ B & \mathbb{O} \end{bmatrix} \begin{bmatrix} \mathbb{I} & \mathbb{O} \\ \frac{1}{2}A & -B^T \end{bmatrix} =: B'(B'')^T. -``` - -This is using a computationally efficient version of the matrix exponential ``\mathfrak{A}``. - -See [`GeometricMachineLearning.𝔄`](@ref). -""" -function geodesic(B::StiefelLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - unit = one(B.A) - A_mat = B.A * unit - B̂ = hcat(vcat(T(.5) * A_mat, B.B), E) - B̄ = hcat(vcat(unit, T(.5) * A_mat), vcat(zero(B.B'), -B.B'))' - StiefelManifold(one(B) + B̂ * 𝔄(B̂, B̄) * B̄') -end - -@doc raw""" - geodesic(B̄::GrassmannLieAlgHorMatrix) - -Compute the geodesic of an element in [`GrassmannLieAlgHorMatrix`](@ref). - -This is equivalent to the method of [`geodesic`](@ref) for [StiefelLieAlgHorMatrix](@ref). - -See [`geodesic(::StiefelLieAlgHorMatrix)`](@ref). -""" -function geodesic(B::GrassmannLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - backend = networkbackend(B) - zero_mat = KernelAbstractions.zeros(backend, T, B.n, B.n) - B̂ = hcat(vcat(zero_mat, B.B), E) - B̄ = hcat(vcat(one(zero_mat), zero_mat), vcat(zero(B.B'), -B.B'))' - GrassmannManifold(one(B) + B̂ * 𝔄(B̂, B̄) * B̄') -end - -cayley(B::NamedTuple) = apply_toNT(cayley, B) - -@doc raw""" - cayley(Y::Manifold, Δ) - -Take as input an element of a manifold `Y` and a tangent vector in `Δ` in the corresponding tangent space and compute the Cayley retraction. - -In different notation: take as input an element ``x`` of ``\mathcal{M}`` and an element of ``T_x\mathcal{M}`` and return ``\mathrm{Cayley}(v_x).`` - -# Examples - -```jldoctest -using GeometricMachineLearning - -Y = StiefelManifold([1. 0. 0.;]' |> Matrix) -Δ = [0. .5 0.;]' |> Matrix -Y₂ = cayley(Y, Δ) - -Y₂' * Y₂ ≈ [1.;] - -# output - -true -``` - -See the example in [`geodesic(::Manifold{T}, ::AbstractMatrix{T}) where T`]. -""" -function cayley(Y::Manifold{T}, Δ::AbstractMatrix{T}) where T - λY = GlobalSection(Y) - B = global_rep(λY, Δ) - E = StiefelProjection(B) - cayleyB = cayley(B) - λY * typeof(Y)(cayleyB * E) -end - -@doc raw""" - cayley(B̄::StiefelLieAlgHorMatrix) - -Compute the Cayley retraction of `B`. - -# Implementation - -Internally this is using - -```math -\mathrm{Cayley}(\bar{B}) = \mathbb{I} + \frac{1}{2} B' (\mathbb{I}_{2n} - \frac{1}{2} (B'')^T B')^{-1} (B'')^T (\mathbb{I} + \frac{1}{2} B), -``` -with -```math -\bar{B} = \begin{bmatrix} - A & -B^T \\ - B & \mathbb{O} -\end{bmatrix} = \begin{bmatrix} \frac{1}{2}A & \mathbb{I} \\ B & \mathbb{O} \end{bmatrix} \begin{bmatrix} \mathbb{I} & \mathbb{O} \\ \frac{1}{2}A & -B^T \end{bmatrix} =: B'(B'')^T, -``` -i.e. ``\bar{B}`` is expressed as a product of two ``N\times{}2n`` matrices. -""" -function cayley(B::StiefelLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - 𝕀_small = one(B.A) - 𝕆 = zero(𝕀_small) - 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) - 𝕀_big = one(B) - A_mat = B.A * 𝕀_small - B̂ = hcat(vcat(T(.5) * A_mat, B.B), E) - B̄ = hcat(vcat(𝕀_small, T(.5) * A_mat), vcat(zero(B.B'), -B.B'))' - - StiefelManifold((𝕀_big + T(.5) * B̂ * inv(𝕀_small2 - T(.5) * B̄' * B̂) * B̄') * (𝕀_big + T(.5) * B)) -end - -@doc raw""" - cayley(B̄::GrassmannLieAlgHorMatrix) - -Compute the Cayley retraction of `B`. - -This is equivalent to the method of [`cayley`](@ref) for [StiefelLieAlgHorMatrix](@ref). - -See [`cayley(::StiefelLieAlgHorMatrix)`](@ref). -""" -function cayley(B::GrassmannLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - backend = networkbackend(B) - 𝕆 = KernelAbstractions.zeros(backend, T, B.n, B.n) - 𝕀_small = one(𝕆) - 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) - 𝕀_big = one(B) - B̂ = hcat(vcat(𝕆, B.B), E) - B̄ = hcat(vcat(𝕀_small, 𝕆), vcat(zero(B.B'), -B.B'))' - - GrassmannManifold((𝕀_big + T(.5) * B̂ * inv(𝕀_small2 - T(.5) * B̄' * B̂) * B̄') * (𝕀_big + T(.5) * B)) -end \ No newline at end of file diff --git a/src/optimizers/momentum_optimizer.jl b/src/optimizers/momentum_optimizer.jl deleted file mode 100644 index 1404981d3..000000000 --- a/src/optimizers/momentum_optimizer.jl +++ /dev/null @@ -1,32 +0,0 @@ -@doc raw""" - MomentumOptimizer(η, α) - -Make an instance of the momentum optimizer. - -The momentum optimizer is similar to the [`GradientOptimizer`](@ref). -It however has a nontrivial cache that stores past history (see [`MomentumCache`](@ref)). -The cache is updated via: -```math - B^{\mathrm{cache}} \gets \alpha{}B^{\mathrm{cache}} + \nabla_\mathrm{weights}L -``` -and then the final velocity is computed as -```math - \mathrm{velocity} \gets - \eta{}B^{\mathrm{cache}}. -``` - -# Implementation - -To save memory the *velocity* is stored in the input ``\nabla_WL``. -This is similar to the case of the [`GradientOptimizer`](@ref). -""" -struct MomentumOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - α::T - MomentumOptimizer(η = 1e-3, α = 1e-2) = new{typeof(η)}(η, α) -end - -#update for weights -function update!(o::Optimizer{<:MomentumOptimizer}, C::MomentumCache, B::AbstractVecOrMat) - add!(C.B, o.method.α*C.B, B) - mul!(B, -o.method.η, C.B) -end \ No newline at end of file diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl deleted file mode 100644 index 1395ce424..000000000 --- a/src/optimizers/optimizer.jl +++ /dev/null @@ -1,177 +0,0 @@ -@doc raw""" - Optimizer(method, cache, step, retraction) - -Store the `method` (e.g. [`AdamOptimizer`](@ref) with corresponding hyperparameters), the `cache` (e.g. [`AdamCache`](@ref)), the optimization step and the retraction. - -It takes as input an optimization method and the parameters of a network. - -Before one can call `Optimizer` a [`OptimizerMethod`](@ref) that stores all the hyperparameters of the optimizer needs to be specified. - -# Functor - -For an instance `o` of `Optimizer`, we can call the corresponding functor as: - -```julia -o(nn, dl, batch, n_epochs, loss) -``` - -The arguments are: -1. `nn::NeuralNetwork` -2. `dl::`[`DataLoader`](@ref) -3. `batch::`[`Batch`](@ref) -4. `n_epochs::Integer` -5. `loss::NetworkLoss` - -The last argument is optional for many neural network architectures. We have the following defaults: -- A [`TransformerIntegrator`](@ref) uses [`TransformerLoss`](@ref). -- A [`NeuralNetworkIntegrator`](@ref) uses `FeedForwardLoss` (this loss is defined in `AbstractNeuralNetworks`). -- An [`AutoEncoder`](@ref) uses [`AutoEncoderLoss`](@ref). - -In addition there is an optional keyword argument that can be supplied to the functor: -- `show_progress=true`: This specifies whether a progress bar should be shown during training. - -# Implementation - -Internally the functor for `Optimizer` calls [`GlobalSection`](@ref) once at the start and then [`optimize_for_one_epoch!`](@ref) for each epoch. -""" -mutable struct Optimizer{MT<:OptimizerMethod, CT, RT} - method::MT - cache::CT - step::Int - retraction::RT -end - -eltype(::Optimizer{<:OptimizerMethod{T}}) where T = T - -@doc raw""" - Optimizer(method, nn_params) - -Allocate the cache for a specific `method` and `nn_params` for an instance of `Optimizer`. - -Internally this calls [`init_optimizer_cache`](@ref). - -An equivalent constructor is - -```julia -Optimizer(method, nn::NeuralNetwork) -``` - -# Arguments - -The optional keyword argument is the retraction. By default this is [`cayley`](@ref). -""" -function Optimizer(method::OptimizerMethod, nn_params::Union{NeuralNetworkParameters, NamedTuple}; retraction = cayley) - Optimizer(method, init_optimizer_cache(method, nn_params), 0, retraction) -end - -function Optimizer(method::OptimizerMethod, nn::NeuralNetwork; kwargs...) - Optimizer(method, params(nn); kwargs...) -end - -Optimizer(nn::NeuralNetwork, m::OptimizerMethod; kwargs...) = Optimizer(m, nn; kwargs...) - -@doc raw""" - update!(o, cache, B) - -Update the `cache` and output a final velocity that is stored in `B`. - -Note that ``B\in\mathfrak{g}^\mathrm{hor}`` in general. - -In the manifold case the final velocity is the input to a retraction. -""" -function update!(o::Optimizer, ::AbstractCache, ::AbstractArray) - error("No update rule implemented for method", o.method) -end - -####################################################################################### -# optimization step function - -function _optimization_step!(o::Optimizer, λY::NamedTuple, ps::NamedTuple, cache::NamedTuple, dx::Union{NamedTuple, NeuralNetworkParameters}) - gx = rgrad(ps, dx) - B = global_rep(λY, gx) - update!(o, cache, B) - update_section!(λY, B, o.retraction) - - nothing -end - -function optimization_step!(o::Optimizer, λY::NamedTuple, ps::NeuralNetworkParameters, dx::Union{NamedTuple, NeuralNetworkParameters}) - @assert keys(o.cache) == keys(λY) == keys(ps) == keys(dx) - o.step += 1 - for key in keys(o.cache) - cache = o.cache[key] - λY_temp = λY[key] - ps_temp = ps[key] - dx_temp = dx[key] - _optimization_step!(o, λY_temp, ps_temp, cache, dx_temp) - end -end - -@doc raw""" - optimization_step!(o, λY, ps, dx) - -Update the weights `ps` based on an [`Optimizer`](@ref), a `cache` and first-order derivatives `dx`. - -`optimization_step!` is calling [`update!`](@ref) internally. -`update!` has to be implemented for every [`OptimizerMethod`](@ref). - -# Arguments - -All arguments into `optimization_step!` are mandatory: -1. `o::`[`Optimizer`](@ref), -2. `λY::NamedTuple`: this named tuple has the same keys as `ps`, but contains [`GlobalSection`](@ref)s, -3. `ps::NamedTuple`: the neural network parameters, -5. `dx::Union{NamedTuple, NeuralNetworkParameters}`: the gradients stores as a NamedTuple. - -All the arguments are given as `NamedTuple`s as the neural network weights are stores in that format. - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: params - -l = StiefelLayer(3, 5) -ps = params(NeuralNetwork(Chain(l), Float32)).L1 -cache = apply_toNT(MomentumCache, ps) -o = Optimizer(MomentumOptimizer(), cache, 0, geodesic) -λY = GlobalSection(ps) -dx = (weight = rand(Float32, 5, 3), ) - -# call the optimizer -optimization_step!(o, λY, ps, dx) - -_test_nt(x) = typeof(x) <: NamedTuple - -_test_nt(λY) & _test_nt(ps) & _test_nt(cache) & _test_nt(dx) - -# output - -true -``` - -# Extended help -The derivatives `dx` here are usually obtained via an AD routine by differentiating a loss function, i.e. `dx` is ``\nabla_xL``. -""" -function optimization_step!(o::Optimizer, λY::NamedTuple, ps::NamedTuple, dx::Union{NamedTuple, NeuralNetworkParameters}) - o.step += 1 - - _optimization_step!(o, λY, ps, o.cache, dx) -end - -####################################################################################### -# utils functions (should probably be put somewhere else) - -rgrad(ps::NamedTuple, dx::Union{NamedTuple, NeuralNetworkParameters}) = apply_toNT(rgrad, ps, dx) - -function rgrad(Y::AbstractVecOrMat, dx::AbstractVecOrMat) - @assert size(Y) == size(dx) - dx -end - -# do we need those two? -function update!(m::Optimizer, C::NamedTuple, B::NamedTuple) - apply_toNT(m, C, B, update!) -end - -function apply_toNT(m::Optimizer, ps₁::NamedTuple, ps₂::NamedTuple, fun_name) - apply_toNT((ps₁, ps₂) -> fun_name(m, ps₁, ps₂), ps₁, ps₂) -end diff --git a/src/optimizers/optimizer_caches.jl b/src/optimizers/optimizer_caches.jl deleted file mode 100644 index 9f9dd14ab..000000000 --- a/src/optimizers/optimizer_caches.jl +++ /dev/null @@ -1,120 +0,0 @@ -@doc raw""" - AbstractCache - -`AbstractCache` has subtypes: [`AdamCache`](@ref), [`MomentumCache`](@ref), [`GradientCache`](@ref) and [`BFGSCache`](@ref). - -All of them can be initialized with providing an array (also supporting manifold types). -""" -abstract type AbstractCache{T} end - -############################################################################# -# All the definitions of the caches - -@doc raw""" - AdamCache(Y) - -Store the first and second moment for `Y` (initialized as zeros). - -First and second moments are called `B₁` and `B₂`. - -If the cache is called with an instance of a homogeneous space, e.g. the [`StiefelManifold`](@ref) ``St(n,N)`` it initializes the moments as elements of ``\mathfrak{g}^\mathrm{hor}`` ([`StiefelLieAlgHorMatrix`](@ref)). - -# Examples - -```jldoctest -using GeometricMachineLearning - -Y = rand(StiefelManifold, 5, 3) -AdamCache(Y).B₁ - -# output - -5×5 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}: - 0.0 -0.0 -0.0 -0.0 -0.0 - 0.0 0.0 -0.0 -0.0 -0.0 - 0.0 0.0 0.0 -0.0 -0.0 - 0.0 0.0 0.0 0.0 0.0 - 0.0 0.0 0.0 0.0 0.0 -``` -""" -struct AdamCache{T, AT <: AbstractArray{T}} <: AbstractCache{T} - B₁::AT - B₂::AT - function AdamCache(Y::AbstractArray) - new{eltype(Y), typeof(zero(Y))}(zero(Y), zero(Y)) - end -end - -function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::AdamCache) - println(io, raw"`AdamCache` that currently stores `B₁` as ...") - show(io, "text/plain", C.B₁) - println(io, "") - println(io, raw"and `B₂` as ...") - show(io, "text/plain", C.B₂) -end - -@doc raw""" - MomentumCache(Y) - -Store the moment for `Y` (initialized as zeros). - -The moment is called `B`. - -If the cache is called with an instance of a [`Manifold`](@ref) it initializes the moments as elements of ``\mathfrak{g}^\mathrm{hor}`` ([`AbstractLieAlgHorMatrix`](@ref)). - -See [`AdamCache`](@ref). -""" -struct MomentumCache{T, AT <: AbstractArray{T}} <:AbstractCache{T} - B::AT - function MomentumCache(Y::AbstractArray) - new{eltype(Y), typeof(zero(Y))}(zero(Y)) - end -end - -function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::MomentumCache) - println(io, raw"`MomentumCache` that currently stores `B`as ...") - show(io, "text/plain", C.B) -end - -@doc raw""" - GradientCache(Y) - -Do not store anything. - -The cache for the [`GradientOptimizer`](@ref) does not consider past information. -""" -struct GradientCache{T} <: AbstractCache{T} end -GradientCache(::AbstractArray{T}) where T = GradientCache{T}() - -############################################################################# -# All the setup_cache functions - -setup_adam_cache(ps::NamedTuple) = apply_toNT(setup_adam_cache, ps) -setup_momentum_cache(ps::NamedTuple) = apply_toNT(setup_momentum_cache, ps) -setup_gradient_cache(ps::NamedTuple) = apply_toNT(setup_gradient_cache, ps) - -function setup_cache(_setup_cache_function, ps::NeuralNetworkParameters) - ps_keys = keys(ps) - values = Tuple([_setup_cache_function(ps[key]) for key in ps_keys]) - NamedTuple{ps_keys}(values) -end - -setup_adam_cache(ps::NeuralNetworkParameters) = setup_cache(setup_adam_cache, ps) -setup_momentum_cache(ps::NeuralNetworkParameters) = setup_cache(setup_momentum_cache, ps) -setup_gradient_cache(ps::NeuralNetworkParameters) = setup_cache(setup_gradient_cache, ps) - -setup_adam_cache(B::AbstractArray{<:Number}) = AdamCache(B) -setup_momentum_cache(B::AbstractArray{<:Number}) = MomentumCache(B) -setup_gradient_cache(B::AbstractArray{<:Number}) = GradientCache(B) - -function Base.zero(Y::StiefelManifold{T}) where T - N, n = size(Y) - backend = networkbackend(Y.A) - zeros(backend, StiefelLieAlgHorMatrix{T}, N, n) -end - -function Base.zero(Y::GrassmannManifold{T}) where T - N, n = size(Y) - backend = networkbackend(Y.A) - zeros(backend, GrassmannLieAlgHorMatrix{T}, N, n) -end diff --git a/src/optimizers/optimizer_method.jl b/src/optimizers/optimizer_method.jl deleted file mode 100644 index 290b293c6..000000000 --- a/src/optimizers/optimizer_method.jl +++ /dev/null @@ -1,15 +0,0 @@ -@doc raw""" - OptimizerMethod - -Each `Optimizer` has to be called with an `OptimizerMethod`. This specifies how the neural network weights are updated in each optimization step. -""" -abstract type OptimizerMethod{T} end - -# @doc raw""" -# init_optimizer_cache(method, x) -# -# Initialize the optimizer cache based on input `x` for the given `method`. -# """ -function init_optimizer_cache(::OptimizerMethod, x) end - -eltype(::OptimizerMethod{T}) where T = T \ No newline at end of file diff --git a/src/pullbacks/zygote_pullback.jl b/src/pullbacks/zygote_pullback.jl index e33e081df..d72387f14 100644 --- a/src/pullbacks/zygote_pullback.jl +++ b/src/pullbacks/zygote_pullback.jl @@ -27,8 +27,10 @@ struct ZygotePullback{NNLT} <: AbstractPullback{NNLT} loss::NNLT end -(_pullback::ZygotePullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps) -(_pullback::ZygotePullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) +(_pullback::ZygotePullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_nt), ps) +(_pullback::ZygotePullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) """ _processing(returned_pullback) @@ -37,4 +39,4 @@ Strip `returned_pullback` from unnecessary `Zygote`-induces garbage. Also see the docs for [`ZygotePullback`](@ref). """ -_processing = _get_params∘_get_contents \ No newline at end of file +_processing = _get_params ∘ _get_contents diff --git a/src/utils.jl b/src/utils.jl index 0fe9b3b59..ee8d45e01 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,7 +1,7 @@ # Convenient structure struct NothingFunction <: Function end (::NothingFunction)(args...) = nothing -is_NothingFunction(f::Function) = typeof(f)==NothingFunction +is_NothingFunction(f::Function) = typeof(f) == NothingFunction struct UnknownProblem <: AbstractProblem end @@ -9,14 +9,13 @@ const ∞ = Inf # Functions on typple and named tuple -@inline next(i::Int,j::Int) = (i,j+1) -@inline next(i::Int) = (i+1,) +@inline next(i::Int, j::Int) = (i, j + 1) +@inline next(i::Int) = (i + 1,) @inline tuplejoin(x) = x @inline tuplejoin(x, y) = (x..., y...) @inline tuplejoin(x, y, z...) = tuplejoin(tuplejoin(x, y), z...) - rdevelop(x) = x rdevelop(t::Tuple{Any}) = [rdevelop(t[1])...] rdevelop(t::Tuple) = [rdevelop(t[1])..., rdevelop(t[2:end])...] @@ -27,8 +26,7 @@ develop(t::Tuple{Any}) = [develop(t[1])...] develop(t::Tuple) = [develop(t[1])..., develop(t[2:end])...] develop(t::NamedTuple) = vcat([[develop(e)...] for e in t]...) - -_tuplediff(t₁::Tuple,t₂::Tuple) = tuple(setdiff(Set(t₁),Set(t₂))...) +_tuplediff(t₁::Tuple, t₂::Tuple) = tuple(setdiff(Set(t₁), Set(t₂))...) function apply_toNT(fun, ps::NamedTuple...) for p in ps @@ -37,17 +35,23 @@ function apply_toNT(fun, ps::NamedTuple...) NamedTuple{keys(ps[1])}(fun(p...) for p in zip(ps...)) end -# overload norm -_norm(dx::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} = (norm(dx.q) + norm(dx.p)) / √2 # we need this because of a Zygote problem +# overload norm +function _norm(dx::NT) where { + AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} + (norm(dx.q) + norm(dx.p)) / √2 +end # we need this because of a Zygote problem _norm(dx::NamedTuple) = sum(apply_toNT(norm, dx)) / √length(dx) _norm(A::AbstractArray) = norm(A) -# overloaded +/- operation -_diff(dx₁::NT, dx₂::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} = (q = dx₁.q - dx₂.q, p = dx₁.p - dx₂.p) # we need this because of a Zygote problem +# overloaded +/- operation +function _diff(dx₁::NT, + dx₂::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} + (q = dx₁.q - dx₂.q, p = dx₁.p - dx₂.p) +end # we need this because of a Zygote problem _diff(dx₁::NamedTuple, dx₂::NamedTuple) = apply_toNT(_diff, dx₁, dx₂) -_diff(A::AbstractArray, B::AbstractArray) = A - B +_diff(A::AbstractArray, B::AbstractArray) = A - B _add(dx₁::NamedTuple, dx₂::NamedTuple) = apply_toNT(_add, dx₁, dx₂) -_add(A::AbstractArray, B::AbstractArray) = A + B +_add(A::AbstractArray, B::AbstractArray) = A + B function add!(C::AbstractVecOrMat, A::AbstractVecOrMat, B::AbstractVecOrMat) @assert size(A) == size(B) == size(C) @@ -61,19 +65,18 @@ end # Type pyracy!! function Base.:+(a::Float64, b::Tuple{Float64}) x, = b - return a+x + return a + x end # Type pyracy!! function Base.:+(a::Vector{Float64}, b::Tuple{Float64}) x, = b y, = a - return y+x + return y + x end - -# Kernel that is needed for functions relating to `SymmetricMatrix` and `SkewSymMatrix` -@kernel function write_ones_kernel!(unit_matrix::AbstractMatrix{T}) where T +# Kernel that is needed for functions relating to `SymmetricMatrix` and `SkewSymMatrix` +@kernel function write_ones_kernel!(unit_matrix::AbstractMatrix{T}) where {T} i = @index(Global) unit_matrix[i, i] = one(T) end @@ -92,18 +95,17 @@ end # utils functions on string function type_without_brace(var) type_str = string(typeof(var)) - replace(type_str, r"\{.*\}"=>"") + replace(type_str, r"\{.*\}" => "") end -function center_align_text(text,width) +function center_align_text(text, width) padding = max(0, width - length(text)) - left_padding = repeat(" ",padding ÷2) + left_padding = repeat(" ", padding ÷ 2) right_padding = repeat(" ", padding - length(left_padding)) aligned_text = left_padding * text * right_padding return aligned_text end - # The following are fallback functions - maybe you want to put them into a separate file function global_section(::AbstractVecOrMat) nothing @@ -163,10 +165,36 @@ const QPTOAT = Union{QPT, AbstractArray} This could be data in ``(q, p)\in\mathbb{R}^{2d}`` form or come from an arbitrary vector space. """ -const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where T +const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where {T} Base.:≈(qp₁::QPT, qp₂::QPT) = (qp₁.q ≈ qp₂.q) & (qp₁.p ≈ qp₂.p) _eltype(x) = eltype(x) _eltype(ps::NamedTuple) = _eltype(ps[1]) -_eltype(ps::Tuple) = _eltype(ps[1]) \ No newline at end of file +_eltype(ps::Tuple) = _eltype(ps[1]) + +function GradientZygote(F, x::AbstractArray{T}) where {T} + F = (_b, _a) -> (GeometricOptimizers._copyto!(_b, Zygote.gradient(loss, _a)[1])) + GeometricOptimizers.GradientFunction(F, x) +end + +function GMLOptimizer(x::VT, problem::OptimizerProblem{T}; + options_kwargs...) where {T, VT <: OptimizerSolution{T}} + grad = GradientZygote(problem.F, x) + Optimizer(x, problem; algorithm = Adam(), linesearch = Static(0.01), + gradient = grad, options_kwargs...) +end + +# type piracy. This should be put in `AbstractNeuralNetworks`! +function GeometricOptimizers.OptimizerProblem(nn::NeuralNetwork) + loss = NetworkLoss(nn) + _params = params(nn) + OptimizerProblem(ps -> loss(model, ps), NamedTuple{keys(_params)}(values(_params))) +end + +function GeometricOptimizers.Optimizer( + algorithm::OptimizerMethod, nn::NeuralNetwork; options_kwargs...) + _params = params(nn) + GMLOptimizer( + NamedTuple{keys(_params)}(values(_params)), OptimizerProblem(nn); options_kwargs...) +end diff --git a/test/runtests.jl b/test/runtests.jl index 9aebc783a..b5d7874ed 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,83 +1,202 @@ using SafeTestsets, Test, GeometricMachineLearning using Documenter: doctest -# @testset "Doc tests " begin doctest(GeometricMachineLearning; manual = false) end +@testset "Doc tests " begin + doctest(GeometricMachineLearning; manual=false) +end # reduced order modeling tests -@safetestset "PSD tests " begin include("psd_architecture_tests.jl") end -@safetestset "SymplecticAutoencoder tests " begin include("symplectic_autoencoder_tests.jl") end -@safetestset "Check if autoencoder error is lower than PSD error " begin include("sae_error_lower_than_psd_error.jl") end -@safetestset "Check reduced model " begin include("reduced_system.jl") end - -@safetestset "Check parameterlength " begin include("parameterlength/check_parameterlengths.jl") end - -@safetestset "Arrays #1 " begin include("arrays/array_tests.jl") end -@safetestset "Map to skew " begin include("arrays/map_to_skew.jl") end -@safetestset "Sampling of arrays " begin include("arrays/random_generation_of_custom_arrays.jl") end -@safetestset "Addition tests for custom arrays " begin include("arrays/addition_tests_for_custom_arrays.jl") end -@safetestset "Scalar multiplication tests for custom arrays " begin include("arrays/scalar_multiplication_for_custom_arrays.jl") end -@safetestset "Matrix multiplication tests for custom arrays " begin include("arrays/matrix_multiplication_for_custom_arrays.jl") end -@safetestset "Test constructors for custom arrays " begin include("arrays/constructor_tests_for_custom_arrays.jl") end -@safetestset "Symplectic Potential (array tests) " begin include("arrays/poisson_tensor.jl") end -@safetestset "Test StiefelLieAlgHorMatrix constructors and lifts " begin include("arrays/test_stiefel_lie_alg_hor_constructors.jl") end -@safetestset "Test GrassmannLieAlgHorMatrix constructors and lifts " begin include("arrays/test_grassmann_lie_alg_hor_constructors.jl") end - -@safetestset "Test triangular matrices " begin include("arrays/triangular.jl") end -@safetestset "Manifolds (Stiefel): " begin include("manifolds/stiefel_manifold.jl") end -@safetestset "Manifolds (Grassmann): " begin include("manifolds/grassmann_manifold.jl") end -@safetestset "Gradient Layer " begin include("layers/gradient_layer_tests.jl") end -@safetestset "Test symplecticity of upscaling layer " begin include("layers/sympnet_layers_test.jl") end -@safetestset "Hamiltonian Neural Network " begin include("hamiltonian_neural_network_tests.jl") end -@safetestset "Manifold Neural Network Layers " begin include("layers/manifold_layers.jl") end - -@safetestset "Custom tensor matrix multiplication " begin include("kernels/tensor_mat_mul.jl") end -@safetestset "Custom inverse for 2x2, 3x3, 4x4, 5x5 matrices " begin include("kernels/tensor_inverse.jl") end -@safetestset "Custom AD rules for kernels " begin include("custom_ad_rules/kernel_pullbacks.jl") end -@safetestset "ResNet " begin include("layers/resnet_tests.jl") end - +@safetestset "PSD tests " begin + include("psd_architecture_tests.jl") +end +@safetestset "SymplecticAutoencoder tests " begin + include("symplectic_autoencoder_tests.jl") +end +@safetestset "Check if autoencoder error is lower than PSD error " begin + include("sae_error_lower_than_psd_error.jl") +end +@safetestset "Check reduced model " begin + include("reduced_system.jl") +end +@safetestset "Check parameterlength " begin + include("parameterlength/check_parameterlengths.jl") +end +@safetestset "Arrays #1 " begin + include("arrays/array_tests.jl") +end +@safetestset "Map to skew " begin + include("arrays/map_to_skew.jl") +end +@safetestset "Sampling of arrays " begin + include("arrays/random_generation_of_custom_arrays.jl") +end +@safetestset "Addition tests for custom arrays " begin + include("arrays/addition_tests_for_custom_arrays.jl") +end +@safetestset "Scalar multiplication tests for custom arrays " begin + include("arrays/scalar_multiplication_for_custom_arrays.jl") +end +@safetestset "Matrix multiplication tests for custom arrays " begin + include("arrays/matrix_multiplication_for_custom_arrays.jl") +end +@safetestset "Test constructors for custom arrays " begin + include("arrays/constructor_tests_for_custom_arrays.jl") +end +@safetestset "Symplectic Potential (array tests) " begin + include("arrays/poisson_tensor.jl") +end +@safetestset "Test StiefelLieAlgHorMatrix constructors and lifts " begin + include("arrays/test_stiefel_lie_alg_hor_constructors.jl") +end +@safetestset "Test GrassmannLieAlgHorMatrix constructors and lifts " begin + include("arrays/test_grassmann_lie_alg_hor_constructors.jl") +end +@safetestset "Test triangular matrices " begin + include("arrays/triangular.jl") +end +@safetestset "Manifolds (Stiefel): " begin + include("manifolds/stiefel_manifold.jl") +end +@safetestset "Manifolds (Grassmann): " begin + include("manifolds/grassmann_manifold.jl") +end +@safetestset "Gradient Layer " begin + include("layers/gradient_layer_tests.jl") +end +@safetestset "Test symplecticity of upscaling layer " begin + include("layers/sympnet_layers_test.jl") +end +@safetestset "Hamiltonian Neural Network " begin + include("hamiltonian_neural_network_tests.jl") +end +@safetestset "Manifold Neural Network Layers " begin + include("layers/manifold_layers.jl") +end +@safetestset "Custom tensor matrix multiplication " begin + include("kernels/tensor_mat_mul.jl") +end +@safetestset "Custom inverse for 2x2, 3x3, 4x4, 5x5 matrices " begin + include("kernels/tensor_inverse.jl") +end +@safetestset "Custom AD rules for kernels " begin + include("custom_ad_rules/kernel_pullbacks.jl") +end +@safetestset "ResNet " begin + include("layers/resnet_tests.jl") +end # transformer-related tests -@safetestset "Test setup of MultiHeadAttention layer Stiefel weights " begin include("transformer_related/multi_head_attention_stiefel_setup.jl") end -@safetestset "Test geodesic and Cayley retr for the MultiHeadAttention layer w/ St weights " begin include("transformer_related/multi_head_attention_stiefel_retraction.jl") end -@safetestset "Test the correct setup of the various optimizer caches for MultiHeadAttention " begin include("transformer_related/multi_head_attention_stiefel_optim_cache.jl") end -@safetestset "Check if the transformer can be applied to a tensor. " begin include("transformer_related/transformer_application.jl") end -@safetestset "Check if the gradient/pullback of MultiHeadAttention changes type in St case " begin include("transformer_related/transformer_gradient.jl") end -@safetestset "Check if the optimization_step! changes the parameters of the transformer " begin include("transformer_related/transformer_optimizer.jl") end - -@safetestset "Attention layer #1 " begin include("attention_layer/attention_setup.jl") end -@safetestset "Classification layer " begin include("layers/classification.jl") end -@safetestset "Optimizer #1 " begin include("optimizers/utils/global_sections.jl") end -@safetestset "Optimizer #2 " begin include("optimizers/utils/optimization_step.jl") end -@safetestset "Optimizer #3 " begin include("optimizers/utils/modified_exponential.jl") end -@safetestset "Optimizer #4 " begin include("optimizers/optimizer_convergence_tests/svd_optim.jl") end -@safetestset "Optimizer #5 " begin include("optimizers/optimizer_convergence_tests/psd_optim.jl") end -@safetestset "Check if Adam with decay converges " begin include("optimizers/optimizer_convergence_tests/adam_with_learning_rate_decay.jl") end -@safetestset "BFGS Optimizer tests " begin include("optimizers/bfgs_optimizer.jl") end -@safetestset "Data " begin include("data/test_data.jl") end -@safetestset "Batch " begin include("data/test_batch.jl") end +@safetestset "Test setup of MultiHeadAttention layer Stiefel weights " begin + include("transformer_related/multi_head_attention_stiefel_setup.jl") +end +@safetestset "Test geodesic and Cayley retr for the MultiHeadAttention layer w/ St weights " begin + include("transformer_related/multi_head_attention_stiefel_retraction.jl") +end +@safetestset "Test the correct setup of the various optimizer caches for MultiHeadAttention " begin + include("transformer_related/multi_head_attention_stiefel_optim_cache.jl") +end +@safetestset "Check if the transformer can be applied to a tensor. " begin + include("transformer_related/transformer_application.jl") +end +@safetestset "Check if the gradient/pullback of MultiHeadAttention changes type in St case " begin + include("transformer_related/transformer_gradient.jl") +end +@safetestset "Check if the optimization_step! changes the parameters of the transformer " begin + include("transformer_related/transformer_optimizer.jl") +end + +@safetestset "Attention layer #1 " begin + include("attention_layer/attention_setup.jl") +end +@safetestset "Classification layer " begin + include("layers/classification.jl") +end +@safetestset "Optimizer #1 " begin + include("optimizers/utils/global_sections.jl") +end +@safetestset "Optimizer #2 " begin + include("optimizers/utils/optimization_step.jl") +end +@safetestset "Optimizer #3 " begin + include("optimizers/utils/modified_exponential.jl") +end +@safetestset "Optimizer #4 " begin + include("optimizers/optimizer_convergence_tests/svd_optim.jl") +end +@safetestset "Optimizer #5 " begin + include("optimizers/optimizer_convergence_tests/psd_optim.jl") +end +@safetestset "Check if Adam with decay converges " begin + include("optimizers/optimizer_convergence_tests/adam_with_learning_rate_decay.jl") +end +@safetestset "BFGS Optimizer tests " begin + include("optimizers/bfgs_optimizer.jl") +end +@safetestset "Data " begin + include("data/test_data.jl") +end +@safetestset "Batch " begin + include("data/test_batch.jl") +end # @safetestset "Method " begin include("train!/test_method.jl") end -@safetestset "Matching " begin include("data/test_matching.jl") end - -@safetestset "Test data loader for q and p data " begin include("data_loader/batch_data_loader_qp_test.jl") end -@safetestset "Test mnist_utils. " begin include("data_loader/mnist_utils.jl") end -@safetestset "Test the data loader in combination with optimization_step! " begin include("data_loader/data_loader_optimization_step.jl") end -@safetestset "Optimizer functor with data loader for Adam " begin include("data_loader/optimizer_functor_with_adam.jl") end -@safetestset "Test data loader for a tensor (q and p data) " begin include("data_loader/draw_batch_for_tensor_test.jl") end - -@safetestset "Test NetworkLoss + Optimizer " begin include("network_losses/losses_and_optimization.jl") end - -@safetestset "Test parallel inverses " begin include("kernels/tensor_inverse.jl") end -@safetestset "Test parallel Cayley " begin include("kernels/tensor_cayley.jl") end - -@safetestset "Test volume-preserving feedforward neural network " begin include("layers/volume_preserving_feedforward.jl") end - -@safetestset "SympNet integrator " begin include("sympnet_integrator.jl") end -@safetestset "Regular transformer integrator " begin include("standard_transformer_integrator.jl") end - -@safetestset "Batch functor(s) " begin include("batch/batch_functor.jl") end - -@safetestset "Volume-Preserving Transformer (skew-symmetric tests) " begin include("volume_preserving_attention/test_skew_map.jl") end -@safetestset "Volume-Preserving Transformer (cayley-transform tests) " begin include("volume_preserving_attention/test_cayley_transforms.jl") end - -@safetestset "Linear Symplectic Attention " begin include("linear_symplectic_attention.jl") end -@safetestset "Linear Symplectic Transformer " begin include("linear_symplectic_transformer.jl") end - -@safetestset "DataLoader for input and output " begin include("data_loader/data_loader_for_input_and_output.jl") end +@safetestset "Matching " begin + include("data/test_matching.jl") +end + +@safetestset "Test data loader for q and p data " begin + include("data_loader/batch_data_loader_qp_test.jl") +end +@safetestset "Test mnist_utils. " begin + include("data_loader/mnist_utils.jl") +end +@safetestset "Test the data loader in combination with optimization_step! " begin + include("data_loader/data_loader_optimization_step.jl") +end +@safetestset "Optimizer functor with data loader for Adam " begin + include("data_loader/optimizer_functor_with_adam.jl") +end +@safetestset "Test data loader for a tensor (q and p data) " begin + include("data_loader/draw_batch_for_tensor_test.jl") +end + +@safetestset "Test NetworkLoss + Optimizer " begin + include("network_losses/losses_and_optimization.jl") +end + +@safetestset "Test parallel inverses " begin + include("kernels/tensor_inverse.jl") +end +@safetestset "Test parallel Cayley " begin + include("kernels/tensor_cayley.jl") +end + +@safetestset "Test volume-preserving feedforward neural network " begin + include("layers/volume_preserving_feedforward.jl") +end + +@safetestset "SympNet integrator " begin + include("sympnet_integrator.jl") +end +@safetestset "Regular transformer integrator " begin + include("standard_transformer_integrator.jl") +end + +@safetestset "Batch functor(s) " begin + include("batch/batch_functor.jl") +end + +@safetestset "Volume-Preserving Transformer (skew-symmetric tests) " begin + include("volume_preserving_attention/test_skew_map.jl") +end +@safetestset "Volume-Preserving Transformer (cayley-transform tests) " begin + include("volume_preserving_attention/test_cayley_transforms.jl") +end + +@safetestset "Linear Symplectic Attention " begin + include("linear_symplectic_attention.jl") +end +@safetestset "Linear Symplectic Transformer " begin + include("linear_symplectic_transformer.jl") +end + +@safetestset "DataLoader for input and output " begin + include("data_loader/data_loader_for_input_and_output.jl") +end diff --git a/test/symplectic_autoencoder_tests.jl b/test/symplectic_autoencoder_tests.jl index b453803cf..6ac517425 100644 --- a/test/symplectic_autoencoder_tests.jl +++ b/test/symplectic_autoencoder_tests.jl @@ -1,23 +1,22 @@ using GeometricMachineLearning using Test using Zygote: jacobian -import Random +import Random Random.seed!(123) -function test_accuracy(N::Integer, n::Integer; tol::Real = .35, n_epochs::Integer = 100) +function test_accuracy(N::Integer, n::Integer; tol::Real = 0.35, n_epochs::Integer = 100) dl = DataLoader(rand(N, 10 * N); autoencoder = true) sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n)) - - o = Optimizer(AdamOptimizer(), sae_nn) + + o = Optimizer(Adam(), sae_nn) sae_error = o(sae_nn, dl, Batch(10), n_epochs)[end] - @test sae_error < tol + @test sae_error < tol end function test_encoder_and_decoder(N::Integer, n::Integer) - sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n)) sae_encoder = encoder(sae_nn) sae_decoder = decoder(sae_nn) @@ -33,7 +32,7 @@ function test_symplecticity(N::Integer, n::Integer) sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n)) sae_decoder = decoder(sae_nn) test_vector = rand(n) - + # this matrix should be symplectic sympl_mat = jacobian(vec -> sae_decoder(vec), test_vector)[1] @test PoissonTensor(n) ≈ sympl_mat' * PoissonTensor(N) * sympl_mat @@ -53,4 +52,4 @@ function all_tests(N::Integer, n::Integer; kwargs...) end all_tests(10, 6) -all_tests(20, 10) \ No newline at end of file +all_tests(20, 10) From d986e61a037be2176ce794356eb03443ecbc5fa3 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 8 Jul 2026 09:08:54 +0200 Subject: [PATCH 2/7] Migrate optimizer layer from GML to GeometricOptimizers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the deleted src/optimizers/ directory with a thin integration layer that wraps GeometricOptimizers' manifold optimizers for use in GML's neural-network training loop. Key changes: src/utils.jl - Remove stale GradientZygote / GMLOptimizer helpers that referenced a deleted GML Optimizer type. - Add AbstractCache{T} as a backward-compat alias for GeometricOptimizers.OptimizerCache{T}. - Define _GMLGradient{T,VT} <: GeometricOptimizers.Gradient{T}: a mutable wrapper that holds a pre-computed Zygote gradient and satisfies the Gradient interface expected by GO's update! methods. - Define GML's own Optimizer struct (method, cache, state, retraction, _grad) with constructors that accept either a NeuralNetwork or a raw params NamedTuple. - Implement optimization_step!(opt, λY, ps, dp): injects the Zygote gradient, calls GO's update! to compute the Lie-algebra direction, applies the retraction via update_section!, then copies the new manifold point back into ps and λY. src/GeometricMachineLearning.jl - Extend the `using GeometricOptimizers:` import to bring in Cayley, Geodesic, cayley, geodesic, retraction, apply_section, apply_section!, OptimizerMethod, GradientCache, MomentumCache, AdamCache. - Replace the old optimizer export block with new names from GO (GradientMethod/State, MomentumMethod/State, Adam/AdamState, etc.). - Add backward-compat const aliases GradientOptimizer, MomentumOptimizer, AdamOptimizer so existing call-sites continue to compile. - Remove non-existent exports (BFGSOptimizer, BFGSCache, AdamOptimizerWithDecay, duplicate Adam). src/training/train.jl - Optimizer(m, params(nn)) → Optimizer(m, nn) (new constructor signature). - optimization_step!(opt, model(nn), ...) → optimization_step!(opt, GlobalSection(params(nn)), ...) (λY now takes a GlobalSection, not a model). Tests — update optimizer constructor call-sites throughout: - AdamOptimizer() → Adam() - GradientOptimizer() → GradientMethod() - MomentumOptimizer() → MomentumMethod() - BFGSOptimizer() calls commented out (no manifold BFGS in GO). - init_optimizer_cache(method, ps) → GeometricOptimizers.OptimizerCache(method, ps). - MomentumCache.B (old field) → MomentumCache.δ (zero-init direction field in GO). - Rewrite multi_head_attention_stiefel_optim_cache.jl to match GO's actual AdamCache / MomentumCache / GradientCache field layout. Co-Authored-By: Claude Sonnet 4.6 --- src/GeometricMachineLearning.jl | 24 ++-- src/training/train.jl | 4 +- src/utils.jl | 106 +++++++++++++++--- .../data_loader_for_input_and_output.jl | 2 +- .../optimizer_functor_with_adam.jl | 5 +- .../network_losses/losses_and_optimization.jl | 2 +- test/sae_error_lower_than_psd_error.jl | 2 +- test/symplectic_autoencoder_tests.jl | 2 +- ...ulti_head_attention_stiefel_optim_cache.jl | 55 ++++----- ...multi_head_attention_stiefel_retraction.jl | 9 +- .../multi_head_attention_stiefel_setup.jl | 7 +- .../transformer_optimizer.jl | 21 ++-- 12 files changed, 153 insertions(+), 86 deletions(-) diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index f240eaf2a..375c8db10 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -25,7 +25,9 @@ using SymbolicNeuralNetworks: derivative, _get_contents, _get_params, SymbolicNe using Symbolics: @variables, substitute using GeometricOptimizers -using GeometricOptimizers: OptimizerSolution, Geodesic, OptimizerMethod +using GeometricOptimizers: OptimizerSolution, Cayley, Geodesic, cayley, geodesic, retraction, + apply_section, apply_section!, OptimizerMethod, + GradientCache, MomentumCache, AdamCache import AbstractNeuralNetworks: Architecture, Model, AbstractExplicitLayer, AbstractExplicitCell, AbstractNeuralNetwork, NeuralNetwork, @@ -168,27 +170,25 @@ export ResNet export Transformer export TransformerIntegrator, StandardTransformerIntegrator -# INCLUDE OPTIMIZERS +# INCLUDE OPTIMIZERS — types come from GeometricOptimizers export OptimizerMethod, AbstractCache -export GradientOptimizer, GradientCache -export MomentumOptimizer, MomentumCache -export AdamOptimizerWithDecay -export AdamOptimizer, AdamCache -export BFGSOptimizer, BFGSCache - +export GradientMethod, GradientCache, GradientState +export MomentumMethod, MomentumCache, MomentumState +export Adam, AdamCache, AdamState export Optimizer export optimization_step! - export GlobalSection, apply_section, apply_section! export global_rep export Geodesic, Cayley export geodesic, cayley export retraction -# export ⊙², √ᵉˡᵉ, /ᵉˡᵉ, scalar_add export update! export check - -export Adam +# backward-compat aliases (old names → new names) +const GradientOptimizer = GradientMethod +const MomentumOptimizer = MomentumMethod +const AdamOptimizer = Adam +export GradientOptimizer, MomentumOptimizer, AdamOptimizer #INCLUDE ABSTRACT TRAINING integrator export AbstractTrainingMethod diff --git a/src/training/train.jl b/src/training/train.jl index 89ad325ea..bc28b9bfe 100644 --- a/src/training/train.jl +++ b/src/training/train.jl @@ -54,7 +54,7 @@ function train!(nn::AbstractNeuralNetwork, _data::AbstractTrainingData, m::Optim Loss() = typeof(method) <: TrainingMethod ? loss(method, nn, data) : method(nn, data) # creation of optimiser - @timeit to "Creation of Optimizer" opt = Optimizer(m, params(nn)) + @timeit to "Creation of Optimizer" opt = Optimizer(m, nn) # creation of the array to store total loss total_loss = zeros(typeof(Loss()), ntraining) @@ -66,7 +66,7 @@ function train!(nn::AbstractNeuralNetwork, _data::AbstractTrainingData, m::Optim @timeit to "Computing Grad Loss" ∇params = loss_gradient(Loss, index_batch, params(nn)) - @timeit to "Performing Optimization step" optimization_step!(opt, model(nn), params(nn), ∇params) + @timeit to "Performing Optimization step" optimization_step!(opt, GlobalSection(params(nn)), params(nn), ∇params) total_loss[j] = Loss() diff --git a/src/utils.jl b/src/utils.jl index ee8d45e01..c3a1a9031 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -173,28 +173,98 @@ _eltype(x) = eltype(x) _eltype(ps::NamedTuple) = _eltype(ps[1]) _eltype(ps::Tuple) = _eltype(ps[1]) -function GradientZygote(F, x::AbstractArray{T}) where {T} - F = (_b, _a) -> (GeometricOptimizers._copyto!(_b, Zygote.gradient(loss, _a)[1])) - GeometricOptimizers.GradientFunction(F, x) +# Backward-compat alias: the old GML name for the abstract cache type. +const AbstractCache{T} = GeometricOptimizers.OptimizerCache{T} + +# A mutable gradient wrapper that returns a precomputed gradient. +# Subtypes SimpleSolvers.Gradient (accessible as GeometricOptimizers.Gradient) so it works +# with GeometricOptimizers' cache update! methods. +mutable struct _GMLGradient{T, VT} <: GeometricOptimizers.Gradient{T} + dp::VT # pre-computed Euclidean gradient (NamedTuple or AbstractArray) end +# When called with parameters, apply rgrad and return the Riemannian gradient. +(g::_GMLGradient{T})(x::GeometricOptimizers.ArrayNamedTuple{T}) where {T} = + GeometricOptimizers.apply_toNT(rgrad, x, g.dp) +(g::_GMLGradient{T})(x::AbstractArray{T}) where {T} = g.dp # euclidean: no rgrad needed -function GMLOptimizer(x::VT, problem::OptimizerProblem{T}; - options_kwargs...) where {T, VT <: OptimizerSolution{T}} - grad = GradientZygote(problem.F, x) - Optimizer(x, problem; algorithm = Adam(), linesearch = Static(0.01), - gradient = grad, options_kwargs...) +""" + Optimizer + +GML's neural-network optimizer. Wraps a GeometricOptimizers method together with its +corresponding cache, state, and retraction. +""" +mutable struct Optimizer{MT <: GeometricOptimizers.OptimizerMethod, CT, ST, RT} + method::MT + cache::CT + state::ST + retraction::RT + _grad::_GMLGradient # mutable gradient reference used by optimization_step! end -# type piracy. This should be put in `AbstractNeuralNetworks`! -function GeometricOptimizers.OptimizerProblem(nn::NeuralNetwork) - loss = NetworkLoss(nn) - _params = params(nn) - OptimizerProblem(ps -> loss(model, ps), NamedTuple{keys(_params)}(values(_params))) +function Optimizer(method::GeometricOptimizers.OptimizerMethod, nn::NeuralNetwork; + retraction = GeometricOptimizers.cayley) + ps = params(nn) + T = eltype(ps[1]) + cache = GeometricOptimizers.OptimizerCache(method, ps) + state = GeometricOptimizers.OptimizerState(method, ps) + grad = _GMLGradient{T, typeof(ps)}(ps) # dummy initial dp (gets overwritten) + Optimizer(method, cache, state, retraction, grad) end -function GeometricOptimizers.Optimizer( - algorithm::OptimizerMethod, nn::NeuralNetwork; options_kwargs...) - _params = params(nn) - GMLOptimizer( - NamedTuple{keys(_params)}(values(_params)), OptimizerProblem(nn); options_kwargs...) +# Convenience constructor that accepts a raw params NamedTuple directly. +function Optimizer(method::GeometricOptimizers.OptimizerMethod, ps::Union{NamedTuple, NeuralNetworkParameters}; + retraction = GeometricOptimizers.cayley) + T = eltype(ps[1]) + cache = GeometricOptimizers.OptimizerCache(method, ps) + state = GeometricOptimizers.OptimizerState(method, ps) + grad = _GMLGradient{T, typeof(ps)}(ps) # dummy initial dp (gets overwritten) + Optimizer(method, cache, state, retraction, grad) end + +""" + optimization_step!(opt, λY, ps, dp) + +Perform one optimizer step given a pre-computed Euclidean gradient `dp`. + +- `λY` — the `GlobalSection` of the current parameters `ps`. +- `ps` — the current parameter NamedTuple (modified in-place). +- `dp` — the Euclidean gradient returned by Zygote. +""" +function optimization_step!(opt::Optimizer, λY, ps, dp) + opt._grad.dp = dp # inject the pre-computed gradient + + # Step 1: update the cache (computes gradient in Lie algebra, then direction) + if opt.method isa GeometricOptimizers.Adam + GeometricOptimizers.update!(opt.cache, opt.state, opt._grad, opt.method, ps) + else + hess = GeometricOptimizers.NoHessian{eltype(ps[1])}() + GeometricOptimizers.update!(opt.cache, opt.state, opt._grad, hess, ps) + end + + # Step 2: apply retraction — update section(cache) from section(state) + direction + GeometricOptimizers.update_section!( + GeometricOptimizers.section(opt.cache), + GeometricOptimizers.section(opt.state), + GeometricOptimizers.direction(opt.cache), + opt.retraction + ) + + # Step 3: copy new manifold point from cache section → cache solution → ps and λY + GeometricOptimizers._copyto!(GeometricOptimizers.solution(opt.cache), + GeometricOptimizers.section(opt.cache)) + GeometricOptimizers._copyto!(ps, GeometricOptimizers.solution(opt.cache)) + GeometricOptimizers._copyto!(λY, GeometricOptimizers.section(opt.cache)) + + # Step 4: advance the state's section to the new position (needed for next step) + GeometricOptimizers.update_section!( + GeometricOptimizers.section(opt.state), + GeometricOptimizers.section(opt.state), + GeometricOptimizers.direction(opt.cache), + opt.retraction + ) + opt.state.iterations += 1 + nothing +end + +# Convenience: allow check(opt) to print nothing (backward compat) +check(::Optimizer) = nothing diff --git a/test/data_loader/data_loader_for_input_and_output.jl b/test/data_loader/data_loader_for_input_and_output.jl index 3e81aa4b4..0ee3c7f29 100644 --- a/test/data_loader/data_loader_for_input_and_output.jl +++ b/test/data_loader/data_loader_for_input_and_output.jl @@ -16,7 +16,7 @@ nepochs = 2000 model = Chain(Dense(1, nwidth), Dense(nwidth, 1)) nn = NeuralNetwork(model, Float32) dl = DataLoader(xsamples, ysamples) -o = Optimizer(AdamOptimizer(), nn) +o = Optimizer(Adam(), nn) batch = Batch(nbatch, 1, 1) loss = FeedForwardLoss() diff --git a/test/data_loader/optimizer_functor_with_adam.jl b/test/data_loader/optimizer_functor_with_adam.jl index 22a5c5c46..283b63a87 100644 --- a/test/data_loader/optimizer_functor_with_adam.jl +++ b/test/data_loader/optimizer_functor_with_adam.jl @@ -25,13 +25,14 @@ function test_optimization_with_adam(;T=Float32, dim₁=6, dim₂=6, n_images=10 # input dim is dim₁ / patch_length * dim₂ / pach_length; the transformer is called with dim₁ / patch_length and two layers model = Chain(Transformer(dl.input_dim, patch_length, 2; Stiefel=true), ClassificationLayer(dl.input_dim, 10, σ)) - ps = NeuralNetwork(model, CPU(), Float32).params + nn_obj = NeuralNetwork(model, CPU(), Float32) + ps = nn_obj.params loss = FeedForwardLoss() loss₁ = loss(model, ps, dl.input, dl.output) - opt = Optimizer(AdamOptimizer(), ps) + opt = Optimizer(Adam(), nn_obj) λY = GlobalSection(ps) loss_average = optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY) diff --git a/test/network_losses/losses_and_optimization.jl b/test/network_losses/losses_and_optimization.jl index 82f14e5f7..9ace7aca4 100644 --- a/test/network_losses/losses_and_optimization.jl +++ b/test/network_losses/losses_and_optimization.jl @@ -17,7 +17,7 @@ function train_network(; n_epochs=10) nn = setup_network(dl) loss = FeedForwardLoss() - o = Optimizer(AdamOptimizer(), nn) + o = Optimizer(Adam(), nn) batch = Batch(5, 1) loss_array = o(nn, dl, batch, n_epochs, loss) T = eltype(dl) diff --git a/test/sae_error_lower_than_psd_error.jl b/test/sae_error_lower_than_psd_error.jl index 2bd56a5a8..952c18a80 100644 --- a/test/sae_error_lower_than_psd_error.jl +++ b/test/sae_error_lower_than_psd_error.jl @@ -12,7 +12,7 @@ function test_accuracy(N::Integer, n::Integer; tol::Real = .35, n_epochs::Intege sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n; n_encoder_layers = 5, n_decoder_layers = 5)) - o = Optimizer(AdamOptimizer(), sae_nn) + o = Optimizer(Adam(), sae_nn) sae_error = o(sae_nn, dl, Batch(10), n_epochs)[end] @test sae_error < psd_error diff --git a/test/symplectic_autoencoder_tests.jl b/test/symplectic_autoencoder_tests.jl index bfd93a2f2..2006ef0f9 100644 --- a/test/symplectic_autoencoder_tests.jl +++ b/test/symplectic_autoencoder_tests.jl @@ -39,7 +39,7 @@ function test_symplecticity(N::Integer, n::Integer) # test if it's still symplectic after training dl = DataLoader(rand(N, 10 * N); autoencoder = true) - o = Optimizer(AdamOptimizer(), sae_nn) + o = Optimizer(Adam(), sae_nn) o(sae_nn, dl, Batch(10), 10) sympl_mat = jacobian(vec -> sae_decoder(vec), test_vector)[1] @test PoissonTensor(n) ≈ sympl_mat' * PoissonTensor(N) * sympl_mat diff --git a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl index 216ba0c64..774c94864 100644 --- a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl +++ b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl @@ -1,59 +1,60 @@ -using GeometricMachineLearning, Test +using GeometricMachineLearning, GeometricOptimizers, Test import Random, LinearAlgebra Random.seed!(1234) @doc raw""" -This checks if the Adam cache was set up in the right way +This checks if the Adam cache was set up in the right way. +AdamCache fields (from GeometricOptimizers): x, g, δ, Δg, m₁, m₂, m̃₂, section. +The direction δ and moments m₁, m₂ should be zero-initialised. """ -function check_adam_cache(C::AbstractCache{T}, tol= T(10) * eps(T)) where T - @test typeof(C) <: AdamCache - @test propertynames(C) == (:B₁, :B₂) - @test typeof(C.B₁) <: StiefelLieAlgHorMatrix - @test typeof(C.B₂) <: StiefelLieAlgHorMatrix - @test LinearAlgebra.norm(C.B₁) < tol - @test LinearAlgebra.norm(C.B₂) < tol -end +function check_adam_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * eps(T)) where T + @test C isa AdamCache + # direction and first/second moments are zero-initialised + @test typeof(C.δ) <: StiefelLieAlgHorMatrix + @test typeof(C.m₁) <: StiefelLieAlgHorMatrix + @test LinearAlgebra.norm(C.δ) < tol + @test LinearAlgebra.norm(C.m₁) < tol +end check_adam_cache(B::NamedTuple) = apply_toNT(check_adam_cache, B) @doc raw""" -This checks if the momentum cache was set up in the right way +This checks if the momentum cache was set up in the right way. +MomentumCache fields (from GeometricOptimizers): x, g, δ, Δg, section. +The direction δ should be zero-initialised. """ -function check_momentum_cache(C::AbstractCache{T}, tol= T(10) * eps(T)) where T - @test typeof(C) <: MomentumCache - @test propertynames(C) == (:B,) - @test typeof(C.B) <: StiefelLieAlgHorMatrix - @test LinearAlgebra.norm(C.B) < tol +function check_momentum_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * eps(T)) where T + @test C isa MomentumCache + @test typeof(C.δ) <: StiefelLieAlgHorMatrix + @test LinearAlgebra.norm(C.δ) < tol end check_momentum_cache(B::NamedTuple) = apply_toNT(check_momentum_cache, B) @doc raw""" -This checks if the gradient cache was set up in the right way +This checks if the gradient cache was set up in the right way. +GradientCache fields (from GeometricOptimizers): x, g, δ, Δg, section. """ -function check_gradient_cache(C::AbstractCache{T}) where T - @test typeof(C) <: GradientCache - @test propertynames(C) == () +function check_gradient_cache(C::GeometricOptimizers.OptimizerCache{T}) where T + @test C isa GradientCache + @test hasproperty(C, :δ) end check_gradient_cache(B::NamedTuple) = apply_toNT(check_gradient_cache, B) @doc raw""" This checks if all the caches are set up in the right way for the `MultiHeadAttention` layer with Stiefel weights. - -TODO: -- [ ] `BFGSOptimizer` !! """ function test_cache_setups_for_optimizer_for_multihead_attention_layer(T::Type, dim::Int, n_heads::Int) @assert dim % n_heads == 0 model = Chain(MultiHeadAttention(dim, n_heads, Stiefel=true)) ps = NeuralNetwork(model, CPU(), T).params - o₁ = Optimizer(AdamOptimizer(), ps) - o₂ = Optimizer(MomentumOptimizer(), ps) - o₃ = Optimizer(GradientOptimizer(), ps) + o₁ = Optimizer(Adam(), ps) + o₂ = Optimizer(MomentumMethod(), ps) + o₃ = Optimizer(GradientMethod(), ps) check_adam_cache(o₁.cache) check_momentum_cache(o₂.cache) check_gradient_cache(o₃.cache) end -test_cache_setups_for_optimizer_for_multihead_attention_layer(Float32, 64, 8) \ No newline at end of file +test_cache_setups_for_optimizer_for_multihead_attention_layer(Float32, 64, 8) diff --git a/test/transformer_related/multi_head_attention_stiefel_retraction.jl b/test/transformer_related/multi_head_attention_stiefel_retraction.jl index 400ab6478..9adc99e2c 100644 --- a/test/transformer_related/multi_head_attention_stiefel_retraction.jl +++ b/test/transformer_related/multi_head_attention_stiefel_retraction.jl @@ -1,7 +1,6 @@ -using GeometricMachineLearning, Test +using GeometricMachineLearning, GeometricOptimizers, Test using GeometricMachineLearning: geodesic using GeometricMachineLearning: cayley -using GeometricMachineLearning: init_optimizer_cache import Random, Test, LinearAlgebra, KernelAbstractions Random.seed!(1234) @@ -15,7 +14,7 @@ function check_retraction_geodesic(A::AbstractMatrix{T}, tol=eps(T)) where T @test LinearAlgebra.norm(A_retracted - StiefelProjection(A_retracted)) < tol end check_retraction_geodesic(cache::NamedTuple) = apply_toNT(check_retraction_geodesic, cache) -check_retraction_geodesic(B::MomentumCache) = check_retraction_geodesic(B.B) +check_retraction_geodesic(B::MomentumCache) = check_retraction_geodesic(B.δ) @doc raw""" This function computes the cayley retraction of an element of `StiefelLieAlgHorMatrix` and then checks if the resulting element is `StiefelProjection`. @@ -26,7 +25,7 @@ function check_retraction_cayley(A::AbstractMatrix{T}, tol=eps(T)) where T @test LinearAlgebra.norm(A_retracted - StiefelProjection(A_retracted)) < tol end check_retraction_cayley(cache::NamedTuple) = apply_toNT(check_retraction_cayley, cache) -check_retraction_cayley(B::MomentumCache) = check_retraction_cayley(B.B) +check_retraction_cayley(B::MomentumCache) = check_retraction_cayley(B.δ) @doc raw""" This is a test for that checks if the retractions (geodesic and Cayley for now) map from `StiefelLieAlgHorMatrix` to `StiefelManifold` when used with `MultiHeadAttention`. @@ -35,7 +34,7 @@ function test_multi_head_attention_retraction(T::Type, dim, n_heads, tol=eps(T), model = Chain(MultiHeadAttention(dim, n_heads, Stiefel=true)) ps = NeuralNetwork(model, backend, T).params - cache = init_optimizer_cache(MomentumOptimizer(), ps) + cache = GeometricOptimizers.OptimizerCache(MomentumMethod(), ps) check_retraction_geodesic(cache) diff --git a/test/transformer_related/multi_head_attention_stiefel_setup.jl b/test/transformer_related/multi_head_attention_stiefel_setup.jl index 67cf438a8..ae5f024ff 100644 --- a/test/transformer_related/multi_head_attention_stiefel_setup.jl +++ b/test/transformer_related/multi_head_attention_stiefel_setup.jl @@ -1,5 +1,4 @@ -using GeometricMachineLearning, Test -using GeometricMachineLearning: init_optimizer_cache +using GeometricMachineLearning, GeometricOptimizers, Test import Random, Test, LinearAlgebra, KernelAbstractions Random.seed!(1234) @@ -22,7 +21,7 @@ function check_grad_setup(B::AbstractMatrix{T}, tol=T(10)*eps(T)) where T @test LinearAlgebra.norm(B) < tol end check_grad_setup(gx::NamedTuple) = apply_toNT(check_grad_setup, gx) -check_grad_setup(B::MomentumCache) = check_grad_setup(B.B) +check_grad_setup(B::MomentumCache) = check_grad_setup(B.δ) @doc raw""" Check if `initialparameters` and `init_optimizer_cache` do the right thing for `MultiHeadAttentionLayer`. @@ -33,7 +32,7 @@ function check_multi_head_attention_stiefel_setup(T::Type, N::Int, n::Int) check_setup(ps) - gx = init_optimizer_cache(MomentumOptimizer(), ps) + gx = GeometricOptimizers.OptimizerCache(MomentumMethod(), ps) check_grad_setup(gx) end diff --git a/test/transformer_related/transformer_optimizer.jl b/test/transformer_related/transformer_optimizer.jl index d6907d577..b612e6847 100644 --- a/test/transformer_related/transformer_optimizer.jl +++ b/test/transformer_related/transformer_optimizer.jl @@ -5,39 +5,36 @@ import Random Random.seed!(1234) @doc raw""" -This function tests if the `GradientOptimzier`, `MomentumOptimizer`, `AdamOptimizer` and `BFGSOptimizer` act on the neural network weights via `optimization_step!`. +This function tests if the `GradientMethod`, `MomentumMethod`, and `Adam` act on the neural network weights via `optimization_step!`. """ function transformer_gradient_test(T, dim, n_heads, L, seq_length=8, batch_size=10) model = Chain(Transformer(dim, n_heads, L, Stiefel=true), ResNetLayer(dim)) model = Transformer(dim, n_heads, L, Stiefel=true) ps = NeuralNetwork(model, KernelAbstractions.CPU(), T).params - + input = rand(T, dim, seq_length, batch_size) - + loss(ps, input) = norm(model(input, ps)) dx = Zygote.gradient(ps -> loss(ps, input), ps)[1] - o₁ = Optimizer(GradientOptimizer(), ps) - o₂ = Optimizer(MomentumOptimizer(), ps) - o₃ = Optimizer(AdamOptimizer(), ps) - o₄ = Optimizer(BFGSOptimizer(), ps) + o₁ = Optimizer(GradientMethod(), ps) + o₂ = Optimizer(MomentumMethod(), ps) + o₃ = Optimizer(Adam(), ps) + # BFGSOptimizer is not available in GeometricOptimizers manifold optimizers ps₁ = deepcopy(ps) ps₂ = deepcopy(ps) ps₃ = deepcopy(ps) - ps₄ = deepcopy(ps) λY₁ = GlobalSection(ps₁) λY₂ = GlobalSection(ps₂) λY₃ = GlobalSection(ps₃) - λY₄ = GlobalSection(ps₄) optimization_step!(o₁, λY₁, ps₁, dx) optimization_step!(o₂, λY₂, ps₂, dx) optimization_step!(o₃, λY₃, ps₃, dx) - optimization_step!(o₄, λY₄, ps₄, dx) - @test typeof(ps₁) == typeof(ps₂) == typeof(ps₃) == typeof(ps₄) == typeof(ps) - @test ps₁[1].PQ.head_1 ≉ ps₂[1].PQ.head_1 ≉ ps₃[1].PQ.head_1 ≉ ps₄[1].PQ.head_1 # ≉ ps[1].PQ.head_1 + @test typeof(ps₁) == typeof(ps₂) == typeof(ps₃) == typeof(ps) + @test ps₁[1].PQ.head_1 ≉ ps₂[1].PQ.head_1 ≉ ps₃[1].PQ.head_1 # ≉ ps[1].PQ.head_1 end transformer_gradient_test(Float32, 10, 5, 4) \ No newline at end of file From 71d07283d7dab39598db694f1869dfe0f90d8a7e Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 8 Jul 2026 12:29:16 +0200 Subject: [PATCH 3/7] Switch GeometricOptimizers to a path dependency GeometricOptimizers is not yet registered in the General registry; it is co-developed alongside GML as a sibling repository. Replace the version-constrained `[compat]` entry with a `[sources]` path entry so that `] instantiate` resolves it from `../GeometricOptimizers` without needing a registry lookup. Also move `[weakdeps]`/`[extensions]` to appear after `[deps]` (the canonical TOML ordering) and sort the `[extras]` block alphabetically. Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index e12f109fb..0f0d53a0d 100644 --- a/Project.toml +++ b/Project.toml @@ -3,12 +3,6 @@ uuid = "194d25b2-d3f5-49f0-af24-c124f4aa80cc" version = "0.4.8" authors = ["Michael Kraus "] -[weakdeps] -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" - -[extensions] -HDF5Ext = "HDF5" - [deps] AbstractNeuralNetworks = "60874f82-5ada-4c70-bd1c-fa6be7711c8a" BandedMatrices = "aae01518-5342-5314-be14-df237901396f" @@ -38,6 +32,15 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" +[weakdeps] +HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" + +[sources] +GeometricOptimizers = {path = "../GeometricOptimizers"} + +[extensions] +HDF5Ext = "HDF5" + [compat] AbstractNeuralNetworks = "0.5, 0.6" BandedMatrices = "1" @@ -48,7 +51,6 @@ Distances = "0.10" ForwardDiff = "0.10, 1" GeometricBase = "0.14" GeometricEquations = "0.21" -GeometricOptimizers = "1.0.0" GeometricSolutions = "0.6" HDF5 = "0.16, 0.17" KernelAbstractions = "0.9" @@ -66,13 +68,13 @@ julia = "1.9" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" GeometricEquations = "c85262ba-a08a-430a-b926-d29770767bf2" GeometricIntegrators = "dcce2d33-59f6-5b8d-9047-0defad88ae06" GeometricProblems = "18cb22b4-ad41-5c80-9c5f-710df63fbdc9" +HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RungeKutta = "fb486d5c-30a0-4a8a-8415-a8b4ace5a6f7" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" From 8f5781d9370cbc37581c38dcbba3024c57a96ad7 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 8 Jul 2026 12:29:30 +0200 Subject: [PATCH 4/7] Replace GML optimizer internals with GeometricOptimizers backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old GML optimizer code assumed a flat, homogeneous parameter structure that matched GeometricOptimizers' OptimizerSolution type. Real networks have nested NamedTuple parameters that mix manifold weights (StiefelManifold, GrassmannManifold) with plain matrices and bias vectors, which GeometricOptimizers cannot handle directly. Key changes: * `_make_optimizer_cache` / `_make_optimizer_state` – recursive tree builders that descend through NamedTuple/NeuralNetworkParameters until they hit a leaf. Leaves that are OptimizerSolution and use a GO-native method (Gradient/Momentum/Adam) get a proper GeometricOptimizers cache; everything else gets a lightweight `GMLEuclideanState`. * `GMLEuclideanState` – mutable struct that stores iteration count plus first and second moments so that gradient, momentum, and Adam updates can be applied to plain AbstractArray parameters that are outside GO's type hierarchy. * `_tree_optim_step!` / `_leaf_optim_step!` – matching recursive dispatchers that call either the GO update path or the plain Euclidean path at each leaf. * `Optimizer` struct – gains `step_size::Float64` and `iterations::Int` fields (replacing the embedded `_GMLGradient` ref) so that decaying step sizes can be computed at call time without touching the method. * `AdamOptimizerWithDecay` – re-implemented (was deleted with the old `src/optimizers/` tree) as a first-class `OptimizerMethod` subtype with exponential decay η(t) = η₁ · γ^t. * `GeometricOptimizers.GlobalSection(ps::NeuralNetworkParameters)` – extension method so that callers can pass NeuralNetworkParameters directly without manually unwrapping the inner NamedTuple. * Export `AdamOptimizerWithDecay` from the top-level module. Co-Authored-By: Claude Sonnet 4.6 --- src/GeometricMachineLearning.jl | 1 + src/utils.jl | 231 +++++++++++++++++++++++--------- 2 files changed, 168 insertions(+), 64 deletions(-) diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 375c8db10..435b92115 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -189,6 +189,7 @@ const GradientOptimizer = GradientMethod const MomentumOptimizer = MomentumMethod const AdamOptimizer = Adam export GradientOptimizer, MomentumOptimizer, AdamOptimizer +export AdamOptimizerWithDecay #INCLUDE ABSTRACT TRAINING integrator export AbstractTrainingMethod diff --git a/src/utils.jl b/src/utils.jl index c3a1a9031..3ea8596ec 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -172,99 +172,202 @@ Base.:≈(qp₁::QPT, qp₂::QPT) = (qp₁.q ≈ qp₂.q) & (qp₁.p ≈ qp₂.p _eltype(x) = eltype(x) _eltype(ps::NamedTuple) = _eltype(ps[1]) _eltype(ps::Tuple) = _eltype(ps[1]) +_eltype(ps::NeuralNetworkParameters) = _eltype(params(ps)[1]) -# Backward-compat alias: the old GML name for the abstract cache type. +# Extend GlobalSection so it works with NeuralNetworkParameters (wraps a NamedTuple). +GeometricOptimizers.GlobalSection(ps::NeuralNetworkParameters) = + GeometricOptimizers.GlobalSection(params(ps)) + +# Backward-compat alias const AbstractCache{T} = GeometricOptimizers.OptimizerCache{T} -# A mutable gradient wrapper that returns a precomputed gradient. -# Subtypes SimpleSolvers.Gradient (accessible as GeometricOptimizers.Gradient) so it works -# with GeometricOptimizers' cache update! methods. +# Gradient wrapper: stores a pre-computed Euclidean gradient and applies rgrad on manifolds. mutable struct _GMLGradient{T, VT} <: GeometricOptimizers.Gradient{T} - dp::VT # pre-computed Euclidean gradient (NamedTuple or AbstractArray) + dp::VT end -# When called with parameters, apply rgrad and return the Riemannian gradient. (g::_GMLGradient{T})(x::GeometricOptimizers.ArrayNamedTuple{T}) where {T} = GeometricOptimizers.apply_toNT(rgrad, x, g.dp) -(g::_GMLGradient{T})(x::AbstractArray{T}) where {T} = g.dp # euclidean: no rgrad needed +(g::_GMLGradient{T})(x::AbstractArray{T}) where {T} = g.dp -""" - Optimizer +# State for Euclidean (non-manifold) parameters. +mutable struct GMLEuclideanState{T, AT<:AbstractArray{T}} + iterations::Int + m₁::AT + m₂::AT +end +GMLEuclideanState(x::AbstractArray{T}) where T = + GMLEuclideanState{T, typeof(x)}(0, zero(x), zero(x)) + +# Adam with exponential learning-rate decay. +struct AdamOptimizerWithDecay{T<:Real} <: GeometricOptimizers.OptimizerMethod + η₁::T; η₂::T; ρ₁::T; ρ₂::T; δ::T; γ::T; n_epochs::Int + function AdamOptimizerWithDecay(n_epochs::Int, η₁=1f-2, η₂=1f-6, + ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8; T=typeof(η₁)) + γ = exp(log(η₂/η₁) / n_epochs) + new{T}(T(η₁), T(η₂), T(ρ₁), T(ρ₂), T(δ), T(γ), n_epochs) + end +end + +_is_go_native_method(::GeometricOptimizers.GradientMethod) = true +_is_go_native_method(::GeometricOptimizers.MomentumMethod) = true +_is_go_native_method(::GeometricOptimizers.Adam) = true +_is_go_native_method(::GeometricOptimizers.OptimizerMethod) = false + +_use_go_cache(method, x) = + _is_go_native_method(method) && x isa GeometricOptimizers.OptimizerSolution + +function _make_optimizer_cache(method, x) + if _use_go_cache(method, x) + GeometricOptimizers.OptimizerCache(method, x) + elseif x isa NamedTuple || x isa NeuralNetworkParameters + NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x))) + else + GMLEuclideanState(x) + end +end + +function _make_optimizer_state(method, x) + if _use_go_cache(method, x) + GeometricOptimizers.OptimizerState(method, x) + elseif x isa NamedTuple || x isa NeuralNetworkParameters + NamedTuple{keys(x)}(Tuple(_make_optimizer_state(method, x[k]) for k in keys(x))) + else + GMLEuclideanState(x) + end +end -GML's neural-network optimizer. Wraps a GeometricOptimizers method together with its -corresponding cache, state, and retraction. -""" mutable struct Optimizer{MT <: GeometricOptimizers.OptimizerMethod, CT, ST, RT} method::MT cache::CT state::ST retraction::RT - _grad::_GMLGradient # mutable gradient reference used by optimization_step! + step_size::Float64 + iterations::Int end +_default_step_size(method::GeometricOptimizers.Adam) = Float64(method.η) +_default_step_size(method::AdamOptimizerWithDecay) = Float64(method.η₁) +_default_step_size(::GeometricOptimizers.OptimizerMethod) = 1e-2 + +_current_step_size(opt::Optimizer, ::Int) = opt.step_size +_current_step_size(opt::Optimizer{<:AdamOptimizerWithDecay}, t::Int) = + Float64(opt.method.η₁ * opt.method.γ^t) + function Optimizer(method::GeometricOptimizers.OptimizerMethod, nn::NeuralNetwork; - retraction = GeometricOptimizers.cayley) + retraction = GeometricOptimizers.cayley, + step_size::Real = _default_step_size(method)) ps = params(nn) - T = eltype(ps[1]) - cache = GeometricOptimizers.OptimizerCache(method, ps) - state = GeometricOptimizers.OptimizerState(method, ps) - grad = _GMLGradient{T, typeof(ps)}(ps) # dummy initial dp (gets overwritten) - Optimizer(method, cache, state, retraction, grad) + Optimizer(method, _make_optimizer_cache(method, ps), _make_optimizer_state(method, ps), + retraction, Float64(step_size), 0) end -# Convenience constructor that accepts a raw params NamedTuple directly. -function Optimizer(method::GeometricOptimizers.OptimizerMethod, ps::Union{NamedTuple, NeuralNetworkParameters}; - retraction = GeometricOptimizers.cayley) - T = eltype(ps[1]) - cache = GeometricOptimizers.OptimizerCache(method, ps) - state = GeometricOptimizers.OptimizerState(method, ps) - grad = _GMLGradient{T, typeof(ps)}(ps) # dummy initial dp (gets overwritten) - Optimizer(method, cache, state, retraction, grad) +function Optimizer(method::GeometricOptimizers.OptimizerMethod, + ps::Union{NamedTuple, NeuralNetworkParameters}; + retraction = GeometricOptimizers.cayley, + step_size::Real = _default_step_size(method)) + Optimizer(method, _make_optimizer_cache(method, ps), _make_optimizer_state(method, ps), + retraction, Float64(step_size), 0) end -""" - optimization_step!(opt, λY, ps, dp) +# Euclidean update rules +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState, ::GeometricOptimizers.GradientMethod, step_size) where T + x .-= T(step_size) .* dx +end +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState{T}, method::GeometricOptimizers.MomentumMethod, step_size) where T + x .-= T(step_size) .* (dx .+ state.m₁) + state.m₁ .+= T(method.α) .* dx +end +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState{T}, method::GeometricOptimizers.Adam, step_size) where T + t = state.iterations; _t = t + 1 + β₁, β₂, δ = T(method.β₁), T(method.β₂), T(method.δ) + fac₁₁ = β₁/(1-β₁^_t); fac₁₂ = (1-β₁)/(1-β₁^_t) + fac₂₁ = β₂/(1-β₂^_t); fac₂₂ = (1-β₂)/(1-β₂^_t) + state.m₁ .= fac₁₁ .* state.m₁ .+ fac₁₂ .* dx + state.m₂ .= fac₂₁ .* state.m₂ .+ fac₂₂ .* dx .^ 2 + x .-= T(step_size) .* state.m₁ ./ (sqrt.(state.m₂) .+ δ) +end +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState{T}, method::AdamOptimizerWithDecay, step_size) where T + t = state.iterations; _t = t + 1 + ρ₁, ρ₂, δ = T(method.ρ₁), T(method.ρ₂), T(method.δ) + fac₁₁ = ρ₁/(1-ρ₁^_t); fac₁₂ = (1-ρ₁)/(1-ρ₁^_t) + fac₂₁ = ρ₂/(1-ρ₂^_t); fac₂₂ = (1-ρ₂)/(1-ρ₂^_t) + state.m₁ .= fac₁₁ .* state.m₁ .+ fac₁₂ .* dx + state.m₂ .= fac₂₁ .* state.m₂ .+ fac₂₂ .* dx .^ 2 + x .-= T(step_size) .* state.m₁ ./ (sqrt.(state.m₂) .+ δ) +end -Perform one optimizer step given a pre-computed Euclidean gradient `dp`. +# GO-managed leaf step (manifolds, vectors, ArrayNamedTuples) +function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, + state::GeometricOptimizers.OptimizerState, + dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) + T = _eltype(ps_leaf) + local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf) + if method isa GeometricOptimizers.Adam + GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) + else + GeometricOptimizers.update!(cache, state, local_grad, + GeometricOptimizers.NoHessian{T}(), ps_leaf) + end + GeometricOptimizers._rmul!(GeometricOptimizers.direction(cache), step_size) + GeometricOptimizers.update_section!(GeometricOptimizers.section(cache), + GeometricOptimizers.section(state), + GeometricOptimizers.direction(cache), + retraction) + GeometricOptimizers._copyto!(GeometricOptimizers.solution(cache), + GeometricOptimizers.section(cache)) + GeometricOptimizers._copyto!(ps_leaf, GeometricOptimizers.solution(cache)) + GeometricOptimizers._copyto!(λY_leaf, GeometricOptimizers.section(cache)) + GeometricOptimizers.update_section!(GeometricOptimizers.section(state), + GeometricOptimizers.section(state), + GeometricOptimizers.direction(cache), + retraction) + if state isa GeometricOptimizers.AdamState + GeometricOptimizers._copyto!(GeometricOptimizers.first_moment(state), + GeometricOptimizers.first_moment(cache)) + GeometricOptimizers._copyto!(GeometricOptimizers.second_moment(state), + GeometricOptimizers.second_moment(cache)) + elseif state isa GeometricOptimizers.MomentumState + GeometricOptimizers._add!(GeometricOptimizers.momentum(state), + GeometricOptimizers._mul(method.α, + GeometricOptimizers.gradient_array(cache))) + end + state.iterations += 1 + nothing +end -- `λY` — the `GlobalSection` of the current parameters `ps`. -- `ps` — the current parameter NamedTuple (modified in-place). -- `dp` — the Euclidean gradient returned by Zygote. -""" -function optimization_step!(opt::Optimizer, λY, ps, dp) - opt._grad.dp = dp # inject the pre-computed gradient +# Euclidean leaf step (plain AbstractArray params) +function _leaf_optim_step!(cache::GMLEuclideanState, state::GMLEuclideanState, + dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) + _euclidean_update!(ps_leaf, dp_leaf, state, method, step_size) + state.iterations += 1 + nothing +end - # Step 1: update the cache (computes gradient in Lie algebra, then direction) - if opt.method isa GeometricOptimizers.Adam - GeometricOptimizers.update!(opt.cache, opt.state, opt._grad, opt.method, ps) +# Recursive dispatcher over the parameter tree +function _tree_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) + if caches isa NamedTuple + for k in keys(caches) + dp_k = dp[k] + dp_k === nothing && continue + λY_k = λY isa NamedTuple ? λY[k] : λY + _tree_optim_step!(caches[k], states[k], dp_k, ps[k], λY_k, + method, retraction, step_size) + end else - hess = GeometricOptimizers.NoHessian{eltype(ps[1])}() - GeometricOptimizers.update!(opt.cache, opt.state, opt._grad, hess, ps) + _leaf_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) end + nothing +end - # Step 2: apply retraction — update section(cache) from section(state) + direction - GeometricOptimizers.update_section!( - GeometricOptimizers.section(opt.cache), - GeometricOptimizers.section(opt.state), - GeometricOptimizers.direction(opt.cache), - opt.retraction - ) - - # Step 3: copy new manifold point from cache section → cache solution → ps and λY - GeometricOptimizers._copyto!(GeometricOptimizers.solution(opt.cache), - GeometricOptimizers.section(opt.cache)) - GeometricOptimizers._copyto!(ps, GeometricOptimizers.solution(opt.cache)) - GeometricOptimizers._copyto!(λY, GeometricOptimizers.section(opt.cache)) - - # Step 4: advance the state's section to the new position (needed for next step) - GeometricOptimizers.update_section!( - GeometricOptimizers.section(opt.state), - GeometricOptimizers.section(opt.state), - GeometricOptimizers.direction(opt.cache), - opt.retraction - ) - opt.state.iterations += 1 +function optimization_step!(opt::Optimizer, λY, ps, dp) + step = _current_step_size(opt, opt.iterations) + _tree_optim_step!(opt.cache, opt.state, dp, ps, λY, opt.method, opt.retraction, step) + opt.iterations += 1 nothing end -# Convenience: allow check(opt) to print nothing (backward compat) check(::Optimizer) = nothing From abe304c9f8520c6212fffa30809797816947ec5e Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 8 Jul 2026 12:29:38 +0200 Subject: [PATCH 5/7] Update optimizer tests to use new GeometricOptimizers API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * `bfgs_optimizer.jl` – BFGSOptimizer no longer exists in GML (it was part of the deleted src/optimizers/ tree and has no GO equivalent yet). Rewrite both test cases as gradient-descent convergence tests using the new `Optimizer(GradientMethod(), ps; step_size=…)` + `optimization_step!` API, covering both Euclidean (plain matrix) and Stiefel-manifold parameters. * `psd_optim.jl`, `svd_optim.jl` – replace the old constructor calls (`GradientOptimizer(η)`, `MomentumOptimizer(η)`, `AdamOptimizer(η)`) with their GeometricOptimizers equivalents (`GradientMethod()`, `MomentumMethod()`, `Adam(η)`). Step size is now passed via the `step_size` keyword for gradient/momentum methods. Co-Authored-By: Claude Sonnet 4.6 --- test/optimizers/bfgs_optimizer.jl | 98 +++++-------------- .../optimizer_convergence_tests/psd_optim.jl | 6 +- .../optimizer_convergence_tests/svd_optim.jl | 6 +- 3 files changed, 33 insertions(+), 77 deletions(-) diff --git a/test/optimizers/bfgs_optimizer.jl b/test/optimizers/bfgs_optimizer.jl index 4f34bd550..e4bdb05fb 100644 --- a/test/optimizers/bfgs_optimizer.jl +++ b/test/optimizers/bfgs_optimizer.jl @@ -1,5 +1,4 @@ using GeometricMachineLearning -using GeometricMachineLearning: update_section! using Zygote using Test using LinearAlgebra: norm @@ -8,87 +7,44 @@ import Random Random.seed!(123) @doc raw""" - bfgs_optimizer(N) - -Test if BFGS optimizer perfroms better than gradient optimizer. - -The test is performed on a simple loss function -```math - \mathrm{loss}(A) = norm(A - B) ^ 3, -``` -where ``B`` is fixed. +Test that gradient descent reduces loss on a Euclidean (plain matrix) problem. +BFGSOptimizer is no longer available; this file retains the structure for future extension. """ -function bfgs_optimizer(N; n_steps = 10, η = 1e-4) +function gradient_optimizer_euclidean(N; n_steps = 20, step_size = 1e-3) B = inv(rand(N, N)) - loss(A) = norm(A - B) ^ (2) - A = randn(N, N) - loss1 = loss(A) - method₁ = GradientOptimizer(η) - o₁ = Optimizer(method₁, (A = A,)) - for _ in 1:n_steps - ∇L = Zygote.gradient(loss, A)[1] - update!(o₁, o₁.cache.A, ∇L) - A .+= ∇L - end - loss2 = loss(A) + loss(ps) = norm(ps.A - B) ^ 2 A = randn(N, N) - method₂ = BFGSOptimizer(η) - o₂ = Optimizer(method₂, (A = A,)) + ps = (A = A,) + loss1 = loss(ps) + o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps - ∇L = Zygote.gradient(loss, A)[1] - update!(o₂, o₂.cache.A, ∇L) - A .+= ∇L + ∇L = Zygote.gradient(loss, ps)[1] + λY = GlobalSection(ps) + optimization_step!(o, λY, ps, ∇L) end - loss3 = loss(A) - @test loss1 > loss2 > loss3 - println(loss2) - println(loss3) - + loss2 = loss(ps) + @test loss1 > loss2 end -bfgs_optimizer(10) - @doc raw""" - bfgs_optimizer(N, n) - -Test if BFGS optimizer perfroms better than gradient optimizer. - -The test is performed on a simple loss function -```math - \mathrm{loss}(A) = norm(AA^T - B) ^ 3, -``` -where ``B = Y_BY_B^T`` for some ``Y\in{}St(n, N)`` is fixed. -``A`` in the equation above is optimized on the Stiefel manifold. +Test that gradient descent reduces loss on the Stiefel manifold. """ -function bfgs_stiefel_optimizer(N, n; n_steps = 10, η = 1e-4) +function gradient_optimizer_stiefel(N, n; n_steps = 20, step_size = 1e-3) YB = rand(StiefelManifold, N, n) - B = YB * YB' - loss(A) = norm(A * A' - B) ^ 2 - Y = rand(StiefelManifold, N, n) - λY = GlobalSection(Y) - loss1 = loss(Y) - method₁ = GradientOptimizer(η) - o₁ = Optimizer(method₁, (A = Y,)) - for _ in 1:n_steps - ∇L = Zygote.gradient(loss, Y)[1] - gradL = global_rep(λY, ∇L) - update!(o₁, o₁.cache.A, gradL) - update_section!(λY, gradL, cayley) - end - loss2 = loss(Y) - Y = rand(StiefelManifold, N, n) - method₂ = BFGSOptimizer(η) - o₂ = Optimizer(method₂, (A = Y,)) + B = YB * YB' + loss(ps) = norm(ps.Y * ps.Y' - B) ^ 2 + Y = rand(StiefelManifold, N, n) + ps = (Y = Y,) + loss1 = loss(ps) + o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps - ∇L = Zygote.gradient(loss, Y)[1] - gradL = global_rep(λY, ∇L) - update!(o₂, o₂.cache.A, gradL) - update_section!(λY, gradL, cayley) + ∇L = Zygote.gradient(loss, ps)[1] + λY = GlobalSection(ps) + optimization_step!(o, λY, ps, ∇L) end - loss3 = loss(Y) - @test loss1 > loss2 > loss3 - println(loss2) - println(loss3) + loss2 = loss(ps) + @test loss1 > loss2 end -bfgs_stiefel_optimizer(10, 5) \ No newline at end of file +gradient_optimizer_euclidean(10) +gradient_optimizer_stiefel(10, 5) diff --git a/test/optimizers/optimizer_convergence_tests/psd_optim.jl b/test/optimizers/optimizer_convergence_tests/psd_optim.jl index 6135a83c2..c926a2154 100644 --- a/test/optimizers/optimizer_convergence_tests/psd_optim.jl +++ b/test/optimizers/optimizer_convergence_tests/psd_optim.jl @@ -38,9 +38,9 @@ function svd_test(A, n, train_steps=1000, tol=1e-1; retraction=cayley) model = Chain(PSDLayer(2*N, 2*n), PSDLayer(2*n, 2*N)) ps = NeuralNetwork(model, CPU(), Float64).params - o₁ = Optimizer(GradientOptimizer(0.01), ps; retraction = retraction) - o₂ = Optimizer(MomentumOptimizer(0.01), ps; retraction = retraction) - o₃ = Optimizer(AdamOptimizer(0.01), ps; retraction = retraction) + o₁ = Optimizer(GradientMethod(), ps; retraction = retraction, step_size = 0.01) + o₂ = Optimizer(MomentumMethod(), ps; retraction = retraction, step_size = 0.01) + o₃ = Optimizer(Adam(0.01), ps; retraction = retraction) U₁, Ũ₁, err₁ = train_network!(o₁, model, deepcopy(ps), A, train_steps, tol) U₂, Ũ₂, err₂ = train_network!(o₂, model, deepcopy(ps), A, train_steps, tol) diff --git a/test/optimizers/optimizer_convergence_tests/svd_optim.jl b/test/optimizers/optimizer_convergence_tests/svd_optim.jl index e260c40ff..378bd002b 100644 --- a/test/optimizers/optimizer_convergence_tests/svd_optim.jl +++ b/test/optimizers/optimizer_convergence_tests/svd_optim.jl @@ -34,9 +34,9 @@ function svd_test(A, n, train_steps=1000, tol=1e-1; retraction=cayley) model = Chain(StiefelLayer(N, n), StiefelLayer(n, N)) ps = NeuralNetwork(model, CPU(), Float64).params - o₁ = Optimizer(GradientOptimizer(0.01), ps; retraction = retraction) - o₂ = Optimizer(MomentumOptimizer(0.01), ps; retraction = retraction) - o₃ = Optimizer(AdamOptimizer(0.01), ps; retraction = retraction) + o₁ = Optimizer(GradientMethod(), ps; retraction = retraction, step_size = 0.01) + o₂ = Optimizer(MomentumMethod(), ps; retraction = retraction, step_size = 0.01) + o₃ = Optimizer(Adam(0.01), ps; retraction = retraction) U₁, Ũ₁, err₁ = train_network!(o₁, model, deepcopy(ps), A, train_steps, tol) U₂, Ũ₂, err₂ = train_network!(o₂, model, deepcopy(ps), A, train_steps, tol) From b7d7edae234405ed5164bca3ee0ac3938eed295f Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 8 Jul 2026 12:29:48 +0200 Subject: [PATCH 6/7] Fix Stiefel MHA tests for new recursive parameter tree MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MultiHeadAttention layer with Stiefel=true now produces an ArrayNamedTuple of StiefelManifold matrices (one per head) rather than a single matrix. Caches therefore contain a NamedTuple of StiefelLieAlgHorMatrix values, not a bare SLAHM. * Add `_check_slahm_zero` with a NamedTuple overload so that `check_adam_cache` and `check_momentum_cache` dispatch correctly regardless of whether a cache field is a bare SLAHM or a NamedTuple of them. Remove the now-incorrect `typeof(C.δ) <: StiefelLieAlgHorMatrix` assertions. * Replace direct `GeometricOptimizers.OptimizerCache(method, ps)` calls in the setup and retraction tests with `Optimizer(method, ps).cache`. The bare GO constructor does not accept NeuralNetworkParameters; the GML Optimizer constructor handles the recursive tree building. Co-Authored-By: Claude Sonnet 4.6 --- .../multi_head_attention_stiefel_optim_cache.jl | 14 +++++++------- .../multi_head_attention_stiefel_retraction.jl | 2 +- .../multi_head_attention_stiefel_setup.jl | 2 +- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl index 774c94864..b0338cac2 100644 --- a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl +++ b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl @@ -8,13 +8,14 @@ This checks if the Adam cache was set up in the right way. AdamCache fields (from GeometricOptimizers): x, g, δ, Δg, m₁, m₂, m̃₂, section. The direction δ and moments m₁, m₂ should be zero-initialised. """ +_check_slahm_zero(A::StiefelLieAlgHorMatrix{T}, tol) where T = + (@test typeof(A) <: StiefelLieAlgHorMatrix; @test LinearAlgebra.norm(A) < tol) +_check_slahm_zero(A::NamedTuple, tol) = foreach(v -> _check_slahm_zero(v, tol), values(A)) + function check_adam_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * eps(T)) where T @test C isa AdamCache - # direction and first/second moments are zero-initialised - @test typeof(C.δ) <: StiefelLieAlgHorMatrix - @test typeof(C.m₁) <: StiefelLieAlgHorMatrix - @test LinearAlgebra.norm(C.δ) < tol - @test LinearAlgebra.norm(C.m₁) < tol + _check_slahm_zero(C.δ, tol) + _check_slahm_zero(C.m₁, tol) end check_adam_cache(B::NamedTuple) = apply_toNT(check_adam_cache, B) @@ -25,8 +26,7 @@ The direction δ should be zero-initialised. """ function check_momentum_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * eps(T)) where T @test C isa MomentumCache - @test typeof(C.δ) <: StiefelLieAlgHorMatrix - @test LinearAlgebra.norm(C.δ) < tol + _check_slahm_zero(C.δ, tol) end check_momentum_cache(B::NamedTuple) = apply_toNT(check_momentum_cache, B) diff --git a/test/transformer_related/multi_head_attention_stiefel_retraction.jl b/test/transformer_related/multi_head_attention_stiefel_retraction.jl index 9adc99e2c..6c2840213 100644 --- a/test/transformer_related/multi_head_attention_stiefel_retraction.jl +++ b/test/transformer_related/multi_head_attention_stiefel_retraction.jl @@ -34,7 +34,7 @@ function test_multi_head_attention_retraction(T::Type, dim, n_heads, tol=eps(T), model = Chain(MultiHeadAttention(dim, n_heads, Stiefel=true)) ps = NeuralNetwork(model, backend, T).params - cache = GeometricOptimizers.OptimizerCache(MomentumMethod(), ps) + cache = Optimizer(MomentumMethod(), ps).cache check_retraction_geodesic(cache) diff --git a/test/transformer_related/multi_head_attention_stiefel_setup.jl b/test/transformer_related/multi_head_attention_stiefel_setup.jl index ae5f024ff..387c8675a 100644 --- a/test/transformer_related/multi_head_attention_stiefel_setup.jl +++ b/test/transformer_related/multi_head_attention_stiefel_setup.jl @@ -32,7 +32,7 @@ function check_multi_head_attention_stiefel_setup(T::Type, N::Int, n::Int) check_setup(ps) - gx = GeometricOptimizers.OptimizerCache(MomentumMethod(), ps) + gx = Optimizer(MomentumMethod(), ps).cache check_grad_setup(gx) end From 5f7e648a1f484aea6f200c28093c4a8f69e78056 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 8 Jul 2026 15:37:17 +0200 Subject: [PATCH 7/7] Add GO bridge methods so GML manifold types work with GeometricOptimizers GML and GO define parallel type hierarchies (both have StiefelManifold, SkewSymMatrix, etc.) that do not share a common supertype, so GO's dispatch methods never fire for GML's types. This commit adds explicit method extensions to bridge the gap: - Base.copy/similar/fill!/copyto! for StiefelManifold, GrassmannManifold, SkewSymMatrix, StiefelLieAlgHorMatrix, GrassmannLieAlgHorMatrix - GeometricOptimizers.global_rep, update_section!, cayley, _copyto! for GML's Stiefel and Grassmann types - GeometricOptimizers._add!, _rac!, _square!, _div! bridges for SkewSymMatrix and StiefelLieAlgHorMatrix (GO dispatches on module hierarchy; GML types need explicit overloads) - GrassmannLieAlgHorMatrix.fill!, copyto! and GO arithmetic bridges - Fix zeros(SkewSymMatrix{T}, n) dispatch: the old Type{SkewSymMatrix{<:Real}} pattern is invariant and never matches concrete SkewSymMatrix{Float64}; changed to the standard where T form - utils.jl: _adapt_method_to_T, _use_go_cache, _make_optimizer_cache/state, _leaf_optim_step!, _tree_optim_step! to route Stiefel/Grassmann parameters through GO's manifold optimisation path Co-Authored-By: Claude Sonnet 4.6 --- .../grassmann_lie_algebra_horizontal.jl | 5 ++ src/arrays/skew_symmetric.jl | 8 +- src/arrays/stiefel_lie_algebra_horizontal.jl | 4 +- src/data_loader/batch.jl | 6 +- src/manifolds/grassmann_manifold.jl | 54 ++++++++++++ src/manifolds/stiefel_manifold.jl | 84 +++++++++++++++++++ src/utils.jl | 15 +++- 7 files changed, 165 insertions(+), 11 deletions(-) diff --git a/src/arrays/grassmann_lie_algebra_horizontal.jl b/src/arrays/grassmann_lie_algebra_horizontal.jl index cb7414c63..0e5e455bb 100644 --- a/src/arrays/grassmann_lie_algebra_horizontal.jl +++ b/src/arrays/grassmann_lie_algebra_horizontal.jl @@ -197,4 +197,9 @@ function _round(B::GrassmannLieAlgHorMatrix; kwargs...) B.N, B.n ) +end + +function Base.copyto!(A::GrassmannLieAlgHorMatrix, B::GrassmannLieAlgHorMatrix) + copyto!(A.B, B.B) + A end \ No newline at end of file diff --git a/src/arrays/skew_symmetric.jl b/src/arrays/skew_symmetric.jl index f3ae66e4d..440928bcb 100644 --- a/src/arrays/skew_symmetric.jl +++ b/src/arrays/skew_symmetric.jl @@ -145,8 +145,8 @@ end Base.:*(α::Real, A::SkewSymMatrix) = A*α -function Base.zeros(ST::Type{SkewSymMatrix{<:Real}}, n::Int) - zeros(CPU(), ST, n) +function Base.zeros(::Type{SkewSymMatrix{T}}, n::Int) where T + zeros(CPU(), SkewSymMatrix{T}, n) end function Base.zeros(backend::KernelAbstractions.Backend, ::Type{SkewSymMatrix{T}}, n::Int) where T @@ -342,7 +342,9 @@ function _round(A::AbstractArray; kwargs...) round.(A; kwargs...) end -# define routines for generalizing ChainRulesCore to SkewSymMatrix +Base.fill!(A::SkewSymMatrix, val) = (fill!(A.S, val); A) + +# define routines for generalizing ChainRulesCore to SkewSymMatrix ChainRulesCore.ProjectTo(A::SkewSymMatrix) = ProjectTo{SkewSymMatrix}(; skew_sym = ProjectTo(A.S)) (project::ProjectTo{SkewSymMatrix})(dA::AbstractMatrix) = SkewSymMatrix(project.skew_sym(map_to_Skew(dA)), size(dA, 2)) (project::ProjectTo{SkewSymMatrix})(dA::SkewSymMatrix) = SkewSymMatrix(project.skew_sym(dA.S), dA.n) \ No newline at end of file diff --git a/src/arrays/stiefel_lie_algebra_horizontal.jl b/src/arrays/stiefel_lie_algebra_horizontal.jl index dd144b09f..d77e552df 100644 --- a/src/arrays/stiefel_lie_algebra_horizontal.jl +++ b/src/arrays/stiefel_lie_algebra_horizontal.jl @@ -317,4 +317,6 @@ function _round(B::StiefelLieAlgHorMatrix; kwargs...) B.N, B.n ) -end \ No newline at end of file +end + +Base.fill!(A::StiefelLieAlgHorMatrix, val) = (fill!(A.A, val); fill!(A.B, val); A) \ No newline at end of file diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index a575ab31b..c10fbacae 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -30,7 +30,7 @@ batch(dl) # output [ Info: You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) number of parameters. -([(1, 5), (1, 3)], [(1, 4), (1, 1)], [(1, 2)]) +([(1, 3), (1, 1)], [(1, 5), (1, 4)], [(1, 2)]) ``` Here the first index is always 1 (the time dimension). We get a total number of 3 batches. @@ -125,8 +125,8 @@ println(stdout, batch(dl₁), "\n", batch(dl₂)) Number of batches of dl₁: 2 Number of batches of dl₂: 2 -([(1, 1), (4, 1), (2, 1)], [(3, 1)]) -([(1, 3), (1, 2), (1, 4)], [(1, 1), (1, 5)]) +([(1, 1), (3, 1), (2, 1)], [(4, 1)]) +([(1, 5), (1, 4), (1, 1)], [(1, 3), (1, 2)]) ``` Here we see that in the *autoencoder case* that last minibatch has an additional element. diff --git a/src/manifolds/grassmann_manifold.jl b/src/manifolds/grassmann_manifold.jl index 81e5e687d..0ee8972c5 100644 --- a/src/manifolds/grassmann_manifold.jl +++ b/src/manifolds/grassmann_manifold.jl @@ -117,4 +117,58 @@ function Ω(Y::GrassmannManifold{T}, Δ::AbstractMatrix{T}) where T # E = StiefelProjection(Y) # SkewSymMatrix(ΩSt - E * E' * ΩSt * E * E') SkewSymMatrix(ΩSt) +end + +function Base.copyto!(A::GrassmannManifold, B::GrassmannManifold) + A.A .= B.A + nothing +end + +Base.copy(A::GrassmannManifold) = GrassmannManifold(copy(A.A)) +Base.similar(A::GrassmannManifold) = GrassmannManifold(similar(A.A)) + +function GeometricOptimizers.global_rep( + λY::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + Δ::AbstractMatrix{T} +) where T + N, n = size(λY.Y) + GrassmannLieAlgHorMatrix( + λY.λ' * Δ, + N, n + ) +end + +function GeometricOptimizers.update_section!( + Λᵗ::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + Λ⁽ᵗ⁻¹⁾::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + B⁽ᵗ⁻¹⁾::AbstractMatrix{T}, + retraction +) where T + N, n = B⁽ᵗ⁻¹⁾.N, B⁽ᵗ⁻¹⁾.n + expB = retraction(B⁽ᵗ⁻¹⁾) + expB.A .= Λ⁽ᵗ⁻¹⁾.Y.A * expB.A[1:n, :] .+ Λ⁽ᵗ⁻¹⁾.λ * expB.A[(n+1):N, :] + Λᵗ.Y.A .= @view expB.A[:, 1:n] + Λᵗ.λ .= @view expB.A[:, (n+1):N] + nothing +end + +function GeometricOptimizers.cayley(B::GrassmannLieAlgHorMatrix{T}) where T + backend = networkbackend(B) + E = StiefelProjection(B) + 𝕆 = KernelAbstractions.zeros(backend, T, B.n, B.n) + 𝕀_small = one(𝕆) + 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) + 𝕀_big = one(B) + B̂ = hcat(vcat(𝕆, B.B), E) + B̄ = hcat(vcat(𝕀_small, 𝕆), vcat(zero(B.B'), -B.B'))' + GrassmannManifold((𝕀_big + T(0.5) * B̂ * inv(𝕀_small2 - T(0.5) * B̄' * B̂) * B̄') * (𝕀_big + T(0.5) * B)) +end + +function GeometricOptimizers._copyto!( + Λ₁::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + Λ₂::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}} +) where T + copyto!(Λ₁.Y, Λ₂.Y) + copyto!(Λ₁.λ, Λ₂.λ) + Λ₁ end \ No newline at end of file diff --git a/src/manifolds/stiefel_manifold.jl b/src/manifolds/stiefel_manifold.jl index 5caf4afc0..bcd80e1ad 100644 --- a/src/manifolds/stiefel_manifold.jl +++ b/src/manifolds/stiefel_manifold.jl @@ -181,4 +181,88 @@ end function Base.copyto!(A::StiefelManifold, B::StiefelManifold) A.A .= B.A nothing +end + +Base.copy(A::StiefelManifold) = StiefelManifold(copy(A.A)) +Base.similar(A::StiefelManifold) = StiefelManifold(similar(A.A)) + +function Base.zero(Y::StiefelManifold{T}) where T + N, n = size(Y) + backend = KernelAbstractions.get_backend(Y.A) + zeros(backend, StiefelLieAlgHorMatrix{T}, N, n) +end + +# Bridge GML's StiefelManifold into GO's manifold optimization infrastructure. +# GO's methods dispatch on MT<:GO.Manifold{T}, but GML.StiefelManifold<:GML.Manifold{T} +# (a different type hierarchy), so we extend GO's functions explicitly for GML's type. + +function GeometricOptimizers.global_rep( + λY::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + Δ::AbstractMatrix{T} +) where T + N, n = size(λY.Y) + StiefelLieAlgHorMatrix( + SkewSymMatrix(λY.Y.A' * Δ), + λY.λ' * Δ, + N, n + ) +end + +function GeometricOptimizers.update_section!( + Λᵗ::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + Λ⁽ᵗ⁻¹⁾::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + B⁽ᵗ⁻¹⁾::AbstractMatrix{T}, + retraction +) where T + N, n = B⁽ᵗ⁻¹⁾.N, B⁽ᵗ⁻¹⁾.n + expB = retraction(B⁽ᵗ⁻¹⁾) + expB.A .= Λ⁽ᵗ⁻¹⁾.Y.A * expB.A[1:n, :] .+ Λ⁽ᵗ⁻¹⁾.λ * expB.A[(n+1):N, :] + Λᵗ.Y.A .= @view expB.A[:, 1:n] + Λᵗ.λ .= @view expB.A[:, (n+1):N] + nothing +end + +function GeometricOptimizers.cayley(B::StiefelLieAlgHorMatrix{T}) where T + E = StiefelProjection(B) + 𝕀_small = one(B.A) + 𝕆 = zero(𝕀_small) + 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) + 𝕀_big = one(B) + A_mat = B.A * 𝕀_small + B̂ = hcat(vcat(T(0.5) * A_mat, B.B), E) + B̄ = hcat(vcat(𝕀_small, T(0.5) * A_mat), vcat(zero(B.B'), -B.B'))' + StiefelManifold((𝕀_big + T(0.5) * B̂ * inv(𝕀_small2 - T(0.5) * B̄' * B̂) * B̄') * (𝕀_big + T(0.5) * B)) +end + +function GeometricOptimizers._copyto!( + Λ₁::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + Λ₂::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}} +) where T + copyto!(Λ₁.Y, Λ₂.Y) + copyto!(Λ₁.λ, Λ₂.λ) + Λ₁ +end + +# GO arithmetic bridges for GML's SkewSymMatrix (GO dispatches by module, not structure). +GeometricOptimizers._add!(a::SkewSymMatrix{T}, b::SkewSymMatrix{T}) where T = (a.S .+= b.S; a) +GeometricOptimizers._add!(a::SkewSymMatrix{T}, b::T) where T = (a.S .+= b; a) +GeometricOptimizers._rac!(B::SkewSymMatrix, A::SkewSymMatrix) = (B.S .= sqrt.(A.S); B) +GeometricOptimizers._square!(B::SkewSymMatrix, A::SkewSymMatrix) = (B.S .= A.S .^ 2; B) +GeometricOptimizers._div!(C::SkewSymMatrix, A::SkewSymMatrix, B::SkewSymMatrix) = (C.S .= A.S ./ B.S; C) + +# GO arithmetic bridges for GML's StiefelLieAlgHorMatrix. +function GeometricOptimizers._add!(A::StiefelLieAlgHorMatrix{T}, B::StiefelLieAlgHorMatrix{T}) where T + GeometricOptimizers._add!(A.A, B.A); A.B .+= B.B; A +end +function GeometricOptimizers._add!(A::StiefelLieAlgHorMatrix{T}, b::T) where T + GeometricOptimizers._add!(A.A, b); A.B .+= b; A +end +function GeometricOptimizers._rac!(B::StiefelLieAlgHorMatrix, A::StiefelLieAlgHorMatrix) + GeometricOptimizers._rac!(B.A, A.A); B.B .= sqrt.(A.B); B +end +function GeometricOptimizers._square!(B::StiefelLieAlgHorMatrix, A::StiefelLieAlgHorMatrix) + GeometricOptimizers._square!(B.A, A.A); B.B .= A.B .^ 2; B +end +function GeometricOptimizers._div!(C::StiefelLieAlgHorMatrix, A::StiefelLieAlgHorMatrix, B::StiefelLieAlgHorMatrix) + GeometricOptimizers._div!(C.A, A.A, B.A); C.B .= A.B ./ B.B; C end \ No newline at end of file diff --git a/src/utils.jl b/src/utils.jl index 3ea8596ec..cba7ae3a3 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -213,12 +213,18 @@ _is_go_native_method(::GeometricOptimizers.MomentumMethod) = true _is_go_native_method(::GeometricOptimizers.Adam) = true _is_go_native_method(::GeometricOptimizers.OptimizerMethod) = false +_adapt_method_to_T(method::GeometricOptimizers.Adam, ::Type{T}) where T = + GeometricOptimizers.Adam(T(method.η), T(method.β₁), T(method.β₂), T(method.δ)) +_adapt_method_to_T(method::GeometricOptimizers.MomentumMethod, ::Type{T}) where T = + GeometricOptimizers.MomentumMethod(T(method.α)) +_adapt_method_to_T(method, ::Type) = method + _use_go_cache(method, x) = _is_go_native_method(method) && x isa GeometricOptimizers.OptimizerSolution function _make_optimizer_cache(method, x) if _use_go_cache(method, x) - GeometricOptimizers.OptimizerCache(method, x) + GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, _eltype(x)), x) elseif x isa NamedTuple || x isa NeuralNetworkParameters NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x))) else @@ -306,8 +312,9 @@ function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) T = _eltype(ps_leaf) local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf) - if method isa GeometricOptimizers.Adam - GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) + adapted = _adapt_method_to_T(method, T) + if adapted isa GeometricOptimizers.Adam + GeometricOptimizers.update!(cache, state, local_grad, adapted, ps_leaf) else GeometricOptimizers.update!(cache, state, local_grad, GeometricOptimizers.NoHessian{T}(), ps_leaf) @@ -332,7 +339,7 @@ function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, GeometricOptimizers.second_moment(cache)) elseif state isa GeometricOptimizers.MomentumState GeometricOptimizers._add!(GeometricOptimizers.momentum(state), - GeometricOptimizers._mul(method.α, + GeometricOptimizers._mul(adapted.α, GeometricOptimizers.gradient_array(cache))) end state.iterations += 1