Iso-Riemannian Optimization on Learned Data Manifolds

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Iso-Riemannian Optimization on Learned Data Manifolds".

Jane: The paper was written by Willem Diepeveen and Melanie Weber from University of California, Los Angeles and Harvard University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back, everyone. Today we're looking at a paper that's been making waves in the optimization community, and it's got a mouthful of a title: "Iso-Riemannian Optimization on Learned Data Manifolds." Jane, what do we even mean when we say "learned data manifolds"?

Jane: Great question, Tom. So imagine your data isn't spread out evenly across all of space, but instead it lives on a kind of curved surface, like a sheet of paper that's been crumpled. That surface is the manifold. And in this paper, they don't just assume the surface is given—they learn it from the data itself using something called a pullback geometry, which is a fancy way of saying they find a transformation that flattens the data out.

Tom: And that's where the "iso" part comes in, right? Because once you flatten things, you want your optimization to respect that flattened geometry?

Jane: Exactly. The "iso" refers to the iso-connection, which is a way of measuring distances and moving along the manifold so that you're moving at constant speed in the original space. The authors—Willem Diepeveen from UCLA and Melanie Weber from Harvard—argue that this constant-speed property is what makes optimization on these learned manifolds actually work.

Lu: And I think that's the key insight. In classical Riemannian optimization, you assume the geometry is fixed and you use the Levi-Civita connection, which is the natural way to measure geodesics. But when the geometry is learned from data, that natural connection can distort things. The iso-connection fixes that distortion by ensuring that when you move along a geodesic, you're moving at a constant Euclidean speed.

Tom: So it's like taking a scenic route versus a highway? The scenic route might be more interesting, but the highway gets you there predictably?

Lu: That's a good way to put it. The iso-connection is the highway. It guarantees that the distance you travel in the ambient space is proportional to the parameter along the curve, which makes convergence analysis much cleaner.

Meng: But let me play devil's advocate here. Does this actually matter in practice? I mean, if you're running gradient descent on a manifold, you just need the exponential map, right?

Jane: That's the thing, Meng. The exponential map under the Levi-Civita connection can be expensive and unstable when the geometry is learned. But under the iso-connection, the exponential map has a closed-form expression in terms of the pullback geometry. That means you can compute it efficiently, even in high dimensions.

Tom: And that's what makes this paper exciting. They're not just proposing a theoretical framework—they're showing it works on real data, including MNIST digits. We'll get into the details of their experiments in a bit, but the headline is that this iso-Riemannian approach gives you convergence guarantees that classical Riemannian optimization can't provide when the geometry is learned.

Lu: And that's the gap they're filling. Previous work on optimization over learned manifolds either assumed the manifold was linear or used chart-based approaches that lose global structure. This paper says, let's use the geometry itself, but with a connection that's actually compatible with the Euclidean structure of the problem.

Jane: So the title really captures it. It's about doing optimization on manifolds that are learned from data, but doing it in a way that respects the isometry—the constant-speed property—of the geometry.

Tom: And that's the hook for the rest of our discussion. Next, we're going to look at what they actually propose in terms of new definitions for monotonicity and Lipschitz continuity, and why those are the right tools for this setting.

Summary: Tom: So we've set the stage with the title. Now let's talk about what this paper actually does. Jane, can you walk us through the core ideas?

Jane: Sure, Tom. The paper introduces three new notions: iso-monotonicity, iso-Lipschitz continuity, and iso-convexity. These are analogues of the classical concepts used in optimization, but they're defined using the iso-connection instead of the Levi-Civita connection.

Tom: And why do we need new definitions? Why can't we just use the old ones?

Jane: Because the old ones assume the geometry is fixed and well-behaved. But when the geometry is learned from data, a function that's convex in the Euclidean sense might not be geodesically convex in the Riemannian sense. The paper shows a simple example: the function f(x) = x2/two under a pullback metric with φ(x) = sinh(x+one) is not geodesically convex. That means Riemannian gradient descent can fail or converge very slowly.

Lu: And that's the fundamental problem. The critical points of a function don't depend on the geometry, but the region of convergence does. So even if you have a nice convex function, the learned geometry can make it look non-convex from the Riemannian perspective.

Meng: So they're essentially saying, let's change the connection so that the geometry doesn't fight against the Euclidean structure of the problem?

Jane: Exactly. And they prove that with these new definitions, you can get linear convergence for a descent algorithm they call iso-Riemannian descent. The algorithm is simple: you take a step along the iso-exponential map in the direction of the negative vector field, with a step size that depends on the iso-monotonicity and iso-Lipschitz constants.

Tom: And they have a theorem that guarantees convergence, right? Theorem three in the paper?

Jane: Yes. If the vector field is α-strongly iso-monotone and L-iso-Lipschitz, and you choose a step size r less than 2α/L2, then the sequence converges linearly to the unique zero of the vector field. That's a direct analogue of the classical result, but now it applies in this learned geometry setting.

Lu: And what's nice is that they show how to verify these conditions locally. You don't have to check every pair of points—you just need to check the iso-connection along geodesics. That's Theorem two and it makes the framework practical.

Meng: But hold on. How do you actually compute these iso-manifold mappings? The iso-exponential map, the iso-logarithmic map—those sound like they could be expensive.

Jane: That's the beauty of the pullback structure. For Euclidean pullback manifolds, the iso-manifold mappings have closed-form expressions in terms of the diffeomorphism φ. So you just apply φ, do Euclidean operations, and apply φ−1. That's cheap and scalable.

Tom: And that's why they can apply this to real data sets like MNIST. But before we get to the experiments, let me ask about the applications. What do they actually use this for?

Jane: Two main applications. First, they define the iso-barycentre, which is a generalization of the Riemannian barycentre. It's the point that minimizes the sum of squared iso-distances to the data points. And they show it has nice properties—for example, on a 1D manifold, the iso-barycentre field is one-strongly iso-monotone and one-iso-Lipschitz, which guarantees convergence of their algorithm.

Lu: And the second application is iso-convex optimization, where they show that the l2-projected gradient field of a Euclidean convex function can be iso-monotone under certain conditions. That's the key to solving inverse problems on learned manifolds.

Tom: So we've got the theory. Next, we need to talk about what they actually improved compared to existing methods, and that's where the experiments come in.

Improvements: Tom: Alright, we've covered the theory. Now let's talk about what this paper actually improves in practice. Jane, what did they show?

Jane: The biggest improvement is in clustering. They introduce something called iso-K-means, which is a variant of K-means that uses iso-barycentres instead of Euclidean or Riemannian barycentres. And the results are striking.

Tom: I remember the figures. On the synthetic river and spiral data sets, Euclidean K-means and even Riemannian K-means give weird clusters. But iso-K-means finds both correct clusters and meaningful centroids.

Jane: Exactly. The reason is that the pullback metric doesn't encode our intuition about Euclidean proximity. So when you use the Riemannian distance for clustering, you can get clusters that look wrong. The iso-distance corrects for that by measuring distances along constant-speed geodesics.

Lu: And that's a real improvement. The paper shows that the iso-barycentre field on a 1D manifold has the same strong monotonicity and Lipschitz constants as the classical Riemannian barycentre field. So you get the same convergence guarantees, but with a more meaningful notion of distance.

Meng: But what about the higher-dimensional case? The paper admits that the constants aren't necessarily close to one in practice. How do they handle that?

Jane: They use a line search. Algorithm one combines iso-Riemannian descent with a backtracking line search, initialized at the Riemannian barycentre. That ensures the starting point is within the strongly convex geodesic submanifold containing the data.

Tom: And it works on MNIST too, right? They compute the iso-barycentre of the digit data and it looks like an eight which is interpretable.

Jane: Yes, and the Euclidean mean is just a blurry mess. The Riemannian and iso-barycentres both give something that looks like an eight but the iso-barycentre is computed with the constant-speed property, which makes it more robust.

Lu: The other improvement is in solving inverse problems. They show that for a 1D manifold, if the function is strongly convex on the tangent bundle and the geometry doesn't introduce too much non-convexity, then the l2-projected gradient field is iso-monotone and iso-Lipschitz. That gives you convergence of their algorithm to the unique minimizer.

Meng: And they demonstrate this on denoising MNIST digits. They constrain the optimization to a twenty-dimensional subspace of the learned manifold and show that their algorithm converges to a nearly noiseless representation in just ten iterations.

Jane: That's the practical payoff. The theory tells you when and why the algorithm works, and the experiments show that it actually works in high-dimensional settings.

Tom: So the improvements are threefold: new theoretical tools, better clustering, and efficient inverse problem solving. But there are also limitations, right? The paper mentions numerical issues with the MNIST experiments.

Jane: Yes, they note that the combination of single precision in PyTorch and numerical approximations for the iso-manifold mappings can cause the algorithm to stall when seeking high accuracy. That's a practical limitation, but it's not a fundamental one.

Lu: And they also note that convergence of iso-K-means is still open. The Lloyd's algorithm adaptation doesn't inherit convergence guarantees from the classical setting. So there's room for future work.

Tom: That's a good segue. Let's bring in Lalam to talk about the bigger picture. Lalam, what do you see as the most impactful vision from this paper?

Conclusion: Tom: We've covered the theory, the applications, and the improvements. Let's wrap up with the big picture. Lalam, what's your take?

Lalam: I see this paper as a bridge between two worlds that have been disconnected for too long. On one side, you have the Riemannian optimization community, which has developed powerful tools for fixed geometries. On the other side, you have the machine learning community, which learns geometries from data but doesn't have the optimization tools to use them effectively. This paper builds that bridge.

Jane: And the bridge is the iso-connection. It's a simple idea—ensure constant-speed geodesics—but it has profound consequences for optimization.

Lalam: Exactly. And the cultural impact is significant. Think about applications like medical imaging, where you want to solve inverse problems on manifolds learned from patient data. Or recommendation systems, where the geometry of user preferences is learned from interaction data. This framework gives you provable convergence guarantees in those settings.

Meng: But what about the practical challenges? The numerical issues they mention—are those going to be a barrier?

Lalam: They're a barrier for high-precision applications, but not for the core methodology. As the paper notes, higher-order numerical approximations and better precision handling can mitigate these issues. And the fact that the algorithm converges in ten iterations on MNIST denoising shows that the approach is already practical.

Tom: So let me try to summarize. The paper "Iso-Riemannian Optimization on Learned Data Manifolds" introduces new notions of convexity, monotonicity, and Lipschitz continuity based on the iso-connection. It proves convergence of iso-Riemannian descent, demonstrates applications in barycentre computation and clustering, and shows how to solve inverse problems on learned manifolds.

Jane: And the key takeaway is that when you learn a geometry from data, you need to be careful about which connection you use for optimization. The iso-connection reconciles the learned geometry with Euclidean convexity, making optimization both theoretically sound and practically efficient.

Lu: I'd add that this opens up a whole research agenda. The paper mentions future work on retractions, higher-order methods, and non-smooth optimization. There's a lot of room to build on these foundations.

Meng: And from an engineering standpoint, the fact that the manifold mappings have closed-form expressions for pullback geometries makes this implementable today. That's a big deal.

Lalam: The cultural impact is that we can now build systems that respect the intrinsic geometry of data while maintaining the guarantees that make optimization reliable. That's a step toward more interpretable and trustworthy AI systems.

Tom: Well said. We've had a great discussion about "Iso-Riemannian Optimization on Learned Data Manifolds" by Willem Diepeveen and Melanie Weber. Thanks to everyone for joining us—Lu, Meng, Lalam. And thanks to our listeners for tuning in. We'll be back with the next paper soon. Until then, keep exploring the geometry of your data.