API Reference

GeometricOptimizers

Optimizers on manifolds, and the structured matrices and manifolds they update.

The package supports real element types only. The structured matrices, the retractions and the symplectic Gram-Schmidt process assume a real field: a complex element type is not rejected, and some operations give a wrong answer for it.

Adam and the methods that build their direction from a pair of bias-corrected moment averages as it does. What the union selects for is that shared character and not a shared implementation: the direction is a moving average that is deliberately allowed not to descend on an individual step, which is what default_linesearch's fixed Static and the ensure_descent! exemption in FirstOrderMethodWithState are about.

The implementation is shared by only some of them, so do not read the union as a promise about the cache:

• AdamWithEuclideanDecay differs only in how the direction is finished off, so it reuses AdamCache and AdamState outright and OptimizerState(AdamWithEuclideanDecay(), x) returns an AdamState.
• ScalarMomentAdam has neither, because its second moment is a number rather than an element of 𝔤ʰᵒʳ and its recursion is therefore not Adam's. It carries ScalarMomentAdamCache and ScalarMomentAdamState of its own.
• DirectScalarMomentAdam and FactoredStiefelAdam step in the tangent space of a StiefelManifold and not in 𝔤ʰᵒʳ, and carry caches and states of their own.

CURVATURE_TOLERANCE

How positive $\delta^T\gamma$ has to be, relative to $\|\delta\|\|\gamma\|$, for the secant pair to be used. See curvature_is_usable.

DEFAULT_STEP_CEILING

How far along its direction a manifold solve may step, in multiples of $2\pi$: the c of step_αmax(c, δ), and so the params.αmax that solver_step! gives the line search through linesearch_parameters. Set it per solve with Optimizer(x, F; step_ceiling = …); Inf switches the ceiling off.

Why the solver sets a ceiling

A line search bounds its step by the merit: it stops when $\varphi$ stops falling. On a compact manifold $\varphi$ is bounded, so a step many turns too long can decrease the merit, and no test of the search rejects it. The retraction of such a lift leaves the manifold. The bound that exists is not a property of $\varphi$: it is the $2\pi$ of a rotation divided by $\|\delta\|$, and it changes at every solver step. The DEFAULT_LINESEARCH_αmax of a line search method is a backstop for every problem, not this bound. The manual section A bounded merit is not a bound on the step gives the measurements on the SVD problem.

Why 1

1 says "never step further along the direction than one full turn". That is the scale that the geometry gives, not a value tuned to one problem.

The three methods that step a StiefelManifold in its tangent space, with cayley_direct! and a vector transport: DirectMomentum, DirectScalarMomentAdam and FactoredStiefelAdam. Their caches, their states and the workspace a TrainingOptimizer builds for them hold $O(Mk + k^2)$ numbers.

FirstOrderMethod

The methods that build their direction from the gradient alone: GradientMethod and the FirstOrderMethodWithState. These are the methods a TrainingOptimizer takes, alone or two of them in a CompositeMethod, and the methods whose update!(cache, state, gradient, method, x) takes the method rather than a SimpleSolvers.Hessian, except the members of GeometricOptimizers.DirectStiefelMethod, which have no update!.

The methods whose update! needs the method rather than a SimpleSolvers.Hessian, because they carry state of their own out of which the direction is built — a momentum or a pair of moments. The (quasi-)Newton methods are the complement: their direction comes from the Hessian and the method object holds nothing.

GradientMethod has no Hessian and no state; with these methods it forms the FirstOrderMethods, whose update! takes the method. The three members of GeometricOptimizers.DirectStiefelMethod have no update!: the training step of TrainingOptimizer builds their direction from the method, the cache and the state.

ensure_descent! has a no-op method for these. It must not apply to them: a momentum and a moment average are deliberately allowed not to descend on an individual step, and the decay of AdamWithEuclideanDecay tilts the direction further away from the gradient still.

That exemption is about the direction before the search, and not about the restart after it — see linesearch_rejected. Declining to overrule the direction is not the same as taking the longest step along it, which is what a rejected search hands back.

OptimizerSolution

A type alias for the solution of an optimizer: an AbstractVector, a Manifold, or a NeuralNetworkParameters.NetworkParameters holding a whole set of network parameters.

A set of parameters enters this package as a container and never as a bare NamedTuple. That is the constraint every other alias in this file is written against, so it is worth saying why the obvious alternative is not available.

A parameter set has to be a type somebody owns. NamedTuple{S,<:Tuple{Vararg{AbstractArray{T}}}} picks out the same values, but it is an alias for Base.NamedTuple rather than a type of its own, and three things follow:

• a method taking one is a method on Base.NamedTuple. Unless the function is also this package's, that is type piracy — and it is global, changing what the function means for every keyed NamedTuple in any session that loads this package, directly or through a dependency;
• it cannot be told apart from any other NamedTuple alias in the same method table. It would overlap GlobalSectionNamedTuple and the bare NamedTuple that a nested section tree is written on, and on the overlap neither method of such a pair is more specific — an ambiguity rather than a choice. optimizers/named_tuple_wrapper.jl is where that would bite;
• its T binds through a Vararg bound on the values, which asserts "flat, and every leaf an AbstractArray{T}". NetworkParameters{T} asserts something else: one element type over leaves at any depth. A union of the two would say one thing about one argument shape and another about the other.

NetworkParameters{T} has none of those properties. It also binds T from a direct type parameter rather than through a Vararg bound on value types, which makes it the cheapest member of this union to intersect.

ScaledAlgorithm

The AbstractExponentialAlgorithms built as a small-argument kernel inside scaling and modified squaring, i.e. ScaledSquaring and NativePade. They share the θ field and the 𝔄! method below, and differ only in their _scaled_kernel!.

VectorStorageMatrix

The structured matrix types that keep their free parameters in one vector, reachable as parent(A): SkewSymMatrix, SymmetricMatrix and the two AbstractTriangulars.

Each of them stores $n(n\pm1)/2$ numbers behind an $n \times n$ interface, so none of the generic AbstractArray methods, which either broadcast or reshape, is usable for them:

• three of the four have no setindex! for a broadcast to write through. SkewSymMatrix and the two triangulars cannot represent an arbitrary matrix at all, so there is no entry to assign. SymmetricMatrix is the exception — it does have one (writing $[A]_{ij}$ writes $[A]_{ji}$ too, which is what a symmetric matrix means), so the broadcast would give the right answer for that one type, at twice the work and only as long as it stays the exception.
• similar has to preserve the type, because the optimizer caches allocate their scratch with it and then require every array to have the same type as the parameter. The AbstractArray fallback returns a dense Matrix.
• flattening has to round trip through the free parameters. A generic AbstractMatrix treatment reshapes the flattened vector back to $n \times n$, and $n(n\pm1)/2$ numbers do not reshape to that. NeuralNetworkParameters asks Base.parent instead, through freeparameters.

For every one of these types the free parameters are the coordinates the optimizer should work in, so each primitive is the corresponding operation on parent. _add!, _rmul! and _difference! live next to the rest of their family in optimizers/named_tuple_wrapper.jl, l2norm in optimizers/optimizer_status.jl; update_section! in global_sections/global_sections.jl splits on this alias for the same reason.

This is what lets a SymmetricMatrix or a triangular matrix be an optimizer parameter — which is what GeometricMachineLearning's SympNet, symplectic-attention and volume-preserving layers need, and what this package could not do before.

The alias lives in a file of its own because it is needed by global_sections.jl, which is included long before the primitives are.

Matrix(λY::GlobalSection)

Put λY into matrix form.

This is not recommended if speed is important!

Use apply_section and global_rep instead!

AbstractExponentialAlgorithm

Supertype of the algorithms that Geodesic can use to evaluate the matrix exponential.

A horizontal lift factors as $\bar{B} = B'(B'')^T$ into two $N\times{}2n$ matrices. Since $(B'(B'')^T)^k = B'X^{k-1}(B'')^T$ for $X=(B'')^TB'$, substituting into the exponential series gives

\[\exp(B'(B'')^T) = \mathbb{I} + B'\,\mathfrak{A}((B'')^TB')\,(B'')^T, \qquad \mathfrak{A}(X) = \sum_{n=1}^\infty \frac{X^{n-1}}{n!},\]

Computing $\mathfrak{A}(X)$ is therefore the central numerical task: it preserves the low-rank factorization and reduces the matrix function from $N\times{}N$ to $2n\times{}2n$. $\mathfrak{A}$ is the notation of this implementation and of [10]; the exponential-integrator literature writes the same function $\varphi_1$ [22, §2.1].

Four subtypes evaluate that $2n\times{}2n$ function, while ProjectedSkew bypasses it and exponentiates the lift in a basis of its range. They compute the same exponential map — so the one-parameter subgroup property Geodesic relies on holds for all of them — but differ in their approximation kernel, recovery strategy, numerical behaviour, cost, and backend requirements.

ScaledSquaring is the default, and the cheapest algorithm with no dense-LAPACK dependency. Reach for another only if you need what it gives you: NativePade an independent direct calculation, ProjectedSkew orthogonality that does not degrade with the size of the lift, AugmentedPade a dense-CPU reference that delegates its numerics to Base.exp. TaylorSeries is the unscaled series, kept as a baseline; it is not a usable retraction.

The Exponential Algorithms page derives all five, measures them against each other, and says which to choose in Choosing one. Its tables are recomputed on every documentation build, so they, and not a figure quoted in a docstring, are what to trust.

A new one has to supply 𝔄(X::AbstractMatrix, ::NewAlgorithm); geodesic and everything above it then follow. An algorithm that does not go through $\mathfrak{A}$ at all — ProjectedSkew is the one such — supplies geodesic(::AbstractLieAlgHorMatrix, ::NewAlgorithm) instead.

See GeometricOptimizers.𝔄 for the implementations.

AbstractLieAlgHorMatrix <: AbstractMatrix

AbstractLieAlgHorMatrix is a supertype for various horizontal components of Lie algebras. We usually call this $\mathfrak{g}^\mathrm{hor}$.

See StiefelLieAlgHorMatrix and GrassmannLieAlgHorMatrix for concrete examples.

AbstractOptimizerProblem

See SimpleSolvers.LinesearchProblem and OptimizerProblem.

AbstractRetraction

Supertype of the retractions, i.e. of the maps that turn an element of a horizontal Lie algebra back into a point of the manifold.

A retraction is how every step in this package is taken: an OptimizerMethod produces a direction $B$ in the horizontal component of $\mathfrak{g}$, and $\Lambda \mapsto \Lambda\cdot\mathrm{retract}(B)$ moves the global section — and with it the point — without ever leaving the manifold. An instance is passed to Optimizer as retraction = Cayley(); the default is Cayley.

Concrete subtypes are Geodesic and Cayley. A new one has to supply retraction(::NewRetraction, x); the callable form R(x) and everything in update_section! then follow.

For every manifold layer one has to specify a retraction that takes the layer and elements of the (global) tangent space.

AbstractTriangular

See StrictlyUpperTriangular and StrictlyLowerTriangular.

AbstractVectorTransport

A vector transport on the Stiefel manifold: it moves a vector $\Xi$ at $Y$ to the tangent space at $Y_\mathrm{new}$, the Cayley retraction of a step $\Delta$ at $Y$. See transport!, ProjectionTransport and DifferentiatedCayleyTransport.

Adam(; β₁, β₂, δ)

The Adam optimizer, with the defaults suggested in [31, page 301].

The cache consists of a first and a second moment, stored in bias-corrected form, i.e. updated in the $t$-th iteration (counted from $t = 1$) as

\[m_1 \gets \frac{\beta_1 - \beta_1^t}{1 - \beta_1^t}m_1 + \frac{1 - \beta_1}{1 - \beta_1^t}\nabla{}L,\]

\[m_2 \gets \frac{\beta_2 - \beta_2^t}{1 - \beta_2^t}m_2 + \frac{1 - \beta_2}{1 - \beta_2^t}\nabla{}L\odot\nabla{}L,\]

from which the direction is computed as $-m_1/(\sqrt{m_2} + \delta)$.

The coefficients carry no element type of the parameters: Optimizer and TrainingOptimizer convert them once, with change_precision, to the element type of the parameters they are given. Adam() optimizes Float32 and Float64 parameters alike.

AdamCache <: OptimizerCache

Cache for the gradient optimizer.

Fields

• x::Manifold: the solution,
• g: the gradient (for the manifold case this is in AbstractLieAlgHorMatrix form),
• δ: the direction,
• Δg: difference in gradients, needed for OptimizerStatus,
• g̃: scratch for latest_gradient; see GradientCache,
• g̃_is_current: whether g̃ is the gradient at x; see store_gradient!,
• section: the GlobalSection.

AdamState <: OptimizerState

State for the gradient optimizer.

AdamWithEuclideanDecay(; β₁, β₂, δ, λ)

Adam with the decoupled weight decay of [34], which is the Euclidean $\lambda{}x$ and therefore acts on the unconstrained weights only.

The moments are the ones of Adam — this method shares its cache and its state — and the decay is applied to the direction instead of to the gradient, i.e. the direction is

\[-\left(\frac{m_1}{\sqrt{m_2} + \delta} + \lambda{}x\right)\]

so that a step with learning rate $\eta$ (see below) is

\[x \gets x - \eta\frac{m_1}{\sqrt{m_2} + \delta} - \eta\lambda{}x.\]

That is what decoupled means: $\lambda{}x$ never enters $m_1$ or $m_2$, so the amount by which a weight is shrunk does not depend on the size of its gradient. Adding $\lambda{}x$ to the gradient instead — $L^2$ regularization, which is what plain Adam on a penalized objective does — is not the same thing, because the second moment then rescales the penalty away.

Why the name says Euclidean

$\lambda{}x$ is the gradient of $\frac{\lambda}{2}||x||^2$, which is constant on the StiefelManifold and the GrassmannManifold — both are compact, with $||Y||_F^2 = \mathrm{tr}(Y^TY) = n$ — so its Riemannian gradient $\mathtt{rgrad}(Y, \lambda{}Y)$ vanishes identically and the decay does nothing at all to a manifold weight. The method therefore decays the ordinary arrays of a NamedTuple of parameters and leaves the Manifolds in it alone — which is the case it exists for, a network that keeps Stiefel weights next to unconstrained ones — and on a bare Manifold it is Adam, for every $\lambda$. Passing a nonzero λ together with parameters that are entirely manifolds is therefore warned about rather than silently ignored.

The derivation and the Grassmann case are in the Weight Decay on Manifolds page. Whether a Riemannian decay is called for is an open question.

Arguments

As for Adam, the coefficients carry no element type of the parameters, and the learning rate $\eta$ is not stored here: it is the α of the line search, i.e. linesearch = η, and default_linesearch is the default. λ is hence multiplied by $\eta$, exactly as in [34], where the two are decoupled from each other but the decay is still scaled by the schedule.

AugmentedPade <: AbstractExponentialAlgorithm

Evaluate $\mathfrak{A}$ as a block of a larger ordinary exponential.

This is an independent dense-CPU reference for testing the direct implementations, not the normal choice for a retraction.

\[\exp\begin{pmatrix} X & \mathbb{I} \\ \mathbb{O} & \mathbb{O} \end{pmatrix} = \begin{pmatrix} \exp(X) & \mathfrak{A}(X) \\ \mathbb{O} & \mathbb{I} \end{pmatrix}\]

so one call to Base.exp on a $4n\times{}4n$ matrix returns $\mathfrak{A}(X)$ in its upper-right block. That hands the numerics to Julia's own dense matrix exponential, which uses a Padé-based scaling-and-squaring algorithm [19, 20], at the cost of exponentiating a matrix four times the size and discarding three quarters of it.

Its value is that it introduces no package-specific approximation, which is what makes it the reference the other algorithms are measured against, both in the test suite and in every accuracy table on the Exponential Algorithms page. Its accuracy is the same order as ScaledSquaring's; what it costs is the work it throws away.

See 4. AugmentedPade for the measurements and AbstractExponentialAlgorithm for the alternatives.

BFGS

The Broyden–Fletcher–Goldfarb–Shanno method, from [33], and the default algorithm of Optimizer. The direction is $-Q\nabla{}f(x)$ for an approximate inverse Hessian $Q$, updated from the secant pair $\delta = x^{(k)} - x^{(k-1)}$, $\gamma = \nabla{}f^{(k)} - \nabla{}f^{(k-1)}$ by

\[Q \gets Q - \frac{\delta\gamma^TQ + Q\gamma\delta^T - \left(1 + \frac{\gamma^TQ\gamma}{\delta^T\gamma}\right)\delta\delta^T}{\delta^T\gamma}\]

Its state is BFGSState and the update itself lives in QuasiNewtonCache, which skips the update when the secant pair fails curvature_is_usable — the condition that keeps $Q$ positive definite. restart!(::BFGSState) discards $Q$ when a line search reports that it could not decrease the merit.

Unlike Newton this builds its curvature up over the iterations rather than evaluating a Hessian, so it takes more of them; unlike DFP it produces a direction already scaled like a Newton step, and accepts $\alpha = 1$ on most of them. Both run on an AbstractVector, a NamedTuple of parameters and a bare Manifold alike: $Q$ is sized by the intrinsic dimension and the secant pair is taken in the horizontal lift, so nothing in the update needs a vector-valued point.

Takes no parameters: like every OptimizerMethod it only produces a direction, and how far the optimizer travels along it is the line search's business.

BFGSState <: OptimizerState

The OptimizerState corresponding to the BFGS method.

Keys

• x: the current iterate,
• x̄: the previous iterate,
• s: stores the previous direction. This needs to be stored in addition to the previous solution because of the manifold case.
• ḡ: the gradient the last step was built from,
• f: the objective at x,
• f̄: the objective at x̄,
• Q

The unbarred fields are the current iterate's and the barred ones the previous iterate's, as update! maintains them, and the accessors are the OptimizerState defaults: value reads f and is exported, GeometricOptimizers.solution, GeometricOptimizers.previous_solution and GeometricOptimizers.previous_value are internal.

CanonicalMetric <: StiefelMetric

The canonical metric of the Stiefel manifold, $g_Y(\Delta_1, \Delta_2) = \mathrm{Tr}(\Delta_1^T(\mathbb{I} - \frac{1}{2}YY^T)\Delta_2)$. The two-argument form of rgrad and the three-argument form of metric use it.

Cayley <: AbstractRetraction

The Cayley transform as a retraction, and the default one.

\[\mathrm{Cayley}(B) = \left(\mathbb{I} - \frac{1}{2}B\right)^{-1}\left(\mathbb{I} + \frac{1}{2}B\right)\]

For $B$ in the horizontal component this maps into the Lie group, so the retracted point stays on the manifold to round-off. cayley never forms the $N\times{}N$ inverse: it factors $B = B'(B'')^T$ into two $N\times{}2n$ matrices with lift_factors and inverts a $2n\times{}2n$ matrix instead.

Examples

using GeometricOptimizers
using GeometricOptimizers: Cayley, check

Y = rand(StiefelManifold, 5, 3)
B = GeometricOptimizers.global_rep(GlobalSection(Y), rand(5, 3))
check(Cayley()(B)) < 1e-14        # the retracted element is still on the manifold

# output

true

See cayley for the implementation and Geodesic for the alternative.

CompositeMethod(; manifold, array)

Two FirstOrderMethods for a parameter set whose leaves are not all of one kind: manifold steps every Manifold leaf, and array every other leaf.

ScalarMomentAdam is the case that needs it: it steps a single StiefelManifold, because a scalar second moment pooled across a set means nothing. So a network with Stiefel weights beside ordinary Matrix and Vector leaves has no single method that covers it:

CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam())

The split is by Manifold and not by StiefelManifold, so a GrassmannManifold leaf under ScalarMomentAdam is a MethodError when its cache is built. Any other method, or a composite, in either place is a MethodError at construction.

A composite is for TrainingOptimizer on a parameter set. Every leaf gets a cache and a state of its own method, in the element type of the set, and moves exactly as it would under a TrainingOptimizer of its own; nothing is pooled across leaves. The leaves share the step size of each step. A composite is not an OptimizerMethod, so Optimizer and solve!, whose line search finds one step length for the direction of the whole set, do not take it.

DFP

The Davidon–Fletcher–Powell method, from [33]: the direction is $-Q\nabla{}f(x)$ for an approximate inverse Hessian $Q$, updated from the secant pair $\delta = x^{(k)} - x^{(k-1)}$, $\gamma = \nabla{}f^{(k)} - \nabla{}f^{(k-1)}$ by

\[Q \gets Q - \frac{Q\gamma\gamma^TQ}{\gamma^TQ\gamma} + \frac{\delta\delta^T}{\delta^T\gamma}\]

(equation 6.15 there). Its state is BFGSState, under the alias DFPState, and the update itself lives in QuasiNewtonCache, which symmetrizes $Q\gamma\gamma^TQ$ explicitly and skips the update when the secant pair fails curvature_is_usable.

Where BFGS is the better-conditioned choice and the default, DFP is here because it is the other classical member of the family and because the two differ in exactly the way that shows what the line search is doing: DFP's direction is systematically under-scaled, so it wants a median $\alpha$ well above 1 and is starved by a search that cannot exceed it. See default_linesearch: DFP converges under the default expanding Backtracking, and StrongWolfe(T; c₂ = 0.1) is the explicit choice for a DFP-heavy workload.

Takes no parameters: like every OptimizerMethod it only produces a direction, and how far the optimizer travels along it is the line search's business.

DFPState

This is equivalent to BFGSState.

DecayingStatic(; η₁, η₂, n)

A SimpleSolvers.LinesearchMethod that takes no search at all, like SimpleSolvers.Static, but whose step size decays geometrically with the iteration number:

\[\alpha(t) = \gamma^t\eta_1, \qquad \gamma = \exp(\log(\eta_2/\eta_1)/n),\]

so that $\alpha(0) = \eta_1$ and $\alpha(n) = \eta_2$. It keeps decaying past n; the decay is not floored at $\eta_2$, because a floor is exactly what stops the iteration from converging (see below).

$\alpha(0)$ is a limit of the schedule and not a step a solve takes: solve! calls increase_iteration_number! before solver_step!, so the first step is $\alpha(1) = \gamma\eta_1$, the $k$-th is $\alpha(k)$, and the horizon is reached exactly at iteration n rather than one before it. $\eta_1$ is therefore an upper bound on the step rather than the first one — by a factor of $\gamma$, which for the defaults is 0.9908. The offset is deliberate: it is also how GeometricMachineLearning's AdamOptimizerWithDecay counts, which is what makes AdamOptimizerWithDecay reproduce it step for step (test/optimizers/decaying_static.jl asserts this against a live solve).

Why this exists

Adam produces a direction of magnitude $\approx{}1$ per component whatever the gradient is — that scale-freeness is the point of it. With a constant step $\alpha$ the iterate therefore does not converge to the minimizer but circles it at a distance of order $\alpha$, so no gradient-based stopping criterion is ever met and the solve runs to max_iterations. A step size that goes to zero is what turns it back into a convergent iteration.

The step size is the business of the line search, not of the OptimizerMethod, so the schedule is a LinesearchMethod; it composes with GradientMethod and MomentumMethod as well.

Examples

ls = DecayingStatic(; η₁ = 1e-2, η₂ = 1e-6, n = 1000)

# output

DecayingStatic from α = 0.01 to α = 1.0e-6 over 1000 iterations.

The schedule carries no element type of the parameters: Optimizer and TrainingOptimizer convert it once, with change_precision, which computes $\gamma$ in the element type of the parameters.

DifferentiatedCayleyTransport <: AbstractVectorTransport

The differential of the Cayley retraction at $\Delta$ in the direction $\Xi$:

\[\mathcal{T}_\Delta(\Xi) = (\mathbb{I}_M - \tfrac{1}{2}W(\Delta))^{-1} W(P_Y\Xi) (\mathbb{I}_M - \tfrac{1}{2}W(\Delta))^{-1} Y, \qquad P_Y\Xi = \Xi - \tfrac{1}{2}Y(Y^T\Xi + \Xi^TY),\]

with $W$ the generator of cayley_direct!. It is $\frac{\mathrm{d}}{\mathrm{d}s}\big|_{s=0}\mathtt{cayley\_direct}(Y, \Delta + s\Xi)$. Any $\Xi$ is accepted, and its part off the tangent space at $Y$ is dropped; at $\Delta = 0$ the result is $P_Y\Xi$.

It is computed in $O(Mk^2)$, with the matrix $M_{2k}$ of the retraction and the factors $U_\Xi V_\Xi = W(P_Y\Xi)$, as

\[\bigl(U_\Xi + \tfrac{1}{2}UM_{2k}^{-1}(VU_\Xi)\bigr)(V_\Xi W_+), \qquad W_+ = (\mathbb{I}_M - \tfrac{1}{2}W(\Delta))^{-1}Y = Y + \tfrac{1}{2}UM_{2k}^{-1}VY.\]

$W_+$ comes from the factorization of the step, so the result does not read the entries of $Y_\mathrm{new}$. Its solve is the lu! of cayley_direct!, with the same backend limit.

DirectCayleyWorkspace(Y::StiefelManifold)

The buffers of cayley_direct! and transport! at a point of the shape of Y.

For $Y \in St(k, M)$ it holds four $M\times{}2k$ matrices, one $M\times{}k$ matrix, two $k\times{}k$ matrices, two $2k\times{}2k$ matrices and three $2k\times{}k$ matrices, all on the backend of Y.A: $O(Mk + k^2)$ storage and no $M\times{}M$ array. A workspace serves every point of its shape and element type.

DirectMomentum(; α = DEFAULT_MOMENTUM_α, metric = CanonicalMetric(), transport = ProjectionTransport())

Riemannian gradient descent with momentum on a StiefelManifold, in the tangent space $T_Y\mathrm{St}(k, M)$ itself: no GlobalSection, no $M\times{}M$ matrix and no orthogonal complement of $Y$.

With $G_\mathcal{R} = \mathtt{rgrad}(Y, G, \mathtt{metric})$ for the Euclidean gradient $G$ and the step size $\eta$ of the line search, a step is

\[p \gets \alpha\,\mathcal{T}(p) + G_\mathcal{R}, \qquad \Delta = -\eta{}p, \qquad Y \gets \mathtt{cayley\_direct}(Y, \Delta),\]

where $\mathcal{T}$ is the vector transport transport (ProjectionTransport or DifferentiatedCayleyTransport) that carries the momentum of the previous step to the tangent space at $Y$. metric is CanonicalMetric or EuclideanMetric.

α satisfies $0 \le \alpha < 1$; another value is an ArgumentError. As for MomentumMethod, it carries no element type of the parameters, and the step size is not stored here: TrainingOptimizer converts α once, with change_precision, and takes $\eta$ as its linesearch. Only a single StiefelManifold is a leaf of this method, and only TrainingOptimizer steps it, with the retraction cayley_direct!.

The step has no cap. A finite gradient whose step has $\lVert\Delta\rVert_F$ above about $\sqrt{\mathrm{floatmax}(T)}$, $1.8\cdot10^{19}$ in Float32, makes the factorization of cayley_direct! overflow, and the step raises ArgumentError: matrix contains Infs or NaNs, as a step of MomentumMethod does. The momentum $p$ and the point are written only after the retraction, with the iteration number, so they stay as they were.

DirectMomentumState <: OptimizerState

The state of DirectMomentum: the iteration number, the point x and the momentum p, a tangent vector at x.

DirectScalarMomentAdam(; β₁ = 9.0e-1, β₂ = 9.9e-1, δ = 1.0e-8, α_max = 1.0, metric = CanonicalMetric(), transport = ProjectionTransport())

Cayley ADAM of [11], with a scalar second moment, in the tangent space $T_Y\mathrm{St}(k, M)$ of a StiefelManifold itself. Where ScalarMomentAdam steps in $\mathfrak{g}^\mathrm{hor}$ through a GlobalSection, this method needs no $M\times{}M$ matrix and no orthogonal complement of $Y$.

With $G_\mathcal{R} = \mathtt{rgrad}(Y, G, \mathtt{metric})$, the step size $\alpha$ of the line search and the iteration $t$, the moments are stored in bias-corrected form, as Adam's are:

\[B_1 \gets \frac{\beta_1 - \beta_1^t}{1 - \beta_1^t}\mathcal{T}(B_1) + \frac{1 - \beta_1}{1 - \beta_1^t}G_\mathcal{R}, \qquad v \gets \frac{\beta_2 - \beta_2^t}{1 - \beta_2^t}v + \frac{1 - \beta_2}{1 - \beta_2^t}\lVert{}G_\mathcal{R}\rVert_F^2,\]

and the step is

\[D = \frac{B_1}{\sqrt{v} + \delta}, \qquad \Delta = -\min\Bigl(\alpha, \frac{\alpha_\mathrm{max}}{\lVert{}D\rVert_F + \delta}\Bigr)D, \qquad Y \gets \mathtt{cayley\_direct}(Y, \Delta).\]

$\mathcal{T}$ is the vector transport transport that carries $B_1$ from the previous point. The $\delta$ is outside the square root, as in Adam and FactoredStiefelAdam; ScalarMomentAdam puts it inside. The cap $\alpha_\mathrm{max}$ bounds the Frobenius norm of the step.

$0 \le \beta_1, \beta_2 < 1$, $\delta > 0$ and $\alpha_\mathrm{max} > 0$; another value is an ArgumentError, also where the conversion of a coefficient to the element type of the parameters takes it out of its range. The coefficients carry no element type of the parameters, and the learning rate is the linesearch of TrainingOptimizer, which is the only optimizer that steps this method. Only a single StiefelManifold is a leaf of it.

$\lVert{}D\rVert_F \approx 1$, so at the step size of Adam, whose direction has $\lVert{}D\rVert_F \approx \sqrt{Mk}$, a step is about $\sqrt{Mk}$ times shorter, as for ScalarMomentAdam. Scale the learning rate before comparing this method with Adam.

The second moment holds the squares of the gradient in the element type T of the parameters, so a gradient with $\lVert{}G_\mathcal{R}\rVert_F$ above about $\sqrt{\mathrm{floatmax}(T)}$, $1.8\cdot10^{19}$ in Float32, makes it infinite. For $\beta_2 > 0$ it stays infinite, so that step and every later step are zero.

DirectScalarMomentAdamState <: OptimizerState

The state of DirectScalarMomentAdam: the iteration number, the point x, the first moment B₁, a tangent vector at x, and the scalar second moment v. Both moments are stored bias-corrected.

DirectStiefelCache <: OptimizerCache

The cache of DirectMomentum and DirectScalarMomentAdam:

• x::StiefelManifold: the point the step is taken from,
• g: the Riemannian gradient at x, an M × k matrix,
• δ: the step $\Delta$, a tangent vector at x.

EuclideanMetric <: StiefelMetric

The metric that the Stiefel manifold inherits from $\mathbb{R}^{N\times{}n}$, $g_Y(\Delta_1, \Delta_2) = \mathrm{Tr}(\Delta_1^T\Delta_2)$.

EventLog(; phases=nothing)

Record optimizer phase notifications as (phase, event) tuples in events.

By default every phase is recorded. Pass a phase name or an iterable of phase names as phases to retain only those phases. Call empty!(log) between runs to reuse an existing log.

Examples

julia> recorder = EventLog();

julia> observe_optimizer_phase(recorder, :gradient) do
           nothing
       end

julia> recorder.events
2-element Vector{Tuple{Symbol, Symbol}}:
 (:gradient, :enter)
 (:gradient, :exit)

FactoredStiefelAdam(; β₁ = 9.0e-1, β₂ = 9.9e-1, δ = 1.0e-8, α_max = 1.0, metric = CanonicalMetric(), transport = ProjectionTransport())

Adam on a StiefelManifold with one adaptive learning rate per mode, in the tangent space $T_Y\mathrm{St}(k, M)$ itself: no GlobalSection, no $M\times{}M$ matrix and no orthogonal complement of $Y$.

A tangent vector $X$ at $Y$ splits into a vertical part $P_V(Y)X = Y\cdot\tfrac{1}{2}(Y^TX - X^TY)$ and a horizontal part $P_H(Y)X = X - Y(Y^TX)$. The first moment $B_1$ is a tangent vector. The second moment is a $k\times{}k$ matrix $V_W$ for the vertical part, and a vector $v_K \in \mathbb{R}^k$ of the squared column norms of the horizontal part. With $G_\mathcal{R} = \mathtt{rgrad}(Y, G, \mathtt{metric})$, the step size $\alpha$ of the line search, the iteration $t$ and the weights $c_{i1} = (\beta_i - \beta_i^t)/(1 - \beta_i^t)$, $c_{i2} = (1 - \beta_i)/(1 - \beta_i^t)$ of the bias-corrected form of Adam, a step is

\[\begin{aligned} B_1 &\gets c_{11}\mathcal{T}_s(B_1) + c_{12}G_\mathcal{R}, & A &= \tfrac{1}{2}(Y^TG_\mathcal{R} - G_\mathcal{R}^TY), & K &= P_H(Y)G_\mathcal{R}, \\ V_W &\gets c_{21}V_W + c_{22}A\odot{}A, & (v_K)_j &\gets c_{21}(v_K)_j + c_{22}\lVert{}K_{:,j}\rVert_2^2, & \varepsilon &= \sqrt{\mathrm{eps}(T)} \Bigl(\textstyle\sum_{ij}(V_W)_{ij} + \sum_j(v_K)_j\Bigr)^{1/2}, \\ \tilde{A} &= \tfrac{1}{2}(Y^TB_1 - B_1^TY) \oslash (\sqrt{V_W} + \delta + \varepsilon), & \tilde{K}_{:,j} &= \frac{(P_H(Y)B_1)_{:,j}}{\sqrt{(v_K)_j} + \delta + \varepsilon}, \end{aligned}\]

\[D = Y\tilde{A} + \tilde{K}, \qquad \Delta = -\min\Bigl(\alpha, \frac{\alpha_\mathrm{max}}{\lVert{}D\rVert_F + \delta}\Bigr)D, \qquad Y \gets \mathtt{cayley\_direct}(Y, \Delta).\]

$\mathcal{T}_s$ carries $B_1$ from the previous point $Y_\mathrm{prev}$ with the vector transport $\mathcal{T}$ of transport, one part at a time, and puts each part back into its own subspace at $Y$:

\[\mathcal{T}_s(B_1) = P_V(Y)\mathcal{T}(P_V(Y_\mathrm{prev})B_1) + P_H(Y)\mathcal{T}(P_H(Y_\mathrm{prev})B_1).\]

$V_W$ is symmetric, so $\tilde{A}$ is skew-symmetric, and every column of $\tilde{K}$ is orthogonal to $Y$: $\Delta$ is a tangent vector with no projection after it. Each $P_H$ is two passes of $K \gets K - Y(Y^TK)$: one pass leaves a part of the order of $\mathrm{eps}(T)\lVert{}X\rVert$ along $Y$, and the scaling of a column with a small norm would make it large. A part of $B_1$ whose own second moment is zero holds only round-off; the floor $\varepsilon$, relative to both second moments and with $T$ the element type of the parameters, keeps it small, and the split transport keeps the one part out of the other. For $G = 0$, $\varepsilon = 0$. Every operation commutes with $Y \mapsto QY$, $G \mapsto QG$ for $Q \in O(M)$, so the iterates are left $O(M)$-equivariant.

The constraints on the coefficients, their element type and the learning rate are those of DirectScalarMomentAdam. Only a single StiefelManifold is a leaf of this method, and only TrainingOptimizer steps it.

Each off-diagonal entry of $\tilde{A}$ and each column of $\tilde{K}$ has a magnitude of about $1$, so $\lVert{}D\rVert_F \approx k$. At the step size of Adam, whose direction has $\lVert{}D\rVert_F \approx \sqrt{Mk}$, a step is about $\sqrt{M/k}$ times shorter.

The second moments hold the squares of the entries of $A$ and of the column norms of $K$ in T, so an entry of $A$ or a column norm of $K$ above about $\sqrt{\mathrm{floatmax}(T)}$, $1.8\cdot10^{19}$ in Float32, makes one of them infinite, and $\varepsilon$ and every denominator with it. For $\beta_2 > 0$ it stays infinite, so that step and every later step are zero. A gradient whose norm is above that limit while every entry of $A$ and every column norm of $K$ is below it does not stop the method: $\varepsilon$ is infinite at most while the sum of the second moments overflows.

FactoredStiefelAdamCache <: OptimizerCache

The cache of FactoredStiefelAdam:

• x::StiefelManifold: the point the step is taken from,
• g: the Riemannian gradient at x, an M × k matrix,
• δ: the step $\Delta$, a tangent vector at x,
• S and A: a $k\times{}k$ product $Y^TX$, and the vertical part $Y^TG_\mathcal{R}$, a skew-symmetric $k\times{}k$ matrix, which the transport overwrites after the step,
• Ã: the scaled vertical part $\tilde{A}$, a skew-symmetric $k\times{}k$ matrix,
• K: the horizontal part $G_\mathcal{R} - Y(Y^TG_\mathcal{R})$, an M × k matrix,
• n² and d: the squared column norms of K, and the factors that scale the columns of the horizontal part of the first moment, two vectors of length $k$,
• A₁, V_W and v_K: the vertical part of the first moment and the two second moments of the step being taken. The state takes them after the retraction.

FactoredStiefelAdamState <: OptimizerState

The state of FactoredStiefelAdam: the iteration number, the point x, the first moment $B_1 = YA_1 + K_1$, a tangent vector at x, in its two parts, the vertical part A₁ = ½(YᵀB₁ - B₁ᵀY), a skew-symmetric $k\times{}k$ matrix, and the horizontal part K₁ = B₁ - Y(YᵀB₁), an M × k matrix, the vertical second moment V_W, a symmetric $k\times{}k$ matrix, and the horizontal second moment v_K, a vector of length $k$. The moments are stored bias-corrected.

Geodesic(algorithm = ScaledSquaring()) <: AbstractRetraction

The exponential map as a retraction, i.e. the true geodesic of the manifold.

\[\mathrm{Geodesic}(B) = \exp(B)\]

Because this is the matrix exponential, $\alpha \mapsto \mathrm{Geodesic}(\alpha{}B)$ is a one-parameter subgroup: it follows the geodesic through the point in the direction $B$ exactly, and $\mathrm{Geodesic}((\alpha+\beta)B) = \mathrm{Geodesic}(\alpha{}B)\mathrm{Geodesic}(\beta{}B)$. That is the property Cayley lacks, and the reason a derivative-based line search is exact here.

geodesic exploits the sparsity of a horizontal lift rather than exponentiating the full $N\times{}N$ matrix: the only matrix function it evaluates is on a $2n\times{}2n$ argument. cayley is built the same way, so the two cost the same order, $O(N^2n)$; see the note on Cayley for the measured difference.

algorithm selects how the exponential is evaluated. All of them compute the same map — the choice is one of accuracy at a large lift, cost, and backend support — and the default ScaledSquaring is the one to use unless there is a reason not to. See AbstractExponentialAlgorithm for the comparison.

Examples

using GeometricOptimizers
using GeometricOptimizers: Geodesic, check

Y = rand(StiefelManifold, 5, 3)
B = GeometricOptimizers.global_rep(GlobalSection(Y), rand(5, 3))
check(Geodesic()(B)) < 1e-14      # the retracted element is still on the manifold

# output

true

A large lift is where the algorithms part company, and the reason the default is what it is:

using GeometricOptimizers
using GeometricOptimizers: Geodesic, ScaledSquaring, TaylorSeries, check
import Random
Random.seed!(1234)

B = 60 * rand(StiefelLieAlgHorMatrix, 20, 3)      # ‖B̄‖ ≈ 393

check(Geodesic()(B)) < 1e-12, check(Geodesic(TaylorSeries())(B)) < 1e-12

# output

(true, false)

See geodesic for the implementation and Cayley for the rational alternative, and Exponential Algorithms for the choice of algorithm.

GlobalSection(Y)

Construct a global section for Y.

A global section $\lambda$ is a mapping from a homogeneous space $\mathcal{M}$ to the corresponding Lie group $G$ such that

\[\lambda(Y)E = Y,\]

Also see apply_section and global_rep.

Implementation

For an implementation of GlobalSection for a custom array (especially manifolds), the function global_section has to be generalized.

GradientCache <: OptimizerCache

Cache for the gradient optimizer.

Fields

• x::Manifold: the solution,
• g: the gradient (for the manifold case this is in AbstractLieAlgHorMatrix form),
• δ: the direction,
• Δg: difference in gradients,
• g̃: scratch for latest_gradient,
• g̃_is_current: whether g̃ is the gradient at x; see store_gradient!,
• section: the GlobalSection.

Implementation

The gradient descent algorithm.

GradientState <: OptimizerState

State for the gradient optimizer.

GrassmannLieAlgHorMatrix(B::AbstractMatrix, N::Integer, n::Integer)

Build an instance of GrassmannLieAlgHorMatrix based on an arbitrary matrix B of size $(N-n)\times{}n$.

GrassmannLieAlgHorMatrix is the horizontal component of the Lie algebra of skew-symmetric matrices (with respect to the canonical metric).

Extended help

The projection here is: $\pi:S \to SE/\sim$ where

\[E = \begin{bmatrix} \mathbb{I}_{n} \\ \mathbb{O}_{(N-n)\times{}n} \end{bmatrix},\]

and the equivalence relation is

\[V_1 \sim V_2 \iff \exists A\in\mathcal{S}_\mathrm{skew}(n) \text{ such that } V_2 = V_1 + \begin{bmatrix} A \\ \mathbb{O} \end{bmatrix}\]

An element of GrassmannLieAlgMatrix takes the form:

\[\begin{pmatrix} \bar{\mathbb{O}} & B^T \\ B & \mathbb{O} \end{pmatrix},\]

where $\bar{\mathbb{O}}\in\mathbb{R}^{n\times{}n}$ and $\mathbb{O}\in\mathbb{R}^{(N - n)\times(N-n)}.$

GrassmannLieAlgHorMatrix(D::AbstractMatrix, n::Integer)

Take a big matrix as input and build an instance of GrassmannLieAlgHorMatrix.

The integer $N$ in $Gr(n, N)$ here is the number of rows of D.

Extended help

If the constructor is called with a big $N\times{}N$ matrix, then the projection is performed the following way:

\[\begin{pmatrix} A & B_1 \\ B_2 & D \end{pmatrix} \mapsto \begin{pmatrix} \bar{\mathbb{O}} & -B_2^T \\ B_2 & \mathbb{O} \end{pmatrix}.\]

This can also be seen as the operation:

\[D \mapsto \Omega(E, DE - EE^TDE),\]

where $\Omega$ is the horizontal lift GeometricOptimizers.Ω.

GrassmannManifold <: Manifold

The GrassmannManifold is based on the StiefelManifold.

Manifold <: AbstractMatrix

A manifold in GeometricOptimizers is a subtype of AbstractMatrix. All manifolds are matrix manifolds and therefore stored as matrices. More details can be found in the docstrings for the StiefelManifold, the GrassmannManifold and the SymplecticStiefelManifold.

MomentumCache <: OptimizerCache

Cache for the gradient optimizer.

Fields

• x::Manifold: the solution,
• g: the gradient (for the manifold case this is in AbstractLieAlgHorMatrix form),
• δ: the direction,
• Δg: difference in gradients (used in OptimizerStatus),
• g̃: scratch for latest_gradient; see GradientCache for why it is separate from g, and note that this is the cache the separation was measured on,
• g̃_is_current: whether g̃ is the gradient at x; see store_gradient!,
• section: the GlobalSection.

MomentumMethod(; α)

The gradient descent algorithm with momentum, i.e. the heavy ball method.

Stores the momentum coefficient α. The momentum is accumulated as

\[ p \gets \alpha{}p + \nabla{}L\]

and the direction is $-p$.

α carries no element type of the parameters: Optimizer and TrainingOptimizer convert it once, with change_precision, to the element type of the parameters they are given.

MomentumState <: OptimizerState

State for the gradient optimizer.

NativePade(θ = 0.5) <: AbstractExponentialAlgorithm

Evaluate $\mathfrak{A}$ with a native degree-6 diagonal Padé approximant.

This scales and squares too. The name distinguishes the kernel and not the algorithm's relationship to scaling: NativePade and ScaledSquaring choose s by the same rule and undo it with the same $s$ applications of $W \mapsto 2W + WXW$, and differ only in what they evaluate at the scaled argument. The two are not alternatives.

Scaling controls the large-argument cancellation that makes TaylorSeries unreliable. At the resulting small argument there is still a choice of approximation kernel, and that choice is what this type makes. Where a degree-6 Taylor polynomial matches the series of $\mathfrak{A}$ through order 6, the $[6/6]$ rational function $p_6/q_6$ derived from the classical $[7/6]$ approximant of $\exp$ matches it through order 12 at the same polynomial degree. The denominator is applied by five Newton–Schulz steps [28] rather than by a dense solve, which is what keeps the algorithm free of LAPACK — matrix products and reductions only, so whether it runs on a given accelerator depends on that backend's support for those. That is what makes it the independent cross-check available where AugmentedPade, the dense-CPU reference, is not.

θ is a ceiling here, not a preference as it is for ScaledSquaring. The Newton–Schulz count is fixed at five, and their residual is below round-off only while the argument is small: the bound $\|Y\|_1 \leq 1/2$ gives $\|I - q_6(Y)\|_1 < 0.257$ and hence a final residual below $1.3\cdot10^{-19}$. Past $θ \approx 1$ the inverse it computes stops being one, and it fails silently — nothing in the result says so. The constructor therefore requires $0 < θ \leq 1/2$, where ScaledSquaring accepts any positive value: the two thresholds are not interchangeable. Lowering this one is safe and adds modified-squaring steps, which is why it is a parameter at all.

3. Padé approximation derives the coefficients, the Newton–Schulz iteration and the residual bound, gives the algorithm step by step, and measures the error at a larger $θ$ at build time in What a large $\theta$ costs NativePade. See AbstractExponentialAlgorithm for the comparison.

Newton

Newton's method: the direction solves $\nabla^2f(x)\delta = -\nabla{}f(x)$ with the exact Hessian, which SimpleSolvers.HessianAutodiff supplies.

Unlike BFGS and DFP this needs no approximation to build up, so it converges in few iterations, but it also inherits the Hessian's indefiniteness: where $\nabla^2f$ is not positive definite the direction ascends, and ensure_descent! substitutes the steepest-descent direction for it.

Runs on an AbstractVector and on nothing else. The exact Hessian is what makes it Newton's method: a bare Manifold has no Riemannian Hessian to build, and for a NeuralNetworkParameters.NetworkParameters the Hessian is not built over the flattening. Both are rejected at Optimizer and OptimizerState construction with a message saying so. BFGS and DFP take all three.

NewtonOptimizerCache <: OptimizerCache

Keys

• x: current iterate (this stores the guess called by the functions generated with linesearch_problem),
• Δx: direction of optimization step (difference between x and x̄); this is obtained by multiplying rhs with the inverse of the Hessian,
• g: gradient value (this stores the gradient associated with x),
• g̃: scratch for latest_gradient — the gradient at a line-search trial point, and at the accepted iterate once solver_step! has refreshed it; see GradientCache for why this is a field of its own and not an alias for g,
• g̃_is_current: whether g̃ is the gradient at x; see store_gradient!,
• Δg: gradient difference (difference between g̃ and g); this is used for computing the OptimizerStatus,
• rhs: the right hand side from which the update is computed,
• H: the Hessian matrix evaluated at x,

Also compare this to SimpleSolvers.NonlinearSolverCache.

NewtonState <: OptimizerState

The optimizer state is needed to update the Optimizer. This is different from OptimizerStatus and OptimizerResult which serve as diagnostic tools.

Keys

• x
• x̄
• g
• ḡ
• f
• f̄

The unbarred fields are the current iterate's and the barred ones the previous iterate's, as update! maintains them.

gradient and value read g and f and are exported. The other four accessors are internal: they are neither exported nor Base.ispublic, so they resolve only under a qualified name — GeometricOptimizers.solution for x, GeometricOptimizers.previous_solution for x̄, GeometricOptimizers.previous_gradient for ḡ and GeometricOptimizers.previous_value for f̄.

NoStepObserver()

The default no-op observer used by Optimizer and TrainingOptimizer.

Pass an EventLog, a PhaseTimer, or any callable observer to Optimizer(...; observer=observer) or TrainingOptimizer(...; observer=observer) to receive observer(phase, event) notifications, where event is :enter or :exit and phase is one of :gradient, :objective, or :retraction_application. Notifications are properly nested and an :exit is emitted from a finally block if the observed operation throws.

Callers that need an exclusive optimizer-state timer can bracket solver_step! and the following update! with observe_optimizer_phase(observer, :optimizer_state_direction). A stack-based observer can then pause the outer phase while any nested phase is active. The optimizer itself reports boundaries rather than choosing a clock; PhaseTimer reads the caller-selected clock and can synchronize a CUDA device immediately before each timestamp.

NoWorkspace()

What GeometricOptimizers.retraction_workspace returns where there is no manifold to retract on.

The vector-space methods of update_section! take a workspace argument and ignore it — the extended retraction on a vector space is addition, which allocates nothing to begin with — so this type is a marker and has no method of its own. It exists only because nothing cannot serve as that marker.

Implementation

nothing is the obvious spelling and is the one thing this may not be. A parameter set's workspace is a tree walked in lockstep with its section tree, and NeuralNetworkParameters.mapparameters! reads a nothing in that position as skip this leaf: the leaf function is then never called, so the section is never transported and the iterate never moves. That is silent — the solve runs to its iteration limit and reports the point it started from. A singleton says the same thing and is an ordinary value to the walk.

(That name in plain code and not an @extref, for the reason iterative_hessians.jl gives about the same choice: an @extref that the published inventory does not carry is a build error rather than a dead link, and mapparameters! is not one of the names this package already links.)

Optimizer

The optimizer that stores all the information needed for an optimization problem.

This problem can be solved by calling solve!(::AbstractVector, ::Optimizer).

Keys

• algorithm::OptimizerState,
• problem::OptimizerProblem,
• gradient::SimpleSolvers.Gradient,
• hessian::SimpleSolvers.Hessian,
• config::SimpleSolvers.Options,
• cache::OptimizerCache,
• linesearch::SimpleSolvers.Linesearch.

Examples

F(x) = sum(sin.(x) .^ 2)
x = ones(3)
algorithm = Newton()
state = OptimizerState(algorithm, x)
optimizer = Optimizer(x, F; algorithm = algorithm, linesearch = Bisection())

solve!(x, state, optimizer)
x

# output

3-element Vector{Float64}:
 0.0
 0.0
 0.0

The choice of line search does not change which minimum this one finds, only how close to it the stopping criterion lets the solve get — Backtracking stops at 1.4e-9 per component, where Bisection above happens to land on 0 exactly:

x = ones(3)
algorithm = Newton()
state = OptimizerState(algorithm, x)
optimizer = Optimizer(x, F; algorithm = algorithm, linesearch = Backtracking())

solve!(x, state, optimizer)
F(x) < 1e-15

# output

true

Optimizer(x, F; ∇F!, mode, algorithm, linesearch, retraction, observer, options_kwargs...)

Build an Optimizer for the objective F at the parameters x.

∇F! is the gradient of F. Without it, mode = :autodiff builds a forward-mode ForwardDiff gradient and mode = :finitediff a central-difference one. A vector and a Manifold take :autodiff when no mode is given. A NeuralNetworkParameters.NetworkParameters takes no default: without ∇F! and mode this throws an ArgumentError. Forward mode costs one pass of F per chunk of twelve parameters, and finite differences two evaluations of F per parameter, so either is for a small set.

linesearch is a SimpleSolvers.LinesearchMethod or a number, which is the fixed step size Static(η). algorithm, and a Static or DecayingStatic linesearch, are converted to the element type of x, so Optimizer(Float32[1, 2, 3], F; algorithm = Adam()) optimizes in Float32.

Implementation

OptimizerCache

See e.g. NewtonOptimizerCache and QuasiNewtonCache.

Extended help

OptimizerMethod <: SolverMethod

The OptimizerMethod is used in Optimizer and determines the algorithm that is used.

OptimizerProblem <: AbstractOptimizerProblem

Used in Optimizer. Also compare this to SimpleSolvers.NonlinearProblem.

Examples

julia> x = ones(3); F(x) = sum(sin.(x) .^ 2)
F (generic function with 1 method)

julia> OptimizerProblem(F, x)
OptimizerProblem{Float64, typeof(F), Missing, Missing}(F, missing, missing)

OptimizerResult

Serves as a diagnostic tool for the Optimizer and is the return argument of solve!.

Keys

• status::OptimizerStatus: current status of the optimization,
• x: solution,
• f: function value at solution,
• trace: one OptimizerTraceEntry per iteration if Options.store_trace was set, and empty otherwise. See trace.

Accessors

• Base.minimum(result) is the objective at the solution, i.e. the f field,
• solution(result) is the point,
• status(result) is the OptimizerStatus, and trace(result) the record.

minimum is Base's. The other three are internal: solution, status and trace are neither exported nor Base.ispublic, so a caller writes GeometricOptimizers.solution(result) and so on.

minimum and not value: value evaluates an OptimizerProblem at a point, which is a different question from reading a number a finished solve already holds.

An OptimizerState is a data structure for the dispatch on different algorithms.

It needs to implement three methods,

initialize!(alg::OptimizerState, ::AbstractVector)
update!(alg::OptimizerState, ::AbstractVector)
solver_step!(::AbstractVector, alg::OptimizerState)

that initialize and update the state of the algorithm and perform an actual optimization step.

Further the following convenience methods should be implemented,

problem(alg::OptimizerState)
gradient(alg::OptimizerState)
hessian(alg::OptimizerState)
linesearch(alg::OptimizerState)

which return the problem to optimize, its gradient and (approximate) Hessian as well as the linesearch algorithm used in conjunction with the optimization algorithm if any.

See NewtonState for a struct that was derived from OptimizerState.

OptimizerStatus

Contains residuals (relative and absolute) and various convergence properties.

This is also used in OptimizerResult.

Examples

x = ones(3)
state = NewtonState(x)
cache = NewtonOptimizerCache(x)
f = 1.
config = Options()
OptimizerStatus(state, cache, f; config = config)

# output

 * Convergence measures

    |x - x'|               = NaN
    |x - x'|/|x'|          = NaN
    |f(x) - f(x')|         = NaN
    |f(x) - f(x')|/|f(x')| = NaN
    |g(x) - g(x')|         = NaN
    |g(x)|                 = NaN

OptimizerTraceEntry

One iteration's worth of the trace solve! records when SimpleSolvers.Options.store_trace is set.

Keys

• iteration: the iteration number the entry was taken at,
• f: the objective there,
• rg: the gradient residual $\|\nabla{}f\|$ there.

See trace.

OrthonormalizationFailure(attempts)

Thrown by orthonormal_columns when attempts independent draws each break CholeskyQR2 down. At the rate measured in that function's docstring this is not bad luck, so what a caller changes in response is the element type and not the seed. The type is distinct so that catching it does not also catch an unrelated ErrorException.

PhaseTimer(; phases=nothing, clock=time_ns, synchronize=_no_synchronize)

Accumulate exclusive phase times in nanoseconds.

The exclusive dictionary stores the accumulated time for each phase and calls stores how often that phase was entered. Nested phases are excluded from their parent. By default every phase is recorded; pass a phase name or an iterable of phase names as phases to record less. Unselected nested phases are still excluded from a selected parent's time.

clock defaults to time_ns and must return a monotonically increasing count of nanoseconds that UInt64 accepts exactly. Base.time is therefore not a usable clock: it returns seconds as a Float64 and the conversion throws an InexactError. Differences are taken in UInt64, so a clock that ever goes backwards underflows rather than reporting a negative interval.

synchronize is called immediately before every clock reading and defaults to a no-op; a GPU caller can pass its device synchronization function. Call empty!(timer) between runs to reset its accumulators.

PrecomputedGradient(x, dp) <: SimpleSolvers.Gradient

A SimpleSolvers.Gradient that holds an already computed Euclidean gradient dp of the parameters x, and returns its Riemannian gradient rgrad(x, dp) when it is called on x.

This is the gradient of a training step: a training loop computes the gradient of its loss on one minibatch, and optimization_step! wraps it in this type, so the optimizer caches build the direction from it exactly as they build it from a gradient they evaluate themselves. For a parameter set, dp is a NeuralNetworkParameters.NetworkParameters of the same shape as x.

Implementation

optimization_step! passes the RiemannianGradientBuffers of its optimizer as a third argument, and the Riemannian gradient of a Stiefel or a Grassmann leaf is then written into them rather than allocated. The result is overwritten by the next call, which store_gradient! makes only after it has read the one before.

ProjectedSkew <: AbstractExponentialAlgorithm

Exponentiate the lift in a basis of its own range, where it is a small skew-symmetric matrix.

$\bar{B}$ is skew-symmetric of rank at most $2n$, so its range and its row space coincide and a thin QR of $B'$ gives an $N\times{}2n$ orthonormal Q with $\bar{B} = QMQ^T$ for $M = Q^T\bar{B}Q$ skew-symmetric and $2n\times{}2n$. Then

\[\exp(\bar{B}) = \mathbb{I} + Q\left(\exp(M) - \mathbb{I}\right)Q^T,\]

and $\exp(M)$ is formed from an eigendecomposition: $iM$ is Hermitian for real skew $M$, so $M = -iV\Lambda{}V^*$ and $\exp(M) = \Re\left(V e^{-i\Lambda} V^*\right)$, which is orthogonal by construction rather than by cancellation.

This is the only algorithm whose check does not degrade with the size of the lift, because its orthogonality is structural — it comes from the eigenvector matrix, not from the accuracy of a series. The gap is widest in Float32, where the other four climb into the $10^{-5}$s while this one stays in the $10^{-6}$s from one end of the norm sweep to the other.

The trade is the forward error against exp(Matrix(B)), the largest of the four at all but the very largest lifts, and a QR plus an eigendecomposition instead of matrix products. Choose it when staying on the manifold matters more than agreeing with the exponential to the last bit — a long Float32 run, for instance, where check accumulates over thousands of steps.

See 5. ProjectedSkew for the measurements and AbstractExponentialAlgorithm for the alternatives.

ProjectionTransport <: AbstractVectorTransport

The transport by orthogonal projection onto the tangent space at $Y_\mathrm{new}$:

\[\Xi_\mathrm{out} = \Xi - \frac{1}{2}Y_\mathrm{new}(Y_\mathrm{new}^T\Xi + \Xi^TY_\mathrm{new}).\]

It reads neither $Y$ nor $\Delta$.

QuasiNewtonCache

The OptimizerCache of BFGS and DFP. The two methods share every field and every step of update! but the update of the inverse Hessian, which _update_inverse_hessian! takes from method.

g̃ is the scratch for latest_gradient and g̃_is_current says whether it is the gradient at x; see GradientCache, which carries the same pair for the same reason, and store_gradient!.

T1, T2, T3, ΔxΔg and ΔxΔx are the $n \times n$ scratch of the update, where $n$ is the intrinsic dimension; DFP needs four of the five, and builds $\gamma\gamma^T$ in T3.

QuasiNewtonOptimizerMethod <: OptimizerMethod

Includes BFGS and DFP.

RetractionWorkspace(x::Manifold, retraction = nothing)
RetractionWorkspace(backend, T, N, n, retraction = nothing)

The buffers update_section! retracts a horizontal lift into, held for the life of an Optimizer rather than rebuilt per call.

Every array a retraction of an $N\times{}n$ point needs is a fixed shape once $N$ and $n$ are known — the two factors of lift_factors, the $2n\times{}2n$ matrix the inverse is taken of, the $N\times{}N$ result, and the frame the result is transported into. An optimizer step applies a retraction once per line-search trial and three times besides, so those shapes are rebuilt six to fifteen times an iteration for a step that changes none of them.

Implementation

The four blocks that do not depend on the lift are written once, here, and never again: the two $n\times{}n$ identity blocks of $B'$ and $(B'')^T$, and the zero blocks beside them. lift_factors! then writes only the four that do. 𝕀_small2 and 𝕀_big are kept because both sums are formed as a five-argument mul! into an identity, which needs the identity to copy from.

The element type and the backend come from the point, so a workspace is on the same backend as the solve it belongs to and no method here names either.

The workspace also holds the scratch of GeometricOptimizers.𝔄!, which a geodesic step evaluates its $\mathfrak{A}$ with, and of retraction_differential!, as much of it as retraction reads: the 𝔄! scratch of its algorithm for a Geodesic, and that of the differential for Cayley on a host Matrix of a LAPACK element type. Every other scratch array is 0 × 0, so the type of a workspace does not depend on its retraction. A workspace built for one retraction serves that retraction only; one built with retraction = nothing holds all the scratch and serves every retraction. Optimizer builds its workspace for its own retraction.

In a workspace, the $\mathfrak{A}$ of ScaledSquaring and NativePade allocates nothing, and that of AugmentedPade allocates what Base.exp of its $4n\times{}4n$ matrix allocates. So in both precisions the two native algorithms are the lighter, and AugmentedPade the heavier by exactly its exp. scripts/in_place_retraction_cost.jl measures the figures.

What this does not remove: the $2n\times{}2n$ inv in cayley and the exp of AugmentedPade. Neither grows with $N$. inv stays an inv because lu! and rdiv! would take it in place on the host only, and a KernelAbstractions backend is under no obligation to supply either — inv is what cayley already runs on Metal through. ProjectedSkew takes its allocating geodesic and copies the answer in.

A solve whose parameters are not on a manifold has no workspace: see GeometricOptimizers.retraction_workspace.

Rfac(Λ::SymplecticHouseholderDecom)

The $R$ factor of a symplectic SR decomposition, as an operator.

$R$ is upper triangular in the symplectic sense: in the $2N\times2M$ block form

\[R = \begin{pmatrix} R_{11} & R_{12} \\ R_{21} & R_{22} \end{pmatrix},\]

$R_{11}$ and $R_{22}$ are upper triangular and $R_{21}$ is strictly upper triangular. The entries come from the packed Λ; nothing is stored twice.

RiemannianGradient(gradient) <: SimpleSolvers.Gradient

A SimpleSolvers.Gradient whose value is projected onto the tangent space of the point it is evaluated at.

gradient is the Euclidean gradient, over the flattened iterate; the functor evaluates it there and then applies rgrad — leaf by leaf for a parameter set, whole for a matrix. Vector calls are forwarded unchanged, so a line search that works in coordinates sees the inner gradient and nothing else.

Implementation

This type exists to own a signature. A method of SimpleSolvers.Gradient itself on a Matrix or on a parameter set would be type piracy: the function is SimpleSolvers' and neither Base.Matrix nor NeuralNetworkParameters.NetworkParameters is this package's. The body cannot go upstream either: rgrad is the Riemannian projection of this package, and SimpleSolvers does not have it. So the gradient is wrapped, as GeometricMachineLearning does for its _GMLGradient.

Optimizer wraps for you, and only where it changes something — see _riemannian_gradient. A Manifold iterate needs no wrapper, because Manifold is this package's type and the method above is therefore already owned; wrapping one is harmless all the same, since the functor below leaves that method to handle it.

SR(S::Sfac, R::Rfac)

The result of sr!: the two factors of $A = SR$, reached as F.S and F.R.

ScalarMomentAdam(; β₁, β₂, δ, ambient_norm)

Cayley ADAM, [11, Algorithm 2], as an experimental Stiefel-only baseline for Adam.

The moments are stored in bias-corrected form, as Adam's are, but the second one is a scalar:

\[m_1 \gets \frac{\beta_1 - \beta_1^t}{1 - \beta_1^t}m_1 + \frac{1 - \beta_1}{1 - \beta_1^t}\bar{G},\]

\[m_2 \gets \frac{\beta_2 - \beta_2^t}{1 - \beta_2^t}m_2 + \frac{1 - \beta_2}{1 - \beta_2^t}\lVert\bar{G}\rVert^2,\]

where $\bar{G}\in\mathfrak{g}^\mathrm{hor}$ is the gradient in the global tangent space representation, and the direction is $-m_1/\sqrt{m_2 + \delta}$.

That one scalar is the entire difference from Adam, and it is what the name says: $\lVert\bar{G}\rVert^2$ is a squared gradient norm where Adam's $m_2$ is a squared gradient, so the second moment carries no direction and the method assigns one adaptive learning rate to the whole matrix instead of one per coordinate. That is a reproduction of a published algorithm and not a straw man — on some objectives it will beat Adam.

Only a single StiefelManifold{T} solution is supported; ordinary arrays, NamedTuples, Grassmann solutions and mixed parameter trees throw a MethodError. As for Adam, the coefficients carry no element type of the parameters, and Optimizer converts them.

Arguments

β₁, β₂ and δ are the source's $\beta_1$, $\beta_2$ and $\varepsilon$, with $0 \le \beta_1, \beta_2 < 1$ and $\delta \ge 0$ validated. ambient_norm selects which $\lVert\cdot\rVert^2$ the second moment accumulates and defaults to false; see GeometricOptimizers._squared_gradient_norm, which is also where the reason the two choices are not interchangeable up to a constant is written down.

Extended help

The source's Algorithm 2, a table mapping each of its symbols onto the code, the derivation that makes its lines 8–10 be global_rep, and the two departures above with their reasons are on the Optimizer Methods page. The third departure is the momentum's transport: the source re-projects $M_k$ onto the tangent space at the new iterate (its equation (6)) where update_section! carries the lift here, which agrees for a Lie group and differs on a proper homogeneous space — the one departure that changes the iterates against the authors' implementation on the same objective and seed.

What is recorded only here is three discrepancies between the source and its own implementation, which are not departures of this port:

• $v_1$. Its line 2 initializes it to $1$, stiefel_optimizer.py to $0$ (v_buffer = torch.zeros([1])), and line 6's $1 - \beta_2^k$ is $0$ at the $k = 0$ its loop starts from. $0$ is followed, which is also Adam's convention here and which makes the first direction a normalized gradient step.
• $\varepsilon$ inside or outside the root. Its line 7 and its implementation (vnew_hat.add(epsilon).sqrt()) both put it inside; Adam here puts it outside, $\sqrt{m_2}+\delta$. The source's placement is used, so this is the one convention ScalarMomentAdam does not share with Adam.
• $s$, the number of fixed-point iterations. Its line 1 gives $s = 2$; Cayley_loop in gutils.py runs for i in range(5). Neither number reaches this method, because the iteration they count is the approximation of the Cayley transform that the admonition above says is not ported — Cayley is evaluated exactly. It is recorded because a reader comparing this port against the authors' code will meet it.

ScalarMomentAdamCache <: OptimizerCache

Cache for ScalarMomentAdam.

The fields are AdamCache's, except that m₂ is a single number rather than an element of 𝔤ʰᵒʳ:

• x::StiefelManifold: the solution,
• g: the gradient, in StiefelLieAlgHorMatrix form,
• δ: the direction,
• Δg: difference in gradients, needed for OptimizerStatus,
• g̃: scratch for latest_gradient; see GradientCache,
• g̃_is_current: whether g̃ is the gradient at x; see store_gradient!,
• m₁: the first moment, a horizontal lift, stored bias-corrected,
• m₂::T: the second moment, the scalar of [11], stored bias-corrected,
• m̃₂::T: √(m₂ + method.δ), the scalar the direction is divided by. The δ in it is ScalarMomentAdam's regularizer and not the δ field above, which is the direction,
• section: the GlobalSection.

ScalarMomentAdamState <: OptimizerState

State for ScalarMomentAdam.

As for ScalarMomentAdamCache, m₂ is a scalar and not a horizontal lift.

ScaledSquaring(θ = 0.5) <: AbstractExponentialAlgorithm

Evaluate $\mathfrak{A}$ with a Taylor kernel and low-rank modified squaring, and the default.

This is not the Padé kernel used by the conventional dense matrix-exponential algorithm. Its improvement over TaylorSeries is that it first scales the argument into a regime where the same Taylor evaluator is accurate, then recovers the original argument with a recurrence specialized to the low-rank factorization [21].

The series for $\mathfrak{A}$ converges for every argument but is only accurate for a small one, so the argument is halved the smallest number $s$ of times with $\|X\|_1/2^s \leq θ$, the series is summed there, and each halving is undone by one application of $W \mapsto 2W + WXW$. That recurrence is the modified squaring of [21]; it stays at $2n\times{}2n$ because the low-rank form is closed under squaring, so nothing is ever squared at $N\times{}N$. The recovery steps are exact identities, so scaling moves where Taylor is evaluated without changing the function computed.

θ is the norm threshold the kernel sees. Smaller values use more scaling steps; larger values evaluate the series at a larger argument. The default 0.5 needs no tuning: across the 32-fold range $θ \in [0.125, 4]$ neither the orthogonality of the retracted point nor the error against exp moves by as much as a factor of six, and neither moves monotonically. Unlike NativePade's, this threshold is a preference and not a ceiling, because the inner series is summed to vanishing terms rather than for a fixed number of steps.

This is the cheapest of the five, and it needs neither dense LAPACK nor scalar indexing in package code — whether it runs on a given accelerator depends on that backend's support for matrix products and reductions. What it gives up is that its orthogonality is an arithmetic outcome rather than a structural property, so check drifts upwards with the size of the lift; only ProjectedSkew avoids that.

2. Scaling and modified squaring gives the algorithm step by step, with the identity behind the recurrence and the measurements behind every claim above. See AbstractExponentialAlgorithm for the alternatives.

Sfac(Λ::SymplecticHouseholderDecom, inverse::Bool = false)

The symplectic factor $S$ of a symplectic SR decomposition, as an operator.

S is never formed: multiplying by it applies the reflectors Λ stores, in $O(N^2)$ per vector. inv(S) is the same object with the reflectors applied in the opposite order, which is why the inverse costs what the forward application costs.

Indexing one entry builds the whole matrix and throws the rest away, because a reflector product has no cheap entrywise form. Use Matrix(S) when more than a few entries are wanted; that is what rand(SymplecticStiefelManifold, …) does before it slices.

SkewSymMatrix(S::AbstractVector, n::Integer)

Instantiate a skew-symmetric matrix with information stored in vector S.

A skew-symmetric matrix $A$ is a matrix $A^T = -A$.

Internally the struct saves a vector $S$ of size $n(n-1)\div2$. The conversion is done the following way:

\[[A]_{ij} = \begin{cases} 0 & \text{if $i=j$} \\ S[( (i-2) (i-1) ) \div 2 + j] & \text{if $i>j$}\\ S[( (j-2) (j-1) ) \div 2 + i] & \text{else}. \end{cases}\]

So $S$ stores a string of vectors taken from $A$: $S = [\tilde{a}_1, \tilde{a}_2, \ldots, \tilde{a}_n]$ with $\tilde{a}_i = [[A]_{i1},[A]_{i2},\ldots,[A]_{i(i-1)}]$.

Also see SymmetricMatrix, StrictlyLowerTriangular and StrictlyUpperTriangular.

Examples

using GeometricOptimizers
S = [1, 2, 3, 4, 5, 6]
SkewSymMatrix(S, 4)

# output

4×4 SkewSymMatrix{Int64, Vector{Int64}}:
 0  -1  -2  -4
 1   0  -3  -5
 2   3   0  -6
 4   5   6   0

SkewSymMatrix(A::AbstractMatrix)

Perform 0.5 * (A - transpose(A)) and store the matrix in an efficient way (as a vector with $n(n-1)/2$ entries).

If the constructor is called with a matrix as input it returns a skew-symmetric matrix via the projection:

\[A \mapsto \frac{1}{2}(A - A^T).\]

Examples

using GeometricOptimizers
M = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]
SkewSymMatrix(M)

# output

4×4 SkewSymMatrix{Float64, Vector{Float64}}:
 0.0  -1.5  -3.0  -4.5
 1.5   0.0  -1.5  -3.0
 3.0   1.5   0.0  -1.5
 4.5   3.0   1.5   0.0

Extended help

Note that the constructor is designed in such a way that it always returns matrices of type SkewSymMatrix{<:AbstractFloat} when called with a matrix, even if this matrix is of type AbstractMatrix{<:Integer}.

If the user wishes to allocate a matrix SkewSymMatrix{<:Integer} then call:

SkewSymMatrix(::AbstractVector, n::Integer)

Note that this is different from StrictlyLowerTriangular and StrictlyUpperTriangular as no projection takes place there.

StiefelLieAlgHorMatrix(A::SkewSymMatrix, B::AbstractMatrix, N::Integer, n::Integer)

Build an instance of StiefelLieAlgHorMatrix based on a skew-symmetric matrix A and an arbitrary matrix B.

An element of StiefelLieAlgMatrix takes the form:

\[\begin{pmatrix} A & B^T \\ B & \mathbb{O} \end{pmatrix},\]

where $A$ is skew-symmetric (this is SkewSymMatrix in GeometricOptimizers).

Also see GrassmannLieAlgHorMatrix.

Extended help

StiefelLieAlgHorMatrix is the horizontal component of the Lie algebra of skew-symmetric matrices (with respect to the canonical metric).

The projection here is: $\pi:S \to SE$ where

\[E = \begin{bmatrix} \mathbb{I}_{n} \\ \mathbb{O}_{(N-n)\times{}n} \end{bmatrix}.\]

The matrix $E$ is implemented under StiefelProjection in GeometricOptimizers.

StiefelLieAlgHorMatrix(D::AbstractMatrix, n::Integer)

Take a big matrix as input and build an instance of StiefelLieAlgHorMatrix.

The integer $N$ in $St(n, N)$ is the number of rows of D.

Extended help

If the constructor is called with a big $N\times{}N$ matrix, then the projection is performed the following way:

\[\begin{pmatrix} A & B_1 \\ B_2 & D \end{pmatrix} \mapsto \begin{pmatrix} \mathrm{skew}(A) & -B_2^T \\ B_2 & \mathbb{O} \end{pmatrix}.\]

The operation $\mathrm{skew}:\mathbb{R}^{n\times{}n}\to\mathcal{S}_\mathrm{skew}(n)$ is the skew-symmetrization operation. This is equivalent to calling of SkewSymMatrix with an $n\times{}n$ matrix.

This can also be seen as the operation:

\[D \mapsto \Omega(E, DE) = \mathrm{skew}\left(2 \left(\mathbb{I} - \frac{1}{2} E E^T \right) DE E^T\right).\]

Also see GeometricOptimizers.Ω.

StiefelManifold <: Manifold

An implementation of the Stiefel manifold [12]. The Stiefel manifold is the collection of all matrices $Y\in\mathbb{R}^{N\times{}n}$ whose columns are orthonormal, i.e.

\[ St(n, N) = \{Y: Y^TY = \mathbb{I}_n \}.\]

The Stiefel manifold can be shown to have manifold structure (as the name suggests) and this is heavily used in GeometricOptimizers. It is further a compact space. More information can be found in the docstrings for rgrad(::StiefelManifold, ::AbstractMatrix) and metric(::StiefelManifold, ::AbstractMatrix, ::AbstractMatrix).

StiefelMetric

A Riemannian metric on the Stiefel manifold, passed as the last argument of rgrad and metric. The two metrics are CanonicalMetric and EuclideanMetric.

StiefelProjection(backend, T, N, n)

Make a matrix of the form $\begin{bmatrix} \mathbb{I} & \mathbb{O} \end{bmatrix}^T$ for a specific backend and data type.

An array that essentially does vcat(I(n), zeros(N-n, n)) with GPU support.

Extended help

For $N \geq n$ an instance of StiefelProjection should technically also belong to StiefelManifold: its columns are orthonormal.

Any other shape is the same matrix Matrix{T}(I, N, n): ones on the diagonal and zeros elsewhere. For $n > N$ that is $\begin{bmatrix} \mathbb{I} & \mathbb{O} \end{bmatrix}$, whose columns are not orthonormal, and a projection with no rows or no columns is the empty matrix of its size. The host and the backend constructors agree on every shape.

StiefelProjection(B::AbstractLieAlgHorMatrix)

Extract necessary information from B and build an instance of StiefelProjection.

Necessary information here refers to the backend, the data type and the size of the matrix.

The size is queried through B.N and B.n.

Examples

using GeometricOptimizers
using GeometricOptimizers: StiefelProjection

B₁ = rand(StiefelLieAlgHorMatrix, 5, 2)
B₂ = rand(GrassmannLieAlgHorMatrix, 5, 2)
E = [1. 0.; 0. 1.; 0. 0.; 0. 0.; 0. 0.]

StiefelProjection(B₁) ≈ StiefelProjection(B₂) ≈ E 

# output

true

StiefelProjection(A::AbstractMatrix)

Extract necessary information from A and build an instance of StiefelProjection.

Necessary information here refers to the backend, the data type and the size of the matrix.

StrictlyLowerTriangular(S::AbstractVector, n::Int)

Build a lower-triangular matrix from a vector.

A lower-triangular matrix is an $n\times{}n$ matrix that has zeros on the diagonal and on the upper triangular.

The data are stored in a vector $S$ similarly to other matrices. See StrictlyUpperTriangular, SkewSymMatrix and SymmetricMatrix.

The struct has two fields: S and n. The first stores all the entries of the matrix in a sparse fashion (in a vector) and the second is the dimension $n$ for $A\in\mathbb{R}^{n\times{}n}$.

adjoint (L') returns a StrictlyUpperTriangular built around the same storage vector, not a copy: parent(L') === parent(L) holds, so writing into the adjoint also writes into L. Reusing the storage transposes without conjugating, so this method is defined for a real element type only; a complex one falls through to LinearAlgebra's lazy Adjoint, which conjugates and does not alias.

Examples

using GeometricOptimizers
S = [1, 2, 3, 4, 5, 6]
StrictlyLowerTriangular(S, 4)

# output

4×4 StrictlyLowerTriangular{Int64, Vector{Int64}}:
 0  0  0  0
 1  0  0  0
 2  3  0  0
 4  5  6  0

StrictlyLowerTriangular(A::AbstractMatrix)

Build a lower-triangular matrix from a matrix.

This is done by taking the lower left of that matrix.

Examples

using GeometricOptimizers
M = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]
StrictlyLowerTriangular(M)

# output

4×4 StrictlyLowerTriangular{Int64, Vector{Int64}}:
  0   0   0  0
  5   0   0  0
  9  10   0  0
 13  14  15  0

StrictlyUpperTriangular(S::AbstractVector, n::Int)

Build an upper-triangular matrix from a vector.

An upper-triangular matrix is an $n\times{}n$ matrix that has zeros on the diagonal and on the lower triangular.

The data are stored in a vector $S$ similarly to other matrices. See StrictlyLowerTriangular, SkewSymMatrix and SymmetricMatrix.

The struct has two fields: S and n. The first stores all the entries of the matrix in a sparse fashion (in a vector) and the second is the dimension $n$ for $A\in\mathbb{R}^{n\times{}n}$.

adjoint (U') returns a StrictlyLowerTriangular built around the same storage vector, not a copy: parent(U') === parent(U) holds, so writing into the adjoint also writes into U. Reusing the storage transposes without conjugating, so this method is defined for a real element type only; a complex one falls through to LinearAlgebra's lazy Adjoint, which conjugates and does not alias.

Examples

using GeometricOptimizers
S = [1, 2, 3, 4, 5, 6]
StrictlyUpperTriangular(S, 4)

# output

4×4 StrictlyUpperTriangular{Int64, Vector{Int64}}:
 0  1  2  4
 0  0  3  5
 0  0  0  6
 0  0  0  0

StrictlyUpperTriangular(A::AbstractMatrix)

Build an upper-triangular matrix from a matrix.

This is done by taking the upper right of that matrix.

Examples

using GeometricOptimizers
M = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]
StrictlyUpperTriangular(M)

# output

4×4 StrictlyUpperTriangular{Int64, Vector{Int64}}:
 0  2  3   4
 0  0  7   8
 0  0  0  12
 0  0  0   0

SymmetricMatrix(S::AbstractVector, n::Integer)

Instantiate a symmetric matrix with information stored in vector S.

A SymmetricMatrix $A$ is a matrix $A^T = A$.

Internally the struct saves a vector $S$ of size $n(n+1)\div2$. The conversion is done the following way:

\[[A]_{ij} = \begin{cases} S[( (i-1) i ) \div 2 + j] & \text{if $i\geq{}j$}\\ S[( (j-1) j ) \div 2 + i] & \text{else}. \end{cases}\]

So $S$ stores a string of vectors taken from $A$: $S = [\tilde{a}_1, \tilde{a}_2, \ldots, \tilde{a}_n]$ with $\tilde{a}_i = [[A]_{i1},[A]_{i2},\ldots,[A]_{ii}]$.

Also see SkewSymMatrix, StrictlyLowerTriangular and StrictlyUpperTriangular.

Examples

using GeometricOptimizers
S = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
SymmetricMatrix(S, 4)

# output

4×4 SymmetricMatrix{Int64, Vector{Int64}}:
 1  2  4   7
 2  3  5   8
 4  5  6   9
 7  8  9  10

SymmetricMatrix(A::AbstractMatrix)

Perform a projection and store the matrix in an efficient way (as a vector with $n(n+1)/2$ entries).

If the constructor is called with a matrix as input it returns a symmetric matrix via the projection:

\[A \mapsto \frac{1}{2}(A + A^T).\]

Examples

using GeometricOptimizers
M = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]
SymmetricMatrix(M)

# output

4×4 SymmetricMatrix{Float64, Vector{Float64}}:
 1.0   3.5   6.0   8.5
 3.5   6.0   8.5  11.0
 6.0   8.5  11.0  13.5
 8.5  11.0  13.5  16.0

Extended help

Note that the constructor is designed in such a way that it always returns matrices of type SymmetricMatrix{<:AbstractFloat} when called with a matrix, even if this matrix is of type AbstractMatrix{<:Integer}.

If the user wishes to allocate a matrix SymmetricMatrix{<:Integer} then call

SymmetricMatrix(::AbstractVector, n::Integer)

Note that this is different from StrictlyLowerTriangular and StrictlyUpperTriangular as no projection takes place there.

SymplecticHouseholderDecom(A, c₁, c₂, ρ, ν, μ)

The packed form of a symplectic SR decomposition, as sr! computes it.

A holds the reflector vectors below the diagonal of each block, in place, exactly as LinearAlgebra's QRCompactWY holds its own. The five vectors are the scalars the reflectors were built with: c₁ and c₂ scale the two reflections of a step, and ρ, ν and μ are the diagonal entries the step produced in R.

This is the storage, not the factors. Sfac and Rfac read S and R off it.

SymplecticStiefelManifold(A)

The symplectic Stiefel manifold, the set of $2N\times2n$ matrices whose columns span a symplectic subspace:

\[Sp(2n, 2N) = \{U \in \mathbb{R}^{2N\times2n} : U^T\mathbb{J}_{2N}U = \mathbb{J}_{2n}\},\]

with $\mathbb{J}$ the canonical Poisson tensor. Compare StiefelManifold, whose constraint is $Y^TY = \mathbb{I}$: the two agree in form and differ in which bilinear form the columns preserve.

The metric is the one of [14], and rand builds a point through the symplectic SR decomposition sr!, which is the construction [36] uses for its retraction. See also [13].

Examples

using GeometricOptimizers
using GeometricOptimizers: check
import Random

Random.seed!(1234)

check(rand(SymplecticStiefelManifold, 6, 4)) < 1e-5

# output

true

TaylorSeries <: AbstractExponentialAlgorithm

Sum the series for $\mathfrak{A}$ directly, without scaling. This is not a usable retraction.

It is kept so that the working algorithms have a baseline to be compared against, in the test suite and in the manual. The series is summed on $X = (B'')^TB'$, and that reduced matrix is far from normal: its norm is quadratic in $\|\bar{B}\|_2$ where its spectral radius is exactly $\|\bar{B}\|_2$. On such an argument the terms cancel, so stopping when a term falls below eps leaves an error of size $\varepsilon\max_m\|S_m\|$ rather than $\varepsilon\|\mathfrak{A}(X)\|$, where $S_m$ are the partial sums. By $\|\bar{B}\| \approx 79$ the partial sums exceed the result by twenty orders of magnitude and the "retracted" point is not on the Stiefel manifold in any sense; by $\|\bar{B}\| \approx 770$ the summation has overflowed.

The failure is silent: nothing reports that the result is off the manifold. check is what detects it, and no optimizer in this package calls it during a run. Making the termination test relative to the partial sum instead of absolute changes none of it: the loss is in the cancellation, not in when the summation stops.

Use ScaledSquaring, which fixes this and is also the cheaper of the two, because the scaled series converges in a handful of terms where the unscaled one grinds through hundreds.

1. Direct Taylor series derives the failure and tabulates it over the whole norm sweep. See AbstractExponentialAlgorithm for the alternatives.

TrainingOptimizer(x; algorithm = Adam(), linesearch = default_step_size(algorithm), retraction = Cayley(), observer = NoStepObserver())

The optimizer of a training loop: one step per minibatch, with optimization_step!.

It holds the first-order algorithm, converted once to the element type of x with change_precision, its cache and its state over the whole of x, the step-size schedule, the retraction, the buffers the retraction works in, and those the Riemannian gradient of a Stiefel or a Grassmann leaf is written into.

algorithm is a GradientMethod, a MomentumMethod, a DirectMomentum or a member of the AdamFamily; ScalarMomentAdam, DirectMomentum, DirectScalarMomentAdam and FactoredStiefelAdam step a single StiefelManifold only, and raise a MethodError for any other x. For a parameter set it may also be a CompositeMethod: then cache, state, workspace and gradient_buffer are trees of the shape of x, with the cache, the state and the buffers of each leaf for the method of that leaf. linesearch is a finite positive number, which is the fixed step size Static(η), a SimpleSolvers.Static or a DecayingStatic; the default is default_step_size(algorithm). retraction is an instance of an AbstractRetraction, Cayley() or Geodesic(). The three direct Stiefel methods retract with cayley_direct! in a DirectCayleyWorkspace, and another retraction than Cayley() is an ArgumentError for them. They form the Riemannian gradient in their own cache, so their leaf has no gradient buffer.

observer is notified of the :retraction_application phase of every step, once per leaf for a CompositeMethod; see Observing the Phases of an Optimizer Step.

Why a training step is not solve!

solve! minimizes an objective it can evaluate: its line search compares the objective along the direction, and it stops when a convergence test on the objective and its gradient passes. A training step has neither. It has the gradient of the loss on one minibatch, which is a stochastic estimate of the gradient of the loss it minimizes; the loss on that minibatch is not the objective, so there is nothing for a line search to search along and nothing for a convergence test to test. So linesearch takes a fixed or a scheduled step size and nothing else, and the loop that calls optimization_step! decides when to stop.

Examples

x = Float32[1, -2, 3]
opt = TrainingOptimizer(x; algorithm = GradientMethod(), linesearch = 0.5)
optimization_step!(x, opt, Float32[2, 2, 2])

# output

3-element Vector{Float32}:
  0.0
 -3.0
  2.0

*(B::AbstractLieAlgHorMatrix, C::AbstractMatrix)
*(C::AbstractMatrix, B::AbstractLieAlgHorMatrix)
*(B::AbstractLieAlgHorMatrix, c::AbstractVector)

The product, taken on the stored blocks rather than through getindex.

Without these the product falls through to the generic AbstractMatrix path, which asks the lift for one entry at a time. That is scalar indexing, so it cannot run on a device at all — the third wrapper in this package to meet that gap, after StiefelProjection and the two GeometricOptimizers.AbstractTriangulars.

A lift is block $\begin{pmatrix} A & -B^T \\ B & \mathbb{O} \end{pmatrix}$, with $A$ absent for a GrassmannLieAlgHorMatrix, so the product against an $N\times{}m$ matrix is three block products — two for a Grassmann lift — and nothing else. No kernel is needed and none is written, which is where this differs from the triangulars: those hold a packed vector their getindex unpacks, while a lift holds ordinary blocks and a SkewSymMatrix that already carries a kernel-backed product of its own. The per-type part is the first $n$ rows, _hor_top_rows, beside each concrete lift.

*(C, B) is written -transpose(B * permutedims(C)), on the identity $B^T = -B$; a row vector x takes transpose(x) in place of permutedims(C), and the vector form goes through the matrix one as a single column. transpose and permutedims, and not adjoint, wherever one appears, for the reason the comment on _rmul(::AbstractMatrix, ::SkewSymMatrix) in special_matrices/skew_symmetric.jl gives at length: the identity is a statement about the transpose, and getindex builds the off-diagonal blocks entrywise without conjugating. The two agree on a real element type and disagree on a complex one.

*(A::AbstractTriangular, B::AbstractMatrix)

The product, read off the packed storage vector by a kernel rather than through getindex.

Without this method the product falls through to the generic AbstractMatrix path, which asks A for one entry at a time. That is scalar indexing, so it cannot run on a device at all. The packed vector holds $n(n-1)/2$ entries and the generic path reads $n^2$ of them, the other $n(n+1)/2$ being zeros that StrictlyLowerTriangular and StrictlyUpperTriangular manufacture — but do not read a host speed-up into that, which is what the arithmetic invites. Against the path this method shadows, the ratio is not monotone in n, and at a small n both products are slower, because the kernel launch has a fixed cost that a small product cannot amortize. scripts/triangular_multiply_cost.jl is the check and CHANGELOG.md carries its table. The device is what this buys.

The kernel is the one each subtype supplies, and it is where the two differ: a lower-triangular row i runs over 1:(i-1) and an upper-triangular one over (i+1):n, reading the same packed vector through different index arithmetic.

*(::AbstractMatrix, ::AbstractTriangular) is written as (A' * B')'. For a real element type that goes through here as well, because adjoint on one of these is then a type swap onto the same storage. A complex one does not: the swap is bound to Real, for the reason the comment on adjoint(::StrictlyLowerTriangular) in upper_triangular.jl gives, and the lazy Adjoint it falls through to is not an AbstractTriangular. So that product stays on the generic path and stays host-only.

λY * Y

Apply the element λY onto Y.

Here λY is an element of a Lie group and Y is an element of a homogeneous space.

*(E::StiefelProjection, A::AbstractMatrix)
*(A::AbstractMatrix, E::StiefelProjection)
*(E::StiefelProjection, b::AbstractVector)

The product, taken on the wrapped array.

StiefelProjection holds its entries in an ordinary array, so unwrapping is all these do — for E and E' alike, and for +, -, mul! and a scalar product as well, with the same methods that a manifold point uses. Without them the product falls through to the generic AbstractMatrix path, which reaches getindex one entry at a time. That is scalar indexing, and it is what stops a retraction on a device: geodesic and cayley each take one product against the projection — expB * E and cayleyB * E — with E built from the horizontal lift and so carrying the point's own backend. Both operands are on the device, and only the wrapper puts the product on the host path.

Both operands have to be on one backend: a pair on two backends raises an ArgumentError naming both, rather than answering on whichever backend the argument order picks.

one(B::AbstractLieAlgHorMatrix)

The $N\times{}N$ identity, built with a KernelAbstractions kernel.

Base.one(::AbstractMatrix) writes the diagonal in a scalar-indexed loop, which is what a GPU array cannot serve; geodesic reaches this on every retraction, for every lift type. GeometricOptimizers.𝔄, whose argument is a bare matrix, builds its $2n\times{}2n$ identities through GeometricOptimizers.unit_matrix for the same reason.

parent(B::AbstractLieAlgHorMatrix)

The tuple of blocks B's free parameters are stored in — (A, B) for a StiefelLieAlgHorMatrix, (B,) for a GrassmannLieAlgHorMatrix — and not the single array every other parent this package defines returns.

Every operation on a lift that is elementwise in the free parameters — as opposed to in the ambient $N\times{}N$ matrix, which has no setindex! and counts each off-diagonal block twice — is written once over this tuple rather than once per lift type. That is assign! below, and l2norm, and the _difference! / _add! / _rmul! family in named_tuple_wrapper.jl and the moment update of Adam, which reach the tuple through freeparameters. So a GrassmannManifold is driven through an Optimizer by the same code as a StiefelManifold.

The docstring is attached to the signature and not to the bare Base.parent: the package also defines parent for Manifold, SkewSymMatrix, SymmetricMatrix and AbstractTriangular, each of which returns the single array it wraps, and a signature-less docstring would be shown as the general meaning of parent for all of them.

rand(backend, manifold_type, N, n)

Draw random elements for a specific device.

Examples

Random elements of the manifold can be allocated on GPU. Call ...

rand(CUDABackend(), StiefelManifold{Float32}, N, n)

... for drawing elements on a CUDA device.

The element type

Naming it, as above, is what fixes it, and a named element type is honoured or refused — never narrowed. rand(MetalBackend(), StiefelManifold{Float64}, N, n) throws Metal's own ErrorException at allocation, whose message names Float64, rather than quietly returning a Float32 point of a type the caller did not ask for.

A call that leaves the element type open — rand(CUDABackend(), StiefelManifold, N, n) — gets default_eltype of the backend: Float64 on the host and Float32 on a device, for the reasons given there.

What the backend has to supply

A matrix product, a cholesky and a triangular solve, all for an array of its own type. The draw orthonormalizes a Gaussian matrix with _cholesky_qr2 and global_section does the same for the complement, so those three are what a backend needs to carry a point, draw one and carry an Optimizer.

qr is deliberately not among them. Metal implements no qr for an MtlArray: a call there raises Cannot access the contents of a private buffer, measured on real hardware. Any path through qr is host-only on such a device, and that path carries GlobalSection(Y) and Optimizer(Y, F) with it. Metal does implement cholesky, so CholeskyQR2 runs where Householder QR cannot, at better orthogonality than the host QR at every size measured — and a CUDA backend, which does supply qr through CUSOLVER, takes the same route as everything else rather than a second one.

The manifolds this draws

StiefelManifold and GrassmannManifold. SymplecticStiefelManifold is drawn on the host alone: its draw is the symplectic SR decomposition sr!, which is a host factorization, so a device spelling would be a host draw followed by a transfer rather than the device-native draw the other two get. Naming a device for it throws an ArgumentError saying so.

rand(manifold_type, N, n)

Draw random elements from the Stiefel and the Grassmann manifold.

Because both of these manifolds are compact spaces we can sample them uniformly [37].

Examples

When we call ...

using GeometricOptimizers
using GeometricOptimizers: _round # hide
import Random
Random.seed!(123)

N, n = 5, 3
Y = rand(StiefelManifold{Float32}, N, n)
_round(Y; digits = 5) # hide

# output

5×3 StiefelManifold{Float32, Matrix{Float32}}:
 -0.27575   0.32991   0.77275
 -0.62485  -0.33224  -0.0686
 -0.69333   0.36724  -0.18988
 -0.09295  -0.73145   0.46064
  0.2102    0.33301   0.38717

... the sampling is done by first allocating a random matrix of size $N\times{}n$ via Y = randn(Float32, N, n).

We then orthonormalize its columns with _cholesky_qr2 and return those. CholeskyQR2 rather than LinearAlgebra.qr — which uses Householder reflections internally — because qr is a host factorization for several of the backends this package supports, and CholeskyQR2 is expressible in products, reductions and triangular solves alone. The two answers differ by the sign of each column, since CholeskyQR2's $R$ has a positive diagonal and Householder's need not.

rand(::Type{SymplecticStiefelManifold{T}}, N2, n2)

Draw a random point of the $2N\times2n$ symplectic Stiefel manifold.

The draw is a Gaussian matrix put through the symplectic SR decomposition sr!; the point is n columns from each half of the symplectic factor. This is the symplectic counterpart of the orthonormalization rand(::Type{StiefelManifold}, …) does, and no orthonormalization will serve here: an orthonormal factor preserves the Euclidean form, not $\mathbb{J}$.

update!(cache::NewtonOptimizerCache, x, g, hes)

Update an instance of NewtonOptimizerCache based on x.

This is used in update!(::OptimizerState, ::AbstractVector).

This sets:

\[\begin{aligned} % \bar{x}^\mathtt{cache} & \gets x, \\ x^\mathtt{cache} & \gets x, \\ g^\mathtt{cache} & \gets g, \\ \mathrm{rhs}^\mathtt{cache} & \gets -g, \\ H^\mathtt{cache} & \gets H(x), \\ \delta^\mathtt{cache} & \gets (H^\mathtt{cache})^{-1}\mathrm{rhs}^\mathtt{cache}, \end{aligned}\]

where we wrote $H$ for the Hessian (i.e. the input argument hes).

update!(state::NewtonState, gradient, x)

Update an instance of NewtonState based on x and gradient, where g is of type SimpleSolvers.Gradient.

Examples

If we only call update! once there are still NaNs for x̄, ḡ and f̄.

f(x) = sum(x.^2)
x = [1., 2.]
state = NewtonState(x)
grad = GradientAutodiff{Float64}(f, length(x))
update!(state, grad, x)

# output

NewtonState{Float64, Vector{Float64}, Vector{Float64}, GlobalSection{Float64, Vector{Float64}, Nothing}}(0, [1.0, 2.0], [NaN, NaN], [2.0, 4.0], [NaN, NaN], 5.0, NaN, GlobalSection{Float64, Vector{Float64}, Nothing}([1.0, 2.0], nothing))

update!(cache::QuasiNewtonCache, state, x, g)

Update the QuasiNewtonCache and the inverse Hessian $Q$ of state based on x and g.

Extended help

Both methods form the secant pair

\[\delta \gets x^{(k)} - x^{(k-1)}, \qquad \gamma \gets \nabla{}f^{(k)} - \nabla{}f^{(k-1)},\]

and update $Q$ from it only where curvature_is_usable holds. The update rules can be found in [38] and [33]. BFGS does

\[\begin{aligned} T_1 & \gets \delta\gamma^TQ, \\ T_2 & \gets Q\gamma\delta^T, \\ T_3 & \gets (1 + \frac{\gamma^TQ\gamma}{\delta^T\gamma})\delta\delta^T,\\ Q & \gets Q - (T_1 + T_2 - T_3)/{\delta^T\gamma}, \end{aligned}\]

and DFP does $Q \gets Q - Q\gamma\gamma^TQ/(\gamma^TQ\gamma) + \delta\delta^T/(\delta^T\gamma)$ (nocedal2006numerical, eq. 6.15).

AdamOptimizerWithDecay(n_epochs; η₁, η₂, kwargs...)

Adam paired with a DecayingStatic line search whose learning rate decays geometrically from η₁ to η₂ over n_epochs, returned as a NamedTuple to splat into Optimizer or TrainingOptimizer:

method = AdamOptimizerWithDecay(1000)
opt    = Optimizer(x, problem; method...)
solve!(x, OptimizerState(method.algorithm, x), opt)

The state has to be built from method.algorithm, so the pairing is worth binding to a name rather than splatting it twice.

There is no new type here and no schedule of its own: this is exactly Adam() and DecayingStatic(; η₁, η₂, n = n_epochs), under the name of GeometricMachineLearning's pairing. Neither carries the element type of the parameters; the optimizer converts both.

Why it exists

The name comes from GeometricMachineLearning, for Adam's ρ₁, ρ₂, δ together with a learning-rate schedule η₁, η₂, n_epochs and the same $\gamma = \exp(\log(\eta_2/\eta_1)/n)$. Here the two halves live in different places — the direction in an OptimizerMethod, the step size in a SimpleSolvers.LinesearchMethod — and this function is the one name that holds both.

Arguments

η₁, η₂ and n_epochs go to the line search; everything else is forwarded to Adam, so β₁, β₂ and δ — GML's ρ₁, ρ₂ and δ — keep Adam's own defaults rather than a second copy of them.

Examples

AdamOptimizerWithDecay(1000).linesearch

# output

DecayingStatic from α = 0.01 to α = 1.0e-6 over 1000 iterations.

_cholesky_qr2(A)

An $N\times{}m$ matrix whose columns are an orthonormal basis of the columns of the $N\times{}m$ matrix A, or nothing when A is too ill-conditioned for its element type.

CholeskyQR2: take $R$ from cholesky(AᵀA), set $Q = AR^{-1}$, and do it again. Every step is a matrix product, a reduction or a triangular solve, so the whole factorization runs wherever A already is. LinearAlgebra.qr! does not: Metal implements no qr for an MtlArray, and a qr on a JLArray cannot rebuild its Q. Measured on an M4 Max with scalar indexing disallowed, this runs on Metal and its $\|Q^TQ - \mathbb{I}\|$ is smaller than the host Householder QR's at every size tried.

The second pass is not a refinement. One pass leaves $\|Q^TQ - \mathbb{I}\|$ orders of magnitude above the round-off of the second, and the manifold tests assert check(Y) < 1e-14 in Float64.

On a device this keeps the point there; it is not the faster route. The manual page Global Tangent Spaces gives the measurements of both statements.

Forming $A^TA$ squares the condition number, which is what the nothing is for. cholesky with check = false reports a Gram matrix that is no longer positive definite rather than throwing, and at this shape in Float32 an ordinary Gaussian draw reaches that about once in a hundred and twenty — see scripts/orthonormalization_breakdown_rate.jl. orthonormal_columns is what answers it.

Squaring the entries is a second hazard, and the scaling below is what answers that one. A Householder QR scales internally and this does not, so an argument whose entries are large enough gives an infinite Gram matrix and no answer at all. It is reachable: optimizer_status_tests.jl builds a point $10^{100}$ off the manifold to test a convergence guard, the projection carries that magnitude into A, and $A^TA$ is then Inf in Float64. Dividing A by $\max_{ij}|A_{ij}|$ — a reduction that cannot overflow, unlike a column norm — leaves $Q$ unchanged, since $(A/s)^T(A/s) = A^TA/s^2$ has Cholesky factor $R/s$.

_dot(a, b)

The inner product of two gradients or directions, taken in the flattened coordinates.

Implementation

For an AbstractVecOrMat this is LinearAlgebra.dot. For a horizontal lift — or a NamedTuple of them — it is emphatically not: dot on an AbstractLieAlgHorMatrix is the ambient Frobenius product, which counts each of the off-diagonal blocks of the lift twice and so comes out exactly twice the product of the free parameters. The intrinsic coordinates are the ones every other quantity in this package is expressed in — Q is sized by the flattening, its outer products are formed there (see _flat_scratch), and the α of a line search parameterizes a curve in them — so pairing a gradient with a direction has to happen there too.

Used by trial_slope for $\varphi'(\alpha)$, by the quasi-Newton caches for $\delta^T\gamma$, whose value has to be consistent with the flattened T₁, T₂ and γ^TQγ it divides, for the predicted decrease $\widetilde{\Delta f}$, so that it is comparable with the measured $\Delta f$, and by ensure_descent!'s descent test, for the same reason.

_flat_mul!(c, A, b, scratch)

$c \gets Ab$ where A is in the flattened coordinates and c, b are in the parameters'.

The nothing method is mul!, for a solution that is already flat. The other flattens b into scratch, multiplies into scratch, and writes the result back through unflatten!, so that no flat vector and no ParameterLayout is allocated.

_flat_scratch(T, g)

The flat buffers a quasi-Newton cache keeps, or nothing where it needs none.

$Q$ lives in the flattened coordinates — it is sized by the length of the flattening, it is where the outer products that build it are formed and where _dot pairs — while the secant pair, the right-hand side and the direction are handed around in the parameters' own representation: a NamedTuple, a container, or a horizontal lift of the ambient shape. These are the buffers that the crossing between the two writes into, through NeuralNetworkParameters' allocation-free flatten! and unflatten!, so that no update! allocates a fresh vector.

A NeuralNetworkParameters.FlatParameters rather than a bare Vector, because it carries its own ParameterLayout and keeps it through similar — so δ is built once and the other three buffers are one similar each, with no parameterlayout call written anywhere below this line. One is still made: δ is the flattening of g on its backend, whose first act is to build the layout. What the similars buy is that it happens once per cache rather than once per product with $Q$.

That layout then goes into the cache's own type, as FlatParameters' third type parameter and so as QuasiNewtonCache's FT. A LeafLayout{N} holds the shape alone, so the cache type does not carry the array types of the leaves, and the cache keeps no reference to the leaves.

The ambient and the intrinsic

This is where the distinction the quasi-Newton methods turn on is written down. $Q$ is sized by the intrinsic dimension of the parameters — the length of their flattening — while the direction and the gradient are handed around in the ambient representation. For a bare StiefelManifold of size $(3, 1)$ those are 2 and $3 \times 3$ respectively, so SimpleSolvers.outer!, which checks the axes of its arguments against those of its destination, would throw a DimensionMismatch. Flattening first is what makes BFGS and DFP run on a bare Manifold at all, and the buffers below hold the flat form once per cache rather than once per call.

Built from g, which callers pass as _zero(x) and not x, for the reason the flatlength(_zero(x)) beside it gives: on a manifold the flattening of the lift is the intrinsic dimension, 12 against 18 for a StiefelManifold(6, 3), and that is the length Q multiplies.

nothing for an AbstractVector solution. There the parameters are the flat coordinates, outer! and mul! reach their own methods on them, and nothing is allocated to begin with — so buffering would add a copy per iteration and buy nothing. _flat_mul! and the secant pair of QuasiNewtonCache's update! take the parameters unchanged in that case.

_is_decayable(x)

Whether Euclidean weight decay can move x at all.

An ordinary array is decayable and a NamedTuple is decayable as soon as one of its entries is, which is the case AdamWithEuclideanDecay exists for. The StiefelManifold and the GrassmannManifold are not: $\lambda{}x$ is the gradient of $\frac{\lambda}{2}||x||^2$, which is constant on them, so the decay vanishes there identically (see the Weight Decay on Manifolds page).

Implementation

This is the only place that geometric fact is recorded — _weight_decay! consults it rather than restating it, so the no-op and the warning in OptimizerCache cannot drift apart.

It is declared on the two concrete manifolds and not on Manifold on purpose. What makes the decay vanish is that these two are compact, with $||Y||_F^2 = \mathrm{tr}(Y^TY) = n$; it is not a consequence of being a manifold, and on a noncompact one Euclidean decay may well do something. A manifold added later therefore has to decide for itself, and gets an error saying so rather than inheriting a silent no-op that may be wrong for it.

_manifold_αmax(solution, direction, c)

The ceiling a step ceiling of c imposes on a solution made of several blocks: the smallest of the per-block step_αmax over the blocks that live on a Manifold, and Inf where none does.

One \alpha is applied to every block, so each manifold block needs $\|\alpha\delta_i\| \leq 2\pi{}c$ and the binding one is the largest $\|\delta_i\|$. A block that is an ordinary array contributes nothing: the $2\pi$ is the turn of a rotation and a vector space has no such scale, so it must neither impose a ceiling of its own nor inflate the norm that sets one for its neighbours. A ceiling of step_αmax(c, direction(cache)) over the whole direction would do both: its l2norm combines the blocks in quadrature.

Implementation

NeuralNetworkParameters.foldparameters over the two trees in lockstep keeps a heterogeneous NamedTuple — a StiefelManifold beside a Matrix beside a Vector — inferable, and its compile time linear in the number of blocks. This function is on the per-iteration path, through linesearch_parameters.

foldparameters and not foldstorage: _block_αmax dispatches on the whole leaf, because a Manifold is the only kind that imposes a ceiling. A descent into the numbers that a leaf stores would make every manifold block an ordinary array and give back Inf.

min is associative and commutative, so the direction of the fold does not matter. A set with no manifold in it gives Inf through the _block_αmax(::Any, …) method, and a set with no blocks through the initial value T(Inf).

The fold pairs the blocks by key and checks the widths: δ follows the solution block by block, but not always as the same kind of branch, because a container solution can be paired with the plain NamedTuple of its GlobalSection tree.

The signature takes y::NetworkParameters and not an unconstrained y. Upstream pairs the children of a keyed branch by key, and the children of a Tuple branch positionally, so _manifold_αmax(values(sol), values(δ), c) would fold crossed keys without an error. Every caller passes a whole set of parameters.

_native_pade_polynomials(X, 𝕀)

Evaluate the degree-6 numerator $p_6(X)$ and denominator $q_6(X)$ used by NativePade. It allocates its result and scratch and calls _native_pade_polynomials!, which writes into given arrays.

If $P^{\exp}_7/Q^{\exp}_6$ is the $[7/6]$ Padé approximant of the exponential, then

\[p_6(z)=\frac{P^{\exp}_7(z)-Q^{\exp}_6(z)}{z}, \qquad q_6(z)=Q^{\exp}_6(z),\]

so $q_6(X)^{-1}p_6(X)$ agrees with $\mathfrak{A}(X)$ through the $X^{12}$ term. The implementation shares $X^2$ and $X^4$ between the two polynomials and groups the remaining terms to avoid forming every matrix power separately. 𝕀 must be the multiplicative identity with the same size, element type, and backend as X.

This is an internal kernel; NativePade supplies scaling, applies the denominator, and undoes the scaling with modified squaring.

_optimizer(x, problem, algorithm, linesearch, gradient, retraction, config)

Convert algorithm and linesearch to the element type of x, build the cache and the Hessian for algorithm, and hand everything to Optimizer's inner constructor.

Every argument is positional; see the implementation note of Optimizer(x, F).

_scaled_kernel!(scratch, X, algorithm)

Evaluate $\mathfrak{A}(X)$ for $\|X\|_1 \leq θ$ into scratch.𝔄X, i.e. the small-argument kernel that ScaledSquaring and NativePade differ in.

ScaledSquaring sums the Taylor series, NativePade evaluates the $[6/6]$ Padé approximant. Everything else the two algorithms do — choosing the number of halvings, and undoing them — is the shared framework in 𝔄!(::Any, ::AbstractMatrix, ::ScaledAlgorithm) below. scratch is as there.

_squared_gradient_norm(method::ScalarMomentAdam, cache, gradient, x)

The ‖·‖² that ScalarMomentAdam's second moment accumulates, selected by method.ambient_norm.

false — the default — squares the horizontal lift already in the cache; true squares [11]'s own ambient Euclidean gradient, at one extra gradient evaluation per step. Gradient applied to a Manifold returns rgrad(Y, ∇L) and the ambient ∇L is not part of the optimizer protocol, so true recovers it by calling the gradient on the flat storage, as GradientAutodiff(F, ::Manifold) and the other Manifold gradients are built for — which is what gradient(vec(parent(x))) is. store_gradient!'s reuse does not help, because what it caches is the lift. A PrecomputedGradient holds its ambient gradient, so its norm is that of dp, at no extra evaluation.

The two are not interchangeable up to a constant, and which to reach for when is on the Optimizer Methods page.

_sumsq_leaves(f, x)

$\sum_i f(x_i)^2$ over the leaves of x, at whatever depth they are.

The quadrature sum in which l2norm and solution_scale both combine their blocks; the two differ only in f. (l2norm is GeometricBase's function, so it has no docstring here to link to.)

Implementation

NeuralNetworkParameters.foldparameters and not map + sum over one level, because a container is a tree of layers and the quantities to combine are at its leaves — map would hand f a whole layer, for which neither of them has a method. foldparameters recurses on the branches and so reaches a leaf at any depth, and it allocates nothing. Its body is generated at literal indices, so its compile time is linear in the number of children of a branch.

foldparameters and not foldstorage, which is the pairing _dot wants and would be wrong here. foldstorage descends past a leaf into the numbers it stores, and solution_scale(Y::Manifold) is √size(Y, 2) — a nominal value that has nothing to do with those numbers.

false and not zero(T) for the initial value: there is no T in scope, and false is the strong zero that takes its type from whatever it is added to, so a one-block set adds it to that block's value and stays a T. A set has one element type, so every leaf gives it the same T. zero(parameter_eltype(x)) would put the generated _promote_eltypes body of upstream on this path, once per branch shape, which costs compile time.

f is annotated ::F where {F}, so that the method specialises on it even where it only passes f on: Julia does not specialise on a function argument that it does not see called. The docstring of foldparameters asks this of every caller that hands on an op.

_weight_decay!(δ, x, λ)

Subtract the decoupled weight decay $\lambda{}x$ from the direction δ, in place.

This is called after the Adam direction has been formed, i.e. δ is $-m_1/(\sqrt{m_2} + \delta)$ on entry and $-m_1/(\sqrt{m_2} + \delta) - \lambda{}x$ on exit. It has to happen before the line search scales the direction, so that the decay is multiplied by the learning rate as [34] prescribes, and it must not touch the moments — that is the whole point of the decoupling, see AdamWithEuclideanDecay.

Implementation

The Manifold method is a genuine no-op and not an omission: δ is then an element of $\mathfrak{g}^\mathrm{hor}$ rather than something that could be added to x at all, and the quantity it would have to be corrected by — the Riemannian gradient of $\frac{\lambda}{2}||x||^2$ — is zero. It asserts _is_decayable rather than restating why, so that the no-op holds only for the manifolds on which that has been established. A Riemannian decay that does something here is an open question. test/manifold_optimizers/adam_with_euclidean_decay_optimizer.jl checks that rgrad(Y, λY) really does vanish rather than taking it on trust.

advance_state!(state, cache, method)

Carry state over to the step cache has just taken, without evaluating anything: the section becomes the cache's, and the method's own memory advances — the momentum $p \gets \alpha{}p + \nabla{}L$ of MomentumMethod, the moments of the AdamFamily.

This is the half of a state update that needs no objective, which is why it is the whole of the state update in optimization_step!. solve!'s update!(state, opt, x, f) calls it after it has recorded the iterate, its gradient and its objective value.

section(cache) is update_section!(section(state), direction, retraction) once the step is taken, so it is copied and not retracted a second time. On a manifold the retraction is $O(N^2n)$, and update_section! then transports the $N\times{}(N-n)$ block of the section by a product with an $(N-n)\times{}N$ block, which is $O(N^2(N-n))$, i.e. $O(N^3)$ for $n \ll N$, where the copy is $O(N^2)$.

apply_S_inverse_left!(B, Λ::SymplecticHouseholderDecom)

Overwrite B with $S^{-1}B$. Three things invert apply_S_left! together: the steps run from the first to the last rather than the last to the first, the two reflectors within a step are applied in the opposite order, and each update adds where the forward one subtracts.

apply_S_inverse_right!(B, Λ::SymplecticHouseholderDecom)

Overwrite B with $BS^{-1}$. Three things invert apply_S_right! together, for the reason given at apply_S_inverse_left!: the steps run from the last to the first, the two reflectors within a step are applied in the opposite order, and each factor is inverted, which for $I + cvv^J$ means negating $c$ because $v^Jv = 0$.

apply_S_left!(B, Λ::SymplecticHouseholderDecom)

Overwrite B with $SB$, applying the reflectors of Λ from the last step to the first.

apply_S_right!(B, Λ::SymplecticHouseholderDecom)

Overwrite B with $BS$. The reflections act on the rows of B, so this is written out rather than expressed through apply_S_left!.

apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}}

Apply λY to Y₂ and store the result in Y.

This is the inplace version of apply_section.

apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT <: StiefelManifold{T}}

Apply λY to Y₂.

Mathematically this is the group action of the element $\lambda{}Y\in{}G$ on the element $Y_2$ of the homogeneous space $\mathcal{M}$.

Internally it calls apply_section!.

cayley(A::AbstractVecOrMat)

Return A unchanged. This is the identity by design, not a missing Cayley transform.

A plain vector or matrix is a parameter in a vector space, and the retraction there is the Euclidean one, $x \mapsto x + A$. That addition happens in update_section!, so the retraction applied to the direction has nothing left to do. geodesic on a plain array is the identity for the same reason. On a horizontal lift, cayley is the Cayley transform.

cayley(B̄::GrassmannLieAlgHorMatrix)

Compute the Cayley retraction of B.

This is equivalent to the method of cayley for StiefelLieAlgHorMatrix.

See cayley(::StiefelLieAlgHorMatrix).

cayley(B̄::StiefelLieAlgHorMatrix)

Compute the Cayley retraction of B.

Implementation

Internally this is using

\[\mathrm{Cayley}(\bar{B}) = \mathbb{I} + B' \left(\mathbb{I}_{2n} - \frac{1}{2} (B'')^T B'\right)^{-1} (B'')^T,\]

with

\[\bar{B} = \begin{bmatrix} A & -B^T \\ B & \mathbb{O} \end{bmatrix} = \begin{bmatrix} \frac{1}{2}A & \mathbb{I} \\ B & \mathbb{O} \end{bmatrix} \begin{bmatrix} \mathbb{I} & \mathbb{O} \\ \frac{1}{2}A & -B^T \end{bmatrix} =: B'(B'')^T,\]

i.e. $\bar{B}$ is expressed as a product of two $N\times{}2n$ matrices.

This is the same shape geodesic(::AbstractLieAlgHorMatrix) has, and for the same reason: the only matrix function either evaluates is on the $2n\times{}2n$ product $(B'')^TB'$. Getting there from the definition $\mathrm{Cayley}(\bar{B}) = (\mathbb{I} - \frac{1}{2}\bar{B})^{-1}(\mathbb{I} + \frac{1}{2}\bar{B})$ takes two steps. Write $C$ for the $2n\times{}2n$ inverse above. Splitting $\mathbb{I} + \frac{1}{2}\bar{B} = (\mathbb{I} - \frac{1}{2}\bar{B}) + \bar{B}$ gives $\mathrm{Cayley}(\bar{B}) = \mathbb{I} + (\mathbb{I} - \frac{1}{2}\bar{B})^{-1}\bar{B}$, and the Woodbury identity turns the remaining inverse into $(\mathbb{I} - \frac{1}{2}\bar{B})^{-1} = \mathbb{I} + \frac{1}{2}B'C(B'')^T$. The factor that leaves on the $2n\times{}2n$ side is then $\mathbb{I}_{2n} + \frac{1}{2}C(B'')^TB'$, which is $C$ itself, because $C$ inverts $\mathbb{I}_{2n} - \frac{1}{2}(B'')^TB'$.

The result is the $N\times{}N$ retraction of the lift, so the cost cannot fall below $O(N^2n)$, and this grouping reaches that floor: nothing larger than $N\times{}2n$ enters a product. Multiplying the factored left-hand side by an unfactored $\mathbb{I} + \frac{1}{2}\bar{B}$ instead multiplies two dense $N\times{}N$ matrices, which is $O(N^3)$ on its own and puts the retraction above geodesic at every size. scripts/cayley_regrouping_cost.jl carries the measurement.

Both sums are written as the five-argument mul! into an identity rather than as 𝕀 + X: the identity has to be built either way, and adding into it is what makes the product's own destination the answer instead of a temporary. scripts/retraction_step_allocations.jl carries that measurement. The $2n\times{}2n$ inverse stays an inv and is the allocation this cannot remove — lu! and rdiv! would take it in place on the host, and neither is something a KernelAbstractions backend is obliged to supply, where inv is what cayley already runs on Metal through.

cayley(Y::Manifold, Δ)

Take as input an element of a manifold Y and a tangent vector in Δ in the corresponding tangent space and compute the Cayley retraction.

In different notation: take as input an element $x$ of $\mathcal{M}$ and an element of $T_x\mathcal{M}$ and return $\mathrm{Cayley}(v_x).$

Examples

using GeometricOptimizers

Y = StiefelManifold([1. 0. 0.;]' |> Matrix)
Δ = [0. .5 0.;]' |> Matrix
Y₂ = GeometricOptimizers.cayley(Y, Δ)

Y₂' * Y₂ ≈ [1.;]

# output

true

See the example in geodesic(::Manifold, ::AbstractMatrix).

cayley_direct!(Y_new::StiefelManifold, Y::StiefelManifold, Δ::AbstractMatrix, ws::DirectCayleyWorkspace)

Write the Cayley retraction of the tangent vector Δ at Y into Y_new, and return Y_new.

The retraction is $(\mathbb{I}_M - \frac{1}{2}W)^{-1}(\mathbb{I}_M + \frac{1}{2}W)Y$ with the generator $W = \Delta{}Y^T - Y\Delta^T + Y\Delta^TYY^T$. It is computed without an $M\times{}M$ matrix and without the orthogonal complement of Y, through the rank-$2k$ factorization $W = UV$ with

\[A = \frac{1}{2}(Y^T\Delta - \Delta^TY), \quad K = \Delta - Y(Y^T\Delta), \quad U = [Y\;\; K], \quad V = \begin{bmatrix} AY^T - K^T \\ Y^T \end{bmatrix},\]

and the Sherman–Morrison–Woodbury identity

\[Y_\mathrm{new} = Y + U M_{2k}^{-1} VY, \qquad M_{2k} = \mathbb{I}_{2k} - \frac{1}{2}VU.\]

The cost is $O(Mk^2)$, and the only solve is the $2k\times{}2k$ system $M_{2k}$. That solve is an in-place LinearAlgebra.lu!, so the retraction runs only on a backend that supplies one; on JLArrays it raises a scalar-indexing error. For a Δ off the tangent space the result is the retraction of its tangent projection $\Delta - \frac{1}{2}Y(Y^T\Delta + \Delta^TY)$.

Y_new may be Y: the result is written into ws first and copied into Y_new.A last. A Δ, Y_new or ws of another size than Y raises a DimensionMismatch before anything is written.

cayley_direct(Y::StiefelManifold, Δ::AbstractMatrix)

The Cayley retraction of the tangent vector Δ at Y, as a new StiefelManifold. See cayley_direct!.

check(Y::Manifold)

Measure how far Y is from the manifold, as $\|Y^TY - \mathbb{I}\|$.

Two of the three manifolds this package provides store a representative whose columns are orthonormal — for StiefelManifold that is the point itself, for GrassmannManifold it is the representative of the equivalence class — so the same expression measures both. A retraction maps onto the manifold by construction, so in exact arithmetic this is zero and what it actually returns is accumulated round-off.

SymplecticStiefelManifold is the third, and its constraint is a different bilinear form, so it carries its own method. Anything defining a further manifold has to decide which of the two it is rather than inherit this one by default.

This is the assertion the manifold tests rest on, for every retraction path.

Examples

using GeometricOptimizers
using GeometricOptimizers: check
import Random
Random.seed!(123)

check(rand(GrassmannManifold, 5, 3)) < 1e-14

# output

true

check(U::SymplecticStiefelManifold)

How far U is from the manifold, as $\|U^T\mathbb{J}_{2N}U - \mathbb{J}_{2n}\|$.

This replaces the generic check(::Manifold), whose residual is the orthonormality one.

contains_nonfinite(a)

Whether a holds any value that is not finite.

NaN is the last thing a diverging solve produces: it reaches Inf first, and before that every finite magnitude on the way, so this tests isfinite and not isnan.

isfinite does not catch 1e200, so it is not the only guard; see convergence_measures for the one that does.

convergence_measures(status, config)

Checks if the optimizer converged.

Here status is an OptimizerStatus object and config is an SimpleSolvers.Options object.

Extended help

curvature_is_usable(ΔxΔg, Δx, Δg)

Whether the secant pair $(\delta, \gamma)$ satisfies the curvature condition, i.e. whether the quasi-Newton update built from it keeps $Q$ positive definite.

Both the BFGS and the DFP update divide by $\delta^T\gamma$ and preserve positive definiteness of $Q$ only for $\delta^T\gamma > 0$ [33]. A test of !iszero(ΔxΔg) admits both signs and denominators that are zero to within round-off, and a rank-two correction divided by such a denominator makes $Q$ badly conditioned or indefinite.

The threshold is relative: an absolute one cannot tell a round-off pairing on a problem scaled to $10^0$ from a legitimate small pairing on a problem scaled to $10^{-15}$. Its exact value matters much less than its existence: 1e-8 and eps(T) behave identically on everything measured here, because the pairs being rejected are the ones that are non-positive rather than merely small.

default_eltype(backend)

The element type a zeros or a rand that names no element type allocates in: Float64 on the host and Float32 on a device.

Neither value is arbitrary, which is the whole reason this is a function rather than a literal in each allocator. Float64 on the host is what zeros(n) and rand(n) already give, so an owned array or manifold allocated without an element type matches every other array allocated without one. Float32 on a device is the width an accelerator is built for, and a device that carries Float64 at all normally carries it at a fraction of the Float32 rate.

The rule deliberately does not ask KernelAbstractions.supports_float64. A backend being able to hold a Float64 is not a reason to hand it one: a caller who has not said which width it wants is better served by the width the device is fast at. An element type the caller does name and the backend cannot hold is refused by the backend's own allocation rather than narrowed, which the rand(backend, manifold_type, N, n) docstring states.

A caller who wants a fixed width names it, in the parametric form rand(backend, StiefelManifold{Float64}, N, n).

default_gradient(problem, x)

Return the SimpleSolvers.Gradient that Optimizer uses if none is supplied.

A NeuralNetworkParameters.NetworkParameters has no default gradient, and this throws an ArgumentError that names ∇F!: forward mode takes one pass per chunk of parameters, which is the wrong default for a network of more than a few hundred parameters. Supply the gradient, or opt in to forward mode or finite differences with the mode of Optimizer(x, F).

Implementation

The Manifold method is built from x: GradientAutodiff(problem.F, x) rebuilds the point before it calls problem.F, where the AbstractArray method would hand problem.F the flat vector.

On a device there is no default gradient, and every method throws an ArgumentError that names ∇F!: ForwardDiff and finite differences compute on the host.

default_linesearch(T, method)

Return the line search that Optimizer uses for method if none is supplied.

Everything except Adam defaults to SimpleSolvers.Backtracking(T; expand = true): the (quasi-)Newton methods because they build a direction with a scale of its own, and GradientMethod and MomentumMethod because a searching line search is what makes them converge rather than merely descend. All produce genuine descent directions, so a backtracking search always has an α to find.

expand = true is what lets that search lengthen a step as well as shorten it, and it is not the SimpleSolvers default — see the tip below for why it is the default here.

The AdamFamily methods are the exception and keep a fixed Static(DEFAULT_LEARNING_RATE). Adam's direction is $-m_1/(\sqrt{m_2} + \delta)$ and ScalarMomentAdam's is $-m_1/\sqrt{m_2 + \delta}$; either way it is a moving average that is deliberately not required to descend on any individual step, so a sufficient-decrease search has nothing to work with and would spend every such step reporting that it found no descent direction.

One caveat on sharing the value DEFAULT_LEARNING_RATE across the family: Adam normalizes componentwise and ScalarMomentAdam normalizes the whole lift at once, so the same $\eta$ is a step $\sqrt{\dim}$ shorter for the latter. See the warning in its docstring.

For AdamWithEuclideanDecay a fixed step is not merely the cheaper choice but the only one under which $\lambda$ means what it is documented to mean. Two reasons, and the first is the one that matters:

1. The merit the line search minimizes is the bare objective $f$. It is not $f + \frac{\lambda}{2}||x||^2$, because this package regularizes by adding $-\lambda{}x$ to the direction rather than by penalizing the objective — that is what decoupled means. So the penalty is invisible to the search, which will spend its $\alpha$ undoing the decay's contribution to $f$ as far as it can. The search is not approximating a line minimum of the function the method is actually descending.
2. The decay per step is $\alpha\lambda$, so a varying $\alpha$ makes the effective regularization strength whatever the search happened to pick on that step.

DecayingStatic is the exception to the exception: it varies $\alpha$ too, but on a schedule rather than in response to the merit, which is exactly AdamW's own learning-rate schedule and leaves $\lambda$ its meaning relative to $\eta$.

Pass linesearch = Static(η) for a fixed learning rate with any method.

The manual page Linesearches on Manifolds gives the evaluation counts on the SVD problem.

default_step_size(method)

The fixed step size a TrainingOptimizer takes for method if its linesearch keyword is not given: 1e-2 for GradientMethod, MomentumMethod and DirectMomentum, and DEFAULT_LEARNING_RATE = 1e-3 for the AdamFamily, whose direction has a magnitude that does not depend on the gradient. That magnitude differs between the members: $\approx{}1$ per component for Adam, $\approx{}1$ in total for ScalarMomentAdam and DirectScalarMomentAdam, and $\approx{}1$ per off-diagonal entry of the vertical part and per column of the horizontal part for FactoredStiefelAdam, $\approx{}k$ in total on $\mathrm{St}(k, M)$. At the same step size, their steps on a StiefelManifold are therefore shorter than Adam's.

A CompositeMethod steps all of its leaves at one step size, so its default is the one its two methods agree on, and an ArgumentError asking for linesearch where they disagree.

ensure_descent!(cache, method, config)

Replace the direction stored in cache by the steepest-descent direction if it does not descend.

A (quasi-)Newton step solves $H\delta = -\nabla{}f$ and is a descent direction only where $H$ is positive definite. Where it is not — an indefinite Hessian near a saddle, or a quasi-Newton $Q$ that has lost positive definiteness — $\delta$ points uphill. Since a line search only ever returns a non-negative step, the iteration then walks to the nearest stationary point, which may be a maximum.

F(x) = sum(sin.(x) .^ 2) started from x = ones(3) is the smallest example: $\sin^2$ has second derivative $2\cos(2x) = -0.83$ there, so the Newton direction ascends and the solve converges to $\pi/2$, where F is maximal. With this safeguard it descends to 0.

Descent is tested as $\nabla{}f\cdot\delta < 0$. rhs stores $-\nabla{}f$, so the test is $\mathrm{rhs}\cdot\delta > 0$; a direction that fails it — including a NaN one, for which every comparison is false — is replaced by rhs, which always descends. The substitution is reported at config.verbosity ≥ 2.

Examples

F(x) = sum(sin.(x) .^ 2)
x = ones(3)
algorithm = Newton()
state = OptimizerState(algorithm, x)
optimizer = Optimizer(x, F; algorithm = algorithm, linesearch = Bisection())

solve!(x, state, optimizer)
F(x) < 1e-20

# output

true

geodesic(B̄::AbstractLieAlgHorMatrix, algorithm = ScaledSquaring())

Compute the geodesic of a horizontal lift, i.e. $\exp(\bar{B})$.

Works for both StiefelLieAlgHorMatrix and GrassmannLieAlgHorMatrix — manifold_type supplies the manifold the result belongs to.

Implementation

Internally this is using

\[\exp(\bar{B}) = \mathbb{I} + B'\mathfrak{A}(B', B'')(B'')^T,\]

with $\bar{B} = B'(B'')^T$ the factorisation of lift_factors. The only matrix function this evaluates is $\mathfrak{A}$ on the $2n\times{}2n$ product $(B'')^TB'$, so the cost is set by $n$ and not by $N$, and so is the accuracy — see AbstractExponentialAlgorithm for the choice of algorithm and GeometricOptimizers.𝔄 for the implementations. The default ScaledSquaring is accurate at every lift norm.

geodesic(Y::Manifold, Δ, algorithm = ScaledSquaring())

Take as input an element of a manifold Y and a tangent vector in Δ in the corresponding tangent space and compute the geodesic (exponential map).

In different notation: take as input an element $x$ of $\mathcal{M}$ and an element of $T_x\mathcal{M}$ and return $\mathtt{geodesic}(x, v_x) = \exp(v_x).$

algorithm selects how the exponential is evaluated and is passed straight through to geodesic(::AbstractLieAlgHorMatrix); see AbstractExponentialAlgorithm for the choice. It matters only for a large $\Delta$ — the default ScaledSquaring is accurate at every step size.

Examples

using GeometricOptimizers

Y = StiefelManifold([1. 0. 0.;]' |> Matrix)
Δ = [0. .5 0.;]' |> Matrix
Y₂ = GeometricOptimizers.geodesic(Y, Δ)

Y₂' * Y₂ ≈ [1.;]

# output

true

Implementation

Internally this geodesic method calls geodesic(::AbstractLieAlgHorMatrix).

global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:GrassmannManifold{T}}

Express Δ (an element of the tangent space of Y) as an instance of GrassmannLieAlgHorMatrix.

The method global_rep for GrassmannManifold is similar to that for StiefelManifold.

Examples

using GeometricOptimizers
using GeometricOptimizers: _round
import Random

Random.seed!(123)

Y = rand(GrassmannManifold, 6, 3)
Δ = rgrad(Y, randn(6, 3))
λY = GlobalSection(Y)

_round(global_rep(λY, Δ); digits = 3)

# output

6×6 GrassmannLieAlgHorMatrix{Float64, Matrix{Float64}}:
  0.0     0.0    0.0     0.981   2.058  -0.4
  0.0     0.0    0.0    -0.424  -0.733   0.919
  0.0     0.0    0.0    -1.815  -1.409  -1.085
 -0.981   0.424  1.815   0.0     0.0     0.0
 -2.058   0.733  1.409   0.0     0.0     0.0
  0.4    -0.919  1.085   0.0     0.0     0.0

global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:StiefelManifold{T}}

Express Δ (an the tangent space of Y) as an instance of StiefelLieAlgHorMatrix.

This maps an element from $T_Y\mathcal{M}$ to an element of $\mathfrak{g}^\mathrm{hor}$.

These two spaces are isomorphic where the isomorphism where the isomorphism is established through $\lambda(Y)\in{}G$ via:

\[T_Y\mathcal{M} \to \mathfrak{g}^{\mathrm{hor}}, \Delta \mapsto \lambda(Y)^{-1}\Omega(Y, \Delta)\lambda(Y).\]

Also see GeometricOptimizers.Ω.

Examples

using GeometricOptimizers
using GeometricOptimizers: _round
import Random

Random.seed!(123)

Y = rand(StiefelManifold, 6, 3)
Δ = rgrad(Y, randn(6, 3))
λY = GlobalSection(Y)

_round(global_rep(λY, Δ); digits = 3)

# output

6×6 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}:
  0.0    -0.525  1.786   0.981   2.058  -0.4
  0.525   0.0    0.298  -0.424  -0.733   0.919
 -1.786  -0.298  0.0    -1.815  -1.409  -1.085
 -0.981   0.424  1.815   0.0     0.0     0.0
 -2.058   0.733  1.409   0.0     0.0     0.0
  0.4    -0.919  1.085   0.0     0.0     0.0

Implementation

The function global_rep does in fact not perform the entire map $\lambda(Y)^{-1}\Omega(Y, \Delta)\lambda(Y)$ but only

\[\Delta \mapsto \mathrm{skew}(Y^T\Delta),\]

to get the small skew-symmetric matrix $A\in\mathcal{S}_\mathrm{skew}(n)$ and

\[\Delta \mapsto (\lambda(Y)_{[1:N, n:N]}^T \Delta)_{[1:(N-n), 1:n]},\]

to get the arbitrary matrix $B\in\mathbb{R}^{(N-n)\times{}n}$.

global_section(Y::GrassmannManifold)

Compute a matrix of size $N\times(N-n)$ whose columns are orthogonal to the columns in Y.

The method global_section for the Grassmann manifold is equivalent to that for the StiefelManifold (we represent the Grassmann manifold as an embedding in the Stiefel manifold).

See the documentation for global_section(Y::StiefelManifold{T}) where T.

global_section(Y::StiefelManifold)

Compute a matrix of size $N\times(N-n)$ whose columns are orthogonal to the columns in Y.

This matrix is also called $Y_\perp$ [6, 8, 9].

Examples

using GeometricOptimizers
using GeometricOptimizers: global_section
import Random

Random.seed!(123)

Y = StiefelManifold([1. 0.; 0. 1.; 0. 0.; 0. 0.])

round.(global_section(Y); digits = 3)

# output

4×2 Matrix{Float64}:
  0.0     0.0
  0.0     0.0
 -0.936   0.353
 -0.353  -0.936

Implementation

Internally we do:

λ = orthonormal_columns() do
    A = randn(N, N - n) # or the gpu equivalent
    A - Y.A * (Y.A' * A)
end
_cholesky_qr2(λ - Y.A * (Y.A' * λ))

The projection and the orthonormalization are done twice. The first projection leaves a rounding error in the span of Y, and the orthonormalization multiplies it by the condition number of the projected draw, which has a heavy tail. Once only, the worst of a few thousand draws gives $\|Y^T\lambda\|$ of 3e-3 to 7e-2 in Float32, depending on the size. The first result is orthonormal, so the second orthonormalization amplifies nothing, and $\|Y^T\lambda\|$ stays at rounding level: measured over 20000 draws at $6\times3$ and 2000 at $50\times3$ and $200\times10$, at most 6.1e-7 in Float32 and 1.1e-15 in Float64, at twice the cost. Projecting twice before one orthonormalization does not do this, because the amplification comes after it.

The orthonormalization is CholeskyQR2 and not LinearAlgebra.qr!, on every backend — see _cholesky_qr2. qr! is a host factorization here: Metal implements no qr for its array type at all, so a qr! in this function puts GlobalSection(Y) and Optimizer(Y, F) out of reach on a device — and this function is on the path geodesic and cayley take on every step.

A is square inside the complement of Y, so its condition number has a square Gaussian's heavy tail and CholeskyQR2 breaks down on it about once in a hundred and twenty in Float32. A draw it cannot orthonormalize is replaced — see orthonormal_columns for the measurement and for why a redraw rather than a repair is the honest answer.

global_section(U::SymplecticStiefelManifold)

A symplectic completion of U: the remaining $2N - 2n$ directions, drawn at random and made symplectic against U. The result is $2N\times(2N - 2n)$, as global_section(::StiefelManifold) is $N\times(N - n)$, and it is what GlobalSection expects.

The counterpart of that method, with the symplectic form in place of the Euclidean one.

This one is host-only, and the other three of this type's operations are not

rgrad, metric and check run wherever the point is — metric as far as the backend supplies an lu, since it forms $\mathrm{inv}(U^TU)$. This does not, and the reason is the same one that makes rand(::AbstractRNG, ::GPU, ::Type{<:SymplecticStiefelManifold}, …) refuse: the completion is orthogonalized by the symplectic SR decomposition sr!, which is a host factorization — _rand_symplectic_stiefel calls Matrix on its factor. A device spelling would be a host computation with two transfers around it, which is a different operation from the device-native section global_section(::StiefelManifold) gives.

A device-backed point is therefore refused with an ArgumentError that says so, rather than left to fail inside sr! with Cannot access the contents of a private buffer, which names neither the manifold nor the call.

gradient(cache)

Return the stored gradient (array) of an OptimizerCache: $\nabla{}f(x_k)$ at the iterate the step is built from.

gradient_difference!(cache, state)

Write the gradient difference ∇f(xᵏ) - ∇f(xᵏ⁻¹) into cache.Δg, for OptimizerStatus.

Implementation

This is the difference of the two gradients the cache holds itself, latest_gradient at $x_{k+1}$ and gradient at $x_k$, which is the successive difference the status prints and needs no state.ḡ. It is not the γ of the quasi-Newton secant pair, which is one step behind the rg reported next to it. So the two g rows of a status are about one step; see convergence_measures for which iterate rg belongs to.

invalidate_latest_gradient!(cache)

Declare that latest_gradient is no longer the gradient at solution(cache).

isconverged(status)

Whether any of the three convergence flags convergence_measures sets is set.

The flags are a disjunction on purpose: x_converged, f_converged and g_converged test different things and a solve is entitled to stop on any one of them. Note that solve! can also stop for reasons that are not convergence — the iteration cap, a non-finite iterate, an increase in f where one is not allowed — and none of those sets a flag here, so this is what tells the two apart; see meets_stopping_criteria.

x_converged is the one to be careful with: "the iterate stopped moving" is a statement about a ratio, and a diverging solve is where the denominator stops meaning anything. See the warning under convergence_measures for the two guards on it and for the one case they do not cover.

latest_gradient(cache)

The array holding the most recently evaluated gradient: a line search trial point while linesearch_problem's $\varphi'$ is being evaluated, and the accepted iterate once solver_step! has called refresh_latest_gradient! on it.

This is the array OptimizerStatus reports as rg, and it is deliberately not gradient: the latter is $\nabla{}f(x_k)$ at the iterate the step is built from, which the direction and the state updates need to keep reading unchanged.

Implementation

Every cache in this package carries a scratch array of its own for this, the field g̃, and a flag g̃_is_current for the pairing store_gradient! relies on. The defaults of this method and of the three that go with it (refresh_latest_gradient!, latest_gradient_is_current, invalidate_latest_gradient!) read those two fields; see GradientCache for why g̃ may not be shared with gradient.

latest_gradient_is_current(cache, state, x)

Whether latest_gradient already holds $\mathrm{global\_rep}(\mathrm{section}(state), \nabla{}f(x))$, i.e. exactly what store_gradient! would otherwise evaluate.

lift_factors!(ws::RetractionWorkspace, B::AbstractLieAlgHorMatrix)

lift_factors written into ws, returning nothing.

Only the lift-dependent blocks are written — four for a Stiefel lift and two for a Grassmann one, whose $A$ block is identically zero. The identity and zero blocks were written when the workspace was built. ws has to have been built for B's shape.

lift_factors(B::AbstractLieAlgHorMatrix)

Factor a horizontal lift into the two $N\times{}2n$ matrices $B'$ and $B''$ with $\bar{B} = B'(B'')^T$:

\[\bar{B} = \begin{bmatrix} A & -B^T \\ B & \mathbb{O} \end{bmatrix} = \begin{bmatrix} \frac{1}{2}A & \mathbb{I} \\ B & \mathbb{O} \end{bmatrix} \begin{bmatrix} \mathbb{I} & \mathbb{O} \\ \frac{1}{2}A & -B^T \end{bmatrix} =: B'(B'')^T.\]

Every retraction in this package is built on this factorisation — it is what makes their cost scale with $n$ rather than with $N$, since the only matrix function either of them evaluates is on the $2n\times{}2n$ product $(B'')^TB'$. Both geodesic and cayley use it, for both manifolds.

For a GrassmannLieAlgHorMatrix this is the same expression with $A \equiv \mathbb{O}$.

Implementation

Both factors are allocated once, zeroed, and written block by block, rather than assembled out of nested vcats and hcats. The blocks are the same ones the formula above names; what goes is the chain of intermediate matrices each nesting level materialised — a vcat of two blocks, then an hcat of two of those, per factor, per call. The two zero blocks are then never written at all. This is on the per-iteration path: an optimizer step applies a retraction once per line-search trial and three times besides, and each of those calls this.

$(B'')^T$ is the matrix that is built, and $B''$ is returned as its adjoint, because every caller — cayley, geodesic, retraction_differential — consumes it as $(B'')^T$ and the two cancel exactly.

lift_from_columns(B, V)

Rebuild a lift of the same type as B from the $N\times{}n$ block V, i.e. from the first $n$ columns of a skew matrix. Used by retraction_differential, which is the only place a lift has to be assembled from its action on $E$ rather than from its own blocks.

Not to be confused with the canonical horizontal lift GeometricOptimizers.Ω, which maps a tangent vector to an element of $\mathfrak{g}^\mathrm{hor}$. This one takes the columns of a matrix that is already in the Lie algebra.

linesearch_parameters(cache, x, state, c)

The parameters solver_step! hands the line search: the iterate x, the state whose section the trial steps retract from, and — where a manifold supplies a scale — the step ceiling αmax.

Implementation

Three methods, chosen on the type of the solution, as trial_iterate!'s two are.

For a Manifold, αmax is step_αmax(c, direction(cache)). It is built again at every solver step, because $\|\delta\|$ changes at every solver step — also between the two searches of one solver_step!, where the second runs on the direction that steepest_descent! has replaced.

For a NamedTuple it is _manifold_αmax over the blocks, i.e. the same quantity derived per block and minimised over the manifold ones. A NamedTuple carrying no manifold block therefore gets Inf, which is what upstream reads as "the caller has no scale of its own" (SimpleSolvers.linesearch_αmax) and is equivalent to the omission below.

For an AbstractVector the field is omitted. Upstream reads it through a hasproperty guard that constant-folds, so a caller that supplies nothing pays nothing. A Euclidean $f(x + \alpha{}p)$ grows with $\alpha$, so the sufficient-decrease test of the search rejects an over-long step. There is no geometric scale to supply, and the method's own SimpleSolvers.DEFAULT_LINESEARCH_αmax is the whole of the bound.

Every branch returns a concrete NamedTuple, so the merit closures stay type-stable. That is why the manifold-free NamedTuple passes Inf rather than joining the AbstractVector branch: deciding the shape of the parameters on the block types makes the return type a Union of two NamedTuples, which the merit closures then pay for on every evaluation.

linesearch_rejected(status)
linesearch_rejected(status, αmax)

Whether a line search reported that it could not decrease the merit along the direction it was given.

The second form takes the ceiling the caller passed as params.αmax and is what solver_step! uses; see A step the caller forbade is not a failed direction below.

SimpleSolvers.solve returns only the step length, so an outcome of LINESEARCH_FLOOR, LINESEARCH_EXHAUSTED or LINESEARCH_NO_DESCENT looks the same as a successful search. The three mean, in order, that the decrease achieved was no larger than the merit's own round-off resolution, that the budget ran out or the merit could not be bracketed, and that $\varphi'(0) \geq 0$. In none of them does the returned $\alpha$ carry a guarantee, and a solve that takes it can walk uphill.

The outcome comes from SimpleSolvers.solve_with_status, which solver_step! calls in place of solve for this reason. (Plain code and not an @extref: SimpleSolvers documents solve_with_status per method and not as a binding, so its inventory has no entry to link to.)

manifold_constructor(x::Manifold)

The one-argument constructor of x's manifold: GrassmannManifold for a GrassmannManifold{Float32, Matrix{Float32}}.

The type name and not typeof(x), because the array the result is applied to may have a different element type from x's: the closure GradientAutodiff differentiates is handed a vector of ForwardDiff.Duals. And x's manifold rather than a hardcoded StiefelManifold, so that a bare GrassmannManifold reaches Optimizer too.

Its callers are GeometricOptimizers._similar(::Manifold), copyto!(::Manifold, ::Manifold) and the reconstruction that GradientAutodiff(F, ::Manifold) and GradientFiniteDifferences(F, ::Manifold) compose F with. The flattening does not call it: NeuralNetworkParameters.rebuild takes a prototype rather than a type, so a round trip keeps a GrassmannManifold a GrassmannManifold.

manifold_type(B::AbstractLieAlgHorMatrix)

The manifold a retraction of B lands on.

$\mathfrak{g}^\mathrm{hor}$ is the horizontal component of the Lie algebra of a specific homogeneous space, so the lift already determines where the retraction maps to. This is what lets geodesic and cayley be written once for both manifolds rather than twice each.

meets_stopping_criteria(status, config, iterations)

Check if the optimizer has converged.

Implementation

meets_stopping_criteria checks if one of the following is true:

• converged (the output of SimpleSolvers.assess_convergence) is true and iterations $\geq$ config.min_iterations,
• if config.allow_f_increases is false: status.f_increased is true,
• iterations $\geq$ config.max_iterations,
• status.rfₐ $>$ config.f_abstol_break,
• any of status.x_nonfinite, status.f_nonfinite, status.g_nonfinite.

Extended help

metric(Y::GrassmannManifold, Δ₁::AbstractMatrix, Δ₂::AbstractMatrix)

Compute the metric for vectors Δ₁ and Δ₂ at Y.

The representation of the Grassmann manifold is realized as a quotient space of the Stiefel manifold.

The metric for the Grassmann manifold is:

\[g^{Gr}_Y(\Delta_1, \Delta_2) = g^{St}_Y(\Delta_1, \Delta_2) = \mathrm{Tr}(\Delta_1^T (\mathbb{I} - Y Y^T) \Delta_2) = \mathrm{Tr}(\Delta_1^T \Delta_2),\]

where we used that $Y^T\Delta_i$ for $i = 1, 2.$

metric(Y::StiefelManifold, Δ₁::AbstractMatrix, Δ₂::AbstractMatrix, ::CanonicalMetric)
metric(Y::StiefelManifold, Δ₁::AbstractMatrix, Δ₂::AbstractMatrix, ::EuclideanMetric)

The dot product of Δ₁ and Δ₂ at Y in the given metric: the canonical one of metric(::StiefelManifold, ::AbstractMatrix, ::AbstractMatrix), or $\mathrm{Tr}(\Delta_1^T\Delta_2)$. A Δ₁ or Δ₂ of another size than Y raises a DimensionMismatch.

metric(Y::StiefelManifold, Δ₁::AbstractMatrix, Δ₂::AbstractMatrix)

Compute the dot product for Δ₁ and Δ₂ at Y.

This uses the canonical Riemannian metric for the Stiefel manifold:

\[g_Y: (\Delta_1, \Delta_2) \mapsto \mathrm{Tr}(\Delta_1^T(\mathbb{I} - \frac{1}{2}YY^T)\Delta_2).\]

metric(U::SymplecticStiefelManifold, Δ₁::AbstractMatrix, Δ₂::AbstractMatrix)

The Riemannian metric of the symplectic Stiefel manifold, taken from [14]:

\[g_U(\Delta_1, \Delta_2) = \mathrm{tr}\left( P \Delta_1^T \left(\mathbb{I}_{2N} - \frac{1}{2}\mathbb{J}_{2N}^TUPU^T\mathbb{J}_{2N}\right)\Delta_2 \right), \qquad P = (U^TU)^{-1}.\]

Implementation

Every factor the trace needs is $2n\times{}2n$, and the expression is evaluated that way:

\[g_U(\Delta_1, \Delta_2) = \mathrm{tr}\left( P \left( \Delta_1^T\Delta_2 - \frac{1}{2} \left(\Delta_1^T\mathbb{J}_{2N}^TU\right) P \left(U^T\mathbb{J}_{2N}\Delta_2\right) \right) \right).\]

Written as the definition reads, the middle factor is a $2N\times{}2N$ matrix, and forming it costs $O(N^3)$ — a matrix of the ambient dimension to reach a number, on every step the optimizer takes. Grouped this way no product has two factors of the ambient dimension; the largest is $\mathbb{J}_{2N}^TU$, and it is formed once, because $U^T\mathbb{J}_{2N}$ is its adjoint. The $2N\times{}2N$ identity is not built at all, and $P$ is inverted once rather than twice. $\mathbb{J}_{2N}$ is still assembled dense by _poisson_tensor, and it is what a call allocates most of at the larger sizes. scripts/symplectic_metric_cost.jl carries the measurement.

observe_optimizer_phase(f, observer, phase)

Run f() between matching observer(phase, :enter) and observer(phase, :exit) notifications. The exit notification is guaranteed even when f throws.

opnorm₁(X)

The induced 1-norm of X, i.e. its largest absolute column sum, as a reduction.

LinearAlgebra.opnorm(X, 1) is the natural spelling and is not used, because LinearAlgebra.opnorm1 is a double loop over X[i, j]. Scalar indexing is precisely what an array on a GPU backend cannot serve, and being free of it is why ScaledSquaring and NativePade have no dense-LAPACK dependency at all — so the one norm they take has to be expressible as sum and maximum. Accelerator execution still depends on the array backend's support for those reductions.

The two algorithms do not call opnorm₁: GeometricOptimizers.𝔄! takes the same reduction with the column sums written into its scratch, and uses the value only to pick the number of halvings s = ⌈log₂(‖X‖₁/θ)⌉. On a real X that value is opnorm₁'s to the bit; on a complex X the two can differ in the last bit.

opnorm₁ and opnorm1 agree to a few eps, not bitwise: opnorm1 accumulates each column sequentially in at least Float64, whereas sum is pairwise and accumulates in eltype(X). So a difference of an ulp can at most shift s by one, and only for an argument that lands exactly on a power of two.

optimization_step!(x, opt::TrainingOptimizer, dp)

Take one step of the training optimizer opt with the Euclidean gradient dp of the loss at x, and write the new parameters into x, which it returns.

x is the object opt was built on, and nothing else changes it between steps: the step starts from the point opt.state carries, which is x after the previous step. dp has the shape and the element type of x. An x or a dp of another element type than opt is an ArgumentError. A dp with arrays of other sizes is a DimensionMismatch, and a parameter set with other keys or another number of leaves is an ArgumentError, and so is a dp with an entry that is not finite for a GeometricOptimizers.DirectStiefelMethod. Each is raised before the iteration number moves.

The step increments the iteration number t of opt.state, reads the step size step_size(opt.linesearch, t), forms the direction of opt.method from the PrecomputedGradient of dp, scales it, retracts it with opt.retraction and copies the result into x. advance_state! then carries the state over to the new parameters. A GeometricOptimizers.DirectStiefelMethod forms the Riemannian gradient of dp in its cache instead, retracts with cayley_direct! and transports its first moment to the new point with the transport of the method. A CompositeMethod takes this step on every leaf, with the method, the cache and the state of that leaf.

For a parameter set, dp is a NeuralNetworkParameters.NetworkParameters of the same shape as x. See TrainingOptimizer for why this is not solve!.

orthonormal_columns(draw)

Orthonormalize draw() with _cholesky_qr2, drawing again while the draw is too ill-conditioned for it.

draw takes no argument and returns an $N\times{}m$ matrix whose entries are a fresh draw each time it is called. Only the entries have to be fresh: nothing here writes into the matrix, so a draw that fills one buffer in place and returns it every time is correct. The result has orthonormal columns and is on whatever backend draw allocated on, so this is the entry point for orthonormalizing on a device: LinearAlgebra.qr! is a host factorization, and Metal implements no qr for its arrays at all. Exhausting the attempts throws OrthonormalizationFailure.

draw is called again rather than its result repaired, so each call must produce new entries and not hand back the previous draw. This package's own callers — rand(backend, manifold_type, N, n) and global_section — each draw a Gaussian matrix, which carries no information, so one CholeskyQR2 cannot orthonormalize is replaced; nothing is repaired and no result is kept that the algorithm did not produce cleanly. A draw that ignores this and returns the same entries every time turns a breakdown into the error below instead of an answer.

The rate is not negligible. global_section factorizes $N\times(N-n)$ Gaussian columns with the span of Y projected out, so the matrix is square inside that complement and its condition number has the heavy tail a square Gaussian's does. Measured in Float32 over 9000 draws at $N = 50, 100, 200$ and $n = 3$, across three seeds: 77 break CholeskyQR2 down, a rate between one in 107 and one in 125. Each is replaced, and three of those 77 replacements fail as well — a higher rate than the draws themselves, which is why the attempt bound is not derived from one failure probability raised to the eighth power. The script is scripts/orthonormalization_breakdown_rate.jl.

A redraw costs the caller determinism, not just time. The number of Gaussian draws global_section consumes in Float32 depends on the draws themselves, so a seeded Float32 computation downstream of one is not reproducible across a change to this function or to the element type. The device sweep in scripts/device_products.jl seeds one, and relies only on the same code repeating its draws.

Shifted CholeskyQR3 measured on the same draws does not close it — one failure in 300 even with the exact $\|A\|_2$ in the shift, because forming $A^TA$ in Float32 loses a singular value that small whatever the shift is.

refresh_latest_gradient!(cache, gradient_instance)

Evaluate the gradient at the iterate cache currently holds and store it in latest_gradient.

solver_step! calls this once the accepted step has been taken, so that rg is a statement about the point the solve is about to report rather than about the one it started the step from.

It is also what establishes the pairing store_gradient! relies on. See latest_gradient.

restart!(state::BFGSState)

Reset the inverse-Hessian approximation to the identity, so that the next direction is the steepest-descent one.

This is the BFGSState method of restart!; DFPState is an alias for BFGSState, so it covers DFP as well. solver_step! calls it when a line search reports that it could not decrease the merit — see linesearch_rejected.

restart!(state)

Discard whatever curvature information state has accumulated and start again from the identity.

Only the quasi-Newton states carry any, so this is a no-op for everything else. See restart!(::BFGSState) and solver_step!.

retraction(R::AbstractRetraction, x)

Apply the retraction R to x, i.e. dispatch on the retraction type rather than calling geodesic or cayley by name.

This is what a caller who has been handed a retraction = … keyword uses: the two shipped types, Geodesic and Cayley, select the two functions, and a Geodesic also carries the AbstractExponentialAlgorithm its geodesic is evaluated with. R(x) is the same thing written as a call.

x is an AbstractLieAlgHorMatrix on a manifold — the retractions this package is about map $\mathfrak{g}^\mathrm{hor}\to{}G$ — a NamedTuple of parameters, or an ordinary array, on which every retraction is the identity because the extended retraction on a vector space is addition and update_section! does the adding.

An R and an x that do not go together is an error that says so, rather than a nothing that fails further downstream.

retraction_differential!(D, workspace, retraction, B, α)

retraction_differential(retraction, B, α) written into D, which has B's type and shape, and returned.

workspace is the RetractionWorkspace of B's point, the tree of them a parameter set has, or nothing. A Cayley lift at $\alpha \neq 0$ with a workspace on a host Matrix of a LAPACK element type, Float32, Float64 or their complex types, is evaluated in that workspace and allocates nothing: its two solves are LAPACK.getrf! and LAPACK.getrs! on workspace buffers, which are the two calls \ makes, so the answer is the allocating method's to the bit. Every other case copies the allocating method's answer into D — at $\alpha = 0$, under Geodesic and for an ordinary array that is a copy of B — including a workspace on any other array type, where \ stays.

retraction_differential(retraction, B, α)

The horizontal lift $D(\alpha)$ that generates the velocity of the curve a line search walks along.

trial_iterate! builds its trial point as $x(\alpha) = \Lambda_0\mathrm{retract}(\alpha\bar{B})E$ with $\Lambda_0$ the section the step starts from, so with $W(\alpha) = \Lambda_0\mathrm{retract}(\alpha\bar{B})$ the frame the cache holds afterwards, this returns the $D(\alpha)$ for which

\[\frac{dx}{d\alpha} = W(\alpha)D(\alpha)E .\]

trial_slope pairs the gradient against it, which is what makes $\varphi'(\alpha)$ the derivative of $\varphi$.

Implementation

For Geodesic this is $\bar{B}$ itself. The exponential is a one-parameter subgroup, so $\frac{d}{d\alpha}\exp(\alpha\bar{B}) = \exp(\alpha\bar{B})\bar{B}$ and the frame factors out whole. The same holds for an ordinary array under either retraction, where the retraction is addition and $\varphi$ is affine in $\alpha$.

Cayley is not a one-parameter subgroup, and this is what that costs. With $M = (\mathbb{I} - \frac{\alpha}{2}\bar{B})^{-1}$,

\[\frac{d}{d\alpha}\mathrm{Cayley}(\alpha\bar{B}) = M\bar{B}M , \qquad D(\alpha) = \mathrm{Cayley}(\alpha\bar{B})^{-1}M\bar{B}M = M^T\bar{B}M ,\]

using $M^T = (\mathbb{I} + \frac{\alpha}{2}\bar{B})^{-1}$ for skew $\bar{B}$. $D$ is skew but not horizontal; only its first $n$ columns ever enter — $D(\alpha)E$ — and the horizontal projection has the same ones, so what lift_from_columns is handed is that projection.

For a GrassmannLieAlgHorMatrix the top $n\times{}n$ block is dropped rather than kept, and that is not an approximation. $M$, $M^T$ and $\bar{B}$ are all rational functions of $\bar{B}$ and so commute, which gives

\[D(\alpha) = \bar{B}\left(\mathbb{I} - \tfrac{\alpha^2}{4}\bar{B}^2\right)^{-1} ,\]

and for a Grassmann lift — where $A \equiv \mathbb{O}$, so $\bar{B}^2$ is block-diagonal — the top-left block of that product is identically zero at every $\alpha$. The drop is exact whatever $D$ is paired against, not merely exact against a gradient that happens to be orthogonal to $Y$.

Both inverses are taken through lift_factors and the Woodbury identity, exactly as cayley does, so the cost is $O(Nn^2 + n^3)$ and no $N\times{}N$ matrix is formed.

α = 0 returns B unchanged, and that is the case that matters for cost: SimpleSolvers.Backtracking — the default of default_linesearch — evaluates $\varphi'$ only at $\alpha = 0$, so it never reaches the general branch.

retraction_matrix!(ws::RetractionWorkspace, R::AbstractRetraction, B::AbstractLieAlgHorMatrix)

The $N\times{}N$ matrix retraction(R, B) wraps in a manifold type, written into ws.retracted and returned bare.

The matrix and not the manifold: update_section! reads the blocks of the result and discards the wrapper, so wrapping it would be one allocation per line-search trial for a value nothing keeps.

An AbstractRetraction this package does not ship has no in-place form here and falls through to the allocating retraction, with the answer copied in. A workspace never changes what a retraction means. A bare callable is not covered — update_section! accepts one as its retraction, and src/utils.jl passes one, but only ever together with no workspace, so there is no retraction type left here to dispatch on and a caller who pairs the two gets a MethodError rather than a silent fallback.

ws has to have been built for B's shape, as for lift_factors!, and every arm asserts it. The assertion is not decoration: the fallback arms reach ws.retracted through copyto!, and copyto! into an oversized destination copies linearly rather than throwing, so a workspace of the wrong shape would scramble the layout of the answer instead of rejecting it.

retraction_workspace(x, retraction = nothing)
retraction_workspace(opt::Optimizer)

The RetractionWorkspace a solve over x with retraction needs, or NoWorkspace where x carries no manifold. On an Optimizer it is the accessor instead, returning the workspace that optimizer was built with.

Optimizer builds one of these for its retraction at construction and hands it to every update_section! on the step path. A parameter set gets a tree of them in the shape its section tree has, a NoWorkspace at every leaf that is an ordinary array — the extended retraction on a vector space is addition, which allocates nothing to begin with.

The two meanings share a name and cannot collide: Optimizer is not a member of OptimizerSolution, which is what the building methods dispatch on.

rgrad(Y::GrassmannManifold, ∇L::AbstractMatrix)

Compute the Riemannian gradient for the Grassmann manifold at Y based on ∇L.

Here $Y$ is a representation of $\mathrm{span}(Y)\in{}Gr(n, N)$ and $\nabla{}L\in\mathbb{R}^{N\times{}n}$ is the Euclidean gradient.

This gradient has the property that it is orthogonal to the space spanned by $Y$.

The precise form of the mapping is:

\[\mathtt{rgrad}(Y, \nabla{}L) \mapsto \nabla{}L - YY^T\nabla{}L.\]

Note the property $Y^T\mathrm{rgrad}(Y, \nabla{}L) = \mathbb{O}.$

Also see rgrad(::StiefelManifold, ::AbstractMatrix).

Examples

using GeometricOptimizers

Y = GrassmannManifold([1 0 ; 0 1 ; 0 0; 0 0])
Δ = [1 2; 3 4; 5 6; 7 8]
rgrad(Y, Δ)

# output

4×2 Matrix{Int64}:
 0  0
 0  0
 5  6
 7  8

rgrad(Y::StiefelManifold, ∇L::AbstractMatrix, ::CanonicalMetric)
rgrad(Y::StiefelManifold, ∇L::AbstractMatrix, ::EuclideanMetric)

The Riemannian gradient at Y for the given metric, from the Euclidean gradient ∇L:

\[\nabla{}L - Y(\nabla{}L)^TY \quad\text{(canonical)}, \qquad \nabla{}L - \frac{1}{2}Y(Y^T\nabla{}L + (\nabla{}L)^TY) \quad\text{(Euclidean)}.\]

The Euclidean form is the orthogonal projection of ∇L onto the tangent space at Y. Both forms satisfy $\mathtt{metric}(Y, \mathtt{rgrad}(Y, \nabla{}L, m), \Delta, m) = \mathrm{Tr}((\nabla{}L)^T\Delta)$ for every tangent vector $\Delta$. A ∇L of another size than Y raises a DimensionMismatch.

rgrad(Y::StiefelManifold, ∇L::AbstractMatrix)

Compute the Riemannian gradient for the Stiefel manifold at Y based on ∇L.

Here $Y\in{}St(N,n)$ and $\nabla{}L\in\mathbb{R}^{N\times{}n}$ is the Euclidean gradient.

The function computes the Riemannian gradient with respect to the canonical metric: metric(::StiefelManifold, ::AbstractMatrix, ::AbstractMatrix).

The precise form of the mapping is:

\[\mathtt{rgrad}(Y, \nabla{}L) \mapsto \nabla{}L - Y(\nabla{}L)^TY\]

Note the property $Y^T\mathtt{rgrad}(Y, \nabla{}L)\in\mathcal{S}_\mathrm{skew}(n).$

Examples

using GeometricOptimizers

Y = StiefelManifold([1 0 ; 0 1 ; 0 0; 0 0])
Δ = [1 2; 3 4; 5 6; 7 8]
rgrad(Y, Δ)

# output

4×2 Matrix{Int64}:
 0  -1
 1   0
 5   6
 7   8

rgrad(U::SymplecticStiefelManifold, ∇L::AbstractMatrix)

The Riemannian gradient of a point of the symplectic Stiefel manifold, for the metric of metric(::SymplecticStiefelManifold, ::AbstractMatrix, ::AbstractMatrix).

$\nabla{}L$ is the Euclidean gradient, i.e. the derivative of the loss with respect to the entries of U read as an unconstrained matrix.

rhs(cache)

Return the right hand side of an instance of NewtonOptimizerCache

rhs(cache)

Return the right hand side of an instance of QuasiNewtonCache

solution_scale(x)

The norm the relative change in the iterate is measured against, i.e. the denominator of $\|x - x'\|/\|x'\|$.

For an ordinary array this is the l2norm of the iterate, which is the only scale a Euclidean problem supplies. For a Manifold it is not: a point of $St(N, n)$ or $Gr(n, N)$ satisfies $Y^TY = \mathbb{I}_n$, so $\|Y\|_F = \sqrt{\mathrm{tr}(Y^TY)} = \sqrt{n}$ exactly, and the scale is a constant the geometry supplies rather than something to be measured off the iterate.

That distinction is what keeps x_converged honest on a solve that has diverged; see convergence_measures. While the iterate is on the manifold the two agree to round-off, so this changes no measurement of any converging solve — it changes what is reported once the iterate is somewhere $\sqrt{n}$ no longer describes.