Characterizing & Finding Good Data Orderings for Fast Convergence of Sequential Gradient Methods

Abstract

While SGD, which samples from the data with replacement is widely studied in theory, a variant called Random Reshuffling (RR) is more common in practice. RR iterates through random permutations of the dataset and has been shown to converge faster than SGD. When the order is chosen deterministically, a variant called incremental gradient descent (IG), the existing convergence bounds show improvement over SGD but are worse than RR. However, these bounds do not differentiate between a good and a bad ordering and hold for the worst choice of order. Meanwhile, in some cases, choosing the right order when using IG can lead to convergence faster than RR. In this work, we quantify the effect of order on convergence speed, obtaining convergence bounds based on the chosen sequence of permutations while also recovering previous results for RR. In addition, we show benefits of using structured shuffling when various levels of abstractions (e.g. tasks, classes, augmentations, etc.) exists in the dataset in theory and in practice. Finally, relying on our measure, we develop a greedy algorithm for choosing good orders during training, achieving superior performance (by more than 14 percent in accuracy) over RR.

GraB: Finding Provably Better Data Permutations than Random Reshuffling

“…HaoChen and Sra [2019], Gürbüzbalaban et al [2021], and Mishchenko et al [2020] discuss extensively on the conditions needed for RR to benefit. Some recent works [Lu et al, 2021a, Mohtashami et al, 2022 suggest constructing better data permutations than RR via a memory-intensive greedy strategy. Concretely, Mohtashami et al [2022] proposes evaluating gradients on all the examples first to minimize Equation (2) before starting an epoch, applied to Federated Learning; Lu et al [2021a] provides an alternative of minimizing Equation (2) using stale gradients from previous epoch to estimate the gradient on each example.…”

“…Some recent works [Lu et al, 2021a, Mohtashami et al, 2022 suggest constructing better data permutations than RR via a memory-intensive greedy strategy. Concretely, Mohtashami et al [2022] proposes evaluating gradients on all the examples first to minimize Equation (2) before starting an epoch, applied to Federated Learning; Lu et al [2021a] provides an alternative of minimizing Equation (2) using stale gradients from previous epoch to estimate the gradient on each example. Rajput et al [2021] introduces an interesting variant to RR by reversing the ordering every other epoch, achieving better rates for quadratics.…”

“…In this section, we study the use of stale gradients-the stochastic gradients for each example from the previous epoch, saved in memory-to construct a data permutation at the start of each epoch. Lu et al [2021a] and Mohtashami et al [2022] propose to use greedy ordering in this way-selecting examples one at a time to minimize Equation (2). Here, we first show that this objective is closely related to the classic herding problem from discrepancy theory.…”

“…Based on the herding framework, a path to SGD with better data permutation becomes clear: when training models during a epoch, we can simply store all the stochastic gradients, and then herd them (with some herding algorithm) offline at the start of the next epoch to obtain the order for that epoch. This stale-gradient approach is formally described in Algorithm 2, and could be contrasted with a fresh-gradient approach (as in Mohtashami et al [2022]) that herded with newly computed stochastic gradients at the start of each epoch-it is desirable to avoid the use of fresh gradients if possible as they double the gradient computations needed each epoch. The greedy ordering of Lu et al [2021a] can be described as running Algorithm 2 using Algorithm 1 as the Herding subroutine.…”

“…Recent studies indicate the possibility of greedily constructing better permutations using stale estimates of each ∇f (w; x j ) from the previous epoch [Lu et al, 2021a, Mohtashami et al, 2022. Concretely, Lu et al [2021a] proved that for any model parameters w ∈ R d , if sums of consecutive stochastic gradients converge faster to the full gradient, then the optimizer will converge faster.…”