From d7e2551ce5eff25b698a67b86cf092dfe821b574 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Fri, 25 Sep 2026 08:41:31 +0200 Subject: [PATCH 1/3] Take one optimizer method per layer, and stop listing methods here A symplectic autoencoder has PSD layers that carry only Stiefel weights beside SympNet layers that carry only Euclidean ones, and `ScalarMomentAdam` accepts exactly one `StiefelManifold` solution. A caller wanting that method on the former and ordinary `Adam` on the latter had to add methods to `_make_optimizer_cache`, `_make_optimizer_state` and `_leaf_optim_step!` from outside this package -- type piracy on four underscore-prefixed internals, which breaks silently at run time on any refactor here, in the middle of whatever it was training. `GeometricOptimizers.CompositeMethod` is now accepted wherever a method is. The three functions resolve the layer's method with `leafmethod` and then behave exactly as they do for that method, so a layer stepped under a composite is bit-for-bit a layer stepped under the method selected for it. Two lists of methods maintained in the package that does not define them go with it. `_leaf_optim_step!`'s `isa AdamState / isa MomentumState` chain becomes one `sync_state!` call -- that chain had the failure mode such a list has, and `ScalarMomentAdam` was simply missing from it -- and `_go_update_leaf!` dispatches on `FirstOrderMethodWithState` rather than naming `Adam` and `MomentumMethod`. `ScalarMomentAdam` is GO-native here as a result; it had been falling through to `GMLEuclideanState`, which is coordinate-wise Adam on the ambient array rather than the method, and off the manifold. Requires GeometricOptimizers 0.8. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 55 ++++++++++ Project.toml | 4 +- src/optimizers/optimizer.jl | 155 ++++++++++++++++++++++------ test/optimizers/composite_method.jl | 141 +++++++++++++++++++++++++ test/optimizers/runtests.jl | 3 + 5 files changed, 323 insertions(+), 35 deletions(-) create mode 100644 test/optimizers/composite_method.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index cbee0e6fa..d03a355ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,61 @@ its leaves as `NetworkParameters{T}`'s first type parameter, which is what lets an `OptimizerSolution{T}` there — and `SymbolicNeuralNetworks` 0.6 permits only the 0.1 container, so it cannot coexist with `GeometricOptimizers` 0.5. +### Added + +- **One optimizer method per layer**, through `GeometricOptimizers.CompositeMethod`. `Optimizer` and + `optimization_step!` take one wherever they take a method: `_make_optimizer_cache`, + `_make_optimizer_state` and `_leaf_optim_step!` resolve the layer's method with + `GeometricOptimizers.leafmethod` and then behave exactly as they do for that method, so a layer + stepped under a composite is bit-for-bit a layer stepped under the method selected for it. + + ```julia + method = GeometricOptimizers.CompositeMethod(; + manifold = GeometricOptimizers.ScalarMomentAdam(Float32), + array = GeometricOptimizers.Adam(Float32)) + optimizer = Optimizer(method, nn; retraction = cayley, step_size = 1.0f-3) + ``` + + A symplectic autoencoder is the shape this exists for: its PSD layers carry only Stiefel weights + and its SympNet layers only Euclidean ones, and `ScalarMomentAdam` accepts exactly one + `StiefelManifold` solution. A caller wanting that split had to add methods to those three + underscore-prefixed functions **from outside this package** — type piracy on four internals, which + breaks silently at run time on any refactor here, in the middle of whatever it was training. There + is now a supported way to say it. + +- **`_as_go_solution(x, method)` and `_as_go_leaf(x, method)`**, which hand a method whose scope is a + single weight that weight rather than the layer holding it — and its gradient and its section with + it. Which methods those are is `GeometricOptimizers.accepts_parameter_set` and not a list of type + names kept here; *when* to unwrap, and the error naming the layer and the method when a layer holds + more than one weight, is this package's, because the grouping into layers is. + +### Changed + +- **`_leaf_optim_step!` no longer carries a list of which methods keep what.** The + `if state isa AdamState … elseif state isa MomentumState` chain that copied the moments or advanced + the momentum is now one call to `GeometricOptimizers.sync_state!`, which dispatches on the state + and the method in the package that defines both. That chain was a list of methods maintained where + the methods are not defined, and it had the failure mode such a list has: `ScalarMomentAdam` was + simply missing from it. + +- **`_go_update_leaf!` dispatches on `GeometricOptimizers.FirstOrderMethodWithState`** rather than on + `Adam` and `MomentumMethod` by name. That union is the split upstream's own `solver_step!` makes — + a method that builds its direction out of state of its own needs the *method*, everything else + needs the Hessian and has none — so a method joining it upstream no longer leaves a stale list + behind here, silently taking the Hessian branch. + +- **`ScalarMomentAdam` is now a GO-native method here**, and `_adapt_method_to_T` converts it as it + converts `Adam`. It previously fell through to `GMLEuclideanState`, which is coordinate-wise Adam + on the ambient array: not the method, and not on the manifold. + +- **`_make_optimizer_cache` and `_make_optimizer_state` ask the layer question before the capability + question.** The structural-before-capability ordering the comments there are about is unchanged — + `_is_layer` is still asked first and a flat set is still one cache — but `_is_go_native_method` is + now asked of the *layer's* method, which for a composite exists only once a layer is in hand. + +- **Requires `GeometricOptimizers` 0.8**, for `CompositeMethod`, `leafmethod`, + `accepts_parameter_set` and `sync_state!`. None of the four is in 0.7. + ### Removed (breaking) - **Nine of the package's twelve type piracies are deleted, and three of them were on `Base`.** diff --git a/Project.toml b/Project.toml index a15297471..b9bb137ce 100644 --- a/Project.toml +++ b/Project.toml @@ -38,7 +38,9 @@ ForwardDiff = "1" GPUArraysCore = "0.2" GeometricBase = "0.14" GeometricEquations = "0.21" -GeometricOptimizers = "0.7" +# 0.8 for `CompositeMethod`, `leafmethod`, `accepts_parameter_set` and `sync_state!`, which the +# optimizer below is written in terms of. None of the four is in 0.7. +GeometricOptimizers = "0.8" GeometricSolutions = "0.6" HDF5 = "0.16, 0.17" InteractiveUtils = "1" diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl index 69d81f32a..1fa502b74 100644 --- a/src/optimizers/optimizer.jl +++ b/src/optimizers/optimizer.jl @@ -58,7 +58,20 @@ end _is_go_native_method(::GeometricOptimizers.GradientMethod) = true _is_go_native_method(::GeometricOptimizers.MomentumMethod) = true _is_go_native_method(::GeometricOptimizers.Adam) = true +# `ScalarMomentAdam` is driven through the same cache as the other three: `_go_update_leaf!` hands +# `update!` the method (it is a `FirstOrderMethodWithState`) and `GeometricOptimizers.sync_state!` +# carries its scalar second moment back into the state. Without this arm it fell through to +# `GMLEuclideanState`, which is coordinate-wise Adam on the ambient array -- not the method, and not +# on the manifold. +_is_go_native_method(::GeometricOptimizers.ScalarMomentAdam) = true _is_go_native_method(::GeometricOptimizers.OptimizerMethod) = false +# A `CompositeMethod` is whatever it selects, and what it selects depends on the leaf -- so this +# question cannot be answered about the composite itself, only about a leaf's method. Every call +# site below resolves the leaf's method with `leafmethod` before asking, and this arm is what says +# so if one is ever added that does not. +_is_go_native_method(::GeometricOptimizers.CompositeMethod) = throw(ArgumentError( + "`_is_go_native_method` is a question about a leaf's method; resolve the composite with " * + "`GeometricOptimizers.leafmethod(method, leaf)` first")) function _adapt_method_to_T(method::GeometricOptimizers.Adam, ::Type{T}) where {T} GeometricOptimizers.Adam(T; β₁ = T(method.β₁), β₂ = T(method.β₂), δ = T(method.δ)) @@ -66,6 +79,13 @@ end function _adapt_method_to_T(method::GeometricOptimizers.MomentumMethod, ::Type{T}) where {T} GeometricOptimizers.MomentumMethod(T(method.α)) end +# As for `Adam`: `ScalarMomentAdam` carries parameters of its own and is not converted by +# `Optimizer`, so it has to be constructed with the element type of the parameters or its cache +# constructor raises. `ambient_norm` is a mode and not a number, so it is carried across unchanged. +function _adapt_method_to_T(method::GeometricOptimizers.ScalarMomentAdam, ::Type{T}) where {T} + GeometricOptimizers.ScalarMomentAdam(T; β₁ = T(method.β₁), β₂ = T(method.β₂), + δ = T(method.δ), ambient_norm = method.ambient_norm) +end _adapt_method_to_T(method, ::Type) = method # Whether one `GeometricOptimizers` cache covers `x` whole. @@ -99,16 +119,60 @@ end """ _as_go_solution(x) + _as_go_solution(x, method) -`x` in the shape `GeometricOptimizers` takes a solution in. +`x` in the shape `GeometricOptimizers` takes a solution in for `method`. A layer given as a bare `NamedTuple` is wrapped; anything already in the right shape — a container, a bare `Manifold`, an `AbstractArray` — is passed through. The wrap **shares the leaf arrays**, so an in-place optimizer step writes through to the network's own weights and nothing has to be copied back. + +The two-argument form additionally *unwraps* for a method whose scope is a single leaf rather than a +set of them. `GeometricOptimizers.ScalarMomentAdam` is the one that is: its second moment is a +scalar, which is a statement about one manifold and means nothing pooled across a set, so it rejects +a container even when that container holds exactly one Stiefel weight. Which methods those are is +upstream's `accepts_parameter_set` and not a list of type names kept here; *when* to unwrap — a +layer of exactly one weight, and an error naming the layer otherwise — is this package's rule, +because the grouping into layers is. """ _as_go_solution(x::NetworkParameters) = x _as_go_solution(x) = _is_layer(x) ? NetworkParameters(x) : x +function _as_go_solution(x, method) + GeometricOptimizers.accepts_parameter_set(method) && return _as_go_solution(x) + _single_weight(x, method) +end + +""" + _as_go_leaf(x, method) + +The *unwrap* half of [`_as_go_solution`](@ref), for the gradient and the section tree that travel +beside a solution. + +Neither wants the wrap — `_gml_rgrad` normalises a plain `NamedTuple` gradient itself, and a section +tree is copied into as it stands — but both have to follow the solution through the unwrap, so that +a method whose scope is a single weight is handed that weight's gradient and that weight's section +rather than the layer's. +""" +function _as_go_leaf(x, method) + GeometricOptimizers.accepts_parameter_set(method) ? x : _single_weight(x, method) +end + +# The unwrap, and the error that stands in for it. A method of single-leaf scope selected for a layer +# that holds more than one weight is a mistake in the *selector*, and saying which layer and which +# method is what makes it findable; without this it surfaces several frames down as upstream's +# scope `ArgumentError`, which knows nothing about layers. +_single_weight(x::Union{NetworkParameters, NamedTuple}, method) = _single_weight(values(x), x, method) +_single_weight(x, ::Any) = x + +function _single_weight(weights, layer, method) + length(weights) == 1 || throw(ArgumentError( + "$(nameof(typeof(method))) takes a single weight as its solution, but the layer it was " * + "selected for holds $(length(weights)) ($(join(keys(layer), ", "))); select a method " * + "that accepts a parameter set for this layer")) + only(weights) +end + # A `NetworkParameters` is always a tree of layers to descend into, never a single # `GeometricOptimizers` leaf, so its branch comes *first* — ahead of `_use_go_cache`. # @@ -135,27 +199,48 @@ _as_go_solution(x) = _is_layer(x) ? NetworkParameters(x) : x # # The tree branch comes second and covers both carriers, because a network is a tree of layers whether # it is wrapped or not. +# Whether `x` is something to descend into rather than something one `GeometricOptimizers` object is +# built for. A flat set is *not* a branch: it is a layer, and one cache is what a layer is for. +_is_branch(x) = (x isa NetworkParameters || x isa NamedTuple) && !_is_layer(x) + +# The method that steps `x`, which for a `GeometricOptimizers.CompositeMethod` depends on `x` and for +# every other method is the method itself. A branch is passed through untouched: it is descended +# into, and its layers are what get asked, so a selector is never shown a whole network. +# +# **This is why the layer test now comes first in the three branches below.** A composite cannot +# answer `_is_go_native_method` about itself -- that is a property of the method it selects, which +# depends on the leaf -- so asking it before establishing that `x` is a layer would raise on every +# tree root. The ordering the original comment is about, structure before capability, is unchanged: +# `_is_layer` is still the first question, and a flat set is still one cache. +_leaf_method(method, x) = _is_branch(x) ? method : GeometricOptimizers.leafmethod(method, x) + function _make_optimizer_cache(method, x) - if _is_go_native_method(method) && _is_layer(x) - GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, parameter_eltype(x)), - _as_go_solution(x)) + leaf_method = _leaf_method(method, x) + if _is_layer(x) && _is_go_native_method(leaf_method) + GeometricOptimizers.OptimizerCache( + _adapt_method_to_T(leaf_method, parameter_eltype(x)), + _as_go_solution(x, leaf_method)) elseif x isa NetworkParameters || x isa NamedTuple NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x))) - elseif _use_go_cache(method, x) - GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, parameter_eltype(x)), x) + elseif _use_go_cache(leaf_method, x) + GeometricOptimizers.OptimizerCache( + _adapt_method_to_T(leaf_method, parameter_eltype(x)), x) else GMLEuclideanState(x) end end function _make_optimizer_state(method, x) - if _is_go_native_method(method) && _is_layer(x) - GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, parameter_eltype(x)), - _as_go_solution(x)) + leaf_method = _leaf_method(method, x) + if _is_layer(x) && _is_go_native_method(leaf_method) + GeometricOptimizers.OptimizerState( + _adapt_method_to_T(leaf_method, parameter_eltype(x)), + _as_go_solution(x, leaf_method)) elseif x isa NetworkParameters || x isa NamedTuple NamedTuple{keys(x)}(Tuple(_make_optimizer_state(method, x[k]) for k in keys(x))) - elseif _use_go_cache(method, x) - GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, parameter_eltype(x)), x) + elseif _use_go_cache(leaf_method, x) + GeometricOptimizers.OptimizerState( + _adapt_method_to_T(leaf_method, parameter_eltype(x)), x) else GMLEuclideanState(x) end @@ -278,13 +363,13 @@ end # Adam — only `step_size` differs, and that comes in as an argument — so the `Adam` method above # serves it too. +# The split upstream's `solver_step!` makes, and made here for the same reason: a +# `FirstOrderMethodWithState` builds its direction out of state of its own and so needs the *method*, +# and everything else needs the Hessian and has none. Written as that union rather than as one arm +# per method, so that a method joining it upstream -- `ScalarMomentAdam` did -- does not leave a +# stale list behind here silently taking the Hessian branch instead. function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.Adam, ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) -end - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.MomentumMethod, ps_leaf) + method::GeometricOptimizers.FirstOrderMethodWithState, ps_leaf) GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) end @@ -304,9 +389,15 @@ function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, state::GeometricOptimizers.OptimizerState, dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) T = parameter_eltype(ps_leaf) - ps = _as_go_solution(ps_leaf) - local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf) - adapted = _adapt_method_to_T(method, T) + adapted = _adapt_method_to_T(GeometricOptimizers.leafmethod(method, ps_leaf), T) + ps = _as_go_solution(ps_leaf, adapted) + # The gradient and the section tree take the *unwrap* half only: `_gml_rgrad` normalises a plain + # `NamedTuple` gradient itself and the section tree is copied into as it stands, so neither wants + # the wrap. Both have to follow the solution through the unwrap, though -- a method whose scope + # is a single weight is handed that weight's gradient and that weight's section. + dp = _as_go_leaf(dp_leaf, adapted) + λY = _as_go_leaf(λY_leaf, adapted) + local_grad = _GMLGradient{T, typeof(dp)}(dp) state.iterations += 1 _go_update_leaf!(cache, state, local_grad, adapted, ps) GeometricOptimizers._rmul!(GeometricOptimizers.direction(cache), T(step_size)) @@ -317,31 +408,27 @@ function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, GeometricOptimizers._copyto!(GeometricOptimizers.solution(cache), GeometricOptimizers.section(cache)) GeometricOptimizers._copyto!(ps, GeometricOptimizers.solution(cache)) - GeometricOptimizers._copyto!(λY_leaf, GeometricOptimizers.section(cache)) + GeometricOptimizers._copyto!(λY, GeometricOptimizers.section(cache)) # `section(cache)` is `update_section!(section(state), direction, retraction)`, so copying it is # the same thing as retracting a second time -- and a retraction on a manifold is `O(N³)` where # the copy is `O(N²)`. GeometricOptimizers._copyto!(GeometricOptimizers.section(state), GeometricOptimizers.section(cache)) - 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 - # `p ← αp + ∇L`; see the note in `_euclidean_update!` for the momentum method. This has to - # match what `update!(::MomentumCache, ...)` anticipated when it formed the direction. - GeometricOptimizers._rmul!(GeometricOptimizers.momentum(state), adapted.α) - GeometricOptimizers._add!(GeometricOptimizers.momentum(state), - GeometricOptimizers.gradient_array(cache)) - end + # What a method carries out of a step -- a momentum, a pair of moments -- and whether the cache's + # copy is the one to keep is the method's own knowledge, and upstream's since 0.8. This was a + # chain of `isa` tests over the state types here, which is a list of methods maintained in the + # package that does not define them: `ScalarMomentAdam` was simply missing from it. + GeometricOptimizers.sync_state!(state, cache, adapted) nothing end # 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) + # The composite is resolved here too: this arm is reached for a leaf whose selected method has + # no `GeometricOptimizers` cache, and `_euclidean_update!` dispatches on that method. + _euclidean_update!(ps_leaf, dp_leaf, state, + GeometricOptimizers.leafmethod(method, ps_leaf), step_size) state.iterations += 1 nothing end diff --git a/test/optimizers/composite_method.jl b/test/optimizers/composite_method.jl new file mode 100644 index 000000000..dc365da9f --- /dev/null +++ b/test/optimizers/composite_method.jl @@ -0,0 +1,141 @@ +# One optimizer method per layer, through `GeometricOptimizers.CompositeMethod`. +# +# A symplectic autoencoder is the shape this exists for: its PSD layers carry only Stiefel weights +# and its SympNet layers only Euclidean ones, and `ScalarMomentAdam` -- whose second moment is a +# scalar, and therefore a statement about *one* manifold -- accepts exactly one `StiefelManifold` +# solution. Before the composite, a caller wanting that method on the Stiefel layers and ordinary +# `Adam` on the rest had to add methods to `_make_optimizer_cache`, `_make_optimizer_state` and +# `_leaf_optim_step!` from outside this package, which is type piracy on four underscore-prefixed +# internals and breaks silently on any refactor here. +# +# What is pinned below is that the selection reaches the cache, the state and the step; that a +# single-weight-scope method is handed the weight and not the layer; that its scalar moment is +# carried back into the state; and that a uniform method is byte-for-byte unaffected. + +using GeometricMachineLearning +using GeometricMachineLearning: params, _as_go_solution, _is_go_native_method, + _make_optimizer_cache, _make_optimizer_state +using Test +import GeometricOptimizers +import Random + +const GO = GeometricOptimizers + +Random.seed!(1234) + +composite(T) = GO.CompositeMethod(; + manifold = GO.ScalarMomentAdam(T), array = GO.Adam(T)) + +@testset "ScalarMomentAdam is driven through a GeometricOptimizers cache" begin + # Without this it fell through to `GMLEuclideanState`, which is coordinate-wise Adam on the + # ambient array: not the method, and not on the manifold. + @test _is_go_native_method(GO.ScalarMomentAdam()) + @test_throws ArgumentError _is_go_native_method(composite(Float64)) +end + +@testset "the cache and state trees follow the selection layer by layer" begin + T = Float32 + ps = params(NeuralNetwork(SymplecticAutoencoder(4, 2; n_encoder_blocks = 1, + n_decoder_blocks = 1, n_encoder_layers = 2, n_decoder_layers = 2), T)) + caches = _make_optimizer_cache(composite(T), ps) + states = _make_optimizer_state(composite(T), ps) + + stiefel_caches = 0 + adam_caches = 0 + function walk(node, layer) + if node isa NamedTuple && !(node isa GO.OptimizerCache) + for key in keys(node) + walk(node[key], layer[key]) + end + return nothing + end + if all(weight -> weight isa StiefelManifold, values(layer)) + @test node isa GO.ScalarMomentAdamCache + stiefel_caches += 1 + else + @test node isa GO.AdamCache + adam_caches += 1 + end + nothing + end + walk(caches, ps) + + @test stiefel_caches > 0 # the PSD layers + @test adam_caches > 0 # the SympNet layers + @test states isa NamedTuple +end + +@testset "a single-weight-scope method is handed the weight, not the layer" begin + Y = rand(StiefelManifold{Float32}, 4, 2) + layer = (weight = Y,) + + @test _as_go_solution(layer, GO.ScalarMomentAdam(Float32)) === Y + @test _as_go_solution(layer, GO.Adam(Float32)) isa + GeometricMachineLearning.NetworkParameters + @test _as_go_solution(Y, GO.ScalarMomentAdam(Float32)) === Y + + # A layer of more than one weight has no single weight to be handed, and the error says which + # layer and which method rather than surfacing as upstream's scope message several frames down. + wide = (weight = Y, bias = rand(Float32, 2)) + @test_throws ArgumentError _as_go_solution(wide, GO.ScalarMomentAdam(Float32)) +end + +@testset "a composite trains a symplectic autoencoder and keeps it symplectic" begin + T = Float32 + architecture = SymplecticAutoencoder(4, 2; n_encoder_blocks = 1, n_decoder_blocks = 1, + n_encoder_layers = 2, n_decoder_layers = 2) + nn = NeuralNetwork(architecture, T) + dl = DataLoader(rand(T, 4, 20, 5); autoencoder = true, suppress_info = true) + + optimizer = Optimizer(composite(T), nn; retraction = cayley, step_size = 1.0f-3) + losses = optimizer(nn, dl, Batch(5), 2; show_progress = false) + + @test length(losses) == 2 + @test all(isfinite, losses) + # Every Stiefel weight is still on the manifold: the composite's Stiefel arm retracts, it does + # not step in the ambient space. + for layer in values(params(nn)), weight in values(layer) + weight isa StiefelManifold && @test GO.check(weight) < 1.0f-4 + end +end + +@testset "the scalar second moment is carried back into the state" begin + # `sync_state!`'s job, and the one the `isa` chain this replaced simply did not have an arm for: + # without it the moment restarts from zero on every step and the method is its first iteration + # forever. + T = Float64 + Y = rand(StiefelManifold{T}, 4, 2) + ps = GeometricMachineLearning.NetworkParameters((L1 = (weight = Y,),)) + method = composite(T) + caches = _make_optimizer_cache(method, ps) + states = _make_optimizer_state(method, ps) + λY = GlobalSection(ps) + dp = (L1 = (weight = rand(T, 4, 2),),) + + @test GO.second_moment(states.L1) == 0 + GeometricMachineLearning._tree_optim_step!(caches, states, dp, ps, λY, method, + cayley, 1.0e-3) + @test GO.second_moment(states.L1) > 0 + @test GO.second_moment(states.L1) == GO.second_moment(caches.L1) + @test states.L1.iterations == 1 +end + +@testset "a uniform method is unaffected" begin + T = Float32 + architecture = SymplecticAutoencoder(4, 2; n_encoder_blocks = 1, n_decoder_blocks = 1, + n_encoder_layers = 2, n_decoder_layers = 2) + + function run(method) + Random.seed!(77) + nn = NeuralNetwork(architecture, T) + dl = DataLoader(rand(Random.Xoshiro(3), T, 4, 20, 5); autoencoder = true, + suppress_info = true) + optimizer = Optimizer(method, nn; retraction = cayley, step_size = 1.0f-3) + optimizer(nn, dl, Batch(5), 2; show_progress = false) + end + + # `Adam` everywhere, and a composite that selects `Adam` everywhere, have to be the same run: + # the composite is a choice of method and not a second implementation of one. + @test run(GO.Adam(T)) == + run(GO.CompositeMethod(; manifold = GO.Adam(T), array = GO.Adam(T))) +end diff --git a/test/optimizers/runtests.jl b/test/optimizers/runtests.jl index ef699c5ae..d4149195c 100644 --- a/test/optimizers/runtests.jl +++ b/test/optimizers/runtests.jl @@ -27,3 +27,6 @@ end @safetestset "The GO-native leaf step scales in the parameter's own element type" begin include("step_size_element_type.jl") end +@safetestset "One method per layer: GeometricOptimizers.CompositeMethod" begin + include("composite_method.jl") +end From 8470206eaab48bc087c4f23ea59ec13b6c16f8ac Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Mon, 5 Oct 2026 17:56:30 +0200 Subject: [PATCH 2/3] Move to AbstractNeuralNetworks 0.9 and the GeometricOptimizers 0.9 generation GeometricOptimizers 0.9 requires AbstractNeuralNetworks 0.9, NeuralNetworkParameters 0.4.2 and GeometricBase 0.15, and no registered release of either GeometricOptimizers or SymbolicNeuralNetworks allows those, so the five bounds move together (JuliaGNI/SymbolicNeuralNetworks.jl#80). - AbstractNeuralNetworks 0.9 defines the three-axis methods of Dense and Linear and save/load for a NeuralNetwork itself: delete the three pirated functors in src/layers/resnet.jl and ext/HDF5Ext.jl. Aqua's piracies and ambiguities checks run again and report none. - LowerTriangular/UpperTriangular are StrictlyLowerTriangular/ StrictlyUpperTriangular; the legacy-file test covers a network with them (needs JuliaGNI/GeometricOptimizers.jl#146). - PoissonTensor answers KernelAbstractions.get_backend, which GeometricOptimizers 0.9 asks of both factors of a product. - Under NeuralNetworkParameters 0.4 a gradient through a NetworkParameters is the storage gradient: the gradient-structure tests assert that, and #319's structure loss is gone. The optimizer is moved to GeometricOptimizers 0.9 in the next commit; this one does not build on its own against GeometricOptimizers 0.9. Co-Authored-By: Claude Opus 5.5 (1M context) --- Project.toml | 13 ++- docs/check_references.jl | 4 +- docs/make.jl | 4 +- .../architectures/abstract_neural_networks.md | 9 +- docs/src/arrays/tensors.md | 6 +- docs/src/index.md | 2 +- .../layers/volume_preserving_feedforward.md | 4 +- ext/HDF5Ext.jl | 87 ------------------- src/GeometricMachineLearning.jl | 4 +- src/arrays/poisson_tensor.jl | 5 ++ .../kernel_ad_routines/mat_tensor_mul.jl | 12 +-- src/kernels/mat_tensor_mul.jl | 16 ++-- src/layers/resnet.jl | 12 --- src/layers/symplectic_attention.jl | 11 +-- src/layers/volume_preserving_feedforward.jl | 16 ++-- test/arrays/triangular.jl | 8 +- .../network_parameters_gradient_structure.jl | 42 ++++----- test/integration/changebackend_tests.jl | 4 +- test/integration/hdf5_support.jl | 16 ++-- ...tiplication_network_parameters_gradient.jl | 8 +- ...c_attention_network_parameters_gradient.jl | 22 ++--- test/layers/volume_preserving_feedforward.jl | 8 +- test/map_to_cpu_tests.jl | 8 +- .../optimizers/structured_array_parameters.jl | 6 +- test/quality/aqua.jl | 62 ++----------- 25 files changed, 114 insertions(+), 275 deletions(-) delete mode 100644 ext/HDF5Ext.jl diff --git a/Project.toml b/Project.toml index 61360a273..2c969d945 100644 --- a/Project.toml +++ b/Project.toml @@ -25,30 +25,27 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" [extensions] GPUArraysCoreExt = "GPUArraysCore" -HDF5Ext = "HDF5" [compat] -AbstractNeuralNetworks = "0.8" +AbstractNeuralNetworks = "0.9" ChainRulesCore = "1" ForwardDiff = "1" GPUArraysCore = "0.2" -GeometricBase = "0.14" +GeometricBase = "0.15" GeometricEquations = "0.21" -GeometricOptimizers = "0.8" +GeometricOptimizers = "0.9" GeometricSolutions = "0.6" -HDF5 = "0.16, 0.17" InteractiveUtils = "1" KernelAbstractions = "0.9" LinearAlgebra = "1" NNlib = "0.8, 0.9" -NeuralNetworkParameters = "0.3, 0.4" +NeuralNetworkParameters = "0.4.2" ProgressMeter = "1" Random = "1" -SymbolicNeuralNetworks = "0.8" +SymbolicNeuralNetworks = "0.9" Symbolics = "7" TimerOutputs = "0.5, 1" Zygote = "0.7" diff --git a/docs/check_references.jl b/docs/check_references.jl index e012f016f..a10e54245 100644 --- a/docs/check_references.jl +++ b/docs/check_references.jl @@ -12,14 +12,12 @@ # julia --project=docs docs/check_references.jl using GeometricMachineLearning -using HDF5 # so that the extension below is loaded, as in make.jl using AbstractNeuralNetworks # signatures in `@docs` blocks name types from here using Documenter using Documenter: DocSystem # keep in sync with `modules` in make.jl -const MODULES = Module[GeometricMachineLearning, - Base.get_extension(GeometricMachineLearning, :HDF5Ext)] +const MODULES = Module[GeometricMachineLearning] const SRC = joinpath(@__DIR__, "src") diff --git a/docs/make.jl b/docs/make.jl index 7309c5aa3..9fad1d468 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -4,6 +4,7 @@ include(joinpath(@__DIR__, "check_references.jl")) using GeometricMachineLearning +# loads `AbstractNeuralNetworks`' HDF5 extension, whose `save` and `load` the tutorials call using HDF5 using AbstractNeuralNetworks using Documenter @@ -387,8 +388,7 @@ makedocs(; # skips *sub*modules — and (ii) pull GeometricOptimizers' own docstrings in, whose internal # `@ref`s resolve in its namespace and point at bindings this manual does not document. # Names owned by GeometricOptimizers are referred to as plain code with a link to its manual. - modules = [ - GeometricMachineLearning, Base.get_extension(GeometricMachineLearning, :HDF5Ext)], + modules = [GeometricMachineLearning], authors = "Michael Kraus, Benedikt Brantner", repo = "https://github.com/JuliaGNI/GeometricMachineLearning.jl/blob/{commit}{path}#L{line}", sitename = "GeometricMachineLearning.jl", diff --git a/docs/src/architectures/abstract_neural_networks.md b/docs/src/architectures/abstract_neural_networks.md index 15a7cc19e..d6ea28316 100644 --- a/docs/src/architectures/abstract_neural_networks.md +++ b/docs/src/architectures/abstract_neural_networks.md @@ -47,11 +47,4 @@ and we see that it consists of two layers: a [`GradientLayerQ`](@ref) and a [`Gr ## Saving and Loading -`GeometricMachineLearning` adds HDF5-backed `save` and `load` methods for `NeuralNetwork` to the generics `NeuralNetworkParameters` defines. Writing and reading the parameter set itself belongs to that package; the structured parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`, …) come back as themselves because `GeometricOptimizers`, which owns them, registers how each is rebuilt. Passing a prototype parameter set to `load` rebuilds against it and needs no registration at all. - -```@docs -save(::HDF5.H5DataStore, ::NeuralNetwork) -save(::AbstractString, ::NeuralNetwork) -load(::Type{NeuralNetwork}, ::HDF5.H5DataStore, ::AbstractNeuralNetworks.Architecture) -load(::Type{NeuralNetwork}, ::AbstractString, ::AbstractNeuralNetworks.Architecture) -``` +`save` and `load` for a `NeuralNetwork` are `AbstractNeuralNetworks`', which defines them in its HDF5 extension on the generics `NeuralNetworkParameters` provides: `using HDF5` loads them. Writing and reading the parameter set itself belongs to `NeuralNetworkParameters`; the structured parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`, …) come back as themselves because `GeometricOptimizers`, which owns them, registers how each is rebuilt. Passing a prototype parameter set to `load` rebuilds against it and needs no registration at all. A network loads on the CPU; `changebackend(backend, nn)` moves it to a device. diff --git a/docs/src/arrays/tensors.md b/docs/src/arrays/tensors.md index a7e0d1eac..0a8d65a9d 100644 --- a/docs/src/arrays/tensors.md +++ b/docs/src/arrays/tensors.md @@ -1,6 +1,6 @@ # Tensors in `GeometricMachineLearning` -We typically store training data as *tensors with three axes* in `GeometricMachineLearning`. This allows for a parallel computation of matrix products, also for the special arrays such as [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular), [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular), [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) and [`SkewSymMatrix`](@extref GeometricOptimizers GeometricOptimizers.SkewSymMatrix) and objects of [`Manifold`](@extref GeometricOptimizers GeometricOptimizers.Manifold) type such as the [`StiefelManifold`](@extref GeometricOptimizers GeometricOptimizers.StiefelManifold). +We typically store training data as *tensors with three axes* in `GeometricMachineLearning`. This allows for a parallel computation of matrix products, also for the special arrays such as [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular), [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular), [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) and [`SkewSymMatrix`](@extref GeometricOptimizers GeometricOptimizers.SkewSymMatrix) and objects of [`Manifold`](@extref GeometricOptimizers GeometricOptimizers.Manifold) type such as the [`StiefelManifold`](@extref GeometricOptimizers GeometricOptimizers.StiefelManifold). ## Library Functions @@ -10,8 +10,8 @@ GeometricMachineLearning.tensor_mat_mul!(::AbstractArray{<:Number, 3}, ::Abstrac GeometricMachineLearning.tensor_mat_mul!(::AbstractArray{<:Number, 3}, ::AbstractArray{<:Number, 3}, ::SymmetricMatrix) GeometricMachineLearning.mat_tensor_mul(::AbstractMatrix, ::AbstractArray{<:Number, 3}) GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::AbstractMatrix, ::AbstractArray{<:Number, 3}) -GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::LowerTriangular, ::AbstractArray{<:Number, 3}) -GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::UpperTriangular, ::AbstractArray{<:Number, 3}) +GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::StrictlyLowerTriangular, ::AbstractArray{<:Number, 3}) +GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::StrictlyUpperTriangular, ::AbstractArray{<:Number, 3}) GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::SkewSymMatrix, ::AbstractArray{<:Number, 3}) GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::SymmetricMatrix, ::AbstractArray{<:Number, 3}) ``` \ No newline at end of file diff --git a/docs/src/index.md b/docs/src/index.md index 00d3dc8d6..7c9fe47f7 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -58,7 +58,7 @@ Everything tried on that device ran. Constructed, and then applied to both a mat Two of those results rest on `GeometricOptimizers` rather than on anything here. -- `VolumePreservingFeedForward` and `VolumePreservingTransformer` applied to a **matrix** multiply a `LowerTriangular` or `UpperTriangular` weight by that matrix. `GeometricOptimizers` supplies a `KernelAbstractions` kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which `GPUArraysCore` refuses. +- `VolumePreservingFeedForward` and `VolumePreservingTransformer` applied to a **matrix** multiply a `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight by that matrix. `GeometricOptimizers` supplies a `KernelAbstractions` kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which `GPUArraysCore` refuses. - `StiefelLayer`, `GrassmannLayer` and `PSDLayer`, and so every architecture that holds one of them, orthonormalize their weight at **construction** through `GeometricOptimizers.orthonormal_columns`. That is CholeskyQR2 — matrix products and triangular solves only — so it runs wherever the draw was allocated. `LinearAlgebra.qr!` is a host factorization and `Metal.jl` implements no `qr` for an `MtlArray`, so a device needs the other one. ## Tutorials diff --git a/docs/src/layers/volume_preserving_feedforward.md b/docs/src/layers/volume_preserving_feedforward.md index d85cac9a8..e52bc26d6 100644 --- a/docs/src/layers/volume_preserving_feedforward.md +++ b/docs/src/layers/volume_preserving_feedforward.md @@ -7,7 +7,7 @@ The *volume-preserving feedforward layers* in `GeometricMachineLearning` are clo ```math \mathtt{VPFF}_{A, b}: x \mapsto x + \sigma(Ax + b), ``` -where ``\sigma`` is a nonlinearity, ``A`` is the weight and ``b`` is the bias. The matrix ``A`` is either a [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular) matrix ``L`` or an [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular) matrix ``U``. We demonstrate volume-preservation of these layers by considering the case ``A = L``. The matrix looks as follows: +where ``\sigma`` is a nonlinearity, ``A`` is the weight and ``b`` is the bias. The matrix ``A`` is either a [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular) matrix ``L`` or a [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular) matrix ``U``. We demonstrate volume-preservation of these layers by considering the case ``A = L``. The matrix looks as follows: ```math L = \begin{pmatrix} @@ -28,7 +28,7 @@ J = \nabla\mathtt{VPFF}_{L, b} = \begin{pmatrix} b_{n1} & \cdots & b_{n(n-1)} & 1 \end{pmatrix}, ``` -and the determinant of ``J`` is 1, i.e. the map is volume-preserving. A similar statement holds if the matrix ``A`` is [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular) instead of [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular). +and the determinant of ``J`` is 1, i.e. the map is volume-preserving. A similar statement holds if the matrix ``A`` is [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular) instead of [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular). ## Library Functions diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl deleted file mode 100644 index 39912a134..000000000 --- a/ext/HDF5Ext.jl +++ /dev/null @@ -1,87 +0,0 @@ -module HDF5Ext - -using HDF5 -using GeometricMachineLearning -import AbstractNeuralNetworks: NeuralNetworkBackend, Architecture -# `save`, `load`, `params` and the parameter container are `NeuralNetworkParameters`' as of -# `AbstractNeuralNetworks` 0.7, which only re-binds them; reach for them where they are defined. -import NeuralNetworkParameters: NetworkParameters, params, save, load - -# --------------------------------------------------------------------------- -# save / load — the entry points that dispatch on this package's `NeuralNetwork`. -# -# The traversal itself is not here. `NeuralNetworkParameters` walks the parameter set and records -# each group's key order; `GeometricOptimizers` says through `freeparameters`/`rebuild` where each -# structured matrix keeps its numbers, and registers the types so a file loads with no prototype. -# --------------------------------------------------------------------------- - -""" - save(h5::HDF5.H5DataStore, nn::NeuralNetwork) - -Save the parameters of `nn` into an already-open HDF5 store. - -Extends `save` with a dispatch on `NeuralNetwork`. The parameters themselves are written by -`NeuralNetworkParameters`, which tags each structured leaf with the type to rebuild it as and -records the key order of every group. -""" -save(h5::HDF5.H5DataStore, nn::NeuralNetwork) = save(h5, params(nn)) - -""" - save(filename::AbstractString, nn::NeuralNetwork) - -Convenience overload: open `filename` for writing, call [`save`](@ref) on the store, and return -`filename`. -""" -function save(filename::AbstractString, nn::NeuralNetwork) - HDF5.h5open(filename, "w") do h5 - save(h5, nn) - end - filename -end - -""" - load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; backend = CPU()) - load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype; backend = CPU()) - -Load network parameters from an already-open HDF5 store and return a `NeuralNetwork` for `arch`. - -The element type is whatever the file holds, so a `Float32` network reloads as `Float32`. - -Structured parameters — `StiefelManifold`, `SymmetricMatrix` and the rest — are rebuilt from the -type each was stored under, which `GeometricOptimizers` registers with -`NeuralNetworkParameters.register_parameter_type!`. Pass `prototype`, a parameter set of the right -shape, to rebuild against it instead and skip the registry altogether. -""" -function load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; - backend::NeuralNetworkBackend = CPU()) - NeuralNetwork(arch, Chain(arch), load(NetworkParameters, h5), backend) -end - -function load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype; - backend::NeuralNetworkBackend = CPU()) - NeuralNetwork(arch, Chain(arch), load(NetworkParameters, h5, prototype), backend) -end - -""" - load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; backend = CPU()) - load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype; backend = CPU()) - -Convenience overload: open `filename` for reading, then call -[`load`](@ref) on the store. -""" -function load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; - backend::NeuralNetworkBackend = CPU()) - HDF5.h5open(filename, "r") do h5 - load(NeuralNetwork, h5, arch; backend = backend) - end -end - -function load( - ::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype; - backend::NeuralNetworkBackend = CPU()) - HDF5.h5open(filename, "r") do h5 - load(NeuralNetwork, h5, arch, prototype; backend = backend) - end -end - -end diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index bb474b69b..000498b31 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -42,7 +42,7 @@ import Symbolics import GeometricOptimizers import GeometricOptimizers: Manifold, StiefelManifold, GrassmannManifold import GeometricOptimizers: SkewSymMatrix, SymmetricMatrix, - LowerTriangular, UpperTriangular, StiefelProjection + StrictlyLowerTriangular, StrictlyUpperTriangular, StiefelProjection import GeometricOptimizers: StiefelLieAlgHorMatrix, GrassmannLieAlgHorMatrix import GeometricOptimizers: rgrad, metric, check, global_section # `orthonormal_columns(draw)` orthonormalizes `draw()` with CholeskyQR2, which runs on whatever @@ -114,7 +114,7 @@ include("arrays/poisson_tensor.jl") # Re-exported from GeometricOptimizers, so that `using GeometricMachineLearning` on its own still # gives a caller the matrix types its layers are parametrized by. export SymmetricMatrix, SkewSymMatrix -export LowerTriangular, UpperTriangular +export StrictlyLowerTriangular, StrictlyUpperTriangular export StiefelLieAlgHorMatrix, GrassmannLieAlgHorMatrix export StiefelProjection # GML's own diff --git a/src/arrays/poisson_tensor.jl b/src/arrays/poisson_tensor.jl index 313276d65..4e7a1eb6d 100644 --- a/src/arrays/poisson_tensor.jl +++ b/src/arrays/poisson_tensor.jl @@ -39,6 +39,11 @@ struct PoissonTensor{T, AT <: AbstractMatrix{T}} <: AbstractMatrix{T} n::Int end +# The backend of the dense matrix it wraps. `GeometricOptimizers` asks it of both factors of a +# product with one of its manifolds or structured matrices, and refuses a product across two +# backends. +KernelAbstractions.get_backend(𝕁::PoissonTensor) = KernelAbstractions.get_backend(𝕁.J) + # The index types are named rather than left open. An untyped index pair claims every index type, # and it then collides with the `getindex` methods BandedMatrices and BlockArrays add to # `AbstractMatrix` for their own: nine ambiguities, of the same kind as the seventeen that the diff --git a/src/kernels/kernel_ad_routines/mat_tensor_mul.jl b/src/kernels/kernel_ad_routines/mat_tensor_mul.jl index aa7970542..34742f6a8 100644 --- a/src/kernels/kernel_ad_routines/mat_tensor_mul.jl +++ b/src/kernels/kernel_ad_routines/mat_tensor_mul.jl @@ -73,15 +73,15 @@ function ChainRulesCore.rrule(::typeof(lo_mat_mul), S::AbstractVector{T}, return C, lo_mat_mul_pullback end -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::LowerTriangular{T}, A::AbstractArray{ - T, 3}) where {T} +function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::StrictlyLowerTriangular{T}, + A::AbstractArray{T, 3}) where {T} @assert size(A, 1) == B.n C = mat_tensor_mul(B, A) function lower_triangular_mul_pullback(dC) dC = unthunk(dC) f̄, dS, dA, _ = rrule(lo_mat_mul, B.S, A, B.n)[2](dC) - return f̄, LowerTriangular(dS, B.n), dA + return f̄, StrictlyLowerTriangular(dS, B.n), dA end return C, lower_triangular_mul_pullback end @@ -144,15 +144,15 @@ function ChainRulesCore.rrule(::typeof(up_mat_mul), S::AbstractVector{T}, return C, up_mat_mul_pullback end -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::UpperTriangular{T}, A::AbstractArray{ - T, 3}) where {T} +function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::StrictlyUpperTriangular{T}, + A::AbstractArray{T, 3}) where {T} @assert size(A, 1) == B.n C = mat_tensor_mul(B, A) function upper_triangular_mul_pullback(dC) dC = unthunk(dC) f̄, dS, dA, _ = rrule(up_mat_mul, B.S, A, B.n)[2](dC) - return f̄, UpperTriangular(dS, B.n), dA + return f̄, StrictlyUpperTriangular(dS, B.n), dA end return C, upper_triangular_mul_pullback end diff --git a/src/kernels/mat_tensor_mul.jl b/src/kernels/mat_tensor_mul.jl index a2bd3d1bc..e0490f465 100644 --- a/src/kernels/mat_tensor_mul.jl +++ b/src/kernels/mat_tensor_mul.jl @@ -119,7 +119,7 @@ function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::SymmetricMatrix, B::A symmetric_mat_mul!(C, A.S, B, A.n) end -########################### LowerTriangular +########################### StrictlyLowerTriangular @kernel function lo_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{T, 3}, ::Int) where {T} @@ -152,21 +152,21 @@ function lo_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where end @doc raw""" - mat_tensor_mul!(C, A::LowerTriangular, B) + mat_tensor_mul!(C, A::StrictlyLowerTriangular, B) Multiply the lower-triangular matrix `A` onto the tensor `B` from the left and store the result in `C`. This performs an efficient multiplication based on the special structure of the lower-triangular matrix `A`. """ -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::LowerTriangular, B::AbstractArray{ - <:Number, 3}) +function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::StrictlyLowerTriangular, + B::AbstractArray{<:Number, 3}) @assert eltype(C) == eltype(A) == eltype(B) @assert A.n == size(C, 1) == size(B, 1) lo_mat_mul!(C, A.S, B, A.n) end -####################### UpperTriangular +####################### StrictlyUpperTriangular @kernel function up_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where {T} @@ -199,14 +199,14 @@ function up_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where end @doc raw""" - mat_tensor_mul!(C, A::UpperTriangular, B) + mat_tensor_mul!(C, A::StrictlyUpperTriangular, B) Multiply the upper-triangular matrix `A` onto the tensor `B` from the left and store the result in `C`. This performs an efficient multiplication based on the special structure of the upper-triangular matrix `A`. """ -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::UpperTriangular, B::AbstractArray{ - <:Number, 3}) +function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::StrictlyUpperTriangular, + B::AbstractArray{<:Number, 3}) @assert eltype(C) == eltype(A) == eltype(B) @assert A.n == size(C, 1) == size(B, 1) diff --git a/src/layers/resnet.jl b/src/layers/resnet.jl index c721c1e55..6a6b109b7 100644 --- a/src/layers/resnet.jl +++ b/src/layers/resnet.jl @@ -60,18 +60,6 @@ function (d::ResNetLayer{M, M, true})(x::AbstractArray{T, 3}, ps::NamedTuple) wh x + d.activation.(mat_tensor_mul(ps.weight, x) .+ ps.bias) end -function (d::Dense{M, N, true})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - d.σ.(mat_tensor_mul(ps.W, x) .+ ps.b) -end - -function (d::Dense{M, N, false})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - d.σ.(mat_tensor_mul(ps.W, x)) -end - -function (d::Linear{M, N})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - mat_tensor_mul(ps.W, x) -end - function (d::Union{Dense{M, N}, ResNetLayer{M, N}})(z::QPT, ps::NamedTuple) where {M, N} @assert iseven(M) == iseven(N) == true @assert size(z.q, 1) * 2 == M diff --git a/src/layers/symplectic_attention.jl b/src/layers/symplectic_attention.jl index a3ba6d7c9..7d877d24d 100644 --- a/src/layers/symplectic_attention.jl +++ b/src/layers/symplectic_attention.jl @@ -128,11 +128,12 @@ function (d::SymplecticAttentionQ{M, M, :arbitrary})( p = z.p) end -# `A = ps.A` once, and then `A` twice: the local binding is what keeps the `SymmetricMatrix` -# structure in the Zygote gradient of this layer. Read `ps.A` twice and the gradient leaf degrades -# to a plain `Matrix` -- `test/parameters/symplectic_attention_network_parameters_gradient.jl` -# asserts both halves of that, and the last testset there covers this method. The binding costs no -# allocation; it is the number of accesses that decides the structure. +# `A = ps.A` once, and then `A` twice. Before `NeuralNetworkParameters` 0.4 the local binding was +# what kept the `SymmetricMatrix` structure in a Zygote gradient taken through a +# `NetworkParameters`: a second `ps.A` degraded the gradient leaf to a plain `Matrix` (issue #319). +# Since 0.4 that gradient is the storage gradient, which keeps the structure however often the leaf +# is read; `test/layers/symplectic_attention_network_parameters_gradient.jl` covers this method. The +# binding costs no allocation and stays. function (d::SymplecticAttentionQ{M, M, :symmetric})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} A = ps.A diff --git a/src/layers/volume_preserving_feedforward.jl b/src/layers/volume_preserving_feedforward.jl index 7822fcda8..ef81bfc5e 100644 --- a/src/layers/volume_preserving_feedforward.jl +++ b/src/layers/volume_preserving_feedforward.jl @@ -4,10 +4,10 @@ Super-type of [`VolumePreservingLowerLayer`](@ref) and [`VolumePreservingUpperLayer`](@ref). The layers do the following: ```math -x \mapsto \begin{cases} \sigma(Lx + b) & \text{where $L$ is }\mathtt{LowerTriangular}, \\ \sigma(Ux + b) & \text{where $U$ is }\mathtt{UpperTriangular}. \end{cases} +x \mapsto \begin{cases} \sigma(Lx + b) & \text{where $L$ is }\mathtt{StrictlyLowerTriangular}, \\ \sigma(Ux + b) & \text{where $U$ is }\mathtt{StrictlyUpperTriangular}. \end{cases} ``` -The functor can be applied to a vector, a matrix or a tensor. The special matrices are implemented as [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular) and [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular). +The functor can be applied to a vector, a matrix or a tensor. The special matrices are implemented as [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular) and [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular). """ abstract type VolumePreservingFeedForwardLayer{M, N, bias} <: AbstractExplicitLayer{M, N} end @@ -36,7 +36,7 @@ x = \begin{pmatrix} 1 \\ 1 \end{pmatrix} using GeometricMachineLearning l = VolumePreservingLowerLayer(2, identity; use_bias=false) -A = LowerTriangular([1,], 2) +A = StrictlyLowerTriangular([1,], 2) ps = (weight = A, ) x = ones(eltype(A), 2) @@ -93,7 +93,7 @@ x = \begin{pmatrix} 1 \\ 1 \end{pmatrix} using GeometricMachineLearning l = VolumePreservingUpperLayer(2, identity; use_bias=false) -A = UpperTriangular([1,], 2) +A = StrictlyUpperTriangular([1,], 2) ps = (weight = A, ) x = ones(eltype(A), 2) @@ -141,7 +141,7 @@ function initialparameters( init_weight!(rng, S) init_bias!(rng, b) - (weight = LowerTriangular(S, M), bias = b) + (weight = StrictlyLowerTriangular(S, M), bias = b) end function initialparameters( @@ -153,7 +153,7 @@ function initialparameters( init_weight!(rng, S) init_bias!(rng, b) - (weight = UpperTriangular(S, M), bias = b) + (weight = StrictlyUpperTriangular(S, M), bias = b) end function initialparameters( @@ -163,7 +163,7 @@ function initialparameters( S = KernelAbstractions.allocate(backend, T, parameterlength(d)) init_weight!(rng, S) - (weight = LowerTriangular(S, M),) + (weight = StrictlyLowerTriangular(S, M),) end function initialparameters( @@ -173,7 +173,7 @@ function initialparameters( S = KernelAbstractions.allocate(backend, T, parameterlength(d)) init_weight!(rng, S) - (weight = UpperTriangular(S, M),) + (weight = StrictlyUpperTriangular(S, M),) end function (d::VolumePreservingFeedForwardLayer{M, M, :bias})( diff --git a/test/arrays/triangular.jl b/test/arrays/triangular.jl index 06226a1b8..abb754f66 100644 --- a/test/arrays/triangular.jl +++ b/test/arrays/triangular.jl @@ -12,8 +12,8 @@ Random.seed!(1234) # third axis of a tensor with `mat_tensor_mul`, and the pullback of that kernel. function triangular_tensor_multiplication_test(T = Float64, n::Int = 5) - Aₗ = rand(LowerTriangular{T}, n) - Aᵤ = rand(UpperTriangular{T}, n) + Aₗ = rand(StrictlyLowerTriangular{T}, n) + Aᵤ = rand(StrictlyUpperTriangular{T}, n) B = rand(T, n, n, n) AₗB = mat_tensor_mul(Aₗ, B) @@ -25,8 +25,8 @@ function triangular_tensor_multiplication_test(T = Float64, n::Int = 5) end function triangular_tensor_multiplication_pullback_test(T = Float64, n::Int = 5) - Aₗ = rand(LowerTriangular{T}, n) - Aᵤ = rand(UpperTriangular{T}, n) + Aₗ = rand(StrictlyLowerTriangular{T}, n) + Aᵤ = rand(StrictlyUpperTriangular{T}, n) B = rand(T, n, n, n) C_diff = rand(T, n, n, n) diff --git a/test/helpers/network_parameters_gradient_structure.jl b/test/helpers/network_parameters_gradient_structure.jl index 114af9229..eda77a2a1 100644 --- a/test/helpers/network_parameters_gradient_structure.jl +++ b/test/helpers/network_parameters_gradient_structure.jl @@ -2,32 +2,25 @@ using Test using GeometricMachineLearning using GeometricMachineLearning: params using LinearAlgebra: norm +using NeuralNetworkParameters: storage_gradient using Zygote: gradient """ -Assert which shape the `SymmetricMatrix` leaf of a gradient has, and that the shape is cosmetic. +Assert the shape and the value of the `SymmetricMatrix` leaf of a gradient taken with respect to a +`NetworkParameters` wrapper. -A gradient taken with respect to a `NetworkParameters` wrapper keeps its `SymmetricMatrix` leaf -only when the loss expression performs a single `getproperty` on that wrapper. A second access -degrades the leaf to a plain `Matrix`. What decides the split is the number of `getproperty` calls -on the wrapper itself -- not how often the leaf is used, and not how often the plain `NamedTuple`s -below the wrapper are accessed. `p -> (A = p.L1.A; sum(A) + sum(A))` uses the leaf twice through one -access and keeps the structure; `p -> sum(p.L1.A) + sum(p.L1.A)` loses it. A `Zygote.gradient` taken -with respect to the bare wrapped `NamedTuple` (`params(ps)`) keeps the structure at any access -count, because `gradient` projects its result and `ChainRulesCore.ProjectTo` restores the leaf for a -`NamedTuple` but is `identity` for a `NetworkParameters`. The underlying `Zygote.pullback` loses the -structure on the second access either way. `CHANGELOG.md` records the experiment. +Since `NeuralNetworkParameters` 0.4 that gradient is the gradient with respect to the leaf's +*storage*: `NeuralNetworkParameters` converts the cotangent of each leaf with `storage_gradient`. A +`SymmetricMatrix` keeps each off-diagonal number once and shows it twice, so its storage gradient +doubles the off-diagonal entries of the cotangent of the dense matrix, and it is a `SymmetricMatrix` +however often the loss reads the leaf. Until 0.4 a second `getproperty` on the wrapper degraded the +leaf to a plain `Matrix` (issue #319). -A wrapper unwrapped with `values` rather than `getproperty` -- which is how `Chain` hands a layer -its parameters -- counts differently: there the accesses on the plain `NamedTuple` below decide. -See the layer testset in -`test/layers/symplectic_attention_network_parameters_gradient.jl`. - -`preserves_structure` states which side of that split `f` sits on. Both gradients are asserted to -agree numerically either way, so losing the structure is a change of type and not of value. +A `Zygote.gradient` taken with respect to the bare wrapped `NamedTuple` (`params(ps)`) is the +cotangent of the dense matrix, projected back onto a `SymmetricMatrix` by +`ChainRulesCore.ProjectTo`. The two are asserted to differ by exactly `storage_gradient`. """ -function network_parameters_gradient_structure_test( - f, ps::NetworkParameters, preserves_structure::Bool) +function network_parameters_gradient_structure_test(f, ps::NetworkParameters) g_wrapped = gradient(_ps -> norm(f(_ps)), ps)[1] g_bare = gradient(_ps -> norm(f(_ps)), params(ps))[1] @@ -35,11 +28,6 @@ function network_parameters_gradient_structure_test( a_bare = g_bare.L1.A @test typeof(a_bare) <: SymmetricMatrix - if preserves_structure - @test typeof(a_wrapped) <: SymmetricMatrix - else - @test_broken typeof(a_wrapped) <: SymmetricMatrix # issue #319 - end - - @test isapprox(Matrix(a_wrapped), Matrix(a_bare)) + @test typeof(a_wrapped) <: SymmetricMatrix + @test isapprox(Matrix(a_wrapped), Matrix(storage_gradient(ps.L1.A, a_bare))) end diff --git a/test/integration/changebackend_tests.jl b/test/integration/changebackend_tests.jl index e816364bf..19c23bf3b 100644 --- a/test/integration/changebackend_tests.jl +++ b/test/integration/changebackend_tests.jl @@ -35,8 +35,8 @@ leaves = ( grassmann = rand(GrassmannManifold{Float64}, N, n), symmetric = SymmetricMatrix(rand(10), 4), skew = SkewSymMatrix(rand(6), 4), - lower = LowerTriangular(rand(6), 4), - upper = UpperTriangular(rand(6), 4), + lower = StrictlyLowerTriangular(rand(6), 4), + upper = StrictlyUpperTriangular(rand(6), 4), stiefhor = StiefelLieAlgHorMatrix(SkewSymMatrix(rand(n, n)), rand(N - n, n), N, n), grasshor = GrassmannLieAlgHorMatrix(rand(N - n, n), N, n) ) diff --git a/test/integration/hdf5_support.jl b/test/integration/hdf5_support.jl index 6e04ea5f8..7e29aa89f 100644 --- a/test/integration/hdf5_support.jl +++ b/test/integration/hdf5_support.jl @@ -14,8 +14,8 @@ _ps_eq(a::AbstractArray, b::AbstractArray) = a ≈ b _ps_eq(a::StiefelManifold, b::StiefelManifold) = a.A ≈ b.A _ps_eq(a::SymmetricMatrix, b::SymmetricMatrix) = a.S ≈ b.S && a.n == b.n _ps_eq(a::SkewSymMatrix, b::SkewSymMatrix) = a.S ≈ b.S && a.n == b.n -_ps_eq(a::LowerTriangular, b::LowerTriangular) = a.S ≈ b.S && a.n == b.n -_ps_eq(a::UpperTriangular, b::UpperTriangular) = a.S ≈ b.S && a.n == b.n +_ps_eq(a::StrictlyLowerTriangular, b::StrictlyLowerTriangular) = a.S ≈ b.S && a.n == b.n +_ps_eq(a::StrictlyUpperTriangular, b::StrictlyUpperTriangular) = a.S ≈ b.S && a.n == b.n function _ps_eq(a::NamedTuple, b::NamedTuple) Set(keys(a)) == Set(keys(b)) || return false all(_ps_eq(a[k], b[k]) for k in keys(a)) @@ -48,7 +48,8 @@ function _write_legacy(h5, Y::StiefelManifold, path::AbstractString) end for (T, name) in ((:SymmetricMatrix, "SymmetricMatrix"), (:SkewSymMatrix, "SkewSymMatrix"), - (:LowerTriangular, "LowerTriangular"), (:UpperTriangular, "UpperTriangular")) + # the tags GML 0.8 and earlier wrote, under the names the two types had then + (:StrictlyLowerTriangular, "LowerTriangular"), (:StrictlyUpperTriangular, "UpperTriangular")) @eval function _write_legacy(h5, A::$T, path::AbstractString) g = HDF5.create_group(h5, path) HDF5.attributes(g)["gml_type"] = $name @@ -131,8 +132,8 @@ end end end -# VolumePreservingFeedForward uses LowerTriangular / UpperTriangular parameters. -@testset "save/load roundtrip: VolumePreservingFeedForward (LowerTriangular/UpperTriangular)" begin +# VolumePreservingFeedForward uses StrictlyLowerTriangular / StrictlyUpperTriangular parameters. +@testset "save/load roundtrip: VolumePreservingFeedForward (the triangular matrices)" begin arch = VolumePreservingFeedForward(4, 4, 1, tanh) nn = NeuralNetwork(arch) x = rand(4) @@ -215,7 +216,10 @@ end @testset "a file in the old gml_type layout still loads" begin for (name, arch, dimin) in (("LASympNet (SymmetricMatrix)", LASympNet(4), 4), ("SymplecticAutoencoder (StiefelManifold)", - SymplecticAutoencoder(10, 4), 10)) + SymplecticAutoencoder(10, 4), 10), + # tagged `LowerTriangular`/`UpperTriangular`, the names before GeometricOptimizers 0.9 + ("VolumePreservingFeedForward (StrictlyLowerTriangular, StrictlyUpperTriangular)", + VolumePreservingFeedForward(4, 4, 1, tanh), 4)) @testset "$name" begin nn = NeuralNetwork(arch) ps = params(nn) diff --git a/test/layers/double_multiplication_network_parameters_gradient.jl b/test/layers/double_multiplication_network_parameters_gradient.jl index d2fb3557c..0079a62dc 100644 --- a/test/layers/double_multiplication_network_parameters_gradient.jl +++ b/test/layers/double_multiplication_network_parameters_gradient.jl @@ -22,12 +22,12 @@ ps = NetworkParameters((L1 = (A = S,),)) t = rand(4, 4) @testset "_custom_mul: NetworkParameters gradient structure" begin - @testset "single_multiplication (one wrapper access: structure kept)" begin + @testset "single_multiplication" begin network_parameters_gradient_structure_test( - _ps -> single_multiplication(t, _ps), ps, true) + _ps -> single_multiplication(t, _ps), ps) end - @testset "double_multiplication (two wrapper accesses: structure lost)" begin + @testset "double_multiplication" begin network_parameters_gradient_structure_test( - _ps -> double_multiplication(t, _ps), ps, false) + _ps -> double_multiplication(t, _ps), ps) end end diff --git a/test/layers/symplectic_attention_network_parameters_gradient.jl b/test/layers/symplectic_attention_network_parameters_gradient.jl index 405794832..518fa56bc 100644 --- a/test/layers/symplectic_attention_network_parameters_gradient.jl +++ b/test/layers/symplectic_attention_network_parameters_gradient.jl @@ -28,28 +28,24 @@ ps = NetworkParameters((L1 = (A = S,),)) t = (q = rand(2, 2), p = rand(2, 2)) @testset "Symplectic attention: NetworkParameters gradient structure" begin - @testset "symplectic_attention (two wrapper accesses: structure lost)" begin + @testset "symplectic_attention" begin network_parameters_gradient_structure_test( - _ps -> symplectic_attention(t, _ps), ps, false) + _ps -> symplectic_attention(t, _ps), ps) end - @testset "symplectic_attention_simplified (one wrapper access: structure kept)" begin + @testset "symplectic_attention_simplified" begin network_parameters_gradient_structure_test( - _ps -> symplectic_attention_simplified(t, _ps), ps, true) + _ps -> symplectic_attention_simplified(t, _ps), ps) end - @testset "symplectic_linear_map (one wrapper access: structure kept)" begin + @testset "symplectic_linear_map" begin network_parameters_gradient_structure_test( - _ps -> symplectic_linear_map(t, _ps), ps, true) + _ps -> symplectic_linear_map(t, _ps), ps) end end -# The layer keeps the structure although its body reads `A` twice. `Chain` reaches the layer -# through `values(_ps)` rather than through a `getproperty` on the wrapper, and on that route the -# count that decides the split is the one on the plain `NamedTuple` the layer is handed. The local -# binding at `src/layers/symplectic_attention.jl` makes that count one, and is what keeps the -# structure here: the same body written without it, reached the same way, degrades the leaf to a -# plain `Matrix`. +# The same through a layer, which `Chain` reaches through `values(_ps)` rather than through a +# `getproperty` on the wrapper. @testset "SymplecticAttentionQ layer: NetworkParameters gradient structure" begin l = SymplecticAttentionQ(4; symmetric = true) nn = NeuralNetwork(Chain(l)) - network_parameters_gradient_structure_test(_ps -> nn(t, _ps), params(nn), true) + network_parameters_gradient_structure_test(_ps -> nn(t, _ps), params(nn)) end diff --git a/test/layers/volume_preserving_feedforward.jl b/test/layers/volume_preserving_feedforward.jl index 673951afc..b8f50cbb4 100644 --- a/test/layers/volume_preserving_feedforward.jl +++ b/test/layers/volume_preserving_feedforward.jl @@ -35,11 +35,11 @@ test_volume_preserving_feedforward(; T = Float32) test_volume_preserving_feedforward(; T = Float64) # A vector in, a vector out. Neither layer shapes the result: the `:no_bias` method is -# `x + d.activation.(ps.weight * x)`, and `LowerTriangular * ::AbstractVector` returns a vector, -# which is what `LinearAlgebra` returns for every other matrix type. +# `x + d.activation.(ps.weight * x)`, and `StrictlyLowerTriangular * ::AbstractVector` returns a +# vector, which is what `LinearAlgebra` returns for every other matrix type. l = VolumePreservingLowerLayer(2, identity; use_bias = false) -A = LowerTriangular([1], 2) +A = StrictlyLowerTriangular([1], 2) @test l(ones(eltype(A), 2), (weight = A,)) == [1, 2] l = VolumePreservingUpperLayer(2, identity; use_bias = false) -A = UpperTriangular([1], 2) +A = StrictlyUpperTriangular([1], 2) @test l(ones(eltype(A), 2), (weight = A,)) == [2, 1] diff --git a/test/map_to_cpu_tests.jl b/test/map_to_cpu_tests.jl index c7d085421..0e9c1ccf5 100644 --- a/test/map_to_cpu_tests.jl +++ b/test/map_to_cpu_tests.jl @@ -12,7 +12,7 @@ using Random using Test import GeometricOptimizers: StiefelManifold, SymmetricMatrix, SkewSymMatrix, - LowerTriangular, UpperTriangular + StrictlyLowerTriangular, StrictlyUpperTriangular Random.seed!(1234) @@ -22,7 +22,7 @@ const N, n = 6, 3 ps = NetworkParameters(( L1 = (Y = rand(StiefelManifold{Float64}, N, n), b = randn(N)), L2 = (S = SymmetricMatrix(randn(n, n)), A = SkewSymMatrix(randn(n, n))), - L3 = (L = LowerTriangular(randn(n, n)), U = UpperTriangular(randn(n, n))) + L3 = (L = StrictlyLowerTriangular(randn(n, n)), U = StrictlyUpperTriangular(randn(n, n))) )) back = map_to_cpu(ps) @@ -35,8 +35,8 @@ const N, n = 6, 3 @test back.L1.Y isa StiefelManifold @test back.L2.S isa SymmetricMatrix @test back.L2.A isa SkewSymMatrix - @test back.L3.L isa LowerTriangular - @test back.L3.U isa UpperTriangular + @test back.L3.L isa StrictlyLowerTriangular + @test back.L3.U isa StrictlyUpperTriangular # the `n` a structured leaf carries is not in its storage; `rebuild` takes it from the prototype @test back.L2.S.n == ps.L2.S.n diff --git a/test/optimizers/structured_array_parameters.jl b/test/optimizers/structured_array_parameters.jl index 4ce1a08d3..a685780c5 100644 --- a/test/optimizers/structured_array_parameters.jl +++ b/test/optimizers/structured_array_parameters.jl @@ -8,7 +8,7 @@ Random.seed!(123) @doc raw""" Every optimizer method has to work with the structured matrix types GML uses as *ordinary* weights: `SymmetricMatrix` (SympNet and symplectic attention layers), `SkewSymMatrix` (volume-preserving -attention) and `LowerTriangular`/`UpperTriangular` (volume-preserving feedforward layers). +attention) and `StrictlyLowerTriangular`/`StrictlyUpperTriangular` (volume-preserving feedforward layers). These have their own storage and no `setindex!`, so they need a type-preserving `similar` and elementwise `_add!`/`_rac!`/`_square!`/`_div!`/`_rmul!` written on the free parameters; without them @@ -41,7 +41,7 @@ end # then fails to match them against the parameter @testset "structured weights keep their type in `similar`" begin for A in (rand(SymmetricMatrix{Float64}, 4), rand(SkewSymMatrix{Float64}, 4), - rand(LowerTriangular{Float64}, 4), rand(UpperTriangular{Float64}, 4)) + rand(StrictlyLowerTriangular{Float64}, 4), rand(StrictlyUpperTriangular{Float64}, 4)) @test typeof(similar(A)) == typeof(A) @test typeof(zero(A)) == typeof(A) @test size(similar(A)) == size(A) @@ -51,7 +51,7 @@ end @testset "`LASympNet` (SymmetricMatrix)" begin optimizer_runs(LASympNet(4), Batch(5), 4) end -@testset "`VolumePreservingFeedForward` (Lower/UpperTriangular)" begin +@testset "`VolumePreservingFeedForward` (the triangular matrices)" begin optimizer_runs(VolumePreservingFeedForward(4), Batch(5), 4) end @testset "`VolumePreservingTransformer` (SkewSymMatrix + triangular)" begin diff --git a/test/quality/aqua.jl b/test/quality/aqua.jl index 8ce09ad3d..00ad00956 100644 --- a/test/quality/aqua.jl +++ b/test/quality/aqua.jl @@ -3,43 +3,20 @@ using ExplicitImports using GeometricMachineLearning using Test -# Aqua's package-level checks. Six of the eight `Aqua.test_all` enables by default run; `ambiguities` -# and `piracies` are switched off by name, and the reason is below rather than in the changelog -# alone, because switching one back on turns the suite red on the spot. +# Aqua's package-level checks, all eight that `Aqua.test_all` enables by default. # # `Aqua.test_all` wraps each of its checks in a `@testset` of its own and adds no enclosing one, and # a testset with no parent finalises as soon as it closes -- so called bare, the first failing check # throws a `TestSetException` and the rest never run. On the suite path the `@safetestset` in # `runtests.jl` is that parent. The `@testset` below is that parent when this file is run on its -# own, and it is also what lets the two counts at the end be further assertions rather than -# top-level `@test`s that abort the file the moment one fails. +# own. # -# WHAT IS SWITCHED OFF, AND WHY IT IS NOT A `broken = true`. -# -# `piracies` reports 3 methods, and all 3 are genuine under Aqua's definition -- the function and -# every argument type belong to other modules. They are not a count waiting to be triaged. All -# three are the same thing: a three-argument functor for a *layer* type that this package does not -# own, `AbstractNeuralNetworks`' `Dense` and `Linear`, added so that the layer accepts a 3-tensor -# input as well as a matrix. -# -# Dense{M, N, true} on an AbstractArray{T, 3} src/layers/resnet.jl:63 -# Dense{M, N, false} on an AbstractArray{T, 3} src/layers/resnet.jl:67 -# Linear{M, N} on an AbstractArray{T, 3} src/layers/resnet.jl:71 -# -# Closing these means either upstream gaining the 3-tensor methods or this package wrapping the two -# layer types, and both are an API change rather than a tidy-up. They are *B7* under -# `## Open Issues` in `CHANGELOG.md` with that reasoning. -# -# `ambiguities` reports 1, and the same 1 whether this package is loaded alone or with everything -# `runtests.jl` loads before this file. It is a `Dense{M, N, true}` functor against -# `AbstractNeuralNetworks.Affine`'s, and it is benign: `Dense` is not a subtype of `Affine` and the -# two have no common instance, so no call can reach the pair. It is the same `resnet.jl:63` method -# as the first piracy above. -# -# The two settings agreeing is what makes the gate below trustworthy, and it holds because no -# method of this package claims an argument position wholesale on `PoissonTensor`. That type is an -# `AbstractMatrix{T}`, so such a method meets every special-array method another package writes -# against `AbstractMatrix`: +# `ambiguities` and `piracies` were switched off until 0.9, with a count gate for each in their +# place: three pirated 3-tensor functors on `AbstractNeuralNetworks`' `Dense` and `Linear`, and one +# ambiguity between the first of them and `Affine`'s. `AbstractNeuralNetworks` 0.9 defines those +# methods itself, so they are deleted here and both checks run. No method of this package claims an +# argument position wholesale on `PoissonTensor`, which is an `AbstractMatrix{T}` and would +# otherwise meet every special-array method another package writes against `AbstractMatrix`: # # `*(𝕁::PoissonTensor{T}, v::Strided…)` takes a `Strided…` right-hand side, which excludes each # special array type `X` that ArrayLayouts, FillArrays, Symbolics and GeometricOptimizers @@ -48,29 +25,8 @@ using Test # `BlockIndex` and `BandRangeType` methods BandedMatrices and BlockArrays add to # `AbstractMatrix`. Neither of those two loads with this package alone, so an untyped index # pair here is a collision visible only from inside the suite. -# -# `ambiguities` still stays off, because 1 is not 0 and Aqua's check has no way to exempt a pair. -# The assertion below replaces it and is strictly better here: it names the number, so it fails -# both when an ambiguity is added and when this one is closed without the comment going with it. -# -# Marking either check `broken = true` would leave a check that reports success while the defect -# stands, which is the failure mode this test suite's guards exist to remove. @testset "Aqua" begin - Aqua.test_all(GeometricMachineLearning; ambiguities = false, piracies = false) - - # A switched-off check detects nothing, so the count it would have reported drifts unobserved, - # and so does every `src` line named above. This is the gate. It fails when a piracy is added, - # and it fails when one is removed without the entry above and in `CHANGELOG.md` going with it. - @test length(Aqua.Piracy.hunt(GeometricMachineLearning)) == 3 - - # The ambiguity count gets the same gate, and it can carry one because the single pair is - # between a method of this package and one of `AbstractNeuralNetworks`'. The count is the same - # in isolation and in the suite, and nothing outside this repository's own `[compat]` moves it - # -- so a red here is a regression, not an unrelated upstream upgrade and not a reordering of - # this suite. A count that moved with ArrayLayouts', FillArrays', Symbolics', - # GeometricOptimizers', BandedMatrices' and BlockArrays' versions, and with *which* of them a - # given process had loaded, could not tell those apart, and could not be asserted. - @test length(Test.detect_ambiguities(GeometricMachineLearning; recursive = true)) == 1 + Aqua.test_all(GeometricMachineLearning) end # `ExplicitImports` answers a question Aqua does not ask, and the reason to gate on it is not From 84d0c4efb3c4bf6c6fba5a056d150b95ce9ba294 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Mon, 5 Oct 2026 17:56:30 +0200 Subject: [PATCH 3/3] Train through GeometricOptimizers' TrainingOptimizer, one method per leaf MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GML's half of G4. `Optimizer` holds one GeometricOptimizers.TrainingOptimizer over the whole parameter set, and `optimization_step!(ps, opt, dp)` is GeometricOptimizers' function extended for it. The per-layer walk, the Euclidean update rules (GMLEuclideanState), _GMLGradient and the λY argument of the step and of optimize_for_one_epoch! are gone. A CompositeMethod (JuliaGNI/GeometricOptimizers.jl#121) now passes straight through: every leaf gets the cache and state of the method selected for it, so ScalarMomentAdam can train the Stiefel weights of a network beside Adam on the rest. test/optimizers/composite_method.jl checks the per-leaf states and that a composite of Adam on every leaf reproduces Adam exactly. The GeometricOptimizers 0.9 API throughout: methods without an element type, Cayley()/Geodesic() as retractions. KNOWN_ISSUES C2 is closed and C1 narrowed to the shared name `Optimizer`. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 64 +++ KNOWN_ISSUES.md | 77 +-- docs/src/introduction.md | 2 +- docs/src/optimizers/optimizer.md | 25 +- docs/src/tutorials/grassmann_layer.md | 9 +- .../tutorials/hamiltonian_neural_network.md | 4 +- .../linear_symplectic_transformer.md | 2 +- docs/src/tutorials/symplectic_autoencoder.md | 2 +- docs/src/tutorials/symplectic_transformer.md | 2 +- docs/src/tutorials/sympnet_tutorial.md | 2 +- .../tutorials/volume_preserving_attention.md | 6 +- ...olume_preserving_transformer_rigid_body.md | 2 +- .../symplectic_autoencoders/training.jl | 2 +- .../reproduction/sympnets/sympnet_pendulum.jl | 3 +- .../double_pendulum.jl | 2 +- .../rigid_body.jl | 2 +- src/GeometricMachineLearning.jl | 19 +- src/data_loader/optimize.jl | 16 +- src/optimizers/optimizer.jl | 449 +++--------------- .../data_loader_optimization_step.jl | 3 +- .../optimizer_functor_with_adam.jl | 3 +- test/layers/gradient_layer_tests.jl | 7 +- test/layers/manifold_layers.jl | 10 +- ...ulti_head_attention_stiefel_optim_cache.jl | 12 +- ...multi_head_attention_stiefel_retraction.jl | 2 +- .../multi_head_attention_stiefel_setup.jl | 2 +- test/layers/transformer_optimizer.jl | 9 +- .../adam_with_learning_rate_decay.jl | 24 +- test/optimizers/composite_method.jl | 55 +++ test/optimizers/gml_gradient_dispatch.jl | 129 ----- test/optimizers/gradient_optimizer.jl | 6 +- test/optimizers/momentum_optimizer.jl | 10 +- test/optimizers/optimization_step.jl | 26 +- test/optimizers/psd_optim.jl | 7 +- test/optimizers/step_size_element_type.jl | 52 -- .../optimizers/structured_array_parameters.jl | 5 +- test/optimizers/svd_optim.jl | 7 +- test/runtests.jl | 3 +- 38 files changed, 296 insertions(+), 766 deletions(-) create mode 100644 test/optimizers/composite_method.jl delete mode 100644 test/optimizers/gml_gradient_dispatch.jl delete mode 100644 test/optimizers/step_size_element_type.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a84549a8..b4919d50b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,11 +17,75 @@ breaking release). ### Added +- **One optimizer method per leaf, through `CompositeMethod`.** `CompositeMethod` and + `CompositeState` are re-exported from `GeometricOptimizers`, with `ScalarMomentAdam` and + `ScalarMomentAdamState`. `Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn)` + gives every leaf of the network's parameters the cache and the state of the method chosen for it, + which is what a network with Stiefel weights beside Euclidean ones needs to use `ScalarMomentAdam`: + that method steps a single `StiefelManifold`. A caller that wanted the split before had to add + methods to `_make_optimizer_cache`, `_make_optimizer_state` and `_leaf_optim_step!` from outside + this package. `test/optimizers/composite_method.jl` checks the per-leaf states and that a composite + of `Adam` on every leaf reproduces `Adam` exactly. - **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and the `tensor_mat_mul!` kernel on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 all five assertions pass. - **A `Metal` workflow runs the Metal tests on GitHub's `macos-15` runner**, for the `min` and `1` Julia versions. The runner's GPU is Apple's paravirtual device, which Metal.jl supports from 1.10 and on macOS 15 or later. The workflow includes a step "Metal is functional" before the tests that fails the job where the device is unavailable, so it cannot pass without having run the tests. The job is not a required check. However, the Metal tests also run in two required checks—CI.yml's `Julia min - macOS-latest - default` and `Julia 1 - macOS-latest - default`—which execute `test/runtests.jl` with no test arguments on Apple silicon, selecting the `metal` group. - **`test/quality/jet.jl` checks the kernel launchers with JET.** It runs in `core`, directly after `quality/aqua.jl`, and takes about 14 s. Each function that launches a kernel of `src/` on the `CPU()` backend, the pullbacks among them, and the `Batch` functor that `test/data_loader/batch_index_set.jl` measures with `@allocated` get one `JET.report_opt` line per element type that the tests use, with `target_modules = (GeometricMachineLearning,)`. Each `@kernel` method also gets a line that analyses its generated `cpu_` function at the launcher's context, because JET drops the reports of a kernel body when it analyses the launcher. 142 lines pass on Julia 1.11.9 (JET 0.9.20) and 1.13.1 (JET 0.12.2). Three lines report a runtime dispatch on both versions and are `@test_broken` with issue #325: `PoissonTensor(::CPU, ::Int, ::DataType)`, and the input-and-output method of `convert_input_and_batch_indices_to_array` at `Float32` and `Float64`. Where JET does not work, the file records one `@test_skip`. JET joins `test/Project.toml` with no `[compat]` bound. `KNOWN_ISSUES.md` K2 records the Revise `EMFILE` messages that JET brings into the test log, and K3 the dispatches that the lines do not see: an unstable argument of a kernel launch, and two barriers in the Poisson-tensor kernel. - **A test calls `iterate` on `(q,)` initial conditions.** No test reached the `NeuralNetworkIntegrator` `iterate` method that takes `ics::NamedTuple{(:q,), Tuple{AT}}`. `test/docstrings/layers_and_architectures.jl` runs the docstring's `ResNet` from `(q = [1, 1, 1],)` for four points and expects the same trajectory as the vector form, `[1 2 4 8; 1 3 9 27; 1 3 7 15]`. +### Changed (breaking) + +- **The training step is `GeometricOptimizers.TrainingOptimizer`.** `Optimizer` holds one, over the + whole parameter set of the network, in its only field `training`; the fields `method`, `cache`, + `state`, `retraction`, `step_size` and `iterations` are gone. One cache and one state cover the + whole set, where there used to be one per layer, and every method steps every kind of leaf the same + way `GeometricOptimizers` does: no leaf is stepped by this package's own Euclidean update rules + (`GMLEuclideanState`) any more. `GeometricOptimizers`' G4 checked that `TrainingOptimizer` + reproduces this package's 0.8 iterates: exactly for `Adam`, and to `4eps(T)` for `GradientMethod` and + `MomentumMethod` on manifold weights. `ScalarMomentAdam` trains the Stiefel weights of a network + through a `CompositeMethod` (see *Added*); the step used to raise a `MethodError`, since this + package had no method of its own for it. +- **`optimization_step!(ps, opt, dp)` replaces `optimization_step!(opt, λY, ps, dp)`**, and + `optimize_for_one_epoch!` loses its last argument `λY`. The section is carried by the optimizer's + state, so the caller no longer builds a `GlobalSection` of the parameters for it. + `optimization_step!` is `GeometricOptimizers`' function, extended for `Optimizer`, so a `Main` that + loads both packages has one `optimization_step!` and not two. `dp` is the gradient a pullback + returns, a `NetworkParameters` or the bare `NamedTuple` of the same shape. +- **`retraction` takes `Cayley()` or `Geodesic()`**, the retraction types, and no longer the + functions `cayley` and `geodesic`. The default is `Cayley()`. +- **The methods carry no element type**, as in `GeometricOptimizers` 0.9: `Adam()`, + `Adam(; β₁, β₂, δ)`, `MomentumMethod(; α)` and `AdamOptimizerWithDecay(n_epochs; η₁, η₂)` in place + of `Adam(Float32)`, `Adam(T; …)`, `MomentumMethod(α)` and `AdamOptimizerWithDecay(n_epochs, T)`. + `Optimizer` converts the method to the element type of the parameters. +- **`LowerTriangular` and `UpperTriangular` are `StrictlyLowerTriangular` and + `StrictlyUpperTriangular`**, the names `GeometricOptimizers` 0.9 gives them, which no longer collide + with `LinearAlgebra`'s. A file saved with one of them by 0.8 or earlier is tagged with the old name + and still loads (JuliaGNI/GeometricOptimizers.jl#146); `test/integration/hdf5_support.jl` loads a + `VolumePreservingFeedForward` network written in the old layout. +- **Requires `GeometricOptimizers` 0.9, `AbstractNeuralNetworks` 0.9, `NeuralNetworkParameters` + 0.4.2, `GeometricBase` 0.15 and `SymbolicNeuralNetworks` 0.9.** No registered `GeometricOptimizers` + before 0.9 allows `NeuralNetworkParameters` 0.4 or `AbstractNeuralNetworks` 0.9, and no registered + `SymbolicNeuralNetworks` before 0.9 allows either (JuliaGNI/SymbolicNeuralNetworks.jl#80), so the + five move together. Under `NeuralNetworkParameters` 0.4 a Zygote gradient taken through a + `NetworkParameters` is the gradient with respect to each leaf's storage: a `SymmetricMatrix` leaf + keeps its type however often the loss reads it, which closes #319, and its off-diagonal entries are + twice those of the cotangent of the dense matrix. `test/helpers/network_parameters_gradient_structure.jl` + asserts both. + +### Removed (breaking) + +- **The three-axis methods of `Dense` and `Linear`, and the HDF5 extension.** `AbstractNeuralNetworks` + 0.9 defines both itself: `Dense`, `Linear` and `Affine` take inputs of rank 3 and more, and its + `HDF5Ext` has `save` and `load` for a `NeuralNetwork`, so `using HDF5` still loads them. Here they + were type piracy, the three methods in `src/layers/resnet.jl` and the four in `ext/HDF5Ext.jl`. + Aqua's `piracies` and `ambiguities` checks now run in `test/quality/aqua.jl` and report none; they + were switched off with counts of 3 and 1 in their place (*B7* in `KNOWN_ISSUES.md`, now closed). + `load(NeuralNetwork, …)` loses its `backend` keyword: the network loads on the CPU, and + `changebackend(backend, nn)` moves it. +- **`GMLEuclideanState`, `_GMLGradient` and the per-layer optimizer walk** (`_make_optimizer_cache`, + `_make_optimizer_state`, `_tree_optim_step!`, `_leaf_optim_step!`, `_as_go_solution`) are gone with + it, and with them `test/optimizers/gml_gradient_dispatch.jl` and + `test/optimizers/step_size_element_type.jl`, which tested two of them. *C2* in `KNOWN_ISSUES.md` + is closed; *C1* is narrowed to the one thing left, the name `Optimizer`, which both packages export. + ### Changed - **The package needs Julia 1.12 or later** (`[compat] julia = "1.12"`, was `"1.11"`). On Julia 1.11.9, the gradient of an `LNNLoss` with respect to the network parameters stops with `StackOverflowError` in type inference ("Internal error: during type inference of … Encountered stack overflow") on macOS and Windows; the `Julia min` job on Linux passes. The overflow comes from SymbolicUtils 4.49.0, whose `rrule` for `SymbolicUtils.Code.create_array` returns an exactly typed tuple of tangents, one per array element; Zygote's reverse pass through the generated Hessian of `lagrangian_acceleration` is then inferred far deeper. With the 4.48.0 rule restored, or in a task with a 64 MB stack, the same gradient passes on 1.11.9; Julia 1.12.7 and 1.13.1 pass it on the root task. This closes #327: the `Julia min` jobs on macOS and Windows had failed in the *Lagrangian Neural Network* testset since 2026-10-01, and every testset after it was never run there. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 7ebbc21f1..cccbe5b90 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -45,37 +45,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. every type parameter, and Aqua's `unbound_args` passes on Julia nightly. See *Fixed* above. The number is left vacant rather than reused.) -### B7 · Three methods are type piracy, and closing them is an API change rather than a deletion. - -- **location:** `src/layers/resnet.jl:63` -- **kind:** defect -- **found:** 2026-09-19 -- **evidence:** - Aqua's `piracies` check reports 3, down from the 12 this entry opened with; the other nine were - deleted in this release and are under *Removed (breaking)* above. All three are genuine under - Aqua's definition — the function and every argument type belong to other modules — and all three - are one thing: a functor on a *layer* type this package does not own, applied to a three-axis - array. - - | method | `src` | without GML | with GML | - |:--|:--|:--|:--| - | `(::Dense{M, N, true})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:63` | `MethodError` | the layer applied along the third axis | - | `(::Dense{M, N, false})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:67` | the same | the same | - | `(::Linear{M, N})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:71` | the same | the same | - - `Dense` and `Linear` are `AbstractNeuralNetworks`', and upstream a 3-tensor argument is a - `MethodError` because `*` cannot take one. So the witness is not a changed answer but a new one: - a call that throws without this package loaded returns a value with it. - - Closing this means either `AbstractNeuralNetworks` gaining the three-axis methods, or this - package wrapping the two layer types in its own. Both are an API change, and neither is a - by-product of a piracy pass — which is why these three stayed when the other nine went. - - `src/layers/resnet.jl:63` is also the one remaining method ambiguity, against - `AbstractNeuralNetworks.Affine`. That pair is benign and has no witness at all: `Dense` is not a - subtype of `Affine` and the two have no common instance, so no call can reach it. `test/quality/aqua.jl` - asserts both counts — 3 piracies and 1 ambiguity — so neither can drift, in either direction. - ### B10 · `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` at `src/data_loader/data_loader.jl:325` has no test and appears unreachable through this package's current dependencies. - **location:** `src/data_loader/data_loader.jl:325` @@ -135,36 +104,18 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ## C. Follow-ups and cleanups -### C1 · The parameter-tree traversal still belongs upstream. +### C1 · `Optimizer` is exported by both this package and GeometricOptimizers. -- **location:** — +- **location:** `src/optimizers/optimizer.jl` - **kind:** upstream - **found:** 2026-08-17 - **evidence:** - `_make_optimizer_cache`, - `_make_optimizer_state`, `_tree_optim_step!`, `_leaf_optim_step!` and the bespoke - `GMLEuclideanState` are GML implementations of what GeometricOptimizers supports natively for a - single parameter. `GMLEuclideanState` in particular duplicates what `GradientState`, - `MomentumState` and `AdamState` already do for a plain array. - - What has to go upstream is *not* a reuse of GeometricOptimizers' `Optimizer`: that one needs an - `OptimizerProblem`, i.e. an objective function, and minibatch training has none — the gradient - arrives from AD one batch at a time. It is a new entry point there, a - gradient-supplied-externally step over a `NamedTuple` parameter tree. GML's `Optimizer` would then - be the `NeuralNetwork` constructor and the training-loop functor, and nothing else. - - `Optimizer` is the one name still exported by both packages, so this is also what closes the last - of B1's class of collision. - -### C2 · Two `isa` branches remain in `_leaf_optim_step!` - -- **location:** — -- **kind:** defect -- **found:** 2026-08-16 -- **evidence:** - (for `AdamState`/`MomentumState`). - Measurement showed the traversal is not implicated in the compile-time problem, so this is tidying, - and it disappears entirely if C1 lands first. + The parameter-tree traversal this entry was opened for went upstream in 0.9: the training step is + `GeometricOptimizers.TrainingOptimizer`, and GML's `Optimizer` only pairs one with a network and + carries the training-loop functor. What remains is the name. GML's `Optimizer` and + GeometricOptimizers' `Optimizer` (the one `solve!` takes) are different types, both exported, so a + `Main` that `using`s both packages cannot write `Optimizer` unqualified. It is the last of B1's class + of collision. Closing it means renaming one of them, which is an API change in either package. ### C6 · Three generated MNIST PDFs are in this branch's history @@ -363,7 +314,7 @@ Not defects — claims this release makes that nothing has actually checked yet. ## Found late -### K1 · Four source comments still send the reader to the Open Issues section of `CHANGELOG.md`, which no longer exists. +### K1 · Two source comments still send the reader to the Open Issues section of `CHANGELOG.md`, which no longer exists. - **location:** `src/reduced_system/reduced_system.jl:109` - **kind:** found late @@ -373,17 +324,11 @@ Not defects — claims this release makes that nothing has actually checked yet. $ grep -rnE 'Open Issues|issue C1[48]|\*C21\*|\*B7\*' src test src/kernels/exponentials/tensor_exponential.jl:2:# name is wider than what it holds; renaming it, and `exponentials/` with it, is issue C18. src/reduced_system/reduced_system.jl:109:# issue C14 in `CHANGELOG.md`. - test/quality/aqua.jl:30:# layer types, and both are an API change rather than a tidy-up. They are *B7* under - test/quality/aqua.jl:31:# `## Open Issues` in `CHANGELOG.md` with that reasoning. test/reduced_system/reduced_system.jl:44:# for four of twelve, with four `NaN` and four in the opposite order. That is *C21* in - $ grep -n 'CHANGELOG' test/quality/aqua.jl - 31:# `## Open Issues` in `CHANGELOG.md` with that reasoning. - 63: # and it fails when one is removed without the entry above and in `CHANGELOG.md` going with it. ``` - The IDs are still valid; the comments at `src/reduced_system/reduced_system.jl:109`, - `test/quality/aqua.jl:31`, `test/quality/aqua.jl:63` and `test/reduced_system/reduced_system.jl:44–45` - name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`. The fix is a follow-up change to those + The IDs are still valid; the comments at `src/reduced_system/reduced_system.jl:109` and + `test/reduced_system/reduced_system.jl:44–45` name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`. The fix is a follow-up change to those comments. ## Upstream diff --git a/docs/src/introduction.md b/docs/src/introduction.md index 7a9178980..08c516cd6 100644 --- a/docs/src/introduction.md +++ b/docs/src/introduction.md @@ -19,7 +19,7 @@ Closely linked to the research presented here is the development of a software p [^0]: This document was produced with `GeometricMachineLearning` `v0.3`. It may be that the interface will slightly change in future versions, but efforts will be made to keep these changes as small as possible. ```@docs; canonical = false -optimization_step! +optimization_step!(::NetworkParameters, ::Optimizer, ::Any) ``` So the docstring shows the name of the type or method, in most cases how to call it and then gives some information explaining what it does and potentially hyperlinks to other similar docstrings ([`Optimizer`](@ref) in this case); all of this information is indented by a tab. Docstrings may include other information under subheaders **Arguments** (showing the arguments the method can be supplied with), **Examples** (giving more detailed examples (including results) of how to use the method) and **Implementation** (giving details on how the method is implemented). When we reference a docstring it is always printed in blue (e.g. [`optimization_step!`](@ref)), indicating a hyperlink. In addition there is an *index of docstrings* showing all docstrings in chronological order with the associated page number. diff --git a/docs/src/optimizers/optimizer.md b/docs/src/optimizers/optimizer.md index ca0b095ee..bed0e14ce 100644 --- a/docs/src/optimizers/optimizer.md +++ b/docs/src/optimizers/optimizer.md @@ -6,18 +6,28 @@ global section — belongs to `GeometricOptimizers` and is described in its docu [Optimization on Homogeneous Spaces](@extref GeometricOptimizers :doc:`manifold_optimizers`) and [Retractions](@extref GeometricOptimizers :doc:`retractions`). -What `GeometricMachineLearning` adds is the part that is about *neural networks*: walking the -parameter tree of a `NeuralNetwork`, applying the right update to each leaf — a retraction -for a weight on a manifold, ordinary arithmetic for a Euclidean one — and driving that from a data -loader over epochs and batches. +The training step itself is `GeometricOptimizers`' `TrainingOptimizer` as well: one cache and one +state over the whole parameter set of a network, and a retraction for every weight on a manifold. +What `GeometricMachineLearning` adds is the part that is about *neural networks*: the +[`Optimizer`](@ref) that pairs a `TrainingOptimizer` with a `NeuralNetwork`, and the loop that drives +it from a data loader over epochs and batches. The gradient comes from automatic differentiation on one batch at a time, so there is no objective function to hand a line search; the step size is a property of the [`Optimizer`](@ref) instead. It is either a number, or a `GeometricOptimizers.DecayingStatic` schedule: ```julia -opt = Optimizer(Adam(Float32), nn; step_size = 1e-3) -opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...) +opt = Optimizer(Adam(), nn; step_size = 1e-3) +opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) +``` + +A network whose weights are not all of one kind can take a different method for each kind. A +`CompositeMethod` chooses the method per leaf of the parameter set, and every leaf then gets the cache +and the state of the method chosen for it. `ScalarMomentAdam` steps a single `StiefelManifold`, so a +network with Stiefel weights beside Euclidean ones needs a composite to use it: + +```julia +opt = Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn) ``` ## Library Functions @@ -25,6 +35,5 @@ opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...) ```@docs Optimizer optimize_for_one_epoch! -optimization_step! -GeometricMachineLearning._as_go_solution +optimization_step!(::NetworkParameters, ::Optimizer, ::Any) ``` diff --git a/docs/src/tutorials/grassmann_layer.md b/docs/src/tutorials/grassmann_layer.md index 1b1acbcd8..3c4eb27c6 100644 --- a/docs/src/tutorials/grassmann_layer.md +++ b/docs/src/tutorials/grassmann_layer.md @@ -105,13 +105,6 @@ nn = NeuralNetwork(model, CPU(), Float64) nothing # hide ``` -We then *lift* the neural network parameters via `GlobalSection`. - -```@example rosenbrock -λY = GlobalSection(params(nn)) -nothing # hide -``` - As the cost function ``c`` for the Wasserstein loss[^3] we simply use: [^3]: For each Wasserstein loss we need to define such a cost function. For this particular choice of ``c`` the Wasserstein distance corresponds to a ``W_2`` loss. @@ -284,7 +277,7 @@ loss_array = zeros(training_steps) for i in 1:training_steps val, dp = compute_gradient(params(nn)) loss_array[i] = val - optimization_step!(optimizer, λY, params(nn), dp) + optimization_step!(params(nn), optimizer, dp) end ``` diff --git a/docs/src/tutorials/hamiltonian_neural_network.md b/docs/src/tutorials/hamiltonian_neural_network.md index 4fe0b0914..5b2c8c42c 100644 --- a/docs/src/tutorials/hamiltonian_neural_network.md +++ b/docs/src/tutorials/hamiltonian_neural_network.md @@ -42,7 +42,7 @@ We can now train the network ```@example hnn batch = Batch(10) n_epochs = 100 -o = Optimizer(Adam(Float64), hnn) +o = Optimizer(Adam(), hnn) loss_array = o(hnn, dl, batch, n_epochs, loss) using CairoMakie # hide lines(loss_array) # hide @@ -84,7 +84,7 @@ nothing # hide We can now train the network: ```@example hnn -o_pairs = Optimizer(Adam(Float64), hnn_pairs) +o_pairs = Optimizer(Adam(), hnn_pairs) loss_array_pairs = o_pairs(hnn_pairs, dl_pairs, batch, n_epochs, loss_pairs) lines(loss_array_pairs) # hide ``` diff --git a/docs/src/tutorials/linear_symplectic_transformer.md b/docs/src/tutorials/linear_symplectic_transformer.md index 7d38af6d6..81010f9f8 100644 --- a/docs/src/tutorials/linear_symplectic_transformer.md +++ b/docs/src/tutorials/linear_symplectic_transformer.md @@ -63,7 +63,7 @@ nn_standard = NeuralNetwork(arch_standard) nn_symplectic = NeuralNetwork(arch_symplectic) nn_sympnet = NeuralNetwork(arch_sympnet) -o_pairing = AdamOptimizerWithDecay(n_epochs, Float64) +o_pairing = AdamOptimizerWithDecay(n_epochs) o_standard = Optimizer(nn_standard; o_pairing...) o_symplectic = Optimizer(nn_symplectic; o_pairing...) diff --git a/docs/src/tutorials/symplectic_autoencoder.md b/docs/src/tutorials/symplectic_autoencoder.md index ab958083c..3a206737e 100644 --- a/docs/src/tutorials/symplectic_autoencoder.md +++ b/docs/src/tutorials/symplectic_autoencoder.md @@ -145,7 +145,7 @@ dl = DataLoader(dl_cpu, backend, Float32) sae_nn_gpu = NeuralNetwork(sae_arch, CUDADevice(), Float32) -o = Optimizer(sae_nn_gpu; AdamOptimizerWithDecay(integrator_train_epochs, Float32)...) +o = Optimizer(sae_nn_gpu; AdamOptimizerWithDecay(integrator_train_epochs)...) # train the network o(sae_nn_gpu, dl, Batch(batch_size), n_epochs) diff --git a/docs/src/tutorials/symplectic_transformer.md b/docs/src/tutorials/symplectic_transformer.md index 75fcf178a..8d1331747 100644 --- a/docs/src/tutorials/symplectic_transformer.md +++ b/docs/src/tutorials/symplectic_transformer.md @@ -63,7 +63,7 @@ nn_standard = NeuralNetwork(arch_standard) nn_symplectic = NeuralNetwork(arch_symplectic) nn_sympnet = NeuralNetwork(arch_sympnet) -o_pairing = AdamOptimizerWithDecay(n_epochs, Float64) +o_pairing = AdamOptimizerWithDecay(n_epochs) o_standard = Optimizer(nn_standard; o_pairing...) o_symplectic = Optimizer(nn_symplectic; o_pairing...) diff --git a/docs/src/tutorials/sympnet_tutorial.md b/docs/src/tutorials/sympnet_tutorial.md index 602f58a2c..7d71eecae 100644 --- a/docs/src/tutorials/sympnet_tutorial.md +++ b/docs/src/tutorials/sympnet_tutorial.md @@ -100,7 +100,7 @@ We have to define an [optimizer](@extref GeometricOptimizers Standard-Neural-Net ```@example sympnet # set up optimizer; for this we first need to specify the optimization method -opt_method = Adam(type) +opt_method = Adam() # we then call the optimizer struct which allocates the cache la_opt = Optimizer(opt_method, la_nn) g_opt = Optimizer(opt_method, g_nn) diff --git a/docs/src/tutorials/volume_preserving_attention.md b/docs/src/tutorials/volume_preserving_attention.md index a4474faec..43abfb902 100644 --- a/docs/src/tutorials/volume_preserving_attention.md +++ b/docs/src/tutorials/volume_preserving_attention.md @@ -117,9 +117,9 @@ Next we train the networks (here we pick a batch size of 30 and train for 1000 e ```@setup volume_preserving_attention function set_up_optimizers(nn_skew, nn_arb, nn_comp) - o_skew = Optimizer(Adam(T), nn_skew) - o_arb = Optimizer(Adam(T), nn_arb) - o_comp = Optimizer(Adam(T), nn_comp) + o_skew = Optimizer(Adam(), nn_skew) + o_arb = Optimizer(Adam(), nn_arb) + o_comp = Optimizer(Adam(), nn_comp) o_skew, o_arb, o_comp end diff --git a/docs/src/tutorials/volume_preserving_transformer_rigid_body.md b/docs/src/tutorials/volume_preserving_transformer_rigid_body.md index 487c659c7..2d60ee879 100644 --- a/docs/src/tutorials/volume_preserving_transformer_rigid_body.md +++ b/docs/src/tutorials/volume_preserving_transformer_rigid_body.md @@ -199,7 +199,7 @@ const n_epochs = 500000 const batch_size = 16384 const feedforward_batch = Batch(batch_size) const transformer_batch = Batch(batch_size, seq_length, seq_length) -const opt_pairing = AdamOptimizerWithDecay(n_epochs, T; η₁ = 1e-2, η₂ = 1e-6) +const opt_pairing = AdamOptimizerWithDecay(n_epochs; η₁ = 1e-2, η₂ = 1e-6) o_vpff = Optimizer(nn_vpff; opt_pairing...) o_vpt = Optimizer(nn_vpt; opt_pairing...) diff --git a/scripts/reproduction/symplectic_autoencoders/training.jl b/scripts/reproduction/symplectic_autoencoders/training.jl index 389369520..064a0ca57 100644 --- a/scripts/reproduction/symplectic_autoencoders/training.jl +++ b/scripts/reproduction/symplectic_autoencoders/training.jl @@ -31,7 +31,7 @@ n_range = 2:1:smoke_size(15, 3) # The learning rate is the `Optimizer`'s `step_size`, not the method's: `Adam` lost the `η` field # it never applied to the direction, and `β₁`, `β₂` and `δ` became keywords so that the old # positional call fails instead of silently reading `η` as `β₁`. -opt = AdamOptimizer(T; β₁ = T(0.9), β₂ = T(0.99), δ = T(1.0e-8)) +opt = AdamOptimizer(; β₁ = T(0.9), β₂ = T(0.99), δ = T(1.0e-8)) const step_size = T(0.001) # The retraction belongs to the `Optimizer` now, not to the layer: `PSDLayer(M, N)` takes no # keyword. diff --git a/scripts/reproduction/sympnets/sympnet_pendulum.jl b/scripts/reproduction/sympnets/sympnet_pendulum.jl index 5fe8fa0f7..abc7cec98 100644 --- a/scripts/reproduction/sympnets/sympnet_pendulum.jl +++ b/scripts/reproduction/sympnets/sympnet_pendulum.jl @@ -70,14 +70,13 @@ method = AdamOptimizer() # initial gradients for calling Cache constructor opt = Optimizer(method, ps) # `optimization_step!` takes the global section of the parameters, not the model. -λY = GlobalSection(ps) # training println("initial loss: ", full_loss(ps, q, p)) training_steps = smoke_size(1000, 2) @showprogress for i in 1:training_steps dp = Zygote.gradient(ps -> loss(ps, q, p), ps)[1] - optimization_step!(opt, λY, ps, dp) + optimization_step!(ps, opt, dp) end println("final loss: ", full_loss(ps, q, p)) diff --git a/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl b/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl index 099bc56eb..3cb904625 100644 --- a/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl +++ b/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl @@ -42,7 +42,7 @@ const T = eltype(dl) # hyperparameters concerning training const n_epochs = smoke_size(1000, 2) const batch_size = 1024 -const opt_method = AdamOptimizer(T) +const opt_method = AdamOptimizer() const resnet_activation = tanh const t_validation = 30 diff --git a/scripts/reproduction/volume_preserving_transformer/rigid_body.jl b/scripts/reproduction/volume_preserving_transformer/rigid_body.jl index 47708f121..f84d93496 100644 --- a/scripts/reproduction/volume_preserving_transformer/rigid_body.jl +++ b/scripts/reproduction/volume_preserving_transformer/rigid_body.jl @@ -51,7 +51,7 @@ const dl = backend == CPU() ? DataLoader(T.(dl₁.input)) : # hyperparameters concerning training const n_epochs = smoke_size(500000, 2) const batch_size = 16384 -const opt_pairing = AdamOptimizerWithDecay(n_epochs, T; η₁ = 1e-2, η₂ = 1e-6) +const opt_pairing = AdamOptimizerWithDecay(n_epochs; η₁ = 1e-2, η₂ = 1e-6) # parameters for evaluation ics_val = [sin(1.1), 0.0, cos(1.1)] diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 000498b31..290c5d1d2 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -6,12 +6,11 @@ using AbstractNeuralNetworks # package also exports `flatten`/`unflatten` and the leaf protocol, none of which this package # extends — `GeometricOptimizers` carries the protocol for the structured matrices. # -# Two things do come from there rather than being written again here: `mapstorage`, which reaches the -# storage of a structured leaf and rebuilds the leaf around the result (`src/map_to_cpu.jl`), and -# `parameter_eltype`, which promotes over the leaves of a set. A plain `NamedTuple` of layers is -# walked with `Base.map`, which needs nothing from anybody. +# One thing does come from there rather than being written again here: `mapstorage`, which reaches +# the storage of a structured leaf and rebuilds the leaf around the result (`src/map_to_cpu.jl`). +# A plain `NamedTuple` of layers is walked with `Base.map`, which needs nothing from anybody. import NeuralNetworkParameters: NetworkParameters -using NeuralNetworkParameters: mapparameters, mapstorage, parameter_eltype +using NeuralNetworkParameters: mapstorage using ChainRulesCore using GeometricBase using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries @@ -55,7 +54,13 @@ import GeometricOptimizers: Geodesic, Cayley, geodesic, cayley, retraction import GeometricOptimizers: OptimizerMethod, GradientMethod, MomentumMethod, Adam, GradientState, MomentumState, AdamState, - AdamOptimizerWithDecay, DecayingStatic + AdamOptimizerWithDecay, DecayingStatic, + ScalarMomentAdam, ScalarMomentAdamState, + CompositeMethod, CompositeState +# The training step is `GeometricOptimizers`' own: `Optimizer` holds a `TrainingOptimizer`, and +# `optimization_step!` gets a method for `Optimizer` rather than a second function of the name, so a +# `Main` that loads both packages has one `optimization_step!` and not two. +import GeometricOptimizers: TrainingOptimizer, optimization_step! import GeometricOptimizers: update! # `solve!` is imported rather than started afresh so that GML's `solve!(::NeuralNetwork{<:PSDArch}, # …)` — solve for the parameters directly, by SVD, instead of training for them — is a method of the @@ -201,6 +206,8 @@ export OptimizerMethod export GradientMethod, GradientState export MomentumMethod, MomentumState export Adam, AdamState +export ScalarMomentAdam, ScalarMomentAdamState +export CompositeMethod, CompositeState export Optimizer export optimization_step! export GlobalSection, global_section, apply_section, apply_section!, update_section! diff --git a/src/data_loader/optimize.jl b/src/data_loader/optimize.jl index 349448b9a..7cfe94fae 100644 --- a/src/data_loader/optimize.jl +++ b/src/data_loader/optimize.jl @@ -1,5 +1,5 @@ @doc raw""" - optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY) + optimize_for_one_epoch!(opt, model, ps, dl, batch, loss) Sample the data contained in `dl` according to `batch` and optimize for these batches. @@ -24,7 +24,6 @@ All the arguments are mandatory (there are no defaults): 4. the data (i.e. an instance of [`DataLoader`](@ref)). 5. `batch`::[`Batch`](@ref): stores `batch_size` (and optionally `seq_length` and `prediction_window`). 6. `loss::NetworkLoss`. -7. the *section* `λY` of the parameters `ps`. # Implementation @@ -51,9 +50,8 @@ function optimize_for_one_epoch!(opt::Optimizer, model, ps::NetworkParameters, dl::DataLoader{T}, batch::Batch, - loss::NetworkLoss, - λY) where {T} - optimize_for_one_epoch!(opt, model, ps, dl, batch, ZygotePullback(loss), λY) + loss::NetworkLoss) where {T} + optimize_for_one_epoch!(opt, model, ps, dl, batch, ZygotePullback(loss)) end function optimize_for_one_epoch!(opt::Optimizer, @@ -61,8 +59,7 @@ function optimize_for_one_epoch!(opt::Optimizer, ps::NetworkParameters, dl::DataLoader{T}, batch::Batch, - _pullback::AbstractPullback, - λY) where {T} + _pullback::AbstractPullback) where {T} count = 0 total_error = T(0) batches = batch(dl) @@ -72,7 +69,7 @@ function optimize_for_one_epoch!(opt::Optimizer, loss_value, pullback = _pullback(ps, model, input_nt_output_nt) total_error += loss_value dp = _processing(pullback(one(loss_value))) - optimization_step!(opt, λY, ps, dp) + optimization_step!(ps, opt, dp) end total_error / count end @@ -83,13 +80,12 @@ function (o::Optimizer)(nn::NeuralNetwork, n_epochs::Integer, loss::NetworkLoss, _pullback::AbstractPullback = ZygotePullback(loss); show_progress = true) - Λ = GlobalSection(params(nn)) progress_object = show_progress ? ProgressMeter.Progress(n_epochs; enabled = true) : nothing loss_array = zeros(float(eltype(dl)), n_epochs) for i in 1:n_epochs loss_array[i] = optimize_for_one_epoch!( - o, nn.model, params(nn), dl, batch, _pullback, Λ) + o, nn.model, params(nn), dl, batch, _pullback) show_progress ? ProgressMeter.next!(progress_object; showvalues = [( :TrainingLoss, loss_array[i])]) : nothing diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl index 69d81f32a..886249676 100644 --- a/src/optimizers/optimizer.jl +++ b/src/optimizers/optimizer.jl @@ -1,423 +1,88 @@ -# Optimizer machinery on top of GeometricOptimizers. -# Kept out of `utils.jl` because it dispatches on `Manifold`, which is defined later. - -# Gradient wrapper: stores a pre-computed Euclidean gradient and applies rgrad on manifolds. -mutable struct _GMLGradient{T, VT} <: GeometricOptimizers.Gradient{T} - dp::VT -end - -_gml_rgrad(x::Manifold, dp) = rgrad(x, dp) -_gml_rgrad(x, dp) = dp -# `mapparameters` and not `map`: it recurses on the branches, so `_gml_rgrad` is only ever called on -# leaves, and it rebuilds in the shape of its *first* argument -- a container, which is what -# `GeometricOptimizers._copyto!(gradient_array(cache), ·)` has a method for. It also normalises its -# trailing arguments, so `dp` may stay the plain `NamedTuple` the pullback produced. -_gml_rgrad(x::NetworkParameters, dp) = mapparameters(_gml_rgrad, x, dp) - -(g::_GMLGradient{T})(x::NetworkParameters{T}) where {T} = _gml_rgrad(x, g.dp) -(g::_GMLGradient{T})(x::AbstractArray{T}) where {T} = g.dp - -# `Manifold` and not `AbstractArray` alone, because the two are otherwise **ambiguous** and not merely -# overlapping. `_GMLGradient{T}` is a `GeometricOptimizers.Gradient{T}`, `Manifold{T}` is an -# `AbstractMatrix{T}`, and `GeometricOptimizers` has `(::Gradient{T})(::Manifold{T})`; so on a bare -# manifold neither `(::_GMLGradient{T})(::AbstractArray{T})` above nor that one is more specific, and -# the call is a run-time `MethodError: ... is ambiguous`. This method is more specific than both. +# The optimizer of a training loop is `GeometricOptimizers.TrainingOptimizer`: one cache and one +# state over the whole parameter set, or one per leaf for a `CompositeMethod`, stepped with the +# gradient of one minibatch at a fixed or scheduled step size. What this package adds is the +# `Optimizer` that pairs one with a network, and the epoch loop in `data_loader/optimize.jl` that +# calls it. # -# The answer it gives is this package's: the gradient is already computed, so the projection is -# `rgrad(x, dp)` and not upstream's `rgrad(x, reshape(grad(vec(x)), ...))`, which would evaluate an -# inner gradient this functor does not have. `_gml_rgrad` is where that lives. -# -# `_make_optimizer_cache` below asks `_is_layer` first and so keeps a bare `Manifold` from arriving -# here at all. That ordering is a correctness fix in its own right -- one cache per layer rather than -# per weight -- but it must not be the *only* thing standing between a caller and an ambiguity. -(g::_GMLGradient{T})(x::Manifold{T}) where {T} = _gml_rgrad(x, g.dp) - -# A bare `AbstractVector`, for the same reason as the `Manifold` method above: `SimpleSolvers` has -# `(::Gradient{T})(::AbstractVector{T})`, `_GMLGradient{T} <: GeometricOptimizers.Gradient{T}`, and -# without a method of its own `(::_GMLGradient{T})(::AbstractArray{T})` above and that one are -# equally specific on a `Vector` -- the call is a run-time `MethodError: ... is ambiguous`. This -# method is more specific than both: `_GMLGradient{T} <: Gradient{T}` on the first argument and -# `AbstractVector{T} <: AbstractArray{T}` on the second. The answer is unchanged from the -# `AbstractArray{T}` method above -- the gradient is already Euclidean, so it is `g.dp`. -(g::_GMLGradient{T})(x::AbstractVector{T}) where {T} = g.dp +# There is no step here, no cache, no state and no list of methods. Each of those used to be written +# here a second time -- a per-layer walk, a Euclidean Adam of its own, a chain of `isa` tests over +# the state types -- and each could disagree with `GeometricOptimizers` without a test noticing. -# State for Euclidean (non-manifold) parameters. -mutable struct GMLEuclideanState{T, AT <: AbstractArray{T}} - iterations::Int - m₁::AT - m₂::AT -end -function GMLEuclideanState(x::AbstractArray{T}) where {T} - GMLEuclideanState{T, typeof(x)}(0, zero(x), zero(x)) -end - -# Adam with a decaying learning rate is GeometricOptimizers', and split the way it belongs: the -# direction is an `Adam` method, the schedule is a `DecayingStatic` line search. Both names are -# imported, and `Optimizer` below takes a `DecayingStatic` as its `step_size`. - -_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 - -function _adapt_method_to_T(method::GeometricOptimizers.Adam, ::Type{T}) where {T} - GeometricOptimizers.Adam(T; β₁ = T(method.β₁), β₂ = T(method.β₂), δ = T(method.δ)) -end -function _adapt_method_to_T(method::GeometricOptimizers.MomentumMethod, ::Type{T}) where {T} - GeometricOptimizers.MomentumMethod(T(method.α)) -end -_adapt_method_to_T(method, ::Type) = method - -# Whether one `GeometricOptimizers` cache covers `x` whole. -# -# A **layer** is a flat `NamedTuple` of arrays, and it is exactly what one cache is for. A `NamedTuple` -# whose values are branches is a *subtree* and gets one cache per layer instead, which is why the test -# is flatness and not merely "a `NamedTuple`". -# -# The rule lives here rather than upstream because a whole set of parameters reaches -# `GeometricOptimizers` only as a `NetworkParameters`, and a container says nothing about how deep the -# tree beneath it is. Splitting a network into layers is this package's decision, so this package makes -# it; [`_as_go_solution`](@ref) then does the wrap. -# -# A layer whose weights do not share an element type is still one cache, with `T` the promotion over -# them. `_adapt_method_to_T` reads that promotion, so the cache is built for the element type the -# layer actually has rather than for one its weights are required to agree on. -# `!isempty` and not `all` alone: `all` over an empty collection is vacuously true, so without it an -# empty set would be one cache's worth -- and `OptimizerCache` would then be asked for the element -# type of a set that has no leaves to promote. An empty set is zero caches' worth, so it descends and -# yields an empty tree of caches, which is what the step below iterates over zero times. -_is_layer(x::NamedTuple) = _all_leaves(values(x)) -_is_layer(x::NetworkParameters) = _all_leaves(values(x)) -_is_layer(_) = false - -_all_leaves(vs) = !isempty(vs) && all(v -> v isa AbstractArray, vs) - -function _use_go_cache(method, x) - _is_go_native_method(method) && - (x isa GeometricOptimizers.OptimizerSolution || _is_layer(x)) -end - -""" - _as_go_solution(x) - -`x` in the shape `GeometricOptimizers` takes a solution in. - -A layer given as a bare `NamedTuple` is wrapped; anything already in the right shape — a container, a -bare `Manifold`, an `AbstractArray` — is passed through. The wrap **shares the leaf arrays**, so an -in-place optimizer step writes through to the network's own weights and nothing has to be copied back. """ -_as_go_solution(x::NetworkParameters) = x -_as_go_solution(x) = _is_layer(x) ? NetworkParameters(x) : x + Optimizer(method, nn; retraction = Cayley(), step_size = default_step_size(method)) + Optimizer(nn; algorithm, linesearch = default_step_size(algorithm), retraction = Cayley()) -# A `NetworkParameters` is always a tree of layers to descend into, never a single -# `GeometricOptimizers` leaf, so its branch comes *first* — ahead of `_use_go_cache`. -# -# That ordering is **load-bearing**. `NetworkParameters` *is* one of the types -# `GeometricOptimizers.OptimizerSolution` unions, so `_use_go_cache` is true at the root: without the -# test above, a whole network would get one cache instead of one per layer -- silently, and with -# `_leaf_optim_step!` handed the entire tree. Asking the structural question before the capability -# question is what makes the shape of the cache depend on the shape of the parameters rather than on -# which types upstream happens to accept. -# -# One cache for the whole network is the better end state, and it is **not** what this does. -# `_GMLGradient` does take a `NetworkParameters`, but that is one wrapped *layer* and not the tree, -# because the container branch above splits the network first. Getting to one cache means dropping -# `_tree_optim_step!` and handing the whole tree over, which is a change of behaviour -- one -# `GlobalSection` tree and one `Q` across every layer -- so it wants its own release. Nothing here -# depends on it. -# -# **"Is this one cache's worth?" is asked before "is this a tree?"**, and a *flat* set answers yes to -# both. `_is_layer` is that question: a set whose values are all leaves is a layer, whether it arrives -# wrapped or as a branch, and one `GeometricOptimizers` cache is exactly what a layer is for. Asking -# the tree question first would descend into such a set and give every individual weight its own cache -# -- which for a manifold weight is not merely wasteful but wrong, since a bare `Manifold` then reaches -# the gradient functor where a whole layer should have. -# -# The tree branch comes second and covers both carriers, because a network is a tree of layers whether -# it is wrapped or not. -function _make_optimizer_cache(method, x) - if _is_go_native_method(method) && _is_layer(x) - GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, parameter_eltype(x)), - _as_go_solution(x)) - elseif x isa NetworkParameters || x isa NamedTuple - NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x))) - elseif _use_go_cache(method, x) - GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, parameter_eltype(x)), x) - else - GMLEuclideanState(x) - end -end - -function _make_optimizer_state(method, x) - if _is_go_native_method(method) && _is_layer(x) - GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, parameter_eltype(x)), - _as_go_solution(x)) - elseif x isa NetworkParameters || x isa NamedTuple - NamedTuple{keys(x)}(Tuple(_make_optimizer_state(method, x[k]) for k in keys(x))) - elseif _use_go_cache(method, x) - GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, parameter_eltype(x)), x) - else - GMLEuclideanState(x) - end -end +The optimizer of a training loop for the parameters of `nn`, a `NeuralNetwork` or a +`NetworkParameters`. -""" - Optimizer(method, nn; retraction, step_size) - Optimizer(nn; algorithm, linesearch, retraction) +It holds one `GeometricOptimizers.TrainingOptimizer` over the whole parameter set, which +[`optimization_step!`](@ref) steps once per minibatch. `method` is a first-order method: +`GradientMethod()`, `MomentumMethod()`, `Adam()`, `ScalarMomentAdam()` or another member of +GeometricOptimizers' `AdamFamily`, or a `CompositeMethod`, which gives every leaf of the set the +method it selects for that leaf: -Optimizer state combining a `GeometricOptimizers` method with the parameters of a neural network. +```julia +method = CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()) +opt = Optimizer(method, nn) +``` -`step_size` is either a number — a fixed learning rate — or a -`GeometricOptimizers.DecayingStatic`, a learning rate that decays geometrically with the iteration -number. The second form is the one `GeometricOptimizers` uses itself, so a method paired with a -schedule splats straight in: +`step_size` is a number — a fixed learning rate — or a `DecayingStatic`, a learning rate that decays +geometrically with the iteration number. The keyword form takes the `(algorithm, linesearch)` pair +that `AdamOptimizerWithDecay` returns, so it splats straight in: ```julia -opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...) +opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) ``` +`retraction` is `Cayley()` or `Geodesic()`. + # Extended help The step size is a property of the optimizer and not of the method: the same `Adam()` trains at any -learning rate. That is the split `GeometricOptimizers` makes — the method supplies a direction, a -`SimpleSolvers.LinesearchMethod` supplies how far to go along it — and `step_size` is GML's half of -it for a training loop, which has no objective function for a real line search to evaluate. +learning rate. That is the split `GeometricOptimizers` makes, where the method supplies a direction +and the step size says how far to go along it. A training loop has no objective for a line search to +evaluate, only the loss of one minibatch, so a step size is a number or a schedule and nothing else. """ -mutable struct Optimizer{MT <: GeometricOptimizers.OptimizerMethod, CT, ST, RT, SST} - method::MT - cache::CT - state::ST - retraction::RT - step_size::SST - iterations::Int +struct Optimizer{TO <: GeometricOptimizers.TrainingOptimizer} + training::TO end -_default_step_size(::GeometricOptimizers.Adam) = 1e-3 -_default_step_size(::GeometricOptimizers.OptimizerMethod) = 1e-2 - -_step_size(η::Real, ::Int) = Float64(η) -# `t` and not `t - 1`: `optimization_step!` increments before it asks, so the first step of a solve -# is `α(1) = γη₁`. That is what `DecayingStatic` means by iteration `t`: `solve!` calls -# `increase_iteration_number!` before `solver_step!`. -_step_size(ls::DecayingStatic, t::Int) = Float64(GeometricOptimizers.step_size(ls, t)) - -_current_step_size(opt::Optimizer, t::Int) = _step_size(opt.step_size, t) - -_optimizer_step_size(η::Real) = Float64(η) -# Anything that is not a plain number goes through the same funnel as the `linesearch` keyword below, -# so that the two entry points accept the same things: `step_size = Static(α)` is a fixed learning -# rate on both, and anything else reports the `ArgumentError` that explains why a real line search -# has nothing to search along here, instead of a `MethodError` naming this helper. -_optimizer_step_size(ls) = _step_size_from_linesearch(ls) - -function Optimizer(method::GeometricOptimizers.OptimizerMethod, nn::NeuralNetwork; - retraction = GeometricOptimizers.cayley, - step_size = _default_step_size(method)) - Optimizer(method, params(nn); retraction = retraction, step_size = step_size) +function Optimizer(method::OptimizerMethod, nn::NeuralNetwork; kwargs...) + Optimizer(method, params(nn); kwargs...) end -function Optimizer(method::GeometricOptimizers.OptimizerMethod, - ps::NetworkParameters; - retraction = GeometricOptimizers.cayley, - step_size = _default_step_size(method)) - Optimizer(method, _make_optimizer_cache(method, ps), _make_optimizer_state(method, ps), - retraction, _optimizer_step_size(step_size), 0) +function Optimizer(method::OptimizerMethod, ps::NetworkParameters; + retraction::GeometricOptimizers.AbstractRetraction = Cayley(), + step_size = GeometricOptimizers.default_step_size(method)) + Optimizer(TrainingOptimizer(ps; algorithm = method, linesearch = step_size, + retraction = retraction)) end -# The keyword form, so that the `(algorithm, linesearch)` pairing `GeometricOptimizers` returns from -# `AdamOptimizerWithDecay` splats in unchanged. `linesearch` is the step size under the name -# upstream gives it; a `Static` carries its own `α`, which is then the fixed learning rate. function Optimizer(nn_or_ps::Union{NeuralNetwork, NetworkParameters}; - algorithm::GeometricOptimizers.OptimizerMethod, - linesearch = nothing, - retraction = GeometricOptimizers.cayley, - step_size = linesearch === nothing ? _default_step_size(algorithm) : - _step_size_from_linesearch(linesearch)) - Optimizer(algorithm, nn_or_ps; retraction = retraction, step_size = step_size) -end - -_step_size_from_linesearch(ls::DecayingStatic) = ls -_step_size_from_linesearch(ls::GeometricOptimizers.Static) = Float64(ls.α) -function _step_size_from_linesearch(ls) - throw(ArgumentError( - "`Optimizer` takes a fixed step size or a `DecayingStatic` schedule, not a $(typeof(ls)). " * - "A training loop evaluates its loss on one batch at a time and has no objective for a line " * - "search to search along; use `GeometricOptimizers.Optimizer` with an `OptimizerProblem` for " * - "that.")) -end - -# 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} - # `p ← αp + ∇L`, the classic momentum recursion. The decay belongs on `p` and not on `∇L`: - # `p ← p + α∇L` is an undamped accumulator that grows without bound for a constant gradient - # instead of saturating at `∇L/(1 - α)`. Same recursion as GO's `update!(::MomentumState, ...)`. - state.m₁ .= T(method.α) .* state.m₁ .+ dx - x .-= T(step_size) .* state.m₁ -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.δ) - # the first factor is `(β - β^t)/(1 - β^t)`, not `β/(1 - β^t)`: the latter is ~100x too large - # at t = 2 and compounds every step, which inflates the second moment until the update vanishes - fac₁₁ = (β₁-β₁^_t)/(1-β₁^_t) - fac₁₂ = (1-β₁)/(1-β₁^_t) - fac₂₁ = (β₂-β₂^_t)/(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 -# `_euclidean_update!` has three methods and not four: Adam with a decaying learning rate *is* -# Adam — only `step_size` differs, and that comes in as an argument — so the `Adam` method above -# serves it too. - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.Adam, ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) -end - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.MomentumMethod, ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) -end - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.OptimizerMethod, ps_leaf) - T = parameter_eltype(ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, - GeometricOptimizers.NoHessian{T}(), ps_leaf) -end - -# GO-managed leaf step (manifolds, vectors, whole layers) -# -# `ps` and not `ps_leaf` from here on: a layer is wrapped for the same reason its cache is, and the -# wrap shares the arrays, so `_copyto!(ps, solution(cache))` below writes into the weights the caller -# holds. A bare `Manifold` or `AbstractArray` passes through untouched. -function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, - state::GeometricOptimizers.OptimizerState, - dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) - T = parameter_eltype(ps_leaf) - ps = _as_go_solution(ps_leaf) - local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf) - adapted = _adapt_method_to_T(method, T) - state.iterations += 1 - _go_update_leaf!(cache, state, local_grad, adapted, ps) - GeometricOptimizers._rmul!(GeometricOptimizers.direction(cache), T(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, GeometricOptimizers.solution(cache)) - GeometricOptimizers._copyto!(λY_leaf, GeometricOptimizers.section(cache)) - # `section(cache)` is `update_section!(section(state), direction, retraction)`, so copying it is - # the same thing as retracting a second time -- and a retraction on a manifold is `O(N³)` where - # the copy is `O(N²)`. - GeometricOptimizers._copyto!(GeometricOptimizers.section(state), - GeometricOptimizers.section(cache)) - 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 - # `p ← αp + ∇L`; see the note in `_euclidean_update!` for the momentum method. This has to - # match what `update!(::MomentumCache, ...)` anticipated when it formed the direction. - GeometricOptimizers._rmul!(GeometricOptimizers.momentum(state), adapted.α) - GeometricOptimizers._add!(GeometricOptimizers.momentum(state), - GeometricOptimizers.gradient_array(cache)) - end - nothing -end - -# 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 - -# Whether `λY` is a tree of sections to descend into, or a single `GlobalSection` standing in for a -# whole subtree. Named rather than written as an `isa` union: it is one question, asked once, and the -# answer decides whether a layer gets its own section or the whole tree. A section tree is a plain -# `NamedTuple` today, because `GeometricOptimizers`' `GlobalSection(::NetworkParameters)` unwraps the -# container; the container arm is here so that a change there is a `MethodError` and not a silent -# hand-off of the whole tree to every layer. -_is_section_tree(::NamedTuple) = true -_is_section_tree(::NetworkParameters) = true -_is_section_tree(_) = false - -# Recursive dispatcher over the *cache* tree. -# -# Deliberately hand-written rather than `NeuralNetworkParameters.foreachparameters`, which walks a -# parameter tree. Two reasons, and both are load-bearing: -# -# - The recursion is keyed on `caches`, and stops where the cache stops. A cache sits at the *layer* -# level, so a layer's `NamedTuple` of weights arrives at `_leaf_optim_step!` whole, which is what -# one `GeometricOptimizers` cache is for. `foreachparameters` recurses on the leaf protocol and -# would descend past the layer into the individual weights, re-pairing every cache with the wrong -# object. -# - `λY` is broadcast, not zipped: a single `GlobalSection` may stand in for a whole subtree, which -# the ternary below expresses. `foreachparameters` has no such rule — it takes `values` of each -# trailing argument, so a bare `GlobalSection` beside a `NamedTuple` of caches is a `MethodError`. -# -# The `nothing` skip is the one thing the two have in common, and it is one line here. -# -# The `λY` test names both container types for the same reason `_make_optimizer_cache` asks the -# structural question first: a section tree is something to descend into, whichever type carries it. -# Only a `NamedTuple` ever arrives: `GeometricOptimizers`' own -# `GlobalSection(ps::NetworkParameters) = GlobalSection(params(ps))` unwraps the container, so a -# section tree is a plain `NamedTuple` whether the parameters are wrapped or not. Should that ever -# return a *container* of sections instead, `isa NamedTuple` alone would be false here and every layer -# would be handed the whole tree instead of its own section. -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 = _is_section_tree(λY) ? λY[k] : λY - _tree_optim_step!(caches[k], states[k], dp_k, ps[k], λY_k, - method, retraction, step_size) - end - else - _leaf_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) - end - nothing + algorithm::OptimizerMethod, + linesearch = GeometricOptimizers.default_step_size(algorithm), + retraction::GeometricOptimizers.AbstractRetraction = Cayley()) + Optimizer(algorithm, nn_or_ps; retraction = retraction, step_size = linesearch) end """ - optimization_step!(opt, λY, ps, dp) + optimization_step!(ps, opt::Optimizer, dp) -Apply one optimization step to the parameters `ps` and their gradient `dp`, with the method and step -size the [`Optimizer`](@ref) `opt` carries. +Take one step of the [`Optimizer`](@ref) `opt` with the gradient `dp` of the loss at the parameters +`ps`, and write the new parameters into `ps`. -`λY` is a `GlobalSection` of `ps` (or a `NamedTuple` of them). Note that it is an *output* -here: the section the optimizer carries from step to step lives in `opt.state`, and `λY` is written -so that callers who inspect it see the updated section. It therefore has to be allocated once and -reused, not rebuilt per step -- rebuilding it costs a QR decomposition per manifold weight. - -The step counter is incremented *before* the step size is read, so the first step of a run is step 1. -This matters for a decaying `step_size` and is how `GeometricOptimizers` counts too. +`ps` is the parameter set `opt` was built for. `dp` is the gradient a pullback returns, a +`NetworkParameters` or the bare `NamedTuple` of the same shape. This is +`GeometricOptimizers.optimization_step!` on the `TrainingOptimizer` that `opt` holds; see there for +what a step does and what it refuses. """ -function optimization_step!(opt::Optimizer, λY, ps, dp) - # The increment comes *first*, so the first step of a run is step 1. It matters only for a - # decaying `step_size`, and there it matters: reading the schedule before incrementing takes - # `α(0) = η₁`, one whole step above what `DecayingStatic` and `GeometricOptimizers.solve!` - # take. `solve!` counts the same way, by calling `increase_iteration_number!` before - # `solver_step!`. - opt.iterations += 1 - step = _current_step_size(opt, opt.iterations) - _tree_optim_step!(opt.cache, opt.state, dp, ps, λY, opt.method, opt.retraction, step) - nothing +function optimization_step!(ps::NetworkParameters, opt::Optimizer, dp) + optimization_step!(ps, opt.training, _as_parameter_set(dp)) + ps end +# A pullback seeded with a bare `NamedTuple` hands its gradient back as one; see `_get_contents`. +_as_parameter_set(dp::NetworkParameters) = dp +_as_parameter_set(dp::NamedTuple) = NetworkParameters(dp) + check(::Optimizer) = nothing diff --git a/test/data_loader/data_loader_optimization_step.jl b/test/data_loader/data_loader_optimization_step.jl index 9a9405bd0..65029a085 100644 --- a/test/data_loader/data_loader_optimization_step.jl +++ b/test/data_loader/data_loader_optimization_step.jl @@ -17,8 +17,7 @@ function test_data_loader(sys_dim, n_time_steps, n_params, T = Float32) dx = Zygote.gradient(ps -> loss(model, ps, dl.input, dl.input), ps)[1] ps_copy = deepcopy(ps) o = Optimizer(GradientOptimizer(), ps) - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) @test ps !== ps_copy end diff --git a/test/data_loader/optimizer_functor_with_adam.jl b/test/data_loader/optimizer_functor_with_adam.jl index ad449788e..b079c1609 100644 --- a/test/data_loader/optimizer_functor_with_adam.jl +++ b/test/data_loader/optimizer_functor_with_adam.jl @@ -50,8 +50,7 @@ function test_optimization_with_adam(; T = Float32, input_dim = 9, n_patches = 4 loss₁ = loss(model, ps, dl.input, dl.output) opt = Optimizer(Adam(), nn_obj) - λY = GlobalSection(ps) - loss_average = optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY) + loss_average = optimize_for_one_epoch!(opt, model, ps, dl, batch, loss) loss₃ = loss(model, ps, dl.input, dl.output) diff --git a/test/layers/gradient_layer_tests.jl b/test/layers/gradient_layer_tests.jl index f0eb7474a..622e6dc1c 100644 --- a/test/layers/gradient_layer_tests.jl +++ b/test/layers/gradient_layer_tests.jl @@ -24,19 +24,18 @@ end function test_gradient_layer_derivative_and_update(T, M, N, batch_size = 10) dummy_model = Chain(GradientLayerP(M, N, tanh), GradientLayerQ(M, N, tanh)) ps = params(NeuralNetwork(dummy_model, CPU(), T)) - o = Optimizer(Adam(T; β₁ = T(0.9), β₂ = T(0.999), δ = T(3e-7)), ps; + o = Optimizer(Adam(; β₁ = T(0.9), β₂ = T(0.999), δ = T(3e-7)), ps; step_size = T(0.1)) # test for vector x = rand(T, M) gs = Zygote.gradient(ps -> sum(dummy_model(x, ps)), ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, gs) + optimization_step!(ps, o, gs) # test for matrix X = rand(T, M, batch_size) gs = Zygote.gradient(ps -> sum(dummy_model(X, ps)), ps)[1] - optimization_step!(o, λY, ps, gs) + optimization_step!(ps, o, gs) end types = (Float32, Float64) diff --git a/test/layers/manifold_layers.jl b/test/layers/manifold_layers.jl index aa51054e5..bc5d7b032 100644 --- a/test/layers/manifold_layers.jl +++ b/test/layers/manifold_layers.jl @@ -6,13 +6,12 @@ Random.seed!(1234) function stiefel_layer_test(T, M, N, tol = 1.0f-1) model = Chain(StiefelLayer(M, N), StiefelLayer(N, N)) ps = NeuralNetwork(model, T).params - o = Optimizer(Adam(T; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; + o = Optimizer(Adam(; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; step_size = T(1.0f0)) dx = (L1 = (weight = rand(T, N, M),), L2 = (weight = rand(T, N, N),)) ps_copy = deepcopy(ps) - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) # check that the new weight is different from the old one @test norm(ps_copy[1].weight - ps[1].weight) > T(tol) # check that the updated elements are on the Stiefel Manifold @@ -23,14 +22,13 @@ end function grassmann_layer_test(T, M, N, tol = 1.0f-1) model = Chain(GrassmannLayer(M, N), StiefelLayer(N, N)) ps = NeuralNetwork(model, T).params - o = Optimizer(Adam(T; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; + o = Optimizer(Adam(; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; step_size = T(1.0f0)) dx = (L1 = (weight = rand(T, N, M),), L2 = (weight = rand(T, N, N),)) ps_copy = deepcopy(ps) - λY = GlobalSection(ps) for i in 1:4 - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) end # check that the new weight is different from the old one @test norm(ps_copy[1].weight - ps[1].weight) > T(tol) diff --git a/test/layers/multi_head_attention_stiefel_optim_cache.jl b/test/layers/multi_head_attention_stiefel_optim_cache.jl index c94822e9a..8a0a29ef5 100644 --- a/test/layers/multi_head_attention_stiefel_optim_cache.jl +++ b/test/layers/multi_head_attention_stiefel_optim_cache.jl @@ -6,9 +6,9 @@ using GeometricMachineLearning, Test import GeometricOptimizers using GeometricOptimizers: GradientCache, MomentumCache, AdamCache import Random, LinearAlgebra -# The cache and the state hold a `NetworkParameters` per layer: a layer is wrapped at the -# `GeometricOptimizers` boundary, and the wrap shares the leaf arrays. `mapparameters` and not `map`, -# because it recurses on the branches and so reaches the leaves whichever shape it is handed. +# The cache holds one `NetworkParameters` over the whole network, the shape of the parameters. +# `mapparameters` and not `map`, because it recurses on the branches and so reaches the leaves +# whichever shape it is handed. # # Each helper below has two methods, because two shapes arrive: a wrapped layer is a # `NetworkParameters` and the tree above it is a plain `NamedTuple`. @@ -73,9 +73,9 @@ function test_cache_setups_for_optimizer_for_multihead_attention_layer(T::Type, o₂ = Optimizer(MomentumMethod(), ps) o₃ = Optimizer(GradientMethod(), ps) - check_adam_cache(o₁.cache) - check_momentum_cache(o₂.cache) - check_gradient_cache(o₃.cache) + check_adam_cache(o₁.training.cache) + check_momentum_cache(o₂.training.cache) + check_gradient_cache(o₃.training.cache) end test_cache_setups_for_optimizer_for_multihead_attention_layer(Float32, 64, 8) diff --git a/test/layers/multi_head_attention_stiefel_retraction.jl b/test/layers/multi_head_attention_stiefel_retraction.jl index 7b38a33ca..9b340e6b9 100644 --- a/test/layers/multi_head_attention_stiefel_retraction.jl +++ b/test/layers/multi_head_attention_stiefel_retraction.jl @@ -53,7 +53,7 @@ function test_multi_head_attention_retraction( model = Chain(MultiHeadAttention(dim, n_heads, Stiefel = true)) ps = NeuralNetwork(model, backend, T).params - cache = Optimizer(MomentumMethod(), ps).cache + cache = Optimizer(MomentumMethod(), ps).training.cache check_retraction_geodesic(cache) diff --git a/test/layers/multi_head_attention_stiefel_setup.jl b/test/layers/multi_head_attention_stiefel_setup.jl index 5a90ea42c..20a55d59c 100644 --- a/test/layers/multi_head_attention_stiefel_setup.jl +++ b/test/layers/multi_head_attention_stiefel_setup.jl @@ -50,7 +50,7 @@ function check_multi_head_attention_stiefel_setup(T::Type, N::Int, n::Int) @test nleaves > 0 @test check_setup_calls[] == nleaves - gx = Optimizer(MomentumMethod(), ps).cache + gx = Optimizer(MomentumMethod(), ps).training.cache check_grad_setup(gx) end diff --git a/test/layers/transformer_optimizer.jl b/test/layers/transformer_optimizer.jl index 1d7ffcfd7..9597b15dc 100644 --- a/test/layers/transformer_optimizer.jl +++ b/test/layers/transformer_optimizer.jl @@ -26,12 +26,9 @@ function transformer_gradient_test(T, dim, n_heads, L, seq_length = 8, batch_siz ps₂ = deepcopy(ps) ps₃ = deepcopy(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!(ps₁, o₁, dx) + optimization_step!(ps₂, o₂, dx) + optimization_step!(ps₃, o₃, dx) @test typeof(ps₁) == typeof(ps₂) == typeof(ps₃) == typeof(ps) @test ps₁[1].PQ.head_1 ≈ ps₂[1].PQ.head_1 @test ps₁[1].PQ.head_1 ≉ ps₃[1].PQ.head_1 diff --git a/test/optimizers/adam_with_learning_rate_decay.jl b/test/optimizers/adam_with_learning_rate_decay.jl index 12c590f94..145467b4c 100644 --- a/test/optimizers/adam_with_learning_rate_decay.jl +++ b/test/optimizers/adam_with_learning_rate_decay.jl @@ -3,6 +3,7 @@ using GeometricMachineLearning using GeometricMachineLearning: ResNetLayer, params using LinearAlgebra: I using Test +import GeometricOptimizers import Random Random.seed!(123) @@ -28,7 +29,7 @@ function train_network(; n_epochs = 2048) # pairing its own `Optimizer` takes, so it splats rather than being passed positionally. The # element type is positional and defaults to `Float64` here, where GML's own version took it # from `η₁` and so defaulted to `Float32`. - o₂ = Optimizer(nn₂; AdamOptimizerWithDecay(n_epochs, eltype(dl))...) + o₂ = Optimizer(nn₂; AdamOptimizerWithDecay(n_epochs)...) batch = Batch(5, 1) loss = FeedForwardLoss() @@ -59,7 +60,7 @@ function train_manifold_network(; n_epochs = 128) arch = Chain(StiefelLayer(1, 20), Dense(20, 20, tanh), Dense(20, 1, identity)) nn = NeuralNetwork(arch, CPU(), eltype(dl)) - o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, eltype(dl))...) + o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) loss_array = o(nn, dl, Batch(5, 1), n_epochs, FeedForwardLoss(); show_progress = false) @test all(isfinite, loss_array) @@ -72,32 +73,31 @@ end @doc raw""" The schedule is walked from ``t = 1``, not from ``t = 0``. -`optimization_step!` increments `opt.iterations` *before* it reads the step size, so the first step -of a run takes ``\alpha(1) = \gamma\eta_1`` and not ``\alpha(0) = \eta_1``. That is how the pre-0.5 -`AdamOptimizerWithDecay` counted — it incremented `o.step` before `update!` — and how +`optimization_step!` increments the iteration number *before* it reads the step size, so the first +step of a run takes ``\alpha(1) = \gamma\eta_1`` and not ``\alpha(0) = \eta_1``. That is how the +pre-0.5 `AdamOptimizerWithDecay` counted — it incremented `o.step` before `update!` — and how `DecayingStatic` counts, because `GeometricOptimizers.solve!` calls `increase_iteration_number!` before `solver_step!`. Reading before incrementing put every step of a run one place early in the schedule, which `test/adam_optimizer_with_decay.jl` upstream asserts does not happen. """ function schedule_starts_at_one(; n_epochs = 100, η₁ = 1e-2, η₂ = 1e-6) - method = AdamOptimizerWithDecay(n_epochs, Float64; η₁ = η₁, η₂ = η₂) + method = AdamOptimizerWithDecay(n_epochs; η₁ = η₁, η₂ = η₂) o = Optimizer(NetworkParameters((weight = zeros(2, 2),)); method...) γ = exp(log(η₂ / η₁) / n_epochs) - @test o.step_size isa DecayingStatic - @test o.iterations == 0 + @test o.training.linesearch isa DecayingStatic + @test GeometricOptimizers.iteration_number(o.training.state) == 0 for t in 1:4 - o.iterations += 1 - @test GeometricMachineLearning._current_step_size(o, o.iterations) ≈ η₁ * γ^t + @test GeometricOptimizers.step_size(o.training.linesearch, t) ≈ η₁ * γ^t end # and the same thing through the public entry point: one step of a Euclidean parameter with a # gradient of `1` moves it by `α₁ / (√1 + δ) ≈ α₁`, so the distance travelled reports the α used ps = NetworkParameters((weight = zeros(2, 2),)) opt = Optimizer(ps; method...) - optimization_step!(opt, GlobalSection(ps), ps, (weight = ones(2, 2),)) - @test opt.iterations == 1 + optimization_step!(ps, opt, (weight = ones(2, 2),)) + @test GeometricOptimizers.iteration_number(opt.training.state) == 1 @test abs(ps.weight[1, 1]) ≈ η₁ * γ rtol = 1e-6 end diff --git a/test/optimizers/composite_method.jl b/test/optimizers/composite_method.jl new file mode 100644 index 000000000..b58e8787e --- /dev/null +++ b/test/optimizers/composite_method.jl @@ -0,0 +1,55 @@ +using GeometricMachineLearning, Test, KernelAbstractions +import GeometricOptimizers +import Random + +# A network with a Stiefel layer beside a Euclidean one is the shape `CompositeMethod` exists for: +# `ScalarMomentAdam` steps a single `StiefelManifold`, so the dense layer needs another method. The +# per-leaf caches and states are `GeometricOptimizers`'; what is pinned here is that `Optimizer` +# hands a composite through to them unchanged. + +function stiefel_beside_dense(T) + Random.seed!(7) + NeuralNetwork(Chain(StiefelLayer(6, 3), Dense(6, 6, tanh)), KernelAbstractions.CPU(), T) +end + +function minibatch_gradient(T, k) + Random.seed!(100 + k) + (L1 = (weight = randn(T, 6, 3),), L2 = (W = randn(T, 6, 6), b = randn(T, 6))) +end + +@testset "a composite gives every leaf the state of the method it selects" begin + nn = stiefel_beside_dense(Float32) + opt = Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn) + states = opt.training.state.states + @test opt.training.state isa CompositeState{Float32} + @test states.L1.weight isa ScalarMomentAdamState{Float32} + @test states.L2.W isa AdamState{Float32} + @test states.L2.b isa AdamState{Float32} + + ps = GeometricMachineLearning.params(nn) + for k in 1:3 + optimization_step!(ps, opt, minibatch_gradient(Float32, k)) + end + @test ps.L1.weight isa StiefelManifold + @test GeometricOptimizers.check(ps.L1.weight) < 1.0f-5 + @test GeometricOptimizers.iteration_number(opt.training.state) == 3 +end + +@testset "a composite of one method is that method, exactly" begin + # Each optimizer is built straight after its network, so both draw their section completions + # from the same position of the seed. + whole = stiefel_beside_dense(Float64) + whole_opt = Optimizer(Adam(), whole) + per_leaf = stiefel_beside_dense(Float64) + per_leaf_opt = Optimizer(CompositeMethod(; manifold = Adam(), array = Adam()), per_leaf) + for k in 1:5 + optimization_step!(GeometricMachineLearning.params(whole), whole_opt, + minibatch_gradient(Float64, k)) + optimization_step!(GeometricMachineLearning.params(per_leaf), per_leaf_opt, + minibatch_gradient(Float64, k)) + end + a, b = GeometricMachineLearning.params(whole), GeometricMachineLearning.params(per_leaf) + @test a.L1.weight.A == b.L1.weight.A + @test a.L2.W == b.L2.W + @test a.L2.b == b.L2.b +end diff --git a/test/optimizers/gml_gradient_dispatch.jl b/test/optimizers/gml_gradient_dispatch.jl deleted file mode 100644 index 85a9bcb6c..000000000 --- a/test/optimizers/gml_gradient_dispatch.jl +++ /dev/null @@ -1,129 +0,0 @@ -# What `_GMLGradient` is called on, and by which method. -# -# `_GMLGradient` is a `GeometricOptimizers.Gradient`, so its functor methods share a method table with -# upstream's. Two of them overlap on a shape this package actually produces, and the overlap is an -# *ambiguity* rather than a choice — which is a run-time error several frames into a solve, not a -# signature a reader would notice. So it is asserted here rather than left to whichever call path -# happens to reach it. - -using GeometricMachineLearning -using GeometricOptimizers: StiefelManifold, GrassmannManifold, rgrad -using NeuralNetworkParameters: NetworkParameters -using Test -import Random - -Random.seed!(1234) - -GML = GeometricMachineLearning - -# A bare `Manifold`. `_GMLGradient{T} <: GeometricOptimizers.Gradient{T}` and `Manifold{T} <: -# AbstractMatrix{T}`, so `(::_GMLGradient{T})(::AbstractArray{T})` and upstream's -# `(::Gradient{T})(::Manifold{T})` are ambiguous here: neither is more specific. Without a method of -# its own this call is `MethodError: ... is ambiguous`. -# -# The value is this package's answer and not upstream's. Upstream's body is -# `rgrad(x, reshape(grad(vec(x)), size(x)...))`, which evaluates an *inner* gradient — and a -# `_GMLGradient` has none, since the whole point of it is that the Euclidean gradient was computed -# already and is being carried in `dp`. -@testset "a bare Manifold reaches this package's method, not upstream's" begin - for M in (StiefelManifold, GrassmannManifold) - Y = rand(M, 4, 2) - dp = randn(4, 2) - g = GML._GMLGradient{Float64, typeof(dp)}(dp) - - @test g(Y) == rgrad(Y, dp) - @test which(g, Tuple{typeof(Y)}).module === GeometricMachineLearning - end -end - -# A bare `AbstractVector`. `_GMLGradient{T} <: GeometricOptimizers.Gradient{T}` and `SimpleSolvers` -# has `(::Gradient{T})(::AbstractVector{T})`, so without a method of its own, -# `(::_GMLGradient{T})(::AbstractArray{T})` and that one are equally specific on a `Vector` -- -# the call is `MethodError: ... is ambiguous`, same class of defect as the `Manifold` case above. -@testset "a bare AbstractVector reaches this package's method, not upstream's" begin - dp = randn(3) - g = GML._GMLGradient{Float64, typeof(dp)}(dp) - - @test g(randn(3)) === dp - @test which(g, Tuple{Vector{Float64}}).module === GeometricMachineLearning -end - -# The two shapes the optimizer actually hands it, so that the method above is not the only one -# covered and a later narrowing of either shows up here. -@testset "a wrapped layer and a plain array" begin - dp = randn(3, 3) - g = GML._GMLGradient{Float64, typeof(dp)}(dp) - # a plain array leaf: the gradient is already Euclidean, so it passes through - @test g(randn(3, 3)) === dp - - ps = NetworkParameters((weight = rand(StiefelManifold, 4, 2),)) - dps = (weight = randn(4, 2),) - gps = GML._GMLGradient{Float64, typeof(dps)}(dps) - out = gps(ps) - # rebuilt in the shape of its *first* argument -- a container, which is what - # `GeometricOptimizers._copyto!` has a method for - @test out isa NetworkParameters - @test out.weight == rgrad(ps.weight, dps.weight) -end - -# `_is_layer` is the rule for how a network is split into caches, so its boundary cases are stated -# rather than inherited. An empty set is *zero* caches' worth and descends; `all` over an empty -# collection would otherwise make it vacuously one. -@testset "_is_layer at the boundary" begin - @test GML._is_layer((W = randn(2, 2), b = randn(2))) - @test GML._is_layer(NetworkParameters((W = randn(2, 2), b = randn(2)))) - # a set whose values are branches is a subtree, and gets one cache per layer - @test !GML._is_layer((L1 = (W = randn(2, 2),), L2 = (W = randn(2, 2),))) - @test !GML._is_layer(NamedTuple()) - @test !GML._is_layer(NetworkParameters(NamedTuple())) - @test !GML._is_layer(randn(2, 2)) - - # a layer whose weights do not share an element type is still one cache, with `T` the promotion - @test GML._is_layer((W = randn(Float32, 2, 2), b = randn(Float64, 2))) -end - -# `_as_go_solution` is the wrap itself, and the property that makes the wrap free: it shares the leaf -# arrays, so the optimizer step writes through to the network's own weights and nothing is copied -# back. That is what makes this a boundary change rather than a change of behaviour, and it is the one -# claim of this release that a reader cannot check by reading a signature. -@testset "_as_go_solution wraps a layer and shares its arrays" begin - W, b = randn(2, 2), randn(2) - layer = (W = W, b = b) - - wrapped = GML._as_go_solution(layer) - @test wrapped isa NetworkParameters - @test wrapped.W === W # shared, not copied - @test wrapped.b === b - - # already in the right shape: passed through untouched - ps = NetworkParameters(layer) - @test GML._as_go_solution(ps) === ps - Y = rand(StiefelManifold, 4, 2) - @test GML._as_go_solution(Y) === Y - v = randn(3) - @test GML._as_go_solution(v) === v - - # a subtree is not a layer, so it is left alone for the tree branch to descend - tree = (L1 = (W = randn(2, 2),), L2 = (W = randn(2, 2),)) - @test GML._as_go_solution(tree) === tree -end - -# `_get_contents` unwraps whatever shape the pullback handed back. Six methods, because a reverse pass -# produces a container or the bare `NamedTuple` it wraps, and `Zygote` may hand either back bare, in a -# one-tuple, or in a one-element vector. Each is exercised here: splitting a union into a method per -# shape means each method needs its own case, or the split trades one tested path for six untested -# ones. -@testset "_get_contents unwraps every shape a pullback returns" begin - nt = (L1 = (W = randn(2, 2), b = randn(2)),) - ps = NetworkParameters(nt) - - for x in (nt, ps) - @test GML._get_contents(x) === x - @test GML._get_contents((x,)) === x - @test GML._get_contents([x]) === x - end - - # more than one set is a caller error and says so, rather than silently taking the first - @test_throws ArgumentError GML._get_contents([ps, ps]) - @test_throws ArgumentError GML._get_contents([nt, nt]) -end diff --git a/test/optimizers/gradient_optimizer.jl b/test/optimizers/gradient_optimizer.jl index 376e946bd..12d68225c 100644 --- a/test/optimizers/gradient_optimizer.jl +++ b/test/optimizers/gradient_optimizer.jl @@ -18,8 +18,7 @@ function gradient_optimizer_euclidean(N; n_steps = 20, step_size = 1e-3) o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps ∇L = Zygote.gradient(loss, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, ∇L) + optimization_step!(ps, o, ∇L) end loss2 = loss(ps) @test loss1 > loss2 @@ -38,8 +37,7 @@ function gradient_optimizer_stiefel(N, n; n_steps = 20, step_size = 1e-3) o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps ∇L = Zygote.gradient(loss, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, ∇L) + optimization_step!(ps, o, ∇L) end loss2 = loss(ps) @test loss1 > loss2 diff --git a/test/optimizers/momentum_optimizer.jl b/test/optimizers/momentum_optimizer.jl index 0351140a1..9359c806f 100644 --- a/test/optimizers/momentum_optimizer.jl +++ b/test/optimizers/momentum_optimizer.jl @@ -16,11 +16,10 @@ growing one that never saturates, which is what this test rules out. """ function momentum_matches_geometric_sum(; α = 0.9, η = 0.1, n_steps = 10) ps = NetworkParameters((A = zeros(2, 2),)) - o = Optimizer(MomentumMethod(α), ps; step_size = η) - λY = GlobalSection(ps) + o = Optimizer(MomentumMethod(; α = α), ps; step_size = η) previous = 0.0 for t in 1:n_steps - optimization_step!(o, λY, ps, (A = ones(2, 2),)) + optimization_step!(ps, o, (A = ones(2, 2),)) step = previous - ps.A[1, 1] previous = ps.A[1, 1] @test step ≈ η * (1 - α^t) / (1 - α) @@ -38,10 +37,9 @@ function momentum_optimizer_stiefel(N, n; n_steps = 50, step_size = 1e-3, α = 0 loss(ps) = norm(ps.Y * ps.Y' - B)^2 ps = NetworkParameters((Y = rand(StiefelManifold, N, n),)) loss1 = loss(ps) - o = Optimizer(MomentumMethod(α), ps; step_size = step_size) - λY = GlobalSection(ps) + o = Optimizer(MomentumMethod(; α = α), ps; step_size = step_size) for _ in 1:n_steps - optimization_step!(o, λY, ps, Zygote.gradient(loss, ps)[1]) + optimization_step!(ps, o, Zygote.gradient(loss, ps)[1]) end @test loss(ps) < loss1 @test ps.Y' * ps.Y ≈ I diff --git a/test/optimizers/optimization_step.jl b/test/optimizers/optimization_step.jl index 781d625a0..b38a068bf 100644 --- a/test/optimizers/optimization_step.jl +++ b/test/optimizers/optimization_step.jl @@ -17,8 +17,7 @@ function optimization_step_test(N, n, T) o = Optimizer(m, ps) ps2 = deepcopy(ps) - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) @test typeof(ps[1].weight) <: StiefelManifold for (layers1, layers2) in zip(values(ps), values(ps2)) for key in keys(layers1) @@ -35,23 +34,16 @@ for N in 4:N_max end end -# The regression net for the branch order in `_make_optimizer_cache`/`_make_optimizer_state`: a -# `NetworkParameters` is a tree to descend into, so it is recognised *structurally*, before the -# capability question `_use_go_cache` asks. The shape that follows is one `GeometricOptimizers` cache -# per *layer* -- not one for the root, which is what putting the capability question first would give -# the moment `GeometricOptimizers` adds the container to `OptimizerSolution`, and not one per weight, -# which is what hoisting the `NamedTuple` branch as well would give. -@testset "one cache and one state per layer, keyed by the network's layers" begin +# The whole network is one `GeometricOptimizers.TrainingOptimizer`: one cache and one state over the +# whole parameter set, and not one per layer as before 0.9. +@testset "one cache and one state over the whole network" begin model = Chain(StiefelLayer(6, 3), Dense(6, 6, tanh)) nn = NeuralNetwork(model, KernelAbstractions.CPU(), Float32) o = Optimizer(AdamOptimizer(), nn) - for tree in (o.cache, o.state) - @test tree isa NamedTuple - @test keys(tree) == keys(GeometricMachineLearning.params(nn)) - end - @test o.cache.L1 isa GeometricOptimizers.OptimizerCache - @test o.cache.L2 isa GeometricOptimizers.OptimizerCache - @test o.state.L1 isa GeometricOptimizers.OptimizerState - @test o.state.L2 isa GeometricOptimizers.OptimizerState + @test o.training isa GeometricOptimizers.TrainingOptimizer + @test o.training.cache isa GeometricOptimizers.AdamCache{Float32} + @test o.training.state isa AdamState{Float32} + @test keys(GeometricOptimizers.solution(o.training.state)) == + keys(GeometricMachineLearning.params(nn)) end diff --git a/test/optimizers/psd_optim.jl b/test/optimizers/psd_optim.jl index 8afad17a7..3ccc27d48 100644 --- a/test/optimizers/psd_optim.jl +++ b/test/optimizers/psd_optim.jl @@ -26,7 +26,7 @@ A = [0.06476993260924702 0.8369280855305259 0.6245358125914054 0.140729967064923 """ This tests if the optimizers can find the optimal PSD solution. """ -function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = cayley) +function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = Cayley()) # Seeded here and not once at the top of the file, for the reason given at the same place in # `svd_optim.jl`: `svd_test` is called twice, and with a single top-level seed the second call # starts from whatever RNG state the first one happened to leave behind, which makes the @@ -69,13 +69,12 @@ function train_network!(o::Optimizer, model::Chain, ps::NetworkParameters, for _ in 1:train_steps dx = Zygote.gradient(error, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) #println(error(ps)) end ps[1].weight, ps[2].weight, error(ps) end -for retraction in (geodesic, cayley) +for retraction in (Geodesic(), Cayley()) svd_test(A, 4, retraction = retraction) end diff --git a/test/optimizers/step_size_element_type.jl b/test/optimizers/step_size_element_type.jl deleted file mode 100644 index a2a13658a..000000000 --- a/test/optimizers/step_size_element_type.jl +++ /dev/null @@ -1,52 +0,0 @@ -# `_leaf_optim_step!` scales `direction(cache)` by `T(step_size)`, matching the three -# `_euclidean_update!` methods a few lines below it. The funnel (`_current_step_size`) always -# hands a `Float64`, no matter what `T` the parameters are, so converting first is what keeps a -# `Float32` layer scaled at `Float32` precision instead of `Float64` precision rounded back on -# write. See `CHANGELOG.md` for the fix this guards. -# -# The test is invariance rather than a fixed expected value: the leaf step must give the *same* -# result whether the funnel hands it a `Float64` step size or one already rounded to `T`. Only -# `T = Float32` can tell the two calls apart -- `T(step_size)` is a no-op when `T` is already -# `Float64`, so a `Float64` combination holds whatever the conversion does and checks nothing. -# `Float64` is therefore not looped over here. -# -# The `seed = 2` and the `step_size = 0.1` are chosen rather than arbitrary. Whether a `Float64`- -# and a `Float32`-scaled direction round to different bit patterns depends on the actual values, -# so for most seeds the two calls agree however the step size is scaled, and the test then passes -# without discriminating. With these values, and with the conversion dropped, both `Adam` and -# `MomentumMethod` separate by `2.9802322f-8` -- one `Float32` ULP at this magnitude -- and each -# method is checked on its own. - -using GeometricMachineLearning -using GeometricOptimizers -using Test -import Random - -GML = GeometricMachineLearning - -function test_leaf_step_matches_prerounded_step_size(method, seed::Int, step_size::Float64) - T = Float32 - Random.seed!(seed) - ps_a = (weight = rand(T, 6, 6),) - dp = (weight = rand(T, 6, 6),) - ps_b = deepcopy(ps_a) - λY_a = GML.GlobalSection(ps_a) - λY_b = GML.GlobalSection(ps_b) - - cache_a = GML._make_optimizer_cache(method, ps_a) - state_a = GML._make_optimizer_state(method, ps_a) - cache_b = GML._make_optimizer_cache(method, ps_b) - state_b = GML._make_optimizer_state(method, ps_b) - - # the funnel's usual `Float64` step size ... - GML._leaf_optim_step!( - cache_a, state_a, dp, ps_a, λY_a, method, GeometricOptimizers.cayley, step_size) - # ... versus the same value pre-rounded to `T`. - GML._leaf_optim_step!( - cache_b, state_b, dp, ps_b, λY_b, method, GeometricOptimizers.cayley, T(step_size)) - - @test ps_a.weight == ps_b.weight -end - -test_leaf_step_matches_prerounded_step_size(GeometricOptimizers.Adam(), 2, 0.1) -test_leaf_step_matches_prerounded_step_size(GeometricOptimizers.MomentumMethod(0.5), 2, 0.1) diff --git a/test/optimizers/structured_array_parameters.jl b/test/optimizers/structured_array_parameters.jl index a685780c5..c9a000d2c 100644 --- a/test/optimizers/structured_array_parameters.jl +++ b/test/optimizers/structured_array_parameters.jl @@ -8,7 +8,8 @@ Random.seed!(123) @doc raw""" Every optimizer method has to work with the structured matrix types GML uses as *ordinary* weights: `SymmetricMatrix` (SympNet and symplectic attention layers), `SkewSymMatrix` (volume-preserving -attention) and `StrictlyLowerTriangular`/`StrictlyUpperTriangular` (volume-preserving feedforward layers). +attention) and `StrictlyLowerTriangular`/`StrictlyUpperTriangular` (volume-preserving feedforward +layers). These have their own storage and no `setindex!`, so they need a type-preserving `similar` and elementwise `_add!`/`_rac!`/`_square!`/`_div!`/`_rmul!` written on the free parameters; without them @@ -31,7 +32,7 @@ function optimizer_runs(architecture, batch, input_dim; T = Float64, n_epochs = # `AdamOptimizerWithDecay` is a `(algorithm, linesearch)` pairing rather than a method, so it # goes in through the keyword constructor nn, dl = nn_and_dl() - o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, T)...) + o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) loss_array = o(nn, dl, batch, n_epochs; show_progress = false) @test length(loss_array) == n_epochs @test all(isfinite, loss_array) diff --git a/test/optimizers/svd_optim.jl b/test/optimizers/svd_optim.jl index 033819dee..47b7b24ed 100644 --- a/test/optimizers/svd_optim.jl +++ b/test/optimizers/svd_optim.jl @@ -25,7 +25,7 @@ A = [0.06476993260924702 0.8369280855305259 0.6245358125914054 0.140729967064923 Random.seed!(1234) -function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = cayley) +function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = Cayley()) # Seeded here and not once at the top of the file. `svd_test` is called twice, and with a single # top-level seed the second call starts from whatever RNG state the first one happened to leave # behind -- so the convergence assertions below turned on how much randomness the optimizer @@ -65,12 +65,11 @@ function train_network!(o::Optimizer, model::Chain, ps::NetworkParameters, for _ in 1:train_steps dx = Zygote.gradient(error, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) end ps.L1.weight, ps.L2.weight, error(ps) end -for retraction in (geodesic, cayley) +for retraction in (Geodesic(), Cayley()) svd_test(A, 3, retraction = retraction) end diff --git a/test/runtests.jl b/test/runtests.jl index 9d1e773bc..5df2e8f0b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -59,8 +59,7 @@ if "core" in GROUPS @safetestset "Gradient optimizer tests" include("optimizers/gradient_optimizer.jl") @safetestset "Momentum optimizer tests" include("optimizers/momentum_optimizer.jl") @safetestset "Optimizers with structured (non-manifold) weights" include("optimizers/structured_array_parameters.jl") - @safetestset "_GMLGradient dispatch" include("optimizers/gml_gradient_dispatch.jl") - @safetestset "The GO-native leaf step scales in the parameter's own element type" include("optimizers/step_size_element_type.jl") + @safetestset "CompositeMethod, one method per leaf" include("optimizers/composite_method.jl") @safetestset "_custom_mul: NetworkParameters gradient structure" include("layers/double_multiplication_network_parameters_gradient.jl") @safetestset "Symplectic attention: NetworkParameters gradient structure" include("layers/symplectic_attention_network_parameters_gradient.jl") @safetestset "map_to_cpu" include("map_to_cpu_tests.jl")