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 divergence-free (volume-preserving) systems. Hamiltonian and Lagrangian neural networks are in HamiltonianNeuralNetworks.jl and LagrangianNeuralNetworks.jl.

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], others have existed before [5].

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 [6]. Optimization is necessary for some neural network architectures such as symplectic autoencoders and can be critical for others such as the standard transformer [7, 8].

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 [9], so its layers and architectures are written against any backend that package supports: CUDA.jl [10], AMDGPU.jl, Metal.jl [11] and oneAPI.jl [12].

Little of that is tested. test/devices/metal.jl runs two testsets on an Apple GPU, in the metal test group, which the Metal job and the macOS-latest job of CI run: the products of _mul — a dense and a SymmetricMatrix weight from the left, a dense one from the right — together with the tensor_tensor_mul kernel; and tensor_cayley with the layer that uses it. A green test matrix says nothing about the rest of the GPU path. What is written below was measured by hand on an Apple M4 Max through Metal.jl 1.11.1, in Float32, because Apple GPUs have no Float64 at all, and against GeometricOptimizers 0.9.0.

Everything tried on that device ran. Constructed, applied to both a matrix and a 3-tensor, and copied back with map_to_cpu: GSympNet, LASympNet, StandardTransformerIntegrator, LinearSymplecticTransformer, SymplecticTransformer at both its default transformer_dim and an upscaling one, VolumePreservingFeedForward, VolumePreservingTransformer, SymplecticAutoencoder, PSDArch and Transformer(…; Stiefel = true). The product of a StrictlyLowerTriangular matrix with a 3-tensor runs. train! trains a GSympNet for two epochs with its data in MtlArrays. ClassificationTransformer was not tried.

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

• VolumePreservingFeedForward and VolumePreservingTransformer make their StrictlyLowerTriangular or StrictlyUpperTriangular weight dense before they multiply, as its product with the identity on the weight's own device. 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.