Score-based Pullback Riemannian Geometry

We introduce two key modifications to the Normalizing Flow framework that allow us to construct Riemannian geometry from unimodal densities \(p: \mathbb{R}^d \rightarrow \mathbb{R}^d \) of the form \[p(\mathbf{x}) \propto e^{-\psi(\varphi(\mathbf{x}))} \] where \(\psi: \mathbb{R}^d \rightarrow \mathbb{R}^d\) is a smooth strongly convex function and \(\phi: \mathbb{R}^d \rightarrow \mathbb{R}^d\) is a diffeomorphism.

1. Anisotropic Base Distribution: Instead of using a fixed isotropic Gaussian as the base distribution, we parameterize the diagonal elements of the covariance matrix \(\mathbf{\Sigma}_{\theta_1}\) of the base distribution. This allows the model to learn an anisotropic Gaussian, capturing variations along different directions of the data manifold.
2. Isometry Regularization: We encourage the flow \(\phi_{\theta_2}\) to be approximately isometric with respect to the Euclidean metric. This is achieved by adding a regularization term to the training objective, promoting the preservation of distances and angles during the transformation. This leads to more stable and interpretable manifold mappings.

In normalizing flows, we model the data density \( p_{\theta_1, \theta_2}(\mathbf{x}) \) using the change-of-variables formula: \[ p_{\theta_1, \theta_2}(\mathbf{x}) = p_{\mathcal{Z}}(\phi_{\theta_2}(\mathbf{x})) \left| \det \left( D_{\mathbf{x}} \phi_{\theta_2} \right) \right|, \] where:

1. \( \phi_{\theta_2}: \mathbb{R}^d \rightarrow \mathbb{R}^d \) is an invertible transformation (diffeomorphism) parameterized by \( \theta_2 \).
2. \( D_{\mathbf{x}} \phi_{\theta_2} \) is the Jacobian matrix of \( \phi_{\theta_2} \) at point \( \mathbf{x} \).
3. \( p_{\mathcal{Z}}(\mathbf{z}) \) is the base density evaluated at \( \mathbf{z} = \phi_{\theta_2}(\mathbf{x}) \).

Base Density

The base density \( p_{\mathcal{Z}}(\mathbf{z}) \) is chosen as the density of an anisotropic Gaussian distribution with a learnable covariance matrix \( \mathbf{\Sigma}_{\theta_1} \): \[ p_{\mathcal{Z}}(\mathbf{z}) = \frac{1}{(2\pi)^{d/2} \det(\mathbf{\Sigma}_{\theta_1})^{1/2}} \exp\left( -\frac{1}{2} \mathbf{z}^\top \mathbf{\Sigma}_{\theta_1}^{-1} \mathbf{z} \right). \]

Training Objective

Our training objective combines the negative log-likelihood of the data under the model with an isometry regularization term: \[ \mathcal{L}(\theta_1, \theta_2) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}} \left[ -\log p_{\theta_1, \theta_2}(\mathbf{x}) \right] + \lambda_{\text{iso}} \, \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}} \left[ \left\| (D_{\mathbf{x}} \phi_{\theta_2})^\top D_{\mathbf{x}} \phi_{\theta_2} - \mathbf{I}_d \right\|_F^2 \right], \] where:

1. Negative Log-Likelihood Term: Encourages the model to fit the data distribution by minimizing the negative log-likelihood. Expanding this term, we get: \[ \begin{align*} \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}} \left[ - \log p_{\theta_1, \theta_2}(\mathbf{x}) \right] &= \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}} \left[ - \log p_{\mathcal{Z}} \left( \phi_{\theta_2}(\mathbf{x}) \right) - \log \left| \det \left( D_{\mathbf{x}} \phi_{\theta_2} \right) \right| \right] \\ &= \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}} \left[ \frac{d}{2} \log (2\pi) + \frac{1}{2} \log \det \left( \mathbf{\Sigma}_{\theta_1} \right) + \frac{1}{2} \phi_{\theta_2}(\mathbf{x})^\top \mathbf{\Sigma}_{\theta_1}^{-1} \phi_{\theta_2}(\mathbf{x}) - \log \left| \det \left( D_{\mathbf{x}} \phi_{\theta_2} \right) \right| \right]. \end{align*} \]
2. Isometry Regularization Term: Penalizes deviations from isometry, encouraging the transformation \( \phi_{\theta_2} \) to be approximately length-preserving. Here \( \lambda_{\text{iso}} > 0 \) is the regularization weight, \( \|\cdot\|_F \) denotes the Frobenius norm, \( ( D_{\mathbf{x}} \phi_{\theta_2} )^\top D_{\mathbf{x}} \phi_{\theta_2} \) is the Gram matrix of the Jacobian and \( \mathbf{I}_d \) is the \( d \times d \) identity matrix.

By training the normalizing flow with an anisotropic base distribution and isometry regularization, we obtain:

1. A learned diffeomorphism \( \phi_{\theta_2}: \mathbb{R}^d \rightarrow \mathbb{R}^d \), representing the transformation learned by the flow.
2. A learned strongly convex function \( \psi: \mathbb{R}^d \rightarrow \mathbb{R} \), associated with the anisotropic Gaussian base distribution.

Definition of \( \psi \): The strongly convex function \( \psi \) corresponds to the negative log-density of the anisotropic Gaussian base distribution. We parameterize only the diagonal elements of the covariance matrix \( \mathbf{\Sigma}_{\theta_1} \), with all off-diagonal elements being zero. Specifically, \( \psi \) is defined as: \[ \psi( \mathbf{z} ) = \frac{1}{2} \mathbf{z}^\top \mathbf{\Sigma}_{\theta_1}^{-1} \mathbf{z}, \] where \( \mathbf{\Sigma}_{\theta_1} = \operatorname{diag}( \sigma_1^2, \sigma_2^2, \ldots, \sigma_d^2 ) \) is a diagonal matrix with positive entries \( \sigma_i^2 \).

Constructing the pullback Riemannian Metric:

We define a pullback metric: on the data space \( \mathbb{R}^d \) using the composition of the gradient of \( \psi \) with the learned diffeomorphism \( \phi_{\theta_2} \): \[ g_{\mathbf{x}}( \mathbf{v}, \mathbf{w} ) = \left( D_{\mathbf{x}} \nabla \psi \circ \phi_{\theta_2}[ \mathbf{v} ],\ D_{\mathbf{x}} \nabla \psi \circ \phi_{\theta_2}[ \mathbf{w} ] \right)_2, \] where:

1. \( D_{\mathbf{x}} \nabla \psi \circ \phi_{\theta_2} \) is the Jacobian of \( \nabla \psi \circ \phi_{\theta_2} \) at \( \mathbf{x} \).
2. \( \mathbf{v}, \mathbf{w} \in T_{\mathbf{x}} \mathbb{R}^d \) are tangent vectors at \( \mathbf{x} \).
3. \( (\cdot, \cdot )_2 \) denotes the standard Euclidean inner product.

The defined pullback metric is related to pullback metric of the score function of the distribution. It enables the computation of geodesics and other manifold mappings that pass through regions of high data density in closed form.

Moreover, the learned covariance matrix \( \mathbf{\Sigma}_{\theta_1} \) from \( \psi \) enables the construction of a Riemannian Auto-encoder (RAE). The latent dimensions associated with significant variance (i.e., large \( \sigma_i^2 \)) are used for the construction of a Riemannian Auto-encoder that captures both the geometry and the dimensionality of the data manifold.

For more details on the connection of our defined metric to the pullback metric of the score function and the construction of the Riemannian Auto-encoder, please refer to sections 3-5 of the paper.

For the Hemisphere(5,20) dataset, the model identified five non-vanishing variances, perfectly capturing the intrinsic dimension of the manifold. This is reflected in the blue reconstruction curve, where the first five latent dimensions, corresponding to the largest variances, are sufficient to reduce the reconstruction error almost to zero. In contrast, the green curve illustrates that the remaining ambient dimensions do not encode useful information about the manifold. The red curve demonstrates improvement only when an important latent dimension is included.

For the more challenging Sinusoid(5,20) dataset, our method still performs very well, though not as perfectly as for the Hemisphere dataset. The first six most important latent dimensions explain approximately \(97\%\) of the variance, increasing to over \(99\%\) with the seventh dimension. This is reflected in the blue reconstruction curve, where the first six latent dimensions reduce the reconstruction error to near zero, and the addition of the seventh dimension brings the error effectively to zero. The slight discrepancy between our results and the ground truth likely arises from increased optimization difficulty, as the normalizing flow must learn a more intricate distribution while maintaining approximate isometry. We believe that with deeper architectures and more careful tuning of the optimization loss, the model will converge to the correct intrinsic dimensionality of five. Currently, it predicts six dimensions at a variance threshold of \(\epsilon = 0.05\) and seven at \(\epsilon = 0.01\), slightly overestimating due to the manifold's complexity.