[Paper Review] Proximal and Federated Random Reshuffling

Konstantin Mishchenko, Ahmed Khaled+1 more|arXiv (Cornell University)|Feb 12, 2021

Stochastic Gradient Optimization Techniques36 references4 citations

TL;DR

This paper proposes ProxRR and FedRR, novel algorithms for proximal and federated learning using random reshuffling, achieving faster convergence by evaluating the proximal operator only once per epoch. The methods outperform Proximal SGD and Local SGD, especially when the proximal operator is expensive, reducing computational cost by up to n times in such cases.

ABSTRACT

Random Reshuffling (RR), also known as Stochastic Gradient Descent (SGD) without replacement, is a popular and theoretically grounded method for finite-sum minimization. We propose two new algorithms: Proximal and Federated Random Reshuffing (ProxRR and FedRR). The first algorithm, ProxRR, solves composite convex finite-sum minimization problems in which the objective is the sum of a (potentially non-smooth) convex regularizer and an average of $n$ smooth objectives. We obtain the second algorithm, FedRR, as a special case of ProxRR applied to a reformulation of distributed problems with either homogeneous or heterogeneous data. We study the algorithms' convergence properties with constant and decreasing stepsizes, and show that they have considerable advantages over Proximal and Local SGD. In particular, our methods have superior complexities and ProxRR evaluates the proximal operator once per epoch only. When the proximal operator is expensive to compute, this small difference makes ProxRR up to $n$ times faster than algorithms that evaluate the proximal operator in every iteration. We give examples of practical optimization tasks where the proximal operator is difficult to compute and ProxRR has a clear advantage. Finally, we corroborate our results with experiments on real data sets.

Motivation & Objective

• To address the inefficiency of existing proximal and federated learning methods that evaluate the proximal operator in every iteration.
• To develop a new algorithm, ProxRR, for composite convex finite-sum minimization with non-smooth regularizers.
• To extend ProxRR to federated learning via a reformulation, resulting in FedRR, suitable for both homogeneous and heterogeneous data.
• To demonstrate that evaluating the proximal operator only once per epoch significantly improves efficiency when the operator is computationally expensive.
• To provide theoretical convergence guarantees and empirical validation on real datasets.

Proposed method

• ProxRR solves composite convex finite-sum problems by combining random reshuffling with a single proximal operator evaluation per epoch.
• The algorithm uses a constant or decreasing stepsize and applies the proximal operator only at the end of each epoch, reducing computational overhead.
• FedRR is derived as a special case of ProxRR applied to a reformulated distributed optimization problem with client-specific data.
• The method leverages the structure of the reformulated problem to maintain convergence guarantees under both homogeneous and heterogeneous data settings.
• Theoretical analysis relies on strong convexity and smoothness assumptions, with convergence bounds derived using variance reduction and recursive inequalities.
• Experiments use real datasets (e.g., 'a1a', 'mushrooms') and compare FedRR with Local SGD under varying local steps and stepsize schedules.

Experimental results

Research questions

• RQ1 Can random reshuffling be effectively extended to proximal finite-sum minimization with non-smooth regularizers?
• RQ2 Does evaluating the proximal operator only once per epoch lead to significant computational savings without sacrificing convergence?
• RQ3 How does FedRR compare to Local SGD in terms of convergence speed and communication efficiency in federated learning?
• RQ4 What is the theoretical convergence rate of ProxRR and FedRR under strong convexity and smoothness assumptions?
• RQ5 In what practical optimization tasks does the reduced proximal evaluation frequency provide a clear advantage?

Key findings

• ProxRR achieves a convergence rate of O(1/T) with a single proximal evaluation per epoch, significantly reducing computational cost when the proximal operator is expensive.
• FedRR inherits the same efficiency and convergence properties, making it suitable for federated learning with both homogeneous and heterogeneous data.
• Experiments show FedRR converges faster than Local SGD, especially when the number of local steps is large, due to its robustness to stepsize selection.
• The method outperforms Proximal SGD and Local SGD in terms of both communication rounds and data passes, particularly on real-world datasets like 'mushrooms'.
• Theoretical analysis confirms that FedRR maintains linear convergence under strong convexity, with error bounds depending on the condition number and data variance.
• The proximal operator for ℓ1 + ℓ2 regularization is efficiently computed via soft-thresholding, enabling practical deployment of ProxRR.