Neural Implicit Manifold Learning for Topology-Aware Generative Modelling
Authors: Brendan Leigh Ross
Gabriel Loaiza-Ganem
Anthony L. Caterini
Jesse C Cresswell
Venue: Submitted to ICLR 2023 (2023)
Keywords: Manifold Learning
Unsupervised Learning
Density Estimation
Topology
Differential Geometry
Generative Modelling

Abstract: Natural data observed in $\mathbb{R}^n$ is often constrained to an $m$-dimensional manifold $\mathcal{M}$, where $m < n$. Current probabilistic models represent this manifold by mapping an $m$-dimensional latent variable through a neural network $f_\theta: \mathbb{R}^m \to \mathbb{R}^n$. Such 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. To learn the data distribution within $\mathcal{M}$, we introduce constrained energy-based models, which use a constrained variant of Langevin dynamics to train and sample within a learned manifold. The resulting model can be manipulated with an arithmetic of manifolds, which allows practitioners to take unions and intersections of model manifolds. In experiments on synthetic and natural data, we show that constrained EBMs can learn manifold-supported distributions with complex topologies more accurately than pushforward models.

Under review as a conference paper at ICLR 2023
NEURAL IMPLICIT MANIFOLD LEARNING FOR TOPOLOGY-AWAL DENSITY ESTIMATION
Anonymous authors Paper under double-blind review
ABSTRACT
Natural data observed in R is often constrained to m-dimensional manifold M, where m < n. Current probabilistic models learn this manifold by mapping m-dimensional latent variable through a neural network fa E. + R.. Such procedures, which we call pushforward models, incur a straightforward limitation: manifolds cannot in general be represented with single parameterization, meaning that attempts to do so will incur either computational instability or the inability learn probability densities within the manifold. To remedy this problem, we propose to model M as neural implicit manifold: the set of zeros of a neural network. To learn the data distribution within M, we introduce constrained energy- based models, which use constrained variant of Langevin dynamics to train and sample within a learned manifold. The resulting model can be manipulated with arithmetic of manifolds which allows practitioners to take unions and intersections of model manifolds. In experiments on synthetic and natural data, we show that constrained EBMs can learn manifold-supported distributions with complex topologies more accurately than pushforward models.
INTRODUCTION
Here we focus on the common statistical task of estimating an unknown probability distribution P' using a set of datapoints {a;} C R" sampled from P". Commonly, the distribution of interest lies on an m-dimensional Riemannian submanifold M embedded in the ambient space R", with m n. For example, data from engineering or the natural sciences can be manifold-supported due to smooth physical constraints (Mardia et al. 2007 Boomsma et al.) 2008 Brehmer & Cranmer 2020). In general, the underlying submanifold M may be unknown priori, which calls for us to design models which learn M in the process of learning p*
The typical paradigm for modelling distributions on learned manifolds is pushforward model: neural parameterization fo I -> IR trained to transform an md-dimensional prior into a flexible distribution the data manifold embedded in R (e.g. Arjovsky al. (2017) Tolstikhin et al. (2018); ;
Arbel et al. (202T)). These techniques can generate high-resolution images, but are insufficiently flexible for learning distributions in settings where the manifold structure is of interest.
Modelling a manifold as the image of a single mapping fo is topologically restrictive. For example, many approaches encourage an encoder go and decoder fo to mutually invert each other at each datapoint (e.g. Donahue et al. (2017); Dumoulin et al. (2017); Xiao et al. (2019)), an objective we can precisely reinterpret as training to become diffeomorphism between M and subset of the latent space R". This specification conflicts with the fact that, in general. M may have a complex topology which is not diffeomorphic to any such subset, exposing fo a frontier of tradeoffs between expressivity and numerical stability (Cornish et al.] 2020) Behrrann et al.) 2021: Salmona et al. 2022p. Even when fo is not a diffeomorphism, its continuity dictates many topological properties of the model manifold, such connectivity and the number of holes (Munkres) 2000).
In this paper we learn data manifolds with a much broader class of topologies using a novel approach outlined in Figure [ We first learn manifold implicitly the zero set of a neural network Fe. controlling the manifold dimension by regularizing the rank of its Jacobian. We then model the density within the manifold using constrained energy-based based model Es,, which uses constrained Langevin dynamics to sample points on the learned manifold. We show that constrained energy- based models on manifolds can be composed with each other akin standard energy-based models (Hinton)

---

Under review as a conference paper at ICLR 2023
Ground Truth
Pushforward EBM
Constrained EBM
Figure In the top row, our method is depicted on simulated circular data from a von Mises distribution. From left to right: ground truth sample of von Mises data, manifold-definine function Fo Iearned from the data, and an ambient energy Es trained with constrained Langevin dynamics on the learned manifold. In the bottom row, manifold learning and density estimation results from the resulting model are juxtaposed with pushforward baseline. From left to right: the ground truth, pushforward energy-based model, and constrained energy-based model (ours). By defining the manifold with a constraint Fx(x) = O, our method can model data with non-trivial topologies.
2002): manifold-defining functions F, along with their energies Egg can be combined take unions and intersections of data manifolds in what we call manifold arithmetic. We demonstrate theoretically and empirically that the proposed model can learn manifold-supp suppooted distributions more accurately than the pushforward paradigm prevalent in the current literature.
2
BACKGROUND AND RELATED WORK
2.1 MODELLING VANIFOLD-SUPPORETT DATA
Manifold structure As above, suppose { is set of samples drawn from probability measure P' supported on M, m-dimensional Riemannian submanifold of R". We focus the case where m < n, so that M is "infinitely thin" in R", meaning P* does not admit a probability density with respect to the standard Lebesgue measure. However, may assume it has density p" (Ir) with respect to the Riemannian measure of M. We elaborate on this setup in Appendix
Models for manifold-supported data have long been of interest in statisties, machine learning, and various applications (Diaconis et al. 2013; McInnes et al.) 2018). In particular, a number of past works have explored Monte Carlo methods
manifolds (Brubaker al., 2012 Byrne & Girolamil 2013 Zappa et al., 2018), which we put to use here. However, the problem of simu Itaneously learning a submanifold and an underlying density has only become interest in tandem with recent advances in deep generative modelling (Brehmer & Cranmer 2020).. To our knowledge, all such models fall under the umbrella of pushforward models.
Density estimation with pushforward models When manifold-supported, P* is most commonly modelled as the pushforward of some latent distribution:
(=),
fo(>),
(1)
2

---

Under review as a conference paper at ICLR 2023
where fo : I - R" is smooth mapping given by neural network and ~~ p2)(2) is a (possibly
trainable) prior on m-dimensional latent space. The resulting model distribution Po.v v. is supported on the model manifold MM == (. This framework encompasses generative adversarial networks (GANs) (Goodtellow al. 2014 Arjovsky et al. 2017), injective flows (Brehmer & Cranmer 2020
Caterini et (0021) and various regularized autoencoders (Makhzani et al. Tolstiihin et al. Ghosh et al. 2020) Kumar et al. (020). Since take the support be an m-dimensional
2018]Ghosh
submanifold, we rule out bijective normalizing flows (Rezende & Mohamed| 2015 Dinh et al.) 2017)
and variational autoencoders (VAEs) (Kingma & Welling 2014), unless po (2) is a point mass.
In recent work, Loaiza-Ganem et al. (2022) outline a general procedure for manifold learning and density estimation with pushforward models, which separates modelling into two components: a generalized autoencoder, which embeds the data manifold into m-dimennional latent space, and a density estimator, which learns the density within the manifold. The generalized autoencoding step treats fo as decoder, pairing it with a smooth encoder J6 R + R and trains them to learn MM by mutually inverting each other on the data d such as by minimizing a reconstruction loss Ez-p= || - fo (1))) The density estimator (=(=) is then fitted to the encoded data (z:))
via maximun--lkeelhood. Given a datapoint æ E M, two-step models estimate p" (1) as follows:
(z) (z))) det (=) 1/2
(2)
where = := ===9(<x) (x) is the encoding of x and Jr. is the Jacobian of So with respect to its inputs 2. The
fidelity of this estimate depends on the condition fo(g)(z)) x for all x E M: in other words. Us
must be right-inverse of fe on M. Injective flow models (Brehmer & CCanmer 2020 Caterini et al.
2021 Kotharii al. 2021 Ross & Cresswell (02T) enforce invertibility on Mg with architectural
constraints; other step models (Xiao et al. 2019 Ghosh et al. 2020 Rombach et al. 2022). like
Loaiza-Ganem al. 2022), achieve this condition at their non-parametric optimum.
Topological challenges Despite the broad applicability of this density estimation procedure, the requisite righttinvertibility condition is effectively impossible to satisfy for general manifolds M. If fe(g>(z)) = for all z E .,, then by definition, 9s smoothly embeds M into E". This condition presents an immediate topological challenge: .M is an m-dimensional manifold, which in general cannot be embedded in m-dimensional Euclidean space. In line with the strong Whimey embedding theorem LLee 2013 pg. 135), M might not be embeddable in Euclidean space of less than 2m dimensions It is thus impossible in the general case for the support of the prior p (z) to match M topologically. y; see the bottom-middle panel of Figure for an example.
In the presence of this topological mismatch, one might hope that Ma can sufficiently approximate M with enough capacity and training. However, Cornish et al. (2020) show that when this is possible, the bi-Lipschitz constant of fe will diverge to infinity, rendering fe either analytically non-injective or numerically unstable, and making density estimates unreliable (Behrmann et al.) 2021). Accordingly. the topological woes of pushforward models cannot be "brute-forced" into submission.
Awareness of the data manifold's '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 2022). In the injective normalizing flows literature in particular, there has been interest in learning manifolds with multiple charts (Kalatzis et al.) 202T (Sidheekh et al.) 2022). which are certainly more expressive than using single chart. Thus far, such approaches require ancillary models for inference, which can complicate density estimation, and must set the number of charts as 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 RT to characterize anything but the simplest of manifolds. A richer class of manifolds can be defined implicitly, as given by the following fact from differential geometry (Lee) 2013 pg. 105):
'M may not formally be a manifold if fo is not an emhedding hecause the resulting image can "self- intersect," but this distinction can be ignored in practice for density estimation models, we will soon justify.
In particular, fo hecomes a left inverse go, and equivalently, Is becomes a right inverse of fo- A naive solution would be to increase the model's 's latent space dimensionality to 2m; however, this would make the encoded data {a)(a.) singular in 2m invalidating density estimates.
3

---

Under review as a conference paper at ICLR 2023
The full-rank zero set theorem Let U C R" be an open subset of R", and let F U - Rn-m be a smooth map. If the Jacobian matrix JE of has full rank on its zero set F ((0)) {I e U: F(x) = 03, then F F ((0}) is a properly embedded submanifold of dimension m in R
In this paper, exploit this theorem by constructing neural network Fg and defining new model manifold MMg == F ({0)). We call Fg the munifold-defining fanntinnniniioon function (MDF) of Ma. We refer to such manifolds as implicitly defined or implicit. These are not be confused with the unrelated term implicil generative model, which has been used to describe both energy-based models (Du & Mordatch 2019) and some types of pushforward models (Mohamed Lakshminarayanal 2016).
The zero sets of neural networks have been employed with great success for 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 Mescheeer et al. 2019 Park et al., 2019) or 2D images of the object (Niemeyer et al. 2020). For our context, Gropp al. (2020) propose the most relevant method, which learns a coherent shape from a point cloud without supervision by regularizing gradients. We can reinterpret this as manifold learning, but it does not readily generalize to submanifolds where > m + 1. Here we propose way to fit Fg to data manifolds of any dimension m embedded in any dimension n 2 m.
2.3 ENERGY-BASED MODELS
Energy-based models (EBMs) have long history in machine learning (LeCun et al. 2006) and physics (Gibbs 1902). but Du & Mordatch (2019) introduced the first deep EBM for generative modelling. Notably, they use Langevin dynamics (Welling Teh 2011], I) a continuous MCMC algorithm, to generate samples. Training strategies and applications for MBMs have since become popular the literature (Grathwohl et al. 2019 2020). In particular, Xiao et al. (2021) I) model an EBM in the latent space of VAE, but its training procedure maximizes full-dimensional likelihoods, making it unsuitable for density estimation on manifolds. [he et al. (2020) and Arbel et al.) (202I)
construct pushforward EBMs by using GAN discriminators refine the generator ' distribution; these models produce distributions on manifolds, but do not admit density estimates.
METHOD
3.1 NEURAL IMPLICIT MANIFOLDS
Let F'g : IR" + pn-m be a smooth neural network with parameters 0; our goal is optimize it to become a manifold-defining function for M, the data manifold. Fg thus needs to meet two conditions:
1. (x) = 0 for all x E M.
Figure Manifold defining functions Fe trained with- out regularizing Jr,- On the left, a regular neural net- work, has become completely flat; ([0)) is the entire space. On the right is the left-inverse of an injec- tive flow, whose Jacobian has full rank analytically, but becomes numerically non-injective. These should be contrasted to the second pane of Figure [
2. JF, (x) has full rank for all
x E M.
Since M is the support of P* condition is encouraged by optimizing toward the constraint E--p- [=))(=)]] = which we achieve by minimizing Ez- [x)
We can interpret the 3D shape objective of Gropp et al. (2020) optimizing for condition by bounding the L2 norm of the gradient JF, away from zero However, this does not generalize to generic numbers of dimensions. Null singular values can still be present when, for example, the Frobenius norm or the operator norm of Jr, is bounded away from zero. To maintain full rank, need bound all singular values away from zero, for which we take inspiration from Kumar et al. (2020). Given their decoder fe, they seek to make the Jacobian Jr (z) injective by bounding [ Z) away from zero for all unit-norm vectors u Rm We can do the same, except by bounding || Jr, (x:]] away from zero for all unit-norm E Ra-m, since seek to make Jr, (a)
4

---

Under review as a conference paper at ICLR 2023
surjective Combining these techniques, we propose the following loss:
(())= = Ey v-Utnit())
[[]((=)) + a (7 - (1=- JF. (2)) (2)1))
(3)
where Unif(S) is the uniform distribution on the unit sphere S == for E R | = 1}. (. )+ is the ReLU function, and and 7) are hyperparameters determining the weight of the rank- regularization term and the minimum singular value of JF. respectively. Empirically speaking, the Jacobian regularization term obviates degeneracy in the MDF. Without regularization, Fo will converge towards degeneracy, even we enforce analytical surjectivity in the Jacobian by structuring Fe as the left-inverse of an injective flow (Kothari et al. 2021) (Figure 2 right). The unregularized flow sends E=- (x))ll 0 by bringing its singular values arbitrarily close zero without learning the manifold. It effectively becomes numerically non-injective, akin to observed instabilities in bijective flows (Cornish et al. 2020 [Behrmann et al. 202T).
Expressivity Making the simplifying assumption that neural networks can embody any smooth function (Homik et al. [989 Csáji et al. 2001), we may compare the expressivity of neural implicit manifolds with pushforward manifolds. Pushforward models can model densities on precisely those manifolds which are diffeomorphic to a subset of I!".
On the other hand, broader class of manifolds can be modelled implicitly. .M can be represented implicitly if and only it satisfies the technical condition that its normal bundle is "trivial" (Lee) 2013 pg.271). Non-trivial normal bundles are not commonly seen in low-dimensional examples except in non-orientable manifolds such as the Möbius strip or Klein bottle. Though it is unclear whether the manifolds of most natural datasets have trivial normal bundles (eg. Carlsson al. (2008) find a dataset of image patches to have the topology of a Klein bottle), it is certainly a much 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 mixture or product of models. Since implicit manifold learning can be interpreted as learning set of constraints, neural implicit manifolds exhibit compos- ability similar to energy-based models (Hinton] 2002; Mnih Hinton| 2005). If F1 and F2 are MDFs for M1 and M2 respectively, then the union M, U M2 is the zero set of the product of functions Fi (x) F2(x). Concatenating outputs into the function x + (Fi( F2)))) instead produces the
I + intersection Mi M222. We note that Mi UM2 and Mi n n2 need not be manifolds anymore, meaning we can combine MDFs to form complex structures that cannot be described with a single manifold. Taking intersections and unions could, for example, be used model conjunctions or
disjunctions of data labelled with multiple overlapping attributes (Du et al. 2020
3.2 CONSTRAINED ENERGY-BASER MODELLING
In this section we introduce the constrained energy- model for density estimation on implicit manifolds MM. Let E.+ R + R be an energy function represented by a neural network and define the corresponding density as follows:
e Ex(
(x) JM, e
x E Mo:
where dy can be equivalently thought of as the Riemannian volume form or Riemannian measure of Ma (see Appendix for details). Let Po.y be the resulting probability measure (we think of Pa.y as a probability distribution characterized by both the manifold Ma and the density po.w)) Since the energy Exp is defined the full ambient space R" but the corresponding model is defined only from its values on Mo, we refer to Pa.v as constrained energy-based model.
Having defined pa, and fixed the manifold Mo, we seek maximize log-likelihood the data via gradient- based optimization of Eg. Since the denominator SM, (v) dy in general an intractable integral, we resort to contrastive divergence (Hinton) 2002):
V 인 log pe (z+) V (x.)) + Ez-Po [x))
(5)
"Note we are here referring to matrix injective (resp. surjective) if has full column (resp. row) rank.
5

---

Under review as a conference paper at ICLR 2023
Importantly, the right-most term in Equation is is an expectation taken over Po.y, samples from the model are required for optimization.
Constrained Langevin Monte Carlo How can one sample from Pa ? Du & Mordatch (2019)
use Langevin dynamics, a continuous MCMC method, to sample from deep EBMs. For constrained EBMs, standard Langevin dynamics is insufficient, as it will produce off manifolld samples from the energy. We need manifold-awer MCMC method.
One such method is constrained Hamiltonian Monte Carlo (CHMC), 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 implicitly defined by neural networks. In particular, we show how to avoid the unstable and sometimes memory-prohibitee operation of explicitly constructing the Jacobian of Fe, which features prominently in CHMC.
We focus on the special case of constrained Langevin Monte Carlo. Fixing a step size e and omitting parameter subscripts for brevity, one iteration from position 2x(*) to ((++)) requires two steps:
1. Sample a momentum r ~ (0, In) conditioned on membership of the tangent space of M at (t) This can be done by sampling r' ~ Io, and projecting onto the null space of JF (x)(t)) (written as JfF for clarity):
r r: (JJJJE) -Jp'''
(6)
2. Solve for the new position ( using constrained Leapfrog step, which entails solving the following system of equations for (t(++) and the Lagrange multiplier A E R"""
2² )
(++ x ee 2
V.E(("))
(7)
F(x(+11)) = 0.
(8)
Now we describe how Equations G can be computed without constructing JF, With access to efficient Jacobian-vector product and vector-Jacobian product routines, such as those available in functorch (He& & Zou 2021), any expression in the form of Jew for w R or JT = (v" Je) for E Rd Is tractable. Furthermore, the inverse term on the right-hand side of Equation 6 can be computed with inspiration from work in injective flows by Caterini et al.] (202T) who overcome a similar expresioon using the conjugate gradients (CG) routine (Nocedal & Wright 2006] Gardner cct al. 2018; Potapczynski et al. 202T) and their forward-bacte backward auto-differentiation trick. CG allows to compute expressions of the form A-16, where A is an (n m) x (n m) matrix. In particular, CG requires access only to the operation v + Av, not the matrix A itself. In our case, b Jer', Jacobian- vector product, and the operation is computable as a vector-Jacobil Jr. this operation is most efficiently performed using backward-modd followed by forward-mode auto-differentiation,, so our method can be termed the backward-fore forwardd variant.
JT. which is again product followed by a Jacobian -vector product. Due to the shape of
Equations 7 and cea be combined into a single minimization problem which can be easily optimized by sxcond-order methods such as L-BFGS (Byrd et al. 1995)
(1() + 2 v (x -- 2 JF. JF. (=())) -
X = arg min
Fe
(9)
where, in computationally challenging contexts, we can settle for suboptimal solutions at the cost of introducing bias. We note that L-BFGS outperformed first-order methods like stochastic gradient descent (Robbins & Monro 1951) or Adam (Kingma & Ba 2014) for this task. Once obtained, X can be plugged back into Equation to directly calculate z!! 1)
The two steps described above constitute single iteration of constrained Langevin dynamics. In practice, many iterations are required to obtain a good approximation to sampling from Po.e (Algorithm]) I.. Following Du & Mordatch (2019), samples to assist convergence during training. To obtain completely new samples, we sample random noise in ambient space and project them to Ma by computing arg min, [x)]())] with L-BFGS.
sample buffer for 95% of generated

---

Under review as a conference paper at ICLR 2023
Algorithm I Constrained Langevin Monte Carlo
Require: manifold manifoldddefining function Fg. energy Eg, step size E, step count initial point xo
< Z0 for 1, 1,.....e. , r' ~ (N, In)
do
f- - J r
[Jo FF.) "1Jrr'
x t arg miny
Po
x + or - - V. - - (x)
x x Er V. (r)
end for return æ
4 EXPERIMENTS
The current literature on density estimation for non-trivial topologies assumes the manifold is known beforehand (Gemici et al. 2016 Mathieu & Nickel 2020) Rezende et al. 2020: De Bortoli et al.) 2022). Here we show that constrained EBMs 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 are incomparable for different learned manifolds. Fortunately, we can compare the following manifolds visually.
As discussed in Section 2 the class of pushforward density estimation models is large, and any can serve as basis of comparison. We focus on the most comparable baseline: a simple pushforward EBM consisting of autoencoder with an EBM in the latent space. We experimented with regulariz- ing the autoencoder by training with a Gaussian VAE objective, but it did not learn the manifold as well as a regular autoencoder (Appendix [ Figure Likewise, one could replace the latent EBM with any density estimator (such as normalizing flow (Brehmer & Cranmer 2020) VAE (Dai & Wipf 2019)), but this would not affect the learned manifold.
Our code is written in PyTorch (Paszke et 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). Hyperparameter settings and other details can be found in Appendix
4.1 SYNTHETIC DATA
Density estimation our first experiment, we evaluate density estimation ability on 1000 points sampled from a mixture of two von Mises distributions on circles embedded in 2D. Results for an ordinary EBM, pushforward EBM, and constrained EBM are visible in Figure Of note the topology of the density learned by the pushforward EBM; it is necessarily connected and appears to be diffcomorphic to the real line except at two points of self-intersection. The constrained EBM, in contrast, captures the manifold even in regions of sparsity. The ordinary EBM is not subject the topological limitations of the pushforward EBM, but it still lacks the inductive bias to learn the low intrinsic dimension of the data.
Manifold arithmetic this experiment, we highlight the ability of constrained EBMs to perform manifold arithmetic. Practical applications of this capability are left to future research. Figure depicts two modes of composition for constrained EBMs. The constrained EBM depicted on the left is learned from 1000 points sampled from balanced mixture of two projected normal distributions. After this, with no additional training, we manipulate to 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 copies is depicted in the middle pane of Figure 4 it consists of the product of their MDFs and balanced mixture of their corresponding energies. Note that the new surface self self-intersects, and is no longer formally embedded submanifold.
7

---

Under review as a conference paper at ICLR 2023
Ground
Figure 3: Manifold learning and density estimation results on balanced, disjoint mixture of two von Mises distributions. Four models depicted: the ground truth, ambient EBM, pushforward EBM, and a constrained EBM (ours).
Figure 4: Manifold arithmetic with an implicitly learned sphere. From left to right: spherical distribution learned by a constrained EBM; the union of two copies of the same model translated in different directions; and the intersection of the same two copies.
Another new model given by the intersection of these copies is visible in the final pane. By concatenating the output of the MDFs and summing the corresponding energies, we arrive circle embedded in three dimensions.
4.2 NATURAL DATA
Geospatial data Following Mathieu Nickel (2020), model dataset of global flood events from the Dartmouth Flood Observatory (Brakenrigee 2010), embedded on a sphere representing the Earth. Despite the relative sparsity of floods compared to previous datasets (they only occur on land). the constrained EBM still 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 constrained and pushforward EB Ms are plotted using a triangular mesh and mesh grid, respectively, due to the difference how they are defined.
Amino acid modelling The structure of amino acids can be characterized by a pair dihedral an- gles and thus possesses toroidal geometry. Designing flexible probabilistic models for corus-supporte data consequently
interest in the bioinformatics literature on protein structure prediction (Singh et al. 2002 Mardia et al.) 20077: (Ameijeiras-Aloss & Ley 2022), and so amino acid angle data is practical candidate for evaluating the density estimation ability of constrained EBMs. Figure 6 we compare constrained EBM against pushforward EBM using open-source amino acid dataset available from the NumPyro software package (Phan et al. 2019). Remarkably, our manifold-definin 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.
8

---

overly Under review as a conference paper at ICLR 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 constrained EBM (ours).
Figure 6: Manifold learning and density estimation results the glycine angle data. From left to right: the ground truth data; a pushforward EBM; and a constrained EBM (ours).
5 CONCLUSION
In this paper we observed that all existing techniques to jointly learn data manifolds and densities can be described as pushforward models. These models must become deadddfffomomphismss. strong topological limitation, in order to provide reliable density estimates. To avoid this limitation, proposed to learn the data manifold implicitly with neural network Fo. We then proposed the constrained EBM, a new type of EBM for modelling data on neural implicit manifolds. In both cases, we showed how the computation of the Jacobian of Fa can be "tamed" using stochastic estimates and automatic differentiation tricks inspired by the injective flows literature (Kumar et al.) 2020 Caterini al. 202T) which frequently grapples with non-square Jacobians. Finally, we demonstrated the qualitative efficacy of constrained EBMs on both synthetic and real-world tasks.
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 representations of data, which have myriad uses such as explainability and artificial reasoning
(Higgins et al. 2016 Mathicu 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 constrained EBMs.
Our model inherits all the difficulties of training EBMs; for example, it relies on the assumption that Langevin dynamics converges, which occurs only with infinite steps. Sampling remains slower than normal EBMs due 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. 2020 score-matching (Song & Kingmal 2021 De Bortoli et al. 2022

---

Under review as a conference paper at ICLR 2023
REFERENCES
Jose Ameijeiras-Alone and Christophe Ley. Sine-skewed toroidal distributions and their application in protein bioinformatics. Biostatistics, (8(2343))685--744 2022.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized Energy Based Models. In International Conference on Learning Representations, 2021.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein Generative Adversarial Networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pp. 214-223, 06-11 Aug 2017.
Jens Behrmann, Paul Vicol, Kuan- Chich Wang. Roger Grosse, and Jörn-Henrik Jacobsen, Understand- ing and mitigating exploding inverses in invertible neural networks. International Conference on Artificial Intelligence and Statistics, pp. 1792-1800. PMLR, 2021.
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh. Paul A. Szerlip, Paul Horsfall, and Noah Goodman. Pyro: Deep universal probabilistic programming. J. Mach. Learn. Res., 20:28:1 2-:6, 2019. URLhttp ///jaa/..rrg/ppppr////////////// 40tmmmmmmmmmm11
Wouter Boomsma, Kanti V Mardia, Charles C Taylor, Jesper Ferkinghoff-Bow Anders Krogh, and Thomas Hamelryck. A generative, probabilistic model of local protein structure. Proceedings of the National Academy of Sciences, 105(26):8932 -937,, 2008
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. In Advances Neural Information Processing Systems, volume 33, 2020.
Mareus Brubaker, Mathieu Salzmann, and Raquel Urtasun. A Family of MCMC Methods on Implicitly Defined Manifolds. In Proceedings of the Fifteenth International Conference on Artificial Intelligence and Statistics, volume 22, pp. 161-172, 21-23 Apr 2012.
Richard H Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. limited memory algorithm for bound constrained optimization. SIAM Journal on scientific computing, 16(5): 1995.
Simon Byrne and Mark Girolami. Geodesic Monte Carlo embedded manifolds. Scandinavian Journal of Statistics, 40(4):825-845 2013.
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 Loaiza-Ganem. Entropic issues in likelihood-based ood detection. In (Still) Can't 't Believe It's Not Better! Workshop at NeurIPS 2021, pp. 21-26. PMLR, 2022.
Anthony L Caterini, Gabriel Laneza-Ganem, Geoff Pleiss, and John Cunningham. Rectangular flows for manifold learning. In Advances in Neural Information Processing Systems, volume 34, 2021.
Tong Che, Ruixiang Zhang, Jascha Sohl-Dickstein, Hugo Larochelle, Liam Paull, Yuan Cao, and Yoshua Bengio. Your gan is secretly energy-based model and you should use discriminator driven latent sampling. Advances in Neural Information Processing Systems, 33:12275-12287, 2020.
Zhiqin Chen and Hao Zhang. Learning implicit fields for generative shape modeling. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattrn Recognition, pp. 5939-59499 5948, 2019.
Rob Comissh, Anthony Caterini, George Deligiannidis, and Arnaud Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International Conference on Machine Learning, pp. 2133-2143. PMLR, 2020.
Balázs Csanád Csáji et al. Approximation with artificial neural networks. Faculty of Sciences. Ervs Lornd University, Hungary, 24(48):7, 2001.
10

---

Under review as a conference paper at ICLR 2023
Bin Dai and David Wipf. Diagnosing and Enhancing VAE Models. In International Conference Learning Representations, 2019.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modeling. arXiv preprint arXiv:2202.0276. 2022.
Persi Diaconis, Susan Holmes, and Mehrdad Shahshahani. Sampling from manifold. In Advances in modern statistical theory and applications: a Festschrift in honor of Morris L. Eaton, pp. 102-125. Institute of Mathematical Statistics, 2013.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. ICLR, 2017.
Jeff Donahue, Philipp Krähenbühl, and Trevor Darrell. Adversarial feature learning. ICLR, 2017.
Yilun Du and Igor Mordatch. Implicit Generation and Modeling with Energy Based Models. Advances in Neural Information Processing Systems, 32:3608-366 -618,, 2019.
Yilun Du, Shuang Li, and Igor Mordatch. Compositional visual generation with energy based models. Advances in Neural Information Processing Systems, 33:6637-6647, 2020.
Vincent Dumoulin, Ishmael Belghazi, Ben Poole, Olivier Mastropietro, Alex Lamb, Martin Arjovsky, and Aaron Courville. Adversarially learned inference. ICLR, 2017.
Jacob Gardner, Gcoff Pleiss, Kilian Q Weinberger, David Bindel, and Andrew G Wilson. Gpytorch: Blackbox matrix-matrin gaussian process inference with gpu acceleration. Advances in neural information processing systems, 31, 2018.
Mevlana C Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on Riemannian manifolds. arXiv preprint arXiv:1611.02304, 2016.
Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. ICLR, 2020.
Josiah Willard Gibbs. Elementary principles in statistical mechanics:
developed with especial reference the rational foundations of thermodynamics. C. Scribner's sons, 1902.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Proceedings of the 27th International Conference on Neural Information Processing Eastems-Volune 2, pp. 1672-26880 2014.
Will Grathwohl, Kuan-Chich Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. ICLR, 2019.
Will Grathwohl, Kuan-Chich Wang, förn-Henrik Jacobsen, David Duvenaud, and Richard Zemel. Learning the stein discrepancy for training and evaluating baaegy models without sampling. In International Conference on Machine Learning, pp. 7747- PMLR, 2020.
Amos Gropp, Lior Yariv, Niv Haim, Matan Atzmon, and Yaron Lipman. Implicit geometric regu- larization for learning shapes. In Proceedings of the 37th International Conference Machine Learning. PMLR, 2020.
Horace He and Richard Zou. functorch: Jax-like composable function transforms for pytorch. https://github. wttp///oooooocc//ftttoooooooo 2021.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
1oX1::1666.0845. 2016.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta- VAE: Learning basic visual concepts with a constrained variational framework. 2016.
11

---

Under review as a conference paper at ICLR 2023
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neurul networks, 2(5):359-366, 1989.
Uyeong Jang, Susmit Jha, and Somesh Jha. On the need for topology-awe aware generative models for manifold-based defenses. ICLR, 2020.
Dimitris Kalatzis, Johan Ziruo Ye, Alison Pouplin, Jesper Wohlert, and Søren Hauberg. Density estimation on smooth manifolds with normalizing flows. arXiv preprint arXiv:2106.03500, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6900 69800, 2014.
Diederik Kingma and Max Welling. Auto-encoding Variational Bayes. ICLR, 2014.
Konik Kothari, AmirEhsan Khorashadizadeh, Maarten de Hoop, and Ivan Dokmanić. Trumpets: Injective flows for inference and inverse problems. In Proceedings of the Sirrt-SSventh Conference Uncertainty in Artificial Intelligence, volume 161, pp. 1269-1278, 2021.
Abhishek Kumar, Ben Poole, and Kevin Murphy. Regularized autoencoders via relaxed injective probability flow. In International Conference Artificial Intelligence and Statistics, pp. 4292- 4301. PMLR, 2020.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
John M Lee. Introduction to Smooth Manifolds. Springer, 2013.
Gabriel Loaiza-Ganem, Brendan Leigh Ross, Jesse Cresswell, and Anthony I. 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 autoencoders. ICLR, 2016.
Kanti Mardia, Charles Taylor, and Ganesh Subramaniam. Protein bioinformatics and mixtures of bivariate mises distributions for angular data. Hiometrics, 63(2):505-5122 2007.
Emile Mathieu and Maximilian Nickel. Riemannian Continuous Normalizing Flows. In Advances in Neural Information Processing Systems, volume 33, 2020.
Emile Mathieu, Tom Rainforth, Nana Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autoencoders. In International Conference on Machine Learning, pp. 1402-4412
PMLR, 2019.
Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and
projection for dimension reduction. Journal of Open Source Software, 3(29):861, 2018.
Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, and Andreas Geiger. Occupancy networks: Learning 3d reconstruction in function space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4460-4470, 2019.
Andriy Mnih and Geoffrey Hinton. Learning nonlinear constraints with contrastive backpropagation. In Proceedings. 2005 IEEE International Joint Conference on Neural Nenworks, 2005., volume 2, pp. 1302-1307. JEEE, 2005.
Shakir Mohamed and Balaji Lakshminarayanan. Learning in implicit generative models. arXiv preprint arXiv:1610004433, 2016.
J.R. Munkres. Topology. Featured Titles for Topology. Prentice Hall, Incorporated, 2000. ISBN 9780131816299, URL htttss:/bbbtk goog ca
12

---

Under review as a conference paper at ICLR 2023
Michael Niemeyer and Andreas Geiger. Giraffe: Representing scenes compositional generative neural feature fields. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11453-114661 2021.
Michael Niemeyer, Lars Mescheder, Michael Oechsle, and Andreas Geiger. Differentiable volumetric rendering: Learning implicit 3d representations without 3d supervision. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 3504-3515, 2020.
Jorge Nocedal and Stephen Wright. Numerical optimization. Springer Science & Business Media, 2006.
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/CVF conference on computer vision and pattern recognition, pp. 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, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. Curran Associates, Inc., 2019. URL http:httppppppppppppopm neur httpppprrrr.rrrrpps...ppeepeerre 9015 pywtpch-t--ho.e an -wwaanimmmeeee imperative style hergh--peffrmaaccece deep learning Library.. pdf
Xavier Pennec. Probabilities and statistics on riemannian manifolds: Basic tools for geometric measurements. In NSIP, volume 3. pp. 194-198. Citeseer, 1999.
Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable effects for flexible and accelerated probabilistic programming in numpyro. NearIPS Workshop on Program Transformations, 2019.
Andres Potapczynski, Luhuan Wu, Dan Biderman, Geoff Pleiss, and John P Cunningham. Bias-free scalable gaussian processes via randomized truncations. International Conference on Machine Learning, to appear, 2021.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint 1rXiv:1710.05941 2017.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, pp. 1530-1538. PMLR, 2015.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racaniere, Michael Albergo, Gurtej Kan- war, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. International Conference Machine Learning, pp. 8083-8092. PMLR, 2020.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVE Confer- ence on Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695, June 2022.
Brendan Leigh Ross and Jesse C Cresswell. Tractable Density Estimation on Learned Manifolds with Conformal Embedding Flows. In Advances in Neural Information Processing Systems, volume 34, 2021.
Walter Rudin. Real and Complex Analysis, 3rd Ed. 0070542341.
McGraw-Hill, Inc., USA, 1987. ISBN
Antoine Salmona. Valentin de Bortoli, Julie Delon, and Agnès Desolneux. Can forh-forward generative models fit multimodal distributions? arXiv preprint arXiv:2206. 14476, 2022.
13

---

Under review as a conference paper at ICLR 2023
Sahil Sidheekh, Chris Barton Dock, Tushar Jain, Radu Balan, and Maneesh Kumar Singh. VQ-flows: Vector quantized local normalizing flows. The 38th Conference on Uncertainty Artificial Intelligence, 2022.
Harshinder Singh, Vladimir Hnizdo, and Eugene Demchuk. Probabilistic model for two dependent circular variables. Biometrika, 89(3):71--723, 2002.
Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.03288, 2021.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. ICLR, 2018
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pp. 681-688. Citeseer, 2011.
Zhisheng Xiao, Qing Yan, and Yali Amit. Generative latent flow. arXiv preprint arXiv: 10485, 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.
Tatsuya Yatagawa. torchmcube: marching cubes for pytorch, 2021. URL https://github
com/ Comn/tattsy/Eorcmcuuees
Emilio Zappa, Miranda Holmes-Cerfon and Jonathan Goodman. Monte Carlo on Manifolds:
Sampling Densities and Integrating Functions. Communications on Pure and Applied Mathematics, 71(12):2009-2647, 2018.
14

---

Under review as a conference paper at ICLR 2023
A FORMAL SETTING
Here we expand on the formal setting in which we seek to perform density estimation.
Geometry Let M be m-dimensional Riemannian submanifold of ambient space R" where m. < n. Formally this refers to the pair (M,g), where . C 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 to M. The metric g, which is typically implied, captures the curvature information we would like to associate with M.
manifold's Riemannian metric gives rise to a unique differential form known as the Riemannian volume form djr, which allows for the integration of continuous, compactly supported, real-valued
functions h over the Riemannian manifold (Lee) 2013):
Ju. dp.
(10)
Probability Let {a;} be observed samples drawn from P a probability measure supported on M. Since M has lower intrinsic dimension than R", is "infinitely thin." In other words, P" (M) = 1 while the (Lebesgue) volume of M is O, meaning no probability density integrated over the ambient space can be used represent P*, Formally stated, p not absolutely continuous with respect to the Lebesgue measure on R".
Instead, we require a new way to define the volumes of subsets of M. We can then formally define a probability density pd over M and integrate with respect to this volume to obtain probabilities. The volume form du on M is the answer; the probability of set S C M can be computed as follows:
P*(() = | p" du.
(11)
We note that the volume form du from differential geometry is not technically a measure in the sense of measure theory. This obstacle is minor: du can be extended to measure by common measure-theoretie tool known as the Thus we may identify du with a measure μ on M which produces volumes of Borel sets in M and which we call the Riemannian measure of M (Pennec 1999).
-ar--Maaroo-Kaaaaaa representation theorem? [Rudin 1987).
Formally, require P* be absolutely continuous with respect to nt, and thus write that p" is the Radon-Nikodym derivative of pa with respect to μc: p² dP- This is the ground-truth density du
function we seek to model this work.
B EXPERIMENT DETAILS
For all experiments, we use feedforward networks with SiLU activations (Hendrycks & Gimpel 2016; Ramachandran et al. 2017). All models are trained with the Adam optimizer with the default PyTorch parameters, except for the learning rate which is set described below (Kingma & Ba 2014). All EBMs, constrained EBMs, and pushforward EBMs are trained with a buffer size of 1000, from which we initialize each Langevin dynamics sample with 95% probability. We do not use spectral normalization for EBMs: we found it harmed the quality of density estimates. Initial noise for the constrained EBM is sampled uniformly from a box ambient space containing the ground truth manifold and then projected to the manifold by solving for arg min, || (Facise)) )]l? with L-BFGS using strong Wolfe line search. Equation 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.
To plot the constrained EBM densities, we estimate the normalizing constants using Monte Carlo. Since the learned MDF's always provide very good approximations of the true manifolds, we estimate
Inn the reference and sometimes in general, this theorem is called the Riesz representation theorem, which also refer to different theorem about Hilhert spaces.
15

---

Under review as a conference paper at ICLR 2023
each normalizing constant using uniform samples from the ground truth manifold for convenience. To plot the pushforward EBM densities, we estimate the normalizing constants latent space with Monte Carlo estimates based on uniform sampling within the elamped bounds. We then compute
pushforward densities with Equation 2
All experiments were performed
Intel Xeon Silver 41 14 CPU.
We provide quantitative results Table [ We estimate the distance of each training point to the manifold using optimization procedure, and report minimum, median, mean, and maximum distances over the training set. Nearest-poin estimates must be computed differently for constrained EBMs and pushforward EBMs,, and therefore estimates for each model are prone to different sources of error, so these metrics should be used only with caution as a basis of comparison. For the constrained EBM with MDF Fg. we compute the nearest point on the manifold to datapoint æ4 as
arg min || x=ll² + FF
where x has been initialized to x,, For the pushforward EBM with encoder-decoder pair (Jo: go), we compute the nearest point as fo(="), where
arg min [[fa(z) -x,113
where z has been initialized to go (x:). On every dataset, the implicit manifold captures the training set with substantially more accuracy.
Table 1: Statistics of estimated dataset distances to the manifold.
CONSTRAINED
PUSHFORWARD
EXPERIMENT
MIN
MEDIAN
MEAN
MAX
MIN
MEMAN
MRAN
MAX
EXAMPLE
10
0.08 10
0.006
0.226 10
10
0.21 10
0.181
DENSITY ENTIMATION
0.639 x 10
0.34 x 10
0.36 x 10
0.013
1.711 x 10
0.53 10
0.65 x 10
0.072
MANIFOLD ARITHMETIC
10
-
GEOSPATIAL DATA
0.158 x 10-3
0.06 x 10
0.08 x 10
0.004
0.780 x 10-5
0.13 10"2
0.16 x 10-2
0.014
AMINO
0.365 10
10
0.77 10-2
0.045
4,296 x 10
2.42 10
3.00 10
0.627
B. SYNTHETIC DATA
Motivating example (Figure 1 We sampled 1000 points from von Mises distribution unit circle centred at 0) with the mode located at (1, 0) and a concentration of 2.
The MDF for the constrained EBM consisted of hidden layers with 8 units per hidden layer. The MDF was trained for 200 epochs with a batch size of 100, a learning rate of 0.01, 7) 1, and a 1. Training took 6.56 seconds.
The energy function for the constrained EBM consisted of 2 hidden layers with 32 units per hidden layer. was trained for 40 epochs with batch size of 100, learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with coefficient of 0.1. Langevin dynamics at each training step were run for 10 steps with E 0.3, step size of 1, and energy gradients clamped to maximum values of 0.1 at each step. Training took 4 minutes, 5 seconds. In Figure 7 we evaluate the effect of the Langevin dynamics step count on training dynamics, where we vary the step size (and remove the training buffer, as this effectively increases the average step count). Fewer steps leads to more peaked mode because the estimated model distribution is overly smooth when estimating the right-hand side of Equation 5
The pushforward EBM's encoder and decoder each had hidden layers with units per hidden layer. They were jointly trained for 300 epochs with batch size of 100, learning rate of 0.001, and gradients clipped to a norm of 1. Training took 16.4 seconds.
The pushforward EBM's energy function had 3 hidden layers and 32 units per hidden layer. It was trained for 200 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 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 each step. Training took 5 minutes, 19 seconds.
16

---

Under review as a conference paper at ICLR 2023
Constrained EBM
Constreined EBM
Constrained EBM
Constrained EBM
Constreined EBM
Constrained EBM
Figure 7: Constrained EBM manifold learning and density estimation results on a von Mises distribution where Langevin dynamics during training has been run (with no 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.
Density estimation We sampled 1000 points from a balanced mixture of two von Mises distributions with concentration on circles unit radius. Respectively, they are centred at (-2, 0) and (2, 0) with modes at (-1,0) 0) and (1, 0) (or, at polar angles of 0 and with respect to the centre of each circle).
The MDF for the constrained EBM consisted of hidden layers with 8 units per hidden layer. The MDF was trained for 1000 epochs with a batch size of 100, a learning rate of 0.01, and a = 1. Training took 31.3 seconds.
The energy function for the constrained EBM consisted of 3 hidden layers with 32 units per hidden layer. was trained for 50 epochs with batch size of 100, learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with a coefficient of 0.3. Langevin dynamics at each training step were run for 10 steps with 0.4, step size of 1, and energy gradients clamped to maximum values of 0.1 at each step. Training took 37.8 seconds.
The (ambient) EBM consisted of 2 hidden layers with 32 units per hidden layer (we found that using only 2 hidden layers gave a smoother density). It was trained for 3 cycles of 200 epochs with step size of 10. We used a batch size of 100, learning rate of 0.01, gradients clipped a norm 1, and energy magnitudes regularized with coefficient of 0.5. Langevin dynamics at each training step were run for steps with e 0.1 and energy gradients clamped to maximum values of 0.03 at each step. Training took minutes, 2 seconds.
The pushforward EBM's encoder and decoder each had hidden layers with units per hidden layer. It was trained for 1000 epochs with batch size of 100, learning rate of 0.001, and gradients clipped to a norm 1. Training took 57.3 seconds. We also tried training the autoencoder using a variational autoencoder loss, but found that learn the manifold properly, the KL term had to be heavily downweighted near the point of nonexistence. In Figure g we show how manifold learning ability deteriorates the KL-weighting is increased.
The pushforward EBM's energy function had 3 hidden layers and 32 units per hidden layer. was trained for 300 epochs with batch size of 100, learning rate of 0.01, gradients clipped to norm of
17

---

Under review as a conference paper at ICLR 2023
Figure 8: Manifold learning and density estimation performance for different weightings 3 on the KL-divergence term of the VAE loss. From left right, top to bottom: 0.01, 3 0.03, / 0.05, B = 0.1, B = 0.2, B 0.4.
1, and energy magnitudes regularized with a coefficient of 0.1. Langevin dynamics at each training step were run for 60 steps with ε 1.0, step size of 10, and energy gradients clamped maximum values of 0.03 at each step. Training took 7 minutes, 5 seconds.
Manifold arithmetic We sampled 1000 points from balanced mixture of two projected normal
distributions the unit sphere. Each component was normal distribution with unit diagonal covariance centred at (1, 0, 0) and (-1, O, 0) respectively before being projected to the sphere.
The MDF for the constrained EBM consisted of 3 hidden layers with 8 units per hidden layer. The MDF was trained for 1500 epochs with a batch size of 100, a learning rate of 0.01, = 1, and a Training took 46.9 seconds.
The energy function for the constrained EBM consisted of 2 hidden layers with 32 units per hidden layer. It trained for rounds of 10 epochs each wherein Langevin dynamics 40, and 50 steps respectively. We used a batch size of 50, learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with coefficient of 1. Langevin dynamics each training step were run for 10 steps with s = 0.1, step size of e², and energy gradients clamped to maximum values of 0.03 each step. Training took 17 minutes, 51 seconds.
for 5, 10, 20,
B.2 NATURAL DATA
Geospatial data We modelled floods from the Dartmouth Flood Observatory's glooa active archive, which is available without charge for research and education purposes.
The MDF for the constrained EBM consisted of 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, 7) 1, and 2. Training took minute. 1 seconds.
The energy function for the constrained EBM consisted of 4 hidden layers with 32 units per hidden layer. It was trained for rounds of 10 epochs each wherein Langevin dynamics was run for 5, 10, 20, and 40 steps respectively. We used batch size of 100, a learning rate of 0.01, gradients clipped to norm of 1, and energy magnitudes regularized with coefficient of 1. Langevin dynamics at each training step was run with 0.1, a step size of e², and energy gradients clamped maximum values of 0.03 at each step. Training took 27 minutes, 43 seconds.
The pushforward EBM's encoder and decoder each had hidden layers with units per hidden layer. They were jointly trained for 500 epochs with a batch size of 100, learning rate of 0.001, and gradients clipped to a norm of 1. Training took minutes, 29 seconds.
The pushforward EBM's energy function had 4 hidden layers and 32 units per hidden layer. It was trained for 50 epochs with batch size of 100, a learning rate of 0.01, gradients clipped to a norm of 1, and energy magnitudes regularized with 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 each step. Training took 4 minutes, 42 seconds.
18

---

Under review as a conference paper at ICLR 2023
Figure 9: Constrained EBM manifold learning and density estimation results on the glycine angle data for different values of 1). the hyperparameter setting the boundary under which singular values will be penalized by the Jacobian regularization term. From left to right: = 0.3, 17 = 1, 17 = 2. 7 3, and 1) = 5.
hyperpaarizatoo
Amino acid modelling The MDF for the constrained EBM consisted of 2 hidden layers with 8 units per hidden layer. The MDF was trained for 500 epochs with batch size of 100, a learning rate of 0.01, 17 = 0.3, and 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 Training took 13.6 seconds.
The energy function for the constrained EBM consisted of hidden layers with 32 units per hidden layer. It was trained for 2 rounds of 10 epochs each wherein Langevin dynamics was run for 5 and 10 steps respectively, We used batch size of 100, learning rate of 0.01, gradients clipped to a norm of 1. and energy magnitudes regularized with a coefficient of 1. Langevin dynamics at each training step was 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 1 minute, 14 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 a batch size of 100, learning rate of 0.001, and gradients clipped to a norm of 1. Training took 33.5 seconds.
The pushforward EBM's s energy function had 3 hidden layers and 32 units per hidden layer. It was trained for 50 epochs with 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 each training step were run for 60 steps with & = 0.5, step size of 10, and energy gradients clamped to maximum values of 0.03 each step. Training took 54 seconds.
19


Reviews:
[{"review_id": "nNOKmvmc2hY", "reviewer": "Reviewer_aqkP", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes an approach to estimate a manifold using an implicit model. The authors use a constrained energy-based model to learn the data distribution within the manifold. A Langevin dynamics is used to sample from the manifold while training and at an inference time.", "summary_of_the_review": "- the idea of using implicit models for manifold representation is interesting, and I have not seen papers explicitly proposing such an idea,\n- the proposed algorithm is not novel; it consists of standard building blocks,\n- the authors did not conduct any ablation study to investigate the influence of different parts of the algorithm and did not perform sufficient experimental testing of the proposed approach for manifold learning.", "strength_and_weaknesses": "Strength:\n- a manifold learning and density estimation on the manifold is an important problem statement for many applications;\n- the proposed approach is simple and easy to follow;\n- the authors applied the approach to several;\n- the idea of modeling a manifold as an implicit function is interesting, although it is not novel. For example, in https://arxiv.org/pdf/2011.12026.pdf (proc. of CVPR, 2021), the authors proposed to represent a manifold of images based on an implicitly defined GAN model,\n- the proposed idea of manifold arithmetic sounds rather interesting, where thanks to implicit representation, we can model various operations on a given set of manifolds.\n\nWeakness:\n- the approach is a combination of several existing and relatively standard building blocks like implicit functions based on neural networks, an energy model, a Langevin dynamics; \nthe authors did not make any ablation study proving that the proposed selections of these building blocks and their hyperparameters are optimal,\n- the proposed approach for the manifold arithmetic requires additional testing. It is not clear for which applied problems it can be used or which existing solutions it can improve,\n- the manifolds considered in the experimental section are low-dimensional. Actually, they resemble some 3D shapes of visual objects. So the difference between the proposed approach and the approaches for shape modeling based on implicit functions, see, e.g., the paper https://arxiv.org/pdf/2008.06520.pdf (Learning Gradient Fields for Shape Generation) is not clear\n- the authors did not test their approach on any standard manifold learning problems, even those proposed in 200x, see, e.g., https://lvdmaaten.github.io/publications/papers/TR_Dimensionality_Reduction_Review_2009.pdf", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TRBj9H4v1y", "reviewer": "Reviewer_sPaT", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper points out some potential issues and limitations using “pushforward” models for density estimation. The main issue is that the general manifold of dimension m may not be effectively embedded in the latent space of the same dimension m. To resolve that challenge, the authors propose to construct an implicit model manifold, called MDF, so that the original data manifold lives on the null set of the MDF. The new implicit model covers a richer class of manifolds. Following this line, the authors further propose to use an energy based model to model the density over the manifold. Also the constrained Langevin Monte Carlois used to sample from the manifold.\n", "summary_of_the_review": "Overall my review is mixed: on one hand I do like the issues/challenges raised by the authors. Also the proposed solution seems very novel. On the other hand, I feel the challenges raised may not be real in applications since people are using dimensions much higher than m in the latent space. Also there are no real-world datasets used in the experimental sections. This makes me question whether the proposed novel but complicated solution is truly better than existing “pushforward” models. ", "strength_and_weaknesses": "Strength: The challenge the authors raised was real and interesting. I feel not many people have realized that. The proposed implicit model, as well the whole procedure of optimization, are novel. \nWeakness: The empirical part of the paper is weak. Only synthetic datasets were used. If the proposed model is so powerful, why not use it on real-world datasets?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "EpQK01QXui", "reviewer": "Reviewer_VPPU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method to learn implicit manifolds and sample from them using energy-based models. In particular, the paper first proposes a regularized method for learning submanifolds of any codimension and then shows how to fit the energy function and sample using a faster version of constrained HMC.", "summary_of_the_review": "Overall, I vote to reject the paper. My reasoning is that, while the paper does deftly utilize preexisting methods to construct manifold-learning flows, the adaptation is somewhat straightforward (e.g. regularize all singular values + use random numerical linear algebra techniques + backprop for CHMC). Importantly, the experiments don't showcase the primary benefits of the developed techniques, calling into question said developments. Either the manifold structure is known a priori or existing methods (e.g. a combination of Gropp et al. + standard CHMC) should suffice since the manifold is a surface in R^3. Of secondary concern, I believe that the current framing as a \"manifold learning flow\" paper is inaccurate, as the proposed method seems to not be addressing the same issues as the field. I would instead encourage the authors to perhaps motivate it from a different lens (perhaps EBM + adding in a manifold constraint).", "strength_and_weaknesses": "Strengths\n-----------\n* The paper brings together a lot of disparate literature to build better manifold-learning flow methods.\n* The results show the first ability to learn topologically nontrivial manifold structures from data.\n\nWeaknesses\n--------------\n* Outside of showing that the methods (implicits, ebms, CHMC) are of practical interest, the technical contributions feel marginal. In particular, for implicits, the geometric regularization term is similar to previous work (some mentioned in the paragraph) except as a generalization to matrices, and, for CHMC, the contribution seems like an application existing random numerical linear algebra techniques (hutchinson + CG) to the standard CHMC algorithm.\n* The experimental section feels very toy. In particular, section 4.1 just shows that the method works (and validates the manifold arithmetic component). Section 4.2 is more interesting, but notably doesn't test on data that is sampled from an unknown manifold (ie it is hard to motivate the method if we know the manifold is a sphere/torus). More importantly, all of these examples are *low dimensional*. Since the developed theory is built specifically to generalize to submanifolds of arbitrary codimension (i.e. the jacobian regularization method)  or for higher dimensional spaces (i.e. the novel CHMC method is built to scale to higher dimensions), I would expect an example showing that the developed theory matters, something that isn't true for the example surfaces in R^3.\n* The proposed methodology presents several important drawbacks. First, we can't evaluate log probabilities (without resorting to evaluating the energy function on all points of the manifold, which is computationally infeasible especially for dimension $> 3$). Second, the manifold and the probability density are two separate models (in particular the EBM is trained as an ambient density, not an intrinsic manifold density). Third, sampling relies on an expensive MCMC technique. Since the purpose of manifold-learning flows was to handle these issues, the current framing as a manifold-learning flow paper is somewhat strange/out-of-place.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject