Riemannian Optimization on Unknown Manifolds via Stochastic Proximal Refinement

Ruixiang Wu; Jiahao Ai; Tinko Sebastian Bartels; Xuanzhuo Liu; Yifei Wu; Yu Mao; Tongxin Li; Costas Courcoubetis

eess.SYOct 30, 2026

Riemannian Optimization on Unknown Manifolds via Stochastic Proximal Refinement

Authors: Ruixiang Wu, Jiahao Ai, Tinko Sebastian Bartels, Xuanzhuo Liu, Yifei Wu, Yu Mao, Tongxin Li, Costas Courcoubetis

Organizations: The Chinese University of Hong Kong, Shenzhen · University of Pennsylvania · Stanford University

Abstract

We study Riemannian optimization on an unknown manifold accessed only through feasible samples. The central challenge is to turn feasible samples into optimization steps, i.e., the samples are drawn from the unknown feasible manifold, but do not directly provide the tangent spaces, projections, or retractions required by classic Riemannian methods. Moreover, regions with low objective values may be sparsely represented in the data, so directly sampling from a distribution fitted to these observations can require many draws to find better feasible solutions. We introduce Proximal Refinement on Implicit Smooth Manifolds (PRISM), which turns feasible-sample access into an iterative optimization procedure by repeatedly sampling an objective-aware proximal distribution conditioned on the current iterate. Theoretically, we show that PRISM recovers the effect of Riemannian-gradient refinement, leading to provable objective descent, feasibility (manifold adherence), and convergence toward first-order stationarity. We demonstrate the effectiveness of our approach through synthetic experiments illustrating its geometric properties and applications to data-driven building control and diverse real world optimization tasks, highlighting its potential as a general framework for data-driven optimization across scientific and engineering domains.

Figures & tables

Figure 1: Illustration of

\mathsf{PRISM}

’s iterative refinement on an unknown manifold. Starting from an ambient point,

\mathsf{PRISM}

generates successive iterates through conditional sampling (left). Small diamonds represent feasible observations, while the peeled cover reveals the underlying manifold, whose description is unavailable to the optimizer. The lower inset illustrates a refinement distribution supported on the entire manifold, with color intensity indicating its density.

\mathsf{PRISM}

’s mean displacement reveals a Riemannian gradient step. The right panel illustrates RGD’s tangent step and retraction using explicitly available geometry for comparison.

Property	RGD	$\mathsf{PRISM}$
Ambient displacement (Thm 4.3 )	$x^{+}=x-\lambda g(x)+O(\lambda^{2})$	$\mathbb{E}\,X^{+}-x=-\lambda g(x)+O(\lambda^{2}+\lambda\tau+\varepsilon)$
One-step descent (Thm 4.4 )	$f(x^{+})\leq f(x)-c\lambda\\|g(x)\\|^{2}$	$\mathbb{E}\,f(p^{+})\leq f(p)-c\lambda\\|g(p)\\|^{2}+O(\lambda\tau+\varepsilon^{2}/\lambda)$
Stationarity (Thm 4.5 )	$\\|g(x_{\infty})\\|=0$	$\mathbb{E}\,\\|g(p_{\infty})\\|=O(\varepsilon/\lambda)$

Table 1: Informal Riemannian optimization properties of

\mathsf{PRISM}

Figure 2: Ideal refinement on the unit circle. (a) Score with

\smash{E_{k}=E^{f}_{\lambda,X_{k}}}

; (b) Reweight

\mu

to obtain

\smash{Q_{k}=Q^{f}_{\lambda,\tau}(x_{k},\cdot)}

; (c) Sample

X_{k+1}\mid X_{k}=x_{k}\sim Q_{k}

and use its realization

x_{k+1}

as the next anchor, at which the energy and conditional law are recomputed. Panels (b) and (c) display the same density of

X_{k+1}

, with darker arcs indicating higher density. Here

\mu

is uniform,

f(y)=y_{1}

Figure 3: Boundary geometry and zero proximal gap. Let

\mathcal{M}

be the closed unit disk,

f(y)=y_{1}

p=(-1,0)

. (a) The outward normal

n(p)=-\nabla f(p)

points outside the feasible tangent cone. The shadowy halfplane represents

p+\mathcal{T}_{\mathcal{M}}(p)

, with

\mathcal{T}_{\mathcal{M}}(p)=\{w\in\mathbb{R}^{2}\mid w_{1}\geq 0\}

; (b) A magnified view shows the strict sublevel

E^{f}_{\lambda,p}(y)<f(p)

E^{f}_{\lambda,p}(y)=f(y)+\|y-p\|^{2}/(2\lambda)

, using the same energy formula on the ambient plane. It is an open amber disk centered at

c_{p}=p-\lambda(1,0)

with radius

\lambda

, disjoint from the blue feasible set

\mathcal{M}

. Its boundary touches

\mathcal{M}

only at

p

, so

r_{\lambda}(p)=0

although

\|\nabla f(p)\|=1

. The vertical dashed lines show the common supporting tangent.

Figure 4: Optimization on Grassmann manifold

\mathrm{Gr}(2,10)

. (a) Fixed

\lambda=0.002

, varying

\tau

, compared with exact RGD (with access to the manifold) with stepsize

0.002

. (b) Fixed

\tau=0.14

, varying

\lambda

. Curves with finite

\lambda

value report current objective gaps, while

\lambda=\infty

reports the gap of the best sampled response so far. The dotted line marks the minimum over the training dataset

\mathcal{D}_{n}

, and zero corresponds to the exact optimum. For each parameter choice, we run

50

trajectories of

240

iterations, and plot their mean, shaded regions denote standard deviation across trajectories.

Figure 5: Building temperature control. (a) Executed room temperatures and the target. (b) Cumulative objective ( 20 ). Green and orange curves show means across five

\mathsf{PRISM}

-MPC and CEM-MPC runs, shading denotes

\pm

standard deviation. Dashed black shows Reference MPC, and blue shows the lowest-cost collected trajectory.

Method	TFBind8	TFBind10	Supercond.	Ant	DKitty	ChEMBL ∗ $(\times 10^{3})$	Average rank $\downarrow$
$\mathsf{PRISM}$	$\mathbf{0.9936}$ $\pm 0.0048$	$0.8766$ $\pm 0.1733$	$\underline{96.28}$ $\pm 6.43$	$517.33$ $\pm 7.99$	$270.75$ $\pm 3.83$	$\underline{396.022}$ $\pm 0.308$	$\mathbf{3.333}$
DiffOPT	$0.9529$ $\pm 0.0349$	$0.8385$ $\pm 0.1401$	$\mathbf{97.99}$ $\pm 3.90$	$530.43$ $\pm 20.17$	$261.24$ $\pm 7.23$	$394.901$ $\pm 1.011$	$\underline{4.333}$
COMs	$0.7960$ $\pm 0.1019$	$\mathbf{1.0152}$ $\pm 0.0000$	$88.91$ $\pm 8.96$	$538.93$ $\pm 15.94$	$\underline{280.77}$ $\pm 6.67$	$\mathbf{396.176}$ $\pm 0.000$	$4.417$
DDOM	$0.9222$ $\pm 0.0348$	$0.7902$ $\pm 0.0761$	$94.65$ $\pm 8.44$	$\underline{553.72}$ $\pm 13.62$	$\mathbf{282.72}$ $\pm 10.59$	$395.407$ $\pm 0.000$	$4.500$
GP-qEI	$0.8309$ $\pm 0.1095$	$\mathbf{1.0152}$ $\pm 0.0000$	$95.05$ $\pm 1.69$	$463.83$ $\pm 0.00$	$209.47$ $\pm 0.00$	$\mathbf{396.176}$ $\pm 0.000$	$5.083$
Gradient ascent	$\underline{0.9810}$ $\pm 0.0111$	$\mathbf{1.0152}$ $\pm 0.0000$	$92.96$ $\pm 3.13$	$98.41$ $\pm 84.47$	$198.30$ $\pm 16.17$	$375.079$ $\pm 41.041$	$5.667$

Table 2: Results on five disjoint confirmation seeds, with 256 submitted designs per task and seed. Each cell gives the mean batch-maximum oracle score and its population standard deviation. Bold and underlined means mark the highest and second-highest distinct means (higher score is better). Average rank uses unrounded means across the six tasks and ten displayed methods, with midranks for exact ties (lower rank is better). ∗ ChEMBL scores are divided by

10^{3}

for readability and invalid designs receive the offline dataset’s minimum label (293,000) before taking the batch maximum. Validity uses the official decoder and a nonempty RDKit 2022.09.5 molecular parse.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Stationarity and feasibility results. (a) Current gradient norm over

1{,}000

updates with

\tau_{k}=\tau_{0}/(k+1)^{2}

. (b) Raw distance to the true manifold over

160

updates at fixed temperature, the zoomed-in figure shows the first three updates. Colors distinguish

\lambda

, while solid and dashed lines denote temperature settings

0.014

and

0.14

, respectively. Lines and shading show means and

\pm

one sample SD across

50

runs.

Method	Total objective	Tracking RMSE ( ${}^{\circ}\mathrm{C}$ )
$\mathsf{PRISM}$ -MPC	$26.969\pm 0.054$	$0.07245\pm 0.00043$
Budgeted CEM-MPC	$32.128\pm 1.102$	$0.10922\pm 0.00725$
Reference MPC	$26.672$	$0.06744$
Lowest-cost collected trajectory	$409.903$	$0.89068$

Appendix

Table 3: Executed 48-hour performance. Stochastic controllers report mean

\pm

sample standard deviation across five runs; the other comparators each represent one trajectory.

Task	Score updates	Reward updates/member	Score noise
TFBind8	9,227	2,632	Broad
TFBind10	8,847	2,561	Mixture
Superconductor	6,728	2,089	Broad
Ant	7,482	2,273	Mixture
DKitty	7,626	2,306	Broad
ChEMBL	4,386	1,355	Broad

Appendix

Table 4: Allocated optimizer updates per seed. Reward updates apply to each of the three ensemble members. Broad score noise is log-uniform on

[0.005,4]

, mixture noise combines that law equally with log-uniform noise on

[0.01,0.5]

. The selected TFBind10, Ant, and ChEMBL samplers do not use the trained neural score model.

Riemannian Optimization

Riemannian optimization develops algorithms for minimizing objectives on curved spaces while respecting geometric structure and constraints.

See all 85 papers

Manifold Sampling via Entropy Maximization
May 12, 2026Cornelius V. Braun, Tilman Burghoff, Marc ToussaintMarkov Chain Monte Carlo Constrained Sampling
Technische Universität Berlin · Robotics Institute Germany
Mirror Descent on Riemannian Manifolds
Mar 18, 2026Jiaxin Jiang, Lei Shi, Jiyuan TanRiemannian Optimization Stochastic Optimization Convergence
School of Mathematical Sciences, Fudan University, Shanghai 200433, China · Shanghai Key Laboratory for Contemporary Applied Mathematics, Fudan University, Shanghai 200433, China
Locally Private Inference for Riemannian Stochastic Optimization
Sep 18, 2026Xiaotian Chang, Yangdi Jiang, Qirui HuRiemannian Optimization Confidence Region Estimation
School of Physical and Mathematical Sciences, Nanyang Technological University · School of Statistics and Data Science, Shanghai University of Finance and Economics.

May 12, 2026J/K move · Enter open · S save

Manifold Sampling via Entropy Maximization

Cornelius V. Braun, Tilman Burghoff, Marc Toussaint

Technische Universität Berlin · Robotics Institute Germany

Markov Chain Monte Carlo Constrained Sampling

Open paper AskPDF arXiv

Sampling from constrained distributions has a wide range of applications, including in Bayesian optimization and robotics. Prior work establishes convergence and feasibility guarantees for constrained sampling, but assumes that the feasible set is connected. However, in practice, the feasible set often decomposes into multiple disconnected components, which makes efficient sampling under constraints challenging. In this paper, we propose MAnifold Sampling via Entropy Maximization (MASEM) for uniform sampling on a manifold with an unknown number of disconnected components, implicitly defined by smooth equality and inequality constraints. The presented method uses a resampling scheme to maximize the entropy of the empirical distribution based on k-nearest neighbor density estimation. We show that, in the mean field, MASEM decreases the KL-divergence between the empirical distribution and the maximum-entropy target exponentially in the number of resampling steps. We instantiate MASEM with multiple local samplers and demonstrate its versatility and efficiency on synthetic and robotics-based benchmarks. MASEM enables fast and scalable mixing across a range of constrained sampling problems, improving over alternatives by an order of magnitude in Sinkhorn distance with competitive runtime.

Abstract

Riemannian Optimization

Explore similar work

Manifold Sampling via Entropy Maximization

Mirror Descent on Riemannian Manifolds

Locally Private Inference for Riemannian Stochastic Optimization