Train on GeometricOptimizers 0.9: NetworkOptimizer and train! (part M1a)

[michakraus]

Summary

Training moves onto GeometricOptimizers 0.9's TrainingOptimizer (part M1a of the 0.9 simplification of GeometricOptimizers and GeometricMachineLearning).

• Optimizer and its callable methods are removed. train! takes GeometricOptimizers' TrainingOptimizer, which GML re-exports and which is built on the parameters of a network: TrainingOptimizer(nn.params; algorithm, linesearch, retraction). One cache and one state cover the whole parameter set. There is no GML wrapper type: src/optimizers/optimizer.jl holds only the re-exports.
• train!(nn, opt, dl, batch, n_epochs; loss, pullback = ZygotePullback(loss), show_progress) in src/data_loader/train.jl replaces src/data_loader/optimize.jl: the four callable methods, optimize_for_one_epoch! as a public name, and λY go. pullback takes any AbstractPullback, so a Hamiltonian network trains with SymbolicPullback(arch). The history has the element type of the loss value. n_epochs < 1 is an ArgumentError before any step, and so is a network with no default loss when loss is not given. A second call with the same opt resumes the first. The epoch loop is a plain loop: a map with a closure made the first call of a transformer compile several times as long, which timed out four reproduction scripts in CI.
• The aliases GradientOptimizer, MomentumOptimizer and AdamOptimizer are removed. GML re-exports TrainingOptimizer, GradientMethod, MomentumMethod, Adam, AdamWithEuclideanDecay, AdamOptimizerWithDecay, DecayingStatic, Cayley and Geodesic. optimization_step! is not re-exported; a hand-written loop calls GeometricOptimizers.optimization_step!(nn.params, opt, dp).
• [compat]: GeometricOptimizers "0.9", AbstractNeuralNetworks "0.9.1", NeuralNetworkParameters "0.5", SymbolicNeuralNetworks "0.9", GeometricBase "0.15", GeometricEquations "0.21.5", GeometricSolutions "0.6.6".
• The 3-tensor Dense/Linear functors and the HDF5 save/load methods move to AbstractNeuralNetworks 0.9. ext/HDF5Ext.jl and the HDF5 weak dependency go. Aqua's piracies check runs with 0 piracies, and the ambiguity count is 0. This relates to Type piracy #145 and closes known issue B7.
• GeometricOptimizers 0.9's renames: LowerTriangular/UpperTriangular become StrictlyLowerTriangular/StrictlyUpperTriangular, in a separate mechanical commit.
• Docs, tests, scripts and the README move to train! and TrainingOptimizer. New and rewritten testsets loop over REAL_ELTYPES (Float32, Float64).
• scripts/verification/all_zero_leaf.jl asserts what a step does to a leaf that the loss does not read.
• The GPU paragraph of the manual is measured again against GeometricOptimizers 0.9.0 on an Apple M4 Max, and train! now runs on the device too.

Relates to JuliaGNI/AbstractNeuralNetworks.jl#15: the resume test (two runs of 5 epochs equal one run of 10) covers it within one process. The issue stays open for a decision by hand.

Commits

• 92ad227 the API, with the old triangular names kept through an aliased import; e8ff170 the rename only, with its CHANGELOG entry.
• 274a7e0, 68c47ff, 69829d2 KNOWN_ISSUES entries and text fixes from the critic and the pre-PR verifier.
• 826ae32 fixes of the Scripts and Documentation CI failures; it changes scripts/verification/all_zero_leaf.jl and docs/make.jl.
• 3529019 scripts import optimization_step! from GeometricOptimizers, and a wider allocation test (superseded by 93409b0, see below).
• 1d661a7, 93409b0, adac904, 20da228 the fixes of the review below.

Review findings and what changed

The review on 826ae32 had 11 findings; each is resolved in the comment that accompanies this update. In short:

• The required Julia min - windows-latest job failed on a 32-byte bound between two @allocated readings of a Stiefel step: one through the wrapper and one direct. With the wrapper removed, that comparison measured a step against itself, so it is gone; the Euclidean == 0 bound and the @inferred checks stay.
• The first train! of a transformer took about 97 s to compile with the map loop and 12–22 s with the plain loop, in fresh processes on Julia 1.13.1, with the same history. The four scripts that timed out at 300 s in CI run in 63–96 s locally.
• SymbolicPullback trains again through the pullback keyword. Its gradient is the gradient of the per-sample sum of losses, not of HNNLoss; this is older than this PR (same numbers on main), is filed as SymbolicPullback(arch) gives the gradient of the per-sample sum, not of HNNLoss #343 and recorded as K13, and the new test marks its comparison with Zygote @test_broken.
• test/optimizers/step_size_element_type.jl and the PrecomputedGradient testset of gml_gradient_dispatch.jl tested GeometricOptimizers, not this package; they are deleted here and move to GeometricOptimizers' suite. The _get_contents testset moves to test/pullbacks/get_contents.jl.

Test runs

• Full suite on 20da228 (the head), Julia 1.13.1, Pkg.test() with the default groups: Pass 4928, Fail 0, Error 0, Broken 6. The broken results: JET 3 of 145 (JET reports a runtime dispatch in PoissonTensor and in one batch method #325), Metal 1 (no GPU in the sandbox), and 2 in train!'s SymbolicPullback testset (SymbolicPullback(arch) gives the gradient of the per-sample sum, not of HNNLoss #343).
• docs/check_references.jl: 0 unresolved references.
• The four reproduction scripts that timed out in CI, in smoke mode: all ran to completion.
• scripts/verification/all_zero_leaf.jl: every assertion holds in Float32 and Float64.
• The GPU paragraph's checks on Metal: 13 of 13 ran.

Not run locally: the full docs build (the Documentation workflow runs it), and the suite on Julia 1.12 and on Windows.

Known issues added

• K6 · The gradient of a SymmetricMatrix leaf with respect to a NetworkParameters differs from the one with respect to the bare NamedTuple (relates to Zygote gradient w.r.t. NetworkParameters loses the SymmetricMatrix leaf on a second access #319).
• K7 · A network with a Stiefel leaf raises in mat_tensor_mul when the data has another element type.
• K8 · A DataLoader of one sample cannot be batched.
• K10 · The two training tests of adam_with_learning_rate_decay.jl run in Float64 only. Pre-existing.
• K11 · A comment in adam_with_learning_rate_decay.jl reports a past experiment. Pre-existing.
• K12 · Batch never draws the last sample of a DataLoader of matrices. Pre-existing.
• K13 · SymbolicPullback(arch) gives the gradient of the per-sample sum, not of HNNLoss (SymbolicPullback(arch) gives the gradient of the per-sample sum, not of HNNLoss #343). Pre-existing.

Removed: B7, C1, C2, D4 (closed by this change). Changed: C16, K1.

Unresolved

• train! infers its return as an abstract Vector, because ZygotePullback's call is annotated ::Tuple and its loss value infers as Any; at run time the history is Vector{L}, and the tests assert it. The pullback is the next part's.
• The tests read opt.state, opt.cache, opt.method and opt.linesearch directly: GeometricOptimizers 0.9 has no public accessor for them.

Pre-existing

• The closed-ID notes in KNOWN_ISSUES.md say "by this release" for B1, B2, B6, B8 and B9, which closed in 0.8.0.
• The tutorials and reproduction scripts pass an epoch count to AdamOptimizerWithDecay, whose n counts minibatch steps. The behaviour is unchanged.
• The docs navigation still labels the training page "Optimizer".

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.49%. Comparing base (cfe4aea) to head (20da228).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #342      +/-   ##
==========================================
+ Coverage   86.64%   88.49%   +1.84%     
==========================================
  Files          76       74       -2     
  Lines        2786     2633     -153     
==========================================
- Hits         2414     2330      -84     
+ Misses        372      303      -69     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[michakraus]

Review of 826ae32

Verdict: request changes (posted as a comment review). A required check is red because of a test that this PR adds. The new train! takes 5 to 8 times longer than the code it replaces to compile its first call on a transformer. That makes the Scripts job time out.

CI, read per job

job	required?	caused by this PR?	red on main?	what the log says
Julia min - windows-latest - default	yes (branch protection and ruleset)	yes: the failing test is new here	no (main run 37980189635: success)	test/optimizers/optimization_step.jl:80, Expression: abs(allocations(optimization_step!, ps, o, dp) - allocations(optimization_step!, ps, o.optimizer, dp)) ≤ 32, Evaluated: 48 ≤ 32; 731 passed, 1 failed
Scripts - ubuntu-latest	no	yes: see finding 2	no (main run 37980189875: all 26 scripts ok)	TIMEOUT 300.1s reproduction/symplectic_transformer/double_pendulum.jl, TIMEOUT 300.2s …/double_pendulum_short_integration.jl, TIMEOUT 300.1s …/symplectic_transformer_vector_softmax.jl, TIMEOUT 300.2s reproduction/volume_preserving_transformer/rigid_body.jl
Downgrade - ubuntu-latest	no	no	yes, the same job is red on main (run 37980189635; the workflow conclusion reads success)	forcedeps check failed: GPUArraysCore resolved to 0.2.1 but lower bound is 0.2.0, the same for Zygote, ChainRulesCore, NNlib and KernelAbstractions, as on main. On main GeometricBase also failed. This PR fixes that one by raising the bound to 0.15

The Scripts job on main is three days old. The attribution does not rest on it: finding 2 compares the base with this PR, and the shipped train! with a rewrite, on one machine.

Findings

#	severity	file:line	claim	evidence
1	blocker	test/optimizers/optimization_step.jl:80-81	The Stiefel half of "allocates nothing of its own" requires two separate @allocated readings of a step to agree within 32 bytes. That step allocates 2 to 4 kB. The bound fails on Windows only, so this test turns a required check red.	The CI log above (48 ≤ 32 on Windows, Julia 1.12.7). On macOS, with --check-bounds=yes, the difference is 0 in 90 of 90 trials per element type on both 1.12.7 and 1.13.1 (Float32 step 2320 B, Float64 3520 B). The thing under test is the one-line forward at src/optimizers/optimizer.jl:33-35. The Euclidean == 0 at :71 already shows that it allocates nothing.
2	bug	src/data_loader/train.jl:33-37 with :41-45	The map(1:n_epochs) do … end around _train_epoch!, which itself holds sum(generator), makes the first train! call compile for 91–98 s. The old functor takes 18.4 s and a plain loop takes 11.2 s. The shipped form and the plain loop give the same history. All six reproduction scripts that slowed or timed out in the Scripts job train transformers for 2 epochs (the smoke size), so their run time is mostly compilation. Two of the six went from 131.7 to 289.3 s and from 145.2 to 293.7 s.	A fresh process for each run, Julia 1.13.1, StandardTransformerIntegrator(4; transformer_dim = 4, n_heads = 1, L = 1, n_blocks = 4, attention_activation = MatrixSoftmax()), Batch(64, 10), 2 epochs, show_progress = false. Base o(nn, dl, batch, 2): 18.4 s. This PR's train!: 98.1, 96.7, 91.0 and 93.5 s. The form below: 11.2 and 11.2 s. History [2.315492636526417, 2.172620338370672] from both PR forms. The parts, measured one at a time: the outer map with a plain inner loop takes 15.2 s, and an outer loop around the shipped _train_epoch! takes 11.2 s. So it is the nesting of the two that costs the time. The load average was 7–18, but the base and PR runs alternated, so the load affected both.
3	quality	src/data_loader/train.jl:31	train! always uses ZygotePullback(loss). The package still exports AbstractPullback and SymbolicPullback (src/GeometricMachineLearning.jl:250), and docs/src/architectures/hamiltonian_neural_network.md presents SymbolicPullback as the alternative gradient for HNNs. Nothing can train with it now. One keyword gives that back.	Read. On main, the functor took _pullback::AbstractPullback = ZygotePullback(loss) (src/data_loader/optimize.jl:85). The CHANGELOG records the removal, but not the fact that the exported pullback is left with no use.
4	quality	test/optimizers/step_size_element_type.jl (whole file), test/optimizers/gml_gradient_dispatch.jl:20-35	These tests test GeometricOptimizers, not this package. One tests the precision conversion of TrainingOptimizer's linesearch, and the other tests PrecomputedGradient(ps, dp)(ps). The only GML code that they reach is the keyword forward in NetworkOptimizer. These tests belong in GeometricOptimizers' suite. The file name gml_gradient_dispatch.jl names the _GMLGradient that this PR deletes.	Read. Neither testset calls a GML function except NetworkOptimizer(ps; kwargs...).
5	quality	scripts/verification/all_zero_leaf.jl:8-13, :71-79	The header states results: a leaf that the loss never reads stays == its start under three methods, and momentum carries a leaf that the loss stops reading. The script only prints figures, so nothing fails when a claim stops being true, and the CHANGELOG cites this script as the evidence. @assert r.never_euclidean == 0 && r.never_stiefel == 0 for GradientMethod, MomentumMethod and Adam, and @assert after_reading(T, m) > 0 for the three momentum methods, make it a check. Also, one = p -> … at :60 shadows Base.one.	Read.
6	quality (an option, for the author to decide)	src/optimizers/optimizer.jl:24-35	NetworkOptimizer is a wrapper with one field. Its only work is the nn constructor and the forward of optimization_step!. GeometricOptimizers' TrainingOptimizer(x::OptimizerSolution{T}; …) already takes a NetworkParameters. The type can go: train!(nn, opt::TrainingOptimizer, …), with TrainingOptimizer re-exported, and callers write TrainingOptimizer(params(nn); algorithm = Adam()). This removes one public type and the .optimizer indirection, which the tests write 17 times. If you want the nn form, keep the wrapper, because TrainingOptimizer(::NeuralNetwork) would be type piracy.	Read: GeometricOptimizers 0.9.0 src/optimizers/training_optimizer.jl, TrainingOptimizer(x::OptimizerSolution{T}; algorithm::FirstOrderMethod = Adam(), linesearch = default_step_size(algorithm), …).
7	nit	src/data_loader/train.jl:54-56	_default_loss has no fallback. If you call train! without loss on any other network, the error names an internal function. A one-line fallback that throws an ArgumentError which tells the caller to pass loss matches the docstring ("any other network needs loss").	train!(nn, NetworkOptimizer(nn), dl, Batch(5), 1; show_progress = false) with nn = NeuralNetwork(Chain(Dense(2, 2, identity))) gives MethodError: no method matching _default_loss(::AbstractNeuralNetworks.UnknownArchitecture, ::Batch{:FeedForward}).
8	nit	src/kernels/kernel_ad_routines/mat_tensor_mul.jl:76-78, :148-150; src/kernels/mat_tensor_mul.jl:161-163, :209-211	When the formatter re-wrapped the rename, it split AbstractArray{T, 3} across two lines (A::AbstractArray{ / T, 3}) where {T}). Join them by hand. Also, the docstrings at :155-158 and :203-206 still say "lower-triangular" and "upper-triangular", not strictly.	Read.
9	nit	test/data_loader/train.jl:44, test/optimizers/optimizer.jl:49-50, test/optimizers/optimization_step.jl:45-47	The tests read GeometricOptimizers' struct fields (opt.optimizer.state.iterations, .method, .linesearch, .cache). This ties GML's suite to GeometricOptimizers' internal layout.	Read.
10	nit	CHANGELOG.md:20, docs/src/index.md:55	The first Removed (breaking) bullet is one paragraph of 400 words. It mixes the migration guide, notes on behaviour and the bookkeeping of known issues. Split it into one bullet per migration. The GPU paragraph reports measurements against GeometricOptimizers 0.8.0, which the package can no longer load. Delete those measurements or measure again.	sed -n 20p CHANGELOG.md | wc -w gives 400.
11	nit	PR body	The body says that "The commits after e8ff170 change no executable line". But 826ae32 changes scripts/verification/all_zero_leaf.jl:21 and docs/make.jl. The suite totals are for e8ff170, not for the head, and at the head the Julia min - windows job fails. The critic and advisor tables are claims. This review did not use them as evidence.	gh pr view 342 --json commits; the CI log above.

The simpler code for finding 2

This form also takes in finding 3. It is measured cold at 11.2 s, against 91–98 s for the shipped form, and gives the same history. A Float32 loss still gives a Vector{Float32}, because the vector takes the type of the first epoch's value:

function train!(nn::NeuralNetwork, opt::NetworkOptimizer, dl::DataLoader, batch::Batch,
        n_epochs::Integer; loss::NetworkLoss = _default_loss(nn.architecture, batch),
        pullback::AbstractPullback = ZygotePullback(loss), show_progress::Bool = true)
    n_epochs ≥ 1 ||
        throw(ArgumentError("`train!` needs at least one epoch; `n_epochs` is $(n_epochs)."))
    progress = ProgressMeter.Progress(n_epochs; enabled = show_progress)
    history = [_train_epoch!(nn, opt, dl, batch, pullback, progress)]
    for _ in 2:n_epochs
        push!(history, _train_epoch!(nn, opt, dl, batch, pullback, progress))
    end
    history
end

function _train_epoch!(nn, opt, dl, batch, pullback, progress)
    batches = batch(dl)
    value = sum(_train_step!(nn, opt, dl, batch, pullback, indices) for indices in batches) /
            length(batches)
    ProgressMeter.next!(progress; showvalues = [(:TrainingLoss, value)])
    value
end

For finding 1, delete test/optimizers/optimization_step.jl:80-81. Keep the @inferred at :79 and the type check at :82. If you keep a bound on the Stiefel step, make it a tolerance well above the 16-byte steps in which Windows measures allocations, for example ≤ 256. Do not use a bound smaller than the noise.

Verified and fine

• src/optimizers/optimizer.jl: the docstring defaults (algorithm = Adam(), linesearch = default_step_size(algorithm), retraction = Cayley()) are the same as those of GeometricOptimizers 0.9.0's TrainingOptimizer.
• The deletion of src/data_loader/optimize.jl, the per-layer caches and _GMLGradient: git grep finds no remaining use of Optimizer(, optimize_for_one_epoch! or the three aliases in src, docs/src, scripts, test or README.md, except the absence test in test/quality/exports.jl:37.
• The warm run time per call is unchanged: the second 2-epoch call takes 0.01 s on both trees.
• KernelAbstractions.get_backend(::PoissonTensor) extends a function of another package on GML's own type, so it is not piracy.
• Project.toml: the Downgrade forcedeps failures do not come from this PR, and the new lower bounds (GeometricOptimizers 0.9.0, AbstractNeuralNetworks 0.9.1, NeuralNetworkParameters 0.5.0, SymbolicNeuralNetworks 0.9.0, GeometricBase 0.15.0, GeometricEquations 0.21.5, GeometricSolutions 0.6.6) all resolve, as the forcedeps step shows.
• closingIssuesReferences is empty. The body writes "relates to #145".

Not checked

• Whether finding 1 always happens on Windows or only sometimes. There is no Windows host here.
• Whether SymbolicPullback's pullback output goes through _get_contents unchanged (finding 3).
• The full test suite, the docs build and the doctests. CI ran them, and they are green apart from the job above.
• Training on a GPU.

[michakraus]

Resolution of the review findings

The review was on 826ae32. The branch is now at 20da228. Each finding by its number:

#	severity	resolution
1	blocker	Fixed, then superseded. 3529019 widened the Stiefel comparison (64 × 64 leaves, 256 B). 93409b0 removes NetworkOptimizer (finding 6), after which that comparison measured one step against itself; it is gone. The Euclidean == 0 bound and the @inferred checks stay in test/optimizers/optimization_step.jl.
2	bug	Fixed in 1d661a7 and kept in 93409b0: train! loops over the epochs instead of map(1:n_epochs) do. First train! of the review's transformer, fresh process, Julia 1.13.1: about 97 s before, 12.2–21.9 s after (three runs, load average 13–39), same history. The four scripts that timed out run locally in 95.9, 68.2, 62.9 and 69.2 s.
3	quality	Fixed in 93409b0: train!(…; loss, pullback = ZygotePullback(loss), show_progress). A new testset trains a Hamiltonian network with SymbolicPullback(arch) and checks that one epoch moves every leaf by -η times the symbolic gradient. Adding it found a defect older than this PR: the symbolic gradient is the gradient of the per-sample sum of losses, not of HNNLoss (finite difference 0.002696, Zygote 0.002696, symbolic 0.004388; the same on main). Filed as #343 and recorded as K13. The comparison with Zygote is @test_broken.
4	quality	Fixed in 93409b0: test/optimizers/step_size_element_type.jl and the PrecomputedGradient testset of gml_gradient_dispatch.jl are deleted here; they test GeometricOptimizers and move to GeometricOptimizers' suite in a pull request of their own. The _get_contents testset that remains moves to test/pullbacks/get_contents.jl, so no file name refers to _GMLGradient.
5	quality	Fixed in 93409b0: all_zero_leaf.jl asserts every claim of its header and stops with an AssertionError at the first that fails; all hold in Float32 and Float64. The closure is now read_only, not one.
6	quality	Fixed in 93409b0: NetworkOptimizer is removed. train! takes a TrainingOptimizer, which GML re-exports; callers write TrainingOptimizer(nn.params; …). The re-exports live in src/optimizers/optimizer.jl. No NetworkOptimizer is left in src, test, docs, scripts, README.md, KNOWN_ISSUES.md or CHANGELOG.md.
7	nit	Fixed in 93409b0: _default_loss has a fallback that raises an ArgumentError that names the architecture and the batch and tells the caller to pass loss, and a test checks it.
8	nit	Fixed in adac904: the four signatures are joined, and the docstrings say strictly lower- and upper-triangular.
9	nit	Not fixed, no public API exists. The tests read opt.state.iterations, opt.cache, opt.method and opt.linesearch. In GeometricOptimizers 0.9.0, iteration_number, solution, section and the moment accessors are not public, and only step_size is; so the field reads stay. Without the wrapper they are one level shorter (opt.state, not opt.optimizer.state).
10	nit	Fixed. The 400-word CHANGELOG bullet is split into five (93409b0). The GPU paragraph's checks are run again on an Apple M4 Max with Metal.jl 1.11.1 and GeometricOptimizers 0.9.0: all ran, and train! now trains a GSympNet on the device (20da228). ClassificationTransformer was not tried, and the paragraph says so.
11	nit	Fixed: the PR body is rewritten; its commit list matches the branch, and its suite totals are for the head 20da228.

Full suite on 20da228, Julia 1.13.1: Pass 4928, Fail 0, Error 0, Broken 6 (JET 3, #325; Metal 1, no GPU in the sandbox; SymbolicPullback 2, #343). docs/check_references.jl: 0 unresolved references.