Wrapped $\beta$-Gaussians with compact support for exact probabilistic modeling on manifolds
Authors: Sergey Troshin
Vlad Niculae
Venue: Accepted by TMLR ()

Abstract: We introduce wrapped $\beta$-Gaussians, a family of wrapped distributions on Riemannian manifolds, supporting efficient reparametrized sampling, as well as exact density estimation, effortlessly supporting high dimensions and anisotropic scale parameters. We extend Fenchel-Young losses for geometry-aware learning with wrapped $\beta$-Gaussians, and demonstrate the efficacy of our proposed family in a suite of experiments on hypersphere and rotation manifolds: data fitting, hierarchy encoding, generative modeling with variational autoencoders, and multilingual word embedding alignment.

Published in Transactions on Machine Learning Research (12/2023)
Wrapped B B-Gaussians with compact support for exact probabilistic modeling on manifolds
Sergey Troshin
8. troshinnuva.nt nt
Language Technology Lab, Informatics Institute, University of Amsterdam
Vlad Niculac
v.niculaeduva.r
Language Technology Lab, Informatics Institute, University of Amsteriam
Reviewed on OpenReview: https: // openreview. net/ formm id- KrequDpist
Abstract
We introduce wrapped B-Gaussians, a family of wrapped distributions on tractable Rie- mannian manifolds, supporting efficient reparametrized sampling, as well as exact density assessment, effortlessly supporting high dimensions and anisotropic scale parameters. We extend Fenchel- Young losses for geometry-avaa learning with wrapped 3-Gaussians, and demonstrate the efficacy of our proposed family in a suite of experiments on hypersphere and rotation manifolds: data fitting, hierarchy encoding, generative modeling with variational autoencoders, and multilingual word embedding alignment.
The source code available at https: com wot.uuu..commttta////aeea
1 Introduction
The Euclidean space is not always sufficiently expressive for modeling the rich data encountered in applied machine learning. Sometimes data naturally lives in another geometry. Examples include geological processes (Curray, 1956) or word embeddings (Meng et al., 2019) where data is represented spheres. Hierarchical data such as knowledge graphs or taxonomies may be better represented in hyperbolic spaces (Nickel & Kiela, 2017), while the manifolds of rotations and of rigid motions are used in applications like robot pose or motion estimation (Rosen et al., 2019) or as domains for parametrizing other models (Artetxe et al., 2016).
Probabilistic modeling on Riemannian manifolds a research area attracting increasing attention. Geometry introduces substantial computational challenges compared the Euclidean case. For instance, whereas the multivariate Gaussian distribution allows efficient calculations over R" for any symmetric positive semidefinite covariance matrix, the Gaussian has no direct equivalence over a Riemannian manifold that is as elficient, exact, and expressive. The possible extensions, further described in $2.2, either require numerical approximations, numerical integration, rejection sampling, sacrifice expressivity of the scale parameter, or abandon geometry in favor of working in a restriction of the ambient Euclidean space.
A promising direction is that of wrapped distributions (Chevallier & Guigui, 2020), defined implicitly in terms of distribution in tangent space. Wrapped distributions support efficient sampling by design, but their exact density is in general intractable. Unless the manifold has non-positive curvature everywhere, like hyperbolic geometry (Nagano et al., 2019), in general, any point on the manifold can be reached by the wrapping of infinitely many tangent points (Figure 1, left). In order to have exact expressions, it is, therefore, necessary to use a compactly-supporte yet expressive distribution in the tangent space, yet no such constructions have been proposed so far.
We fill this gap with 3-Gaussians (Martins al., 2022), a recent family generalizing Gaussians, which provide controllable compact support, tractable sampling, and natural loss functions. Our main contributions

---

Published in Transactions on Machine Learning Research (12/2023)
63
02
-20 19.10 19.1 -20
22 o0
22 13.18 13.18
-05
00
-05
1.0
-20
-20
Figure 1: Wrapped G-Gaussian on S with (a) 3 1 (wrapped Gaussian), (b) = 0.5 and (c) 3 0. The 8- G-Guussian (orange), defined over the tangent plane at μ (-1,0), wrapped over the circle x² + yy obtain the wrapped B-Gaussian distribution (blue). In (b,c), wrapping injective on the support.
the wrapped 8-Gaussian, a distribution family on Riemannian manifolds, providing exact densily assessment and reparametrized sampling even high dimensions and with anisotropic scale;
an efficient embedded parametrization of scale matrices for the case of hyperspheres;
geomelry-aware Fenchel-Young losses for parameter fitting;
experimental demonstration of the usefulness of wrapped 3-Gaussians in syuthetic and real modeling lasks on hyperspherical S and special orthogonal S0(n) manifolds ($4), including hierarchy modeling, variational autoencoders, and multilingual word embedding alignment tasks, in high dimensions.
The wrapped ß-Ganssian distribution is applicable whenever standard toolbox is available and tractable for the manifold: exponential and logarithmic mapping, the determinant of their Jacobian, and the injectivity radius; a fairly broad (but not universal) class of manifolds. Particularly, it includes the class of symmetric Riemannian manifolds (Chevallier et al., 2022), including hyperspheres, hyperbolic spaces, and matrix manifolds such as SO(n) and SEn)).
Wrapped B-Gaussians tend toward wrapped Gaussians in the limit of 3 + 1, and the proposed losses tend toward log-probability and Kullback-Leibbe divergences, but for any 3 < 1 the support of the distribution is compact, allowing for efficient and exact computation, even on manifolds with finite injectivity radius.
2 Background
2.1
Differential geometry
A smooth manifold M is a topological space that locally resembles Euclidean space. Common examples include the sphere, the torus, the real projective plane, as well as various matrix manifolds such as the special orthogonal group. The tangent spaces allow to generalize calculus to manifolds: at each point x E M, the tangent space at x, denoted T,M, manifold. A Riemannian manifold is a manifold in which every langent space TaM is equipped with an inner product Sx, and thus also with an induced norm ||=||/-
real vector space isomorphic to Rd where d is the dimension of the
Every tangent vector uniquely defines a geodesic. The exponential mapping Exp,, TaM + M maps a langent vector v to point on the manifold by traveling along the corresponding geodesic for unit of time (thus ||||| controls the velocity). The exponential mapping is not injective in general. For instance, on the circle $1, two tangent vectors e and 2n + map to the same point: unit-time journey along the geodesic defined by 2n v means travelling in the same direction as except faster by 2n, and so will wrap around and end up at the same destination. However, Exp, is injective on small enough ball B r) (v E T.M ||o]| r).
2

---

Published in Transactions on Machine Learning Research (12/2023)
The injectivity radius of M at is the radius of the largest such ball:
inj, M = sup > 0: Exp, is injective on B)),r))
The canonical mapping in the opposite direction is known as the logarithmic map. If we denote by Exp= the set-valued inverse (preimage), then Log, : : + TMM defined as?
(vg arg min f : v Exp, (v)).
When the minimizer above is not unique, (e. for the antipodal point on circle), Log, is undefined. With these definitions, on the set B(0, inj, M) E M,M, Exp, and Log, form a pair of inverse diffeomorphisms.
Finally, we discuss the special case of embedded manifolds, which covers all manifolds used in our experiments. In many cases of interest a manifold can be represented as a subset of R". In this case, the tangent space can also be taken as a subset of R", specifically T.MM = {u E R" : (u,x) = 0}, and the inner product and Riemannian metric are inherited from R
2.2 Statistics on manifolds
Standard probability and statistics on Euclidean vector spaces enjoy the benefit of linearity which are no longer available when generalizing to manifolds. In particular, on manifolds, there is no unique generalization of the Gaussian distribution. We brielly present the few alternatives and their tradeoffs.
The intrinsic approach (Pennec, 2006). A natural and principled way define probability distributions manifolds is in terms of the intrinsic geodesic distances d(z, v). The intrinsic standard Gaussian centered at μ) takes the form p(x) x (x,(-d(xx,)/2)) with its extension to full scale matrices:
p(a) x exp (-2 Log, (x) - Lo),
Log
(1)
This formulation can be derived from a maximum entropy principle the manifold, but it generally does not lead lo computationally friendly reparametrized sampling, and while unnormalized densities are tractable, the normalization constant might not be (see Hauberg (2018) for the instantiation on S", where the normalizing constant is only available in the isotropic case ).
The embedded approach. For manifolds embedded in a vector space IR, the embedded approach relies on defining distribution in the ambient space and conditioning (renormalizing) on the manifold. This approach sometimes computationally friendlier: in fact, conditioned S", p(x) x exp(-x|læ* - -ull)) yields Mises-Fisher (Mardia & Jupp, 1999, $9.3.2). Still,
the celebrated Langevin distribution, also known more complex anisotropic constructions are generally intractable. Moreover, since distances in ambient space are not necessarily related to distances on the manifold, this approach may misrepresent the geometry.
The wrapped approach (Mardia & Jupp, 1999, $3.5.7). The wrapped distribution has a long history in directional statistics (Mardia & Jupp, 1999) with most early instances targeting the circle (Stephens, 1963). Since the tangent. space of a Riemannian manifold is isomorphic to a Fuclidean space, a natural idea is to pick a location E M and define a zero-mean distribution in T,M. Any mapping from T,M then induces implicit distribution over M, but in particular, the Exp mapping is appealing choice due to its relationship to geodesic distances ||v|| = d (1,, Exp,, (v)). Sampling from wrapped distributions amounts to sampling from the tangent distributions and applying the Exp mapping. Assessing probabilities is possible for manifolds where Exp is invertible and it's Jacobian tractable, e. on the hyperbolic space, where Nagano et al. (2019) successfully used wrapped normal distributions. However, for many important manifolds, Exp is not invertible due to curvature, and many tangent points can map to the same manifold location, leading to
We urge caution with the different types of logarithms and exponentials used throughout the paper: Logg, Exp, refer the mappings between manifold and tangent space. We use exp and log for the basc-c natural exponential and logarithm on real and complex numbers, and by extension also the matrix exponential and logarithm. The notation logs, expy, introduced in 82.3, are generalizations of log and exp from nonextensive statistical mechanics, and not base-8 operations.
3

---

Published in Transactions on Machine Learning Research (12/2023)
035
0 = 0.9
A 0.5
mm
0.0 3
815
0000
exp(c)
Figure Left: the Tsallis expa recovers the usual exp as 3 + Right: isotropic 3-Gaussian density.
generally intractable infinite summations. Constraining or renormalizing an arbitrary tangent distribution over the injectivity domain of Exp would solve this problem, at the cost complicating both sampling and normalizing constant calculation, outside of simple isotropic cases. Instead, we show how careful choice of compactly-suppooteed B-Gaussian distributions in tangent space can keep the support within the injectivity domain with straightforward calculations, leading to computationally friendly wrapped distribution that enjoys all the benefits mentioned in this section: flexible, anisotropic, geomelry-aware distribution with efficient reparametrized sampling and tractable exact probability assessment.
2.3 Sparse continuous distributions
To ensure that the injectivity constraint is satisfied, the tangent distribution should have bounded, controllable support. Martins et al. (2022) study such distributions and propose the 3-Gaussian family, a generalization of the Gaussian (normal) distribution. We summarize their construction here and refer readers to Martins et al. (2022, $6) for more details.
Tsallis generalized entropies. Whereas the Gaussian is deeply connected to the Shannon-Boltzman Boltzmann- Gibbs differential entropy Hy) Epl- log p()), 3-Gaussians are generated by the Tsallis entropies (Tsallis, 1988), from nonextensive statistical mechanics. Consider, for 3 * 1, and denoting [++ = max(0,
etp = [1 (1 3
log = -12-9 -112--1111 1)/(1- - B).
(2)
Notice that lim (t) = cxp(t) and lim 11 (t) = log(t), and so we extend both functions by continuity at B = 1 as the usual exp and log. The Tsallis negative entropy defined as
Galp] = (1/(2 - 3)) p(()].
(3)
3- Gaussians. Given location parameter E R and symmctric positive semidefinite scale parameter Σ, the multivariate Gaussian (u,, Σ) has density p(u) = exp(- u) E--(2 u) - A(E)), where A(E) log / ((+r)*(() is normalizing constant that ensures the density integrates to The multivariate B- G-Gaussian (Martins et al., 2022) is instead defined by the following density function:
p(u) = exp, (-2 (o(-- u) TI (v T
0 N~Nu,
-
(4)
where B is the Tsallis entropy parameter and the normalization constant is given by
1-g
3-(1-8jm
R²2 2
+ 3-1/2
2 H
Ax(E) = 3 3-1
Σ (=+2/(1-0)) where R =
5
T
3-1/73
1 - B
In the limit of + 1, exp, recovers exp and lim3-11 As(E) = A(E), and so, by continuity, the 1-Gaussian is the usual Gaussian. For & 0, we obtain truncated paraboloid; for B 000, the B-Gaussian converges to
4

---

Published in Transactions on Machine Learning Research (12/2023)
the uniform distribution an ellipsoid. Introducing convenience parameter Σ, Martins et al. (2022) show that 8-Gaussians are elliptical distributions satisfying
(u,
y y= where ~ No(0,,1) E= |
Σ,
(6)
and that the support (,.e. the set of points with nonzero density) of the general B- B-Gaussian is an ellipsoid: sup Σ)) {o: w) u uv R2),
where the radius R (Eq. (5)) only depends / and the dimension of the space.
Fenchel- Young losses. For v & supp(Na p(e) = 0 and therefore gradient-based fitting cannot be used to fit the parameters by maximizing log-likeli leglliklihood. Instead, Blondel et al. (2020) propose the framework of Youch (FY) losses, derived as a natural learning objective for distributions inducod by the Tsallis entropies. In particular, when p is a 3-Gaussian and q is an arbitrary densily, the FY loss between p and g is:
flg (q:p) = Qala + 2g151 Bo(w) [f()),
(8)
where (-)* denotes the Fenchel conjugate (Borwein & Lewis, 2010, $3.3), and f(v) = (v) T u) is the function generating the 3-Gaussian (Martins et al., 2022, Definition 3). When q is empirical distribution, [g) is infinite, but constant w.r. I. the learnable parameters of f, motivating the cross-FY loss:
(q p) = 13lf Egto) [f(v)].
The functions Ωg and Sg have closed-form expressions for B-Gaussians (Appendix A.1). As 3 + 1, L recovers the KL divergence and f* recovers the cross-entropy loss. Importantly, Elg : p) is finite when the support of q and p don't match, and (x p) is finite even when g has infinite entropy outside of the support p. This allows learning and gradient-based modeling.
3
The wrapped B-Gaussian distribution
3.1
Construction
We propose a tractable distribution general Riemanniin manifolds by wrapping suitably-paramettizel 3- 3-Gaussian defined in tangent space.
Definition 1 (Wrapped 3-Gaussian). random variable has wrapped 3-Goussian distribution with location p E M and scale ②, if y- = (v), where e ~ N0, E). We write y ~ (N,
As a wrapped distribution, we have that Cov(yi = Cov(o] and the Karcher mean is p (Chevallier et al., 2022), which is further discussed in $5.1.
For general wrapped distributions, including the wrapped Gaussian (B= = 1), assessing the probability density value p(y) is intractable. This due to the general non-injectivity of Exp,,, and is often sidestepped using approximations that assume small |E|] With wrapped B B-Gaussians, we can have exactly tractable density estimation, alongside all other benefits of wrapped distributions, by parametrizing Σ to ensure the support of is within the injectivity radius at μ.
To derive the necessary parametrization we use the following result on the ß-Gaussians:
Lemma 1. Let Σ = Σ -42/0-0) E, and R as in Eq. (26). If the marimal eigenvalue Amaa (Σ) then supp.Ng Σ) C B(0,r)
r²
<
R2
The proof can be found in Appendix A.2. In practice, it is often more numerically convenient to parameterize 3-Gaussians directly using Σ, so may use the two interchangeably.
Equipped with the result of Lemma 1, we derive the following tractable expression for densities, provided the scale is within a generous constraint set. By varying B, μ, and E, we can obtain an expressive family with adaptive support within the injectivity domain around any point.
5

---

Published in Transactions on Machine Learning Research (12/2023)
Proposition 1. Let y WNo(H,)) be a wrapped 3- 3-Gaussian on M, with Σ as in Lemma
If Amaxe (2) < (M) then the density of y has the exact expression
12
ply) expx (-2 (v) E (
θ (og
Log, (1) -
(10)
dy
Proof. From Lemma with r = ini, (M), we have that the support of (0, Σ) lies strictly in the injectivity domain of M at μe, on which Exp, and Log, are bijections. The Sfange-of-denssee formula (push-forward) yields the desired result.
Like B-Gaussians, wrapped B-Gaussians enjoy efficient tractable reparametrized sampling, since the Exp mapping is differentiable. This means all necessary building blocks for deep generative models are available.
3.2
Younh losses
Due to the compact support of B-Gaussians, p(y) is zero on points outside of the support. For fitting the distribution parameters data, therefore, the usual gradient-based marimum likelihood approach does not apply. In Euclidean space, Fenchel- Young losses ($2.3) address this concern. In this section, study their extension to wrapped distributions.
Losses in tangent space. On TyM, the natural learning objective for a 3-Gaussian the cross-Fenchel. Young loss, which for a single target point v E T,M takes the value
ex (6, : :) = -1-30 [[--)pppp]1 + A + 300 T
(11)
The derivation is provided by Martins et al. (2022, Proposition 18). For x E M, can therefore directly define the tangent loss
e (x) = 1xx (OLog, :) -1- + As(E) + Log, (x) (x)T=> T Logg (EgLL,,
(12)
Wrapped Manchel-Yuung losses. While the tangent space loss is adequate and we find it to perform well in practice, it does not account for the distortion induced by wrapping. Notably, while in tangent space lim L* (6 p) - log p(v), the wrapped case lim3-11 (xx) * - logg(()) as the Jacobian term in Eq. (10) is not accounted for. The correct way to incorporate a Jacobian term J(x) from a change of density, under the Tsallis non-extensive system, is not straightforward from Eq. (11). We next show how to relate the Fenchel- Younh losses with the probability measure directly by proposing a new rearrangement.
Notice that, in tangent space, if p(u) then
g (8, p) = (1- 8)0b(p) log p(v),
(13)
where p(v) = exp3 (f(v) Ag(E)) = exp (h(v)) Substituting p(v) = p(Log(r)) [J(a)) we can apply
logs(ab) 1og (a) + log, (6) (1 3) (a) (bg =-- logs(a) + log,(6).
(14)
Then, introducing h(v) = - 1/2 T (E), get the following the wrapped loss:
(x(x) = (1 -())),))-))))))"(X(([((())))))) (og)((
(15)
Contrasting Eq. (15) with E% (x) = (1 (1-3)0s(pp) (L),
(z)) reveals that the Jacobian has both multiplicative
and additive contributions; morcover, the loss satisfies the desired property limæ-11 L (z) log p(x).
Outside of the support, log, and exp, are not inverses of each other, but we can write I T,x(6,, p) =
(1 [)] log, Po(v) where Po(v) [1 + (1- /)((v)) is the "unclipped" version of p. continuous and agreeing with p(u) the support. Since log, (v) h(v), may Eq. (15) on all of M.

---

Published in Transactions on Machine Learning Research (12/2023)
We stress that even outside of the compact support of the wrapped 3-Gaussians, the proposed FY losses are finite and effective in learning on the entire domain of Log, (2) (e.g., for spheres, this excludes only the antipode -u).
3.3 Parametrizations of the scale parameter
Instantiating (and modeling with) a wrapped B-Gaussian ( Σ) requires the specification of B-Gaussian Ns(0, Σ) such that (i) its support of (0, Σ) is on the tangent space TyM and (ii) the spectrum of its scale is bounded following Proposition 1. In this section, detail several ways to achieve both desiderata efficiently.
To ensure that a langenl space G-Gaussian distribution remains within the injectivity radius, we parametrize the cigenvalues of E to have values between and Amax ini(M) from Proposition 1, using the sigmoid R
function in log- fog-domain: we fit real parameter s E Rm and set
log Aÿ = log Amax log(1 + exp(-88))-
(16)
We may then use == diag(exp \) feasible diagonal scale matrix. While we do not report experiments with full-scale matrices, the same stralegy may be used, introducing additional orthogonal parameter matrix for the eigenvectors.
Local coordinates. The scale parameter can be defined as n-dimensional matrix in some chosen basis of T.M. This requires an arbitrary choice of basis which should depend on f continnously, leads to complicated gradients, and makes the comparison between distributions at different points challenging. We propose identifying a point p E .M (a pole) and parametrizing a matrix E (or Σ) in TyM. Then we may implicitly define En via parallel transport map such that
PTP-PU ~ (0, Σp) where v ~~ N2(
(17)
For Riemannian manifold M embedded in R", parallel transport is isometry, therefore the two densities are equivalent and we can assess the probability of any vector in Tp by first applying PT pT,--p" In other words, in embedded coordinates PT has an orthogonal matrix representation Quc and |Σ 1Q, = [Σ]
Ambient coordinates. Alternatively, we may specify E as an n-by-n matrix in ambient space and denote Py the rank-d orthogonal projection matrix onto T,,,, and set En, P.EPPA This removes the dependency on an arbitrary pole and leads to a more interpretable paranctrization. In general, carrying out calculations with P,EP, requires a costly factorization even when is already decomposed or diagonal. In $4.2 we show efficient expressions that avoid costly factorization when working with the sphere manifold S", Since projection matrices have cigenvalues and 1, it is sufficient to constrain the ambient as described above, to obtain feasible scale matrix after projection.
4 Instantiations on specific manifolds
4.1 Geometric toolbox
Table 1 summarize all the necessary geometric tools to instantiate the wrapped 8-Gaussian on two manifolds: the hypersphere (8") and the manifold of rotations, also known as the special orthogonal group (SO(n)).
4.2
Hypersphere g embedded in R
Relative coordinates. Without loss of generality, we choose the north pole p = (O,...,0,) 1) as a reference. The corresponding tangent space is T,M = {(v, 0)lu æ=-1) where we can easily parametrize a scale parameter by removing the unused last coordinate. We rotate the space via parallel transport, and implement it composition of two Householder reflections in R"; PTp-- Ry o R ptp (Algorithm 1). R ett_ maps be considered as identity.
point from T T1 to T,S"-1, and R2 maps any point T,8n-1 to itself,
7

---

Published in Transactions on Machine Learning Research (12/2023)
Figure 3: Wrapped 3-Gaussian density on 32 for 3 = 0.9, 3 0.5, (from left to right).
Table 1: Geometrie toolbox for the sphere and rotation manifolds. For the Jacobian on SO(n), til, the
cigenvalues of VV,, m |n/2], and y mod 2.
Hypersphere 8" (Appendix B)
Special orthogonal group SO(n) (Appendix D)
Tangent space
T28" {u E R", (z, 2) = 0}
TxSO(n) = {V E ER"X" V x =0}
Injectivity radius
inj g" = π
inj SO(n) = TV2
Expoment
(x),(r)
cos((|o|l2)ır sin([|ol) ािा a (y (),))=)
Expx (V) X' Xexp(XTV)
Logarithm
Log, (3)
Logx (Y) = XX Xlog(X (og(XT))
V1 (u.=)2
Jacobian
Ex)
(sinc((lell2))) "-
Expy V 11
sinc?
( (1=02
sinc?
de
OV
Ambient parametrization. On the sphere, we provide result that allows the ambient space parametriza- tion can be used efficiently. We may posit scale matrix Σ E R=*= and observe that the projection operator Pμ = pup is rank-one matrix. The projected scale Σn P.E has the subspace Ry in its null space, and we have Amax(E,) I Amax(E). Thanks to its structure, we can derive exact expressions for the pseudoinverse Σt and the pseudodetermiiant [3ul+ in terms the ones of Σ.
Proposition 2. Let S be a n-by-n positive semidefinite matric parameter and P = (I xx ) be the projection operator onto the hyperplane orthogonal to a unit vector x. Let projection onto the kernel of S, and x Rx. Then, we have:
1 - S (+++ be the
SS| x S x
B= 0,
(i)
PSP
|S| . B,
8 O.
S+ Stan st
3=
(it)
(PSP)*
(1 1
Stx
S+ (1- REx "
B + 0.
The proof, and further implementation details, can be found in Appendix C. If E is diagonal, this result allows us to perform all calculations needed for assessing densities or FY losses in O(n), and if is stored in a factorized form, in O(n²). Without this result, a cubic-cost eigendecomposition I is guaranteed full-rank (e.g., by parametrization) then only the B = case is noeded.
E. would be required. If
4.3 The Special Orthogonal manifold of rotation matrices 50(n)
The manifold S0(n) is well-studied Lic group. As an embedded submanifold of Igax", its elements are identified with orthogonal matrices with determinant. 1, i.e., rolation matrices. This makes it valuable in practical applications involving rotations and alignment between spaces. The tangent space at the identity matrix the vector space of skew-symmetric matricos, with dimensionality n(n - 1)/2. On this manifold, we only employ a pole-based parametrization. The Exp mapping involves the matrix exponential, and its Jacobian has efficient expression which we present, along with the rest of the toolbox, in Appendix D.
8

---

Published in Transactions on Machine Learning Research (12/2023)
tangent
20
1.1
wrapped L
15
1.0
Lo5s
10
0.9
5
0.8
0
-2
0
2
0.10
-0.05
0.00
0.05
0.10
Aμ
A
(a) Fenchel-Young losses w. t. shifting the true location μe: x² = (PTT
(b) Fenchel-Youre losses w. r. shifting the true σ² of the WNa(u,
0)).
Figure Visualization of wrapped L* and tangent e losses f. shifts of (a) location, and (b) 8² parameter of (N, z random the $ and the trace of Cov,, 0.3. Losses averaged over 10000 samples from the true WWg distribution (3 = 0.9). Additionally, in Appendix B.4, we provide the plot for 8 where we observe the higher influence of the Jacobian correction term from Eq. (15).
5 Experiments
We conduct a series of experiments to showcase learning with wrapped 8- 8-Gaussians.
In our experiments, we use 3 = 0.9 as a default value, which the distribution shape is relatively close to a Gaussian distribution (B 1.0), but with finite support size.
5.1 Parameter estimation: moment matching and FY loss minimization
As Chevallier Guigui (2020) note, one of the important advantages of wrapped distributions is a more straightforward relationship between the moments and the parameters of the distribution.
For Riemannian manifolds, the Euclidean definition of means does not apply, since manifold points cannot be summed. A natural definition instead comes from a center of mass optimization problem:
J Mean[æ] = arg min [d(, arg min yEM
(18)
N
BEM
1=1
Over IR" this recovers the expression 1 [,, [r]. For general Riemannian manifold M the solution may not be unique. A global minimizer of Eq. (18) is known a Fréchet mean, while local minimizer is known as Karcher mean.
The covariance is defined in tangent space and thus matches the one of the tangent distributions (Chevallier et al., 2022, eq. 2.6):
N Log, (2x) Log, (x;)
Cov = E,[Log, (x) (og
[I)
(19)
N
In Euclidean space, the asymptotics of estimators from empirical measurements is well studied; on manifolds the situation more difficult. Some central limit theorems for Frechét means (Bhattacharya Patrangenaru, 2005; Pennec, 2019) apply under the Karcher-Kendall concentration conditions (Karcher, 1977; Kendall, 1989), which imply existence and uniqueness of the Frechét mean when the support is strictly contained inside geodesic ball B(,, & min(inj.: m/)K)), where is an upper bound of the sectional curvature. Wrapped 3- Gauusssiass can be readily constrained to satisfy this condition.
Although we can use the moment matching method described above to dircetly search for a Karcher mean 2 and estimate empirical covariance using Eq. (19), applying gradient optimization to Fenchel- Younel losses is
We the geomstats package (Miolane al., 2020) find Karcher the sphere.

---

Published in Transactions on Machine Learning Research (12/2023)
animal
2-
animal
lion animal
chihuahua
bird seaguli
dog
cat
eguar
9-
corgl
parrot
Figure 6: Support of the embeddings trained on lexical entailment task: pole parametrization. For the visibility split the plot into three subplots according to the subtrees of nodes dog, cat, and bird (from left. to right). The supports are sent to the tangent space of the words dog, cat, bird via logarithm map.
alternative which also works part of end-to-end deep learning modeling. We visualize the Fenchel- Young loss in the Figure 4. In Appendix B.4 we demonstrate fitting a G-Gaussian to samples, and show that indeed minimizing FY losses converges to the true distribution.
5.2 Modeling hierarchies on the sphere with FY losses
In this experiment, remonstrate Tenchel-Young losses for modeling hierarchical relationships between distributional em- beddings. We employ a synthetic hierarchy inspired by lexical entailment (Figure 5) and embed each word as wrapped g-Gaussian. Vilnis & McCallum (2015) learn distributional em- beddings with KL divergences between R" Gaussians. Inspired by Muzellee & Cuturi (2018), we consider compact-support elliptical distributions, that may visualize the embedding of word the support of the distribution. Motivated by the relationship between the FY loss and the KL divergence, assign a WNg distribution Pw for each word w, and minimize:
animal
cat
dog
bird
corgi
lion jaguer tiger
dalmatian chihuahua
seagull parrot
Figure 5: Synthetic dataset emulating lox- ical entailment.
L(w)= = L(pw Pa(w)) log Σ exp L(pw Pw
(20)
where the term L(Pw Pr(m)) is calculated for words and its parent (w),, and L(p Pw:) is loss for negative pair (no entailment edge), where w' is negative sample, that is not ancestor of w. The k = 3 negatives are sampled with a replacement for each w in batch, and L(w) is averaged over the batch. We train for 30 iterations with a learning rate of 0.05 and use full-batch training. We use 0.1 to prevent. L(pw PP')) from dominating the learning objective when positive L(w : p) becomes close to zero. The deep connection between 3-Gaussians and FY losses allows us to learn with finite gradients even when supports are disjoint, unlike the KL divergence; we can then more directly generalize the approach of Vilnis & McCallum (2015) without switching to the Wasserstein geometry like Muzellec & Cuturi (2018).
Results. In Figure 6, visualize the learned probabilistic embeddings. We observe, that the learned embeddings tend to organize the hierarchical structure: support of the superordinate incorporating the support of hyponym.
5.3 Anisotropic hyperspherical VAE on MNIST
In this experiment, we evaluate the wrapped 8 GaGusssian on dynamically binarized MNIST (Salakhutdinov & Murray, 2008) reconstruction task following Davidson et al. (2018). We train variational auto-encoder
10

---

Published in Transactions on Machine Learning Research (12/2023)
(Kingma & Welling, 2014, VAE:), optimizing the evidence lower bound (ELBO) objective:
logp(x) 2 (x)pe - KL]go(z|o) p(2)].
(21)
We explore modeling the latent vectors on the n-dimensional manifold s". We define the encoder output (0(2(z) as wrapped B-Gaussian, and use a uniform spherical prior p(z) o 1. Even if the proposal has compact support, as long as its support is included in the support of the prior, the ELBO in Eq. (21) is finite and well-defined, therefore we may use standard probability theory and do not need to resort to Fenchel- Young losses, for which corresponding nonextensive ELBO would be nontrivial.
We estimate the gradients of the ELBO using Monte Carlo methods. For the reconstruction part of the ELBO (the first term), we apply the reparametrization trick (Kingma & Welling, 2014) to W.2,, as compose Eq. (6) with the differentiable Exp mapping. We estimate the KT. term with the common random numbers strategy (Blundell et al., 2015; Owen, 2013), reusing the same sample as in the first term.
We follow the hyperparameter setting Davidson et al. (2018) using MLP with 2 hidden layers for both the encoder and the decoder: [256, 128) hidden units for the encoder and [128, 256) hidden units for the decoder. We trained for 1000 epochs using the Adam optimizer (Kingma & Ba, 2015) with mini-batches of size 64, and with linear warm-up for 100 epochs and maximum learning rate of 0.001.
Results. The results in Table 2 indicate that wrapped 8-Gaussian VAEs perform slightly worse than the isotropic S-VAE baseline in higher dimensions n 2 20, but can outperform it for smaller latent di- mensions = 2, 5, 10. We obsorve that the embedded parametrization has overall lower variance compared to the relative to the pole parametrization and scales better with dimension.
Summary of the results for MNIST exper- iment for different latent space dimensions S". We report the averaged log-likehoods estimated via im- portance sampling with 500 samples (Burda et al., 2016), well the standard deviations. Results are averaged over 10 runs.
Table 2:
relative WNg-VAB
embedded WWg-VAB VAE
5.4 Multilingual embedding alignment via Bayesian orthogonal Procrustes
dim
S-VAE
2
132.50+0.73
-130.10+1.0G
-131.3640.91
5
-188.4310.09
-107.18
-107.16
For most natural language processing tasks it is sub- stantially easier to obtain monolingual data rather
10
-9..1640.31
-92.60
-92.7840.25
20
-8..02
-97.94
90.5810.3s
than cross-lingually aligned data. Word embeddings (Turian et al., 2010) provide the opportunity for find- ing alignments between the implicit Euclidean spaces occupied by two languages. It is common to restrict the scarch to linear aligmments (Mikolov et al., 2013a; Dinu et al., 2015; Lazaridou et al., 2015): given a dataset of paired embeddings (u,, vz) E D, obtained using bilingual dictionary, we seek
40
-90.87
+0.34
-100.4944.03
-95.62+0.36
arg min
Σ ||4 Xv, 2/2 12:XE CE Rax").
1.
(22)
A common choice of constraint is M SO(n), which restricts the space of alignments to rotations. As rotations preserve angles, this orthogonal constraint leads to good performance in number of NLP tasks (Xing et al., 2015; Arteixe et al., 2016; Hamilton et al., 2016). Equation (22) is known the orthogonal Procrustes problem, and a direct solution is available from the SVD of the matrix Σ, V.u," (Schönemann, 1966). However, in some situations that require more advanced analysis, a Bayesian treatment can be more informative than point estirattion. For probabilistic treatment, note that Eq. (22) is a constrained regression with unit-variance Gaussian observations:
p(u x) x exp (Ilu X2)l3/2).
(23)
We treat a random variable over SO(n), with known prior distribution p(X): for simplicity, we pick the uniform distribution under the Haar measure. Under this measure, the volume of S0(n) is the product
11

---

Published in Transactions on Machine Learning Research (12/2023)
prognoses, prognosi
cooking, cottura
pregnancy, gravidanza
animation, animazione
photograph, fotografia
billingham, billingham
head, capo
3.00
3.25
3.50
Distance between rotated english and Italian word
Figure 7: Euclidean distances between Italian embedding and rotated English embedding. English emboddings are rotated by 30 samples from the (X). The smooth distributions are 1d Gaussian kernel density estimation plots (original samples are present well).
of the sphere surface areas for dimensions ranging from 1 to n, leading to: (Zhang, 2015, Theorem 2.24)
1 2 2n-1ma(x11/1
p(X) Vol(SO(n))
(24)
Under the uniform prior, the marimum a posteriori point estimate is the same as the standard orthogonal Procrustes solution Eq- (22), but the full posterior distribution p(X (D) is intractable. We perform variational inference with wrapped 8-Gaussian approximate posterior
(x(X) Exx,)x,Ex),
where μх SO(n) is trained location parameter parametrized px is a covariance operator which is diagonal in our chosen basis of TiSO(() (per Proposition 5), and thus parametrized as a log-domain n(n - 1)/2-d/ dimmeensio vector. We then maximize an evidence lover-bound:
Rex" and Ex
logp(u | v) 2 [locx)]logp(u X, - (L(g(()))) : p(X)],
(25)
using reparametrized samples from (x( and the tractable expression of the density log (x()), like in $5.3.
We this probabilistic orthogonal Procrustes model on the English-Ttalian alignment dataset from Dinu et al. (2015), where the source and target spaces are independently trained continuous bag-of-words representations (Mikolov et al., 2013b) of dimensionality encounter no numerical issue in optimizing, starting with random initialization, and
300. (The dimension of SO(n) is then 44 850.) Despite the very
high dimension, converge to probabilistic with precision@1 test set performance just slightly worse than the point estimate (42.3 + 0.3% vs. 43.3%). Figure 7 demonstrates the kind of analysis of some individual word translation made possible by the probabilistic treatment: some word pairs with more ambiguity (e-g., pregnancy gravidanza) get mapped to a wider range of values, as measured by the distance from the expected translation in the dictionary, compared to other words, like the toponym billingham, which much less ambiguous. In addition to the analysis, the probabilistic treatment. also allows to incorporate further domain--pecific knowledge through the prior p(X), which need not be uniform.
6 Related work
Prescribed distributions on spheres. The sphere manifold S" has in particular attracted much attention from the field of directional statistics (Mardia & Jupp, 1999). Possibly the most widely used distribution on S" is the Langevin distribution, obtained by conditioning an isotropic Gaussian. De Cao & Aziz (2020) propose the Power Spherical distribution, which is similarly shaped to Langevin but more numerically stable in high dimensions. For anisotropic covariances, the counterpart to Langevin is the Fisher- Biihh-- distribution
3Te theorem gives the volume of O(n), which consists of two disconneceed subsets, each isomorphic SO(n).
12

---

Published in Transactions on Machine Learning Research (12/2023)
(Kent, 1982), which in general is not tractable, requiring mumerical integration (Chen & Tanaka, 2021). The Spherical Normal (Hauberg, 2018) is an intrinsical Riemannian construction supporting full covariances, but neither densities nor reparametrized samples are available in general dimensions. The wrapped 8 8-Gaussian supports full covariances without sacrificing exact and efficient density assessment and sampling.
Wrapped distributions. Chevallier & Guigui (2020) study wrapped distributions within full generality, and Chevallier et al. (2022) provide a general form for the required Jacobian in symmetric spaces (including spheres and SO())-) Tf the injectivity radius is infinite (e-g-, for hyperbolic space), wrapping a Gaussian leads to tractable and well-performing models in large-scale experiments (Nagano al., 2019). For finite injectivity radii, however, the infinite support of the Gaussian leads to intractable integrals, which require either ignoring the tails (an approximation only valid for high concentrations), or truncating the distribution, which is generally not tractable (Griaz-Garcia et al., 2022). Our construction, based on 8-Gaussians, leads to a practical solution for exact wrapping on manifolds with finite injectivity radii.
Implicit distributions on manifolds For statistical inference or modeling, more flexibility is required than even anisotropic distributions can provide. Approaches based on neural networks have been extended to Riemannian geomelry; including normalizing flows (Rezende al., 2020; Mathieu & Nickel, 2020), score-based methods (Bortoli et al., 2022), and diffusson methods (Huang et al., 2022; Okhotin et al., 2023). Sampling from general energies on manifolds is achievable by methods like Geodesic Monte Carlo (Byrne & Girolami, 2013) and Riemannian Stein Variational Gradient Descent (Liu & Zhu, 2018) (Liu & Zhu (2022) provide a review). The flexibility of these methods comes at a higher computational cost.
7 Conclusion
We introduce wrapped B-Gaussians: a flexible family of distributions on Riemannian manifolds, with efficient expressions for sampling, learning, and exact assessment of density. We adapt the Fenchel- Young losses, the natural learning criterion for 3-Gaussians, to take into account the manifold curvature. We instantiate our construction on spheres and rotations, deriving new expressions for an ambient-space parametrization on spheres, and experimentally validate the utility of 3-Gaussians for learning, variational auto-encoding, and digh-ddmensional variational inference for orthogonal Procrustes. We release an open-source implementation that we hope will aid rescarchers and practitioners exploring new applications and other manifolds.
8 Limitations
In manifold statistics, there is a classical trade-off between the wrapped and intrinsic approaches. Intrinsic distributions have energies defined directly by geodesic distances and have more intuitive contours and optimization behavior when far from the location. Wrapped distributions incur a distortion from tangent. space, but in exchange are algorithmically friendlicr. The advantage of 3-Gaussian is that by restricting the support size we may keep all probability mass in regions of low distortion, desired.
Our work is directly applicable a limited but fairly wide choice of Riemannian manifolds, where standard tools for wrapped distribution are available: Exp, Log maps, determinant of their Jacobian, and an expression for the injectivity radius. This includes a widely used class of symmetric Riemannian manifolds (Chevallier et al., 2022), of which S", SE(n), SO(n) are particular examples.
By construction, the support size of the wrapped ß-Gaussian is bounded by the injectivity radius at a particular point on the manifold. Depending on the injectivily radius, there may be an empirical distribution that cannot be fit by single wrapped 3 B-Gaussian. In such a situation, mixture distributions may a be promising extension of our approach.
Acknowledgements
This publication is part of the project VI. 21ni.222.2 of the research programme 'Veni', which is financed by the Dutch Research Council (NWO); and is part of 'Hybrid Intelligence: augmenting human intellect!
13

---

Published in Transactions on Machine Learning Research (12/2023)
with project number 024.004.022 of the research programme 'Gravitation' which is (partly) financed by the Dutch Research Council (NWO). We thank the action editor and the anonymous reviewers for their helpful comments and discussion, which have improved this text. We are thankful to Wilker Aziz, Caio Corro, Mario Figuciredo, André Martins, Dmitry Molchanov, and all members of the Language Technology Lab for helpful leedback for their valuable leedback on earlier drafts of this work. This work built on open-source software; we acknowledge the scientific Python stack (Van Rossum & Drake, 2009; Oliphant, 2006; Walt al., 2011; Virtanen al., 2020; Paszke al., 2019).
References
Mikel Artetxe, Gorka Labaka, and Eneko Agirre. Learning principled bilingual mappings of word embeddings while preserving monolingual invariance. In EMNLP, 2016. doi: 10. 10.18633/v1/D16-111 URL https:
//aclanthology.co 1ttpaaattooooog.oo///-----1- org/D16-1250.
Seth Axen. The injectivity radii of the unitary groups. online, 2023. URL httpe://sethaaeen.com/bboo/ 2023/02/the- http://tttteeiiiiiiivvrrrraaaaaa wwwoofftehutttee eee--nnttrrggrupppp Accessed 2023-05-24.
Gary Becigneul and Octavian-Eugen Ganea. Riemannian adaptive optimization methods. In ICLR, 2019.
URL https:://ppenrevies..o http:///wwwt.aa....omoooottttttd
Rabi Bhattacharya and Vic Patrangenaru. Large sample theory of intrinsic and extrinsic sample means manifolds- II. Anm Stal, 33(3), 2005. doi: 10. 10.1214/000033000000000000
Mathieu Blondel, André F. T. Martins, and Vlad Niculae. Learning with Fenchel-Young losses. J Mach Learn Res, 21(1), jan 2020. ISSN 1532-4435. URL https://jmlr http:p:///hll..oogapprrrr//////s org/pappers/v211110021 html.
Charles Blundoll, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In ICML, 2015. URL https://proceedings. .aa.g......pppp........ pttppeee///37/bllmmm1111.. http:///ppoccooig......rrrpessss pdf.
Valentin De Bortoli, Emile Mathieu, Michael John Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score- based generative modelling. In NewrIPS, 2022. URL met:/forumriideooD0G666177.
Jonathan Borwein and Adrian S. Lewis. Conver Analysis and Nonlinear Optimization: Theory and Examples. Springer Science & Business Media, 2010.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In ICLR, 2016. URL http:/ attp://arxiv. arxiv. org/abs/1509 htt/a////15590005.0/0 00519.
Simon Byrne and Mark Girolami. Geodesic Monte Carlo on embedded manifolds. Soand Stat, 40(4):825 845, 2013. URL https://wwar http://ww..orr.r//o/////.o/oo..t jstor.
Yici Chen and Ken'ichiro Tanaka. Maximum likelihood estimation of the Fisher Fisher-Binghhee Bingham distribution via efficient calculation of its normalizing constant. Stat Comput, 31:1-12, 2021. doi: https://doo.org/10..00 1001100/ 911222-021-100159.
Emmanuel Chevallier and Nicolas Guigui. Wrapped statistical models manifolds: motivations, the case SE(n), and generalization to symmetric spaces. In Workshop Joint Structures and Common Foundations of Statistical Physics, Information Geometry and Inference for Learning, 2020. URL https:
Emmanuel Chevallier, Didong Li, Yulong Lu, and David Dunson. Exponential-wrappel distributions on symmetric spaces. S1AM J Math Data Sci, 4(4):1347-133 1368, 2022. doi: 10.1137/21M1461551 URL
https://doi.org: orwpii..rro1..... 115///0001137/2114615555
Joseph R. Curray. The analysis of two-dimensional orientation data. J Geol, 64(2):117-131, March 1956. doi: //tp:::////o.eeee htttps//ddo..org/000/ 6266/9
10. 10.1086/626329 URL
14

---

Published in Transactions on Machine Learning Research (12/2023)
Tim R. Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M. Tomezak. Hyperspher- ical variational auto-encoders. In UAI, 2018. URL https://research. vu. al/en/publications/ encoders.
Nicola De Cao and Wilker Aziz. The power spherical distribution. In JCML Workshop on Invertible Neural Networks, Normalizing Flous, and Explicit Likelihood Models, 2020. URL attp:://awvertbbeeoxxxhpp github.ion 2020/
Georgiana Dinu, Angeliki Lazaridou, and Marco Baroni. Improving zero-shot learning by mitigating the hubness problem. In ICLR workshop, 2015. URL https://aassa http:://axxvv.orrr//// borg/bss/ae 1ate./111.... 14422.5668.
Fernando Galaz-Garcia, Marios Papamichalis, Kathryn Turnbull, Simon Lunagomez, and Edoardo Airoldi. Wrapped distributions on homogencous Riemannian manifolds. preprint 2iiu:2204009700 2022. doi: https://doi.orr/lul 10. (at.a//orgg/../55//a///i//i/ 2204.09790.
Felix R. Gantmacher. The Theory of Matrices, vol. I. AMS Chelsea Publishing, 1959.
Matthias Gelbrich. On a formula for the 1,2 Wasserstein metric between measures on Euclidean and Hilbert. spaces. Mathematische Nachrichten, 147(1):185-22)) 203, 1990. doi: http:://l6..rg//1..00/////a 0tt.:///gg/11.000///mmmm00111011 19901470121.
Ben Grossmann. Pseudo determinant of product of square matrices. Mathematics Stack Exchange, 2022.
URL httpss//mat..stacceechhnggg com/ 3732640.
William L. Hamilton, Jure Leskovec, and Dan Jurafsky. Diachronic word embeddings reveal statistical laws of 1141. URL https:
semantic change. In ACL, 2016. doi:
Søren Hauberg. Directional statistics with the spherical normal distribution. In Fusion, pp. 704-711, 2018. doi: 10.23919/ICIF 8015.2455.4..
Nicholas J. Higham. Functions of Matrices: Theory and Computation. STAM, 2008
Chin-Wa Huang, Milad Aghajohari, Jocy Bose, Prakash Panangaden, and Aaron Courville. Ricmannian diffusion models. In NewrIPS, 2022. URL
Karcher. Riemannian center of mass and mollifier smoothing. Commun Pure Aypl Math, 30(5) :509 541, September 1977. doi: 00.1022/cpa33160300002 URL sttps:///oi... 1ttp://0..00///aa.1.000000000000
David G. Kendall. A survey of the statistical theory of shape. Statistical Science, 4(2), May 1989. doi: 10.1214/ss/117701258) URL attps://doi. org/10.../. http:////oooooooo...1/1///111110
John T. Kent. The Fisher-Bingham distribution on the sphere. J R Stat Soc Series B Stat Methodol, 44(1):
80, 1982. URL attps://wwy, jstor .arr..or www. httgsttabbe/2984112
Diederick P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Diederik P. Kingma and Max Welling, Auto-encoding variational Bayes. In ICLR, 2014. URL https:
// //openrevier wtt.//ormmiddd333eett.t..
Angeliki Lazaridou, Georgiana Dinu, and Marco Baroni. Hubness and pollution: Delving into cross-space
mapping for zero-shot learning. In ACI-IJCNLP. Association for Computational Linguistics, 2015. doi: 10.3115/v1/P15-107, URL https://acllanthology http:///lmmtholooooorr////11112
Chang Liu and Jun Zhu. Riemannian Stein variational gradient descent for Bayesian inference. In AAAI, 2018. doi: https://doi.org: wttp.//rgg.....///aaiii.aaom v32.:/a.....2..11110
Chang Liu and Jun Zhu. Geometry in sampling methods: review on manifold MCMC and particle- based
variational inference methods. In Advancements Bayesian Methods and Implementation, volume 47 of Handbook of Statistics, pp. 239-293. Elsevier, 2022. doi: https://doi.org/10..000 http://bsshott.22..00. URL
https: //t/://wwww... http:p/c.e..ttttt............eae
15

---

Published in Transactions on Machine Learning Research (12/2023)
Jan R. Magnus, Henk G.J. Pijls, and Enrique Sentana. The Jacobian of the exponential function. J Econ Dym Control, 127:104122, 2021. doi: http:///0..ang/00000000//o0000.0
Anton Mallasto and Aasa Feragen. Optimal transport distance between wrapped Gaussian distributions. In International Workshop on Bayesian Inference and Mazimum Entropy Methods in Science and Engineering, 2018. URL http:://driiee.ggggle.ooo////// http:////....ee//aaaa/oa/taaaaa.
Satpi/vvMMeFee/vieeeeeeeeeeee Fe/view.
Kanti V. Mardia and Peter F.. Jupp. Directional Statistics. John Wiley & Sons, Inc., 1999. doi: 10.1002/ 9780470316979. URL http::///oo..ogg/0..000/99000000
André F. T. Martins, Marcos Treviso, António Farinhas, Pedro M. Q. Aguiar, Mário A. T. Figuciredo, Mathieu Blondel, and Vlad Niculae. Sparse continuous distributions and Fenchel-Yone Young losses. Mach Learn Res, 23227) 2022. URL http://jmlr.... httt:://jilr.orrppppppp... arg/papors/v23/21-0877 html.
Emile
Mathicu and Maximilian Nickel. NewrIPS, 2020. 1aa38lc66e6f2277eetee50001111110 wwt.rssttbttttt.hhca html.
Riemannian continuous normalizing flows. URL http::/prroeednggg.eeen www.ipu.rrp...m ottp://ddng...eereessccc//pp///p cccppaper/202///aaae /hahh/
In
Yu Meng, Jiaxin Huang, Guangyuan Wang, Chao Zhang, Honglei Zhuang, Lance Kaplan, and Jiawei Han. Spherical text embedding. In NewrIPS, 2019. URL https: //pppere hips.cc/papert filess/appee///19///////// html.
Carl D. Meyer, Jr. Generalized inversion of modified matrices. STAM J Appl Math, 24(3):315-32 1973. doi: 10.1137/0124033, URL https://doi.orr/1... 11111111/////1033
Tomas Mikolov, Quoc V Le, and Ilya Sutskever. Exploiting similarities among languages for machine translation. preprinl 130::309.44660 2013a.
Tomas Mikolov, 1lya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In NewrIPS, 2013b. URL p wtwpoceeeiinnnnnuuiiir..eee. neurips. pdff.
Nina Miolane, Nicolas Guigui, Alice Le Brigant, Johan Mathe, Benjamin Hou, Yann Thanwerdas, Stefan Heyder, Olivier Peltre, Niklas Koep, Hadi Zaatiti, Hatem Hajri, Yann Cabanes, Thomas Gerald, Paul Chauchat, Christian Shewmake, Daniel Brooks, Bernhard Kainz, Claire Donnat, Susan Holmes, and Xavier Pennec. Geomstats: A python package for Riemannian geometry in machine learning. J Mach Learn Res, 21(223):1 9, 2020. URL http://jmlr org Saapers/v1//11 0000-0277.tmm.. html.
Boris Muzellec and Marco Cuturi. Generalizing point embeddings using the Wasserstein space of elliptical distributions. In NewrIPS, 2018. URL https::/proceeddnngg neurips. cccpaper_fttteeeeeeeeeeeeeeepeee files wallespaape httpp/ppprr//1/ filc//ssrrrtrreeeeeraasssso00000
pdf.
Yoshihiro Nagano, Shoichiro Yamaguchi, Yasuhiro Fujita, and Masanori Koyama. A wrapped normal distribution on hyperbolic space for gradient-based learning. ICML, 2019. URL http://prrocedding.
mlr. mlr.preas/veel wttpppsss////naaaooooo1 html.
Maximilian Nickel and Douwe Kiela. Poincaré embeddings for learning hierarchical representa- tions. In NewrIPS, 2017. URL https: https::/proceeddngg neurips peapper_files/paper/00771 999t1221122000331110000022222220
.df..
Andrey Okhotin, Dmitry Molchanov, Vladimir Arkhipkin, Grigory Bartosh, Aibek Alanov, and Dmitry
Vetrov, Star-shaped denoising diffusion probabilistic models. In NewrIPS, 2023. URL https://arxiv.
org/abs/2302.0m 05259.
Travis E Oliphant. A guide to NumPy, volume 1. Trelgol Publishing USA, 2006.
Art B. Owen. Monte Carlo Theory, Methods and Feamples. https://artowen. www.dddmmaan/////om..mmm /mc/ 2013.
16

---

Published in Transactions on Machine Learning Research (12/2023)
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 NewrIPS, 2019. URL http://papers. http://paaprs..uuipppppppp neurips. neurips.cc/paper.. higer
.pdf.
Xavier Pennec. Intrinsic statistics on Riemannian manifolds: Basic tools for geometric measure- ments. Math Imaging Vis, 25:127-154, 2006. URL https://linke https://Ankk.spingge..oomm s10851-006- 6228-4.
Xavier Pennec. Curvature effects the empirical mean in Riemannian and affine manifolds: non- asymptotic high concentration expansion in the mall-sample regime. preprint 0tfi::9000074111 2019. doi: 10.48550/AR 00055/0//ARXV11990.07747
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racaniere, Michael Albergo, Gurtej Kanwar, Phiala Shanahan, and Kyle Cranmer. Normalizing flows tori and spheres. In ICML, 2020. URL
httpss//proccedingg. mlr.
htmme.
David M. Rosen, Luca Carlone, Afonso S. Bandeira, and John Leonard. SE-Syne: certifiably correct algorithm for synchronization over the special Euclidean group. Int J Rob Res, 38(2-3):95-125, 2019. doi: 10.1177/00273364918778383 URL http::///bb..org/1... /11.11111777777766961778885
Ruslan Salakhutdinov and lain Murray. On the quantitative analysis of deep belief networks. Tn ICML, 2008. /tt.//199915111900 URL attps:://doi.orr//000 http:///.....//11111111111 1390266.
ISBN 9781605582054. doi:
Peter Schönemann. A generalized solution of the orthogonal Procrustes problem. Psychometrika, 31(1):1-10, 1966. URL https: butps://line /tt////inkkk.//// springer..
Michael A. Stephens. Random walk circle. Biometrika, 50(3-4):385 390, 12 1963. ISSN 0006-3444. doi:
10.1093 1..0099/biomt///00
bit./m/ttt...4..33.. 5-.35- URL https: //op::://ddi orgg//0..19333/1mmt//555.. 3-...et...........555
Constantino Tsallis. Possible gencralization of Boltzmann- statisties. Stat Phys, 52(1-2):479-487,
1988. doi: 10.1007/b001016422 URL https: http:s///aas.orr//000000//fffff5
Joseph Turian, Lev-Aric Ratinov, and Yoshua Bengio. Word reprosentaaiio: A simple and general method for semi-supervised learning. In ACL, 2010. URL attpaa//aclaambbooog..or//O....0
Guido Van Rossum and Fred L. Drake. Pylhon 3 Reference Manual. CreateSpace, Scotts Valley, CA, 2009. ISEN 1441412697.
Luke Vilnis and Andrew McCallum. Word representations via Gaussian embedding. Tn ICLR, 2015. URL http://arxiv. org/ abs/ 1411.66233.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, and Haberland et al. SciPy 1.0: Fundamental Algorithms for Scientific Computing Python. Nature Methods, 2020. doi: http:///oo..or/1010033333 (0008//41192001-000005
Stéfan van der Walt, S Chris Colbert, and Gael Varoquaux. The NumPy array: a structure for efficient nunerical computation. Computing in Science & Engineering, 13(2):22 30, 2011.
Chao Xing, Dong Wang, Chao Liu, and Yiye Lin. Normalized word embedding and orthogonal transform for bilingual word translation. In NAACL-HLT, 2015. doi: 10.3115/v1N115-104. URL https://
acllmtthologg..or//bBB-110100
Lin Zhang. Volumes of orthogonal groups and unitary groups preprint arXio:1509. 00537, 2015. doi: 00-45500/ARXTVV15000000 URL https://arxiv. org/ orgg/bs//509.000 00537.
17

---

Published in Transactions on Machine Learning Research (12/2023)
Appendix
A Sparse continuous distributions
A.1 Expressions
For completeness, we cite here some expressions given by Martins et al. (2022).
The radius of standard 3-Gaussian is:
T
8-
B
1/2
2
1-
21
R =
(26)
π
1-1
B
The normalizing constant of -aaaussiinn has expression
1
R²
As(E)=
ΣΣ
122 23.
(27)
3-1
3
1
The Tsallis regentropy of 3-Gaussian p is
1 x127 37
2g1p)
(28)
(2 (2- (3)
2(2 B) n(1 3)
If f is the function that induces the sparse distribution p, then have the relationship (Martins al., 2022, Proposition 10)
pflf = (1 p) + Az(E).
(29)
In all expressions, if Σ is rank-deficient, we may use the pseudodetermiinat order to obtain a distribution in lower-dimensional subspace with respect the appropriate Lebesgue measure (Gelbrich, 1990)
A.2 Eigenvalues and support: Proof of Lemma 1
When the scale of ß-Gaussian is parametrized by E0 it easy to control for the support size of 3-Gaussian to strictly less than a desired radius r. Let Σ = PAPT for PPT 1, then:
sup 1412
sup ||TTI|
|1 R² Amaxx ($.) <
sup || RP
ETA MER2
B The wrapped B-Gaussian on the sphere
B.1 Inverse exponential map
In this section, we derive the inverse of the exponential map,
(x) cos([|ull2)>= sin((1ol12) Tollz
We seek v E Exp= (v). The first step is to identify the possible values for ||uil22-
(v, = cos([[=|12) (x, + sin((1)122)
110/ (v. 2)
(30)
||llla E farccos (v.x) + 2km 1: E Z).
18

---

Published in Transactions on Machine Learning Research (12/2023)
At this point, it is apparent why the Exp mapping is not injective. Since the Log mapping seeks the smallest-norm solution, we take ||u||2 = arcoos (). Finally, solving for the direction of v,
||0|| (y cos([|]|l))z)= = arccos (v. x)
Log (v) = sin |11|
(y (y, x)x).
(31)
V (v,
B.2 Determinant of Jacobian of exponential map Proposition 3. Let x E 8" and TgS". Then,
0 Expx(() sin ||||| ||l|
m-1
dr
=
Proof. We adapt the proof of Nagano et al. (2019) for determinant of the exponential map hyperbolic space. We start by noting that 8 Exp,(() is a function, which acts from T,S" to (o) We are free [Exp,(u) Honce, we choose a basis by including T and any other of orthogonal unit vectors in TyS" {T. ,,.,.,,................ - Next, we compute a set of directional derivatives w.r. l. to each of the basis vectors.
du
to select an orthonormal basis in which we compute the Jacobian
Hu
Exp,,
Exp,, + cō)
OE
c=0
0
v + e ||e eūll
cos([|ull2 e)x + sin([ + E)
Oe
c=0
sin([|=|l2) + cos([Iv|12)
Here we used de [[++
step
since the norm of the vector does not change w. r.t.
The norm of this directional derivative w.r. t. v is equal to 1:
|| sin(l[]))g) + [[ = v sin((1)))))) )² + cos([lu|1,) =
Next, we calculate the directional derivatives w. t. the other basis vectors u/....,, 2
i
d EExp, (v/)
Exp, (v + tv())
Oe
==
cos( cos((|o|l2)< + sin(())(l2) x + Evy
||:l12
c==0
sin
||0||2
r/6-
|11//2
The norm of the directional derivative u. n.r.t. 1. v/g is thus
sin sinl|vllz 1ll2
sin sinllell2
T
Finally, the determinant is the product of norms of directional derivatives since the chosen basis is orthonormal:
i (xEpa(t)) det
((snn
sin ||0|12 n-1
dr
||0///2
Note that here the dimension of the sphere and not the dimension of the ambient space.
19

---

Published in Transactions on Machine Learning Research (12/2023)
B.3 Wasserstein-2 2 distances and interpolation
Mallasto & Feragen (2018) provides analytical expression for pullback Wasserstein-2 squared distance between two wrapped normal distributions defined over the tangent bundle, which we can directly adapt to calculate pullback Wasserstein-2 squared distance between two WNg distributions WNg(M, E1) and
P2 = WWa En) transported to a shared tangent plane at some point
W2 = d' (2)
P2) + W PT (2)
(32)
For pole parametrization, when tangent B-Gaussians are already defined over common tangent space the expression simplifies to
R
(2 (2(,), +
2(2-9) 9) (2),),),
B²(E
(33)
N
1-8
where (A(A, = TT
(Martins al., 2022; Gelbrich, 1990).
In Figure 8, we plot the pullback Wasserstein-2 interpolation between two wrapped G-Gaussian distributions.
Figure 8: Pullback Wasserstein-2 interpolation between two wrapped B- G-Gaussians over the $
B.4 Fitting data with L loss
In Figure 9, we provide additional visualization for tangent and wrapped losses for higher dimensionality ($29), where the influence of wrapping higher.
Additionally, we experiment with fitting dataset of points on g² using the Fench- Fench-Young losses. To fit the wrapped B-Gaussian to the dataset X with Jacobian correction (Eq. (15)). For the simple experiment, from true wrapped B-Gaussian (initialized randomly to have non-isotropic covariance). We initialize the wrapped 8-Gaussian with a random location and scale parameters and optimize the location parameter with Riemannian Adam (Becigneul & Ganca, 2019), with the scale parameterized as in Eq. (16), with learning rate 0.01 (other hyperparameters are default) for 5000 iterations.
the tangent fx loss from Eq. (11), and compare to L loss
create dataset of size 10000 by sampling
Results In Figure 10, we show that by optimizing the Fenchel- Fenchel-Young loss, we can successfully recover the true location and covariance parameters. Both tangent loss and wrapped loss recover true parameters, while the tangent loss converges faster in experiments. In this experiment, the accuracy of finding ns,
20

---

Published in Transactions on Machine Learning Research (12/2023)
tangent ?
600
200
wrapped
400
Loss
200
0
-100
-1.0
-0.5
0.0
0.5
1.0
-0.02
0.00
0.02
Au
AØ
(a) Loss w.r. f. shifting the location μ from the true location: Log,, (PTpss 0)).
(b) Loss w.r.t. shifting the 8² of WNg(u,
Figure 9: Visualization of wrapped and tangent Fenchel- Young losses w. (a) location, and (b) scale shift from the true parameters of (N, 821. where μ is random on the g29, and the trace of Cov 0.3. Losses are averaged over 10000 samples from the true WNg distribution (B 0.9).
3.0
Langent L",
Wragped 0.99
Langent 13, 20,000000
2.5
wrapped 15, 8.000
Distance to the true distribution
tangent 13,
wrapped L2, 0.75
2.0
tangent 1 8 0.50
Wrapped
0.5
0.0
10a
10
10a
10a
Iteration
Figure 10: Fitting data on S?. For different 3 parameters, we fit the wrapped 8- 3-Gaussian the samples from the true wrapped B-Gaussian with the same B. For each run, we plot pullback W distance (Appendix B.3) between true wrapped Gaunssia and the fitted one (with the same $). Solid colored lines represent fitting with tangent. Œ and dashed colored lines represent fitting with wrapped I (with Jacobian correction). Different colors represent different B. On the x axis is the number optimization steps.
S is comparable for moment matching and FY loss optimization (< 0.02 pullback W2 error, as defined in A Appendix B.3). When available, the iterative process based Karcher mean is empirically preferrable.
C Extrinsic parametrization of covariances the sphere: Proof of Proposition 2
For full-rank covariance Σ defined in the ambient space, the following result allows calculate pseudo- determinants and pseudoinverses efficiently. For narrative purposes, split the proposition into separate results for the full-rank and rank-deficient cases.
21

---

Published in Transactions on Machine Learning Research (12/2023)
Lemma 2. Let S be n-by-n posilive definite scale parameler and P (I xx.) be the projection operator onto the hyperplane orthogonal the unit vector z. We have:
[PSP|+ = |5| T S--
(PSP)* = = 5 5-1xx rx S-1
(ii)
#S-- 5-11
(x S S)
(Mii)
T. (PSP))+== TS
S--SS1pp 1
x. 5-1x
Proof. (i) For any two matrices A, B, AB and BA have the same nonzero eigenvalues. Therefore, |PSP|+ = |PPS|+ = [PS| since projection operators are idempotent. Following Grossmann (2022), since PS has rank n - 1, its adjugate adj(PS) has a single nonzero eigenvalue equal to [PS|+ and therefore:
PPS = tr tr(adj (S)) = tr(adj(S) adj] ()) = |S| adj(P))
Since P has non-null eigenvalues equal to 1, adj(P) = xæ and so |SSP+ = =|S|=* T S-1x..
(ii) Let p be a square root of PSP, so M M PSP. For any matrix, (M) -M++++++ so we just need to find M Remark that
M = S - - = S1/2 + ed
with c = S and æ. M is thus a rank-one update to 51/2, Meyer (1973) categorizes such updates depending on whether e and d are in row (column) space of 51/2 and on the value 3 c S-1//222 T 51/2g-1/22 1-21 xx O. Since S is full-rank, its row and column spaces contain all non-null vectors, leaving under the scope of Meyer (1973, Theorem 6) which reads:
M S-1/2
... T 5-1/2
S-1x S-1zz
S-1 xx T S-1/2
x T S x
(-xx S S-1/2)
S-1x
The second and fourth term cancel out, leaving
M+ = S-1/2
1
S-1 S-1/22 S-112 -- (iu1--zT] T27).
x 5-11
where we define
1/2 [8-174m]
such that (I - 1-2zz] is an orthogonal projection, and thus idempotent. Finally,
(PSP)+= = =M+(M+)T 5-1/ 1-1-222 2z 22 5--1/2 1/2
S-' S
xTS-1n
Pre- and molt-mulg by any vector gives the relationship (iii).
In practice, covariances are usually enforced to full rank in parametrization. Nevertheless, for completeness we derive in the next proposition the corresponding expressions for low-rank Σ.
Lemma 3. Let S be an n-by-n low-rank matrir and R = ST be the orthogonal projection onto the kernel of S. Let x be a unit vector and P - xx be its tangent projection, and let B = x. Ra.
If 3 then Lemma 2 applies, replacing intermminants/inveers with pseudo versions.
22

---

Published in Transactions on Machine Learning Research (12/2023)
If 3 0, we have
(i) (ii) (PSP)* = 3 Rxx T) 8t (1- Rex
|PSP|+ = 18]+
T
Proof. For the 3 0 case, the condition implies x E Span S. As both P and S have no action on the kernel of S, we can apply a rotation that moves those dimensions to the end, and perform all calculations on the full rank top-left block.
When B + things dramatically change.
(i) We make use of the limit definition of psendodeterminants:
A = lim |A + all/ n n-rank a-+
(34)
In the full-rank case, rank PSP have ruled out, otherwise rank PSP rank S. Since P is a projection onto subspace of dimension n 1, we have P = UTT where U E R and U = In U has orthogonal columns, so the spectrum of PSP is equal to the one of U PSPU = SU. We therefore have
1. In general rank PSP rank S 1 only if æ E Span S, which we
|PSP| = UU PSPU|+
lim a-1
UU PSPU + al (a
lim
JPSP + aP aPl+/// (n-1))
(left and right multiply by and U
a-1
lim IP(S+al] al) p (1aan-1))
(35)
a-10
lim al r T + al)- x/ ( (n-1))a a-1
| lim ar (S al) a+0
= |S Rx,
where the final step comes from noticing that the eigenvalucs of a(S+al)"* are of the form 2Ú a-2, where Xy are the eigenvalues of S. In the limit, A if X and 0 otherwise, therefore lim, a(S + al) R, finishing this part of the proof.
(ii) We once again turn to Meyer (1973), finding ourselves under the auspices of their Thcorem 3, since 3 F 0. We define the same symbols as in the previous proof, except 5-1/2 is replaced by (51/2)+, and therefore the terms that become an identity matrix the full-rank case now become (S1/2)+(81/2) = T R. The qi term in Meyer's theorem is null. Carrying out the calculation leads to
-T-1 Rxx (82/2y+. M+ (r (r-} 1-RRxx"
Using again (PSP)+ (M++ yields the desired result. Efficient Mahalanobis distances by calculating vM+ without materializing M+ are left. as an exercise.
Combining the two lemmas yields a proof of the proposition.
D The wrapped B-Gaussian over Rotation Matrices
Our construction not. limited to spheres but is applicable and useful in other manifolds, as long as we know the injectivity radius, the Log and Exp mappings, and parallel transport.
In this section, we instantiate wrapped 3-Gaussians over the manifold of rotation matrices, also known as the special orthogonal group SO(n).
23

---

6m Published in Transactions on Machine Learning Research (12/2023)
D.1 Preliminary: properties of skew-symmetric matrices
The set of real skew-symmetric matrices is defined as
Skow (n) == {A R
xx A -A}.
The next result recaps some classical results about skew-symmetric matrices.
Lemma 4 (Properties of skew-symmetric matrices). Let A E Skew(n) with m [n/2].
(i) The nonzero eigenvalues of A are purely imaginary and come in pairs Xt₃ = i 1.8.
(ii) The eigenvectors of A have the form Paj 02 + vz with || = 1 for all j, and (g., qt) 0 for any s + t.
iii A admits the canonical decomposition
m 0 - 9-19t))-
Σ
(36)
(tv) For even n, forming the malric Q with columns (4+1,9-1,--- 1....... ............ mm q-m), and the block-diagonal matriz
01
0
0
0 0g
B
- 02 0
(37)
8m 0
Q orthogonal and we have = QBQT. For odd n, eatend with an additional zero row and column, and Q with unit vector from the null space of A.
(v) A admils a (potentially (runcaled) singular value decomposition with singulor triplets {(0), (+s- (-1))--),...... U {(0) 9-j. --1))
(vi) exp(A) is special orthogomal malrix and is equal to Qexp(B)QT.
Proof. A proof of properties (iv) can be found in Gantmacher (1959, $1X.12). For completeness, since the proofs are compact, re-derive them here.
(i) The complex eigenvalues and eigenvectors of any real matrix come in conjugate pairs: if Ap Ap, taking conjugates on both sides gives Ap = Ap = Ap. For skew-symmetric A, we have on one hand
(Ap,p) (Xp.p) 2(p.p)
and on the other hand
(Ap.p) =
- ->Xppp)
Therefore is purely imaginary, and for convenience we index the pairs as Xt₁ =i0, for 0 R.
(ii) If p is eigenvector of A associated with nonzero eigenvalue, then so is and since their eigenvalues are distinct they must be orthogonal. Writing b+ id,
(p.)) pT (b+ id) (b + id) (b a d) + iib d d b).
(38)
It follows that 6Tb= d d. But since eigenvectors have length 1, we have
1 ||||2 - =bb===
24

---

Published in Transactions on Machine Learning Research (12/2023)
Combining this with the fact that the real part of Eq. (38) is zero, we get ||6|| ||d] v2 set b 1 9+
and d = -- where ||0=||=1. From sctting the imaginary part of Eq. (38) to zero we get b = 0 so q q- = 0. For two distinct j, we use orthogonality to obtain:
-- 0 = 2/p-,. P+1) = -- + -- - + +i4
0 2(p+s. P+k) 9-,9+h
iq_
Adding and subtracting the two relationships reveals the orthogonality of q=s, g=k, and qrk.
(ii) First, note that, p = b b+ id.
pp H pp "1 = (b + (b) id) - (b (d) id) = -2i(db" - bd bdl).
(39)
We rearrange the eigendeccomposition:
A= Σ
Σ [i0,(P+sp" - P- 11
Σ i8 (2i(9 - )-
Σ 10,(a+59I, --1-t:)l.
(iv) and (iv). Both follow from (iii) and the fact that all qu are unit length and all distinct-index pairs are orthogonal.
(vi) We may start from the decomposition A = QBQT. Note that A = QB*QT for natural powers, since
QTQ = T. Since the matrix exponential is defined through a power series, we have exp( (A) = Q exp B)QT. cos 0, sin 0 exp(A) B) is a block-diagonal matrix with each block corresponding to rotation by f,: sin f, cos 8 and therefore is orthogonal with determinant 1. Therefore so is the product of QBQT.
D.2 Standard toolbox for 50(n)
SO(n) can be seen as an embedded manifold of Raxx defined as follows:
SO(n) = (X E B T X In. det(X) 1}.
The tangent space is characterized as
TxSO(n) = (V E Raxn : X 0}:
which can be seen by differentiating the orthogonality constraint X = I, in the ambient space.
The identity matrix In a distinguished element of S0(n). Its langent space T,50 is the space of skew-symmetric matrices Skew(n), which can be seen to be a vector space of dimension n(n 1in-10 11//2, and in this context it is often called the Lie algebra so(n). At the identity matrix, the Exp and Log mappings between (O(n) and (n)O(n) sn(n) are exactly the matrix exponential and matrix logarithm:
Ex), exp(A),
(Y) log(Y).
Now remark that any rotation X SO(n) is also an invertible linear operator that maps In to X. Moreover, therefore X A E SOn). Since orthogonal matrices
if A E (nnO(n) then X A+ AT
25

---

Published in Transactions on Machine Learning Research (12/2023)
are isometries, rotate the ambient space by multiplying by x-1 in order to identify Tx with so(n) (since the rotation takes X to In) and thus we have
TxSO(n) = (XA: : AE Skew(n)).
and we compute:
Expx X exp((XXA) LY) X log(XX'x Y).
From the same observation it also follows that PTx-Y YX-1. Remark. Since all orthogonal matrices are invertible, all tangent spaces are isomorphic to s0)).. It is common in practice to identify them and avoid some of the redundant parallel transports. With this representation we would write Expx(A) = exp( (A) and Logx (Y) = log(X"
The matrix exponential and logarithm are not univerally inverses of each other; the following result gives the injectivity radius.
Proposition The injectivity radius of SO(n) is \/2.
Proof. We recap the proof given by Axen (2023), with a few more details. By symmetry we need only focus on the injectivity of Exp, exp. Let A PAP" be the spectral docomposition E Skew(n), with cigenvalues of the form X+j i. j. A is in the injectivity domain iff A log(exp(AA) for all eigenvalues of A. The injectivity domain therefore constrains 03 E (-T,, σ). Let denote the radius of
equivalently À = log(exp(A))
the largest ball contained in the injectivity domain, and notice that || 1 tr(A A) 2 2Σ, 08. Consider the matrix A1 defined by 01 domain and || || = 222, so r S x 2 On the other hand, for an arbitrary we have || 2 20, for any j. and so || < 7V2 implies 10,1 < क for any j, so r 2 \V2.
and all other f, = O: this matrix is just on the boundary of the injectivity
Finally, for completeness, we give an explicit orthonormal basis of Skew(n), allowing us to parametrize wrapped distributions in this space.
Proposition 5. Let jjk E Ruxn be an indicator matrix with all elements set to zero except the element at inder j, k, which is set to one. Let denote the linear operator : Ren(n1)/2 - Skew(n) C Rexn with c
Djs
V2 (ejk e.p)
Then, is an orthogonal matrix.
"Given vector v E 1n(n-1)/2, Do (V V where V is lower triangular matrix with the lower triangle given by the elements of v.
Proof. Verify that the columns of D are by construction skow-symmetric, pairwise orthogonal, and have unit norm, as: [|D 1/ 2 2 =1 They therefore form an orthonormal basis of Skew (n) and we have DTD 1)/2
We can therefore construct wrapped 3-Gaussian on SO(n) as:
X ~ WN&(M,E) + x -xPxx (V) = M exp(VV), V = Dv, ~ Na E).
D.3 Jacobian Determenemnn
To compute the (log)-density of WNg random variables, we apply the change of density formula. Multiplication by an hotoogl matrix is an isometry, and we have shown that so is the embedding D. It. remains to calculate the change of volume induced by the matrix exponential.
26

---

Published in Transactions on Machine Learning Research (12/2023)
Proposition 6. Let A (n)w(n), [n/2], and denote the eigenvalues of A as メトل =10 for ......,.................... m, and for odd additional Ao = 80 0. Let J denote the Jacobion of the matrix exponential at A, i.e., JijMl 0(exp Alsg, Then,
BA
[8] (B)_0x
0k 2
18y + 0k
det J
II sinc 3m7
sinc²
("2 2
(40)
where the range of (but not of k) starts at yt- 0 if is odd, and at 7 otherwise,
Proof. First, we note the result is similar the case of Grassmanians given by Chevallier al. (2022, eq. 4.4).
We proceed to give complete proof.
Per Higham (2008, Theorem 3.9), exp is analytic, the eigenvalues of are
exp A -opp A. A--
8 # 4
Vst
exp Am:
8 = t.
For general characterization of the spectral decomposition of J see Magmus et al. (2021).
The eigenvalues of J are complex, but we next show that for skew-symmetric A the determinant always real. Let us assume for now that n is even. The eigenvalues skew-symmetric matrix are purely imaginary and come in conjugate pairs, so we may use the indexing A_y = i0, and X_, = -if for 1....,n/2 In this notation, the indices and take values in {+1,..., tm].
We notice that the eigenvalues of the Jacobian can be grouped two by two:
W4jthU-sーk exp i6, exp i0k
exp 100--16 i00 exp i0
i0, i0
-i0, + i0,
exp i(0; 0x) exp -i(0 0%)
i(0, - 02) (--i)(0; - 0x)
2 cos(f)) (0) - 0k)
(0, 02)
exp 10, exp - i00 exp -i0 exp i0
W4j-kM-j4K
i0 i0
-i0, i0
2 exp i(0, + 0%) exp -i(0, 0%)
i(0, + 0k) (-i) + 0k)
2 2 cos(() + 0k)
(0, + 02)²
Using the fact that 2 2 cos Φ = 4sin z we rewrite both terms above in the form sinc
yielding the desired result. We point out that further simplificatin can be obtained by recognizing symmetry (Vee = vts)
and noticing that the diagonal terms cancel out because Uostsk-i-i 1. Calculating the (log) determinant. with a minimum number of operations can be achieved by considering only the upper triangle.
In the case of odd n, there is an additional unpaired zero eigenvalue which we index with s 0. Then, we have to consider the additional entries:
exp 0 =
exp de -
Po:
V40
d
Using a similar calculation above we get that P04:70-1 V420V-10 sinc" (%) which for brevity we can absorb by extending the range of j in the product.
27

---

Published in Transactions on Machine Learning Research (12/2023)
Numerical details. By contimuity, sinc(0) 1. For stable computation directly in log-domain, use
log (inc(TT)) log [ 2) log [(1 + 2).
(41)
In general, just like in the spherical case, the log-determinant go that within the injectivity domain of (n),(n), we have [6] < n, and so all the sinc terms in Eq. (40) take arguments in (- and are therefore finite.
However, Proposition 4 implies
Computing the rotation angles. Learning with a wrapped distribution over SO(n) requires evaluating the determinant of the Jacobian (which requires the eigenvalues of A) as well as calculating the matrix exponential itself. This in contrast to applications where just the matrix exponential needed, where truncated series approximations are widely applicable. Obtaining an cigendccomposition of would allow us to compule both quantities needed; however, complex decompositions are more expensive. The following propositions shows how to efficiently use SVD of A to obtain both the rotation angles 0 and the value of the matrix exponential.
Proposition (Computation). Let Skew(n) and USV T be a singular value decomposition. Then,
exp(A) = U cos(S)UT + +U sin(SSV..
Proof. Assume for now that n even. From Lemma 4, item (vi), we have
exp(A) = Σ Qs exp(B;) where Q2 = 9+9 ) and exp(B,) (-- cos 0 sin 0 sin 0 cos 02 )
(42)
9-3
Explicitly calculating one such rank-2 term gives:
2, cxp( = =os (,000(00=994 + 9-
+ { 9-19-5)
From Lemma item (v), one possible SVD is given by the triplets (0), 9+ (0,, 9-j. (-s-)))-----). A straight- forward calculation shows that the part of U cos(S)UT U sin(S)V corresponding to the two copics of 0j is exactly Qj exp(B;)QJ. While SVD is only unique up to simultaneous sign flip of singular vector pair (y, v), the rank-one terms are invariant, since (-u)(-u)T uu." and (-u)(-u)T wo Therefore, the desired expression holds for any SVD. If i odd, there an additional row go the canonical decomposition, with a corresponding 1-d block B, (0) exp(B,) (1) and exp(A) has an additional rank-1 1 term 90% In this case, complete SVD of A has an unpaired zero singular triplet (0p = 0, t + Regardless of the sign choice, since sin(0) 0 and cos(O) = 1, we get an additional term 9o% as expected. Other than the sign ambiguity, there is no ambiguity about the direction of the singular vector corresponding to the unpaired singular value, since must be orthogonal to the 2m distinct Q+ vectors, even if some of them may correspond other zero singular values.
E
Algorithms for S
Here we provide the algorithm for the parallel transport for the hypersphere.
Algorithm 1: Parallel transport for gn-1 embedded in R...
Input: p. e RR, llelly || = 1; x * -P; the point the T,
z' (p+
== T
11 z defines reflection plane R utp
l|+pp]
return u 2(u, 2)2 (volated vector)
28


Reviews:
[{"review_id": "ZrgZiRiGIV", "reviewer": "Reviewer_L93q", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n* The paper makes a very clear contribution, which is practically useful for manifold statistics.\n* The paper is well-written and generally easy to follow.\n* The paper is fairly complete (I do not miss anything significant from the paper as is).\n* References are generally appropriate.\n\n## Weaknesses\n* In general, I miss illustrative plots. Specifically, I would like to see plots of (to guide intuitions)\n  * The $\\beta$-Gaussian density for different values of $\\beta$.\n  * The wrapped $\\beta$-Gaussian on a sphere for different values of $\\beta$.\n  * The functions $\\exp_{\\beta}$ and $\\log_{\\beta}$\n* I miss a discussion of the limitations of the model. At least, I see that the model is quite limited on manifolds with a small injectivity radius. Also, distributions defined in tangent spaces often end up behaving quite differently when mapped to the manifold than one would expect (the tangent space can give a highly distorted view of the manifold).\n\n## Minor comments\n* On page 3, Chevallier & Guigui are cited as the creators of 'The wrapped approach', which strikes me as misleading. That's a much older strategy; I think the book by Mardia & Jupp discusses wrapped Gaussians, for example.\n* 'FY-losses' does not appear to be defined (I can guess the meaning, but it should be defined nonetheless).\n* In Eq. 2, the notation $[\\cdot]_+$ is not defined. I assume it's the usual hinge $\\max(\\cdot, 0)$, but it should be defined.\n\n## Questions\n* Is the wrapped $\\beta$-Gaussian density on the sphere continuous and/or differentiable at the antipodal point?\n* As far as I could tell, the wrapped $\\beta$-Gaussian is used with VAEs to give the latent space a pre-specified geometric structure. Is it easy to use the same distribution as the conditional likelihood $p(x|z)$?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper is concerned with tractable distributions on Riemannian manifolds. The main idea is to consider $\\beta$-Gaussians, which are Euclidean distributions with compact support. These distributions can then be placed in the tangent space of a manifold. As long as the compact distribution is fully supported within the injectivity radius, you avoid all the difficulties of wrapping on manifolds.\n\nThis is rather nice, I think.\n\nThis is one of those papers that are obvious if you know both about $\\beta$-Gaussians and manifold statistics. In practice, most readers will not be familiar with both (I was not familiar with $\\beta$-Gaussians, for example), and the paper provides a very nice contribution to such readers.\n\nAll in all, I recommend acceptance of the manuscript subject to the requested changes being made.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It would be good to fix the weaknesses above (should all be easily done).", "broader_impact_concerns": "The paper is on the 'theory side', so I do not see any ethical issues.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XF3uX6qymp", "reviewer": "Reviewer_NTsP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The exposition and organization is clear, with precise statements, intuition explained, and proofs and details pushed to appendices.\n- Much work has been put into coming up with usable algorithms in the sphere and rotation manifold cases, ensuring that distribution supports are within the injectivity radius and incorporating the nontrivial exponential/logarithm/parallel transport maps.\n\nWeaknesses:\n- The method perhaps overclaims in its generality, as utility of the method requires analytical exponential, logarithm, and parallel transport maps. This is a small class of manifolds, as illustrated with just two example (classes) in their experiments.\n- It could be argued that the experimental results are relatively weak. \n   - The first experiment on moment matching is merely a proof-of-concept, and one would typically use simple sampling to determine these quantities. \n   - The second experiment is quite anecdotal, and there is not commentary on whether this method improves upon previous embeddings into Euclidean space.\n   - The third experiment only outperforms the baseline in lower dimensions, and the increase in performance seems relatively small.\n   - The fourth experiment again does slightly worse than the baseline, so merely serves as an example that the basic idea works for $SO(n)$. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper considers wrapped $\\beta$-Gaussians on embedded spheres and rotation matrices. This generalizes a compactly-supported generalization of Gaussians to these manifolds, allowing for efficient sampling and exact density evaluation. Furthermore, they provide Fenchel-Young losses, which are $\\beta$-generalized analogues of the KL divergence and cross-entropy loss. Proof-of-principle experiments are done with moment matching, hierarchy modelling, VAEs, and multilingual embedding alignment.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I feel the work is just above the bar, due to its careful theoretical contributions on (narrow-ish) domains of interest. The work would be significantly strengthened by addressing the two weaknesses I mentioned above. The first may be dealt with via some rephrasing and limitation acknowlegement in the introduction. The second is more challenging, as it involves additional experiments, and perhaps example applications that are entirely novel. Without a clear application that shows the advantages of using these wrapped $\\beta$-Gaussians, it becomes less likely that others in the community will consider them for use.", "broader_impact_concerns": "I do not have any concerns in this direction.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "sYgTz3u0Ik", "reviewer": "Reviewer_LAzL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**:\n\n1. Originality: A family of wrapped $\\beta$-Gaussian distributions is a novel distribution family on Riemannian manifolds which provides exact density evaluation and reparametrized sampling in high dimensions.   \n\n2. Significance: The proposed family of distributions is efficient for sampling, parameter learning and density evaluation in comparison with other distributions on Riemannian manifolds.\n\n3. Quality: Theoretical results are solid with rigorous proof but I do not spend time double-checking all of them. Moreover, applications of wrapped $\\beta$-Gaussian distributions in parameter estimation, generative modeling (variational autoencoder) and multilingual word embedding alignment.\n\n4. Clarity: The paper is well organized and written.\n\n**Weakness**: \n\n1. (Statistical) bounds for the difference between parameter estimates in equations (17) and (18) and their true values are missing.\n    \\item A discussion on optimization of the Fenchel-Young loss function in equation (14) is missing.\n\n 2.  Some notations have not been defined thoroughly (see Requested Changes).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors propose a novel family of distributions defined on Riemannian manifolds named wrapped $\\beta$-Gaussians, which can be seen as an extension $\\beta$-Gaussian family in the literature. Additionally, for the sake of fitting the distribution parameters to data, since gradient-based methods do not work in this scenario, they construct a wrapped Fenchel-Young loss function based on the learning objective function for $\\beta$-Gaussian distribution. Next, they introduce an efficient embedded parametrization of scale matrices for the case of hyperspheres. Finally, they conducted several experiments to demonstrate the effectiveness of the wrapped $\\beta$-Gaussians on many tasks, namely data fitting, generative modeling and multilingual word embedding alignment.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Regarding Proposition 1, the authors should include the formulation of $A_{\\beta}(\\Sigma)$ in the main text. Additionally, compared to Gaussian distributions in the Euclidean space, the appearance of the term $A_{\\beta}(\\Sigma)$ in the exponential operator is new, so could the authors discuss more about this term? \n\n2. Given the density formulation in Proposition 1, are we able to compute explicitly the corresponding moment generating function? If not, then what are the difficulties?  \n\n3. In equation (26), since the term $A_{\\beta}(\\Sigma)$ is invalid when $\\beta=1$, I wonder how the authors evaluate the density function given in equation (4) when $\\beta=1$. \n\n4. In Appendix B.3, the authors compute the discrepancy between two $\\beta$-Gaussian distributions of the same dimension under the Wasserstein distance. I wonder if we still have a closed-form expression for that of two $\\beta$-Gaussian distributions defined on Riemannian manifolds of different dimensions.\n\n5. Under Figure 1, the notations $\\mathbb{S}^n$ and $SO(n)$ are not defined. Please add their definitions before using or call their names such as hyperspheres and special orthogonal manifolds.\n\n6. Typos: (1) In the second paragraph of the introduction section, 'equivalent' should be 'equivalence'. (2) In Appendix B.3, 'analitical' should be 'analytical', and $WN_{\\beta}(\\mu_1,\\Sigma_2)$ should be $\\mathcal{WN}_{\\beta}(\\mu_1,\\Sigma_1)$.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The support for the paper was strong overall, and the authors were satisfied with the rebuttal.
Decision: Accept (poster)