The leaf protocol

freeparameters(x)

The differentiable storage of a leaf x — the numbers that a flat parameter vector should contain, and the coordinates an optimizer should work in.

For an ordinary array this is the array itself. For a structured type it is whatever the type keeps its degrees of freedom in: an $n \times n$ symmetric matrix stores $n(n+1)/2$ numbers, and it is those that belong in the flat vector — not the $n^2$ entries of the dense interface, which do not even have the right length.

The return value may be

• the leaf itself, which marks it as terminal — the recursion stops and the numbers are copied straight out of it;
• another array or number;
• a Tuple or NamedTuple of either, for a type whose freedom lives in several blocks.

Together with rebuild this is the whole extension protocol: define the two for a type and it can be flattened, walked and written to HDF5, with nothing in this package knowing about it.

Extending

NeuralNetworkParameters.freeparameters(A::SymmetricMatrix) = A.S
NeuralNetworkParameters.rebuild(A::SymmetricMatrix, data)  = SymmetricMatrix(data, A.n)

Those two are the whole protocol for a leaf that is an AbstractArray subtype, which every structured parameter type in the ecosystem is. A leaf that keeps its numbers behind some other interface needs one method more, because parameter_eltype must not raise and so cannot ask an arbitrary type where its storage is:

NeuralNetworkParameters.parameter_eltype(x::MyLeaf) = parameter_eltype(freeparameters(x))

Without it the leaf contributes no element type, and flatten says so rather than guessing one.

GeometricOptimizers already exposes exactly this relation as Base.parent for its manifolds, VectorStorageMatrixes and horizontal lifts, so one delegating method covers all of them:

NeuralNetworkParameters.freeparameters(x::Union{Manifold, VectorStorageMatrix, AbstractLieAlgHorMatrix}) =
    parent(x)

Base.parent is deliberately not used as the protocol here: parent(::SubArray) is the whole underlying buffer rather than the view's own entries, so Base's relation and this one are not the same.

Examples

using NeuralNetworkParameters: freeparameters

A = [1.0 2.0; 3.0 4.0]
freeparameters(A) === A

# output

true

rebuild(prototype, data)

Rebuild a leaf shaped like prototype from the storage data. The inverse of freeparameters, and for an ordinary array simply data.

prototype is the leaf the storage came from, which is what carries the non-differentiable information along: the n of a SymmetricMatrix, the N and n of a horizontal lift. Taking them from a prototype rather than from the type is what lets data have a different element type from prototype — ForwardDiff.Duals, when a flattened parameter set is differentiated through — and it is also why the concrete type comes back unchanged, where reconstructing from a type name can quietly return a sibling type instead.

See freeparameters for how to extend this.

storage_gradient(leaf, Δ)

The gradient with respect to the storage of leaf — its freeparameters — from the cotangent Δ that reverse-mode differentiation gives for it.

A structured leaf presents one interface and stores another, and AD differentiates the interface. For a symmetric matrix with storage $S$, the cotangent is the natural one, a matrix $G$ paired with the dense interface; the parameter gradient is $\partial L/\partial S$, which counts each stored off-diagonal entry twice: $G_{ij} + G_{ji}$. The two differ by a factor of two off the diagonal, so a type whose storage is not its interface defines a method that converts one to the other.

This package calls it exactly once per top-level leaf of a set, on the accumulated cotangent, through map_cotangent, at the two places where an AD cotangent becomes a parameter gradient: the gradient that the ZygoteRules extension returns for a NetworkParameters, and the flat gradient that the reverse rule of unflatten returns. The flat gradient is the structured one, flattened. It dispatches on the primal leaf, because the type of Δ depends on how the loss used the leaf: a leaf used twice gets a dense Matrix, whatever its own type is. It is never called with a structural zero (nothing or a ChainRulesCore.AbstractZero): an untouched leaf gets a zero leaf of its type instead. It is not called for a leaf whose cotangent is part of a NetworkParameters, which is what the reverse rule of flatten returns: its leaves hold the storage gradient already.

The extension converts the gradient of Zygote.pullback(f, ps) and Zygote.gradient(f, ps) for a Function f and a set ps that is the only argument, which is how a loss is differentiated with respect to a set. Zygote's own methods answer every other call: a further argument, a callable struct for f, or a set held inside another argument. Those return the structural tangent (params = …,) with the natural cotangent at each leaf, and do not call this function.

A method converts an array cotangent only, and returns a leaf of the same type, whose storage the flat path reads with freeparameters. Its element type is the leaf's, parameter_eltype(leaf), or float of it for an Integer leaf, whatever the precision of the cotangent: a loss that multiplies a Float32 leaf by a Float64 constant gives that leaf a Float64 cotangent. The default method converts an array or a number cotangent of an AbstractFloat or Integer element type to that element type, with no copy where the types match, as ChainRulesCore.ProjectTo does. A cotangent of another number type, such as a ForwardDiff.Dual from forward-over-reverse differentiation, passes through unchanged. A gradient whose leaves end up with two element types raises the ArgumentError of NetworkParameters for two element types.

A method reads the type, the shape and the metadata of leaf, never its numbers. The sum of a gradient and a structural tangent of the same set converts the tangent with the gradient's leaves in place of the primal ones.

A structural tangent over the leaf's fields, a NamedTuple or a ChainRulesCore.Tangent, comes from a loss that read the storage field directly and holds $\partial L/\partial S$ already. The default method for it rebuilds the leaf with rebuild around the components of the tangent that belong to its storage — matched by key for a NamedTuple of blocks, by field name for a Tuple of blocks, whose i-th block must be the leaf's i-th field, and for one block by the one field it is — each converted by this function in turn, with a zero block where the tangent has none. A storage that meets none of these raises an ArgumentError. So does a leaf whose one field is a Tuple of blocks that is its storage: give such a leaf a storage_gradient method of its own, or keep its blocks in separate fields.

Extending

function NeuralNetworkParameters.storage_gradient(A::SymmetricMatrix, G::AbstractMatrix)
    # a `SymmetricMatrix` whose storage is the lower triangle of `G + transpose(G)`, with the
    # diagonal counted once
end

map_cotangent(f, ps, Δ)

Walk the branches of the primal ps and of its cotangent Δ together, and return the cotangent with f(leaf, Δleaf) at each leaf. The result has the shape of the primal: keyed by the primal's keys for a NamedTuple, positional for a Tuple, and a NetworkParameters for a NetworkParameters.

A hole is nothing, any flavour of zero, or a Tangent with no fields. Where Δ as a whole is a hole, the result is nothing. A hole inside Δ, at a leaf or a whole branch, comes back as a zero leaf of each leaf's own type, and f never sees it. Its element type is that of the leaves f returned, so a gradient of ForwardDiff.Duals has Dual zeros; where f returned no numbers, it is the leaf's own, or float of it for an Integer leaf. A leaf with no numbers, whose parameter_eltype is Union{} — nothing, or a function — has no zero and stays nothing. So the result is complete wherever Δ touched the set at all.

Δ is the cotangent of ps itself, in the shape Zygote gives it, and a cotangent of another shape raises an ArgumentError at the first branch it does not fit. At a NamedTuple branch the cotangent is a NamedTuple, or a Tangent over one, whose keys are keys of the branch; a key it leaves out is a hole, as a field left out of a Tangent is. At a Tuple branch it is a Tuple of the same length, or a Tangent over one. At a NetworkParameters it is the set's structural tangent (params = …,), as a NamedTuple or a Tangent, or a NetworkParameters. So (params = …,) passed against params(ps) raises, and so does a NetworkParameters gradient passed against params(ps).

With f =storage_gradient, this converts the natural cotangent Zygote gives a parameter set to its gradient with respect to the storage of each leaf. A package that writes its own ZygoteRules.pullback for a function of a parameter set calls it on the cotangent of the set:

y, pb = ZygoteRules.pullback(g, x, values(ps))
function g_pullback(ȳ)
    x̄, p̄ = pb(ȳ)
    p̄ = map_cotangent(storage_gradient, values(ps), p̄)
    x̄, p̄ === nothing ? nothing : NetworkParameters(NamedTuple{keys(ps)}(p̄))
end

A NetworkParameters inside Δ holds the storage gradient at every leaf already — it is what the reverse rule of flatten returns — so map_cotangent(storage_gradient, …) returns it as it is.

parameter_eltype(ps)

The element type to flatten ps into: the element type of its leaves, which is one type. Leaves of two numeric element types raise an ArgumentError that names both; nothing is promoted.

Float32 parameters therefore flatten to a Vector{Float32}. Defaulting to Float64 instead — as ParameterHandling.flatten does — silently doubles the width of every single-precision network that passes through, and every quantity computed from the flat vector downstream with it.

Examples

using NeuralNetworkParameters: parameter_eltype

parameter_eltype((a = Float32[1, 2], b = Float32[3;;]))

# output

Float32

Implementation

This function raises only for two element types, where freeparameters raises for a leaf without the protocol: every NetworkParameters runs its constructor through here, including sets that hold no numbers at all — a tree with nothing in place of a leaf, or SymbolicNeuralNetworks wrapping generated functions in one. A leaf this package cannot read numbers out of contributes no element type, exactly as an empty set does, and both report Union{}. A set that cannot be flattened is still a set with an element type; the protocol error comes from parameterlayout, which decides something with it.

The recursion follows freeparameters for an AbstractArray leaf — every structured parameter type in the ecosystem — and a leaf that keeps its numbers behind another interface opts in with one method more, described under freeparameters. Asking an arbitrary type where its storage is would mean raising, which this function must not do.

It follows it at the level of values, freeparameters(x) === x, and that is a guarantee rather than an implementation detail. A structured leaf reports the element type of its storage, which need not be the eltype of the interface it presents: a leaf that is an AbstractMatrix{Float64} over a Vector{Float32} flattens into a Vector{Float32}, because that is what its numbers are. Deciding the element type from the leaf's type instead would read the interface and be wrong about such a leaf, which is why it is not decided there — and there is no cost to buy with it.

For a NetworkParameters the answer is already on the type, put there by its constructor, so nothing is recomputed and the call folds to a constant. A bare NamedTuple or Tuple is read in place at literal indices, so it too costs nothing at any width, depth or shape, and flatten(ps) costs what flatten(T, ps) costs.