Manifold Density Estimation via Generalized Dequantization
Abstract
Density estimation is an important technique for characterizing distributions given observations.
Much existing research on density estimation has focused on cases wherein the data lies in a Euclidean space.
However, some kinds of data are not well-modeled by supposing that their underlying geometry is Euclidean.
Instead, it can be useful to model such data as lying on a manifold with some known structure.
For instance, some kinds of data may be known to lie on the surface of a sphere.
We study the problem of estimating densities on manifolds.
We propose a method, inspired by the literature on “dequantization,” which we interpret through the lens of a coordinate transformation of an ambient Euclidean space and a smooth manifold of interest.
Using methods from normalizing ﬂows, we apply this method to the dequantization of smooth manifold structures in order to model densities on the sphere, tori, and the orthogonal group.
1. Introduction
This material appears in greater detail in our long-form version on arXiv; please see Brofos et al. (2021) for full details.
Certain kinds of data are not well-modeled under the assumption of an underlying Euclidean geometry.
Examples include data with a fundamental directional structure, data that represents transformations of Euclidean space (such as rotations and reﬂections), data that has periodicity constraints or data that represents hierarchical structures.
In such cases, it is important to explicitly model the data as lying on a manifold with a suitable structure; for instance
Third workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models (ICML 2021).
Copyright s a sphere would be appropriate for directional data, the orthogonal group for rotations and reﬂections, and the torus captures structural properties of periodicity.
The contribution of this work is to express density estimation on manifolds as a form of dequantization.
Given a probability density in an ambient Euclidean space, one can obtain the density on the manifold by performing a manifold change-ofvariables in which the manifold structure appears and then projecting out any auxiliary structures.
This marginalization can be viewed as analogous to “quantization” where, for instance, continuous values are discarded and only rounded integer values remain.
In this view the auxiliary structure deﬁnes how the manifold could be “dequantized” into the ambient Euclidean space.
By marginalizing along these auxiliary dimensions, one obtains the marginal distribution on the manifold.
In practice, however, one has only the manifold-constrained observations from an unknown distribution on the manifold.
A second contribution of this work is to formulate the density estimation as a learning problem on the ambient Euclidean space.
We show how to invoke the manifold change-of-variables, and then perform the marginalization along the auxiliary dimensions to ob-
Spaces
[Equation]
Quantization
[Equation]
Dequantization
Densities
[Equation]
DequantizationDensity
[Equation]
[Table]
Spherical
[Table]
( y, r ) (cid:55)→ ry r m −
T m
R + × · · · × R +
R 2 m
[Table]
polar coordinates ( y i , r i )
(cid:55)→ r i y i (cid:81) mi =1 r
i
Stiefel( m, n )
Tri + ( n )
R m × n
QR decomposition ( Y , R )
(cid:55)→ YR
R m − 1· · · R m − n nn
Stiefel( m, n )
PD( n )
R m × n
Matrix polar decomposition ( Y , R ) (cid:55)→ YR
Automatic differentiation tain effective estimates of the density on the manifold.
An advantage of our dequantization approach is that it allows one to utilize any expressive density directly on the ambient
Euclidean space (e.g., RealNVP (Dinh et al., 2017), neural
ODEs (Chen et al., 2018; Grathwohl et al., 2018) or any other normalizing ﬂow (Kobyzev et al., 2020)); the dequantization approach does not require a practitioner to construct densities intrinsically on the manifold.
2. Theory
Theorem 1.
Let Y and Z be smooth manifolds embedded in R n and R p , respectively.
Let G : R m → Y × Z be a smooth, invertible transformation.
Let π R
m be a density on R m .
Under the change-of-variables G , the corresponding density on Y × Z is given by,
[Equation]
where x = G − 1 ( y, z ) .
Even when G is not an invertible mapping, it may be possible to compute the change-of-variables when G is invertible on partitions of R m .
Corollary 1.
Let O 1 , . . .
, O l be a partition of R m .
Let
G : R m → Y ×Z be a function and suppose that there exist smooth and invertible functions G i : O i
→ Y ×Z such that
G i =
G |O i for i = 1 , . . .
, l .
Then, if x ∼ π R m , the density of ( y, z ) =
G ( x ) is given by
[Equation]
G i ( x i )) .
(2) where x i =
G − 1 i ( y, z ) .
How does theorem 1 relate to the dequantization of smooth manifolds?
2.1. Dequantization
Manifolds of interest (such the sphere, the torus, or the orthogonal group) can be introduced as elements of a new coordinate system for an ambient Euclidean space.
In each case, the manifold appears with an auxiliary manifold which may not be of immediate interest.
Namely, (i)
The sphere appears with set of positive real numbers when deﬁning a coordinate system for R m \ { 0 } ∼ =
S m − 1 × R + ; (ii)
The torus appears the product manifold of m copies of the positive real numbers when deﬁning a coordinate system for R 2 m \ { 0 } ∼ = T m × R + × . . .
× R + ; (iii) the Stiefel manifold appears with the set of lower-triangular matrices with positive diagonal entries when deﬁning a coordinate system of full-rank matrices: FR( n, p ) ∼ =
Stiefel( n, p ) ×
Tri + ( p ) .
We would like to marginalize out these “nuisance manifolds” so as to obtain distributions on the manifold of primary interest.
A convenient means to achieve this is to introduce an importance sampling distribution over the nuisance manifold.
Formally, we have the following result, which is an immediate consequence of theorem 1.
Corollary 2.
Let Y , Z , G , and π Y×Z be as deﬁned in theorem 1.
Let
˜ π
Z be a non-vanishing density on Z .
To obtain the marginal density on Y , it sufﬁces to compute,
[Equation]
where x = G − 1 ( y, z ) .
3. Discussion
We investigate the problem of density estimation given observations on a manifold using the dequantization procedure described in section 2.
Problem.
Let Y be a manifold embedded in R n and let π Y be a density on Y .
Given observations of π Y , construct an estimate ˆ π Y of the density π Y .
We apply eq.
(3) in order to obtain the density estimate on Y .
Generating samples from π Y may be achieved by ﬁrst sampling x ∼ π X , applying the transformation G ( x ) =
( y, z ) , and taking y as a sample from the approximated distribution ˆ π Y .
3.1. Densities on R m
As R m is a Euclidean space, we have available a wealth of possible mechanisms to produce ﬂexible densities in the ambient space.
One popular choice is RealNVP (Dinh et al., 2017).
An alternative is neural ODEs which parameterizes a vector ﬁeld in the Euclidean space; the change in probability density under the vector ﬁeld ﬂow is obtained by integrating the instantaneous change-of-variables formula (Chen et al., 2018; Grathwohl et al., 2018).
3.2. Objective Functions
We consider two possible objective functions for density estimation.
The ﬁrst is the evidence lower bound of the observations { y 1 , . . .
, y n obs } :
[Equation]
This follows as a consequence of Jensen’s inequality applied to eq.
Experimental results using this objective function are denoted with the sufﬁx (ELBO).
The second is the loglikelihood computed via importance sampling:
[Equation]
Because the calculation of eq.
(5) requires an importance sampling estimate, experimental results using this objective function are denoted with the sufﬁx (I.S.).
4. Experimental Results
To demonstrate the effectiveness of the approach, we now show experimental results for density estimation on three different manifolds: the sphere, the torus and the orthogonal group.
In our comparison against competing algorithms, we ensure that each method has a comparable number of learnable parameters.
Our evaluation metrics are designed to test the ﬁdelity of the density estimate to the target distribution.
In all of our examples we use rejection sampling in order to draw samples from the target distribution.
4.1. Sphere and Hypersphere
Our ﬁrst experimental results concern the sphere S 2 where we consider a multimodal distribution with four modes.
We consider performing density estimation using the ELBO (eq.
and log-likelihood objective functions (eq.
; we construct densities in the ambient space using RealNVP and neural ODEs.
As baselines we consider the M ¨ obius transform approach described in (Rezende et al., 2020), which is a specialized normalizing ﬂow method for tori and spheres, and the neural manifold ODE applied to the sphere as described in (Lou et al., 2020).
We give a comparison of performance metrics between these methods in table 2.
In these experiments, we ﬁnd that parameterizing a neural ODE model in the ambient space gave the better
KL-divergence and effective sample size (ESS) metrics than
RealNVP when our dequantization approach is used.
We ﬁnd that our dequantization algorithm minimizing either eq.
(4) or eq.
(5) achieves similar performance in the ﬁrst and second moment metrics.
However, when using eq.
(5), slightly lower KL-divergence metrics are achievable as well as slightly larger effective sample sizes.
In either case, dequantization tends to outperform the M¨obius transform on this multimodal density on S 2 .
The manifold ODE method is outperformed by the ODE dequantization algorithms with both eq.
(4) and eq.
We next consider a multimodal density S 3 ∼ = SU(3) (the special unitary group).
As before, we compare dequantization to M¨obius ﬂow transformations and manifold neural ODEs and present results in table 3.
Similar to the case of the
[Table] [Table]
[Table]
multimodal density on S 2 , we ﬁnd that dequantization with an ambient neural ODE model is most effective, with ELBO maximization giving the smallest KL-divergence metrics.
All dequantization algorithms out-performed the M¨obius transformation on the sphere but only dequantization with an ambient ODE and ELBO minimization outperformed the manifold neural ODE method.
4.2. Orthogonal Group
The previous two examples focused on manifolds composed of spheres and circles.
We now examine density estimation on the orthogonal group, where we consider inference in a probabilistic variant of the orthogonal Procrustes problem; we seek to sample orthogonal transformations that transport one point cloud towards another in terms of squared distance.
We consider parameterizing a distribution in the ambient Euclidean space using RealNVP in these experiments.
Results are presented in table 4.
We observe that optimizing the ELBO objective function (eq.
(4)) tended to produce better density estimates than the log-likelihood (eq.
Nevertheless, we ﬁnd that either dequantization algorithm is highly effective at matching the target density.
We may also leverage corollary 1 so as to apply our method to the “dequantization” of SO( n ) .
As an example, we consider a multimodal density on SO(3) .
Results of applying our method to sampling from this distribution are also shown in table 4.
In this example we ﬁnd that minimizing the negative log-likelihood using importance sampling tended to produce the best approximation of the ﬁrst- and second-moments of the distribution, in addition to smaller
KL-divergence metrics.