Injective flows for star-like manifolds

TL;DR

This work tackles density estimation when targets lie on lower-dimensional manifolds by introducing injective flows tailored for star-like manifolds. The key advance is an exact, efficient computation of the Jacobian determinant with complexity $O(d^2)$, matching the cost of standard normalizing flows, enabling reliable density evaluation in variational inference without samples. The authors demonstrate two impactful applications: an Objective Bayesian penalized-likelihood framework that places priors on level-set manifolds, and a variational-inference approach for probabilistic mixing models on the simplex. Empirical results show substantial speedups over brute-force determinants, improved density reconstruction, and effective handling of multi-modal and sparse posterior structures, highlighting the practical significance for Bayesian inference on manifolds. Overall, the paper broadens the applicability of density estimation on manifolds by combining exact geometry-aware Jacobians with flexible injective-flow architectures.

Abstract

Normalizing Flows (NFs) are powerful and efficient models for density estimation. When modeling densities on manifolds, NFs can be generalized to injective flows but the Jacobian determinant becomes computationally prohibitive. Current approaches either consider bounds on the log-likelihood or rely on some approximations of the Jacobian determinant. In contrast, we propose injective flows for star-like manifolds and show that for such manifolds we can compute the Jacobian determinant exactly and efficiently, with the same cost as NFs. This aspect is particularly relevant for variational inference settings, where no samples are available and only some unnormalized target is known. Among many, we showcase the relevance of modeling densities on star-like manifolds in two settings. Firstly, we introduce a novel Objective Bayesian approach for penalized likelihood models by interpreting level-sets of the penalty as star-like manifolds. Secondly, we consider probabilistic mixing models and introduce a general method for variational inference by defining the posterior of mixture weights on the probability simplex.

Figures (12)

• Figure 1: 1D star-like manifold parametrized in spherical coordinates.
• Figure 2: Architecture of the proposed injective flows for star-like manifolds, see Theorem \ref{['thm:spherical_flow']}.
• Figure 3: (a) Runtime comparison of Jacobian determinant computation with proposed approach in Eq \ref{['eq:jacobian_determinant_starlike_manifolds']} and with the brute-force computation Eq. \ref{['eq:manifold_flow_change_variable']}. (b) MSE of learned log-density for an injective flow trained to learn a uniform distribution on a $l_p$-(pseudo) norm ball with $p=0.5$. We compare the same model trained with our exact Jacobian and with the Hutchinson trace estimator with $n=1,10,50,100$ number of Gaussian samples. Note that the number of samples is upper bounded by the dimensionality of the problem and that at evaluation the exact Jacobian is used.
• Figure 4: Learned ("ours") and ground truth ("gt") density for proposed injective flow trained via reverse KL divergence (no samples). Only the (unnormalized) target is given: (a) von Mises-Fisher ($\kappa = 5$), (b) mixture of 50 von Mises-Fisher ($\kappa = 50$) and (c) sinusoidal density on deformed sphere
• Figure 5: $95\%$ posterior C.I. for 3 subjective priors with the same MAP. The choice of prior affects the posterior.

Theorems & Definitions (12)

• Definition 1
• Theorem 1
• Definition 2
• Remark 1
• proof
• Theorem 2
• Lemma 1
• proof
• Lemma 2
• proof