NeuralNetworkParameters
The parameters of a neural network, in two shapes, with conversions between them.
The package serves the networks defined in AbstractNeuralNetworks, GeometricMachineLearning and SymbolicNeuralNetworks, and holds nothing but the parameters: it depends on none of them, and adds one dependency of its own (ChainRulesCore).
It exists for two reasons. The first is that a parameter set needs to be a type somebody owns. A bare NamedTuple belongs to Base, so a package that wants to give the parameter set its own behaviour — writing it to a file, flattening it, stepping an optimizer over it — has to write methods on a signature in which it owns nothing. Every such method is type piracy, and two packages doing it can silently disagree.
The second is that the traversal which goes with the container was being written once per package. Recurse into the NamedTuples, do something at each leaf, put the result back in the same shape: that is flatten/unflatten, h5save/h5load, changebackend, map_to_cpu and the elementwise optimizer primitives, each re-declaring a method per structured parameter type. Written once against the leaf protocol below, one implementation covers all of them and needs to know none of the types.
A note on the name: the package is NeuralNetworkParameters, the type it exports is NetworkParameters. A package cannot export a type sharing its own name — the module binding wins at the using site, so NeuralNetworkParameters(nt) would try to call a Module. AbstractNeuralNetworks 0.7 dropped the name rather than aliasing it, so that one type has one name across the ecosystem: code written against NeuralNetworkParameters as a type reaches for NetworkParameters here instead.
The two shapes
NetworkParameters follows the architecture — a NamedTuple of NamedTuples of arrays, one entry per layer:
using NeuralNetworkParameters
ps = NetworkParameters((L1 = (W = [1.0 2.0; 3.0 4.0], b = [5.0, 6.0]),
L2 = (W = [7.0 8.0], b = [9.0])))
ps.L1.W
# output
2×2 Matrix{Float64}:
1.0 2.0
3.0 4.0FlatParameters is the same numbers as one vector, which is the shape a derivative, a linear solver or a quasi-Newton method wants:
fp = FlatParameters(ps)
collect(fp)
# output
9-element Vector{Float64}:
1.0
3.0
2.0
4.0
5.0
6.0
7.0
8.0
9.0Neither is derived from the other on the fly: flatten and unflatten convert between them through a ParameterLayout that is built once and then reused.
Differentiating with respect to the flat form
The point of the flat shape is that ForwardDiff — or any method that wants a vector — can work on it, while the answer still comes back laid out like the network:
using ForwardDiff
v, layout = flatten(ps)
loss(p) = sum(model(x, p))
g = ForwardDiff.gradient(w -> loss(unflatten(layout, w)), v)
unflatten(layout, g) # the gradient, one entry per layerunflatten is generic in the element type of its vector, which is what makes this work: a Dual-valued vector produces Dual-valued parameters. Reverse mode is covered too — there are ChainRulesCore rules for both conversions, so Zygote can differentiate through them, and with ZygoteRules loaded the gradient of a NetworkParameters is a NetworkParameters rather than a tangent nobody can consume.
Structured parameters
A parameter need not be a plain array. A symmetric matrix keeps $n(n+1)/2$ numbers behind an $n \times n$ interface; a manifold element or a horizontal lift keeps its own. Those numbers — not the dense entries — are what belongs in the flat vector, and what an optimizer should move.
Two methods teach the package about such a type, and everything else follows: flattening, the tree walks, and HDF5. See The leaf protocol.