Riemannian Stochastic Optimization for Sufficient Dimension Reduction

Technical Summary: Riemannian Stochastic Optimization for Sufficient Dimension Reduction

Problem Statement
Sufficient Dimension Reduction (SDR) aims to find a low-dimensional linear projection $B_*^\top X$ of high-dimensional covariates $X \in \mathbb{R}^p$ that preserves the conditional mean of the response $Y$, i.e., $E[Y|X] = E[Y|B_*^\top X]$. While this reduces the intractable problem of estimating a regression function in $\mathbb{R}^p$ to one in $\mathbb{R}^d$ (where $d \ll p$), existing gradient-based estimators face a fundamental trade-off between statistical flexibility and computational scalability.

Inverse regression methods (e.g., SIR, SAVE) are computationally efficient but rely on restrictive linearity or constant variance assumptions. Gradient-based methods like the Outer Product of Gradients (OPG) and Minimum Average Variance Estimation (MAVE) offer greater flexibility by estimating the subspace directly from local derivative information without distributional assumptions. However, these methods suffer from the "curse of dimensionality" when operating in the ambient space $X \in \mathbb{R}^p$. Refined variants like RMAVE (Refined MAVE) address this by localizing in the projected space $\mathbb{R}^d$, but they incur a computational cost of at least $O(n^2)$ per outer iteration due to the need to compute dense pairwise kernel weights and solve coordinate-descent subproblems for $n$ samples.

Methodology: SMAVE
The authors propose SMAVE (Stochastic MAVE), an algorithm that combines Riemannian stochastic optimization with adaptive sparse localization to resolve the tension between statistical efficiency and computational cost.

1. Theoretical Unification: The paper establishes that the minimizers of the population MAVE risk approximate the same Grassmannian target as the OPG method (the span of the regression function's gradients). However, MAVE achieves this through projected-space local regression rather than ambient-space gradient aggregation. This insight motivates treating the MAVE objective as a smooth function on the Grassmann manifold.
2. Riemannian Reformulation: The empirical MAVE criterion is recast as a maximization problem on the Stiefel manifold $\text{St}(p, d) = \{B \in \mathbb{R}^{p \times d} : B^\top B = I_d\}$. The authors derive a closed-form expression for the Riemannian gradient of this objective, which allows for optimization directly on the manifold without explicit orthogonality constraints.
3. Scalable Algorithm:
   • Sparse Localization: Instead of computing dense kernel weights for all $n^2$ pairs, SMAVE uses sparse $k$-nearest neighbor ($k$-NN) weights in the projected space $\mathbb{R}^d$. The neighborhood definition depends on the current subspace estimate $B$, and the $k$-NN index is periodically refreshed (every $\tau$ iterations) to adapt to the evolving subspace.
   • Stochastic Optimization: The algorithm employs mini-batch Riemannian stochastic gradient ascent. At each iteration, a mini-batch of anchors is sampled, and the gradient is computed using only the local $k$-NN neighbors.
   • Initialization: Unlike RMAVE, which warm-starts from an OPG estimate (which degrades in high dimensions), SMAVE initializes with a uniform random draw on the Stiefel manifold, decoupling the optimization from the potentially poor OPG basin.

Key Contributions

1. Theoretical Characterization: The authors prove that MAVE minimizers target the same subspace as OPG (the top eigenspace of the gradient outer product matrix) but via a different mechanism (local regression vs. gradient aggregation). This unifies two major SDR families and reveals the intrinsic Grassmannian nature of the MAVE objective.
2. Riemannian Reformulation: They provide a closed-form Riemannian gradient for the MAVE objective on the Stiefel manifold, enabling efficient manifold optimization.
3. Scalable Algorithm with Guarantees: SMAVE reduces the per-iteration complexity from the $O(n^2)$ scaling of RMAVE to $O(n)$ (linear in sample size) by leveraging sparse $k$-NN weights and stochastic gradients. The paper provides:
   • An almost-sure convergence proof to a critical point for a simplified version of the algorithm.
   • A non-asymptotic convergence rate of $O(1/\sqrt{T})$, matching the optimal rate for non-convex stochastic first-order methods.
4. Empirical Validation: Extensive experiments demonstrate that SMAVE matches or improves upon RMAVE's subspace recovery accuracy at moderate-to-high ambient dimensions while running 10 to 35 times faster.

Results

• Synthetic Benchmarks: In simulations with $p \in \{50, 100, 200\}$ and $n \in \{1000, 2000, 5000\}$, SMAVE reduced subspace estimation error by a factor of 2–9 compared to RMAVE. The performance gap widened as the ambient dimension $p$ increased, primarily because RMAVE's initialization (based on OPG) degrades exponentially with $p$, whereas SMAVE's random initialization and stochastic exploration allow it to escape suboptimal basins.
• Real-World Datasets: On four datasets (Wine Quality, Seoul Bike, Pumadyn, Gas Sensor) with dimensions ranging from $p=11$ to $p=128$, SMAVE uniformly outperformed OPG. It was competitive with or outperformed RMAVE in prediction accuracy (measured by test MSE) while achieving orders-of-magnitude reductions in runtime. Notably, on the high-dimensional Gas Sensor dataset ($p=128$), SMAVE significantly outperformed RMAVE.
• Efficiency: SMAVE achieved the best accuracy-efficiency trade-off, occupying the "favorable lower-left region" of the Pareto frontier (low error, low runtime).

Significance and Claims
The paper claims that SMAVE unlocks both statistical and computational gains for nonparametric statistics by applying scalable Riemannian stochastic optimization to a foundational problem. The primary significance lies in demonstrating that the computational bottleneck of gradient-based SDR (the quadratic cost of dense localization and coordinate descent) can be overcome without sacrificing statistical consistency. By shifting the optimization to the Stiefel manifold and utilizing stochastic gradients with sparse localization, SMAVE makes sufficient dimension reduction tractable for high-dimensional settings where previous methods were either too slow or statistically unreliable due to initialization issues. The authors note that their approach is general-purpose and could potentially be extended to other objective functions, such as those for non-Gaussian regression (MADE) or classification, though these extensions are left for future work.