Source-linked AI summary
Reparameterizing Distributions on Lie Groups
Luca Falorsi, Pim de Haan, Tim R. Davidson, Patrick Forré
TL;DR
Reparameterization tricks offer efficient gradient estimates, but lacked a general formulation for distributions on Lie groups. This paper develops ReLie by mapping reparameterizable algebra samples through the exponential map with analytically computed density corrections. The framework supports complex multimodal SO(3) distributions, and experiments show LI-Flows are necessary for uncertainty estimation with discrete or continuous symmetries.
Problem
Existing reparameterization methods did not provide a general formulation for distributions on non-trivial manifolds such as Lie groups, limiting their use for these spaces.
Method
ReLie constructs Lie-group densities by pushing reparameterizable samples from the Lie algebra through the exponential map and analytically accounting for volume change.
Results
The framework creates complex and multimodal SO(3) distributions, and experiments show LI-Flows are necessary for uncertainty estimation with discrete or continuous symmetries.
Takeaways & Limitations
ReLie extends reparameterizable density modeling to arbitrary Lie groups and supports uncertainty estimation in applications involving Lie-group symmetries.
Takeaways & Limitations
When a location element is predicted by a neural network, care is required to avoid homeomorphism conflicts.
Abstract
from arXiv · showhide
Reparameterizable densities are an important way to learn probability distributions in a deep learning setting. For many distributions it is possible to create low-variance gradient estimators by utilizing a `reparameterization trick'. Due to the absence of a general reparameterization trick, much research has recently been devoted to extend the number of reparameterizable distributional families. Unfortunately, this research has primarily focused on distributions defined in Euclidean space, ruling out the usage of one of the most influential class of spaces with non-trivial topologies: Lie groups. In this work we define a general framework to create reparameterizable densities on arbitrary Lie groups, and provide a detailed practitioners guide to further the ease of usage. We demonstrate how to create complex and multimodal distributions on the well known oriented group of 3D rotations, $\operatorname{SO}(3)$, using normalizing flows. Our experiments on applying such distributions in a Bayesian setting for pose estimation on objects with discrete and continuous symmetries, showcase their necessity in achieving realistic uncertainty estimates.
1 INTRODUCTION
The paper addresses the lack of a general reparameterization trick for distributions on non-trivial manifolds, especially Lie groups. It introduces a framework that constructs reparameterizable Lie-group densities and applies them to complex distributions and inference problems.
- Motivation: Lie groups describe differentiable symmetry transformations including rotations, translations, and scaling, making them relevant to robotics and other applications.
- Motivation: Reparameterization methods lacked a general formulation for distributions on non-trivial manifolds such as Lie groups.Existing extensions primarily targeted Euclidean or other settings without explicitly addressing Lie groups.
- Contribution: The paper introduces a general framework for defining reparameterizable densities on Lie groups, recovering the Gaussian construction of Kingma and Welling (2013) as a special case.
- Contribution: The framework pushes samples from the Lie algebra through the exponential map and analytically accounts for the resulting density change.
- Scope: The authors derive constructions for the n-Torus, SO(3), and SE(3), then create complex multimodal SO(3) densities with normalizing flows for supervised and unsupervised applications.
2 PRELIMINARIES
The preliminaries define Lie groups, Lie algebras, exponential maps, and the reparameterization trick. They motivate the paper by noting that existing reparameterization methods do not provide a general solution for Lie groups.
- Lie Groups and Lie Algebras: A Lie group is both a group and a smooth manifold, representing continuously traversable symmetries such as rotations and translations.
- Lie Groups and Lie Algebras: The Lie algebra is the tangent space at the identity, whose elements act as infinitesimal generators for constructing group elements.
- Lie Groups and Lie Algebras: The exponential map sends algebra elements to group elements, is surjective for compact connected groups such as SO(3), and is often non-injective.
- Reparameterization Trick: The reparameterization trick represents z ∼ q(z, θ) as z = T(ϵ; θ), with parameter-independent noise and a transformation differentiable in θ, enabling backpropagation through expectations.
- Reparameterization Trick: No general reparameterization scheme was available for arbitrary distributions on spaces with non-trivial topologies such as Lie groups.
3 REPARAMETERIZING DISTRIBUTIONS ON LIE GROUPS
ReLie reparameterizes Lie-group distributions by sampling in the Lie algebra, mapping through the exponential map, correcting the density, and applying group-specific location transformations. Its theory handles non-injective maps through partitioning and establishes a density for the pushforward measure.
- 3.1 Reparameterization Steps: ReLie samples a reparameterizable distribution on the Lie algebra, applies the exponential map, and left-multiplies by a group element to set the distribution’s location.
- 3.1 Reparameterization Steps: The exponential-map step requires correcting for volume change and possible many-to-one mappings because the map is generally neither volume-preserving nor injective.
- 3.2 Theory: Non-injective maps can still yield densities when their domain is partitioned into regions where the map is injective, while measure-zero sets are ignored.
- 3.2 Theory: The Lie-group generalization partitions the algebra into a measure-zero set and a countable union of open regions where the exponential map is a diffeomorphism.
- 3.2 Theory: Theorem 3.1 establishes that the exponential pushforward is absolutely continuous with respect to the group’s Haar measure and has an explicit density.
- 3.2 Theory: Left multiplication preserves Haar volume and leaves the distribution entropy invariant, while neural-network-predicted locations require care to avoid homeomorphism conflicts.
4 IMPLEMENTATION
The implementation derives reparameterization details for the n-Torus, SO(3), and SE(3), focusing on Jacobian-based volume changes under exponential-map pushforwards.
- Implementation details: Worst-case reparameterization complexity for an n-dimensional matrix Lie group is O(n^3), although closed forms can substantially reduce it.The implementation covers the n-Torus, SO(3), and SE(3).
- Computing J(x): The Jacobian determinant gives the volume change when pushing a density from the Lie algebra to the group.The Lie-algebra basis provides corresponding tangent-space bases throughout the group, enabling analytical computation for matrix Lie groups.
- Computing J(x): The general Jacobian term J(x) is computable for any Lie group through the theorem’s expression.For matrix Lie groups, derivatives of the matrix exponential provide an analytical route to J(x).
- The n-Torus, TN: For the n-Torus, the exponential map has a tractable preimage and its pushforward density preserves volume.The n-Torus is the product of n copies of S1 and is an abelian Lie group.
- The Special Orthogonal Group, SO(3): For SO(3), the exponential map uses the Rodrigues rotation formula, is surjective but non-injective, and requires accounting for multiple preimages in the pushforward density.The Lie algebra is represented by skew-symmetric 3 × 3 matrices identified with R3.
- The Special Euclidean Group, SE(3): SE(3) extends SO(3) with translations, using a Lie algebra formed by concatenating a skew-symmetric matrix and an R3 vector.Its exponential map and pushforward density are handled analogously to SO(3), including the relevant preimage structure.
5 RELATED WORK
Prior reparameterization methods broaden available variational families, but the paper identifies a missing treatment of volume changes and non-injective exponential maps on arbitrary Lie groups.
- Reparameterization methods: Existing approaches include surrogate distributions, implicit reparameterization, and generalized reparameterizations for expanding variational families.These approaches can be combined with different distributions over the Lie algebra.
- Relation to prior work: Unlike the proposed approach, prior techniques do not provide tools to account for volume changes when pushing densities from R^N to arbitrary Lie groups.The paper also highlights that the exponential map is generally non-injective, complicating change-of-variable computations.
- Normalizing flows: Normalizing flows are closely related because they increase density complexity through injective maps, but Lie-group transformations also change the underlying space and may use non-injective exponential maps.The paper combines normalizing flows with its framework to construct complex distributions on Lie groups.
- Distributions on Lie groups: The related-work discussion includes prior distributions on homogeneous spaces and Lie groups, including work focused on SE(3), but notes their restriction to neighborhoods of the identity.A cited example reparameterizes the von Mises–Fisher distribution on S^M, with S1 isomorphic to SO(2).
6 EXPERIMENTS
The experiments use Locally Invertible Flows to model complex SO(3) distributions in variational inference and supervised pose estimation, including multimodal symmetry-aware posteriors.
- Normalizing Flow: The flow uses an invertible neural network, a norm-applied tanh transformation, and a unit Gaussian base distribution to construct multimodal distributions.The hyperparameter r controls whether the exponential-map composition is invertible or locally invertible with a finite inverse set.
- Normalizing Flow: For sufficiently small r the flow is invertible but may not be surjective, whereas larger r yields a non-injective flow requiring backward traversal over each preimage.The image of r · tanh must remain within the regular region of the exponential map.
- 6.1 Variational Inference: The variational-inference experiment correctly infers symmetrical SO(3) group actions, matching Markov Chain Monte Carlo samples.The model uses a uniform latent-group prior and optimizes the ELBO with an LI-Flow approximate posterior.
- 6.2 Maximum Likelihood Estimation: LI-Flow successfully learns a multimodal conditional SO(3) distribution for supervised pose estimation with 2π/3 rotational symmetry.The model is trained by maximum likelihood on pose-corrupted objects and the resulting samples are shown in Fig. 5.2.
7 CONCLUSION
The paper presents ReLie as a general framework for reparameterizable distributions on Lie groups and extends it with LI-Flows for complex, multimodal uncertainty modeling. Experiments show these distributions are necessary for realistic uncertainty estimates in problems with discrete or continuous symmetries.
- 7 CONCLUSION: ReLie extends reparameterizable distributions to arbitrary Lie groups, while LI-Flows provide complex and multimodal distributions on SO(3).The framework is intended for probabilistic deep learning settings where Lie groups are important.
- 7 CONCLUSION: LI-Flows are empirically necessary for estimating uncertainty in problems containing discrete or continuous symmetries.
- 7 CONCLUSION: The SO(3) variational-inference experiment uses two ordered points and compares learned samples with a Metropolis-Hastings posterior.The setup targets approximate posterior estimation over rotations for a symmetrical object.
- 7 CONCLUSION: The learned SO(3) posterior almost perfectly matches the true posterior, whereas a centered Gaussian pushforward would concentrate mass around one rotation and struggle with the observed topology.
A.3 Supplementary Details on MLE Experiment
The MLE experiment models noisy pose observations with symmetry-induced multimodal posteriors on SO(3). LI-Flows successfully learn the resulting complicated conditional distributions.
- Setup: For an object invariant under a subgroup D, the pose posterior should have modes corresponding to the subgroup elements.The experiment generates noisy observations by applying a perturbed group action to a random pose.
- Results: The SO(3) object is made symmetric under the cyclic group of order 3, representing rotations of 2π/3 along one axis.
- Results: LI-Flows successfully learn the resulting complicated conditional distributions on SO(3).
C PREREQUISITES
The prerequisites introduce measures, densities, pushforwards, and standard measures needed to describe distributions on manifolds. They also illustrate how pushing a Normal distribution to SO(2) produces a density assembled from algebra pre-images.
- C PREREQUISITES: A density exists when a measure can be represented by a measurable nonnegative function with respect to a reference measure.
- C PREREQUISITES: For SO(2), the pushforward density at angle θ sums the algebra densities of all pre-images of θ under the exponential map.The circular visualization draws radius 1 + q(θ) at each angle.
- C PREREQUISITES: A pushforward measure is obtained by transporting a measure through a measurable map.
- C PREREQUISITES: On a Riemannian manifold, the standard measure is defined from the metric tensor and its determinant in local coordinates.
- C PREREQUISITES: The prerequisites frame probability distributions on Riemannian manifolds as measures absolutely continuous with respect to the standard measure.
D CHANGE OF VARIABLES
The change-of-variables treatment equips a Lie group with a left-invariant metric and derives the volume correction for the exponential map. For matrix Lie groups, this correction is represented by the left Jacobian and can be computed from matrix derivatives.
- D CHANGE OF VARIABLES: A scalar product on the Lie algebra induces a left-invariant Riemannian metric on the Lie group.The construction uses differentials of left translation to transport algebra bases to tangent spaces.
- D CHANGE OF VARIABLES: The volume change under the exponential map is the determinant of d(L_exp(x)^−1)_exp(x) ∘ d exp_x and is independent of the chosen basis for the endomorphism.
- D CHANGE OF VARIABLES: The induced measures are independent of the scalar product and basis up to scaling, but the pushforward density generally depends on those choices.
- D CHANGE OF VARIABLES: The formal change-of-variables result applies on open algebra regions where the exponential map is a diffeomorphism.The general derivation combines the exponential differential with the induced measures on the algebra and group.
- D CHANGE OF VARIABLES: For matrix Lie groups, the volume correction is represented by the left Jacobian, constructed by differentiating the matrix exponential and left action.Theorem D.6 gives its matrix representation relative to a Lie-algebra basis.
- D CHANGE OF VARIABLES: The singular set of the exponential map is characterized through the adjoint representation, and polynomial zero sets used in the analysis have Lebesgue measure zero unless the polynomial vanishes identically.
E.2 Main Theorem
The section establishes that pushing a density from the Lie algebra through the exponential map yields a density on the Lie group with respect to Haar measure, despite the map's non-injectivity. It also discusses approximation and computational costs for evaluating this density.
- E.2 Main Theorem: Theorem E.6 guarantees that the pushforward measure exp*(m) is absolutely continuous with respect to Haar measure ν and has an explicit density.The proof partitions the Lie algebra into regions where the exponential map is a diffeomorphism, applies change of variables, and combines the resulting contributions almost everywhere.
- E.2 Main Theorem: Non-injective exponential maps are handled by partitioning the Lie algebra into countably many diffeomorphic regions while discarding measure-zero sets.This permits the change-of-variables formula to be applied locally and the local densities to be summed.
- Complexity: The exponential map and its differential generally require O(n^3) computation for n-dimensional matrix Lie groups.The matrix exponential involves an infinite series, while differential computation through the left-Jacobian or adjoint representation is also generally cubic.
- Complexity: Group-specific analytic formulas can substantially reduce computational cost compared with the generic cubic procedures.Closed forms may be available for the exponential map and adjoint eigenvalues for particular Lie groups.
- Approximation: The density expression may contain an infinite summation, but truncation gives arbitrarily accurate approximations under suitable conditions.Exponentially decaying algebra densities often require only a handful of terms empirically, whereas compact support or support within the injectivity radius yields finite or single-term sums.
G CHOICE OF BASIS AND SCALAR PRODUCT IN THE LIE ALGEBRA
The section explains how representing the Lie algebra as R^n requires choosing a basis and scalar product, and how those choices can affect the induced geometry and pushed-forward distributions. It identifies invariant distributions and bi-invariant metrics as ways to reduce this dependence.
- Basis representation: A concrete basis identifies the abstract Lie algebra g with R^n, enabling Euclidean reparameterization before mapping samples back to g and then to G.One approach starts from an intrinsic measure on g and uses the basis only as a computational representation.
- Basis dependence: The representation-based construction depends on the chosen basis, whereas starting from an intrinsically specified measure on g is representation-independent.The directly specified R^n approach is easier in practice but can change when the basis changes.
- Basis dependence: The induced Riemannian geometry and pushed-forward distributions can change with the basis, especially for skewed distributions on R^n.Even orthonormal bases can produce different results unless the base distribution is invariant under orthonormal transformations.
- Scalar products and metrics: Choosing a scalar product determines orthonormal bases and the induced metric, while bi-invariance and the negative Killing form provide structured choices in suitable groups.Bi-invariant metrics exist for groups isomorphic to a compact group times R^n, and the negative Killing form induces one for compact semisimple groups.
- Invariant choices: For simple compact Lie groups, O(n)-invariant distributions pushed through the exponential map are independent of the chosen basis and bi-invariant metric up to scale.For an isotropic Normal p(x) = N(x|0, σ2 · I), the remaining choice effectively reduces to the variance σ2.
- Scalar products and metrics: Rescaling the scalar product can normalize the induced bi-invariant Haar measure when the total measure is computable.The construction uses an orthonormal basis with respect to the selected scalar product before applying the R^n-to-g isomorphism.