[Paper Review] Avoiding Communication in Proximal Methods for Convex Optimization Problems

Saeed Soori, Aditya Devarakonda+3 more|arXiv (Cornell University)|Oct 24, 2017

Sparse and Compressive Sensing Techniques1 references3 citations

TL;DR

This paper proposes communication-avoiding variants of stochastic FISTA (CA-SFISTA) and proximal Newton methods (CA-SPNM) by unrolling k iterations to reduce latency and bandwidth costs in distributed convex optimization. By reformulating the algorithms to communicate only every k iterations, the method achieves up to 10× speedup on 1024 nodes while preserving convergence rates and arithmetic costs, significantly improving scalability on modern distributed architectures.

ABSTRACT

The fast iterative soft thresholding algorithm (FISTA) is used to solve convex regularized optimization problems in machine learning. Distributed implementations of the algorithm have become popular since they enable the analysis of large datasets. However, existing formulations of FISTA communicate data at every iteration which reduces its performance on modern distributed architectures. The communication costs of FISTA, including bandwidth and latency costs, is closely tied to the mathematical formulation of the algorithm. This work reformulates FISTA to communicate data at every k iterations and reduce data communication when operating on large data sets. We formulate the algorithm for two different optimization methods on the Lasso problem and show that the latency cost is reduced by a factor of k while bandwidth and floating-point operation costs remain the same. The convergence rates and stability properties of the reformulated algorithms are similar to the standard formulations. The performance of communication-avoiding FISTA and Proximal Newton methods is evaluated on 1 to 1024 nodes for multiple benchmarks and demonstrate average speedups of 3-10x with scaling properties that outperform the classical algorithms.

Motivation & Objective

• To address the scalability limitations of classical stochastic FISTA and proximal Newton methods in distributed settings, where communication costs dominate performance.
• To reduce communication overhead—particularly latency—without increasing arithmetic or bandwidth costs in large-scale convex optimization.
• To reformulate FISTA and proximal Newton methods using a k-step unrolling technique to communicate data only every k iterations.
• To evaluate whether the communication-avoiding reformulations maintain the same convergence behavior and stability as classical formulations.
• To demonstrate strong scaling and performance improvements on large datasets across 1 to 1024 processors.

Proposed method

• Proposes k-step unrolling of classical SFISTA and SPNM algorithms to compute k iterations' worth of updates before communicating, reducing communication frequency by a factor of k.
• Uses randomized sampling to overlap computation and communication, enabling efficient data movement and load balancing in distributed environments.
• Reformulates the proximal mapping and update rules to maintain arithmetic equivalence with classical methods, ensuring identical convergence behavior.
• Applies data partitioning strategies to minimize communication costs while preserving numerical stability and residual error.
• Employs a modified stopping criterion and tolerance parameter to fairly compare performance across different k values and hardware scales.
• Evaluates the algorithms using both synthetic and real-world datasets (abalone, covtype, susy) on distributed clusters with up to 1024 nodes.

Experimental results

Research questions

• RQ1 Can unrolling k iterations in FISTA and proximal Newton methods reduce communication costs without altering convergence rates or arithmetic complexity?
• RQ2 How does increasing k affect the latency, bandwidth, and overall performance of distributed optimization on large-scale datasets?
• RQ3 To what extent do the communication-avoiding variants preserve the stability and residual error of classical algorithms?
• RQ4 Does the k-step reformulation enable better strong scaling on distributed platforms compared to classical implementations?
• RQ5 What is the maximum achievable speedup of the communication-avoiding algorithms across different datasets and hardware configurations?

Key findings

• CA-SFISTA and CA-SPNM achieve up to 10× speedup on 1024 processors compared to classical SFISTA and SPNM, with consistent performance gains across datasets.
• Speedup increases with both k and the number of processors, demonstrating effective reduction in latency costs by a factor of k.
• The k-step algorithms scale strongly, with execution times decreasing on more processors, unlike classical methods that degrade due to communication overhead.
• For abalone, a speedup of 9.63× is achieved on 64 nodes with k=32, showing significant performance improvement on smaller datasets.
• The convergence rate and residual error remain stable across k values up to 128, confirming that the reformulation does not compromise algorithmic behavior.
• Performance is eventually bounded by bandwidth when k and the number of processors increase, indicating that bandwidth becomes the limiting factor at scale.