Loopless Proximal Riemannian Gradient EXTRA for Distributed Optimization on Compact Manifolds

This paper proposes the Proximal Riemannian Gradient EXTRA (PR-EXTRA) algorithm, a loopless distributed optimization method for composite problems on compact manifolds that achieves an $\mathcal{O}(1/K)$ sublinear convergence rate using only single-round communication per iteration.

The Gist

Imagine you are leading a team of explorers scattered across a vast, curved landscape (like the surface of a giant sphere or a twisted mountain range). Your goal is for the whole team to find the absolute lowest valley (the optimal solution) together. However, there are two major challenges:

1. The Terrain is Weird: Unlike a flat map (Euclidean space), this landscape is curved. If two explorers try to meet in the middle by simply averaging their coordinates, they might end up floating in the air or sinking underground, off the actual ground. They need special rules to stay on the surface.
2. The Path is Rocky: The path to the lowest point isn't just a smooth slide; it has sudden cliffs and jagged rocks (nonsmooth regularizers) that make it hard to calculate the exact direction to step.

This paper introduces a new, super-efficient way for these explorers to work together to find that lowest point. Here is the breakdown using simple analogies.

1. The Problem: The "Curved Map" Dilemma

In the past, most algorithms for distributed teams were designed for flat, grid-like worlds (like a chessboard). But in the real world, data often lives on "manifolds"—curved shapes like spheres (think of GPS coordinates on Earth) or complex shapes used in AI.

When explorers on a curved surface try to share information, they can't just say, "Let's meet halfway." If they do, they might end up off the map. They need to use "retractions" (a fancy way of saying "bouncing back onto the ground") to stay on the surface.

Furthermore, the goal isn't just to find a low point; it's to find a low point that also follows specific rules (like "keep your shape simple" or "be sparse"). This adds a "rocky" obstacle to the path.

2. The Old Way: The "Endless Huddle"

Previous methods for solving this were like a team that had to huddle up, discuss, calculate, and then huddle up again just to agree on the next step.

• Too many meetings: They needed multiple rounds of communication per step to make sure everyone stayed on the same page.
• Stuck in the mud: Because of the curvature and the rocky obstacles, they often got stuck in a "steady-state neighborhood"—close to the answer, but never quite exactly there.

3. The New Solution: PR-EXTRA (The "Loopless" Messenger)

The authors propose PR-EXTRA (Proximal Riemannian Gradient EXTRA). Think of this as a highly efficient, loopless communication protocol.

Here is how it works, step-by-step:

• The "Loopless" Magic:
  Imagine each explorer has a notebook. In the old days, they had to pass the notebook around the circle three times to agree on a move. With PR-EXTRA, they pass the notebook only once per step.

  • Analogy: It's like a relay race where the baton is passed exactly once per lap, rather than having the runners run back and forth to double-check the route. This saves massive amounts of time and energy (communication cost).

• The "Proximal" Step (Handling the Rocks):
  When an explorer hits a "rocky" part of the path (the nonsmooth regularizer), they don't try to calculate a perfect slope. Instead, they use a "proximal mapping."

  • Analogy: Imagine you are walking in the dark and hit a wall. Instead of trying to calculate the exact angle of the wall, you just take a small step, feel the wall, and adjust your path to slide along it. The algorithm does this mathematically to handle the jagged parts of the problem without getting stuck.

• The "Correction" (The Ghost of Gradients Past):
  This is the secret sauce. The algorithm keeps a "ghost" variable (a correction term) that remembers the mistakes made in previous steps.

  • Analogy: If you are walking in a foggy forest and you realize you drifted slightly left, you don't just turn right. You remember how much you drifted and adjust your future steps to compensate for that specific drift. This ensures that even though everyone is on a curved surface, they eventually converge to the exact same spot, not just a spot "near" the answer.

4. The Result: Speed and Precision

The paper proves mathematically that this new method:

1. Converges Fast: It finds the solution at a rate of $O(1/K)$. In plain English, if you double the number of steps, you get twice as close to the perfect answer. This matches the best performance of algorithms on flat ground, which is a huge achievement for curved ground.
2. Saves Bandwidth: By only requiring one round of communication per step, it's much faster for teams with slow internet or limited battery (like sensor networks).
3. Works on Curved Ground: It successfully handles the geometry of spheres and complex shapes without falling off the map.

Summary

Think of PR-EXTRA as a new set of instructions for a team of hikers on a curved mountain. Instead of constantly stopping to argue about where to meet (which wastes time) or getting lost because the map is curved, they use a clever "memory trick" to correct their path and a "one-pass" communication rule to stay efficient. They navigate the rocky terrain and the curved surface to find the exact lowest point together, faster and more accurately than ever before.

Technical Summary

Here is a detailed technical summary of the paper "Loopless Proximal Riemannian Gradient EXTRA for Distributed Optimization on Compact Manifolds."

1. Problem Formulation

The paper addresses the challenge of distributed composite optimization over compact Riemannian manifolds. Unlike standard Euclidean optimization, this setting involves constraints where data naturally resides on curved geometries (e.g., orthogonality constraints in PCA or low-rank matrix completion).

The specific problem is formulated as minimizing the sum of local smooth functions and a global nonsmooth regularizer:
$$\min_{x \in \mathcal{M}} h(x) = \frac{1}{n} \sum_{i=1}^n f_i(x) + r(x)$$
Where:

• $\mathcal{M} \subset \mathbb{R}^{d \times r}$ is a compact smooth Riemannian manifold.
• $f_i(x)$ are local smooth cost functions held by node $i$.
• $r(x)$ is a common convex but nonsmooth regularizer (e.g., $\ell_1$ or $\ell_{2,1}$ norms).
• The network consists of $n$ nodes communicating over a fixed, connected, undirected graph.

Key Challenges:

1. Non-Euclidean Geometry: Standard linear consensus (weighted averaging) fails because the average of points on a manifold may not lie on the manifold.
2. Nonsmoothness: The presence of $r(x)$ requires proximal operators, which are computationally difficult to define on manifolds using geodesic distances.
3. Communication Efficiency: Existing Riemannian distributed algorithms often require multiple communication rounds per iteration or complex gradient tracking mechanisms, leading to high overhead.
4. Exact Convergence: Many distributed methods with constant stepsizes only converge to a neighborhood of the solution, not the exact stationary point.

2. Methodology: PR-EXTRA Algorithm

The authors propose PR-EXTRA (Proximal Riemannian Gradient EXTRA), a "loopless" algorithm that extends the EXTRA framework to Riemannian manifolds. The algorithm avoids nested loops for consensus and handles nonsmooth terms via a specialized proximal step.

Core Mechanism:
The algorithm maintains local variables $x_{i,k}$ (primal), $y_{i,k}$ (auxiliary), and $s_{i,k}$ (gradient correction). The update rules per iteration $k$ are:

1. Gradient Correction (History Accumulation):
   Node $i$ updates a correction term $s_{i,k}$ to track the global gradient direction, compensating for the steady-state error common in distributed settings:
   $$s_{i,k} = s_{i,k-1} + \sum_{j=1}^n (w_{ij} - \tilde{w}_{ij})x_{j,k-1} - \alpha [\text{grad} f_i(x_{i,k}) - \text{grad} f_i(x_{i,k-1})]$$
   Here, $\text{grad}$ denotes the Riemannian gradient, and $W, \tilde{W}$ are mixing matrices.

2. Consensus and Projection:
   Nodes aggregate neighbor information and project the result back onto the manifold to ensure feasibility:
   $$y_{i,k} = P_{\mathcal{M}} \left( \sum_{j=1}^n w_{ij} x_{j,k} + s_{i,k} \right)$$
   $P_{\mathcal{M}}$ is the projection operator onto the manifold.

3. Proximal Step (Tangent Space):
   To handle the nonsmooth term $r(x)$, the algorithm solves a minimization subproblem in the tangent space $T_{y_{i,k}}\mathcal{M}$ rather than the manifold itself. This avoids the computational cost of geodesic-based proximal mappings:
   $$\eta_{i,k} = \arg\min_{\eta \in T_{y_{i,k}}\mathcal{M}} \left\{ \frac{1}{2\tau} \|\eta\|^2 + r(y_{i,k} + \eta) \right\}$$

4. Primal Update:
   The new iterate is obtained by moving along the descent direction $\eta_{i,k}$ and projecting back to the manifold:
   $$x_{i,k+1} = P_{\mathcal{M}}(y_{i,k} + \eta_{i,k})$$

Key Features:

• Single-Round Communication: Each iteration requires only one round of neighbor communication.
• Loopless: It does not require inner loops to solve consensus or projection subproblems iteratively.
• Projection-Based: Uses Euclidean projection operators ($P_{\mathcal{M}}$) which are often computationally cheaper than exponential maps or retractions.

3. Key Contributions

1. Algorithmic Innovation: The first "loopless" distributed algorithm for composite optimization (smooth + nonsmooth) on Riemannian manifolds. It generalizes the EXTRA algorithm to handle manifold constraints and nonsmooth regularizers simultaneously.
2. Theoretical Guarantees:
   • Convergence Rate: Proves a sublinear convergence rate of $O(1/K)$ to a stationary point using a constant stepsize. This matches the best-known rates for proximal gradient EXTRA in Euclidean spaces.
   • Exact Convergence: Unlike many distributed Riemannian methods that converge to a neighborhood, PR-EXTRA achieves exact convergence to a stationary point (where $0 \in P_{T_x\mathcal{M}}(\nabla f(x) + \partial r(x))$).
   • Boundedness: Establishes the boundedness of iterates and the sufficient descent property of the objective function.
3. Efficiency: Reduces communication and computational overhead by eliminating the need for multi-step consensus loops or expensive geodesic calculations in the proximal step.

4. Experimental Results

The authors evaluated PR-EXTRA against state-of-the-art baselines: DR-ProxGT (Distributed Riemannian Proximal Gradient Tracking) and DRSM (Distributed Riemannian Subgradient Method).

Test Problems:

• Distributed Sparse PCA (SPCA): Minimizing variance with an $\ell_1$-regularizer on the Stiefel manifold.
• Coordinate-Independent Sparse Estimation (CISE): Minimizing with an $\ell_{2,1}$-regularizer.

Performance Metrics:

• KKT Violation: Measures stationarity (optimality gap).
• Consensus Error: Measures agreement among nodes.

Findings:

• Faster Convergence: PR-EXTRA converged significantly faster than DR-ProxGT and DRSM. In the SPCA experiment, PR-EXTRA stabilized within ~1,000 iterations, whereas DR-ProxGT required ~3,000.
• Lower Error: PR-EXTRA achieved lower KKT violations and consensus errors in fewer iterations.
• Robustness: The algorithm effectively handled the nonsmooth regularizers ($\ell_1$ and $\ell_{2,1}$) while maintaining feasibility on the manifold.

5. Significance and Impact

• Bridging the Gap: This work fills a critical theoretical gap between Euclidean distributed optimization (where composite problems are well-studied) and Riemannian optimization (which has largely focused on smooth objectives).
• Practical Applicability: By using projection operators and tangent space proximal steps, the algorithm is computationally feasible for real-world applications like distributed PCA, dictionary learning, and sensor networks where data has geometric constraints.
• Communication Efficiency: The "loopless" design is crucial for large-scale networks where communication bandwidth is a bottleneck, offering a communication-efficient alternative to gradient tracking methods.
• Theoretical Benchmark: The $O(1/K)$ rate sets a new benchmark for distributed nonconvex optimization on manifolds, demonstrating that exact convergence is achievable without sacrificing communication efficiency.

In summary, PR-EXTRA provides a robust, efficient, and theoretically sound framework for solving complex, constrained, distributed learning problems on curved data spaces.