Parallel-in-iteration optimization using multigrid reduction-in-time

1. Problem Statement

Standard gradient-based optimization algorithms (e.g., Gradient Descent, Proximal Gradient, FISTA) are widely used for convex and nonsmooth optimization problems. However, these methods often suffer from slow convergence rates, particularly for ill-conditioned problems, requiring tens of thousands of sequential iterations to reach a desired tolerance.

• The Bottleneck: Because these iterations are computed sequentially ($u_{k+1}$ depends on $u_k$), they present a significant computational bottleneck in large-scale simulations.
• The Goal: The authors aim to reduce the wall-clock time required to solve these optimization problems by parallelizing across the iterations (parallel-in-iteration), rather than just parallelizing across spatial domains or time steps in PDEs.

2. Methodology

The core innovation of this work is the re-purposing of Parallel-in-Time (PinT) methods, specifically Multigrid Reduction-in-Time (MGRIT), to solve optimization problems.

A. Theoretical Connection: Optimization as ODE Discretization

The authors establish a mathematical link between iterative optimization algorithms and the time-discretization of Ordinary Differential Equations (ODEs):

• Gradient Descent is interpreted as the Forward Euler discretization of the gradient flow ODE: $\frac{du}{dt} = -\nabla f(u)$.
• Proximal Gradient is interpreted as an Implicit-Explicit (IMEX) discretization of the gradient flow for $f+g$.
• Nesterov's Accelerated Gradient corresponds to a second-order ODE (hyperbolic type).

This connection allows the authors to treat the sequence of optimization iterates $\{u_0, u_1, \dots, u_K\}$ as a time-evolution problem, where the iteration index $k$ acts as the time index.

B. The MGRIT Framework for Optimization

MGRIT is a multilevel method that solves time-dependent systems in parallel by introducing a hierarchy of coarse time grids.

• Fine Grid: Uses the standard sequential optimization operator (e.g., Gradient Descent step $\Phi = I - s\nabla f$).
• Coarse Grid: Uses a computationally cheaper, implicit approximation of the fine operator.
  • For MP1 (Quadratic Minimization): The fine operator is Gradient Descent ($I - sA$). The coarse operator is the Proximal Point method ($I + msA)^{-1}$, which corresponds to a Backward Euler step. This provides stability for large step sizes ($ms$) on the coarse grid.
  • For MP2 (Elastic Obstacle/Nonsmooth): The fine operator is Proximal Gradient. The coarse operator is Alternating Proximal Mappings, which approximates the composition of steps implicitly.
• Algorithm: The method employs FCF-relaxation (F-relaxation, C-relaxation, F-relaxation) to propagate corrections across the time (iteration) domain. The "all-at-once" system of iterations is solved iteratively in parallel.

C. Convergence Theory

The paper provides theoretical bounds (Lemma 5.1) linking the convergence of the MGRIT residual to the convergence of the optimization gradient. It proves that if the MGRIT residual is reduced sufficiently, the gradient of the approximate solution at the final iteration will also be reduced to the desired tolerance, provided the coarse-grid operator is a good approximation of the fine-grid operator's $m$-step evolution.

3. Key Contributions

1. Novel Framework: Introduction of a "parallel-in-iteration" framework that applies MGRIT to optimization problems, treating iterations as time steps.
2. Operator Design: Development of specific coarse-grid operators (Proximal Point and Alternating Proximal Mappings) that serve as stable, implicit approximations of explicit optimization steps, enabling large step sizes on coarse levels.
3. Theoretical Analysis: Derivation of error bounds connecting the MGRIT residual reduction to the gradient norm reduction, and theoretical speedup estimates for 2-level and 3-level MGRIT schemes.
4. Handling Nonsmoothness: Successful application of the method to nonsmooth problems (Elastic Obstacle Problem) using $L_1$ penalty formulations.

4. Numerical Results

The authors tested the method on two model problems:

• MP1: Quadratic minimization (1D Laplacian).
• MP2: Elastic Obstacle Problem (1D and 2D), a nonsmooth convex problem.

Key Findings:

• Convergence Speed: MGRIT converged in a very small number of iterations (typically 8–13 iterations) regardless of the problem size ($N$) or the number of sequential iterations required ($N_t$). This demonstrates $h$-independent convergence, a hallmark of effective multigrid methods.
• Residual vs. Gradient: The reduction in the MGRIT space-time residual directly correlated with the reduction in the optimization gradient, validating the theoretical bounds.
• Nonsmooth Challenges: For the Elastic Obstacle Problem, convergence was slightly slower near the "contact set" (where the membrane touches the obstacle) due to nonsmoothness, but the method remained robust.
• Momentum Issues: The authors attempted to apply MGRIT to Nesterov's Accelerated Gradient. It failed to converge because Nesterov's method corresponds to a hyperbolic PDE, whereas MGRIT is optimized for parabolic (diffusion) problems. Standard coarse operators failed to capture the wave-like nature of the accelerated method.

5. Significance and Future Work

• Parallel Speedup: Theoretical estimates suggest significant wall-clock speedups (e.g., factors of 10–80 in ideal scenarios) for large-scale problems where sequential iterations are the bottleneck. The speedup depends heavily on the ratio of coarse-to-fine operator costs ($\alpha = t_c/t_f$) and the number of processors.
• Scalability: The method shows promise for large-scale machine learning and physics simulations where optimization requires massive numbers of iterations.
• Future Directions:
  • Hyperbolic Optimization: Developing specialized coarse-grid operators to handle momentum-accelerated methods (Nesterov, FISTA) which behave like hyperbolic PDEs.
  • Complex Models: Extending the framework to non-quadratic objectives (e.g., logistic regression), constrained optimization (ADMM), and stochastic gradient descent.
  • Implementation: The current study uses a serial implementation to verify convergence; future work will implement the full parallel version to demonstrate actual wall-clock time reductions.

In summary, this paper successfully bridges the gap between time-parallel computing and iterative optimization, offering a promising pathway to accelerate the solution of large-scale ill-conditioned optimization problems by treating the iteration count as a parallelizable time dimension.