Geometric Machine Learning

GeometricMachineLearning is a package for structure-preserving scientific machine learning. It contains models that can learn dynamical systems with geometric structure, such as Hamiltonian (symplectic) or Lagrangian (variational) systems.

In that regard its aim is similar to traditional geometric numerical integration [1, 2] in that it models maps that share properties with the analytic solution of a differential equation:

Installation

GeometricMachineLearning and all of its dependencies can be installed via the Julia REPL by typing

]add GeometricMachineLearning

Architectures

Some of the neural network architectures in GeometricMachineLearning [3, 4] have emerged in connection with developing this package^[1], other have existed before [5, 6].

New architectures include:

• symplectic autoencoder,
• volume-preserving transformers,
• linear-symplectic transformer.

Existing architectures include:

• SympNets,
• standard transormer.

Manifolds

GeometricMachineLearning supports putting neural network weights on manifolds such as the Stiefel manifold and the Grassmann manifold and Riemannian optimization.

When GeometricMachineLearning optimizes on manifolds it uses the framework introduced in [7]. Optimization is necessary for some neural network architectures such as symplectic autoencoders and can be critical for others such as the standard transformer [8, 9].

Special Neural Network Layer

Many layers have been adapted in order to be used for problems in scientific machine learning, such as the attention layer.

GPU Support

GeometricMachineLearning allocates and computes through KernelAbstractions.jl [10], so its layers and architectures are written against any backend that package supports: CUDA.jl [11], AMDGPU.jl, Metal.jl [12] and oneAPI.jl [13].

None of that is tested. There is no GPU test under test/ and no GPU job in CI, so a green test matrix says nothing about the GPU path. What is written below was measured by hand on an Apple M4 Max through Metal.jl, in Float32, because Apple GPUs have no Float64 at all, and against GeometricOptimizers 0.8.0.

Everything tried on that device ran. Constructed, and then applied to both a matrix and a 3-tensor: GSympNet, LASympNet, StandardTransformerIntegrator, LinearSymplecticTransformer, SymplecticTransformer at both its default transformer_dim and an upscaling one, VolumePreservingFeedForward, VolumePreservingTransformer, SymplecticAutoencoder, PSDArch and Transformer(…; Stiefel = true). ClassificationTransformer was constructed but not applied, because its input is an image. The tensor kernels and map_to_cpu run, and so does training with Optimizer, which returns a Float32 history and leaves the parameters on the device — for PSDArch that includes the manifold weights, which stay a StiefelManifold over an MtlMatrix.

Two of those results rest on GeometricOptimizers rather than on anything here.

• VolumePreservingFeedForward and VolumePreservingTransformer applied to a matrix multiply a LowerTriangular or UpperTriangular weight by that matrix. GeometricOptimizers supplies a KernelAbstractions kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which GPUArraysCore refuses.
• StiefelLayer, GrassmannLayer and PSDLayer, and so every architecture that holds one of them, orthonormalize their weight at construction through GeometricOptimizers.orthonormal_columns. That is CholeskyQR2 — matrix products and triangular solves only — so it runs wherever the draw was allocated. LinearAlgebra.qr! is a host factorization and Metal.jl implements no qr for an MtlArray, so a device needs the other one.

Tutorials

There are several tutorials demonstrating how GeometricMachineLearning can be used.

These tutorials include:

• a tutorial on SympNets that shows how we can model a flow map corresponding to data coming from an unknown canonical Hamiltonian system,
• a tutorial on symplectic Autoencoders that shows how this architecture can be used in structure-preserving reduced order modeling,
• a tutorial on the volume-preserving attention mechanism which serves as a basis for the volume-preserving transformer,
• a tutorial on training a transformer with manifold weights for image classification, in the companion package GMLDatasets, to show that manifold optimization is also useful outside of scientific machine learning.

Data-Driven Reduced Order Modeling

The main motivation behind developing GeometricMachineLearning is reduced order modeling, especially structure-preserving reduced order modeling. For this purpose we give a short introduction into this topic.