Neural Implicit Manifold Learning for Topology-Aware Density Estimation
Authors: Brendan Leigh Ross
Gabriel Loaiza-Ganem
Anthony L. Caterini
Jesse C. Cresswell
Venue: Accepted by TMLR ()

Abstract: Natural data observed in $\mathbb{R}^n$ is often constrained to an $m$-dimensional manifold $\mathcal{M}$, where $m < n$. This work focuses on the task of building theoretically principled generative models for such data. Current generative models learn $\mathcal{M}$ by mapping an $m$-dimensional latent variable through a neural network $f_\theta: \mathbb{R}^m \to \mathbb{R}^n$. These procedures, which we call pushforward models, incur a straightforward limitation: manifolds cannot in general be represented with a single parameterization, meaning that attempts to do so will incur either computational instability or the inability to learn probability densities within the manifold. To remedy this problem, we propose to model $\mathcal{M}$ as a neural implicit manifold: the set of zeros of a neural network. We then learn the probability density within $\mathcal{M}$ with a constrained energy-based model, which employs a constrained variant of Langevin dynamics to train and sample from the learned manifold. In experiments on synthetic and natural data, we show that our model can learn manifold-supported distributions with complex topologies more accurately than pushforward models.

Published in Transactions on Machine Learning Research (12/2023)
Neural Implicit Manifold Learning for Topology-Aware Density Estimation
Brendan Leigh Ross
brendon @layer6.at
Layer 6 AT
Gabriel Loaiza-Ganem
gabricl@layer6..t
Lager 6 AT
Anthony L. Caterini Layer 6 AT
anthony@leyer6..ot
Jesse C. Cresswell
jesse @laaye6. ai
Layer 6
Reviewed on OpenReview: https: //
// openreview.ww..... net/ forum? id- 1TOku8362y
Abstract
Natural data observed R" is often constrained to an m-dimensional manifold M, where m < This work focuses on the task building theoretically principled generative models for such data. Current generative models learn M by mapping an m-dimensional latent variable through neural network fo R - R Those procedures, which call pushfor- word models, incur straightforward limitation: manifolds cannot in general be represented with a single parameterization, meaning that attempts to do so will incur either compu- tational instability or the inability to learn probability densities within the manifold. To remody this problem, we propose to model M as a newral implicit manifold: the set of zeros of a neural network. We then learn the probability density within M with construined eneryy-base model, which employs a constrained variant of Langevin dynamics to train and sample from the learned manifold. In experiments synthetic and natural data, show that our model can learn manifold-supported distributions with complex topologies more accurately than pushforward models.
1 Introduction
Here undertake the broad statistical task of generative modelling: estimating unknown probability distribution Pa from a sample {a:} C IR" of datapoints. Real-world distributions are diverse, complex, and often high-dimensional. Despite the apparent difficulty of modelling such data (Cacoullos) 1966), generative modelling methods have been resoundingly successful at synthesizing photorealistic images (Rombach et al. 2022 among other achievements. However, the empirical success of generative modelling has come with unexpected phenomena, such as exploding inverses (Behrmann al., 2021 and high densitics placod on out-of-distribution data (Nalisnick et al. 2019), possibly as a result of its complex geometric (Loaiza-Ganee eet al. 2022) topological (Cornish et al., 2020) structure. These pathologies underscore the necessity of designing mathematically principled generative models whose inductive biases encode reasonable assumptions about the data.
In this work we focus on building a density estimator that correctly specifies the topology of the data manifold. With a thorough analysis of related literature, we highlight that most models for densitics on manifolds are pushforward models: neural networks fo R + R trained transform an m-dimensionnl prior into flexible distribution the data manifold embedded in R As we show, such methods likely

---

Published in Transactions on Machine Learning Research (12/2023)
Data
Fg
E
Ground Truth
Pushforward
EBIM
Figure In the top row, our EBIM method is depicted on simulated circular data from a von Mises distribution. From left right: ground truth sample of von Mises data; a manifold-defning function Fo learned from the data; and an ambient energy Eg trained with constrained Langevin dynamics on the learned manifold. The energy takes the lowest values in areas of high density, precisely one would expect. In the bottom row, manifold learning and density estimation results from the resulting model are juxtaposed with a pushforward bascline. From left to right: the ground truth; pushforward energy-basee model; and an energy-based implicit manifold (ours). By defining the manifold with a constraint Fo(x) = 0, our method can accurately model data with non-trivial topologies.
to misspecify all but the simplest of data topologies, exposing fo to frontier of tradeoff's between expressivity and numerical stability (Salmona et al. 2022 This situation is a potential source of the aforementioned pathologies.
To model a much broader class of topologies, we propose a new generative model, the energy-based implicit manifold (EBIM). Our approach consists of two steps, which are outlined in Figure 1 Armed with a funda- mental theorem from geometry, first describe a novel way model data manifolds: implicitly, the zero set of neural network Fg. Next, model the density within the manifold with constrained energy-bascd model, likelihood-based model which uses constrained Langevin dynamics to sample points the learned manifold. We show that EBIMs can be composed with each other akin to standard energy-based models (Hinton 2002): manifold-defining functions Fe along with their energies Eg can be combined to take unions and intersections of data manifolds in what we call manifold arithmetic. We demonstrate theoretically and empirically that FBIMs can learn manifold-supported densities more accurately than pushforward models. Given its success modelling low- and high-dimensional data alike, EBIMs represent fruitful direction for mathematically grounded modelling of complex datascts.
2 Background, Related Work, and Motivation
2.1 Modelling Manifol--Supporter Data
Manifold structure Commonly, the distribution pa of interest is supported on a submanifold M embed- ded in the ambient space RT. For example, the manifold hypothesis states that real-world high-dimensional data tends to have low-dimensional submanifold structure (Bengio al. 2013). Elsewhere, data from engi-
neering or the natural sciences can be manifold-supported due to smooth physical constraints (Mardia al.
2007 Boomsma al. 2008 Brehmer & Cranmer 2020 Cresswell al.
2022
2

---

Published in Transactions on Machine Learning Research (12/2023)
We will seek to estimate probability density this manifold. Formally, suppose {z,} is set of samples drawn from probability measure P supported on M, an m-dimensional Riemannian submanifold of RR. We focus on the case where m < n, so that M is "infinitely thin" in IR", meaning P does not admit a probability density with respect to the standard Lebesgue measure. However, we may assume it has a density (x) with respect. to the Riemannian measure of M. We elaborate formally on this setting in Appendix [A]
Models for manifold-supported data have long been of interest in statisties, machine learning, and various applications (Pless & Souvenir, 2009 Diaconis et al. 2013 McInnes al.) 2018 A common theme in machine learning has been to account for attribute performance to underlying manifold structure in data (Ozakin & Gray, 2009 Rifai et al. 2011). In particular, a number of past works have explored Monte Carlo methods on manifolds (Brubaker et al. [2012] Byrne & Girolami 2013 Zappa et al.) 2018), which we put to use here. However, the problem of simultaneously learning a submanifold and an underlying density has only become of interest in tandem with recent advances deep generative modelling (Brehmer
& Cranmer 2020). To our knowledge, all such models fall under the umbrella of pushforward models.
Density estimation with pushforward models elled as the pushforward of some latent distribution:
When manifold-supportaal, P is most commonly mod-
2 ~~ p>), x fo(),
(1)
where ʃo: : Imm + RR" is a smooth mapping given by a neural network and 2 ~ Po(2) is a (possibly trainable)
prior on m-dimensional latent space. The resulting model distribution Pe.v supported on the model
manifold f Ma == (I))). This framework encompasses generative adversarial networks (GANs) Goodfellow
ct al. 2014 Arjovsky et al., 2017), injective flows (Brehmer & Cranmer
2020:
Caterini et al. 2021 ; and
various regularized autoencoders (Makhzani el al. 2015 Tolstikhin et al. 2018
Ghosh et al.]
2020 Kumar
et al.
2020). Since we take the support to be an m-dimensional submanifold, we rule out models with full-
dimensional support, such as bijective normalizing flows (contimuous or otherwise) (Rezende & MMhhmmdd
2015 Dinh et al., 2017; Chen et al. 2018), diffusion models (Sohl-Dickstein et al. 2015 Ho et al. 2020 Song
et al. 2021b), and variational autoencoders (VAEs) (Kingma & Welling 2014: Rezende et al. 2014).
However, not all pushforward models have explicitly computable densities. In recent work, Loniza-Ganem et al. (2022) outline general procedure for density estimation with pushforward models, which sepa- rates modelling into two components: a generalized autoencoder, which embeds the data manifold into m-dimensional latent space, and a density estimator, which learns the density within the manifold. The generalized autoencoding step treats fo as a decoder, pairing it with a smooth encoder 9s : :" + R", and trains them to learn M by mutually inverting each other on the data@ as by minimizing a reconstruc- tion loss P ||æ - (x)g) The density cstimator (2) is then fitted to the encoded data {go(x:)} via pexximum-likelihood Given a datapoint M,, two-step models estimate p² (z) as follows:
(x.v(z) p=(z) |det JJ. (=))))))) 1/2
(2)
where =: (2) is the encoding of x and Ji is the Jacobian of fo with respect to its inputs z. The fidelity of this estimate depends on the condition fx() = x for all x E M: in other words, 9s must be a right-
inverse of fe on M. Injective flow models (Brehmer & Cranmer 2020 Caterini et al.] 2021 Kothari et al.,
2021 Ross & Cresswell 2021) enforee invertibilityy on Me with architectural constraints; other two-step
models (Xiao al. 2019 Ghosh el al., 2020 Rombach et al. 2022 like Loaiza-Ganem al. (2022), achieve
this condition at their pon-paraett optimum.
Topological challenges Despite the broad applicability of this density estimation procedure, the requisite condition is effectively impossible to satisfy for general manifolds M. If fe(gs(x)) = r for all x E M, then by definition, Ie smoothly embeds M into R.. This condition presents an immediate topo- logical challenge: M is an mn-diennsional manifold, which in general cannot be embedded in m-dimensional Euclidean space. In line with the strong Whitney embedding theorem
(Lee 2012 pg. 1351, M might not be
Mg may formally be a manifold if fo is not an embedding because the resulting image can "self-intersect" but this distinction can he ignored in practice for density estimation models, we will soon justify. 2I particular, le becomes a left inverse of 96 on M, and go becomes right inverse of Ja on M.
3

---

Published in Transactions on Machine Learning Research (12/2023)
embeddable in Euclidean space of less than 2m dimensions It is thus impossible in the general case for the support of the prior (=) to match M topologically; see the botiom-middle of Figure [ for an example.
In the presence of this topological mismatch, one might optimistically hope that Ma can sufficiently ap- proximate M with enough capacity and training. However, Cornish et al. (2020) show that when this is possible, the bi-Lipschitz constant of fo will diverge to infinity; rendering fo either analytically non-injective or numerically unstable, and making density estimates unreliable (Behrmann al., 2021]. Accordingly, the topological woes of pushforward models cannot be "brute-forced" into submission.
Awareness of the data manifold's topology may be necessary for downstream applications such as defending against adversarial examples (Jang et al. 2020) or out-of-distribution detection (Caterini & Loaiza-Ganem, 2021). In the injective normalizing flows literature in particular, there has been interest in learning manifolds with multiple charts (Kalatzis et al.) 2021 Sidheekh et al. 2022), which are certainly more expressive than using a single chart. Thus far, such approaches require ancillary models for inference, which can complicate density estimation, and must set the number of charts as a hyperparameter. Multiple charts also may not overlap perfectly, misspecifying the manifold.
2.2 Implicitly Defined Manifolds
The aforementioned limitations of pushforward models stem from the inability of smooth embeddings of T. to characterize anything but the simplest of manifolds. A richer class of manifolds can be defined implicitly, given by the following fact from differential geometry (Lee 2012 pg. 105):
The full-rank zero set theorem Let U S R be an open subset of R and let F : U - R= be a smooth map. If the Jacobian matrix JF of F has full rank on its zero set ([0)) == {x E U F(x) = -}, then F-1 ({(}) is a properly embedded submanifold of dimension m in R
In this paper, we exploit this theorem by constructing neural network Fo and defining a new model manifold Me F.' We call Fg the manifold-defining function (MDF) of Mg. We refer to such manifolds implicilly defined or implicil. These are not to be confused with the unrelated term implicil generalive model, which has been used to describe both energy-based models (Du & Mordatch 2019) and some types of pushforward models (Mohamed & Lakshminarayanan 2016
The zero sets of neural networks have been employed with great success for one special type of manifold: 3D shapes (Niemeyer& & Geiger| 2021). An active subcommunity has formed around learning implicit 3D shapes with varying types of supervision, such as priori shape information (Chen & Zhang 2019 Meschoder
2019 Park et al., 2019) or 2D images of the object (Niemeyer al.) 2020). For our context. Gropp ct al.
et al.
2020 propose the most relevant method, which learns coherent shape from a point cloud without
supervision by regularizing gradients. We can reinterpret this as manifold learning, but it can only be applied in the restrictive setting where m = n-1. Here we propose a way to fit F9 to data manifolds of any dimension m embedded in any dimension n 2 m.
2.3 Energy-Based Models
Energy-based based models (EBMs) have a long history in machine learning (Lecun et al. 2006) and physics (Gibbs. 1902). The energy-basce model represents probability density on R" with an energy model Eg : IR" + R by way of the relation
(x) = (x(x)
(3)
JRm dz)
EBMs can be trained for maximum likelihood using contrastive divergence Hiitoo 2002). If Pe represents the sampling distribution of the EBM with energy E,,, the likelihood with respect to single observation has the following gradient:
V v log -V Ev EE IV &
(4)
A naive solution would be increase the model's latent space dimensionality 2m; however, this would make the encoded data (9a(x,)) singular in R22, invalidating density estimates.
4

---

Published in Transactions on Machine Learning Research (12/2023)
The first gradient term minimizes the energy the true (positive) observations, while the second maximizes the energy on (negative) samples generated by the model over the course of training. As a result, EBMs can be trained by minimizing
Ep-P- Pe] E (z) - Eo(x'),
(5)
where sg|-] denotes the stop gradient operator commonly available automatic differentiation libraries such PyTTorch (Paszke et al. 2019), meaning that gradients not propagated through the sampling process during training.
Xie et al. (0166) introduced the first deep EBM for generative modelling. Notably, their method uses Langevin dynamics [Welling & Teh 2011), a continuous MCMC algorithm, to generate samples from the model Pe during both training and test time. One iteration Langevin dynamics from point x!1) to x(++1)
consists of the following update:
E² - 2
((+11)= x(() er
(6)
where is hyperparameter for the step size and iterations of Langevin dynamics samples necessary to obtain samples close to the stationary distribution,
N(0, I..) is Gaussian noise. In practice, many
Pg. Strategies to assist convergence have since become popular in the literature Nijkamp et al. 2019 Du
& Mordatch 2019
Grathwohl et al. 2020a. Nijkamp et al. 2020).
Some past works incorporate energy-based models with pushforward models. Xiao ot al. 2021 model
an EBM the latent space of a VAE, but its training procedure maximizes full-dimensional likelihoods, making it unsuitable for density estimation on manifolds. Che al. (2020) and Arbel al.
(2021) construct
pushforward EBMs by using GAN discriminators to refine the generator's distribution; these models produce distributions manifolds, but do not admit density estimates. Yoon et al. (2021) propose normalized autoencoders (NAFs), which treat. the reconstruction error of an autoencoder as an energy function with the goal of improving its out-of-distribution detection capability, but like Xiao et al. (2021), it is trained as a full-dimensional model, making it unsuitable for manifold-suppont density estimates.
3 Method
Our method comprises two steps. In the first step (Section 3.1 we implicitly capture the data manifold using a neural network Fg-. In the second step (Section 3.2), we train a second neural network Eg- to capture the data density within the manifold as an energy. Together, these neural networks form a single distributional model, which call an energy-based implicil manifold (EBIM).
3.1 Neural Implicit Manifolds
In this section, we describe a practical procedure for modelling data manifolds implicitly using neural net- works. Let Fe R" + Rn-m be smooth neural network with parameters 0; our goal is to optimize it become manifold-defining function for M, the data manifold. According to the full-rank zero set thcorem of Section 2.2 Fa needs to meet three conditions:
Fx( = for all x E M.
2. Fo(a') + O for all x M.
3. JF, (x) has full rank for all x M.
Since M is the support of P*, condition can be encouraged by the constraint Ez-p- [F)(z)] = 0, which we achieve by minimizing Ex-p- [x)()))]
Condition Condition&Jrqqureess requires that Fg has no zeros outside of the true data manifold. We can satisfy this constraint by locating off-monifold points x' R for which [[F((z')] is small and maximizing this norm these points. This goal bears resemblance to the negative sample term in contrastive divergence, where here [o(()))
5

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 2: Manifold defining functions Fo trained without regularizing negative samples or Jr.- On the left, a regular neural network, has become completely flat; ((0)) is the entire space. On the right is the left- inverse of an injective flow, whose Jacobian has full rank analytically, but becomes numerically non-injective. These should be contrasted to the second pane of Figure
acts energy function. Taking inspiration from the EBM literature, thus use Langevin dynamics to sample these negative points from the distribution Po given by the energy [o)(-)) Since [F)) is a non-negative scalar equal to zero only if æ is on the manifold, it is analogous to the reconstruction loss of an autoencoder. Treating it like an energy thus parallels the training procedure of NAEs (Yoon et al, 2021). We emphasize, however, that maximizing || [>a)]l on negative samples is a regularizer to satisfy the full-rank zero thcorem which bears useful resemblance to contrastive divergence. Its purpose to fit a manifold, not to train an EBM.
Condition is equivalent to ensuring all singular values of JF, (z) are nonzero for x E M. To achieve this condition, we take inspiration from Kumar et al. (2020). Given their decoder fo and 2 E R", they make their Jacobian (=) injective by bounding (z)ul away from zero for all unit-norm vectors u E Ik.. To
implement this bound for all unit-norm u E I2m in practice, they sample u uniformly from the unit sphere and minimize the following regularizer with respect to 0:
EgsPe U(Sm-11
[ - (1
(7)
where U(Sm-1) is the uniform distribution the unit sphere 5 = {u R" || 1), ( )+ denotes the ReLU function, and 7) is a hyperparameter that determines the minimum singular value of Jr (z). We can do the same, except by bounding ||" Jr, (x)]l away from zero for all unit-norm E R=- since we seek to make JF, (x) surjective This serves a similar purpose to the 3D) shape learning objective of Gropp et al. (2020), which bounds the L2 norm of the gradient Jr, (z) away from zero. However, their process docs not readily generalize to higher dimensionalities..
Combining these techniques, we propose the following loss:
E EgpP po [IF((=)] allFa(2')1 8(n 1 1|0 Jr, (=)]
(8)
where here U(5 is the uniform distribution on the unit sphere in (d-m)-dimensiooal space Sn-m {o E IR"-m, ||a||| = 1} and a, 3, and are hyperparameters determining the negative sample weighting, the rank-regularizztiion weighting, and the minimum singular value of JF,, respectively. The three terms of this loss respectively encourage conditions 1, 2, and described above. In high- diggh-immnsional settings, we find it useful to constrain the marimum singular value to n as well; this implemented by replacing the ReLU (-)4 with the identity function. Otherwise, the MDF has a propensity to become very steep during training, destabilizing the negative sample generation process and the resulting gradients of the model weights.
Empirically speaking, the two regularization terms obviate degeneracy in the MDF. Without regularization, Fg will converge towards the zero-functon, even if we enforce analytical surjectivity in the Jacobian by
Note we here referring matrix injective (resp. surjective) it has full column (resp. row) rank.

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 3: Manifold arithmetic with an implicitly learned sphere. From left to right: a spherical distribution learned by an EBIM; the union of two copies of the same model translated in different directions; and the intersection of the same two copies.
right). The unregularized flow sends E- [x)[] + 0 by bringing its singular values arbitrarily close to zero without learning the manifold. It effectively becomes numerically non-injective, akin observed instabilities bijective flows
structuring Fo as the left-inverse of an injective flow Kothari et al. 2021| (Figure 2
(Cornish al. 2020 Behrmann al. 2021).
Expressivity Making the simplifying assumption that neural networks can embody any smooth func- tion (Hornik et al. 1989; Csáji 200T), we may compare the expressivity of neural implicit manifolds with pushforward manifolds. Pushforward models can model densities on preciscly those manifolds which are homeomorphic to subset of R"
On the other hand, broader class of manifolds can be modelled implicitly. .M can be represented implicitly if and only if it satisfies the technical condition that its normal bundle is "trivial" (Lec) 2012 27.. Non-trivial normal bundles are not commonly seen in low-dinensional examples except in non-orientable manifolds such as the Möbius strip or Klein bottle. Though is unclear whether the manifolds of most natural datasets have trivial normal bundles (e.g. Carlsson et al. (2008) find a dataset of image patches to have the topology of Klein bottle), is certainly a broader class than pushforward models can capture.
Manifold arithmetic Some datasets might satisfy multiple constraints, which one might want to learn separately before combining into a mixture or product of models. Since implicit manifold learning can interpreted as learning a set of constraints, neural implicit manifolds exhibit composability similar to energy- based models (Hinton 2002 Mnih & Hinton 2005). If Fi and F2 are MDFs for Mj and M2 respectively, then the union M1 UM2 is the zero set of the product of functions x Fi (2)) (x:). Concatenating outputs into the function + (z),, F2(x)) instead produces the intersection Mi M22 We note that Mi U UM2 and Mj n M2 need not be manifolds anymore, meaning we can combine MDFs to form complex structures that cannot be described with a single manifold (Figure 3). Taking intersections and unions could, for example, be used to model conjunctions or disjunctions of data labelled with multiple overlapping attributes (Du al. 2020). or model manifolds with components different. dimensionalities (Brown al. 2023
3.2 Constrained Energy-Based Modelling
Suppose we have used the techniques from the last section to come up with a fully trained neural implicit manifold Mg which is henceforth fixed. In this section we introduce the second step in our procedure, the constrained energy-based model (CEBM), for density estimation on Mø.. Let Egp Ra - IR be an energy function represented by a neural network and define the corresponding density as follows:
c (x)
E M
(9)
MM (=E(t")du(o")

---

Published in Transactions on Machine Learning Research (12/2023)
where du can be equivalently thought of as the Riemannian volume form or Riemannian measure of Mg. (see Appendix A for details). Let Por., be the resulting probability measure (we can think of Po as a probability distribution characterized by both the manifold Mg. and the density por,)). Since the energy Eg is defined on the full ambicnt space I" but the corresponding model is defined only from its values on Ma', we refer to pe-,y s as a constrained energy-based model.
Having defined per, .b and fixed the manifold Mag, we seek to maximize log-likelihood on the data via gradient- based optimization of Ev. Since the denominator SMp- e -Fa(a') du(x)) is an intractable integral, we resort to contrastive divergence:
V log por = Fip + Po Ra(')))
(10)
Importantly, the right-most most term in Equation 10 is an expectation taken over Po so samples from the model are required for optimization.
Constrained Langevin Monte Carlo How can one sample from p Du & Mordatch (2019) use
Langevin dynamics, continuous MCMC method, to sample from deep EBMs. For constrained EBMs, standard Langevin dynamics is insufficient, as it will produce off-manifold samples from the energy. We need manifold-aware MCMC method.
One such method is constrained Hamiltonian Monte Carlo (CHMC), a family of Markov chain Monte Carlo models for implicitly defined manifolds proposed by Brubaker et al.] (2012). Our main contribution in this section, aside from defining constrained EBMs, is to show that CHMC - which is typically applied to analytically known manifolds can be adapted to manifolds that are defined by neural networks. Tn particular, we show how to avoid the unstable and sometimes memory-prohiiiiive operation of explicitly constructing the Jacobian of Fgr, which features prominently in CHMC.
We focus on the special case of constrained Langevin Monte Carlo (CLMC). Fixing step size and omitting most parameter subscripts for brevity, one iteration from position (11) to ((++1) requires two steps:
1. Sample a momentum r ~ I,) conditioned on membership of the tangent space of MM- at a,(0) This can be done by sampling r' N(O, In) and projecting onto the null space of JFg- (x(1) (written as Jp for clarity):
r:=r' (J Jrr'.
(11)
2. Solve for the new position (+++)) using a constrained Leapfrog step, which entails solving the following system of equations for (((-1) and the Lagrange multiplier E R=-m:
((+1)) (1) + gr - =² V.E( - =² 2
2 - )
(12)
(++(+))) 0.
(13)
Now describe how Equations 11 and 12 can be computed without constructing Jp. With access to efficient Jacobian-vector product (jvp) and vector-Jacobian product (vjp) routines, such as those available in functorch (He & Zou 2021 any expression in the form of Jrw = (F,))))) for w E IR" or JT =
Jo) vip(u, J) for e E gn-m is tractable. Furthermore, the inverse term on the right-hand side of Equation 11 can be computed with inspiration from work in injective flows by [Caterini et al. (0021)
who overcome a similar expression using the conjugate gradients (CG) routine (Nocedal & Wright, 2006 Gardner et al., 2018 Potapczynski et al. 2021) and their forward- allows us to compute expressions of the form A-1 particular, CG requires access only to the operation + Av, not the matrix A itself. In our case, b= Jpr', Jacobian-vector product, and the operation is e ++ JrJJJ u, which is again computable as a vector-Jacobian product followed by Jacobian-vector product. Since Jp is a wide matrix, this operation is most efficiently performed using backward-mode followed by forward-mode autoodifferentiation, so our method can be termed the backward-forward variant.
auto-differentiation trick. CG
b), where A is an (n (n-11) m) x (n m) matrix. In
Equations 12 and 13 can be combined into a single minimization problem which can be easily optimized by first-order methods like stochastic gradient descent (Robbins & Monro 1951p or second-order methods like
8

---

Published in Transactions on Machine Learning Research (12/2023)
Algorithm Efficient Constrained Langevin Monte Carlo
Require: trained manifold-dcfning function For, energy Ev,, step size e, step count k, initial point x(0)
s' 4 x0x(0)
for == = ,k..,, do
r' ~'~ 1'~NNO (N, In).
// Project r' to tangent space
mVp_JJT(.) - jvp( vip(:, vip(: Fo-)") +-r vip (CG
ivp(For,r")) Fo-)
// Take a constrained Leapfrog step along r X + arg minx Fi- (I'+er - (A, Fo)) "Il
r' x-x' r' ter (.E((r)) - 2 vip Fo-)
return x'
L-BFGS
Byrd et al.
1995
x arg min Fg. 219 + cr - - - e²
=²
(14)
In computationally challenging contexts, we can settle for suboptimal solutions the cost of introducing bias. Once obtained, X can be plugged back into Equation 12 to directly calculate z(++1),
The two steps described above constitute single iteration of CLMC. In practice, many iterations are required to obtain a sample resembling Par.v (Algorithm [] Following Du & Mordatch (2019). we use a sample buffer for 95% of generated samples to assist convergence during training. To obtain completely new samples initialize the CLMC chain, we sample random noise ambient space and project them to Ma. by computing arg ming Fo (ar))l2.
3.3 Method Summary
Here we summarize the procedure to fit our entire model, which we call an energy-based implicit manifold (EBIM). We start with dataset {ri} sampled from pa and belonging to some ground truth manifold M.
We first train a network F& to satisfy the full-rank zero set theorem by optimizing Equation Once fully trained, Fø- is an MDF for the neural implicit manifold Ma..
We next learn the density within MM by fitting a constrained EBM, E... The likelihood of Exp can be maximized using the gradient in Equation [0] but computing this gradient requires sampling from Egg To do so, Once trained, refer to the energy as Egg.
efficient Constrained Langevin Monte Carlo outlined in Algorithm
A trained MDF Fo- paired with a trained energy Eng- together define the manifold and density of the model, Po- w- We call this (Fa-, Ege) pair EBIM.
4 Experiments
In this section we demonstrate the efficacy of EBIMs on diverse range of topologically non-trivial data. Our code is written in PyTorch (Pasuke el. al. 2019 We use GPyTorch (Gardner et. al. 2018 for conjugate gradients and the marching cubes algorithm of Yatagawa (2021) to plot 2D implicit manifolds in 3D. We generate synthetic data with Pyro (Bingham et al. 2019). Network architectures, hyperparameter settings, and further experimental details can be found in Appendix
4.1 Modelling Non- Trivial Topologies
The current literature on density estimation for non-trivial topologies assumes the manifold is known be-
forehand (Gemici et al. 2016; Mathieu & Nickel 2020 Rezende et. al. 2020 De Bortoli al. 2022 Here
9

---

Published in Transactions on Machine Learning Research (12/2023)
show that EBIMs are the best choice for such distributions in the absence of priori knowledge of the manifold. We reiterate that all manifolds learned in these experiments are determined only based on samples, without additional knowledge. Quantitative comparisons of density estimates are challenging when manifolds are unknown: likelihood values, the usual way to compare donsity estimators, are uninformative for different learned manifolds. Any lest datapoint that is not within the model manifold would result in a model likelihood of because model densities are strictly constrained to their manifolds. Instcad, we grade model densities on the basis of the Wasserstein-1 distance to the ground truth (Table D which is a rigorous way measure the distance between two distributions possibly non-overlapping submanifolds (Arjovsky et al. 2017
As discussed in Section 2 there are many pushforward density estimation models that could serve as a basis of comparison. We focus on a simple pushforward EBM consisting of an autoencoder with an EBM in the latent space. We experimented with regularizing the autoencoder by training with a Gaussian VAE objective, but it did not learn the manifold as well as a regular autoencoder (Appendix B Figure 10). Likewise, one could replace the latent EBM with any density estimator (such a normalizing flow (Brehmer Cranmer 2020) or VAE (Dai & Wipf 2019)), ) but this would not affect the learned manifold. The pushforward EBM is thus a sufficient basclinc for evaluating manifold-lcariing ability on low-dimensional examples.
Von Mises mixture In our first experiment, we evaluate density estimation ability on 1000 points sampled from a mixture two von Mises distributions on circles embedded in 2D. Results for ordinary FBM,, pushforward FBM,, and EBIM are visible in Figure 4 Of note is the topology of the density learned by the pushfor- ward EBM; it is necessarily connected and ap- pears to be homeomopphic to the real line except at two points of self-intersection. The EBIM, in contrast, captures the manifold even in regions of sparsity. The ordinary EBM is not subject to the topological limitations of the pushforward EBM, but still lacks the inductive bias to learn the low intrinsic dimension of the data.
Ground Baats
EDM
Pustfonward
EBIM
Figure 4: Manifold learning and density estimation re- sults on a balanced, disjoint mixture of two von Mises distributions. Four models are depicted: the ground truth, an ambient EBM, a pushforward EBM, and an EBIM (ours).
Gcospatial data Following Mathicu & Nickel 2020 we model a dataset of global flood events from the Dartmouth Flood Observatory [Brakenridge 2010), embedded on a sphere representing the Earth. Despite previous datasets (they only occur on land), the EBIM still the relative sparsity of floods compared perfectly learns the spherical shape of the Earth (Figure 5 The pushforward EBM represents the densities fairly well, but struggles to learn the sphere and places some density off of the true manifold. Note that the EBIMs and pushforward FBMs are plotted using a triangular mesh and mesh grid, respectively, due to the difference in how they are defined.
Amino acid modelling The structure of some amino acids can be characterized by a pair of dihedral angles and thus possesses toroidal geometry. Designing flexible probabilistic models for torus-supported data is consequently of interest the bioinformatics literature on protein structure prediction (Singh et al. 2002 Mardia al., 2007 Ameijeiras-Alouse & Ley 2020), and amino acid angle data. is practical candidate for evaluating the density estimation ability of EBIMs. In Figure pushforward EBM using open-source amino acid dataset available from the NumPyro software package (Phan et al.) 2019). Remarkably, our manifold-definnng function learns the torus well in the presence of sparse data. We postulate this is because the torus is the simplest manifold matching the data's curvature. On the other hand, the pushforward EBM was unable to reliably model the manifolds. This stark drop in performance concerning because one might reasonably expect higher-dimensional datasets to have more complex topologies than a simple torus, but the corresponding misbehaviour of the pushforward model would be impossible to visualize and difficult to detect.
compare an EBIM with
10

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 5: Manifold learning and density estimation results on flood location data. From left to right with two different viewpoints (top and bottom): the ground truth data; pushforward EBM; and an EBIM (ours).
Figure 6: Manifold learning and density estimation results on the glycine angle data. From left to right: the ground truth data; a pushforward EBM; and an EBIM (ours).
Table 1: Mean Wasserstein distances with standard errors over runs (lower is better).
DATASET
PEBM
EBIM (ours)
VoN MISES MIXTURE
0.013 = 0.001
0.002 + 0.001
GEOSPATIAL
0.056 0.017
0.014 0.003
AMINO ACID MODELLING
0.042 = 0.013 0.026 0.001
Numerical results for the above three topologically interesting datasets are given in Table ! In all cases the EBIM more closely replicates the data distribution than the pushforward EBM, and usually with less variation over runs.
4.2 Modelling Image Manifolds
Image generation In this section show that EBIMs can be scaled to digner-dimensional data manifolds:
MNIST (LeCun et al.) 1998) and Fashion MNIST Xiao et al. 2017). For the manifold dimension we select 16 a priori; this value is close lo intrinsic dimension istimates of MNIST and Fashion MNIST in the literature (Pope et al. 2021 Zheng et al., 2022 Brown et al. 2023). We provide three baseline comparisons: a vanilla autocncoder with an EBM in the latent space (AE+EBM), a variational autoencoder with an EBM in the latent space (VAE+EBM), and normalized autoencoder (NAE) YYoon et al. 2021). In short, NAEs are EBMs whose energies defined by the reconstruction loss of an autoencoder and which its latent
11

---

Published in Transactions on Machine Learning Research (12/2023)
3 1
6
7
2
5
1
8
1
3
1
4
9
3
/
7
2
1
5
9
6
0
6
9 1 8
/
7
2 7
0
7
5
0
2
7
2 3
9
/
7
R
4 4 1 11
7
3
9
3
4
[
6
4
4
3
4
9
S
8
4
9
4
0
0
8
5 0 3
0
5
9 1
4
/
0
4
0
9
4
2
1
2
9
8
4
2
3
J
9 O
9
6
2
9
3
4
0
2 8
2
$ } 4 8
4
2
7
22022
2
5
3
7
1
3
7
8
7
V
5 7 7
9
7
1
0
2
9
9
9
5 7 8 4
3 ú 16696621
2
8
8
3
1
4
3
3
2
2
2
7
7
7
7
0
5
Q
6
3
8
6
2
3
3
9 / 3
7 0 500000000
2
4
3
8
2
3
9
3
0
7
2
3
7
2
7
6
9
6
4
6
4
S
0
6
4
6
H
A 3 2 3
1 9
6
8
1
9
6
7
3
3
2
2
7
2
7
8
J
/
7
7
6 5
9
s
5 3 1) / 8
5 6
3
7
7
O
3
S
6
3
4
>
L
7
8
7
2
2
6
3
7
7
9
7
9
3
2
2
(
A
Figure 7: MNIST (top) and Fashion MNIST (bottom) samples. Columns from left to right: AE + EBM, VAE + EBM, NAE, IM (ours, ablation), and EBIM (ours).
Table 2: Mean FID scores with standard errors over 5 runs (lower is better).
DATASET
AE+EBM
VAE+EBM
NAE
IM (ours, ablation)
EBIM (ours)
MNIST
19.8
17.9 0.6
17.0 + 0.3
19.2 + 0.5
15.3 + 1.0
FMNIST
0.3
54.1 + 0.5
28.2 + 2.6
39.3 + 0.3
30.7 + 0.6
space to perform efficient on-manifold initialization of Langevin dynamics. They were designed for out-of- distribution detection and cannot. provide density estimates on manifolds, our stated goal, but we include them for their training similaritics to EBIMs. We also report the quality of samples generated from the
implicit manifolds (LMs) prior to fitting the density with CEBM (i.e. we treat (-))-()) || as energy
function in I. Samples from all models are provided in Figure 7 with FID scores (Heusel et al. 2017 in
Table for reference!!
The EBIM outperforms the two pushforward models, the AE+EBM and VAE+EBM, on both datasets, as well as the NAE on MNIST. However, the EBIM slightly underperforms the NAE on Fashion MNIST. One cause is that Langevin dynamics can be run more efficiently with NAÉs;; it is partly run in low-dimensional latent space, and requires only one network pass per step. In contrast, constrained Langevin Monte Carlo must be run in high-dimensional ambient space and requires many network passes per step (Algorithm D. As a result, we only run one step per training step for image data, whereas NAÉs are trained with 10 latent steps and 50 ambient steps per training step. Even so, NAEs take about 0.5 seconds per training step, while CEBMs take about 3 seconds. Another possible reason for the strength of the NAEs is that they can correct for topological constraints through the NAE's ambient energy function, which allows it to sample off off-manifold points. This topological correction also explains its outperformance of the two pushforward models.
The fact that EBIMs generate images of comparable quality suggests that they make up for the inefficiency of constrained Langevin dynamics with other advantages: namely, an improved ability to learn manifold topologies. In particular, the implicit manifold (IM) by itself was able to generate samples of decent quality using Langevin dynamics as described in Section 3.1 whereas samples we produced from training NAÉs without on-manifold initialization in latent space (i.c. by just training the manifold without using the density within) were not digit- or clothing-like at. all. This result indicates EBIMs may be learning manifolds better, but the sampling procedure within their manifolds is simply less efficient.
$FID scores measure some distance between two probability distributions, but is unclear whether they measure the fidelity of learned manifold and the density within (Stein et 2023). We report as it is the most popular metric, though its suitability for simple greyscale datasets like MNIST and Fashion MNIST is questionahle.
12

---

Published in Transactions on Machine Learning Research (12/2023)
Despite its inefficiency, however, constrained Langevin Monte Carlo demonstrates clear value in the fact that FBIMs produce better samples than the original IMs. This suggests that the CFBM step does indeed refine the density of the 1M and that this refinement translates measurably into sample quality.
Manifold arithmetic In this section, we apply manifold arith- metic to image manifolds. We denote by M; the manifold corre- sponding to class of MNIST. Here we consider two overlapping image manifolds based on subsets of the MNIST manifold: (1) M1 U M, and (2) M, U M7. In our experiment, we train two MDFs; one fitted to cach of Mı U Ma and M4 U UM7 We then take the intersection of the two learned manifolds by concatenat- ing the outputs of the MDFs as outlined in Section [.1 This new MDF should define the manifold (M1 U Ma) n (Ma My). We sample from this intersection manifold using Langevin dynamics described in Section 3.1] An uncurated collection of samples is visible in Figure [ Most samples belong to MM,, expected, but interestingly, some samples are ambiguous and appear to belong to M1 O My as well.
64
4
2
4
4 4
4
4
Figure 8: Uncurated samples from (M1 U.M M7).
We also tried directly computing the union (M, (M, but Langevin dynamics was unstable for the product of two MDFs. Instead, to generate images from unions of manifolds, we recommend sampling from a mixture of the two image manifolds: by first sampling from a categorical distribution indicating which model to sample from, and then from the corresponding manifold. Still, the manifold arithmetic approach to unions is unique in providing implicit representation of the new manifold (eg. for visualization, as in Figure
5 Conclusion
In this paper we observed that most, existing techniques to jointly learn data manifolds and densities can be described as pushforword models. These models must become near-homeomorphisus an overly strong topological limitation, in order to provide reliable density estimates. To circumvent this limitation, we introduced the energy-based implicit model, which was outlined in two parts. First, we proposed to learn the data manifold implicitly with a neural network Fe. We then proposed the constrained EBM, new type of EBM for modelling data neural implicit manifolds. In both cases, we showed how the computation of the Jacobian of Fe can be "tamed" using stochastic estimates and automatic differentiation tricks inspired by the injective flows literature (Kumar et al. 2020 Caterini et al. 2021) which frequently grapples with non-square Jacobians. We used these techniques to model distributions with complex topologics; the resulting efficiency gains allowed us to scale the resulting model up to image data.
Although we have covered the limitations of pushforward models when used for density estimation, we highlight here some of their advantages over our model. Primarily, pushforward models come with latent ions of data, which have myriad uses such explainability and artificial reasoning (Higgins et al. 2017 Mathieu et al.) 2019 and efficient density estimation in the latent space. A promising direction for future work is to combine these benefits with those of EBIMs
Another direction would be to sidestep pushforward models entirely and investigate whether densities on manifolds can be extracted from diffusion models. Diffusion models are known to provide full-dimensional densities once converted into continuous normalizing flows (Song et al., 2021a b). However, these are not densities manifolds. On the other hand, Pidsttrgaac/ (2022) has shown that, if a diffusion model's score function is allowed to explode towards infinity as the timestep approaches O, it is capable of modelling manifold-supported data. Being able to extract densities on the manifold under these conditions would be interesting, though to our knowledge, there is no known way to do so.
Future work might also study extra applications of the manifold recovered by Fg, such as unsupervised out- of- of-distribution detection. Informally, since [[F(())]] measures distance from the data manifold, one might expect its value to be larger for out-of-distribution points. However, anecdotally, we occasionally observed in-
13

---

Published in Transactions on Machine Learning Research (12/2023)
distribution (Fashion-MNISST) points being assigned larger [1)(()) values than out-of-distribution (MNIST) points, mirroring the pathologies observed by Nalisnick et al. (2019) for density estimators. This behaviour suggests that our network for (z)]] is inhibited in this task by similar inductive biases (Kirichenko et al, 2020 An interesting direction would be to understand this problem in the context of manifold learning.
Our model also inherits all the limitations of training FBMs; for example, relies on the assumption that Langevin dynamies converges, which occurs only with infinite steps. Sampling remains slower than normal EBMs due to the complexity of constrained Langevin dynamics. Constrained EBMs might thus benefit from training methods that do not involve sampling, such as the Stein discrepancy (Grathwohl et al. 2020b) score-matching [Hyvärinen 2005 Song & Kingmal 2021; De Bortoli et al. 2022). This might help the model scale to larger datasets like CIFAR-10 (Krizhevsky & Hinton 2009), which we were able to train an implicit manifold on, but on which we found tuning of the CEBM intractable.
Acknowledgments
We thank Emile Mathieu for assisting with the global flood event dataset. We also thank Kin Kwan Leung for our discussions about which submanifolds can be modelled implicitly.
References
Jose Ameijeiras-Aloned and Christophe Ley. Sine-skewed toroidal distributions and their application in protein bioinformatics. Biostatistics, 23(3):685-7700 2020.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. In International Con- ference on Learning Representations, 2021.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pp. 214-223, 2017.
Jens Behrmann, Paul Vicol, Kuan-Chieh Wang, Roger Grosse, and Jörn-Henrik Jacobsen. Understanding and mitigating exploding inverses in invertible neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 1792- 1800, 2021.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspec- tives. IEFE Transactions on Pattern Analysis and Machine Intelligence, 35(8):1798 1198811898-888,, 2013.
Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D Goodman. Pyro: Deep universal probabilistic programming. Journal of Machine Learning Research, 20(28):1-6, 2019.
Wouter Boomsma, Kanti V Mardia, Charles C Taylor, Jesper Ferkinghoff-Borg Anders Krogh, and Thomas Hamelryck. A generative, probabilistic model of local protein structure. Proceedings of the National Academy of Sciences, 105(26):890)) 88332 8937, 2008.
G Robert Brakenridge. Global active archive of large flood events. Dartmouth Flood Observatory, University of Colorado, 2010.
Johann Brehmer and Kyle Cranmer. Flows for simultaneous manifold learning and density estimation. Advances in Newral Information Processing Systems, 2020.
Bradley CA Brown, Anthony I, Caterini, Brendan Leigh Ross, Jesse C Cresswell, and Gabriel Loaiza Ganiz. Verifying the union of manifolds hypothesis for image data. In International Conference on Learning Representations, 2023.
Marcus Brubaker, Mathieu Salzmann, and Raquel Urtasun. A family of MCMC methods implicitly defined manifolds. In International Conference Arlificial Intelligence and Slatistics, pp. 161-172, 2012.
Richard Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. A limited memory algorithm for bound constrained optimization. STAM Journal Scientific Computing, 16(5):1 (6(5):1190-199) 190- 1995.
14

---

Published in Transactions on Machine Learning Research (12/2023)
Simon Byrne and Mark Girolami. Geodesic Monte Carlo on embedded manifolds. Scandinavian Journal of Stalistics, 40(4):825 845, 2013.
Theophilos Cacoullos. Estimation of multivariate density, Annals of the Institute of Statistical Mathematics, 18(1):179-189, 1966.
Gunnar Carlsson, Tigran Ishkhanov, Vin De Silva, and Afra Zomorodian. On the local behavior of spaces of natural images. International Journal of Computer Vision, 76(1):1-12, 2008.
Anthony L Caterini and Gabriel (Stat) Can't Believe It's Not Better! Workshop at NewrIPS 2021, pp. 21 21-26, 2021.
Entropic issues in likelihoddbased OOD detection. In
Anthony L Caterini, Gabriel Loaiza-Ganent Geoff Pleiss, and John P Cunningham. Rectangular flows for manifold learning. In Aduances in Neural Information Processing Systems, 2021.
Tong Che, Ruixiang Zhang, Jascha Sohl-Dickstein, Hugo Larochelle, Liam Paull, Yuan Cao, and Yoshua Bengio. Your GAN is secretly an energy-based model and you should use discriminator driven latent sampling. In Advances in Neural Information Processing Systems, 2020.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenand. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, 2018.
Zhiqin Chen and Hao Zhang. Learning implicit fields for generative shape modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 5939-5948, 2019.
Rob Cornish, Anthony L Caterini, George Deligiannidis, and Arnaud Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International Conference on Machine Learning, pp. 2133- 2143, 2020.
Jesse C Cresswell, Brendan Leigh Ross, Gabriel Loaiza- Ganem, Humberto Reyes- Gonzalez,, Marco Letizia, and Anthony L Caterini. Caloman: Fast generation of calorimeter showers with density estimation on learned manifolds. arXiv preprint arXiv:2211. 1.3380, 2022.
Balázs Csanád Csáji. Approximation with artificial neural networks. Facully of Sciences, Fölvös Loránd University, Hungory, 24(48):7, 2001.
Bin Dai and David Wipf. Diagnosing and enhancing VAE models. In International Conference on Learning Representations, 2019.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thortoon, Yee Whye Teh, and Arnand Doucet. Ricmannian score-basod generative modeling. In Advances in Newral Information Processing Systems, 2022.
Persi Diaconis, Susan Holmes, and Mehrdad Shahshahani. Sampling from a manifold. In Advances in Modern Statistical Theory and Applications: A Festschrift Honor of Morris L. Eaton, pp. 102-125. Institute of Mathematical Statisties, 2013.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In International Conference on Learning Representations, 2017.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. In Advances in Neural Information Processing Systems, 2019.
Yilun Du, Shuang Li, and Igor Mordatch. Compositional visual generation with energy based models. Advances in Newral Information Processing Systems, 2020.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Cham- bon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Rodko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python optimal transport. Journal of Machine Learning Research, 22(78):1 8, 2021.
15

---

Published in Transactions on Machine Learning Research (12/2023)
Jacob Gardner, Geoff Pleiss, Kilian Weinberger, David Bindel, and Andrew G Wilson. GPyTorch: Black- box matrix-matrix Gaussian process inference with GPU acceleration. In Advances in neural information processing systems, 2018.
Mevlana C Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on Riemannian manifolds. wwXiv:611.00330 2016.
Partha Ghosh, Mchdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. In International Conference on Learning Regresentations, 2020.
Josiah Willard Gibbs. Elementary Principles in Statistical Mechanies: Developed with Special Reference to the Ralional Foundalions of Thermodynamics. C. Scribner's Sons, 1902.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets, In Advances in Neural Information Processing Systems, 2014.
Will Grathwohl, Kuan-Chich Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier secretly an energy based model and you should treat it like one. International Conference on Leurning Representations, 2020a.
Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, and Richard Zemel. Learning the stein discrepancy for training and evaluating energy-based models without sampling. In International Conference on Machine Learning, pp. 3732-3747, 2020b.
Amos Gropp, Lior Yariv, Niv Haim, Matan Atzmon, and Yaron Lipman. Implicit geometric regularization for learning shapes. In International Conference on Machine Learning, pp. 3789 3799, 2020.
Horace He and Richard Zou, functorch: JAX-like composable function transforms for Py Porch.. https:
//githuu..com/ytorreeeee httpa//ccccopyoooo////ttotoorddd 2021.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXi::1666.08// 2016.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to local Nash equilibrium. In Advances in Neural Information Processing Systems, 2017.
Trina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mo- hamed, and Alexander Lerchner. beta VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2017.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Newral Computation, 14(8):1711-1800) 18000 2002.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Newral Networks, 2(5):359 366, 1989.
Aapo Hyvärinen. Estimation of non normaizal statistical models by score matching. Journal of Machine Learning Research, 6(24):695-709, 2005.
Uyoong Jang, Susmit Jha, and Somesh Jha. On the nood for topology-awand aware generative models for manifold- based defenses. In International Conference on Learning Representations, 2020.
Varuna Jayasiri and Nipun Wijerathne. labml.i annotated paper implementations, 2020.
Dimitris Kalatzis, Johan Ziruo Ye, Alison Pouplin, Jesper Wohlert, and Søren Hauberg. Density estimation smooth manifolds with normalizing flows. Nix:21066005000 2021.
16

---

Published in Transactions on Machine Learning Research (12/2023)
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Represenlations, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
Polina Kirichenko, Pavel Izmailov, and Andrew G Wilson. Why normalizing flows fail to detect out-of- distribution data. In Advances in Neural Information Processing Systems, 2020.
Konik Kothari, AmirEhsan Khorashadizadeh, Maarten de Hoop, and Ivan Dokmanić. Trumpets: Injective flows for inference and inverse problems. In Uncertainty in Artificial Intelligence, pp. 1269-1278, 2021.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Abhishek Kumar, Ben Poole, and Kevin Murphy, Regularized autoencoders via relaxed injective probability flow. In International Conference on Artificial Intelligence and Statistics, pp. 4292-4301 2020.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to docu- ment recognition. Proceedings of the IEEE, 86(11))2278-2333 2324, 1998.
Yann Lecun, Sumit Chopra, Raia Hadsell, Mare Aurelio Ranzato, and Fu Jie Huang. A tutorial on energy- based learning. Predicting Structured Data, 2006.
John M Lee. Introduction to Smooth Manifolds. Springer New York, 2012.
Gabricl Loaiza-Ganem, Brendan Leigh Ross, Jesse C Cresswell, and Anthony Caterini. Diagnosing and fixing manifold overfitting in deep generative models. Transactions on Machine Learning Research, 2022.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoen- 2015.
coders. or
Kanti V Mardia, Charles C Taylor, and Ganesh K Subramaniam. Protein bioinformatics and mixtures of bivariate von Mises distributions for angular data. Biometrics, 63(2):505- 512,, 2007.
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. In Advances in Neural Information Processing Systems, 2020.
Emile Mathieu, Tom Rainforth, Nana Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autocncoders. In International Conference on Machine Learning, pp. 4402-4412, 2019.
Leland McInnes, John Healy, and James Melville. UMAP: Uniform manifold approximation and projection for dimension reduction. Journal of Open Source Software, 3(2)):861, 2018.
Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, and Andreas Geiger. Occupancy networks: Learning 3D reconstruction in function space. Proceedings of the IEEE/CVF conference on compuler vision and pallern recoguition, pp. 4460 4470, 2019.
Andriy Mnih and Geoffrey Hinton. Learning nonlincar constraints with contrastive backpropagation. In Proceedings of the IEFE International Joint Conference on Neural Networks, pp. 1302,1307, 2005.
Shakir Mohamed and Balaji Lakshminarayanan. Learning in implicit generative models. arXiv:1610.09488, 2016.
Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan. Do deep generative models know what they don't know? In International Conference on Learning Representations, 2019.
Michael Niemeyer and Andrcas Geiger. GIRAFFE: Representing scenes as compositional generative neural feature fields. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11453 11464, 2021.
17

---

Published in Transactions on Machine Learning Research (12/2023)
Michael Niemeyer, Lars Mescheder, Michael Oechsle, and Andreas Geiger. Differentiable volumetric ren- dering: Learning implicit. 3D representations without 3D supervision. In Proceedings of the IEEE/ IEEE/CVP Conference on Computer Vision and Pattern Recognition, pp. 3504 3515, 2020.
Erik Nijkamp, Mitch Hill, Song- Chun Zhu, and Ying Nian Wu. Learning non-convergent non-persistent short-run MCMC loward energy-based model. In Advances Nearul Information Processing Systems, 2019.
Erik Nijkamp, Mitch Hill, Tian Han, Song- Chun Zhu, and Ying Nian Wu. On the anatomy of MCMC-based maximum likelihood learning of suergy-based models. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 5272 5280, 2020.
Jorge Nocedal and Stephen Wright. Numerical Optimizalion. Springer Science & Business Media, 2006.
Arkadas Ozakin and Alexander Gray. Submanifold density estimation. In Advances in Neurol Information Processing Systems, 2009.
Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. DeepSDF: Learning continuous signed distance functions for shape representation. In Proceedings of the IEEE/CVI conference computer vision and pattern recognition, pp. 165 165-174, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjic Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Newral Information Processing Systems, 2019.
Xavier Pennec. Probabilitics and statistics Riemannian manifolds: Basic tools for geometric measure- ments. In Workshop Nonlinear Signal and Image Processing, pp. 194-198, 1999.
Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable effects for flexible and accelerated proba- bilistic programming in NumPyro. arXiv preprint arXic:1912.1155; 2019.
Jakiw Pidstrigach. Score-based generative models detect manifolds. In Advances in Neural Information Processing Systems, 2022.
Robert Pless and Richard Souvenir. A survey of manifold learning for images. IPSJ Transactions on Computer Vision and Applications, 1:83 94, 2009.
Phillip Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. In International Conference on Learning Representations, 2021.
Andres Potapczynski, Lubuan Wu, Dan Biderman, Geoff Pleiss, and John P Cunningham. Bias-free scalable Gaussian processes via randomized truncations. In International Conference on Machine Learning, 2021.
Prajit. Ramachandran, Barret Zoph, and Quoe Le. Searching for activation functions. arXiv: 111.00091.9911 2017.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Con- ference Machine Learning, pp. 1530 1538, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pp. 1278-1286, 2014.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. In International Conference on Machine Learning, pp. 8083 8092, 2020.
18

---

Published in Transactions on Machine Learning Research (12/2023)
Salah Rifai, Yann N Dauphin, Pascal Vincent, Yoshua Bengio, and Xavier Muller. The manifold tangent classifier. In Advances in Neural Information Processing Systems, 2011.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, 22(3):400- (13):400-407, 1951.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resohttin image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 10684 10695, 2022.
Brendan Leigh Ross and Jesse C Cresswell. Tractable density cstimation on learned manifolds with conformal embedding flows. In Advances in Neural Information Processing Systems, 2021.
Walter Rudin. Real and Complez Analysis. McGraw-Hill Inc., third edition, 1987.
Antoine Salmona, Valentin de Bortoli, Julie Delon, and Agnès Desolneux. Can push-forward generative models fit multimodal distributions? In Advances in Neural Information Processing Systems, 2022.
Sahil Sidheekh, Chris B Dock, Tushar Jain, Radu Balan, and Maneesh K Singh. VQ-flows: Vector quantized local normalizing flows. In Uncertainty in Artificial Intelligence, pp. 1835-1845, 2022.
Harshinder Singh, Vladimir Hnizdo, and Eugene Demchuk. Probabilistic model for two dependent circular variables. Biometrika, 89(3):719-77) 723, 2002.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathaa,, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2265-2265, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a.
Yang Song and Diederik P Kingma. How train your energy-based models. arXiv:2101.03285 2021.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Frmon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
George Stein, Jesse C Cresswell, Rasa Hosseinzadeh, Yi Sui, Brendan Leigh Ross, Valentin Villecroze, Zhaoyan Liu, Anthony Caterini, J Eric T Taylor, and Gabriel Loaiza-Ganem. Exposing flaws of gen- erative model evaluation metrics and their unfair treatment of diffusion models. In Advances in Neural Information Processing Systems, 2023.
llya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. In International Conference on Learning Representations, 2018.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference on Machine Learning, pp. 681-688, 2011.
Han Xiao, Kashif Rasul, and Roland Vollgraf. MaShion-MNIST: A novel image dataset for benchmarking machine learning algorithms. arXiv: 708.078..0007 07747, 2017.
Zhisheng Xiao, Qing Yan, and Yali Amit. Generative latent flow. vrXiw:1905.10488 2019.
Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. VAEBM: A symbiosis between variational autoencoders and energy-based models. In International Conference on Learning Representations, 2021.
Jianwen Xie, Yang Lu, Song- Chun Zhu, and Yingnian Wu. A theory of generative ConvNet. In International Conference Machine Learning, pp. 2635 2644, 2016.
Tatsuya Yatagawa. torchmeubes: marching cubes for pytorch, 2021. URL attps://ggttuub.com/aaaas torchmcubes
19

---

Published in Transactions on Machine Learning Research (12/2023)
Sangwoong Yoon, Yung-Kyun Noh, and Frank Park. Autoencoding under normalization constraints. In International Conference on Machine Learning, pp. 12087-12097, 2021.
Emilio Zappa, Miranda Holmes-Cerfon, and Jonathan Goodman. Monte Carlo on manifolds: Sampling
densities and integrating functions. Communécations Pure and Applied Mathematics, (161))2600-- 2647, 2018.
Yijia Zheng, Tong He, Yixuan Qiu, and David Wipf. Learning manifold dimensions with conditional variational autoencoders. In Advances in Neural Information Processing Systems, 2022.
20

---

Published in Transactions on Machine Learning Research (12/2023)
A Formal setting
Here we expand on the formal sctting in which we seek to perform density estimation.
Geometry Let M be an m-dimensional orientable Riemannian submanifold of ambient space R" where m < n. Formally this refers to the pair (M,g), where M S R" is a manifold and g is the Riemannian metric inherited from ambient Euclidean space. In other words, g is the restriction of the canonical Euclidean metric, which is characterized by the standard dot product between vectors, to vectors which are tangent M. The metric g, which is typically implied, captures the curvature information would like to associate with M.
A manifold's Riemannian metric gives rise to a unique differential form known as the Riemannian volume form du, which allows for the integration of continuous, compactly supported, real-valued functions h over the Riemannian manifold Lee
2012):
J&L h du.
(15)
Probability Let { be observed samples drawn from Pt, probability measure supported on M. Since M has a lower intrinsic dimension than IR",, it is "infinitely thin." In other words, P*(M) = while the (Lebesgue) volume M 0, meaning no probability density integrated over the ambient space can be used represent P Formally stated, p is not absolutely continuous with respect to the Lebesgue measure on R
Instead, require a new way to define the volumes of subsets of M. We can then formally define a probability density p' over M and integrate with respect to this volume to obtain probabilities. The volume form du on M is the answer; the probability of a set S C M can be computed as follows:
P*(S) | (x) dp(x).
(16)
We note that the volume form dp from differential geometry is not technically a measure in the sense of measure theory. This obstacle is minor: du can be extended to a true measure by a common measure- theoretic tool known as the Riesz-Markov identify du with measure μe on MM which produces volumes of Borel sets in M and which we call the Riemannian measure of M (Pennec 1999
representation theoren (Rudin 1987). Thus we may
Formally, we require P to be absolutely continuous with respect to μr, and we thus write that p" is the Radon-Nikodym derivative of pa with respect to p" dPa This is the ground-truth density function we seek to model in this work.
B Experimental Details
For all experiments, we use feedforward networks with SiLU activations (Hendrycks & Gimpel 2016 Ra- machandran et al. 2017). All models are trained with the Adam optimizer (Kingma & Ba 2015) with the default PyTorc parameters, except for the learning rate which is set as described below.
B.1 Low-Dimensional Data
All EBMs, EBIMs, and pushforward EBMs in this subsection are trained with a buffer size 1000, from which we initialize each Langevin dynamics sample with 95% probability. We do not use spectral normaliva- tion for EBMs: we found harmed the quality of density estimates. Initial noise for the EBIM is sampled
We focus orientable manifolds because, discussed in Section 3.1 manifolds with non-trivial normal bundles, which include non-orientable mauifolds, canuot be modelled implicitly. This formal problem setting be expanded to non-orientable submanifolds too, but requires slightly different construction than is used in this section. In the reference and sometimes in general, this theorem called the Riesz representation theorem, which can also refer different theorem about Hilhert spaces.
21

---

Published in Transactions on Machine Learning Research (12/2023)
Table 3: Statistics of estimated dataset distances to the manifold.
EBIM
PUSHFORWARD EBM
EXPERIMENT
MIN
MEDIAN
MEAN
MAX
MIN
MEDIAN
MEAN
MAX
MOTIVATING EXAMPLE
0.059 x 10-6
0.30 x 10-2
0.34 x 10-2
0.012
0.376 x 10"5
0.27 x 10-2
0.33 x 10-2
0.153
MANIFOLD ARITHMETIC
2.387 x 10-5
0.77 x 10-2
0.79 x 10-2
0.018
VON MISES MIXTURE
0.006 x 10-5
0.73 x 10-2
0.98 x 10-2
0.045
0.810 x 10-5
0.47 x 10
0.51 x 10-2
0.095
GROSPATIAL DATA
0.132 x 10-6
0.24 x 10-2
0.24 x 10-2
0.009
20 x 10-3
0.16 x 10-2
0.17 x 10-2
0.016
AMINO ACID MODELLING
9.176 x 10-6
2.80 x 10-2
2.86 x 10-2
0.16
6,96 x 10-6
0.92 x 10-2
1.72 x 10-2
0.271
uniformly from a box in ambient space containing the ground truth manifold and then projected to the manifold by solving for arg ming || (a²) with L-BFGS using strong Wolfe line search. Equation 14 is also optimized using a single step of L-BFGS with strong Wolfe line search. All models were tuned by hand for visual performance. Training times are reported below, but we caution that models were not tuned for runtime, so the raw times should not be compared between models to evaluate efficiency. In general the constrained EBM is the slowest, followed by the MDF and ordinary EBMs, then the two stages of the pushforward EBM.
To plot the EBIM densities, we estimate the normalizing constants using Monte Carlo. Since the learned MDFs always provide very good approximations of the true manifolds, we estimate each normalizing constant using uniform samples from the ground Iruth manifold for convenience. To plot the pushforward EBM densities, we estimate the normalizing constants in latent space with Monte Carlo estimates based uniform sampling within the clamped bounds. We then compute pushforward densities with Equation 2
All low-dimensional experiments were performed on an Intel Xeon Silver 4114 CPU.
To evaluate Wasserstein-1 distances, we discretize the space into cells (with a granularity of 100 x 100 for the 2D example and 30 30 x 30 for the 3D examples) and use a large quantity of samples from the model manifold to determine whether the model manifold exists in each cell. We then use these samples to evaluate the average density value per cell. These density values are normalized over all cells and used to compute discrete Wasserstein-1 distance with the pot library (Flamary et al. 2021 where distances between pairs of cells are encoded as the distance between their centres.
We provide additional quantitative results in Table 3 Here we estimate the distance of each training point to the manifold using an optimization procedure, and report minimum, median, mean, and maximum distances over the training set.. Nearest-point estimates must. be computed differently for EBIMs and pushforward EBMs, and therefore estimates for each model are prone to different sources of error, so these metries should be used only with caution basis for comparison. For the EBIM with MDF Fø-, we compute the nearest point on the manifold to dalapoint 24
arg mig -xil 10'0 (2')11
where x' has been initialized Tt. For the pushforward EBM with encoder-decoder pair (fo- go.). ) we compute the nearest point as so- (2"), where
2 = arg
min =rg
|| (2) 1,,12
where z has been initialized to gø- (r:).
Motivating example (Figure 1
We sampled 1000 points from a von Mises distribution unit circle centred at (0,0) with the mode located at (1, and concentration of 2.
The MDF for the EBIM consisted of 3 hidden layers with 8 units per hidden layer. The MDF was trained for 300 epochs with a batch size of 50, a learning rate of 0.01, 1, a = 0.3, and B = 10. Langevin dynamics with run with e = 0.1 and a step size of 10. Training took minutes, 8 seconds.
The energy function for the EBIM consisted of hidden layers with 32 units per hidden layer. It was trained for 20 epochs with a batch size of 50, a learning rate of 0.01, gradients clipped to norm of 1, and energy
22

---

Published in Transactions on Machine Learning Research (12/2023)
EBIM
EBIM
EBIM
EBIM
EBIM
EBIM
Figure 9: EBIM manifold learning and density estimation results on von Mises distribution where Langevin dynamics during training has been run (with replay buffer) with different step counts. From left to right and top to bottom, step counts per training step: 1, 3, 5, 10, 20, 40. The setting of 20 langevin dynamics steps is sufficient for convergence.
magnitudes regularized with a coefficient 0.1. Langevin dynamics at each training step were run for 10 steps with & 0.3, a step size of 1, and energy gradients clamped to maximum values of 0.1 at each step. Training took minutes, 51 seconds. In Figure g we evaluate the effect of the Langevin dynamics step count training dynamics, where vary the step size (and remove the training buffer, this effectively increases the average step count). Fewer steps leads distribution is overly smooth when estimating the right-hand side of Equation 10
peaked mode because the estimated model
The pushforward EBM's encoder and decoder each had 3 hidden layers with 32 units per hidden layer. They were jointly trained for 300 epochs with a batch size 50, a learning rate 0.001, and gradients clipped to a norm of 1. Training took 31.6 seconds.
The pushforward EBM's energy function had hidden layers and 32 units per hidden layer. It trained for 200 epochs with batch size of 50, learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with a coefficient of 0.1. Langevin dynamics at each training step were run for 20 steps with E = 0.5, a step size of 10, and energy gradients clamped to maximum values of 0.03 at each step. Training took 3 minutes, 5 seconds.
Manifold arithmetic Figure D depicts two modes of composition for EBIMs. The EBIM depicted on the left is learned from 1000 points sampled from a balanced mixture of two projected normal distributions. Each component was a normal distribution with unit diagonal covariance centred at (1,0,0) and (-1, 0, 0), respectively, before being projected to the sphere. After this, with no additional training, we manipulato it create new probability models. First, two copies of the learned model are translated by 0.5 units in opposite directions.
A new model given by the union of these two copics is depicted in the middle pane of Figure it consists of the product of their MDFs and a balanced mixture of their corresponding energies. Note that the new surface self-intersects, and is no longer formally an embedded submanifold.
Another new model given by the intersoction of these two copies visible in the final panc. By concatenating the output of the MDFs and summing the corresponding energies, we arrive at circle embedded in three dimensions.
23

---

Published in Transactions on Machine Learning Research (12/2023)
Pushforward
Pushforward
Pushforward
Pushforward EBM
Pushforward EBM
Pushforward EBM
Figure 10: Manifold learning and density estimation performance for different weightings 8 on the KL- divergence term of the VAE loss. From left to right, top to bottom: 3 = 0.01, 3 = 0.03, B = 0.05, 3 = 0.1, 0.2, B 0.4.
The MDF for the EBIM consisted of 3 hidden layers with 8 units per hidden layer. The MDF was trained for 300 epochs with a batch size of 100, a learning rate of 0.01, 17 = 0.1, a = 0.3, and S = 10. Langevin dynamics each training step were run for steps with = 0.1, a step size of e², and energy gradients clamped to maximum values of 0.03 at each step. Training took 2 minutes 11 seconds.
The energy function for the EBIM consisted of hidden layers with 32 units per hidden layer. It trained for 20 epochs. We used batch size of 100, learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with a coefficient of 1. Langevin dynamics at each training step was run for 10 steps with = 0.3, a step size of e², and energy gradients clamped to maximum values of 0.03 at each step. Training took 2 minutes, seconds.
Von Mises mixture We sampled 1000 points from a balanced mixture of two von Mises distributions with concentration on circles of unit radius, Respectively, they are centred at (-2, 0) and (2, 0) with modes at (-1,0) and (1,0) (or, polar angles of and with respect to the centre of each circle).
The MDF for the EBIM consisted of 3 hidden layers with 8 units per hidden layer. The MDF was trained for 500 epochs with batch size of 100, learning rate of 0.01, = 1, = 0.3, and 5 = 1. Langevin dynamics at each training step was run for 20 steps with & = 0.1, step size of 10, and energy gradients clamped to maximum values of 0.03 at each step. Training took 3 minutes, 47 seconds.
The energy function for the EBIM consisted of hidden layers with 32 units per hidden layer. It was trained for 10 epochs with batch size of 100, a learning rate of 0.01, gradients clipped a norm of 1, and energy magnitudes regularized with coefficient of 0.3. Langevin dynamics at each training step were run for 10 steps with & 0.5, step size of 1, and energy gradients clamped to maximum values of 0.1 at each step. Training took 1 minute, 23 seconds.
The (ambient) EBM consisted of hidden layers with 32 units per hidden layer. was trained for 200 epochs with a step size of 10. We used a batch size of 100, a learning rate of 0.01, gradients clipped a norm of 1, and energy magnitudes regularized with a cocfficient of 0.5. Langevin dynamics at cach training step were run for 10 steps with 0.1 and energy gradients clamped maximum values of 0.03 at each step. Training took 1 minute, 14 seconds at each step.
The pushforward EBM's encoder and decoder each had 3 hidden layers with 32 units per hidden layer. It was trained for 300 epochs with batch size of 100, a learning rate of 0.001, and gradients clipped to a norm of 1. Training took 34.1 seconds. We also tried training the autoencoder using a variational autoencoder loss, but found that learn the manifold properly, the KL term had be heavily downweighted near the point of nonexistence. In Figure 10 we show how manifold learning ability deteriorates as the KL-weighting is increased.
24

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 11: EBIM manifold learning and density estimation results on the glycine angle data for different values of n, the hyperparameter setting the boundary under which singular values will be penalized by the Jacobian regularization term. From left to right: = 0.3, - 1, 17 = 2, 17 = 3, and 5.
The pushforward EBM's energy function had 3 hidden layers and 32 units per hidden layer. It was trained for 300 epochs with a batch size of 100, a learning rate of 0.01, gradients clipped to a norm of 1, and energy magnitudes regularized with a coefficient of 0.1. Langevin dynamics at each training step was run for 20 steps with E 1.0, a step size of 10, and energy gradients clamped to maximum values of 0.03 at each step. Training took 2 minutes, 21 seconds.
Geospatial data We modelled floods from the Dartmouth Flood Observatory's global active archive, which is available without charge for research and education purposes.
The MDF for the EBIM consisted of 3 hidden layers with 8 units per hidden layer. The MDF was trained for 300 epochs with batch size of 50, a learning rate of 0.01, 17 1, a 0.3, and , = 10. Langevin dynamics at each training step was run for 20 steps with 0.1, step size of 10, and energy gradicnts clamped to maximum values of 0.03 at each step. Training took 19 minutes, 43 seconds.
The energy function for the EBIM consisted of 4 hidden layers with 32 units per hidden layer. It was trained for 100 epochs. We used a batch size of 50, a learning rate of 0.01, gradients clipped to a norm of 1, and energy magnitudes regularized with a coelficient of 1. Langevin dynamics at each training step was run for 5 steps with æ = 0.1, a step size of g2, and energy gradients clamped to maximum values of 0.03 at each step. Training took 1 hour, 52 seconds.
The pushforward EBM's encoder and decoder each had 4 hidden layers with 32 units per hidden layer. They were jointly trained for 500 epochs with batch size of 100, learning rate of 0.001, and gradients clipped to norm of 1. Training took 2 minutes, 33 seconds.
The pushforward EBM's energy function had 4 hidden layers and 32 units per hidden layer. It was trained for epochs with a batch size of 50, a learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with a coefficient of 0.1. Langevin dynamics at each training step were run for 60 steps with E = 0.5, step size of 10, and energy gradients clamped to maximum values of 0.03 at each step. Training took 9 minutes, 9 seconds.
Amino acid modelling The MDF for the EBIM consisted of hidden layers with 8 units per hidden layer. The MDF was trained for 500 epochs with a batch size of 50, a lcarning rate of 0.01, ry = 0.3, a O, and 3 1. We found that increasing 7, the smallest singular value required of Jr, by the regularization term, made the implicit manifold harder to optimize. This occasionally yielded plateaus in the loss function and resulted in incorrect manifolds, depicted in Figure I1 Training took 9.5 seconds.
The energy function for the EBIM consisted of hidden layers with 32 units per hidden layer. It was trained for 10 epochs each wherein Langevin dynamics. We used a batch size of 100, a learning rate of 0.01, gradients clipped to a norm of 1, and energy magnitudes regularized with a cocfficient of Langevin dynamics at each training step was run for 10 steps with & 0.1, a step size of E², and energy gradients clamped to maximum values of 0.03 at each step. Training took 1 minute, 8 seconds.
25

---

Published in Transactions on Machine Learning Research (12/2023)
The pushforward EBM's encoder and decoder each had 3 hidden layers with 32 units per hidden layer. They were jointly trained for 500 epochs with batch size of 50, learning rate of 0.001, and gradients clipped to a norm of 1. Training took 33.5 seconds.
The pushforward EBM's energy function had 3 hidden layers and 32 units per hidden layer. It was trained for 50 epochs with a batch size of 100, a learning rate of 0.01, gradients clipped to a norm of 1, and energy magnitudes regularized with a coefficient of 0.1. Langevin dynamics at each training step were run for 60 steps with E 0.5, step size of 10, and energy gradients clamped to maximum values of 0.03 at each step. Training took 1 minute, 53 seconds.
B.2 Image Data
We parameterized the MDF in our EBM with a small UNet architecture modifed from the implementation in the labml.ai Python package (Jayasiri & Wijerathne 2020) with layer widths scaled down by 75%. We chose
UNet. because its skip connections give it full rank with a large outpui dimennionnliity (28 x 28 16 = 768). All other image model architectures were based on the Conv28 architecture of Yoon et al. (2021). All encoders and decoders were identical to that of Conv28. The constrained energy function for the EBIM models was also identical, except with an output dimension of 1. The energy functions for the AE+EBM, VAE+EBM, and EBIM, as well as the MDF, were trained with spectral normalization.
Nearly the same hyperparameters were used for both datasets, with the exception that a manifold dimension of 30 was used for the AE + EBM on Fashion MNIST (while 16 was used everywhere else).
The autoencoder in the AE + EBM was trained with a learming rate of 1 x 10 and a batch size of 128 for 100 epochs. Gradients were clipped to a maximum norm of 10. The EBM in the AE + EBM was trained with a learning rate of 1 x 10-5 and a batch size of 128 for 200 epochs. Gradients were also clipped to a maximum norm of 10. Energy norms were regularized with a coefficient of 0.1. During training, samples were initialize from the buffer 95% of the time. Langevin dynamics was run for 60 steps with step size of 10, noise of 0.005, and energy gradients were clamped to maximum entries of 0.03.
The autoencoder in the VAE + EBM was trained using the same hyperparameters as in the AE + EBM, and the KL divergence term was trained with a weight of 1 x 10-5, The energy was also trained with the same hyperparameters, except with a learning rate of 1 x 10-4
The hyperparameters used in Yoon et al. (021) for the NAE were chosen for OOD detection; we changed them slightly for gencrative performance. For the latent on-manifold initialization step, we used a step size of 10 and noise standard deviation of 0.005. In ambient space, we used noise with standard deviation 0.005. Models were pre-trained with the reconstruction loss alone for 100 epochs and then trained using the NAE loss for 100 more epochs.
The MDF in the EBIM was trained for 100 epochs with a learning rate of 0.0001 and batch size of 128. Empirically, we found it most effective to sample from and penalize negative samples for [=')) while minimizing [x))l (unsquared) for ground truth data. The negative sample square-norm magnitude was regularized with a coefficient of 0.1. We did not clip optimization gradients. Langevin dynamics gradients were clamped to 0.01 during, and we ran Langevin dynamics for 20 steps per training step. We regularized for minimum and maximun singular value of 0.01, with regularization coefficient of 100000. We sampled from the buffer during training with 95% probability.
The constrained EBM was trained for 100 epochs with a batch size of 64 and learning rate of 0.00001. Enorgy values were regularized with a coefficient of 1. Equation 14 was optimized with stochastic gradient descent with a learning rate of for 60 steps per Langevin dynamics step, 1 of which was run per training step. Langevin dynamics steps were run with step size 10, noise with a standard deviation of 0.005, and energy gradients clamped to values of 0.005. We sampled from the buffer during training with 95% probability.
26


Reviews:
[{"review_id": "Xf98W7tJIZ", "reviewer": "Reviewer_uUEQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "S1. The paper includes a detailed discussion of the background, related work, and motivation, which is very helpful.\n\nS2. The experiments are accompanied by figures illustrating the outcome of the generated images.\n\nS3. Energy-based modeling is a promising direction in light of recent developments in generative modeling.\n\nW1. In section 3, the description lacks a proper problem setup and problem formulation, thus, it is not completely clear what exactly is the problem the proposed algorithm is tackling exactly.\n\nW2. Some of the steps in the algorithm are not completely justified. I'm listing a few here:\n\nIn section 3.1, condition 3 is stated as taking inspiration from Kumar et al. (2020). However, it is not clear to me where this condition is used later in the approach, or why it is needed here.\n\nIn Eq (7),  it is not clear to me where this loss is used later in the experiments. Some notations are not defined here, such as \"sg,\" \"P*.\" It is not clearly explained why U(S) (the uniform distribution) would be used here. Thus, I was not able to understand this loss.\n\nThe discussion concerning \"Expressivity\" does not follow naturally from the paragraph above; the logic is not clear either.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper develops an energy-based implicit manifold approach, which is a generative modeling framework for generating geometric data such as Geospatial data including global flood events, Amino acid modeling, as well as image generation including MNIST digits and Fashion MNIST.\n\nThe main idea of this approach is applying a stochastic gradient descent like method on the energy based model, which assigns a probabilistic density to each data point in the manifold. A key point here is the constrained projection step.\n\nThen, experiments are conducted to validate the constrained langevin monte carlo algorithm on the defined manifold.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- I find the \"Method\" section to be difficult to understand due to a lack of properly defined problem setup. Many of the derivations are not explained in enough detail to understand them step-by-step.\n- The contribution of this paper is unclear to me because a) there is a lack of comparative studies in the experiments; b) the problem that Algorithm 1 is trying to solve is not clearly defined. Thus, a clear summary of contributions would be needed.\n- The examples used in the experiments are fairly simple geometric objects. I wonder how would the new algorithm compare against (say) diffusion models or other more recent generative modeling (e.g., stable diffusion). This would be helpful for appreciating the significance of the results.", "broader_impact_concerns": "This section is not applicable", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "P5JbvqEwKZ", "reviewer": "Reviewer_DZcr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The idea of learning the manifold as a nullset of $F_{\\theta}$ is\n  attractive.\n- The numerical experiments are convincing and seem to be the state of\n  the art of energy-based models.\n\n\n**Weaknesses**\n1. The main drawback of the method is the large computational cost. One\n  has to train two neural networks, both training procedures involving\n  interleaving gradient steps with MCMC. Furthermore, one of the training procedures\n  requires second-order derivatives of the neural network, while\n  the other one requires intermediate steps between each of the MCMC steps to\n  project the MC trajectory back to the manifold. This makes the comparisons\n  to the other algorithms not informative. As detailed in the appendix, it\n  seems like most of the other parameters (amount of training epochs, number\n  of Langevin steps during training and generation) are chosen the same,\n  which means that the samples generated by the other EBM-based models are\n  computationally much cheaper to obtain. This should be at least clearly\n  stated in the main text.\n2. There should be a comparison to diffusion models for manifold data, as they are\n  the state of the art for learning distributions as push-forward measures; it\n  has been shown empirically as well as theoretically that diffusion models\n  are able to sample manifold distributions. The authors should at least motivate\n  why they think that their arguments about pushforward models do not apply here,\n  or if they do. The reviewer thinks this could have to do with the fact that the drift of\n  the reverse ODE in diffusion models will explode, making the optimal ODE flow $f_\\theta$\n  not continuous (and therefore not smooth).\n3. It is unsatisfactory that even in the toy examples (von mises), where\n  the authors had access to the true density/likelihood, they did not use it\n  to compare model performance. An easy check would have been to compare\n  the true likelihood on the training samples to the learned ones.\n4. An experiment about out-of-distribution detection would have been very\n  beneficial. The reviewer sees learning $F_{\\theta}$ and using it to describe\n  the manifold as the main contribution of the paper, and this would have been\n  the most direct application of this part of the work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors of the paper propose a method to do generative modelling for\ndistributions supported on manifolds. They assume that the data lies on an\n$m$ dimensional submanifold $\\mathcal{M} \\subset \\mathbb{R}^n$ and learn the\ndistribution in a two-step procedure.\n\n**Learning the Manifold**\n\nThe main novel idea is to characterize the manifold as the zero-set of a\nfunction $F_{\\theta}$,\n\n$$ M_{\\theta} = F_{\\theta}^{- 1} (\\{ 0 \\}),$$\n\nwhere $F_{\\theta} : \\mathbb{R}^n \\rightarrow \\mathbb{R}^{n - m}$. The\ndimension $m$ of the manifold has to be known beforehand. An ad-hoc loss is\nintroduced to find a candidate $F_{\\theta}$:\n\n$$ \\text{Loss} (\\theta) =\\mathbb{E} [\\| F_{\\theta} \\| -\n   \\alpha \\| F_{\\theta} (x') \\| + \\beta (\\eta - \\| v^T J_{F_{\\theta}} (x)\n   \\|)^2_+], $$\nwhere $x \\sim P^{\\ast}, x' \\sim \\text{sg} [P_{\\theta}], v \\sim \\text{Unif} (\\text{Sphere})$.\nTo train such an $F_{\\theta}$ one has to run MCMC to generate samples from\n$P_{\\theta} \\propto \\exp (- \\| F_{\\theta} \\|^2)$ at each step and perform two\ngradient operations through $F_{\\theta}$ (once w.r.t. its inputs, and\nthen w.r.t. its parameters).\n\n**Learning the Distribution on the Manifold**\n\nAfter finding an optimal $\\theta^{\\ast}$, the distribution is learned as an\nEnergy-based, model, by performing gradient steps on\n\n$$ \\nabla_{\\psi} p_{\\psi} (x) = - \\nabla_{\\psi} E_{\\psi} (x_i)\n   +\\mathbb{E}[\\nabla_{\\psi} E_{\\psi}(x')], $$\n\nwhere\n\n$$ x' \\sim P_{\\psi } \\propto \\exp (- E_{\\psi}) $$\nis normalized to be a probability distribution on\n$\\mathcal{M}_{\\theta^{\\ast}}$.\n\nTo generate samples from $P_{\\psi }$ for training or testing, one has to run\na constrained version of the Langevin algorithm. To that end, one has to solve\nan optimization problem involving $F_{\\theta}$ over an $n - m$ dimensional\nspace, to project the dynamics back onto the manifold. However, since this\noptimization should be initialized close to its optimal value for small\nstep-sizes of the Langevin-algorithm, it should not take many iterations to converge.\n\nFor sampling, one has to run the constrained Langevin algorithm again.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the section named \"Weaknesses\" in \"Strengths and Weaknesses\".\n\nThe main point for me to answer \"No\" for Claims and Evidence is $2.$. The statement \"such methods are likely to specify all but the simplest of data topologies\" is clearly not backed by evidence, as the most prominent example of push forward models, i.e. diffusion models and now flow models are both able to do so. The $f_\\theta$ there is learned implicitly as the flow of an ODE. \n\nEither discussing this point in more detail, showing how diffusion models indeed fail to learn simple data distributions, or removing statements of the kind would change my answer to that question.", "broader_impact_concerns": "The reviewer does not have any ethical concerns concerning this work, except the ones that apply to any work on generative modeling or manifold detection.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YreiQXggyQ", "reviewer": "Reviewer_DBcR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strenghts:\n* The method is theoretically motivated and can provably model a wider class of manifolds than prior *pushforward models*.\n* The framework makes it straightforward to model unions and intersections of neural implicit manifolds and generate samples from them.\n* The experiments demonstrate that the framework is indeed successful at capturing the topology of the data manifold.\n\n## Weaknesses:\n* As the authors have also pointed out, the constrained Langevin algorithm can take a long time to converge and be inefficient in high dimensions.\n* As a result of this inefficiency, the experiments are only run on relatively simple datasets.\n* The method requires choosing the dimension of the data manifold *a priori*, and an unsuitable choice might lead to missing the topology of the manifold, which is the main advantage of this work in comparison with the literature.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a new topology-aware framework for constructing generative models when the data is supported on a lower-dimensional manifold. In this framework, the manifold is represented by the zero set of a neural network, named a neural implicit manifold, and a constrained energy-based model is introduced for density estimation on this manfiold. This method allows modelling more general manifolds, which is in contrast to prior works who define the manifold as the image of a smooth mapping on a low-dimensional latent space, thus can only capture manifolds that are almost diffeomorphic to a Euclidean space. Furthermore, the framework allows a straightforward representation of unions and intersections of neural implicit manifolds. Experiments on low-dimensional datasets with intriguing topologies and image datasets demonstrate the supreriority of the framework, in comparison with prior works, at capturing the topology of the data manifold, as well as generating images with better or comparable quality.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Major Comments:\n* The authors have mentioned that they believe a possible explanation for why NAE outperforms EBIM on Fashion MNIST is due to the computational complexity of the constrained Langevin algorithm. Is it possible to test this hypothesis by showing that with enough compute EBIM will outperform NAE on Fashion MNIST?\n\n* As a recent popular generative model, denoising diffusion probabilistic models don't seem to fit in the framework of pushforward models. Is there a heuristic on why these models may or may not be able to properly capture the topology of the data in comparison with the proposed methods?\n\n### Minor Comments:\n* I am not an expert in this field, but at this point the experiments don't seem to be sufficient to prove the method's applicability in practical settings when data is high-dimensional. Having experiments on more complex image datasets such as CIFAR10 or stating that such experiments are not computationally tractable might help readers better understand the limitations of the current work.", "broader_impact_concerns": "I do not see any immediate implications which would require a Broader Impact Statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: I think the paper is fine as it is. It does not oversell its contributions, the full rank zero set theorem is interesting, and its implementation through (8) is also a nice read. The overall algorithm is sound, though has limitations (slowness of Langevin for that setup) as outlined by reviewers.
Decision: Accept (poster)