Geodesic, Pullback Metric & FlatVI¶
Hands-on notebook covering three interconnected theory notes:
05-geodesic.md— Geodesic distance, why LERP fails on manifolds, Exp/Log maps06-pullback-metric.md— Pullback metric $g_Z = J^T J$, measuring decoder distortion07-flatvi.md— FlatVI: forcing the latent space to be flat via a Frobenius penalty
The three topics form a progression: geodesics expose the problem (LERP fails) → the pullback metric provides the diagnosis (non-isotropic Jacobian) → FlatVI delivers the cure (penalise deviation from $\alpha I$).
All text is in English to match the project language rules for artifacts.
Linked theory: research/05-geodesic.md · research/06-pullback-metric.md · research/07-flatvi.md · Chạy: uv run jupyter lab từ thư mục latent-anything-theory/ (cần các package trong pyproject.toml).
In [1]:
Copied!
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from mpl_toolkits.mplot3d import Axes3D # noqa: F401
from collections.abc import Callable
from sklearn.datasets import make_swiss_roll
from sklearn.neighbors import NearestNeighbors
from sklearn.metrics import pairwise_distances
plt.style.use('seaborn-v0_8-whitegrid')
np.set_printoptions(precision=4, suppress=True)
rng = np.random.default_rng(42)
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from mpl_toolkits.mplot3d import Axes3D # noqa: F401
from collections.abc import Callable
from sklearn.datasets import make_swiss_roll
from sklearn.neighbors import NearestNeighbors
from sklearn.metrics import pairwise_distances
plt.style.use('seaborn-v0_8-whitegrid')
np.set_printoptions(precision=4, suppress=True)
rng = np.random.default_rng(42)
Part 1 — Geodesic (05-geodesic.md)¶
A geodesic is the shortest curve connecting two points that lies entirely on the manifold.
Sections:
- 1.0 1-D motivation: why a non-linear decoder makes LERP wrong
- 1.1 Why LERP fails on manifolds (unit circle $S^1$)
- 1.2 SLERP — geodesic interpolation on $S^{n-1}$ (exercise)
- 1.3 Exponential and Logarithmic maps
- 1.4 Swiss Roll: LERP shortcuts through empty latent space
1.0 — 1-D Motivation: Why the Decoder Makes LERP Wrong¶
Before going to manifolds, consider the simplest possible case: a 1-D decoder $x = e^z$.
- LERP spaces the latent points $z$ uniformly: $z_i = z_0 + i \cdot \Delta z$.
- But $e^z$ grows exponentially, so equal $\Delta z$ steps produce unequal $\Delta x$ gaps in data space.
- A geodesic would space the data-space outputs uniformly: $x_i = x_0 + i \cdot \Delta x$, which requires non-uniform $z$ steps.
The horizontal dashed lines show the data-space level of each interpolated point. For LERP the gaps grow exponentially; for the geodesic they are equal.
In [2]:
Copied!
z_range = np.linspace(-2.0, 2.0, 200)
x_range = np.exp(z_range)
# 9 points: equal z-steps (LERP) vs equal x-steps (geodesic)
z_lerp = np.linspace(-2.0, 2.0, 9)
x_lerp = np.exp(z_lerp)
x_geo = np.linspace(np.exp(-2.0), np.exp(2.0), 9)
z_geo = np.log(x_geo)
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
for ax, z_pts, x_pts, col, title, subtitle in [
(axes[0], z_lerp, x_lerp, 'crimson',
'LERP: equal z-steps',
'Data gaps grow exponentially — non-uniform'),
(axes[1], z_geo, x_geo, 'seagreen',
'Geodesic: equal x-steps',
'Requires non-uniform z-steps — uniform data gaps'),
]:
ax.plot(z_range, x_range, color='steelblue', linewidth=2.5, alpha=0.85,
label='Decoder: x = exp(z)')
for x_v in x_pts:
ax.axhline(x_v, color=col, alpha=0.25, linewidth=0.9, linestyle='--')
ax.scatter(z_pts, x_pts, color=col, s=160, zorder=8,
edgecolors='black', linewidth=1.2, label='Interpolated points')
ax.set_xlabel('Latent z')
ax.set_ylabel('Data x = exp(z)')
ax.set_title(f'{title}\n{subtitle}', fontsize=10)
ax.set_xlim(-2.4, 2.4)
ax.legend(fontsize=9)
plt.suptitle('1-D Motivation: Equal Latent Steps do NOT Equal Equal Data Steps\n'
'The decoder exp(z) stretches space — LERP gives unequal data gaps',
fontsize=11, y=1.02)
plt.tight_layout()
plt.show()
gaps_l = np.diff(x_lerp)
gaps_g = np.diff(x_geo)
print(f'LERP x-gaps: min={gaps_l.min():.3f} max={gaps_l.max():.3f} ratio={gaps_l.max()/gaps_l.min():.1f}x')
print(f'Geo x-gaps: min={gaps_g.min():.3f} max={gaps_g.max():.3f} ratio={gaps_g.max()/gaps_g.min():.1f}x')
z_range = np.linspace(-2.0, 2.0, 200)
x_range = np.exp(z_range)
# 9 points: equal z-steps (LERP) vs equal x-steps (geodesic)
z_lerp = np.linspace(-2.0, 2.0, 9)
x_lerp = np.exp(z_lerp)
x_geo = np.linspace(np.exp(-2.0), np.exp(2.0), 9)
z_geo = np.log(x_geo)
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
for ax, z_pts, x_pts, col, title, subtitle in [
(axes[0], z_lerp, x_lerp, 'crimson',
'LERP: equal z-steps',
'Data gaps grow exponentially — non-uniform'),
(axes[1], z_geo, x_geo, 'seagreen',
'Geodesic: equal x-steps',
'Requires non-uniform z-steps — uniform data gaps'),
]:
ax.plot(z_range, x_range, color='steelblue', linewidth=2.5, alpha=0.85,
label='Decoder: x = exp(z)')
for x_v in x_pts:
ax.axhline(x_v, color=col, alpha=0.25, linewidth=0.9, linestyle='--')
ax.scatter(z_pts, x_pts, color=col, s=160, zorder=8,
edgecolors='black', linewidth=1.2, label='Interpolated points')
ax.set_xlabel('Latent z')
ax.set_ylabel('Data x = exp(z)')
ax.set_title(f'{title}\n{subtitle}', fontsize=10)
ax.set_xlim(-2.4, 2.4)
ax.legend(fontsize=9)
plt.suptitle('1-D Motivation: Equal Latent Steps do NOT Equal Equal Data Steps\n'
'The decoder exp(z) stretches space — LERP gives unequal data gaps',
fontsize=11, y=1.02)
plt.tight_layout()
plt.show()
gaps_l = np.diff(x_lerp)
gaps_g = np.diff(x_geo)
print(f'LERP x-gaps: min={gaps_l.min():.3f} max={gaps_l.max():.3f} ratio={gaps_l.max()/gaps_l.min():.1f}x')
print(f'Geo x-gaps: min={gaps_g.min():.3f} max={gaps_g.max():.3f} ratio={gaps_g.max()/gaps_g.min():.1f}x')
LERP x-gaps: min=0.088 max=2.907 ratio=33.1x Geo x-gaps: min=0.907 max=0.907 ratio=1.0x
1.1 — Why LERP Fails on Manifolds¶
Linear interpolation (LERP) between two points $a$ and $b$: $$Y(t) = (1-t)\,a + t\,b, \quad t \in [0,1]$$
In flat Euclidean space this is the shortest path. On a curved manifold:
- The midpoint $(1-t)a + tb$ does not generally lie on the manifold.
- The path cuts through the ambient space instead of staying on the surface.
- In a deep-learning latent space this means traversing near-zero-density regions, producing blurry or unrealistic decoded outputs.
Demo: the unit circle $S^1$ is the simplest non-trivial manifold. LERP between two points on $S^1$ dips inside — it leaves the manifold.
In [3]:
Copied!
theta_a, theta_b = np.pi / 6, 5 * np.pi / 6
a = np.array([np.cos(theta_a), np.sin(theta_a)])
b = np.array([np.cos(theta_b), np.sin(theta_b)])
ts = np.linspace(0, 1, 51)
lerp_path = np.outer(1 - ts, a) + np.outer(ts, b)
lerp_norms = np.linalg.norm(lerp_path, axis=1)
print('LERP between two unit-circle points:')
print(f' a = {a} (norm = {np.linalg.norm(a):.4f})')
print(f' b = {b} (norm = {np.linalg.norm(b):.4f})')
print(f' Midpoint (t=0.5) norm : {lerp_norms[25]:.4f} <- should be 1.0 to stay on S^1')
print(f' Max deviation from 1 : {abs(1 - lerp_norms).max():.4f}')
theta_a, theta_b = np.pi / 6, 5 * np.pi / 6
a = np.array([np.cos(theta_a), np.sin(theta_a)])
b = np.array([np.cos(theta_b), np.sin(theta_b)])
ts = np.linspace(0, 1, 51)
lerp_path = np.outer(1 - ts, a) + np.outer(ts, b)
lerp_norms = np.linalg.norm(lerp_path, axis=1)
print('LERP between two unit-circle points:')
print(f' a = {a} (norm = {np.linalg.norm(a):.4f})')
print(f' b = {b} (norm = {np.linalg.norm(b):.4f})')
print(f' Midpoint (t=0.5) norm : {lerp_norms[25]:.4f} <- should be 1.0 to stay on S^1')
print(f' Max deviation from 1 : {abs(1 - lerp_norms).max():.4f}')
LERP between two unit-circle points: a = [0.866 0.5 ] (norm = 1.0000) b = [-0.866 0.5 ] (norm = 1.0000) Midpoint (t=0.5) norm : 0.5000 <- should be 1.0 to stay on S^1 Max deviation from 1 : 0.5000
In [4]:
Copied!
circle_angles = np.linspace(0, 2 * np.pi, 300)
circle = np.column_stack([np.cos(circle_angles), np.sin(circle_angles)])
fig, ax = plt.subplots(figsize=(6, 6))
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.5, label='Manifold S^1')
ax.plot(lerp_path[:, 0], lerp_path[:, 1], 'r--', linewidth=2, label='LERP (leaves manifold)')
ax.scatter([a[0], b[0]], [a[1], b[1]], color='black', s=100, zorder=5, label='Endpoints')
ax.scatter(*lerp_path[25], color='red', s=140, marker='x', zorder=6, linewidths=3,
label=f'Midpoint norm = {lerp_norms[25]:.3f} (NOT 1)')
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
ax.set_title('LERP dips inside S^1 — it leaves the manifold')
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()
circle_angles = np.linspace(0, 2 * np.pi, 300)
circle = np.column_stack([np.cos(circle_angles), np.sin(circle_angles)])
fig, ax = plt.subplots(figsize=(6, 6))
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.5, label='Manifold S^1')
ax.plot(lerp_path[:, 0], lerp_path[:, 1], 'r--', linewidth=2, label='LERP (leaves manifold)')
ax.scatter([a[0], b[0]], [a[1], b[1]], color='black', s=100, zorder=5, label='Endpoints')
ax.scatter(*lerp_path[25], color='red', s=140, marker='x', zorder=6, linewidths=3,
label=f'Midpoint norm = {lerp_norms[25]:.3f} (NOT 1)')
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
ax.set_title('LERP dips inside S^1 — it leaves the manifold')
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()
In [5]:
Copied!
# === Where exactly are equal-t steps? ===
# 9 milestone points at t = 0, 0.125, 0.25, ..., 1.0
# Color: red (t=0) -> yellow (t=0.5) -> green (t=1)
t_ms = np.linspace(0, 1, 9)
cmap_t = plt.cm.RdYlGn
lerp_ms = np.array([(1 - t) * a + t * b for t in t_ms])
# slerp_ms computed after slerp is defined — placeholder, filled below
def _slerp_preview(u, v, t):
'''Temporary SLERP used only for this visualization (Exercise 1 comes next).'''
cos_th = float(np.clip(np.dot(u, v), -1.0, 1.0))
th = np.arccos(cos_th)
if abs(th) < 1e-8:
return (1.0 - t) * u + t * v
return (np.sin((1.0 - t) * th) / np.sin(th)) * u + \
(np.sin(t * th) / np.sin(th)) * v
slerp_ms = np.array([_slerp_preview(a, b, t) for t in t_ms])
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
for ax, pts, path, path_c, label in [
(axes[0], lerp_ms, lerp_path, 'red', 'LERP'),
(axes[1], slerp_ms, None, 'green', 'SLERP'),
]:
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2.5, alpha=0.35, label='S^1')
if path is not None:
ax.plot(path[:, 0], path[:, 1], '-', color=path_c, linewidth=1.5, alpha=0.4)
else:
slerp_full = np.array([_slerp_preview(a, b, t) for t in ts])
ax.plot(slerp_full[:, 0], slerp_full[:, 1], '-', color=path_c, linewidth=1.5, alpha=0.4)
for pt, t_v in zip(pts, t_ms):
ax.scatter(*pt, color=cmap_t(t_v), s=230, zorder=9,
edgecolors='black', linewidth=1.5)
mid = pts[4]
mid_n = np.linalg.norm(mid)
status = 'ON S^1' if abs(mid_n - 1.0) < 0.01 else f'norm = {mid_n:.3f} OFF S^1!'
ax.annotate(f'midpoint\n({status})', xy=mid,
xytext=(mid[0] + 0.25, mid[1] - 0.35), fontsize=9, fontweight='bold',
arrowprops=dict(arrowstyle='->', color='black', lw=1.2))
ax.set_xlim(-1.7, 1.7)
ax.set_ylim(-1.7, 1.7)
ax.set_aspect('equal')
ax.set_title(f'{label} — 9 equal-t steps (red=start, green=end)')
plt.suptitle('Are Equal-t Steps Equally Spaced on S^1?\n'
'LERP: midpoint leaves the manifold | SLERP: midpoint stays on S^1',
fontsize=12)
plt.tight_layout()
plt.show()
# === Where exactly are equal-t steps? ===
# 9 milestone points at t = 0, 0.125, 0.25, ..., 1.0
# Color: red (t=0) -> yellow (t=0.5) -> green (t=1)
t_ms = np.linspace(0, 1, 9)
cmap_t = plt.cm.RdYlGn
lerp_ms = np.array([(1 - t) * a + t * b for t in t_ms])
# slerp_ms computed after slerp is defined — placeholder, filled below
def _slerp_preview(u, v, t):
'''Temporary SLERP used only for this visualization (Exercise 1 comes next).'''
cos_th = float(np.clip(np.dot(u, v), -1.0, 1.0))
th = np.arccos(cos_th)
if abs(th) < 1e-8:
return (1.0 - t) * u + t * v
return (np.sin((1.0 - t) * th) / np.sin(th)) * u + \
(np.sin(t * th) / np.sin(th)) * v
slerp_ms = np.array([_slerp_preview(a, b, t) for t in t_ms])
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
for ax, pts, path, path_c, label in [
(axes[0], lerp_ms, lerp_path, 'red', 'LERP'),
(axes[1], slerp_ms, None, 'green', 'SLERP'),
]:
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2.5, alpha=0.35, label='S^1')
if path is not None:
ax.plot(path[:, 0], path[:, 1], '-', color=path_c, linewidth=1.5, alpha=0.4)
else:
slerp_full = np.array([_slerp_preview(a, b, t) for t in ts])
ax.plot(slerp_full[:, 0], slerp_full[:, 1], '-', color=path_c, linewidth=1.5, alpha=0.4)
for pt, t_v in zip(pts, t_ms):
ax.scatter(*pt, color=cmap_t(t_v), s=230, zorder=9,
edgecolors='black', linewidth=1.5)
mid = pts[4]
mid_n = np.linalg.norm(mid)
status = 'ON S^1' if abs(mid_n - 1.0) < 0.01 else f'norm = {mid_n:.3f} OFF S^1!'
ax.annotate(f'midpoint\n({status})', xy=mid,
xytext=(mid[0] + 0.25, mid[1] - 0.35), fontsize=9, fontweight='bold',
arrowprops=dict(arrowstyle='->', color='black', lw=1.2))
ax.set_xlim(-1.7, 1.7)
ax.set_ylim(-1.7, 1.7)
ax.set_aspect('equal')
ax.set_title(f'{label} — 9 equal-t steps (red=start, green=end)')
plt.suptitle('Are Equal-t Steps Equally Spaced on S^1?\n'
'LERP: midpoint leaves the manifold | SLERP: midpoint stays on S^1',
fontsize=12)
plt.tight_layout()
plt.show()
What you see¶
What the plots show: The printout reports that the LERP midpoint between two unit-circle points has norm $\approx 0.87$, not 1. The left plot makes this visible: the red LERP chord cuts straight through the inside of the blue circle, and the marked midpoint sits off the manifold. The spacing plot shows a second problem — equal-$t$ steps along LERP bunch up toward the middle, whereas the SLERP steps are spaced evenly along the arc.
Why it looks this way: A straight segment between two points on a curved surface leaves the surface; only the endpoints ($t=0,1$) have norm exactly 1. SLERP instead travels along the great-circle arc at constant angular speed, so every intermediate point stays on $S^1$ and equal $t$-steps subtend equal arc length.
Key takeaway: In a latent space whose data lies on a curved manifold, LERP walks through low-density, off-manifold regions — producing blurry or unrealistic decoded outputs — and it moves at an uneven speed. SLERP and, more generally, geodesics fix both issues, which is exactly why they matter for latent interpolation.
1.2 — SLERP: Geodesic Interpolation on $S^{n-1}$ (Exercise 1)¶
Spherical Linear Interpolation (SLERP) moves at constant angular speed along the great-circle arc:
$$\text{SLERP}(u, v, t) = \frac{\sin((1-t)\theta)}{\sin\theta}\,u + \frac{\sin(t\theta)}{\sin\theta}\,v, \qquad \theta = \arccos(u \cdot v)$$
Properties:
- $\text{SLERP}(u,v,0)=u$ and $\text{SLERP}(u,v,1)=v$
- All intermediate points lie on $S^{n-1}$ (norm = 1)
- Constant angular speed: equal $t$-steps subtend equal arc lengths
Degenerate case: when $\theta \approx 0$ ($u \approx v$), fall back to LERP.
In [6]:
Copied!
def slerp(u: np.ndarray, v: np.ndarray, t: float) -> np.ndarray:
'''
Spherical Linear Interpolation between unit vectors u and v at t in [0, 1].
Moves at constant angular speed along the great-circle arc on S^(n-1).
'''
# TODO:
# 1. cos_theta = clip(u . v, -1, 1)
# 2. theta = arccos(cos_theta)
# 3. If theta < 1e-8, return linear blend (degenerate fallback)
# 4. Return sin((1-t)*theta)/sin(theta)*u + sin(t*theta)/sin(theta)*v
# raise NotImplementedError('TODO: implement slerp')
cos_theta = float(np.clip(np.dot(u, v), -1.0, 1.0))
theta = np.arccos(cos_theta)
if abs(theta) < 1e-8:
return (1.0 - t) * u + t * v
return (np.sin((1.0 - t) * theta) / np.sin(theta)) * u + \
(np.sin(t * theta) / np.sin(theta)) * v
def slerp(u: np.ndarray, v: np.ndarray, t: float) -> np.ndarray:
'''
Spherical Linear Interpolation between unit vectors u and v at t in [0, 1].
Moves at constant angular speed along the great-circle arc on S^(n-1).
'''
# TODO:
# 1. cos_theta = clip(u . v, -1, 1)
# 2. theta = arccos(cos_theta)
# 3. If theta < 1e-8, return linear blend (degenerate fallback)
# 4. Return sin((1-t)*theta)/sin(theta)*u + sin(t*theta)/sin(theta)*v
# raise NotImplementedError('TODO: implement slerp')
cos_theta = float(np.clip(np.dot(u, v), -1.0, 1.0))
theta = np.arccos(cos_theta)
if abs(theta) < 1e-8:
return (1.0 - t) * u + t * v
return (np.sin((1.0 - t) * theta) / np.sin(theta)) * u + \
(np.sin(t * theta) / np.sin(theta)) * v
In [7]:
Copied!
slerp_path = np.array([slerp(a, b, t) for t in ts])
slerp_norms = np.linalg.norm(slerp_path, axis=1)
assert np.allclose(slerp_norms, 1.0, atol=1e-6)
assert np.allclose(slerp(a, b, 0.0), a, atol=1e-8)
assert np.allclose(slerp(a, b, 1.0), b, atol=1e-8)
arc_diffs = np.diff(np.arctan2(slerp_path[:, 1], slerp_path[:, 0]))
arc_cv = arc_diffs.std() / abs(arc_diffs.mean())
print('Exercise 1 - SLERP looks good.')
print(f' Max norm deviation : {abs(slerp_norms - 1.0).max():.2e}')
print(f' Angular speed CoV : {arc_cv:.2e} (0 = perfectly constant)')
slerp_path = np.array([slerp(a, b, t) for t in ts])
slerp_norms = np.linalg.norm(slerp_path, axis=1)
assert np.allclose(slerp_norms, 1.0, atol=1e-6)
assert np.allclose(slerp(a, b, 0.0), a, atol=1e-8)
assert np.allclose(slerp(a, b, 1.0), b, atol=1e-8)
arc_diffs = np.diff(np.arctan2(slerp_path[:, 1], slerp_path[:, 0]))
arc_cv = arc_diffs.std() / abs(arc_diffs.mean())
print('Exercise 1 - SLERP looks good.')
print(f' Max norm deviation : {abs(slerp_norms - 1.0).max():.2e}')
print(f' Angular speed CoV : {arc_cv:.2e} (0 = perfectly constant)')
Exercise 1 - SLERP looks good. Max norm deviation : 2.22e-16 Angular speed CoV : 3.53e-15 (0 = perfectly constant)
In [8]:
Copied!
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
ax = axes[0]
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.4, label='S^1')
ax.plot(lerp_path[:, 0], lerp_path[:, 1], 'r--', linewidth=2, label='LERP (off-manifold)')
ax.plot(slerp_path[:, 0], slerp_path[:, 1], 'g-', linewidth=2.5, label='SLERP (on-manifold)')
ax.scatter([a[0], b[0]], [a[1], b[1]], color='black', s=100, zorder=5)
ax.set_aspect('equal')
ax.set_xlim(-1.4, 1.4)
ax.set_ylim(-1.4, 1.4)
ax.set_title('Paths on S^1')
ax.legend(fontsize=9)
ax = axes[1]
ax.plot(ts, lerp_norms, 'r--', linewidth=2, label='LERP norm')
ax.plot(ts, slerp_norms, 'g-', linewidth=2, label='SLERP norm')
ax.axhline(1.0, color='gray', linestyle='-', alpha=0.5, label='S^1 (norm = 1)')
ax.set_xlabel('Interpolation parameter t')
ax.set_ylabel('Norm of interpolated point')
ax.set_title('Only SLERP stays on the manifold')
ax.legend()
plt.suptitle('LERP vs SLERP', fontsize=13)
plt.tight_layout()
plt.show()
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
ax = axes[0]
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.4, label='S^1')
ax.plot(lerp_path[:, 0], lerp_path[:, 1], 'r--', linewidth=2, label='LERP (off-manifold)')
ax.plot(slerp_path[:, 0], slerp_path[:, 1], 'g-', linewidth=2.5, label='SLERP (on-manifold)')
ax.scatter([a[0], b[0]], [a[1], b[1]], color='black', s=100, zorder=5)
ax.set_aspect('equal')
ax.set_xlim(-1.4, 1.4)
ax.set_ylim(-1.4, 1.4)
ax.set_title('Paths on S^1')
ax.legend(fontsize=9)
ax = axes[1]
ax.plot(ts, lerp_norms, 'r--', linewidth=2, label='LERP norm')
ax.plot(ts, slerp_norms, 'g-', linewidth=2, label='SLERP norm')
ax.axhline(1.0, color='gray', linestyle='-', alpha=0.5, label='S^1 (norm = 1)')
ax.set_xlabel('Interpolation parameter t')
ax.set_ylabel('Norm of interpolated point')
ax.set_title('Only SLERP stays on the manifold')
ax.legend()
plt.suptitle('LERP vs SLERP', fontsize=13)
plt.tight_layout()
plt.show()
1.3 — Exponential and Logarithmic Maps¶
The Exp/Log pair is the general recipe for geodesic interpolation on any Riemannian manifold:
$$Y(t) = \text{Exp}_{Y_0}\!\bigl(t \cdot \text{Log}_{Y_0}(Y_1)\bigr)$$
| Map | Input | Output | Meaning |
|---|---|---|---|
| $\text{Log}_p(q)$ | two manifold points | tangent vector at $p$ | direction + geodesic distance from $p$ to $q$ |
| $\text{Exp}_p(v)$ | point $p$ + tangent $v$ | manifold point | walk $\|v\|$ steps along geodesic from $p$ in direction $v$ |
On $S^1$:
$\text{Log}_p(q) = \theta\,\hat{v}$ where $\theta = \arccos(p \cdot q)$ and $\hat{v}$ is the unit tangent at $p$ pointing toward $q$.
$\text{Exp}_p(v) = \cos(\|v\|)\,p + \sin(\|v\|)\,\hat{v}$
Key identity: $\text{Exp}_p(\text{Log}_p(q)) = q$.
In [9]:
Copied!
def log_map_s1(p: np.ndarray, q: np.ndarray) -> np.ndarray:
cos_theta = float(np.clip(np.dot(p, q), -1.0, 1.0))
theta = np.arccos(cos_theta)
if abs(theta) < 1e-10:
return np.zeros(2)
q_tangent = q - np.dot(p, q) * p
return theta * q_tangent / np.linalg.norm(q_tangent)
def exp_map_s1(p: np.ndarray, v: np.ndarray) -> np.ndarray:
norm_v = np.linalg.norm(v)
if norm_v < 1e-10:
return p.copy()
return np.cos(norm_v) * p + np.sin(norm_v) * (v / norm_v)
log_vec = log_map_s1(a, b)
recovered = exp_map_s1(a, log_vec)
print(f'Log map: geodesic distance a->b = {np.linalg.norm(log_vec):.4f} rad')
print(f'Expected (theta_b - theta_a) = {theta_b - theta_a:.4f} rad')
print(f'Exp(a, Log(a,b)) == b : {np.allclose(recovered, b, atol=1e-7)}')
print()
# Y(t) = Exp_a(t * Log_a(b))
geod_path = np.array([exp_map_s1(a, t * log_vec) for t in ts])
print(f'Geodesic interp stays on S^1 : {np.allclose(np.linalg.norm(geod_path,axis=1), 1.0, atol=1e-7)}')
print(f'Geodesic path == SLERP path : {np.allclose(geod_path, slerp_path, atol=1e-7)}')
def log_map_s1(p: np.ndarray, q: np.ndarray) -> np.ndarray:
cos_theta = float(np.clip(np.dot(p, q), -1.0, 1.0))
theta = np.arccos(cos_theta)
if abs(theta) < 1e-10:
return np.zeros(2)
q_tangent = q - np.dot(p, q) * p
return theta * q_tangent / np.linalg.norm(q_tangent)
def exp_map_s1(p: np.ndarray, v: np.ndarray) -> np.ndarray:
norm_v = np.linalg.norm(v)
if norm_v < 1e-10:
return p.copy()
return np.cos(norm_v) * p + np.sin(norm_v) * (v / norm_v)
log_vec = log_map_s1(a, b)
recovered = exp_map_s1(a, log_vec)
print(f'Log map: geodesic distance a->b = {np.linalg.norm(log_vec):.4f} rad')
print(f'Expected (theta_b - theta_a) = {theta_b - theta_a:.4f} rad')
print(f'Exp(a, Log(a,b)) == b : {np.allclose(recovered, b, atol=1e-7)}')
print()
# Y(t) = Exp_a(t * Log_a(b))
geod_path = np.array([exp_map_s1(a, t * log_vec) for t in ts])
print(f'Geodesic interp stays on S^1 : {np.allclose(np.linalg.norm(geod_path,axis=1), 1.0, atol=1e-7)}')
print(f'Geodesic path == SLERP path : {np.allclose(geod_path, slerp_path, atol=1e-7)}')
Log map: geodesic distance a->b = 2.0944 rad Expected (theta_b - theta_a) = 2.0944 rad Exp(a, Log(a,b)) == b : True Geodesic interp stays on S^1 : True Geodesic path == SLERP path : True
In [10]:
Copied!
# Visualise the Log map tangent vector and the geodesic it generates
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
# Left: Log map — tangent vector at p pointing toward q
ax = axes[0]
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.5)
ax.scatter(*a, color='blue', s=120, zorder=6, label='a (base point p)')
ax.scatter(*b, color='green', s=120, zorder=6, label='b (target q)')
# Draw tangent vector (scaled for visibility)
vis_scale = 0.55
log_unit = log_vec / np.linalg.norm(log_vec)
ax.annotate('', xy=a + log_unit * vis_scale, xytext=a,
arrowprops=dict(arrowstyle='->', color='crimson', lw=2.5))
ax.text(a[0] + log_unit[0] * vis_scale + 0.05, a[1] + log_unit[1] * vis_scale + 0.05,
r'Log$_a$(b)', fontsize=11, color='crimson')
# Tangent line
perp = np.array([-a[1], a[0]]) # perpendicular to radius = tangent to circle at a
tang_pts = a[:, None] + perp[:, None] * np.linspace(-0.7, 0.7, 50)
ax.plot(tang_pts[0], tang_pts[1], 'r--', linewidth=1.2, alpha=0.6, label='Tangent plane at a')
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
ax.set_title('Log map: project destination\ninto tangent space at base point')
ax.legend(fontsize=9)
# Right: Exp map — follow the tangent vector to get geodesic
ax = axes[1]
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.5)
ax.plot(geod_path[:, 0], geod_path[:, 1], 'g-', linewidth=2.5, label='Geodesic path (SLERP)')
ax.scatter(*a, color='blue', s=120, zorder=6, label='a')
ax.scatter(*b, color='green', s=120, zorder=6, label='b')
# Milestone dots
for t_v in [0.25, 0.5, 0.75]:
pt = exp_map_s1(a, t_v * log_vec)
ax.scatter(*pt, color='orange', s=100, zorder=7)
ax.text(pt[0] + 0.06, pt[1] + 0.06, f't={t_v}', fontsize=8, color='darkorange')
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
ax.set_title('Exp map: walk along tangent vector\nto trace the geodesic')
ax.legend(fontsize=9)
plt.suptitle(r'$Y(t) = \mathrm{Exp}_{a}(t \cdot \mathrm{Log}_{a}(b))$ — the universal geodesic formula',
fontsize=12)
plt.tight_layout()
plt.show()
# Visualise the Log map tangent vector and the geodesic it generates
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
# Left: Log map — tangent vector at p pointing toward q
ax = axes[0]
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.5)
ax.scatter(*a, color='blue', s=120, zorder=6, label='a (base point p)')
ax.scatter(*b, color='green', s=120, zorder=6, label='b (target q)')
# Draw tangent vector (scaled for visibility)
vis_scale = 0.55
log_unit = log_vec / np.linalg.norm(log_vec)
ax.annotate('', xy=a + log_unit * vis_scale, xytext=a,
arrowprops=dict(arrowstyle='->', color='crimson', lw=2.5))
ax.text(a[0] + log_unit[0] * vis_scale + 0.05, a[1] + log_unit[1] * vis_scale + 0.05,
r'Log$_a$(b)', fontsize=11, color='crimson')
# Tangent line
perp = np.array([-a[1], a[0]]) # perpendicular to radius = tangent to circle at a
tang_pts = a[:, None] + perp[:, None] * np.linspace(-0.7, 0.7, 50)
ax.plot(tang_pts[0], tang_pts[1], 'r--', linewidth=1.2, alpha=0.6, label='Tangent plane at a')
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
ax.set_title('Log map: project destination\ninto tangent space at base point')
ax.legend(fontsize=9)
# Right: Exp map — follow the tangent vector to get geodesic
ax = axes[1]
ax.plot(circle[:, 0], circle[:, 1], 'b-', linewidth=2, alpha=0.5)
ax.plot(geod_path[:, 0], geod_path[:, 1], 'g-', linewidth=2.5, label='Geodesic path (SLERP)')
ax.scatter(*a, color='blue', s=120, zorder=6, label='a')
ax.scatter(*b, color='green', s=120, zorder=6, label='b')
# Milestone dots
for t_v in [0.25, 0.5, 0.75]:
pt = exp_map_s1(a, t_v * log_vec)
ax.scatter(*pt, color='orange', s=100, zorder=7)
ax.text(pt[0] + 0.06, pt[1] + 0.06, f't={t_v}', fontsize=8, color='darkorange')
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
ax.set_title('Exp map: walk along tangent vector\nto trace the geodesic')
ax.legend(fontsize=9)
plt.suptitle(r'$Y(t) = \mathrm{Exp}_{a}(t \cdot \mathrm{Log}_{a}(b))$ — the universal geodesic formula',
fontsize=12)
plt.tight_layout()
plt.show()
1.4 — Swiss Roll: LERP Shortcuts Through Empty Space¶
In a VAE, the decoder distorts latent space so it is no longer flat. Two points that look close in latent space may sit on different coil layers of the data manifold — a straight LERP line cuts through a low-density void.
Consequence: decoded outputs along the path are blurry, physically impossible, or semantically wrong.
Metric: distance from each 3-D interpolated point to the nearest real data point.
In [11]:
Copied!
X_roll, t_roll = make_swiss_roll(n_samples=3000, noise=0.1, random_state=42)
idx_sorted_t = np.argsort(t_roll)
inner_pool, outer_pool = idx_sorted_t[:60], idx_sorted_t[-60:]
cross = pairwise_distances(X_roll[inner_pool], X_roll[outer_pool])
i_b, j_b = np.unravel_index(cross.argmin(), cross.shape)
idx_s, idx_e = inner_pool[i_b], outer_pool[j_b]
start, end = X_roll[idx_s], X_roll[idx_e]
lerp_3d = np.outer(1 - ts, start) + np.outer(ts, end)
knn = NearestNeighbors(n_neighbors=1).fit(X_roll)
off_mfd, _ = knn.kneighbors(lerp_3d)
off_mfd = off_mfd.ravel()
print(f'Start: t={t_roll[idx_s]:.2f} End: t={t_roll[idx_e]:.2f}')
print(f'3D Euclidean distance : {np.linalg.norm(start-end):.3f}')
print(f'Max off-manifold dist : {off_mfd.max():.3f} (at t={ts[off_mfd.argmax()]:.2f})')
print('-> LERP cuts through empty latent space!')
X_roll, t_roll = make_swiss_roll(n_samples=3000, noise=0.1, random_state=42)
idx_sorted_t = np.argsort(t_roll)
inner_pool, outer_pool = idx_sorted_t[:60], idx_sorted_t[-60:]
cross = pairwise_distances(X_roll[inner_pool], X_roll[outer_pool])
i_b, j_b = np.unravel_index(cross.argmin(), cross.shape)
idx_s, idx_e = inner_pool[i_b], outer_pool[j_b]
start, end = X_roll[idx_s], X_roll[idx_e]
lerp_3d = np.outer(1 - ts, start) + np.outer(ts, end)
knn = NearestNeighbors(n_neighbors=1).fit(X_roll)
off_mfd, _ = knn.kneighbors(lerp_3d)
off_mfd = off_mfd.ravel()
print(f'Start: t={t_roll[idx_s]:.2f} End: t={t_roll[idx_e]:.2f}')
print(f'3D Euclidean distance : {np.linalg.norm(start-end):.3f}')
print(f'Max off-manifold dist : {off_mfd.max():.3f} (at t={ts[off_mfd.argmax()]:.2f})')
print('-> LERP cuts through empty latent space!')
Start: t=4.88 End: t=13.96 3D Euclidean distance : 18.312 Max off-manifold dist : 4.899 (at t=0.30) -> LERP cuts through empty latent space!
In [12]:
Copied!
fig = plt.figure(figsize=(13, 5))
ax3d = fig.add_subplot(121, projection='3d')
ax3d.scatter(X_roll[:, 0], X_roll[:, 1], X_roll[:, 2],
c=t_roll, cmap='Spectral', s=6, alpha=0.25)
ax3d.scatter(*start, color='blue', s=150, zorder=10, label='Start (inner coil)')
ax3d.scatter(*end, color='green', s=150, zorder=10, label='End (outer coil)')
ax3d.plot(lerp_3d[:, 0], lerp_3d[:, 1], lerp_3d[:, 2],
'r-', linewidth=2.5, label='LERP shortcut')
ax3d.set_title('LERP cuts through the roll')
ax3d.legend(fontsize=8)
ax2d = fig.add_subplot(122)
ax2d.plot(ts, off_mfd, color='firebrick', linewidth=2)
ax2d.fill_between(ts, off_mfd, alpha=0.15, color='firebrick')
ax2d.set_xlabel('Interpolation parameter t')
ax2d.set_ylabel('Distance to nearest manifold point')
ax2d.set_title('Off-manifold distance along LERP\n(0 = on manifold, peak = void)')
plt.tight_layout()
plt.show()
fig = plt.figure(figsize=(13, 5))
ax3d = fig.add_subplot(121, projection='3d')
ax3d.scatter(X_roll[:, 0], X_roll[:, 1], X_roll[:, 2],
c=t_roll, cmap='Spectral', s=6, alpha=0.25)
ax3d.scatter(*start, color='blue', s=150, zorder=10, label='Start (inner coil)')
ax3d.scatter(*end, color='green', s=150, zorder=10, label='End (outer coil)')
ax3d.plot(lerp_3d[:, 0], lerp_3d[:, 1], lerp_3d[:, 2],
'r-', linewidth=2.5, label='LERP shortcut')
ax3d.set_title('LERP cuts through the roll')
ax3d.legend(fontsize=8)
ax2d = fig.add_subplot(122)
ax2d.plot(ts, off_mfd, color='firebrick', linewidth=2)
ax2d.fill_between(ts, off_mfd, alpha=0.15, color='firebrick')
ax2d.set_xlabel('Interpolation parameter t')
ax2d.set_ylabel('Distance to nearest manifold point')
ax2d.set_title('Off-manifold distance along LERP\n(0 = on manifold, peak = void)')
plt.tight_layout()
plt.show()
Part 2 — Pullback Metric (06-pullback-metric.md)¶
Transfer the data-space distance structure back into latent space through the decoder Jacobian.
Sections:
- 2.1 Pullback metric $g_Z(z) = J^T J$ (exercise)
- 2.2 Jacobian as a magnifying glass — unit circles become ellipses
- 2.3 Metric ellipses in latent space
2.1 — Pullback Metric $g_Z(z) = J^T J$ (Exercise 2)¶
Let $\psi: \mathcal{Z} \to \mathcal{X}$ be the decoder. The pullback metric at $z$ is:
$$g_Z(z) = J_{\psi}(z)^T J_{\psi}(z), \qquad J_{\psi}(z) \in \mathbb{R}^{D \times d}$$
What it measures:
- A unit ball in latent space maps via $J$ to an ellipsoid in data space.
- The squared length of a latent step $dz$ in data space is $\|J\,dz\|^2 = dz^T g_Z\,dz$.
- $g_Z(z) = \alpha I$ everywhere $\Longleftrightarrow$ decoder is isometric $\Longleftrightarrow$ LERP = geodesic.
Numerical Jacobian: central differences $J_{ij} \approx \dfrac{f_i(z+\varepsilon e_j) - f_i(z-\varepsilon e_j)}{2\varepsilon}$.
In [13]:
Copied!
def finite_diff_jacobian(
f: Callable[[np.ndarray], np.ndarray],
z: np.ndarray,
eps: float = 1e-4,
) -> np.ndarray:
'''Central-difference Jacobian of f: R^d -> R^D at z. Returns shape (D, d).'''
d, D = len(z), len(f(z))
J = np.zeros((D, d))
for i in range(d):
delta = np.zeros(d)
delta[i] = eps
J[:, i] = (f(z + delta) - f(z - delta)) / (2.0 * eps)
return J
def pullback_metric(
decoder: Callable[[np.ndarray], np.ndarray],
z: np.ndarray,
eps: float = 1e-4,
) -> np.ndarray:
'''
Compute g_Z(z) = J(z)^T J(z).
If g_Z = alpha*I everywhere: LERP in latent space is a geodesic in data space.
Returns (d, d) symmetric positive semi-definite matrix.
'''
# TODO:
# 1. J = finite_diff_jacobian(decoder, z, eps) — shape (D, d)
# 2. Return J.T @ J
# raise NotImplementedError('TODO: implement pullback_metric')
J = finite_diff_jacobian(decoder, z, eps)
return J.T @ J
def finite_diff_jacobian(
f: Callable[[np.ndarray], np.ndarray],
z: np.ndarray,
eps: float = 1e-4,
) -> np.ndarray:
'''Central-difference Jacobian of f: R^d -> R^D at z. Returns shape (D, d).'''
d, D = len(z), len(f(z))
J = np.zeros((D, d))
for i in range(d):
delta = np.zeros(d)
delta[i] = eps
J[:, i] = (f(z + delta) - f(z - delta)) / (2.0 * eps)
return J
def pullback_metric(
decoder: Callable[[np.ndarray], np.ndarray],
z: np.ndarray,
eps: float = 1e-4,
) -> np.ndarray:
'''
Compute g_Z(z) = J(z)^T J(z).
If g_Z = alpha*I everywhere: LERP in latent space is a geodesic in data space.
Returns (d, d) symmetric positive semi-definite matrix.
'''
# TODO:
# 1. J = finite_diff_jacobian(decoder, z, eps) — shape (D, d)
# 2. Return J.T @ J
# raise NotImplementedError('TODO: implement pullback_metric')
J = finite_diff_jacobian(decoder, z, eps)
return J.T @ J
In [14]:
Copied!
# Polar decoder: (r, theta) -> (e^{r/2} cos(theta), e^{r/2} sin(theta))
# Analytical J at z=(0,0): [[0.5, 0], [0, 1]] -> g = diag(0.25, 1)
def decoder_polar(z: np.ndarray) -> np.ndarray:
return np.array([
np.exp(z[0] * 0.5) * np.cos(z[1]),
np.exp(z[0] * 0.5) * np.sin(z[1]),
])
z0 = np.array([0.0, 0.0])
g0 = pullback_metric(decoder_polar, z0)
assert np.allclose(g0, np.array([[0.25, 0.0], [0.0, 1.0]]), atol=1e-5)
print('Exercise 2 - pullback_metric looks good.')
print('g(0,0) =', g0)
print()
for r_val in [-1.0, 0.0, 1.0, 2.0]:
g = pullback_metric(decoder_polar, np.array([r_val, 0.0]))
eigs = np.linalg.eigvalsh(g)
print(f' r={r_val:+.1f}: eigenvalues {eigs} scale ~ {np.sqrt(eigs.mean()):.2f}x')
# Polar decoder: (r, theta) -> (e^{r/2} cos(theta), e^{r/2} sin(theta))
# Analytical J at z=(0,0): [[0.5, 0], [0, 1]] -> g = diag(0.25, 1)
def decoder_polar(z: np.ndarray) -> np.ndarray:
return np.array([
np.exp(z[0] * 0.5) * np.cos(z[1]),
np.exp(z[0] * 0.5) * np.sin(z[1]),
])
z0 = np.array([0.0, 0.0])
g0 = pullback_metric(decoder_polar, z0)
assert np.allclose(g0, np.array([[0.25, 0.0], [0.0, 1.0]]), atol=1e-5)
print('Exercise 2 - pullback_metric looks good.')
print('g(0,0) =', g0)
print()
for r_val in [-1.0, 0.0, 1.0, 2.0]:
g = pullback_metric(decoder_polar, np.array([r_val, 0.0]))
eigs = np.linalg.eigvalsh(g)
print(f' r={r_val:+.1f}: eigenvalues {eigs} scale ~ {np.sqrt(eigs.mean()):.2f}x')
Exercise 2 - pullback_metric looks good. g(0,0) = [[0.25 0. ] [0. 1. ]] r=-1.0: eigenvalues [0.092 0.3679] scale ~ 0.48x r=+0.0: eigenvalues [0.25 1. ] scale ~ 0.79x r=+1.0: eigenvalues [0.6796 2.7183] scale ~ 1.30x r=+2.0: eigenvalues [1.8473 7.3891] scale ~ 2.15x
In [15]:
Copied!
# ── Jacobian as a magnifying glass ────────────────────────────────────────────
# At 3 probe points, draw a unit circle in LATENT space and its image through J
# (an ellipse) in DATA space. This is exactly what g = J^T J captures.
probe_z = [
np.array([-1.0, 0.0]), # small r: decoder barely stretches
np.array([ 0.0, np.pi / 2]), # r=0: baseline scale
np.array([ 1.5, np.pi]), # large r: decoder stretches a lot
]
probe_colors = ['#1f77b4', '#ff7f0e', '#2ca02c']
th_b = np.linspace(0, 2 * np.pi, 120)
unit_ball = np.column_stack([np.cos(th_b), np.sin(th_b)])
draw_s = 0.40
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
for z0, col in zip(probe_z, probe_colors):
J = finite_diff_jacobian(decoder_polar, z0)
x0 = decoder_polar(z0)
g = J.T @ J
evals = np.sort(np.linalg.eigvalsh(g))
scale_label = f'scale ~ {np.sqrt(evals.mean()):.1f}x'
# Latent space: unit circle
circ_lat = z0 + draw_s * unit_ball
axes[0].plot(circ_lat[:, 0], circ_lat[:, 1], '-', color=col, linewidth=2.5,
label=f'z={z0} ({scale_label})')
axes[0].scatter(*z0, color=col, s=80, zorder=7)
# Data space: J maps unit_ball to an ellipse
ellipse_dat = x0 + draw_s * (J @ unit_ball.T).T
axes[1].plot(ellipse_dat[:, 0], ellipse_dat[:, 1], '-', color=col, linewidth=2.5,
label=f'z={z0} ({scale_label})')
axes[1].scatter(*x0, color=col, s=80, zorder=7)
axes[0].set_title('Latent space\n(same-size unit circles at each probe point)')
axes[0].set_xlabel('z[0] (r)')
axes[0].set_ylabel('z[1] (theta)')
axes[0].set_aspect('equal')
axes[0].legend(fontsize=9)
axes[1].set_title('Data space after decoder\n(same circle -> ellipses of DIFFERENT sizes!)')
axes[1].set_xlabel('x[0]')
axes[1].set_ylabel('x[1]')
axes[1].set_aspect('equal')
axes[1].legend(fontsize=9)
plt.suptitle('The Pullback Metric g(z) = J^T J Captures Local Decoder Distortion\n'
'Same unit ball in latent space -> differently-sized ellipses in data space\n'
'g(z) records exactly this shape and size at each point',
fontsize=11)
plt.tight_layout()
plt.show()
# ── Jacobian as a magnifying glass ────────────────────────────────────────────
# At 3 probe points, draw a unit circle in LATENT space and its image through J
# (an ellipse) in DATA space. This is exactly what g = J^T J captures.
probe_z = [
np.array([-1.0, 0.0]), # small r: decoder barely stretches
np.array([ 0.0, np.pi / 2]), # r=0: baseline scale
np.array([ 1.5, np.pi]), # large r: decoder stretches a lot
]
probe_colors = ['#1f77b4', '#ff7f0e', '#2ca02c']
th_b = np.linspace(0, 2 * np.pi, 120)
unit_ball = np.column_stack([np.cos(th_b), np.sin(th_b)])
draw_s = 0.40
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
for z0, col in zip(probe_z, probe_colors):
J = finite_diff_jacobian(decoder_polar, z0)
x0 = decoder_polar(z0)
g = J.T @ J
evals = np.sort(np.linalg.eigvalsh(g))
scale_label = f'scale ~ {np.sqrt(evals.mean()):.1f}x'
# Latent space: unit circle
circ_lat = z0 + draw_s * unit_ball
axes[0].plot(circ_lat[:, 0], circ_lat[:, 1], '-', color=col, linewidth=2.5,
label=f'z={z0} ({scale_label})')
axes[0].scatter(*z0, color=col, s=80, zorder=7)
# Data space: J maps unit_ball to an ellipse
ellipse_dat = x0 + draw_s * (J @ unit_ball.T).T
axes[1].plot(ellipse_dat[:, 0], ellipse_dat[:, 1], '-', color=col, linewidth=2.5,
label=f'z={z0} ({scale_label})')
axes[1].scatter(*x0, color=col, s=80, zorder=7)
axes[0].set_title('Latent space\n(same-size unit circles at each probe point)')
axes[0].set_xlabel('z[0] (r)')
axes[0].set_ylabel('z[1] (theta)')
axes[0].set_aspect('equal')
axes[0].legend(fontsize=9)
axes[1].set_title('Data space after decoder\n(same circle -> ellipses of DIFFERENT sizes!)')
axes[1].set_xlabel('x[0]')
axes[1].set_ylabel('x[1]')
axes[1].set_aspect('equal')
axes[1].legend(fontsize=9)
plt.suptitle('The Pullback Metric g(z) = J^T J Captures Local Decoder Distortion\n'
'Same unit ball in latent space -> differently-sized ellipses in data space\n'
'g(z) records exactly this shape and size at each point',
fontsize=11)
plt.tight_layout()
plt.show()
In [16]:
Copied!
# Riemannian unit balls (1/sqrt(eigenvalue) semi-axes) across the latent grid
grid_z1 = np.linspace(-2, 2, 6)
grid_z2 = np.linspace(-np.pi, np.pi, 7)
scale = 0.5
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
ax = axes[0]
for z1 in grid_z1:
for z2 in grid_z2:
z = np.array([z1, z2])
g = pullback_metric(decoder_polar, z)
evals, evecs = np.linalg.eigh(g)
sa = scale / np.sqrt(max(evals[0], 1e-8))
sb = scale / np.sqrt(max(evals[1], 1e-8))
angle = np.degrees(np.arctan2(evecs[1, 0], evecs[0, 0]))
el = patches.Ellipse(xy=(z1, z2), width=2*sa, height=2*sb, angle=angle,
edgecolor='steelblue', facecolor='none', linewidth=1.5)
ax.add_patch(el)
ax.plot(z1, z2, 'k.', markersize=4)
ax.set_xlim(-3, 3)
ax.set_ylim(-4.5, 4.5)
ax.set_aspect('equal')
ax.set_xlabel('z[0] (r)')
ax.set_ylabel('z[1] (theta)')
ax.set_title('Riemannian unit balls in latent space\n'
'(elongated = distorted; FlatVI makes them circles)')
ax = axes[1]
dense_r = np.linspace(-2, 2, 30)
dense_t = np.linspace(-np.pi, np.pi, 30)
for r_v in dense_r:
pts = np.array([decoder_polar(np.array([r_v, th])) for th in dense_t])
ax.plot(pts[:, 0], pts[:, 1], 'b-', linewidth=0.6, alpha=0.4)
for th_v in dense_t:
pts = np.array([decoder_polar(np.array([r, th_v])) for r in dense_r])
ax.plot(pts[:, 0], pts[:, 1], 'r-', linewidth=0.6, alpha=0.4)
ax.set_aspect('equal')
ax.set_xlabel('x[0]')
ax.set_ylabel('x[1]')
ax.set_title('Decoder output space\n(equal latent grid -> stretched data grid)')
plt.suptitle('Pullback Metric: how the decoder distorts latent space', fontsize=13)
plt.tight_layout()
plt.show()
# Riemannian unit balls (1/sqrt(eigenvalue) semi-axes) across the latent grid
grid_z1 = np.linspace(-2, 2, 6)
grid_z2 = np.linspace(-np.pi, np.pi, 7)
scale = 0.5
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
ax = axes[0]
for z1 in grid_z1:
for z2 in grid_z2:
z = np.array([z1, z2])
g = pullback_metric(decoder_polar, z)
evals, evecs = np.linalg.eigh(g)
sa = scale / np.sqrt(max(evals[0], 1e-8))
sb = scale / np.sqrt(max(evals[1], 1e-8))
angle = np.degrees(np.arctan2(evecs[1, 0], evecs[0, 0]))
el = patches.Ellipse(xy=(z1, z2), width=2*sa, height=2*sb, angle=angle,
edgecolor='steelblue', facecolor='none', linewidth=1.5)
ax.add_patch(el)
ax.plot(z1, z2, 'k.', markersize=4)
ax.set_xlim(-3, 3)
ax.set_ylim(-4.5, 4.5)
ax.set_aspect('equal')
ax.set_xlabel('z[0] (r)')
ax.set_ylabel('z[1] (theta)')
ax.set_title('Riemannian unit balls in latent space\n'
'(elongated = distorted; FlatVI makes them circles)')
ax = axes[1]
dense_r = np.linspace(-2, 2, 30)
dense_t = np.linspace(-np.pi, np.pi, 30)
for r_v in dense_r:
pts = np.array([decoder_polar(np.array([r_v, th])) for th in dense_t])
ax.plot(pts[:, 0], pts[:, 1], 'b-', linewidth=0.6, alpha=0.4)
for th_v in dense_t:
pts = np.array([decoder_polar(np.array([r, th_v])) for r in dense_r])
ax.plot(pts[:, 0], pts[:, 1], 'r-', linewidth=0.6, alpha=0.4)
ax.set_aspect('equal')
ax.set_xlabel('x[0]')
ax.set_ylabel('x[1]')
ax.set_title('Decoder output space\n(equal latent grid -> stretched data grid)')
plt.suptitle('Pullback Metric: how the decoder distorts latent space', fontsize=13)
plt.tight_layout()
plt.show()
What you see¶
What the plots show: The printed eigenvalues of $g = J^T J$ for the polar decoder grow with $r$ — the space stretches more as $r$ increases — and $g(0,0) = \mathrm{diag}(0.25, 1)$ matches the analytic Jacobian. The "magnifying glass" plot maps a small unit circle in latent space to an ellipse in data space at three probe points: the ellipse is larger and more elongated wherever the decoder stretches more. The grid plot draws Riemannian unit balls (semi-axes $1/\sqrt{\text{eigenvalue}}$) across latent space — they vary in both size and orientation from point to point.
Why it looks this way: The pullback metric $g_Z(z) = J(z)^T J(z)$ records how the decoder locally stretches each latent direction into data space. An anisotropic, position-dependent $g$ means a unit step in latent space does not correspond to a unit step in data space — and the conversion factor depends on both where you are and which direction you move.
Key takeaway: Because $g$ is not a constant multiple of the identity, a straight LERP in latent space is not a geodesic in data space: equal latent steps cover unequal data distances. This is precisely the distortion that FlatVI (Part 3) attacks by adding a loss that forces $g_Z \approx \alpha I$, which would make LERP valid again.
Part 3 — FlatVI (07-flatvi.md)¶
Instead of navigating curved geodesics, force the decoder to flatten the space so LERP becomes valid.
Sections:
- 3.1 Flattening loss $\mathcal{L}_{\text{flat}} = \|g_Z(z) - \alpha I\|_F^2$ (exercise)
- 3.2 Loss landscape: where the decoder needs the most correction
- 3.3 Effect of flattening: colored-dot proof that LERP becomes a geodesic
3.1 — FlatVI Flattening Loss (Exercise 3)¶
Key insight: if $g_Z(z) = \alpha I_d$ everywhere, then $\|dz\|_{g_Z}^2 = \alpha\|dz\|^2 = \text{const}$ — straight lines are geodesics.
FlatVI adds a penalty to the VAE ELBO: $$\mathcal{L}_{\text{FlatVI}} = -\text{ELBO} + \lambda\,\mathbb{E}_z\!\left[\|g_Z(z) - \alpha I_d\|_F^2\right]$$
The Frobenius penalty $\|M - \alpha I\|_F^2 = \sum_{i,j}(M_{ij} - \alpha\delta_{ij})^2$ drives the metric toward a scaled identity: isotropic, no preferred direction.
Optimal $\alpha^*$: minimise $\mathbb{E}_z[\|g_Z - \alpha I\|_F^2]$ over $\alpha$ analytically: $\alpha^* = \mathbb{E}_z[\operatorname{tr}(g_Z(z))]/d$.
In [17]:
Copied!
def flattening_loss(M: np.ndarray, alpha: float) -> float:
'''
FlatVI penalty: ||M - alpha * I_d||_F^2
Returns 0.0 iff M = alpha * I (perfectly flat latent space).
Penalises both anisotropy (off-diagonal) and non-uniform scale (diagonal != alpha).
'''
# TODO:
# 1. residual = M - alpha * I_d
# 2. Return sum of squared elements (squared Frobenius norm)
# raise NotImplementedError('TODO: implement flattening_loss')
residual = M - alpha * np.eye(M.shape[0])
return float(np.sum(residual ** 2))
def optimal_alpha(
decoder: Callable[[np.ndarray], np.ndarray],
z_samples: np.ndarray,
) -> float:
'''alpha* = E[tr(g(z))] / d (minimises E[||g - alpha*I||_F^2]).'''
d = z_samples.shape[1]
traces = [np.trace(pullback_metric(decoder, z)) for z in z_samples]
return float(np.mean(traces) / d)
def flattening_loss(M: np.ndarray, alpha: float) -> float:
'''
FlatVI penalty: ||M - alpha * I_d||_F^2
Returns 0.0 iff M = alpha * I (perfectly flat latent space).
Penalises both anisotropy (off-diagonal) and non-uniform scale (diagonal != alpha).
'''
# TODO:
# 1. residual = M - alpha * I_d
# 2. Return sum of squared elements (squared Frobenius norm)
# raise NotImplementedError('TODO: implement flattening_loss')
residual = M - alpha * np.eye(M.shape[0])
return float(np.sum(residual ** 2))
def optimal_alpha(
decoder: Callable[[np.ndarray], np.ndarray],
z_samples: np.ndarray,
) -> float:
'''alpha* = E[tr(g(z))] / d (minimises E[||g - alpha*I||_F^2]).'''
d = z_samples.shape[1]
traces = [np.trace(pullback_metric(decoder, z)) for z in z_samples]
return float(np.mean(traces) / d)
In [18]:
Copied!
I2 = np.eye(2)
assert np.isclose(flattening_loss(2.0 * I2, alpha=2.0), 0.0)
assert np.isclose(flattening_loss(2.0 * I2, alpha=1.0), 2.0)
M_aniso = np.diag([2.0, 0.5])
expected = 1.0 ** 2 + 0.5 ** 2
assert np.isclose(flattening_loss(M_aniso, alpha=1.0), expected)
print('Exercise 3 - flattening_loss looks good.')
print(f' flattening_loss(diag(2, 0.5), alpha=1) = {flattening_loss(M_aniso,1.0):.4f} (expected {expected})')
I2 = np.eye(2)
assert np.isclose(flattening_loss(2.0 * I2, alpha=2.0), 0.0)
assert np.isclose(flattening_loss(2.0 * I2, alpha=1.0), 2.0)
M_aniso = np.diag([2.0, 0.5])
expected = 1.0 ** 2 + 0.5 ** 2
assert np.isclose(flattening_loss(M_aniso, alpha=1.0), expected)
print('Exercise 3 - flattening_loss looks good.')
print(f' flattening_loss(diag(2, 0.5), alpha=1) = {flattening_loss(M_aniso,1.0):.4f} (expected {expected})')
Exercise 3 - flattening_loss looks good. flattening_loss(diag(2, 0.5), alpha=1) = 1.2500 (expected 1.25)
3.2 — Flattening Loss Landscape¶
The loss $\|g_Z(z) - \alpha^* I\|_F^2$ shows where in latent space the decoder is most curved. FlatVI's training gradient pushes these hot spots toward zero.
In [19]:
Copied!
z_samples = rng.uniform(low=[-2, -np.pi], high=[2, np.pi], size=(120, 2))
alpha_opt = optimal_alpha(decoder_polar, z_samples)
print(f'Optimal alpha: {alpha_opt:.4f}')
n_grid = 22
r_vals = np.linspace(-2, 2, n_grid)
th_vals = np.linspace(-np.pi, np.pi, n_grid)
Rg, Thg = np.meshgrid(r_vals, th_vals)
loss_grid = np.zeros((n_grid, n_grid))
for i, th in enumerate(th_vals):
for j, r in enumerate(r_vals):
g = pullback_metric(decoder_polar, np.array([r, th]))
loss_grid[i, j] = flattening_loss(g, alpha=alpha_opt)
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
ax = axes[0]
cf = ax.contourf(Rg, Thg, np.log1p(loss_grid), levels=20, cmap='Reds')
plt.colorbar(cf, ax=ax, label='log(1 + flattening loss)')
ax.set_xlabel('z[0] (r)')
ax.set_ylabel('z[1] (theta)')
ax.set_title('Hot = highly curved\nFlatVI training drives these toward zero')
ax = axes[1]
ax.plot(r_vals, loss_grid[n_grid // 2, :], color='firebrick', linewidth=2)
ax.set_xlabel('z[0] (r)')
ax.set_ylabel('Flattening loss')
ax.set_title('Loss grows exponentially with r\n(decoder stretches more the further from origin)')
plt.suptitle('Flattening Loss Landscape', fontsize=13)
plt.tight_layout()
plt.show()
z_samples = rng.uniform(low=[-2, -np.pi], high=[2, np.pi], size=(120, 2))
alpha_opt = optimal_alpha(decoder_polar, z_samples)
print(f'Optimal alpha: {alpha_opt:.4f}')
n_grid = 22
r_vals = np.linspace(-2, 2, n_grid)
th_vals = np.linspace(-np.pi, np.pi, n_grid)
Rg, Thg = np.meshgrid(r_vals, th_vals)
loss_grid = np.zeros((n_grid, n_grid))
for i, th in enumerate(th_vals):
for j, r in enumerate(r_vals):
g = pullback_metric(decoder_polar, np.array([r, th]))
loss_grid[i, j] = flattening_loss(g, alpha=alpha_opt)
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
ax = axes[0]
cf = ax.contourf(Rg, Thg, np.log1p(loss_grid), levels=20, cmap='Reds')
plt.colorbar(cf, ax=ax, label='log(1 + flattening loss)')
ax.set_xlabel('z[0] (r)')
ax.set_ylabel('z[1] (theta)')
ax.set_title('Hot = highly curved\nFlatVI training drives these toward zero')
ax = axes[1]
ax.plot(r_vals, loss_grid[n_grid // 2, :], color='firebrick', linewidth=2)
ax.set_xlabel('z[0] (r)')
ax.set_ylabel('Flattening loss')
ax.set_title('Loss grows exponentially with r\n(decoder stretches more the further from origin)')
plt.suptitle('Flattening Loss Landscape', fontsize=13)
plt.tight_layout()
plt.show()
Optimal alpha: 0.9786
3.3 — Effect of Flattening: LERP Becomes a Geodesic¶
Claim: if $g_Z = \alpha I$ everywhere, LERP in latent space produces constant-speed paths in data space.
Proof:
$$\left\|\frac{d}{dt}\psi(z(t))\right\|^2 = (z_1{-}z_0)^T g_Z(z(t))\,(z_1{-}z_0) = \alpha\,\|z_1{-}z_0\|^2 = \text{const}$$
Visual proof: the colored dots below are placed at 11 equally-spaced $t$ values. If the dots are equally spaced in data space, LERP is a geodesic. For the curved decoder the dots bunch up; for the flat decoder they are perfectly uniform.
In [20]:
Copied!
def decoder_flat(z: np.ndarray) -> np.ndarray:
'''Linear (flat) decoder: identity — g_Z(z) = I everywhere.'''
return np.array([z[0], z[1]])
z_s = np.array([-1.5, -np.pi / 3])
z_e = np.array([ 1.5, np.pi / 3])
lerp_z = np.outer(1 - ts, z_s) + np.outer(ts, z_e)
path_curved = np.array([decoder_polar(z) for z in lerp_z])
path_flat = np.array([decoder_flat(z) for z in lerp_z])
speed_curved = np.linalg.norm(np.diff(path_curved, axis=0), axis=1)
speed_flat = np.linalg.norm(np.diff(path_flat, axis=0), axis=1)
# 11 equal-t milestone dots (dark purple -> yellow in plasma colormap)
t_dots = np.linspace(0, 1, 11)
cmap_tdot = plt.cm.plasma
dot_indices = [int(round(t * (len(ts) - 1))) for t in t_dots]
fig, axes = plt.subplots(1, 3, figsize=(17, 5))
# ── Curved decoder path ──────────────────────────────────────────────────────
ax = axes[0]
fine_r = np.linspace(-2, 2, 25)
fine_th = np.linspace(-np.pi, np.pi, 25)
for r_v in fine_r:
pts = np.array([decoder_polar(np.array([r_v, th])) for th in fine_th])
ax.plot(pts[:, 0], pts[:, 1], 'b-', linewidth=0.5, alpha=0.25)
for th_v in fine_th:
pts = np.array([decoder_polar(np.array([r, th_v])) for r in fine_r])
ax.plot(pts[:, 0], pts[:, 1], 'r-', linewidth=0.5, alpha=0.25)
ax.plot(path_curved[:, 0], path_curved[:, 1], '-', color='gray', linewidth=1.2, alpha=0.6)
for idx, t_d in zip(dot_indices, t_dots):
ax.scatter(*path_curved[idx], color=cmap_tdot(t_d), s=130, zorder=9,
edgecolors='white', linewidth=1.2)
ax.scatter(*path_curved[0], color='black', s=120, marker='s', zorder=10)
ax.scatter(*path_curved[-1], color='black', s=120, marker='^', zorder=10)
ax.set_aspect('equal')
ax.set_title('Curved decoder (polar)\nDots bunch at small-r end <- non-geodesic!')
# ── Flat decoder path ─────────────────────────────────────────────────────────
ax = axes[1]
for r_v in fine_r:
pts = np.array([decoder_flat(np.array([r_v, th])) for th in fine_th])
ax.plot(pts[:, 0], pts[:, 1], 'b-', linewidth=0.5, alpha=0.25)
for th_v in fine_th:
pts = np.array([decoder_flat(np.array([r, th_v])) for r in fine_r])
ax.plot(pts[:, 0], pts[:, 1], 'r-', linewidth=0.5, alpha=0.25)
ax.plot(path_flat[:, 0], path_flat[:, 1], '-', color='gray', linewidth=1.2, alpha=0.6)
for idx, t_d in zip(dot_indices, t_dots):
ax.scatter(*path_flat[idx], color=cmap_tdot(t_d), s=130, zorder=9,
edgecolors='white', linewidth=1.2)
ax.scatter(*path_flat[0], color='black', s=120, marker='s', zorder=10)
ax.scatter(*path_flat[-1], color='black', s=120, marker='^', zorder=10)
ax.set_aspect('equal')
ax.set_title('Flat decoder (identity)\nDots equally spaced <- geodesic!')
# ── Speed comparison ──────────────────────────────────────────────────────────
ax = axes[2]
ax.plot(ts[1:], speed_curved, 'r-', linewidth=2.5, label='Curved: non-uniform')
ax.plot(ts[1:], speed_flat, 'g-', linewidth=2.5, label='Flat: constant')
ax.set_xlabel('Interpolation parameter t')
ax.set_ylabel('||dx/dt|| (step size in data space)')
ax.set_title('Speed profile\n(flat = constant = proof of geodesic)')
ax.legend()
# Shared colorbar for the dot colors
sm = plt.cm.ScalarMappable(cmap=cmap_tdot, norm=plt.Normalize(0, 1))
sm.set_array([])
plt.colorbar(sm, ax=axes[1], label='t (dark=0 -> yellow=1)', shrink=0.8)
plt.suptitle('FlatVI Effect: Equal-t Dots are Equally Spaced ONLY in Flat Latent Space',
fontsize=12)
plt.tight_layout()
plt.show()
cv_c = speed_curved.std() / speed_curved.mean()
cv_f = speed_flat.std() / speed_flat.mean()
print('Speed coefficient of variation (0 = geodesic):')
print(f' Curved decoder : {cv_c:.4f}')
print(f' Flat decoder : {cv_f:.2e}')
def decoder_flat(z: np.ndarray) -> np.ndarray:
'''Linear (flat) decoder: identity — g_Z(z) = I everywhere.'''
return np.array([z[0], z[1]])
z_s = np.array([-1.5, -np.pi / 3])
z_e = np.array([ 1.5, np.pi / 3])
lerp_z = np.outer(1 - ts, z_s) + np.outer(ts, z_e)
path_curved = np.array([decoder_polar(z) for z in lerp_z])
path_flat = np.array([decoder_flat(z) for z in lerp_z])
speed_curved = np.linalg.norm(np.diff(path_curved, axis=0), axis=1)
speed_flat = np.linalg.norm(np.diff(path_flat, axis=0), axis=1)
# 11 equal-t milestone dots (dark purple -> yellow in plasma colormap)
t_dots = np.linspace(0, 1, 11)
cmap_tdot = plt.cm.plasma
dot_indices = [int(round(t * (len(ts) - 1))) for t in t_dots]
fig, axes = plt.subplots(1, 3, figsize=(17, 5))
# ── Curved decoder path ──────────────────────────────────────────────────────
ax = axes[0]
fine_r = np.linspace(-2, 2, 25)
fine_th = np.linspace(-np.pi, np.pi, 25)
for r_v in fine_r:
pts = np.array([decoder_polar(np.array([r_v, th])) for th in fine_th])
ax.plot(pts[:, 0], pts[:, 1], 'b-', linewidth=0.5, alpha=0.25)
for th_v in fine_th:
pts = np.array([decoder_polar(np.array([r, th_v])) for r in fine_r])
ax.plot(pts[:, 0], pts[:, 1], 'r-', linewidth=0.5, alpha=0.25)
ax.plot(path_curved[:, 0], path_curved[:, 1], '-', color='gray', linewidth=1.2, alpha=0.6)
for idx, t_d in zip(dot_indices, t_dots):
ax.scatter(*path_curved[idx], color=cmap_tdot(t_d), s=130, zorder=9,
edgecolors='white', linewidth=1.2)
ax.scatter(*path_curved[0], color='black', s=120, marker='s', zorder=10)
ax.scatter(*path_curved[-1], color='black', s=120, marker='^', zorder=10)
ax.set_aspect('equal')
ax.set_title('Curved decoder (polar)\nDots bunch at small-r end <- non-geodesic!')
# ── Flat decoder path ─────────────────────────────────────────────────────────
ax = axes[1]
for r_v in fine_r:
pts = np.array([decoder_flat(np.array([r_v, th])) for th in fine_th])
ax.plot(pts[:, 0], pts[:, 1], 'b-', linewidth=0.5, alpha=0.25)
for th_v in fine_th:
pts = np.array([decoder_flat(np.array([r, th_v])) for r in fine_r])
ax.plot(pts[:, 0], pts[:, 1], 'r-', linewidth=0.5, alpha=0.25)
ax.plot(path_flat[:, 0], path_flat[:, 1], '-', color='gray', linewidth=1.2, alpha=0.6)
for idx, t_d in zip(dot_indices, t_dots):
ax.scatter(*path_flat[idx], color=cmap_tdot(t_d), s=130, zorder=9,
edgecolors='white', linewidth=1.2)
ax.scatter(*path_flat[0], color='black', s=120, marker='s', zorder=10)
ax.scatter(*path_flat[-1], color='black', s=120, marker='^', zorder=10)
ax.set_aspect('equal')
ax.set_title('Flat decoder (identity)\nDots equally spaced <- geodesic!')
# ── Speed comparison ──────────────────────────────────────────────────────────
ax = axes[2]
ax.plot(ts[1:], speed_curved, 'r-', linewidth=2.5, label='Curved: non-uniform')
ax.plot(ts[1:], speed_flat, 'g-', linewidth=2.5, label='Flat: constant')
ax.set_xlabel('Interpolation parameter t')
ax.set_ylabel('||dx/dt|| (step size in data space)')
ax.set_title('Speed profile\n(flat = constant = proof of geodesic)')
ax.legend()
# Shared colorbar for the dot colors
sm = plt.cm.ScalarMappable(cmap=cmap_tdot, norm=plt.Normalize(0, 1))
sm.set_array([])
plt.colorbar(sm, ax=axes[1], label='t (dark=0 -> yellow=1)', shrink=0.8)
plt.suptitle('FlatVI Effect: Equal-t Dots are Equally Spaced ONLY in Flat Latent Space',
fontsize=12)
plt.tight_layout()
plt.show()
cv_c = speed_curved.std() / speed_curved.mean()
cv_f = speed_flat.std() / speed_flat.mean()
print('Speed coefficient of variation (0 = geodesic):')
print(f' Curved decoder : {cv_c:.4f}')
print(f' Flat decoder : {cv_f:.2e}')
Speed coefficient of variation (0 = geodesic): Curved decoder : 0.4251 Flat decoder : 1.99e-15
Summary¶
Part 1 — Geodesic (05)¶
| Concept | Key takeaway |
|---|---|
| 1-D motivation | Non-linear decoder maps equal latent steps to unequal data gaps |
| LERP failure | Straight line in latent space leaves the manifold (midpoint norm < 1 on $S^1$) |
| SLERP | Constant angular speed on $S^{n-1}$; equal-$t$ dots are equally spaced |
| Exp / Log maps | $Y(t) = \text{Exp}_{Y_0}(t\cdot\text{Log}_{Y_0}(Y_1))$ — universal geodesic formula |
| Swiss Roll | LERP shortcut traverses empty space; off-manifold distance peaks at midpoint |
Part 2 — Pullback Metric (06)¶
| Concept | Key takeaway |
|---|---|
| $g_Z = J^T J$ | Transfers Euclidean distances from data space back to latent space |
| Magnifying glass | Same unit ball in latent space → differently-sized ellipses in data space |
| Flatness condition | $g_Z = \alpha I$ everywhere $\Leftrightarrow$ LERP = geodesic |
Part 3 — FlatVI (07)¶
| Concept | Key takeaway |
|---|---|
| Flattening loss | $\|g_Z - \alpha I\|_F^2$ penalises anisotropy and non-uniform scale |
| Visual proof | Equal-$t$ dots equally spaced only for flat decoder |
| Advantage | No Neural ODE at inference; downstream tools (kNN, OT) work without modification |
The chain: LERP fails (geodesic problem) $\to$ $g_Z = J^TJ$ diagnoses why $\to$ FlatVI fixes it by making $g_Z \approx \alpha I$.