Library

NeuralNetworkParameters

The parameters of a neural network, in two shapes and with conversions between them.

• NetworkParameters is the structured shape: a NamedTuple of NamedTuples of arrays following the architecture, wrapped in a type of its own so that the parameter set is something a package can dispatch on without piracy.
• FlatParameters is the flat shape: one AbstractVector of every number in the set, which is what a derivative, a linear solver or a quasi-Newton method wants to work with.

flatten and unflatten convert between them through a ParameterLayout that is built once and reused, so a derivative taken with respect to the flat vector comes back in the shape of the network:

v, layout = flatten(ps)
g = ForwardDiff.gradient(w -> loss(unflatten(layout, w)), v)
unflatten(layout, g)          # the gradient, laid out like `ps`

Structured leaves — a symmetric matrix keeping $n(n+1)/2$ numbers behind an $n \times n$ interface, a manifold element, a horizontal lift — plug in through two methods, freeparameters and rebuild. Everything else in the package, including the HDF5 support, is written against that protocol rather than against a list of types.

Loading HDF5 brings in save and load through a package extension.

GeometricBase.L2norm of a whole parameter set is here rather than behind an extension, because most of this ecosystem depends on GeometricBase already and its own sole dependency is Unicode. See L2norm.

LeafLayout(range, size)

A leaf whose numbers are copied straight into the flat vector: it occupies range, and comes back reshaped to size (() for a scalar).

This is the terminal case — freeparameters hands back something of the leaf's own type — so there is nothing to rebuild and the layout is the shape alone. Two leaves of the same shape therefore share one layout type whatever they hold, and a stored layout keeps no reference to the parameters it was built from. WrappedLayout is the other case, and it does keep a prototype.

WrappedLayout(prototype, inner)

A leaf whose freeparameters are themselves structured — a SymmetricMatrix storing a vector, or a horizontal lift storing two blocks. inner is the layout of that storage; unflattening runs it and then calls rebuild on prototype.

NestedLayout(children, range)

The layout of a NamedTuple: one child layout per key, together spanning range.

TupleLayout(children, range)

The layout of a Tuple, positional counterpart of NestedLayout.

ParametersLayout(inner)

The layout of a NetworkParameters: the layout of the NamedTuple it wraps, tagged so that unflattening returns a NetworkParameters again rather than a bare NamedTuple.

L2norm(ps::NetworkParameters)

$\sum_i \mathrm{l2norm}(x_i)^2$ over the leaves of ps, at whatever depth they are.

l2norm(ps) follows from GeometricBase's generic l2norm(x) = sqrt(L2norm(x)), which is the quantity the callers want: the blocks of a parameter set combine in quadrature. Summing their norms instead overestimates the ℓ² norm by up to $\sqrt{k}$ for $k$ blocks, and thereby every stopping criterion computed from it.

Implementation

This method is here rather than in GeometricBase because what it has to get right is this package's, not that one's: it walks the leaf protocol with foldparameters and it dispatches on NetworkParameters. A change to either breaks it, and this is the package where such a change is made and tested. GeometricBase supports Julia 1.10 while this package requires 1.11, so a test environment there could not resolve this package at all — a method a package cannot exercise is a method it cannot keep correct.

Ownership does not decide the question, because it admits both: a method is type piracy only when the function and every dispatched argument type belong elsewhere, and NetworkParameters is this package's.

L2norm is the method and l2norm is what it calls on the leaves, which is deliberately not symmetric. A leaf is entitled to its own notion of norm over its free parameters, and downstream packages define exactly that: GeometricOptimizers' l2norm(::AbstractLieAlgHorMatrix) folds over the lift's blocks, and its l2norm(::VectorStorageMatrix) over the stored vector. Both would be wrong if this recursed through L2norm instead, because the generic L2norm(::AbstractArray) reads the dense $n \times n$ interface and so counts a skew-symmetric block's entries twice.

abs2(l2norm(x)) and not L2norm(x) at each leaf, which is the same asymmetry from the other side: L2norm is what a leaf would have to override to be counted correctly, and a structured leaf overrides l2norm. It costs a sqrt and a square per leaf, which is a rounding step and not a term.

foldparameters and not map + sum: a parameter set is a tree of layers and the quantities to combine sit at its leaves, so map would hand l2norm a whole layer, for which there is no method. The fold recurses into the branches, reaches a leaf at any depth, and allocates nothing.

false and not zero(T) as the initial value: there is no T in scope here, 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.

isparametertree(x)

Whether x is a branch of a parameter set — a NetworkParameters, NamedTuple or Tuple that the walks recurse into — as opposed to a leaf.

isterminal(x)

Whether the numbers of the leaf x can be copied straight out of it, i.e. whether freeparameters returns x itself.