Stochastic Langevin Optimization on Stiefel Products

• The paper presents a novel approach that simulates noisy Riemannian gradient flows on Stiefel products to handle nonconvex orthogonality constraints.
• It employs a Cayley-transformation-based integrator to maintain exact orthogonality while ensuring convergence to global minimizers under a probabilistic-annealing regime.
• Empirical results demonstrate significant performance gains over traditional random-start solvers in challenges like high-degree polynomial minimization, graph stability, and cryo-EM structure recovery.

Stochastic Langevin optimization on Stiefel product manifolds addresses global optimization under one or more nonconvex orthogonality constraints by simulating noisy Riemannian gradient flows. The Stiefel manifold $$M_{n,p} = \{X \in \mathbb{R}^{n \times p}\ |\ X^\top X = I_p\}$$ encodes the orthogonality, and its product generalizes to multiple such constraints. Optimization is formulated as a stochastic differential equation (SDE) on (products of) Stiefel manifolds, endowed with the canonical metric, with an explicit SDE representation and a Cayley-transformation-based numerical integrator, which ensures exact feasibility. This scheme yields provable convergence to global minimizers in a probabilistic-annealing regime, outperforming traditional random-start local solvers on classically hard nonconvex problems including high-degree polynomial minimization, graph stability number estimation, and structure recovery in cryo-EM data (Yuan et al., 2017).

1. Formulation on the Stiefel Manifold

The optimization task is to minimize a smooth objective $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$ over $M_{n,p}$. The tangent space at $X$ is

$$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$

The canonical metric is given by

$$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$

The Riemannian gradient is

$\nabla_M \F(X) = G - X G^\top X,$

with $G$ the Euclidean gradient of $\F(X)$.

A stochastic Langevin (Stratonovich) flow on $M_{n,p}$ is given by

$\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$0

with Brownian motion $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$1 on $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$2 and diffusion parameter $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$3. The extrinsic construction projects ambient noise into the tangent space using operators

$\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$4

so that

$\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$5

The Ito formulation includes drift correction: $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$6

2. Stochastic Diffusion on Product Stiefel Manifolds

For $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$7 blocks, each $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$8 with $\F(X): \mathbb{R}^{n \times p} \to \mathbb{R}$9, the feasible set is the product manifold

$M_{n,p}$0

with tangent space the direct sum $M_{n,p}$1 and metric the sum of metrics. The natural Langevin diffusion is a system of coupled SDEs (gradient coupling but noise for each $M_{n,p}$2 independent): $M_{n,p}$3 Gradient terms couple the blocks, while the diffusion is blockwise decoupled.

3. Numerical Integration via Cayley Transform

Integration of the SDE is achieved by alternately taking a Riemannian gradient step and adding tangent-space noise, followed by a Cayley transformation retraction which ensures exact orthogonality. For time-step $M_{n,p}$4 and $M_{n,p}$5:

1. Compute Riemannian gradient: $M_{n,p}$6,
2. Sample Gaussian noise $M_{n,p}$7, i.i.d. $M_{n,p}$8,
3. Project and combine:

$M_{n,p}$9

1. Skew-symmetrize:

$X$0

2. Cayley step:

$X$1

Each iteration costs $X$2 flops, and orthogonality $X$3 is preserved exactly.

4. Theoretical Guarantees and Convergence

Under Lipschitz and growth conditions on $X$4, the Cayley integrator achieves strong order-$X$5 accuracy: $X$6 In the constant-$X$7 case, the corresponding Fokker–Planck PDE for the density $X$8 is

$X$9

and, under a log-Sobolev inequality, $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$0 as $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$1.

With a sequence of diminishing diffusions $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$2, alternating SDE and local Riemannian-gradient phases, the probability of landing in a basin $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$3 of the global minimum becomes arbitrarily high after $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$4 independent cycles, attaining a point within $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$5 of the optimum with probability $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$6.

5. Algorithmic Implementation and Empirical Performance

Recommended step sizes are $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$7–$$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$8 and initial $$T_XM_{n,p} = \{Z \in \mathbb{R}^{n \times p} : Z^\top X + X^\top Z = 0\}.$$9 (adaptable based on problem). Diffusion schedules either decay as $$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$0 or are piecewise constant.

Empirical studies on several classes of nonconvex problems demonstrate consistent outperformance of independent diffusion-annealed descent with diffusion and drift mixing (“IDDM”) over random-start local minimization:

Problem Class	IDDM Performance Relative to Local Methods	Dimension Range
Homogeneous polynomial	Lower objectives by up to order of magnitude	$$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$1 up to 200
Biquadratic forms	Lower mean/best objectives	$$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$2 from 6 to 25
Graph stability number	Larger maximum stability found	Standard benchmarks
Cryo-EM orientation	Lower residual and MSE (at low noise)	Multiple $$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$3

In all cases, exact feasibility is maintained, with per-cycle cost comparable to local solvers.

6. Advantages, Limitations, and Extensions

Advantages include exact orthogonality throughout by the Cayley integrator, probabilistic guarantees of global optimum convergence, and flexibility for multi-block (product manifold) constraints. Unlike retraction-based SGD, the stochastic diffusion allows for global exploration, not merely local descent.

Limitations are the need for careful tuning of $$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$4 and $$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$5, increased per-iteration cost as $$g^c_X(Z_1, Z_2) = \operatorname{tr}\left(Z_1^\top (I - \tfrac{1}{2} X X^\top) Z_2\right).$$6 grows, and potentially large mixing times for highly nonconvex landscapes.

Possible extensions include adaptive diffusion schedules leveraging energy-barrier estimations, use of other metrics or quotient manifolds such as the Grassmannian, and development of variance-reduced or higher-order SDE integrators respecting manifold structure (Yuan et al., 2017).