Flow matching on the sphere (Riemannian flow matching)

Before this lesson: flow matching in $\R^d$: the marginalization trick, the CFM loss and Euler sampling (lesson 14); coordinates on the sphere (lesson 09); matrix manifolds, tangent spaces and projections (linear algebra 19). Time: about 80 minutes, plus the notebook.

---

Why this matters #

• Events on the Earth. Earthquake epicentres, volcanic eruptions, floods and wildfires are points on the sphere. Chen and Lipman (ICLR 2024) test Riemannian flow matching on exactly these four Earth-science catalogues (compiled by Mathieu and Nickel, 2020). A density model on the sphere gives hazard maps and synthetic catalogues.
• Other curved spaces. Wind directions live on a circle, protein backbone angles on a torus, and the orientation of a rigid body (or of a fault plane with its slip direction) in the rotation group $SO(3)$. The same tools (exp, log, tangent projections, retractions) build flow models on all of them.
• No new training cost. On the sphere, geodesics have a closed form, so training stays simulation-free, as in lesson 14. In flat space the method is exactly the CFM of lesson 14.

Do not confuse this lesson with lesson 09. Lesson 09 is about fields on the sphere (a value at every point, such as temperature). This lesson is about probability distributions of points on the sphere (where events happen).

1. Predict first #

2. The picture #

The simplest case: the circle #

A point on the circle is an angle $\varphi$, and a velocity is one number. The shortest way from $\varphi_0$ to $\varphi_1$ is the shorter arc, of length $\lvert\operatorname{wrap}(\varphi_1 - \varphi_0)\rvert$, where $\operatorname{wrap}$ maps an angle to $(-\pi, \pi]$. Moving with velocity $v$ for unit time gives $\varphi + v$, and the velocity that reaches $\varphi_1$ in unit time is $\operatorname{wrap}(\varphi_1 - \varphi_0)$. So flow matching on the circle is the flow matching of lesson 14 with one change: the difference $x_1 - x_0$ becomes the wrapped difference.

At one point, $\varphi_1 = \varphi_0 + \pi$, the two arcs are equally short and the wrapped difference jumps from $-\pi$ to $+\pi$. That point is the trouble spot on every sphere.

The sphere: tangent planes and great circles #

A point is a unit vector $x \in \R^3$. The velocities of curves through $x$ form the tangent plane: the vectors orthogonal to $x$. A straight line $x + sv$ leaves the sphere at once. The "straight lines" of the sphere are the great circles: walk a distance $s$ from $x$ in the unit tangent direction $\hat w$ and you arrive at $\cos(s)\,x + \sin(s)\,\hat w$.

Interactive figure: loading…

The recipe #

Riemannian flow matching is lesson 14 with three replacements:

Lesson 14 ($\R^d$)	On the sphere
$x_t = (1 - t)\,x_0 + t\,x_1$	$x_t = \exp_{x_0}(t\log_{x_0}x_1)$, the great-circle path
target velocity $x_1 - x_0$	$\log_{x_t}(x_1)/(1 - t)$, tangent at $x_t$
Euler step $x + hv$	$\exp_x(hv)$, or $(x + hv)/\lVert x + hv\rVert$

and one constraint: the network's velocity must be tangent, so we project its output onto the tangent plane.

3. Definitions and main results #

The sphere as a Riemannian manifold #

• $S^2 = \{x \in \R^3 : \lVert x\rVert = 1\}$. In this track (lesson 09), a point with colatitude $\theta \in [0, \pi]$ and longitude $\phi$ is $x = (\sin\theta\cos\phi,\ \sin\theta\sin\phi,\ \cos\theta)$. With latitude $90° - \theta$, this is the usual map coordinate.
• Tangent plane. For a curve $c(t)$ on $S^2$, differentiate $c \cdot c = 1$: $2\,c \cdot \dot c = 0$. So $T_xS^2 = \{v \in \R^3 : x \cdot v = 0\}$, and the orthogonal projection onto it is $P_x = I - xx^\top$, that is, $P_xw = w - (x \cdot w)\,x$.
• Metric. The inner product of two tangent vectors is their dot product in $\R^3$ (the sphere is a Riemannian submanifold of $\R^3$). The length of a curve is $\int\lVert\dot c\rVert\,dt$. The geodesic distance is the length of the shortest curve: $d(x, y) = \arccos(x \cdot y)$, the angle between $x$ and $y$ (times the radius, 6371 km, on the Earth).
• Geodesics are curves with no tangential acceleration. For a unit $x$ and a tangent $v \ne 0$, $c(t) = \cos(t\lVert v\rVert)\,x + \sin(t\lVert v\rVert)\,v/\lVert v\rVert$ has $\ddot c = -\lVert v\rVert^2 c$. This is normal to the sphere, so $P_c\,\ddot c = 0$: great circles at constant speed are the geodesics.
• The exponential map $\exp_x(v)$ is the point reached at time 1 by the geodesic that starts at $x$ with velocity $v$. The logarithm map $\log_x(y)$ is its inverse: the shortest tangent vector $v$ with $\exp_x(v) = y$. In $\R^d$: $\exp_x(v) = x + v$ and $\log_x(y) = y - x$.

Theorem 1 (exponential and logarithm maps of the sphere) #

For $x \in S^2$ and $v \in T_xS^2$,

For $y \in S^2$ with $y \ne -x$, let $\omega = \arccos(x \cdot y) \in [0, \pi)$ and $u = P_xy = y - (x \cdot y)\,x$. Then

and $\exp_x(\log_xy) = y$, $\lVert\log_xy\rVert = d(x, y)$, and $\log_x(\exp_xv) = v$ whenever $\lVert v\rVert < \pi$.

1. $\exp_x(v)$ is $c(1)$ for the great circle above, with $c(0) = x$ and $\dot c(0) = v$. It has norm 1 because $x \perp v$: $\cos^2 + \sin^2 = 1$.
2. Let $y = \exp_x(v)$ with $0 < \lVert v\rVert < \pi$. Take the dot product with $x$: $x \cdot y = \cos\lVert v\rVert$. Since $\arccos$ is one-to-one on $[0, \pi]$, $\lVert v\rVert = \arccos(x \cdot y) = \omega$.
3. The tangent part of $y$ is $u = y - (x \cdot y)x = \sin\lVert v\rVert\;v/\lVert v\rVert$. Since $\sin\lVert v\rVert > 0$, the unit vectors agree: $u/\lVert u\rVert = v/\lVert v\rVert$. With step 2, $v = \omega\,u/\lVert u\rVert$, which is the formula for $\log_x(y)$, and $\lVert u\rVert = \sin\omega$.
4. Conversely, $\exp_x(\omega\,u/\lVert u\rVert) = \cos\omega\,x + \sin\omega\,u/\sin\omega = (x \cdot y)\,x + u = y$. $\blacksquare$

Where log fails. At $y = -x$, $u = 0$, and every tangent direction gives a geodesic of length $\pi$ to $-x$: $\log_x$ is not defined there. The set of points where the shortest geodesic from $x$ stops being unique is the cut locus of $x$; on the sphere it is the single point $-x$. Near it, the direction of $\log_x(y)$ jumps when $y$ moves a tiny distance.

Interactive figure: loading…

Geodesic interpolation. Put $u/\lVert u\rVert = (x_1 - \cos\omega\,x_0)/\sin\omega$ into $\exp_{x_0}(t\log_{x_0}x_1)$ and use $\sin(\omega - t\omega) = \sin\omega\cos t\omega - \cos\omega\sin t\omega$:

This is slerp (spherical linear interpolation). It runs along the great circle at the constant speed $\omega$, from $x_0$ at $t = 0$ to $x_1$ at $t = 1$.

Theorem 2 (the velocity of the geodesic path) #

For $x_1 \ne -x_0$ and $0 \le t < 1$:

Why it is true, in words. At time $t$ the point still has the distance $(1 - t)\omega$ to go along the same great circle, and it must cover it in the time $1 - t$. The log map points along the circle and has the length of the remaining distance; dividing by the remaining time gives the speed $\omega$.

Proof of Theorem 2

1. $x_t = \exp_{x_0}(tv)$ with $v = \log_{x_0}(x_1)$ is a geodesic with constant speed $\lVert v\rVert = \omega$.
2. The rest of the path, $s \mapsto x_{t + s(1 - t)}$ for $s \in [0, 1]$, is again a geodesic. It starts at $x_t$, ends at $x_1$, and its length is $(1 - t)\omega < \pi$.
3. By Theorem 1, the only tangent vector at $x_t$ of norm below $\pi$ whose geodesic reaches $x_1$ at time 1 is $\log_{x_t}(x_1)$. So the initial velocity of the curve in step 2 is $\log_{x_t}(x_1)$.
4. By the chain rule, that initial velocity is $(1 - t)\,\dot x_t$. Hence $\dot x_t = \log_{x_t}(x_1)/(1 - t)$. $\blacksquare$

This is the conditional velocity $u_t(x \mid x_1) = \log_x(x_1)/(1 - t)$ of Chen and Lipman. In $\R^d$, $\log_{x_t}(x_1)/(1 - t) = (x_1 - x_t)/(1 - t) = x_1 - x_0$: the CFM target of lesson 14. The geodesic path is the sphere's version of a straight path: one exponential-map step from $x_0$ with the velocity $\log_{x_0}(x_1)$ lands exactly on $x_1$.

Interactive figure: loading…

Theorem 3 (Riemannian conditional flow matching) #

Setup. A source $p$ on $S^2$ (for example uniform), data $q$ on $S^2$, a coupling of the two (independent by default), $t \sim U[0, 1]$, the path $x_t = \exp_{x_0}(t\log_{x_0}x_1)$, and the conditional target $u_t(x_t \mid x_1) = \log_{x_t}(x_1)/(1 - t)$. Let $p_t$ be the law of $x_t$, and $u_t(x) = \mathbb E[u_t(x_t \mid x_1) \mid x_t = x]$ the marginal velocity. Norms of tangent vectors are Euclidean norms in $\R^3$.

• Marginalization trick on manifolds (Flow Matching Guide, Theorem 10, under its regularity assumptions): $u_t$ generates $p_t$, so the ODE $\dot x = u_t(x)$ started from $p$ ends at $q$. The proof is Theorem 1 of lesson 14 with test functions on the sphere.
• RCFM (Flow Matching Guide, Theorem 11): for a velocity network $v_\theta$, let $L_{\mathrm{RFM}}(\theta) = \mathbb E_{t,\,x \sim p_t}\lVert v_\theta(x, t) - u_t(x)\rVert^2$ and $L_{\mathrm{RCFM}}(\theta) = \mathbb E_{t,\,x_0,\,x_1}\lVert v_\theta(x_t, t) - u_t(x_t \mid x_1)\rVert^2$. Then $L_{\mathrm{RCFM}}(\theta) = L_{\mathrm{RFM}}(\theta) + \text{const}$: the same gradients and the same minimizer $u_t$.

Proof of the RCFM statement. It is the proof of Theorem 2 in lesson 14, word for word:

1. Expand $\lVert v_\theta - u_t(x_t \mid x_1)\rVert^2 = \lVert v_\theta\rVert^2 - 2\langle v_\theta, u_t(x_t \mid x_1)\rangle + \lVert u_t(x_t \mid x_1)\rVert^2$.
2. $\mathbb E\lVert v_\theta(x_t, t)\rVert^2$ is the same in both losses, because $x_t \sim p_t$.
3. Cross term, by the tower property: $\mathbb E\langle v_\theta(x_t, t), u_t(x_t \mid x_1)\rangle = \mathbb E\langle v_\theta(x_t, t), u_t(x_t)\rangle$. All vectors here lie in the same tangent plane $T_{x_t}S^2$, where the inner product is the one of $\R^3$.
4. The remaining terms do not contain $\theta$. $\blacksquare$

Nothing in the proof uses flatness. Two practical conditions:

1. The model must be tangent: $v_\theta(x, t) = P_x\operatorname{net}_\theta(x, t)$ for any network $\operatorname{net}_\theta$ with output in $\R^3$. Otherwise the learned ODE pushes samples off the sphere.
2. Avoid antipodal pairs: $x_1 = -x_0$ has probability zero for continuous $p$ and $q$, but near-antipodal pairs give unstable targets; skip or clamp them.

Simulation-free. On the sphere, $x_t$ and its target have closed forms, so a training step needs no ODE solve: sample $(x_0, x_1, t)$, compute $x_t$ and the target, regress. On a general manifold (a triangle mesh, say) geodesics have no closed form. Chen and Lipman then build a distance from the eigenfunctions of the Laplace–Beltrami operator (on the sphere these would be the spherical harmonics of lesson 09), and they must simulate the conditional path during training.

Chen and Lipman's general formula (premetrics)

Chen and Lipman define the conditional flow through a premetric $d(x, x_1)$ (zero only at $x = x_1$) and a decreasing scheduler $\kappa(t)$ with $\kappa(0) = 1$, $\kappa(1) = 0$. The flow is built so that $d(x_t, x_1) = \kappa(t)\,d(x_0, x_1)$, and its velocity is

$$u_t(x \mid x_1) = \frac{d\log\kappa(t)}{dt}\; d(x, x_1)\, \frac{\nabla d(x, x_1)}{\lVert\nabla d(x, x_1)\rVert^2}.$$

With $\kappa(t) = 1 - t$ and the geodesic distance, $\nabla_xd(x, x_1) = -\log_x(x_1)/d(x, x_1)$ has norm 1, and the formula gives $\frac{-1}{1 - t}\,d \cdot \big(-\log_x(x_1)/d\big) = \log_x(x_1)/(1 - t)$: Theorem 2 again. With a spectral distance on a mesh the same formula still works, but $x_t$ has no closed form.

Sampling on the sphere #

The exact flow of a tangent field stays on the sphere: $\frac{d}{dt}\lVert x\rVert^2 = 2\,x \cdot v = 0$. Numerical steps of size $h$ behave differently:

• Plain Euler $x + hv$: since $x \cdot v = 0$, $\lVert x + hv\rVert^2 = 1 + h^2\lVert v\rVert^2 > 1$. Every step moves outward, and the error adds up.
• Exponential-map Euler $x_{k+1} = \exp_{x_k}\big(h\,v_\theta(x_k, t_k)\big)$: a step along the great circle.
• Projected Euler $x_{k+1} = (x_k + hv_k)/\lVert x_k + hv_k\rVert$: a retraction, that is, a cheap map from the tangent plane back to the manifold that matches the exponential map at least to first order. This one is even better. It moves by the angle $\arctan(h\lVert v_k\rVert)$ instead of $h\lVert v_k\rVert$, and since $\arctan s = s - s^3/3 + \dots$, it matches the exponential map to second order.

Both keep the samples on the sphere up to round-off, and both are first-order methods like Euler. The RiemannianODESolver of Meta's flow_matching library projects the velocity onto the tangent plane and the point back onto the sphere after each Euler, midpoint or RK4 step.

Interactive figure: loading…

Answers to section 1. (1) At the origin: $x_{1/2} = (x_0 + x_1)/2 = 0$. In general, $\lVert(1 - t)x_0 + tx_1\rVert^2 = 1 - 2t(1 - t)(1 - x_0 \cdot x_1) < 1$ for $0 < t < 1$ and $x_1 \ne x_0$. So the paths run through the interior of the ball, the network must learn a 3D flow for 2D data, and nothing forces the final samples onto the sphere. (2) Outside: $\lVert x + hv\rVert^2 = 1 + h^2\lVert v\rVert^2$.

4. Worked example #

(a) From the North Pole to the equator. $x_0 = (0, 0, 1)$ and $x_1 = (1, 0, 0)$ (latitude 0, longitude 0).

1. Distance: $\omega = \arccos(x_0 \cdot x_1) = \arccos 0 = \pi/2 \approx 1.5708$. On the Earth that is $6371 \times \pi/2 \approx 10\,008$ km.
2. Tangent part: $u = x_1 - (x_0 \cdot x_1)\,x_0 = (1, 0, 0)$, with $\lVert u\rVert = 1 = \sin\omega$. So $\log_{x_0}(x_1) = \tfrac{\pi}{2}(1, 0, 0)$.
3. Path: $x_t = \cos(t\pi/2)\,x_0 + \sin(t\pi/2)\,(1, 0, 0) = (\sin(t\pi/2),\ 0,\ \cos(t\pi/2))$. At $t = \tfrac12$: $x_{1/2} = (0.7071, 0, 0.7071)$, at latitude 45°. The straight-line midpoint $(0.5, 0, 0.5)$ has norm $0.7071$: on the Earth it lies $6371 \times (1 - 0.7071) \approx 1866$ km below the surface.
4. Velocity by differentiation: $\dot x_t = \tfrac{\pi}{2}(\cos(t\pi/2),\ 0,\ -\sin(t\pi/2))$, so $\dot x_{1/2} = (1.1107, 0, -1.1107)$, with speed $\pi/2$ at every $t$.
5. Check Theorem 2 at $t = \tfrac12$. The distance left is $d(x_{1/2}, x_1) = \arccos(0.7071) = \pi/4$. The tangent part of $x_1$ at $x_{1/2}$ is $x_1 - 0.7071\,x_{1/2} = (0.5, 0, -0.5)$, with unit vector $(0.7071, 0, -0.7071)$. So $\log_{x_{1/2}}(x_1) = \tfrac{\pi}{4}(0.7071, 0, -0.7071) = (0.5554, 0, -0.5554)$. Dividing by $1 - t = 0.5$ gives $(1.1107, 0, -1.1107)$, as in step 4.

(b) One step of each sampler from $x_0$, with the velocity $u_0 = (\pi/2, 0, 0)$.

1. Plain Euler: $(1.5708, 0, 1)$, of norm $\sqrt{1 + \pi^2/4} = 1.8621$: far off the sphere.
2. Projected Euler: $(1.5708, 0, 1)/1.8621 = (0.8436, 0, 0.5370)$. It moved by the angle $\arctan(\pi/2) = 1.0039$ (57.5°), so it is still $\pi/2 - 1.0039 = 0.5669$ (32.5°) away from $x_1$.
3. Exponential map: $\cos(\pi/2)\,x_0 + \sin(\pi/2)\,(1, 0, 0) = x_1$, exactly.

With $N$ steps along the conditional field (section 5): plain Euler ends at norms $1.8621$, $1.5607$, $1.2556$, $1.1283$ for $N = 1, 2, 5, 10$; projected Euler ends $0.5669$, $0.1694$, $0.0118$, $0.0014$ away from $x_1$; the exponential map is exact for every $N$.

(c) Near the antipode. $x = (0, 0, 1)$ and $y = (\pm 10^{-8}, 0, -1)$ (normalized). The two points $y$ are $2 \times 10^{-8}$ apart, but $\log_x(y) = (\pm\pi, 0, 0)$: opposite directions. A target that depends on $\log$ is discontinuous there.

5. In code #

SciPy's rotations and geometric_slerp check the formulas of section 3:

import numpy as np
from scipy.spatial import geometric_slerp
from scipy.spatial.transform import Rotation

def exp_map(x, v):
    n = np.linalg.norm(v)
    return x if n == 0 else np.cos(n) * x + np.sin(n) * v / n

def log_map(x, y):
    u = y - (x @ y) * x                                   # P_x y, the tangent part of y
    return np.arccos(np.clip(x @ y, -1, 1)) * u / np.linalg.norm(u)

rng = np.random.default_rng(0)
x, y = rng.standard_normal((2, 3))
x, y = x / np.linalg.norm(x), y / np.linalg.norm(y)
v = log_map(x, y)
print(abs(x @ v), np.linalg.norm(exp_map(x, v) - y))      # 2.8e-17 1.6e-16: tangent, and exp(log y) = y
print(np.linalg.norm(v), np.arccos(x @ y))                # 0.7616 0.7616: |log_x y| = geodesic distance

# exp_x(w) is a rotation of x by the angle |w| about the axis x × w / |w|: rotvec = x × w.
w = np.cross(x, rng.standard_normal(3))
w = 2.5 * w / np.linalg.norm(w)                           # a tangent vector at x with |w| = 2.5 < pi
print(np.linalg.norm(exp_map(x, w) - Rotation.from_rotvec(np.cross(x, w)).apply(x)),
      np.linalg.norm(log_map(x, exp_map(x, w)) - w))      # 1.6e-16 4.4e-16: log(exp w) = w for |w| < pi

# The geodesic path exp_x(t log_x y) is SciPy's spherical interpolation (slerp).
t = np.linspace(0, 1, 7)
path = np.array([exp_map(x, s * v) for s in t])
print(np.abs(path - geometric_slerp(x, y, t)).max())       # 2.5e-16

# Its velocity, by central differences, equals log_{x_t}(y) / (1 - t), with constant speed d(x, y).
s, h = 0.3, 1e-6
xt = exp_map(x, s * v)
vel = (exp_map(x, (s + h) * v) - exp_map(x, (s - h) * v)) / (2 * h)
print(np.linalg.norm(vel - log_map(xt, y) / (1 - s)), np.linalg.norm(vel) - np.arccos(x @ y))  # 5.8e-11 -4.3e-11

Then (same session) the three samplers on the conditional field of worked example (b), and the antipode of (c):

x0, x1 = np.array([0.0, 0.0, 1.0]), np.array([1.0, 0.0, 0.0])
u = lambda x, t: log_map(x, x1) / (1 - t)
angle = lambda z: np.arccos(np.clip(z @ x1 / np.linalg.norm(z), -1, 1))   # angle between z and x1
for N in (1, 2, 5, 10):
    h = 1 / N
    xe, xp, xr = x0.copy(), x0.copy(), x0.copy()
    for k in range(N):
        t = k * h
        xe = xe + h * u(xe / np.linalg.norm(xe), t)         # plain Euler (field read at the radial projection)
        xp = xp + h * u(xp, t); xp /= np.linalg.norm(xp)    # projected Euler (a retraction)
        xr = exp_map(xr, h * u(xr, t))                      # exponential-map Euler
    print(N, round(np.linalg.norm(xe), 4), round(angle(xp), 4), round(angle(xr), 12))
# N  |x| plain Euler  angle to x1 (projected)  angle to x1 (exp map)
# 1  1.8621           0.5669                   0.0
# 2  1.5607           0.1694                   0.0
# 5  1.2556           0.0118                   0.0
# 10 1.1283           0.0014                   0.0

x = np.array([0.0, 0.0, 1.0])
for e in (1e-8, -1e-8):
    y = np.array([e, 0.0, -1.0]); y /= np.linalg.norm(y)
    print(log_map(x, y).round(4))                          # [ 3.1416 0. 0.], then [-3.1416 0. 0.]

6. Where it shows up #

• Earth-science event data. Chen and Lipman's sphere benchmarks are catalogues of volcanic eruptions, earthquakes, floods and fires (compiled by Mathieu and Nickel, 2020, from NOAA, the Dartmouth Flood Observatory and NASA fire data). They report test log-likelihoods for each.
• Seismology and weather. Epicentre densities for seismic hazard; cyclone genesis points and lightning on $S^2$; wind directions on the circle; focal mechanisms (strike, dip and rake) as rotations in $SO(3)$, up to the symmetries of the double couple.
• Proteins and robotics. Flow matching on $SE(3)$ and $SO(3)$ (for example FrameFlow and FoldFlow) generates protein backbones as rigid frames; Chen and Lipman also model protein dihedral angles on tori.
• Software. Meta's open-source flow_matching library has a Sphere manifold (expmap, logmap, projx, proju), a geodesic probability path and a RiemannianODESolver, with a sphere example notebook.

7. Common mistakes #

• A velocity model that is not tangent. Without $P_x$, normal components push samples off the sphere. Project the output, and project the input point too if it may have drifted.
• Plain Euler steps. They leave the sphere at every step; use the exponential map or project.
• arccos without clipping. Round-off can give $x \cdot y = 1.0000000000000002$ (the next float above 1), and arccos returns NaN. Clip to $[-1, 1]$.
• Dividing by $\lVert u\rVert \approx 0$. For $y$ close to $x$, $\log_x(y) \approx P_x(y - x)$; switch to it below a small threshold.
• A "uniform" source made from uniform latitude and longitude. The area element is $\cos(\text{latitude})\,d(\text{latitude})\,d(\text{longitude})$, so uniform latitudes pile up mass near the poles. Normalize a standard 3D Gaussian instead: its density depends only on $\lVert z\rVert$, so its direction is uniform.
• Different base points. Chen and Lipman write $x_t = \exp_{x_1}\big((1 - t)\log_{x_1}x_0\big)$; other code computes $\exp_{x_0}(t\log_{x_0}x_1)$. Both are the same great-circle path when the geodesic is unique.
• Fields versus points. SFNO (lesson 09) maps fields on the sphere to fields; Riemannian flow matching models where points fall on the sphere.

8. Check yourself #

1. Compute $d(x, y)$ and $\log_x(y)$ for $x = (0, 0, 1)$ and $y = (0, \sqrt3/2, 1/2)$. How far apart are the two points on the Earth?
   Answer

   $x \cdot y = 1/2$, so $\omega = \pi/3 \approx 1.0472$. $u = y - \tfrac12 x = (0, \sqrt3/2, 0)$, with $\lVert u\rVert = \sin\omega$. $\log_x(y) = (0, \pi/3, 0) \approx (0, 1.0472, 0)$. On the Earth: $6371 \times 1.0472 \approx 6672$ km.
2. A projected Euler step and an exponential-map step use the same tangent step of length $s = 0.1$. By how much do the angles they move differ?
   Answer

   The projected step moves $\arctan(0.1) = 0.099669$, the exponential map $0.1$: a difference of $3.3 \times 10^{-4} \approx s^3/3$.
3. Show that RCFM with the geodesic path is the CFM of lesson 14 when the manifold is $\R^d$.
   Answer

   In $\R^d$, $\exp_x(v) = x + v$ and $\log_x(y) = y - x$. So $x_t = x_0 + t(x_1 - x_0)$, and the target $\log_{x_t}(x_1)/(1 - t) = (x_1 - x_t)/(1 - t) = x_1 - x_0$. The tangent projection is the identity.
4. Why is training simulation-free on the sphere but not on a general triangle mesh?
   Answer

   On the sphere, $\exp$ and $\log$ have closed forms, so $x_t$ and the target are computed directly. On a mesh, geodesics have no closed form; Chen and Lipman use a spectral distance instead, and then $x_t$ must be found by solving the ODE of the conditional velocity.
5. Your trained model's samples have $\big\lvert\lVert x\rVert - 1\big\rvert \approx 10^{-2}$. Name two likely causes.
   Answer

   Plain Euler steps (each adds $h^2\lVert v\rVert^2$ to $\lVert x\rVert^2$), or a velocity network whose output is not projected onto the tangent plane. With exp-map or projected steps and a tangent model, the error is at round-off level.
6. Is $\log_x(\exp_x v) = v$ for $\lVert v\rVert = 4$? Explain with a picture.
   Answer

   No. Walking 4 > π along the great circle passes the antipode $-x$; the end point is closer to $x$ the other way round, at distance $2\pi - 4 \approx 2.28$. The log gives that shorter way: a vector of length $2\pi - 4$ pointing opposite to $v$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is the projection onto the tangent plane of $S^2$ at $x$?

   Show answer

2. What is the geodesic distance between unit vectors $x$ and $y$ on $S^2$?

   Show answer

3. Write the exponential map of the sphere.

   Show answer

4. Write the logarithm map of the sphere.

   Show answer

5. Why is $\log_x$ undefined at $y = -x$?

   Show answer

6. What is the velocity of the geodesic path $x_t = \exp_{x_0}(t\log_{x_0}x_1)$?

   Show answer

7. Write the Riemannian CFM loss.

   Show answer

8. One plain Euler step $x + hv$ with tangent $v$: what is $\lVert x + hv\rVert^2$?

   Show answer

9. Why use a retraction such as $(x + hv)/\lVert x + hv\rVert$ when sampling?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How would you build a generative model of earthquake epicentres?" Convert latitude and longitude to unit vectors, use a uniform source on $S^2$, the geodesic path, a tangent-projected MLP velocity, the RCFM loss, and an exponential-map or projected ODE solver. Evaluate the test log-likelihood, as Chen and Lipman do.
• "Why not just normalize the output of a Euclidean model?" The training paths and the model live in $\R^3$ (through the Earth), and the normalization is not part of the learned flow. It can work, but it wastes capacity, and the density on the sphere is not modelled directly.
• "Diffusion or flow matching on manifolds?" Manifold diffusion models need simulation during training and approximations of the score on the manifold; Riemannian flow matching is simulation-free on simple manifolds such as the sphere.
• "What goes wrong numerically?" Points leaving the manifold, NaNs from arccos, division by $\lVert u\rVert \approx 0$ near $x$, and unstable directions near the antipode.
• "What is the analogue of the OT straight path on a manifold?" The geodesic interpolation; in $\R^d$ it is the conditional OT path of lesson 14.

10. Practice #

The notebook for this lesson runs in Colab or in the browser. It has five exercises, each with a check cell and a folded solution:

1. The exponential and logarithm maps of $S^2$ and the round-trip test.
2. The geodesic path equals slerp (scipy.spatial.geometric_slerp).
3. The velocity of the geodesic path by central differences equals $\log_{x_t}(x_1)/(1 - t)$.
4. Plain, projected and exponential-map Euler on the conditional velocity.
5. A data set of three clusters on the sphere: the exact marginal velocity (a weighted average of conditional velocities) and sampling from the uniform distribution to the clusters.

11. Watch and read #

Flow Matching: Simplifying and Generalizing Diffusion Models | Yaron LipmanProxima · 60 min

Flow matching and its generalizations by its first author; compare with Theorem 3 of this lesson.

Leonel Rozo - Riemannian Generative Models Part 2Noémie Jaquier · 25 min

A lecture on generative models on Riemannian manifolds, the general setting of this lesson.

MIT 6.S184: Flow Matching and Diffusion Models - Lecture 02: Flow Matching (2026)Peter Holderrieth · 83 min

The Euclidean flow matching of lesson 14, which this lesson moves to the sphere.

• Chen and Lipman, Flow Matching on General Geometries (ICLR 2024): §3.1 (Riemannian flow matching and the RCFM objective), §3.2 (premetrics and the geodesic distance), §3.3 (spectral distances), §5 (the Earth-science data sets).
• Lipman et al., Flow Matching Guide and Code (2024): §5.3–5.6 (paths on manifolds, Theorems 10 and 11, the geodesic conditional flow, the antipodal singularity).
• Boumal, An Introduction to Optimization on Smooth Manifolds (2023, free): §3.6 (retractions on the sphere), §5.8 (great circles are geodesics), §7.2 (the unit sphere), §10.2 (exp and log of the sphere).
• Holderrieth and Erives, An Introduction to Flow Matching and Diffusion Models (2025): the MIT 6.S184 lecture notes.
• Meta, flow_matching sphere example: Riemannian flow matching on the sphere in PyTorch.