🤖 machine learning

Smooth Reparameterizations of Functions on Simplicial Product Spaces: Applications to Probabilistic Tensor Decomposition and Functional Data Registration

This paper introduces a smooth, strictly convex reparameterization of product simplex spaces that transforms constrained optimization problems into unconstrained manifold problems, enabling a Riemannian Gradient Descent algorithm that outperforms Projected Gradient Descent in applications like probabilistic tensor decomposition and functional data registration.

Original authors: Shashwat Kumar, Arafat Rahman, Anuj Srivastava, P. -A. Absil

Published 2026-08-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shashwat Kumar, Arafat Rahman, Anuj Srivastava, P. -A. Absil

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant, multi-dimensional puzzle where every piece has to fit perfectly into a specific shape. In the world of data science, this puzzle often involves organizing information into "simplexes." Think of a simplex as a strict rulebook for a team of players: every player must have a non-negative score, and when you add up all their scores, the total must equal exactly one. It's like a pie chart where the slices can't be negative, and the whole pie must always be 100%. This rulebook shows up everywhere, from figuring out how different minerals mix in a rock to aligning medical scans of moving body parts.

The challenge is that these strict rules make the puzzle incredibly hard to solve with standard tools. It's like trying to walk a tightrope while juggling; if you take a step in the wrong direction, you fall off the edge, and you have to be dragged back to the rope before you can try again. This "dragging back" process, known as projection, is slow and can distort the path you're trying to take. Scientists have long wondered if there was a way to smooth out the tightrope itself, turning the jagged, rule-bound path into a gentle, rolling hill where you can just roll forward without ever falling off. This paper explores exactly that idea: can we reshape the rules of the game so the math becomes easier, without changing the actual answer we're looking for?

The authors of this paper, Shashwat Kumar and his colleagues, say "yes," but with a very specific twist. They propose a clever trick called "smooth reparameterization." Instead of forcing the data to stay on the strict simplex (the pie chart with fixed rules), they invent a new set of variables that live on a smooth, round sphere. Imagine taking that flat, jagged pie chart and stretching it out over the surface of a perfect ball. On this ball, there are no sharp edges or hard walls; you can move in any direction, and the math flows naturally.

The paper demonstrates that this transformation is safe. It proves that if you find a "sweet spot" (a mathematical optimum) on this smooth ball, it corresponds perfectly to a valid solution on the original, strict simplex. They show that the "second-order" conditions—which are like checking if a hill is truly a valley and not just a flat spot—work just as well on the smooth ball as they do on the strict simplex. Specifically, they prove that the second-order critical points on the smooth manifold map to the weak second-order KKT points on the simplex, ensuring the solutions align correctly.

To test this, the team applied their new method to two real-world problems. First, they tackled "tensor decomposition," which is like trying to break down a complex 3D block of data (think of a stack of pie charts) into its simplest, underlying ingredients. They found that their new method, called Riemannian Gradient Descent (RGD), solved this puzzle much faster and more accurately than the old "drag-and-drop" method (Projected Gradient Descent). In their simulations, the new method often outperformed the old one by several orders of magnitude, meaning it reached the solution with far fewer steps.

Second, they used the method for "functional data registration," which is like trying to line up a group of people running a race so you can compare their strides, even if some are running fast and others slow. The goal is to stretch or shrink the time axis of each runner so they all match up. The old method often produced jerky, unnatural alignments that looked like a robot trying to dance. The new smooth method, however, produced fluid, natural-looking alignments that preserved the true shape of the data.

The paper doesn't just suggest this works; it provides mathematical proofs that the critical points (the best solutions) on the smooth sphere map directly to the valid solutions on the simplex. They also show that while the old method sometimes gets stuck or creates jagged results, the new method maintains the smoothness of the original data shapes. The authors conclude that by trading the rigid rules of the simplex for the smooth freedom of a sphere, we can solve these complex data puzzles more efficiently and with greater fidelity, making it a powerful new tool for anyone working with probability distributions or aligning time-based data.

Problem Statement
The paper addresses a class of optimization problems defined on product spaces of simplices and Euclidean components, denoted as $F_{\text{simplex}} = \Delta_{n_1} \times \dots \times \Delta_{n_L} \times \mathbb{R}^{n_{L+1}} \times \dots \times \mathbb{R}^{n_{L+Q}}$ . Such problems arise in diverse domains including simplex-constrained matrix factorization, probabilistic tensor decomposition (learning low-rank discrete multivariate distributions), and functional data registration using the Square Root Velocity Function (SRVF) representation.

Standard approaches to solving these problems, such as Projected Gradient Descent (PGD), involve moving along the gradient direction and projecting back onto the feasible set $F_{\text{simplex}}$ . While effective, projection can introduce discontinuities and may not preserve the smoothness of the underlying function shapes, particularly in applications like curve registration.

Methodology
The authors propose replacing the constrained optimization problem on the product simplex with an unconstrained optimization problem on a smooth manifold via a smooth, elementwise reparameterization.

Reparameterization: The feasible set is mapped from a smooth manifold $F_{\text{smooth}}$ to $F_{\text{simplex}}$ using a transformation $\phi(\theta)$ . For the simplex components, this transformation is defined elementwise by a function $\psi(\theta)$ , such that $v = \psi(\theta)$ . The function $\psi: \mathbb{R} \to \mathbb{R}_{\geq 0}$ must satisfy specific properties: $\psi(0)=0$ , $\psi'' > 0$ (strictly convex), and $\psi'(\theta)=0 \iff \theta=0$ .
- This generalizes the well-known Hadamard reparameterization ( $\psi(\theta) = \theta^2$ ) to include other functions like even-degree polynomials with positive coefficients or $\psi(\theta) = \cosh(\theta) - 1$ .
- The resulting problem is $\min_{\theta \in F_{\text{smooth}}} f(\phi(\theta))$ , where $F_{\text{smooth}}$ is a Riemannian manifold defined by the constraints $1^T \psi(\theta_l) = 1$ .
Theoretical Analysis (KKT Conditions): The core theoretical contribution is a rigorous analysis of the relationship between the Karush–Kuhn–Tucker (KKT) points of the original constrained problem ( $P_{\text{simplex}}$ ) and the reparameterized unconstrained problem ( $P_{\times\text{smooth}}$ ).
- The authors distinguish between "weak" and "strong" second-order KKT conditions for the constrained problem.
- They prove that if $v^*$ is a first-order KKT point of the original problem, all corresponding $\theta^*$ are first-order KKT points of the reparameterized problem.
- Crucially, they show that if $v^*$ is a weak second-order KKT point of the original problem, then all corresponding $\theta^*$ are second-order KKT points of the reparameterized problem.
- Conversely, if $\theta^*$ is a second-order KKT point of the reparameterized problem, then $v^* = \phi(\theta^*)$ is a first-order KKT point and a weak second-order KKT point of the original problem.
- The paper notes that the correspondence does not hold for "strong" second-order KKT points in the reverse direction, providing a counterexample where a vertex of the simplex is a weak but not strong second-order KKT point, despite its pre-image being a strong second-order KKT point on the manifold.
Algorithmic Implementation: Based on these theoretical guarantees, the authors implement Riemannian Gradient Descent (RGD) algorithms. They utilize product retractions and vector transports on the sphere (for the simplex components) and Euclidean space. Two specific algorithms are proposed:
- RGD with backtracking update (Armijo condition).
- RGD with Barzilai-Borwein (BB) step size selection (RGD BB).

Key Contributions

Generalization of Reparameterization: The work extends existing results (specifically Li et al. [21]) from the Hadamard map ( $\theta^2$ ) to a broader class of smooth, strictly convex transformations $\psi(\theta)$ .
Generalization of Feasible Sets: The results are extended from single simplex constraints to product spaces involving multiple simplices and Euclidean components, covering applications like tensor factorization and functional data registration.
KKT Correspondence: The paper establishes the precise mapping between weak second-order KKT points on the product simplex and second-order KKT points on the smooth manifold, clarifying the limitations regarding strong second-order conditions.
Legendre Transform: The authors show that the Legendre transform of the reparameterization function $\psi$ also yields a valid reparameterization preserving these KKT properties.

Results
The authors benchmarked their RGD and RGD BB algorithms against PGD and PGD BB on two specific applications:

Simplex Constrained Tensor Decomposition:
- Task: Learning low-rank discrete multivariate probability distributions via Canonical Polyadic Decomposition (CPD) using Kullback-Leibler (KL) divergence.
- Performance: RGD BB consistently outperformed PGD and PGD BB, often by several orders of magnitude in terms of final solution cost (KL divergence).
- Robustness: RGD BB was found to be less sensitive to the choice of initial step size compared to other methods, making it suitable for large-scale problems where extensive parameter tuning is infeasible.
Functional Data Registration (SRVF):
- Task: Aligning functional data (curves) by finding optimal time-warping functions represented on a simplex.
- Performance: RGD and RGD BB achieved faster convergence and lower registration errors compared to PGD variants.
- Qualitative Improvement: The reparameterized approach produced smoother warping functions. In contrast, PGD often resulted in "step-like" or irregular warping functions due to the projection operation forcing derivatives to zero, which distorted the aligned shapes (e.g., turning sine waves into square waves). The Riemannian approach preserved the natural smoothness of the curves.

Significance and Claims
The paper claims that smooth reparameterization is a feasible and superior alternative to projection-based methods for optimization on product simplex spaces. By mapping the problem to a smooth manifold, the authors demonstrate that:

One can avoid the computational and geometric complexities of projection.
The optimization landscape is better behaved, leading to faster convergence (fewer iterations).
The solutions are more faithful to the original problem's geometry, particularly in preserving the smoothness of functional data, which is critical for applications like shape analysis and curve registration.
The theoretical framework ensures that finding critical points on the manifold is sufficient to find weak second-order critical points on the original constrained space.

The authors conclude that their Riemannian Barzilai-Borwein algorithm is robust and effective, offering a "plug-and-play" solution for these classes of problems.

More like this