A Retraction-Free EXTRA Method for Decentralized Optimization on the Stiefel Manifold

Abstract: Decentralized optimization provides a fundamental framework for large-scale learning and signal processing with distributed data. We study decentralized optimization with orthogonality constraints on the Stiefel manifold and propose RF-EXTRA, a distributed retraction-free primal-dual method on static undirected networks. The method combines an approximate gradient mapping for orthogonality-constrained optimization with an EXTRA-based decentralized recursion, thereby avoiding retractions while preserving a simple communication pattern. On the theoretical side, the analysis considers \revise{the joint error} $(\mathbf{X}_k-\overline{\mathbf X}_k,\mathbf{s}_k-\overline{\mathbf s}_k)$ in the local variables and local directions, and establishes a contractive recursion for the joint error. This contractivity ensures that the joint error can be controlled using small yet constant step sizes, thus leading to an exact $\mathcal{O}(1/K)$ convergence rate of RF-EXTRA to a stationary point. Experiments on PCA and low-rank matrix completion show that RF-EXTRA compares favorably with the reported decentralized baselines and exhibits strong communication efficiency on the tested tasks on the Stiefel manifold.

• The paper introduces RF-EXTRA, a retraction-free EXTRA method that achieves decentralized Stiefel manifold optimization with exact O(1/K) convergence.
• It employs a penalty term for feasibility control and surrogate gradients to eliminate costly retractions while leveraging efficient BLAS routines.
• Experimental results across decentralized PCA, MNIST, and low‐rank matrix completion demonstrate robust consensus, fast stationarity, and improved communication efficiency.

Retraction-Free EXTRA for Decentralized Optimization on the Stiefel Manifold: An Expert Essay

Problem Setting and Motivation

The paper "A Retraction-Free EXTRA Method for Decentralized Optimization on the Stiefel Manifold" (2604.23754) addresses the critical challenge of decentralized optimization under orthogonality constraints, where the feasible set is the Stiefel manifold $$\mathrm{St}(d, r) = \{ X \in \mathbb{R}^{d \times r} : X^\top X = I_r \}$$. This setting arises in distributed PCA, subspace learning, federated estimation of low-rank matrices, and deep learning with orthogonal layer constraints.

Decentralized scenarios, e.g., sensor or edge networks, enforce that agents only access local data and communicate over sparse networks—with solution variables required to remain globally consistent, both in consensus and in manifold feasibility. State-of-the-art primal-dual frameworks such as EXTRA achieve exact decentralized consensus and fast convergence in Euclidean settings, but fundamentally rely on unconstrained operations or expensive retraction/projection steps when generalized to manifolds. Retraction operations (e.g., QR or polar) induce substantial computational overhead and break compatibility with highly optimized BLAS routines on modern accelerators.

The central question addressed is: Can we design an EXTRA-style correction method for decentralized Stiefel manifold optimization that eliminates per-iteration retractions while retaining strong theoretical guarantees and empirical efficiency?

RF-EXTRA: Algorithmic Design

The authors introduce RF-EXTRA, a decentralized algorithm combining three key innovations:

1. Retraction-Free Feasibility Control: Instead of explicit retractions/projections to the Stiefel manifold at each iteration, feasibility is softly enforced using an embedded penalty term in the update, controlling $\|X^\top X - I_r\|_F^2$.
2. EXTRA-Based Primal-Dual Correction: The protocol retains a two-state recursion (primal and correction/auxiliary state), structurally preserving the distinctive consensus correction of EXTRA—ensuring unbiased consensus convergence across heterogeneous agents with only local communication.
3. Stiefel-Manifold-aware Gradient Surrogates: Direct Riemannian gradients are approximated by a surrogate operator $H_i(X)$, constructed via a single gradient evaluation per agent (as opposed to the two required by exact extra-manifold mappings). This sidesteps prohibitive manifold computations and aligns with efficient linear algebra kernels.

The update reads (for agent $i$ at iteration $k$):

$$\begin{aligned} X_{i,k+1} &= \sum_j w_{ij}X_{j,k} + s_{i,k} \ s_{i,k+1} &= s_{i,k} + \sum_j (w_{ij} - v_{ij})X_{j,k} - \alpha [H_i(X_{i,k+1}) - H_i(X_{i,k})] \end{aligned}$$

where $w_{ij}$ are graph-defined mixing weights, $v_{ij}$ defines the correction matrix $V$, and $H_i(X)$ encapsulates both the local (surrogate) Riemannian gradient and penalty for infeasibility. This recursion leverages only matrix additions and multiplications, dispensing with retractions entirely.

Theoretical Guarantees

A central technical challenge arises from the mixture of nonconvex manifold constraints, inexact primal-dual correction, and maintaining decentralized consensus. The main theoretical contributions are:

• Joint Error Analysis: The evolution of $\|X^\top X - I_r\|_F^2$0—quantifying consensus disagreement and correction error—is captured via a contractive recursion, but not under the usual Frobenius norm. Instead, an equivalent operator norm (adapted to the system matrix $\|X^\top X - I_r\|_F^2$1) is constructed to rigorously establish a contraction property, essential for bounding the perturbations from network inconsistency.
• Neighborhood Propagation: The algorithm admits an explicit controllable region (a neighborhood of the Stiefel manifold) where feasibility violations remain bounded, provided that stepsizes $\|X^\top X - I_r\|_F^2$2 and penalty $\|X^\top X - I_r\|_F^2$3 are chosen according to concrete, computable thresholds.
• Exact $\|X^\top X - I_r\|_F^2$4 Convergence Rate: Combining averaged descent (on a penalized objective) with joint error arguments, RF-EXTRA is shown to converge at rate $\|X^\top X - I_r\|_F^2$5—matching optimal rates for nonconvex decentralized problems—without requiring vanishing stepsizes or projections, and under only mild regularity assumptions.

Experimental Results

Extensive empirical validation compares RF-EXTRA to DPRGD, DPRGT, DESTINY, and REXTRA across synthetic and real data for decentralized PCA and low-rank matrix completion (LRMC), including both network and computational efficiency metrics.

Synthetic Decentralized PCA

The method demonstrates:

• Robustness to Topology and Tuning: Performance is stable across diverse network graphs (ring, star, and Erdős–Rényi topologies), and a wide range of internal penalty parameter ($\|X^\top X - I_r\|_F^2$6) values.

  Figure 1: RF-EXTRA maintains effective feasibility and consensus across various networks and $\|X^\top X - I_r\|_F^2$7 scalings in synthetic decentralized PCA.

• Superior Communication Efficiency: On synthetic PCA, RF-EXTRA achieves comparable or superior convergence in both stationarity and consensus, while reducing wall-clock execution time compared to retraction-based REXTRA and matching or outperforming gradient-tracking baselines.

  Figure 2: RF-EXTRA exhibits favorable stationarity and consensus behaviors as a function of communication, outperforming or rivaling retraction-based and gradient-tracking algorithms.

Real-World MNIST Experiments

On MNIST, with $\|X^\top X - I_r\|_F^2$8 agents each holding a data shard, RF-EXTRA:

• Achieves parity or slight improvements in both consensus and stationarity over DESTINY and DPRGT within matched communication budgets using fixed stepsizes (no line search).

  Figure 3: On MNIST PCA, RF-EXTRA and DESTINY outperform gradient-tracking schemes in both stationarity and consensus as a function of communication.

Decentralized Low-Rank Matrix Completion (LRMC)

Under ring communication and synthetic low-rank data:

• Fast Consensus and Stationarity: RF-EXTRA yields accelerated reduction in both consensus error and manifold stationarity compared to DESTINY and DPRGT at selected stepsizes, with less sensitivity to penalty tuning.

  Figure 4: For decentralized LRMC, RF-EXTRA rapidly achieves low stationarity and consensus under communication constraints.

  

  Figure 5: Step size selection analysis shows RF-EXTRA sustains consistent advantages across choices in decentralized LRMC.

Implications, Limitations, and Future Directions

RF-EXTRA delivers a principled answer to integrating retraction-free feasibility control with primal-dual decentralized optimization on matrix manifolds. Key implications include:

• Computational Practicality: Avoiding retractions enables full utilization of hardware-accelerated BLAS, essential for scaling decentralized manifold learning to large models and datasets.
• Theoretical Maturity: The explicit, contractive error analysis via an equivalent norm confirms that exactness and stationarity guarantees—previously reserved to Euclidean or retraction-based methods—are feasible in a retraction-free protocol.
• Decentralized Matrix Manifold Learning: The success with the Stiefel manifold suggests extensibility to other compact matrix manifolds such as the Grassmannian or flag manifolds, potentially impacting decentralized training of deep models with structure.

Limitations and open questions:

• The analysis assumes static, undirected networks; extensions to time-varying or directed graphs require additional work, particularly in error control.
• The contractive analysis depends on explicit graph spectral gaps through $\|X^\top X - I_r\|_F^2$9, which may become weak in poorly connected networks.
• The method's performance under severe heterogeneity or adversarial data splits remains to be fully characterized.

Future research may address generalization to composite objectives (e.g., with nonsmooth or stochastic terms), asynchronous and robust communication, and decentralized large-scale learning scenarios where manifold constraints encode real-world inductive biases.

Conclusion

RF-EXTRA establishes a computationally and theoretically efficient solution to decentralized optimization with Stiefel manifold constraints, leveraging a novel synergy of retraction-free ambient updates and primal-dual EXTRA-type corrections. It attains robust $H_i(X)$0 convergence with strong empirical performance across diverse networked matrix learning tasks. This approach positions retraction-free consensus optimization as a scalable blueprint for the next generation of decentralized manifold learning frameworks.