NOMAD: Nonlinear Manifold Decoders for Operator Learning
Applied
Abstract
Supervised learning in function spaces is an emerging area of machine learning research with applications to the prediction of complex physical systems such as ﬂuid ﬂows, solid mechanics, and climate modeling.
By directly learning maps (operators) between inﬁnite dimensional function spaces, these models are able to learn discretization invariant representations of target functions.
A common approach is to represent such target functions as linear combinations of basis elements learned from data.
However, there are simple scenarios where, even though the target functions form a low dimensional submanifold, a very large number of basis elements is needed for an accurate linear representation.
Here we present NOMAD, a novel operator learning framework with a nonlinear decoder map capable of learning ﬁnite dimensional representations of nonlinear submanifolds in function spaces.
We show this method is able to accurately learn low dimensional representations of solution manifolds to partial differential equations while outperforming linear models of larger size.
Additionally, we compare to state-of-the-art operator learning methods on a complex ﬂuid dynamics benchmark and achieve competitive performance with a signiﬁcantly smaller model size and training cost.
Introduction
[Figure]
Machine learning techniques have been applied to great success for modeling functions between ﬁnite dimensional vector spaces.
For example, in computer vision (vectors of pixel values) and natural language processing (vectors of word embeddings) these methods have produced stateof-the-art results in image recognition [18] and translation tasks [48].
However, not all data has an obvious and faithful representation as ﬁnite dimensional vectors.
In particular, functional data is mathematically represented as a vector in an inﬁnite dimensional vector space.
This kind of data appears naturally in problems coming from physics, where scenarios in ﬂuid dynamics, solid mechanics, and kinematics are described by functions of continuous quantities.
Supervised learning in the inﬁnite dimensional setting can be considered for cases where we want to map functional inputs to target functional outputs.
For example, we might wish to predict the velocity of a ﬂuid as function of time given an initial velocity ﬁeld, or predict the pressure ﬁeld across the surface of the Earth given temperature measurements.
This is similar to a ﬁnite dimensional regression problem, except that we are now interested in learning an operator between spaces of functions.
We refer to this as a supervised operator learning problem : given a data-set of N pairs of functions { ( u 1 , s 1 ) , . . .
, ( u N , s N ) } , learn an operator F which maps input functions to output functions such that F ( u i )
=
s i , 8 i .
One approach to solve the supervised operator learning problem is to introduce a parameterized operator architecture and train it to minimize a loss between the model’s predicted functions and the true target functions in the training set.
One of the ﬁrst operator network architectures was presented in [8] with accompanying universal approximation guarantees in the uniform norm.
These results were adapted to deep networks in [30] and led to the DeepONet architecture and its variants [51, 32, 21].
The Neural Operator architecture, motivated by the composition of linear and nonlinear layers in neural networks, was proposed in [27].
Using the Fourier convolution theorem to compute the integral transform in Neural Operators led to the Fourier Neural Operator [28].
Other recent architectures include approaches based on PCA-based representations [2], random feature approaches [36], wavelet approximations to integral transforms [15], and attention-based architectures [23].
A common feature shared among many of these approaches is that they aim to approximate an operator using three maps: an encoder, an approximator, and a decoder, see Figure 1 and Sectionfor more details.
In all existing approaches embracing this structure, the decoder is constructed as a linear map.
In doing so, the set of target functions is being approximated with a ﬁnite dimensional linear subspace in the ambient target function space.
Under this setting, the universal approximation theorems of [8, 24, 25] guarantee that there exists a linear subspace of a large enough dimension which approximates the target functions to any prescribed accuracy.
However, as with ﬁnite dimensional data, there are scenarios where the target functional data concentrates on a low dimensional nonlinear submanifold.
We refer to the phenomenon of data in function spaces concentrating on low dimensional submanifolds as the Operator Learning Manifold
Hypothesis , see Figure 1.
For example, it is known that certain classes of parametric partial differential equations admit low dimensional nonlinear manifolds of solution functions [9].
Although linear representations can be guaranteed to approximate these spaces, their required dimension can become very large and thus inefﬁcient in capturing the true low dimensional structure of the data.
In this paper, we are motivated by the Operator Learning Manifold Hypothesis to formulate a new class of operator learning architectures with nonlinear decoders.
Our key contributions can be summarized as follows.
• Limitations of Linear Decoders: We describe in detail the shortcomings of operator learning methods with linear decoders and present some fundamental lower bounds along with an illustrative operator learning problem which is subject to these limitations.
• Nonlinear Manifold Decoders (NOMAD):
This motivates a novel operator learning framework with a nonlinear decoder that can ﬁnd low dimensional representations for ﬁnite dimensional nonlinear submanifolds in function spaces.
• Enhanced Dimensionality Reduction: A collection of numerical experiments involving linear transport and nonlinear wave propagation shows that, by learning nonlinear submanifolds of target functions, we can build models that achieve state-of-the-art accuracy while requiring a signiﬁcantly smaller number of latent dimensions.
• Enhanced Computational Efﬁciency: As a consequence, the resulting architectures contain a signiﬁcantly smaller number of trainable parameters and their training cost is greatly reduced compared to competing linear approaches.
We begin our presentation in Section 2 by providing a taxonomy of representative works in the literature.
In Section 3 we formally deﬁne the supervised operator learning problem and discuss existing approximation strategies, with a focus on highlighting open challenges and limitations.
In
Section 4 we present the main contributions of this work and illustrate their utility through the lens of a pedagogical example.
In Section 5 we provide a comprehensive collection of experiments that demonstrate the performance of using NOMAD against competing state-of-the-art methods for operator learning.
Section 6 summarizes our main ﬁndings and discusses lingering limitations and broader impact.
Additional details on architectures, hyperparameter selection, and training details are provided in the Supplemental Materials.
The code and data accompanying this manuscript are available at https://github.com/PredictiveIntelligenceLab/NOMAD .  Related Work in Dimensionality Reduction
Low Dimensional Representations in Finite Dimensional Vector Spaces
: Finding low dimensional representations of high dimensional data has a long history, going back to 1901 with the original formulation of principal components analysis (PCA) [39].
PCA is a linear method that works best when data concentrates on low dimensional subspaces.
When data instead concentrates on low dimensional nonlinear spaces, kernelized PCA [42] and manifold learning techniques such as Isomap and diffusion maps [46, 11, 35] can be effective in ﬁnding nonlinear low dimensional structure, see [47] for a review.
The recent popularity of deep learning has introduced new methods for ﬁnding low dimensional structure in high dimensional data-sets, most notably using auto-encoders [52, 6] and deep generative models [12, 22].
Relevant to our work, such techniques have found success in approximating submanifolds in vector spaces corresponding to discretized solutions of parametric partial differential equations (PDEs) [43, 41, 14], where a particular need for nonlinear dimension reduction arises in advection-dominated problems common to ﬂuid mechanics and climate science [26, 33].
Low Dimensional Representations in Inﬁnite Dimensional Vector Spaces: The principles behind PCA generalize in a straightforward way to functions residing in low dimensional subspaces of inﬁnite dimensional Hilbert spaces [50].
In the ﬁeld of reduced order modeling of PDEs this is sometimes referred to as proper orthogonal decomposition [7] (see [29] for an interesting exposition of the discrete version and connections to the Karhunen-Loève decomposition).
Afﬁne representations of solution manifolds to parametric PDEs and guarantees on when they are effective using the notion of linear n -widths [40] have been explored in [9].
As in the case of ﬁnite dimensional data, using a kernel to create a feature representation of a set of functions, and then performing PCA in the associated Reproducing Kernel Hilbert Space can give nonlinear low dimensional representations [45].
The theory behind optimal nonlinear low dimensional representations for sets of functions is still being developed, but there has been work towards deﬁning what “optimal” should mean in this context and how it relates to more familiar geometric quantities [10].
Operator Learning
Notation: Let us ﬁrst set up some notation and give a formal statement of the supervised operator learning problem.
We deﬁne C ( X ; R d ) as the set of continuous functions from a set X to R d .
When
X ⇢ R n , we deﬁne the Hilbert space,
[Equation]
This is an inﬁnite dimensional vector space equipped with the inner product h f, g i = [Equation] f
[Equation] x [Equation] g ( x ) dx .
When X is compact, we have that C ( X ; R d ) ⇢ L 2 ( X ; R d ) .
We now can present a formal statement of the supervised operator learning problem.
Problem Formulation: Suppose we are given a training data-set of N pairs of functions ( u i , s i ) , where u i 2 C ( X ; R d u ) with compact X ⇢
R d x , and s i 2 C
( Y ; R d s ) with compact Y ⇢ R d y .
Assume there is a ground truth operator G : C ( X ; R d u ) !
C ( Y ; R d s ) such that G ( u i ) =
s i and that the u i are sampled i.i.d.
from a probability measure on C ( X ; R d u ) .
The goal of the supervised operator learning problem is to learn a continuous operator F : C ( X ; R d x ) ! C ( Y ; R d s ) to approximate G .
To do so, we will attempt to minimize the following empirical risk over a class of operators F ✓ , with parameters ✓ 2 ⇥ ⇢ R d ✓ ,
[Equation]
An Approximation Framework for Operators: A popular approach to learning an operator
G : L 2 ( X ) ! L 2 ( Y ) acting on a probability measure µ on L 2 ( X ) is to construct an approximation out of three maps [25] (see Figure 1),
[Equation]
The ﬁrst map, E : L 2 ( X ) ! R m is known as the encoder.
It takes an input function and maps it to a ﬁnite dimensional feature representation.
For example, E could take a continuous function to its point-wise evaluations along a collection of m sensors, or project a function onto m basis functions.
The next map A : R m ! R n is known as the approximation map.
This can be interpreted as a ﬁnite dimensional approximation of the action of the operator G .
Finally, the image of the approximation map is used to create the output functions in L 2 ( Y ) by means of the decoding map
D : R n ! L 2 ( Y ) .
We will refer to the dimension, n , of the domain of the decoder as the latent dimension .
The composition of these maps can be visualized in the following diagram.
[Equation]
Linear Decoders: Many successful operator learning architectures such as the DeepONet [30], the (pseudo-spectral) Fourier Neural Operator in [24], LOCA [23], and the PCA-based method in [2] all use linear decoding maps D .
A linear D can be deﬁned by a set of functions ⌧ i 2 L 2 ( Y ) , i = 1 , . . . , n , and acts on a vector  2 R n as
[Equation]
For example, the functions ⌧ i can be built using trigonometric polynomials as in the -FNO [24], be parameterized by a neural network as in DeepONet [30], or created as the normalized output of a kernel integral transform as in LOCA [23].
Limitations of Linear Decoders: We can measure the approximation accuracy of the operator F with two different norms.
First is the L 2 ( µ ) operator norm,
[Equation]
Note that the empirical risk used to train a model for the supervised operator learning problem (see (1)) is a Monte Carlo approximation of the above population loss.
The other option to measure the approximation accuracy is the uniform operator norm.
For a compact set U ⇢ L 2 ( X ) we can deﬁne the error
[Equation]
When a linear decoder is used for F = D  A  E , a data-dependent lower bound to each of these errors can be derived.
L 2 lower bound: When the pushforward measure has a ﬁnite second moment, its covariance operator  : L 2 ( Y ) ! L 2 ( Y ) is self-adjoint, positive semi-deﬁnite, and trace-class, and thus admits an orthogonal set of eigenfunctions spanning its image, {  1 ,  2 , . . . } with associated decreasing eigenvalues  1   2  ... .
The decay of these eigenvalues indicates the extent to which samples from G # µ concentrate along the leading ﬁnite-dimensional eigenspaces.
It was shown in [25] that for any choice of E and A , these eigenvalues give a fundamental lower bound to the expected squared Lerror of the operator learning problem with architectures as in (3) using a linear decoder D ,
[Equation]
This result can be reﬁned to show that the optimal choice of functions ⌧ i (see equation (4)) for a linear decoder are the leading n eigenfunctions of the covariance operator {  1 , . .
,  n } .
The interpretation of this result is that the best way to approximate samples from G # µ with an n -dimensional subspace is to use the subspace spanned by the ﬁrst n “principal components” of the probability measure G # µ .
The error incurred by using this subspace is determined by the remaining principal components, namely the sum of their eigenvalues P [Equation] k .
The operator learning literature has noted that for problems with a slowly decaying pushforward covariance spectrum (such as solutions to advectiondominated PDEs) these lower bounds cause poor performance for models of the form (3) [25, 13].
Uniform lower bound:
In the reduced order modelling of PDEs literature [9, 10, 26] there exists a related notion for measuring the degree to an n -dimensional subspace can approximate a set of functions S ⇢ L 2 ( Y ) .
This is known as the Kolmogorov n -width [40], and for a compact set S is deﬁned as
[Equation]
This measure of how well a set of functions can be approximated by a linear subspace in the uniform norm leads naturally to a lower bound for the uniform error (6).
To see this, ﬁrst note that for any u 2 U , the error from F ( u ) to G ( u ) is bounded by the minimum distance from G ( u ) to the image of
F .
For a linear decoder D : R n ! L 2 ( Y ) , deﬁne the (at most) n -dimensional V n = im( D ) ⇢ L 2 ( Y ) .
Note that im(
F ) ✓ V n , and we may write
[Equation]
Taking the supremum of both sides over u 2 U , and then the inﬁmum of both sides over all n -dimensional subspaces V n gives
[Equation]
The quantity on the right is exactly the Kolmogorov n -width of G ( U ) .
We have thus proved the following complementary statement to (7) when the error is measured in the uniform norm.
Proposition 1 Let U 2 L 2 ( X ) be compact and consider an operator learning architecture as in (3) , where D : R n ! L 2 ( Y ) is a linear decoder.
Then, for any E : L 2 ( X ) ! R p and A : R p ! R n , the uniform norm error of [Equation]
satisﬁes the lower bound
[Equation]
Therefore, we see that in both the L 2 ( µ ) and uniform norm, the error for an operator learning problem with a linear decoder is fundamentally limited by the extent to which the space of output functions “ﬁts” inside a ﬁnite dimensional linear subspace.
In the next section we will alleviate this fundamental restriction by allowing decoders that can learn nonlinear embeddings of R n into L 2 ( Y ) .
Nonlinear Decoders for Operator Learning
A Motivating Example: Consider the problem of learning the antiderivative operator mapping functions to their ﬁrst-order derivative
[Equation]
acting on a set of input functions
[Equation] [Equation]
The set of output functions is given by G ( [Equation] 0 < t 0 < t < T } .
This is a onedimensional curve of functions in L 2 ([0 , 1]) parameterized by a single number t .
However, we would not be able to represent this set of functions with a one-dimensional linear subspace.
In Figure 2b [Figure] we perform PCA on the functions in this set evaluated on a uniform grid of values of t .
We see that the ﬁrst 20 eigenvalues are nonzero and relatively constant, suggesting that an operator learning architecture with a linear or afﬁne decoder would need a latent dimension of at least 20 to effectively approximate functions from G ( U ) .
Figure 2a gives a visualization of this curve of functions projected onto the ﬁrst three PCA components.
We will return to this example in Section 5, and see that an architecture with a nonlinear decoder can in fact approximate the target output functions with superior accuracy compared to the linear case, using a single latent dimension that can capture the underlying nonlinear manifold structure.
Operator Learning Manifold Hypothesis: We now describe an assumption under which a nonlinear decoder is expected to be effective, and use this to formulate the NOMAD architecture.
To this end, let µ be a probability measure on L 2 ( X ) and G : L 2 ( X ) ! L 2 ( Y ) .
We assume that there exists an n -dimensional manifold M ✓ L 2 ( Y ) and an open subset O ⇢ M such that
[Equation]
[Equation]
[Figure]
[Equation] [Figure] In connection with the manifold hypothesis in deep learning [5, 4], we refer to this as the Operator Learning Manifold Hypothesis .
There are scenarios where it is known this assumption holds, such as in learning solutions to parametric PDEs [33].
This assumption motivates the construction of a nonlinear decoder for the architecture in (3) as follows.
For each u , choose v ( u ) {+
+}
O such that
[Equation]
Let  : O !
R n be a coordinate chart for
O ⇢ M .
We can represent v ( u ) 2 O by its coordinates  ( v ( u )) 2 R n .
Consider a choice of encoding and approximation maps such that
A ( E ( u )) gives the coordinates for v ( u ) .
If the decoder were chosen as D :=   1 then by construction, the operator F := D
A  E will satisfy
[Equation]
Therefore, we interpret a learned decoding map as attempting to give a ﬁnite dimensional coordinate system for the solution manifold.
Consider a generalized decoder of the following form
[Equation]
[Equation] This induces a map from D : R n ! L 2 ( Y ) , as D (  ) =
˜ D (  , · ) .
If the solution manifold M is a ﬁnite dimensional linear subspace in L 2 (
Y ) spanned by { ⌧ i } ni =1 , we would want a decoder to use the coefﬁcients along the basis as a coordinate system for M .
A generalized decoder could learn this basis as the output of a deep neural network to act as
[Equation]
However, if the solution manifold is not linear, then we should learn a nonlinear coordinate system given by a nonlinear D .
A nonlinear version of ˜ D can be parameterized by using a deep neural network f : R n ⇥ Y !
R which jointly takes as arguments (  , y ) ,
[Equation]
When used in the context of an operator learning architecture of the form (3), we call a nonlinear decoder from (17) NOMAD (NOnlinear MAnifold Decoder).
Figure 3 presents a visual comparison between linear and nonlinear decoders.
Summary of NOMAD: Under the assumption of the Operator Learning Manifold Hypothesis , we have proposed a fully nonlinear decoder (17) to represent target functions using architectures of the form (3).
We next show that using a decoder of the form (17) results in operator learning architectures which can learn nonlinear low dimensional solution manifolds.
Additionally, we will see that when these solution manifolds do not “ﬁt” inside low dimensional linear subspaces, architectures with linear decoders will either fail or require a signiﬁcantly larger number of latent dimensions.
Results
In this section we investigate the effect of using a linear versus nonlinear decoders as building blocks of operator learning architecture taking the form (3).
In all cases, we will use an encoder E which takes point-wise evaluations of the input functions, and an approximator map A given by a deep neural network.
The linear decoder parametrizes a set of basis functions that are learned as the outputs of an MLP network.
In this case, the resulting architecture exactly corresponds to the DeepONet model from [30].
We will compare this against using NOMAD where the nonlinear decoder is built using an MLP network that takes as inputs the concatenation of  2 R n and a given query point y 2 Y .
All models are trained with by performing stochastic gradient descent on the loss function in (1).
The reported errors are measured in the relative L 2 ( Y ) norm by averaging over all functional pairs in the testing data-set.
More details about architectures, hyperparameters settings, and training details are provided in the Supplemental Materials.
Learning the Antiderivative Operator: First, we revisit the motivating example from Section 4, where the goal is to learn the antiderivative operator (10) acting on the set of functions (11).
In
Figure 2c we see the performance of a model with a linear decoder and NOMAD over a range of latent dimensions n .
For each choice of n , 10 experiments with random initialization seeds were performed, and the mean and standard deviation of testing errors are reported.
We see that
the
NOMAD architecture consistently outperforms the linear one (by one order of magnitude), and can even achieve a 10% relative prediction error using only n = 1 .
Solution Operator of a Parametric Advection PDE: Here we consider the problem of learning the solution operator to a PDE describing the transport of a scalar ﬁeld with conserved energy, [Equation]
[Equation]
over a domain ( x, t ) 2 [0 , 2] ⇥ [0 , 1] .
The solution operator maps an initial condition s ( x, 0)
=
u ( x ) to the solution at all times s ( x, t ) which satisﬁes (18).
We consider a training data-set of initial conditions taking the form of radial basis functions with a very small ﬁxed lengthscale centered at randomly chosen locations in the interval [0 , 1] .
We create the output functions by evolving these initial conditions forward in time for 1 time unit according to the advection equation (18) (see
Supplemental Materials for more details).
Figure 4a gives an illustration of one such solution plotted over the space-time domain.
Performing PCA on the solution functions generated by these initial conditions shows a very slow decay of eigenvalues (see Figure 4b), suggesting that methods with linear decoders will require [Figure] a moderately large number of latent dimensions.
However, since the data-set was constructed by evolving a set of functions with a single degree of freedom (the center of the initial conditions), we would expect the output functions to form a solution manifold of very low dimension.
In Figure 4c we compare the performance of a linear decoder and NOMAD as a function of the latent dimension n .
Linear decoders yield poor performance for small values of n , while NOMAD appears to immediately discover a good approximation to the true solution manifold.
Propagation of Free-surface Waves:
As a more challenging benchmark we consider the shallowwater equations; a set of hyperbolic equations that describe the ﬂow below a pressure surface in a ﬂuid [49].
The underlying PDE system takes the form @
U
[Equation]
F @
[Equation]
@ G @
y
[Equation]
[Equation]
U = ⇢⇢ v⇢ v
[Equation]
F
[Equation]
{+@ ⇢ v⇢ v 21 + 12 g ⇢⇢ v 1 v A , G =@ ⇢ v⇢ v 1 v⇢ v 22 + 12 g ⇢ 
 A , (19)+} where ⇢ ( x, y, t ) the ﬂuid height from the free surface, g is the gravity acceleration, and v 1 ( x, y, t ) , v 2 ( x, y, t ) denote the horizontal and vertical ﬂuid velocities, respectively.
We consider reﬂective boundary conditions and random initial conditions corresponding to a random droplet falling into a still ﬂuid bed (see Supplemental Materials).
In Figure 5a we show the average testing error of a model with a linear and nonlinear decoder as a function of the latent dimension.
Figure 5b shows snapshots of the predicted surface height function on top of a plot of the errors to the ground truth for the best, worst, median, and a random sample from the testing data-set.
We additionally use this example to compare the performance of a model with a linear decoder and
NOMAD to other state-of-the-art operator learning architectures.
For this experiment we include a harmonic feature expansion in the DeepONet, as it was shown in [32] that this can signiﬁcantly improve its performance.
We do not include this in the NOMAD architecture for this experiment (see Supplemental Materials for more details).
In Table 1, we present the mean relative error and its standard deviation for different operator learning methods, as well as the prediction that provides the worst error in the testing data-set when compared against the ground truth solution.
For each method we also report the number of its trainable parameters, the number of its latent dimension n , and the training wall-clock time in minutes.
Since the general form of the FNO [28] does not neatly ﬁt into the architecture given by (3), there is not a directly comparable measure of latent dimension for it.
We also observe that, although the model with NOMAD closely matches the performance of LOCA [23], its required latent dimension, total number of trainable parameters, and total training time are all signiﬁcantly smaller.
Discussion
Summary: We have presented a novel framework for supervised learning in function spaces.
The proposed methods aim to address challenging scenarios where the manifold of target functions has [Figure] [Table] low dimensional structure, but is embedded nonlinearly into its associated function space.
Such cases commonly arise across diverse functional observables in the physical and engineering sciences (e.g. turbulent ﬂuid ﬂows, plasma physics, chemical reactions), and pose a signiﬁcant challenge to the application of most existing operator learning methods that rely on linear decoding maps, forcing them to require an excessively large number of latent dimensions to accurately represent target functions.
To address this shortcoming we proposed a fully nonlinear framework that can effectively learn low dimensional representations of nonlinear embeddings in function spaces, and demonstrated that it achieves competitive accuracy to state-of-the-art operator learning methods while using a signiﬁcantly smaller number of latent dimensions, leading to lighter model parametrizations and reduced training cost.
Limitations: Our proposed approach relies on the Operator Learning Manifold Hypothesis (see equation (12)), suggesting that cases where a low dimensional manifold structure does not exist will be hard to tackle (e.g. target function manifolds with fractal structure, solutions to evolution equations with strange attractors).
Moreover, even when the manifold hypothesis holds, the underlying effective latent embedding dimension is typically not known a-priori, and may only be precisely found via cross-validation.
Another direct consequence of replacing linear decoders with fully nonlinear maps is that the lower bound in (9) needs to be rephrased in terms of a nonlinear n -width, which in general can be difﬁcult to quantify.
Finally, in this work we restricted ourselves to exploring simple nonlinear decoder architectures such as an MLPs with the latent parameters  and query location y concatenated as inputs.
Further investigation is needed to quantify the improvements that could be brought by considering more contemporary deep learning architectures, such as hypernetworks [16] which can deﬁne input dependent weights for complicated decoder architectures.
One example of this idea in the context of reduced order modeling can be found in Pan et.
al.
[37], where the authors propose a hypernetwork based method combined with a Implicit Neural Representation network [44].
Potential Negative Societal Impacts: Our contributions enables advances in accelerating the modeling, simulation, and optimization of physical and engineering systems.
As with any tool that furthers our understanding and ability to predict the outcomes of complex systems, there may be ill-intentioned use cases, but we do not expect any speciﬁc negative impact from this work. Acknowledgments and Disclosure of Funding
G.K. and P.P. would like to acknowledge support from the US Department of Energy under under the
Advanced Scientiﬁc Computing Research program (grant DE-SC0019116), the US Air Force (grant
AFOSR FA9550-20-1-0060), and US Department of Energy/Advanced Research Projects Agency (grant DE-AR0001201).
J.S. and G.P. would like to acknowledge support from the AFOSR under grant FA9550-19-1-0265 (Assured Autonomy in Contested Environments) and the NSF Simmons
Mathematical and Scientiﬁc Foundations of Deep Learning (grant 2031985).
We also thank the developers of the software that enabled our research, including JAX [3], Kymatio [1], Matplotlib [20], Pytorch [38] and NumPy [17].