Neural Network Integrators

NeuralNetworkIntegrator is a super type of various neural network architectures such as SympNet and ResNet.

The purpose of such neural networks is to approximate the flow of an ordinary differential equation (ODE).

NeuralNetworkIntegrators can be seen as modeling traditional one-step methods with neural networks, i.e. for a fixed time step they perform:

\[ \mathtt{NeuralNetworkIntegrator}: z^{(t)} \mapsto z^{(t+1)},\]

to try to approximate the flow of some ODE:

\[ || \mathtt{Integrator}(z^{(t)}) - \varphi^h(z^{(t)}) || \approx \mathcal{O}(h),\]

where $\varphi^h$ is the flow map of the ODE for a time step $h$.

ResNet(dim, n_blocks, activation)
ResNet(dim, n_blocks, width, activation)

Make an instance of a ResNet.

A ResNet is a neural network that realizes a mapping of the form:

\[ x = \mathcal{NN}(x) + x,\]

so the input is again added to the output (a so-called add connection). In GeometricMachineLearning the specific ResNet that we use consists of a series of simple ResNetLayers, or of WideResNetLayers where width differs from dim: that is the one structural choice the architecture makes, and a width equal to dim keeps the plain ResNetLayer.

Arguments

• dim::Integer: the system dimension.
• n_blocks::Integer: the number of blocks before the linear output layer.
• width::Integer: the width of the hidden layer of each block, which is dim in ResNet(dim, n_blocks, activation) and the third positional argument of ResNet(dim, n_blocks, width, activation) and of ResNet(dl, n_blocks, width), where it defaults to dl.input_dim. See WideResNetLayer.
• activation: the third positional argument of ResNet(dim, n_blocks, activation) and the fourth of ResNet(dim, n_blocks, width, activation), where ResNet(dl, n_blocks, width) takes it as the keyword activation = tanh. It applies to each of the n_blocks blocks; the linear output layer applies the identity.

Constructor

ResNet can also be called with the constructor:

ResNet(dl, n_blocks, width)

where dl is an instance of DataLoader and width defaults to dl.input_dim.

See integrate for an example of this.

For a ResNet whose blocks also take the parameters of the system, see ParametricResNet.

WideResNetLayer(dim, width, activation)

A ResNetLayer whose hidden layer is width wide, independent of dim.

ResNetLayer applies a single $\mathrm{dim} \times \mathrm{dim}$ weight, so its capacity is tied to the system dimension. This one upscales to width, applies the activation, and downscales again:

\[ x \mapsto x + \sigma(W_\mathrm{down}\sigma(W_\mathrm{up}x + b_\mathrm{up}) + b),\]

with $W_\mathrm{up}\in\mathbb{R}^{\mathrm{width}\times\mathrm{dim}}$ and $W_\mathrm{down}\in\mathbb{R}^{\mathrm{dim}\times\mathrm{width}}$. ResNet uses it whenever the width it is given differs from the system dimension, and keeps the plain ResNetLayer where the two agree.

The input is a vector, a matrix of samples, a tensor of samples, or a (q, p) pair, and the map is the same in every form.

Also see ParametricResNetLayer, which additionally takes the parameters of the system.

ResNetLayer(dim)

Make an instance of the resnet layer.

The ResNetLayer is a simple feedforward neural network to which we add the input after applying it, i.e. it realizes $x \mapsto x + \sigma(Ax + b)$.

Arguments

The ResNet layer takes the following arguments:

1. dim::Integer: the system dimension.
2. activation = identity: The activation function.

The following is a keyword argument:

• use_bias::Bool = true: This determines whether a bias $b$ is used.

ParametricResNet(dim; width, n_blocks, activation, parameters)

A ResNet whose blocks also take the parameters of the system, built from ParametricResNetLayers.

The architecture without structure preservation that the generalized Hamiltonian neural network of HamiltonianNeuralNetworks.jl is compared against: it maps $(x, \mu) \mapsto x'$ with the same information available to it, but nothing in it makes the map symplectic. Unlike ResNet it uses a ParametricResNetLayer at every width, and the linear output layer applies the identity activation.

Keyword arguments

• width = dim: the width of the hidden layer of each block.
• n_blocks = 1: the number of blocks before the linear one.
• activation = tanh: the activation of every block.
• parameters = NullParameters(): a NamedTuple of system parameters, used only for its shape.

Constructor

ParametricResNet can also be called with the constructor

ParametricResNet(dl, n_blocks, width)

where dl is an instance of DataLoader and width defaults to dl.input_dim, or an instance of ParametricDataLoader. The latter gives the system dimension, the rows of its data or of q and p together, which is also the default of width, and the system parameters, whose first μ the architecture keeps for its shape; it therefore takes no parameters keyword.

The forward pass is apply_parametric, or the functor nn(x, μ) of the NeuralNetwork built around it.

ParametricResNetLayer(dim, width, activation; parameters, return_parameters)

A WideResNetLayer whose hidden layer also sees the parameters of the system.

The flattened system parameters are appended to the input of the upscaling weight, so the layer computes

\[ x \mapsto x + \sigma(W_\mathrm{down}\sigma(W_\mathrm{up}[x; \mu] + b_\mathrm{up}) + b),\]

where $\mu$ are the system parameters. $W_\mathrm{up}$ is therefore $\mathrm{width}\times(\mathrm{dim} + |\mu|)$ wide, and without system parameters the map is that of a WideResNetLayer with the same width and weights. The input is a vector, a matrix of samples, a tensor of samples, or a (q, p) pair of any of these, with one μ — the forms that WideResNetLayer takes.

Keyword arguments

• parameters = NullParameters(): a NamedTuple of system parameters, used only for its shape — the layer stores the resulting NeuralNetworkParameters.ParameterLayout and the flattened length, and the values are supplied per call.
• return_parameters = false: whether to pass the system parameters on to the next layer beside the output, which is what lets a Chain of these thread them through. A Chain of one layer only takes the default.

This is the building block of ParametricResNet.

flatten_system_parameters(μ::NamedTuple)
flatten_system_parameters(::NullParameters)
flatten_system_parameters(::Type{T}, μ::NamedTuple)
flatten_system_parameters(::Type{T}, ::NullParameters)

Return the tuple (v, layout) of the flat Vector v of the system parameters μ and the ParameterLayout that puts such a vector back into the shape of μ. It is the flatten of NeuralNetworkParameters of μ, and of the empty NamedTuple() for NullParameters().

With T, the element type of v is T; without it, it is the one element type of the leaves of μ, and Union{} for NullParameters() and (;). NullParameters() flattens to an empty vector. The layout is a value and not a closure, so a layer can store it in a field.

The function throws no error of its own; the errors of flatten propagate. It throws an ArgumentError for a leaf of μ that flatten cannot copy, such as a String, and, without T, for leaves of two element types. With T, it throws an InexactError for a value that T cannot hold, such as 0.5 for T = Int.

concatenate_array_with_parameters(x::AbstractArray, ::NullParameters)
concatenate_array_with_parameters(x::AbstractVector, μ::NamedTuple)
concatenate_array_with_parameters(x::AbstractMatrix, μ::NamedTuple)
concatenate_array_with_parameters(x::AbstractArray{<:Any, 3}, μ::NamedTuple)
concatenate_array_with_parameters(x::AbstractMatrix, μs::AbstractVector)

Return the input of a parametric layer: the state x with the flattened system parameters below it, in the element type of x.

• With NullParameters(), return x itself.
• With one μ, flatten it with flatten_system_parameters(eltype(x), μ), and put it below a vector, below every column of a matrix, and below every column of every slice of a 3-tensor. The result has as many more rows than x as the flat μ has entries.
• With a vector μs, put μs[i] below column i of the matrix x, and return the columns side by side.

Throws a DimensionMismatch when the number of columns of x is not length(μs), an ArgumentError for an empty μs with a matrix of no columns, and an InexactError when a parameter has no value in the element type of x, such as 0.5 below an integer state. An error of flatten_system_parameters propagates. An array of more than three dimensions with a NamedTuple, and a vector or an array of three or more dimensions with a vector of parameters, have no method.

integrate(nn, ics)

This function computes a trajectory for a NeuralNetworkIntegrator that has already been trained for valuation purposes.

It takes as input:

1. nn: a NeuralNetwork (that has been trained).
2. ics: initial conditions (a NamedTuple of two vectors)

Examples

To demonstrate integrate we use a simple ResNet that does:

\[\mathrm{ResNet}: x \mapsto \begin{pmatrix} 1 & 0 & 0 \\ 0 & 2 & 0 \\ 0 & 0 & 1\end{pmatrix}x + \begin{pmatrix} 0 \\ 0 \\ 1 \end{pmatrix}\]

and we iterate three times with

\[ \mathtt{ics} = \begin{pmatrix} 1 \\ 1 \\ 1 \end{pmatrix}.\]

using GeometricMachineLearning

model = ResNet(3, 0, identity)
weight = [1 0 0; 0 2 0; 0 0 1]
bias = [0, 0, 1]
ps = NetworkParameters((L1 = (weight = weight, bias = bias), ))
nn = NeuralNetwork(model, Chain(model), ps, CPU())

ics = [1, 1, 1]
integrate(nn, ics; n_points = 4)

# output

3×4 Matrix{Int64}:
 1  2  4   8
 1  3  9  27
 1  3  7  15

Arguments

The optional keyword argument is

• n_points = 100

The number of integration steps that should be performed.

TransformerIntegrator <: Architecture

Encompasses various transformer architectures, such as the VolumePreservingTransformer and the LinearSymplecticTransformer.

The central idea behind this is to construct an explicit multi-step integrator:

\[ \mathtt{Integrator}: [ z^{(t - \mathtt{sl} + 1)}, z^{(t - \mathtt{sl} + 2)}, \ldots, z^{(t)} ] \mapsto [ z^{(t + 1)}, z^{(t + 2)}, \ldots, z^{(t + \mathtt{pw})} ],\]

where sl stands for sequence length and pw stands for prediction window, so the numbers of input and output vectors respectively.

integrate(nn, ics)

Iterate the neural network of type TransformerIntegrator for initial conditions ics.

The initial condition is a matrix $\in\mathbb{R}^{n\times\mathtt{seq\_length}}$ or NamedTuple of two matrices).

This function computes a trajectory for a Transformer that has already been trained for valuation purposes.

Parameters

The following are optional keyword arguments:

• n_points::Int=100: The number of time steps for which we run the prediction.
• prediction_window::Int=size(ics.q, 2): The prediction window (i.e. the number of steps we predict into the future) is equal to the sequence length (i.e. the number of input time steps) by default.