Riemannian stochastic optimization for sufficient dimension reduction
Abstract
Sufficient dimension reduction (SDR) makes high-dimensional regression tractable by projecting the covariates onto a low-dimensional subspace that preserves the conditional mean of the response. Existing gradient-based estimators either operate in the ambient space and suffer from the curse of dimensionality, or localize in the reduced space at a per-outer-iteration cost at least quadratic in the sample size. We show that minimizers of the population Minimum Average Variance Estimation (MAVE) risk approximate the same Grassmannian target as the Outer Product of Gradients (OPG), and recast the empirical criterion as a smooth maximization on the Stiefel manifold with closed-form Riemannian gradient. The resulting algorithm, SMAVE, combines sparse projected-space nearest-neighbor localization with Riemannian stochastic gradient ascent. A simplified version comes with almost-sure convergence and a non-asymptotic rate matching the standard non-convex stochastic first-order scaling. Empirically, SMAVE matches or improves on RMAVE's synthetic subspace recovery at moderate-to-high ambient dimension, and on four real datasets it uniformly improves over OPG and is competitive with or outperforms RMAVE at orders of magnitude lower runtime.
Lay Summary
When predicting an outcome from many variables, most variables are often irrelevant, and the challenge is to automatically identify the few directions in the data that matter. This is the problem of dimension reduction, and a natural approach is to estimate how the prediction changes as input variables vary — that is, to estimate the gradient of the prediction function. If the gradient consistently points in the same few directions, those directions capture everything relevant for prediction. The difficulty is that estimating gradients reliably in high dimensions is both statistically fragile and computationally expensive — a problem known as the curse of dimensionality. Existing methods either compute gradients in the full high-dimensional space, which is fast but inaccurate, or refine estimates iteratively in a lower-dimensional space, which is accurate but very slow. We show that a well-known method called MAVE targets exactly the same directions as gradient-based approaches, but operates in the reduced space, offering the best of both worlds. By reformulating MAVE as an optimization problem on a curved geometric space called the Stiefel manifold, we derive a simple and explicit update formula. This enables efficient optimization using small random subsets of the data — a stochastic gradient approach — giving rise to our algorithm SMAVE. SMAVE automatically adapts its local neighborhood as the subspace estimate improves, combining statistical efficiency with low computational cost. On simulated and real datasets, SMAVE matches or outperforms existing methods in accuracy while running 10 to 50 times faster on standard benchmarks and over 1000 times faster on real regression tasks.