A Computer Scientist's Guide to Riemannian Manifold
This blog post is still under construction. I will finish the 3D rotation group section soon.
Preface
A lot of my CS friends have asked me about various Riemannian manifolds in generative modeling, especially their geometric properties that induce important concepts like exponential map and logarithm map utilized in Riemannian flow matching
Generally, I will try to avoid using too much mathematical jargon and focus on the intuition behind the concept. Nonetheless, for those who are interested in the mathematical details (like me!), I include them in the collapsible. Feel free to skip them as the main text is designed to be self-contained. As a self-taught non-mathematician in Riemannian geometry, it will be inevitable that I might make some mistake. Please email me if you find any errors or have any suggestions.
Riemannian Manifold
A manifold \(\mathcal{M}\) is a space that locally looks like a flat Euclidean space everywhere. Imagine you are a tiny ant on the manifold, you will feel like you are living in a flat world, like our human being on the spherical Earth. For a smooth manifold, the tangent space \(T_x\mathcal{M}\) at a specific point \(x\in\mathcal{M}\) can be understood as the vector space of all possible directions (i.e., tangent vectors) one can move from \(x\). The tangent bundle \(T\mathcal{M}:=\bigcup_{x\in \mathcal{M}}T_x\mathcal{M}\) is defined as the union of all tangent spaces at each point. See for an illustration of the tangent space.
A Riemannian manifold \(\mathcal{M}\) is a smooth manifold equipped with a Riemannian metric \(g\), which is a smooth positive-definite inner product on the tangent space of each point: \(\begin{equation} g_x: T_x\mathcal{M}\times T_x\mathcal{M}\to \mathbb{R}. \end{equation}\) As \(g\) is an inner product, we also write \(\langle u,v\rangle_g:=g(u,v)\) and sometimes omit the subscript \(x\) if it can be inferred from the context. However, it is important to keep in mind that the Riemannian metric depends on the point. The Riemannian norm of a tangent vector \(u\in T_x\mathcal{M}\) is defined as \(\|u\|_g:=\sqrt{\langle u,u\rangle_g}\). In this way, the Riemannian metric allows us to define the length of tangent vectors, which is essential for defining the distance between two points on the manifold. Recall from your calculus class for defining the infinitesimal arc length of a smooth curve as \(\mathrm{d}s=\|f^\prime(t)\|_2\,\mathrm{d}t\) for a smooth curve \(f:[0,T]\to\mathbb{R}^n\), the arc length of a smooth curve \(\gamma(t):[0,T]\to\mathcal{M}\) on the Riemannian manifold can be calculated as: \(\begin{equation} \mathrm{d}s=\|\dot\gamma(t)\|_g\,\mathrm{d}t=\sqrt{\langle\dot\gamma(t),\dot\gamma(t)\rangle_g}\,\mathrm{d}t. \end{equation}\) Here, we follow the convention to use Newton’s notation to denote derivatives with respect to \(t\).
Geodesic, Exponential Map & Logarithm Map
In Euclidean geometry, there is one specific family of curves that are crucial for establishing the standard geometry: the straight lines. In Riemannian manifold, we naturally want to extend the idea of straight lines to curved spaces. So again, imagine you are a tiny ant on the manifold. As the manifold is locally flat, you can follow the direction of some tangent vector and go straight while maintaining the same direction and constant speed on the manifold. The route you have traced is a straight line on the manifold, which is called a geodesic.
One fundamental results for geodesics is the local existence and uniqueness theorem, which states that for any point \(x\in\mathcal{M}\) and any tangent vector \(u\in T_x\mathcal{M}\), there exists a unique geodesic \(\gamma:[0,T]\to\mathcal{M}\) such that \(\gamma(0)=x\) and \(\dot\gamma(0)=u\). There are some interesting properties of geodesics as the straight lines on the Riemannian manifold. One of the most mentioned properties is that the geodesic is the locally distance minimizing curve. This can be understood as the generalization of the well-known result from elementary school that, in the Euclidean space, the shortest path between two points is the straight line segment. This property can be also understood as the kinetic energy minimizing curve, with the energy functional defined as: \(\begin{equation} E(\gamma)=\frac{1}{2}\int_0^Tg_{\gamma(t)}(\dot\gamma(t),\dot\gamma(t))\,\mathrm{d}t. \end{equation}\) Using the classical calculus of variations techniques, one can derive the geodesic equation for explicitly solving the geodesic \(\gamma(t)\), which we will omit here. The geodesic distance \(d_g(x,y)\) between two points \(x,y\in\mathcal{M}\) is defined as the length of the globally length-minimizing curve that connects \(x\) and \(y\). As the globally optimum always ensures the local optimum, the geodesic distance is always the curve length of some geodesic. However, the inverse is not necessarily true, i.e., the curve length of a geodesic connecting two points is not necessarily the geodesic distance.
Length-minimizing property
One common misunderstanding is that the geodesic is the globally shortest path between any two points on the manifold. This is not true in general. For example, consider the sphere whose geodesics are the great circles. For a pair of non-antipodal points, the great circle is cut into a shorter arc and a longer arc, both of which are geodesics. However, only the shorter path is the globally shortest path. Another example is provided below. Nonetheless, globally optimum is always guaranteed for geodesic distance as a result of its definition.
More general definition
Contrary to other standard approach that deduce the geodesic from the locally length-minimizing property, I adopt the definition that relies on a dynamics in which the point travels at a constant speed. It turns out that the definition of constant tangent vectors is non-trivial, as the tangent spaces are different for different points. Eventually, this idea of defining constant tangent vectors leads to the concepts of parallel transport, affine connection, and the general definition of affine geodesic. For general affine geodesics, the locally length-minimizing property no longer holds. In contrast, the standard locally length-minimizing geodesic is a special case of the affine geodesic, which is known as the Levi-Civita geodesic. The Levi-Civita geometry is uniquely defined by the Riemannian metric and is therefore also referred to as the canonical Riemannian structure.The exponential map \(\exp:\mathcal{M}\times T\mathcal{M}\to \mathcal{M}\) answers the questions that, after traveling for a unit time with a constant direction and speed, what point will you reach on the manifold. Formally, for a point \(x\in\mathcal{M}\) and a tangent vector \(u\in T_x\mathcal{M}\), the exponential map is defined as: \(\begin{equation} \exp_x(u):=\gamma(1), \end{equation}\) where \(\gamma\) is the unique geodesic that satisfies \(\gamma(0)=x\) and \(\dot\gamma(0)=u\). The logarithm map \(\log:\mathcal{M}\times\mathcal{M}\to T\mathcal{M}\) is the inverse of the exponential map, which answers the question that, given two points on the manifold, what is the direction and distance one need to travel to reach one point from the other. Formally, for two points \(x,y\in\mathcal{M}\), the logarithm map is defined such that the following identity always holds: \(\begin{equation} \exp_x(\log_x y)\equiv y . \end{equation}\)
Based on the idea of constant speed dynamics, we have \(\|\dot\gamma(t)\|_g\equiv\|\dot\gamma(0)\|_g=\|u\|_g\). In this scenario, we have \(\begin{equation} \|\log_x (y)\|_g=L(\gamma_{x,y}), \end{equation}\) where \(L(\gamma_{x,y})\) is the length of the geodesic \(\gamma_{x,y}\) that connects \(x\) and \(y\). If the uniqueness of the geodesic is globally guaranteed, we further have \(\|\log_x (y)\|_g=d_g(x,y)\). Simiarly, based on the idea of dynamic, the geodesic interpolation between two points \(x,y\in\mathcal{M}\) is defined as \(\gamma_{x,y}(t)\) for \(t\in[0,1]\). Using the exponential and logarithm map, the geodesic interpolation can be compactly written as: \(\begin{equation} \gamma(t)=\exp_x(t\log_x y). \end{equation}\)
When does the assumption hold?
For the canonical Riemannian structure induced by the Levi-Civita connection, the assumption $\|\dot\gamma(t)\|_g\equiv\|\dot\gamma(0)\|_g$ always holds. For general affine connection, it is not guaranteed. In most cases, we work with the canonical Riemannian structure in this post.Some caveats
- There might be scenarios when $T<1$ for the existence and uniqueness theorem to hold. This means that the exponential map may not be well-defined for all tangent vectors. Nonetheless, we can always scale the tangent vector accordingly for a larger $T$.
- The global uniqueness of geodesic between any two points on the manifold is not guaranteed. For example, consider a pair of antipodal points on a sphere, there are infinite many great circles that have the same length. In this scenario, the logarithm map is not well-defined.
The Riemannian norm and geodesic interpolation are the crucial concepts in Riemannian flow matching, in which the probability paths are defined along the geodesics. From the next section, I will dig into some commonly encountered Riemannian manifolds in machine learning tasks.
Euclidean Space
Of course, as Riemannian manifold is modeling as a generalization of the Euclidean space \(\mathbb{R}^n\), the first thing we will do it naturally test its compatibility with our common intuition in the latter. The Riemannian metric is the standard Euclidean inner product: \(\begin{equation} g_x(u,v)=\langle u,v\rangle=\sum_{i=1}^n u_i v_i. \end{equation}\) Similarly, the Riemannian norm is the standard Euclidean norm \(\|\cdot\|_2\). As we can travel in the Euclidean space without obstacles, the exponential map on the Euclidean space simple translate the base vector with the tangent vector. The exponential and logarithm maps, therefore, can be easily computed as: \(\begin{align} \exp_x(u)&=x+u,\\ \log_x(y)&=y-x. \end{align}\) Similarly, the geodesic interpolation is the linear interpolation between two points: \(\begin{equation} \gamma(t)=tx+(1-t)y. \end{equation}\)
Wait ... How about the exponential function?
If you came up with this question, then probably you are a math enthusiast! The exponential map reads a linear function of $u$, whereas the standard exponential function reads $\exp u=e^u$. Wait! We've got a conflict even on the Euclidean space!
It turns out that we need to consider an alternative Riemannian manifold to resolve this conflict. Consider the manifold of positive real numbers $\mathbb{R}_+$. At each point $x\in\mathbb{R}_+$, the tangent space is identified with $T_x\mathbb{R}_+\cong\mathbb{R}$. We can use the Riemannian metric: $$ g_x(u,v)=\frac{uv}{x^2},\quad x\in \mathbb{R}_+,u,v\in T_x\mathbb{R}_+. $$ We now verify that the following geodesic $$ \gamma(t)=e^{ut} $$ satisfied the constant speed requirement with the initial condition of $\gamma(0)=1,\dot\gamma(0)=u$. We first note that $\dot\gamma(t)=u e^{ut}=u\gamma(t)$, which satisfies the initial condition. Calculating the Riemannian norm, we have: $$ g_{\gamma(t)}(\dot\gamma(t),\dot\gamma(t))=g_{\gamma(t)}(u\gamma(t),u\gamma(t))=\frac{u^2\gamma(t)^2}{\gamma(t)^2}=u^2. $$ In this way, we always have $\|\dot\gamma(t)\|^2_g=u^2$, which is independent of the time $t$. Therefore, the curve $\gamma(t)$ is indeed a geodesic. Setting $t=1$ recovers the exponential map as $\exp_1(u)=e^u$, which coincides with the standard exponential function.
Why this metric?
You probably wonder why we are using such an exotic Riemannian metric on $\mathbb{R}_+$. It turns out that, as $\mathbb{R}_+$ is a Lie group — a differentiable manifold that is also a group, — there is a canonical way of choosing the metric such that it should be bi-invariant. We know the group structure of $\mathbb{R}_+$ is well-defined with the standard multiplication, which means the bi-invariant metric on $\mathbb{R}_+$ is invariant to scaling: $$ g_{\lambda x}(\lambda u,\lambda v)=g_x(u,v),\quad x\in\mathbb{R}_+,u,v\in T_x\mathbb{R}_+. $$ It is easy to verify that the metric $g_x(u,v)=\frac{uv}{x^2}$ satisfies the bi-invariant property, whereas the standard Euclidean metric does not. In this case, we are essentially using the logarithmic coordinates to represent the positive real numbers. Therefore, at each point $x\in\mathbb{R}_+$, the local coordinate in the target space should be $u/x$, which represents the relative scaling under the group operation of multiplication and also leads to the bi-invariant metric.
This bi-invariant property is important for drawing the connection between the Lie group, the Riemannian manifold, and the Lie algebra. The exponential map $\exp:\mathbb{R}\to \mathbb{R}_+$ converts an element in the Lie algebra to the Lie group, and the logarithm map $\log:\mathbb{R}_+\to \mathbb{R}$ converts an element in the Lie group to its Lie algebra. We will also see this important relationship in the 3D rotation group, another well-known Lie group.
Alternative interpretation
We can also provide an alternative interpretation to resolve the conflict between the exponential map and the exponential function by considering general affine geodesics. Recall that affine geodesics are defined by the dynamics in which the point travels at a constant speed. Let's fix our start point to $x=1$ and consider the tangent vector $u\in\mathbb{R}$. Now consider the following two definitions of constant speed:
- At each time step, the speed should also be $u$.
- At each time step, the speed should be proportional to the current position.
In the second definition, the speed of traveling is $u\gamma(t)$ to be considered constant, — weird but valid! The first definition gives the following differential equation: $$ \dot\gamma(t)=u,\quad \gamma(0)=1, \dot\gamma(0)=u, $$ which gives the solution $\gamma(t)=1+ut$ and the exponential map of $\exp(u)=1+u$. The second definition gives the following differential equation: $$ \dot\gamma(t)=u\gamma(t),\quad \gamma(0)=1, \dot\gamma(0)=u, $$ which gives the solution $\gamma(t)=e^{ut}$ and the exponential map of $\exp(u)=e^u$. This indicates that the exponential function can also interpret as the exponential map of some affine geodesic.
Flat Torus
We now continue to explore the flat torus \(\mathbb{T}^n\). For simplicity and better visualization, let’s first consider \(n=2\), which is often visualized as a doughnut embedded in the 3D Euclidean space. This, however, is actually not an accurate representation when we consider their geometric properties. As its name indicated, a flat torus have zero Gaussian curvature everywhere, meaning it is essentially a flat plane without any bending or stretching. In contrast, the embedded 2D torus has non-zero curvatures at different points, which can be verified by gluing opposite edges of a square paper. One can easily roll a flat paper into a cylinder which also has zero curvature, but it is impossible to further roll the cylinder into a doughnut without stretching or bending the paper.
Therefore, it is better to understand a flat torus as a quotient space of the Euclidean space \(\mathbb{R}^n\) by the lattice \(\mathbb{Z}^n\) as \(\mathbb{T}^n=\mathbb{R}^n/\mathbb{Z}^n\). In other words, in the 2D case, we consider the following equivalence class by identifying points that differ by integer translations in two independent directions: \(\begin{equation} (x,y)\sim (x+m,y+n),\quad m,n\in\mathbb{Z}. \end{equation}\) Such an identification can be understood visually as gluing the opposite sides of a square with the same direction, as demonstrated in . A better example to understand the flat torus is the crystal structure. In crystallography, the crystal structure is often represented as a unit cell, the smallest repeating unit. With periodic boundary conditions, the unit cell can be extended to the whole crystal structure, which is essentially a 3D flat torus (see ).
It is also possible to define the quotient space with respect to other values, such as the cell length of a unit cell, or \(2\pi\) for a torsion angle as a rotation of \(2\pi\) is equivalent to no rotation. Nonetheless, it is simpler to use fractional coordinates such that any coordinate always lies in \([0, 1)\). We will use this assumption in the following discussion.
About the different tori
We distinguish the flat torus from the embedded torus because we care about their geometric properties in differential geometry. If we consider topological properties, these two tori are indeed equivalent (more precisely, they are homeomorphic). Homeomorphism allows us to bend and stretch manifolds when drawing equivalence.One choice to identify the tangent space of the flat torus is also to limit each component of the tangent vector to \([0, 1)\). In this way, we have \(T_x\mathbb{T}^n\cong \mathbb{T}^n\). As a quotient space of the Euclidean space, the flat torus inherits the Euclidean metric and Euclidean norm. Therefore, the common straight lines are geodesics on the flat torus. The exponential and logarithm maps can be easily computed as: \(\begin{align} \exp_x(u)&=x+u \mod 1,\\ \log_x(y)&=y-x \mod 1, \end{align}\) which is essentially the same as the Euclidean space except for the additional modulo operation. Simiarly, the geodesic interpolation only differs with the additional modulo. However, we should be careful when computing the Riemannian norm in the tangent space. Consider the target vector field of \(u=0.1\). For two predictions \(v_1=0.5,v_2=0.9\), it turns out that \(v_2\) is closer to the target vector due to the periodicity in the exponential map. We shall follow the minimum image convention to select the shortest path between two points. In this way, the Riemannian norm should be calculated as: \(\begin{align} \|u\|_\mathbb{T}=\left\|\min \{u,1-u\}\right\|_2,\quad u\in T\mathbb{T}^n=[0,1)^n. \label{eqn:torus-norm} \end{align}\) where the minimum operation is element-wise.
Alternative parameterization
There exists an alternative parameterization of the flat torus that models the tangent vector as angles $u\in(-\pi,\pi]^n$. Essentially, this parameterization try to find the direction of the straight line that connects $x$ to $y$. The exponential and logarithm maps can be computed as: $$ \begin{align*} \exp_x(u)&=x+u/(2\pi)\mod 1,\\ \log_x(y)&=\operatorname{atan2}(\sin 2\pi(y-x), \cos 2\pi(y-x)), \end{align*} $$ where $\mathrm{atan2}$ is the 2-argument arctangent function that automatically wraps the direction in the Cartesian plane to $(-\pi,\pi]$. Such a parameterization makes sure that the vector field returned by the logarithm map follows the minimum image convention, i.e., it always represents the shortest direction between two points. Still, we need to follow the similar procedure to wrap the difference of the angles to $(-\pi,\pi]$ when calculating the norm.
Spherical Manifold
The previous two manifolds are flat manifolds with relatively intuitive and simple geometries. We now move to the spherical manifold \(S^n=\{x\in\mathbb{R}^{n+1}:\|x\|_2=1\}\). \(S^n\) is an \(n\)-dimensional manifold embedded in the \((n+1)\)-dimensional Euclidean space. As its name indicates, we often visualize \(S^2\) as a sphere, the surface of a 3D ball (see ).
Intuitively, the tangent space at a point on the sphere \(x\in S^n\) can be identified as the hyperplane perpendicular to the vector \(x\): \(T_xS^n=\{u:\langle u,x\rangle=0\}\). As an embedded manifold, \(S^n\) inherits the ambient Euclidean metric. Therefore, the spherical metric is exactly the same as the Euclidean one: \(g_x(u,v)=\langle u,v\rangle.\)
By definition, the geodesic distance between two points \(x,y\in S^n\) is the length of the shortest curve connecting them along the sphere. This curve follows the shorter arc of a great circle, obtained by intersecting the sphere with a 2D plane through its center (see ). As the sphere has unit radius, the arc length is exactly the angle between \(x\) and \(y\), measured in radians: \(\begin{equation} d_S(x,y)=\theta:=\arccos\langle x,y\rangle,\quad \theta\in[0,\pi]. \end{equation}\)
For example, two perpendicular unit vectors are separated by a distance of \(\pi/2\), while two opposite points are separated by \(\pi\). In the latter case, there are multiple shortest arcs, all with the same length. The exponential map moves along a great circle in the direction of a tangent vector \(u\in T_xS^n\), which can be obtained by 2D geometry. We have: \(\begin{equation} \exp_x(u)=x\cos\|u\|_2+\frac{\sin \|u\|_2}{\|u\|_2}u. \end{equation}\) As before, the length of \(u\) determines how far we travel in a unit time; on the unit sphere, it also determines the angle through which we rotate.
Obtaining the exponential map with 2D geometry
For any dimension \(n\geq 2\), we can focus on the 2D plane spanned by \(x\) and a nonzero tangent vector \(u\). Since \(\langle x,u\rangle=0\), the vectors \(x\) and \(v:=u/\|u\|_2\) form an orthonormal basis of this plane. Its intersection with the sphere is a unit circle, and the starting point \(x\) has coordinates \((1,0)\) in this basis.
Now recall that an angle of one radian sweeps out an arc of length one on a unit circle. To travel at speed \(\theta=\|u\|_2\), we therefore increase the angle at rate \(\theta\). At time \(t\), the point has 2D coordinates \((\cos(t\theta),\sin(t\theta))\).
To obtain its coordinates in the original Euclidean space \(\mathbb{R}^{n+1}\), we resolve these two components along the basis vectors \(x\) and \(v\):
\[\begin{equation} \begin{aligned} \gamma(t)&=x\cos(t\theta)+\sin(t\theta)v\\ &=x\cos(t\theta)+\sin(t\theta)\frac{u}{\|u\|_2}. \end{aligned} \end{equation}\]Evaluating at \(t=1\) gives the exponential map above. This expression remains valid even when \(\theta>\pi\), although the arc we travel then goes beyond half a circle and is no longer a shortest path to its endpoint.
Conversely, the logarithm map gives the initial tangent vector that reaches \(y\) along the shortest arc from \(x\), which is essentially finding a perpendicular direction to in the 2D plane. Its expression is:
\(\begin{equation} \log_x(y)=\frac{\theta}{\sin\theta}(y-\langle x,y\rangle x)=\frac{\arccos\langle x,y\rangle}{\sqrt{1-\langle x,y\rangle^2}}(y-\langle x,y\rangle x), \end{equation}\) where \(\theta=\arccos\langle x,y\rangle\in(0,\pi)\).
Obtaining the logarithm map with 2D geometry
In the 2D intersection, it is easy to see that the vector \(y-\langle x,y\rangle x\) is the projection of \(y\) onto the tangent space at \(x\). It points toward \(y\) along the great circle and has length \(\sin\theta\). Dividing by this length gives a unit direction, and multiplying by \(\theta\) gives the required travel distance. This choice of the logarithm is called the principal logarithm, as it selects the shortest arc.
In practice, we also need to handle different numerical stability issues when we implement the spherical manifold with PyTorch (see the discussions below). A detailed standalone implementation can be found in my codebase.
Numerical stability
The following numerical stability must be handled properly when implemented with PyTorch:
- Exponential map has singularity at \(u\approx0\). Use
torch.sinc(theta / torch.pi)(torch uses normalized sinc). - Logarithm map has singularity at \(\theta\approx0\). Use the Taylor expansion \(\theta/\sin\theta\approx 1 + \theta^2/6\) if needed.
- Logarithm map has infinite many possible values for antipodal points: two points sitting at the two poles. Usually choosing any one of them is fine.
- Arccos (in distance calculation and logarithm map) has gradient singularity around 1 and -1.