GraB: Finding Provably Better Data Permutations than Random Reshufﬂing
Abstract
Random reshufﬂing, which randomly permutes the dataset each epoch, is widely adopted in model training because it yields faster convergence than withreplacement sampling.
Recent studies indicate greedily chosen data orderings can further speed up convergence empirically, at the cost of using more computation and memory.
However, greedy ordering lacks theoretical justiﬁcation and has limited utility due to its non-trivial memory and computation overhead.
In this paper, we ﬁrst formulate an example-ordering framework named herding and answer afﬁrmatively that SGD with herding converges at the rate O ( T  2 / 3 ) on smooth, non-convex objectives, faster than the O ( n 1 / 3 T  2 / 3 ) obtained by random reshufﬂing, where n denotes the number of data points and T denotes the total number of iterations.
To reduce the memory overhead, we leverage discrepancy minimization theory to propose an online Gradient Balancing algorithm (GraB) that enjoys the same rate as herding, while reducing the memory usage from O ( nd ) to just O ( d ) and computation from O ( n 2 ) to O ( n ) , where d denotes the model dimension.
We show empirically 1 on applications including MNIST, CIFAR10,
WikiText and GLUE that GraB can outperform random reshufﬂing in terms of both training and validation performance, and even outperform state-of-the-art greedy ordering while reducing memory usage over 100 ⇥ .
Introduction
Many machine learning problems can be formulated as minimizing a differentiable (loss) function f : R d ! R for a set of data examples { x i } ni =1 .
To train a parameterized model, the goal is to obtain a set of target parameters w ⇤ =
arg min f ( w ) , where f ( w ) = 1 n
P ni =1 f ( w ; x i ) , and f ( w ; x i ) denotes the loss incurred on the i -th data example x i (usually a mini-batch of images, sentences, etc.) with model parameters w .
A typical model training, or optimization process, is to iteratively update the model parameter w starting from some initial w (1) by running
[Equation]
where ↵ denotes the step size, and  : { 1 , . . . , n } ! { 1 , . .
. , n } denotes a permutation (ordering) from which the examples are chosen to compute the stochastic gradients 2 .
A widely adopted ordering protocol is Random Reshufﬂing (RR), where the optimizer scans in an order drawn at random without replacement over the entire training dataset multiple times during training (a common name for a single such scan is an “epoch”).
RR allows the optimizer to converge faster empirically and enjoys a better convergence rate in theory [1].
Despite RR’s theoretical improvement, it has been proven that
RR does not always guarantee a good ordering [2, 3]; in fact a random permutation is far from being optimal even when optimizing a simple quadratic objective [4].
In light of this, a natural research question is:
Can we ﬁnd provably better orderings than Random Reshufﬂing—and do so efﬁciently?
Recent studies indicate the possibility of greedily constructing better permutations using stale estimates of each r f ( w ; x j ) from the previous epoch [5, 6].
Concretely, Lu et al.
[5] proved that for any model parameters w 2 R d , if sums of consecutive stochastic gradients converge faster to the full gradient, then the optimizer will converge faster.
Formally, in one epoch, any permutation  minimizing the term (named average gradient error )
[Equation]
leads to fast convergence.
Leveraging this insight, Lu et al.
[5] proposes to greedily select  ( t ) one at a time for all t 2 { 1 , . .
. , n } at the beginning of the ( k + 1) -th
epoch, using the stochastic gradients computed during the k -th epoch as an estimate.
This strategy works well empirically, but it remains an open question how its convergence rate compares to that of RR.
Moreover, the greedy ordering method intensively consumes O ( nd ) memory to store the gradients and O ( n 2 ) computation to order, which signiﬁcantly hinders its usefulness in practice.
For instance, training a simple logistic regression on MNIST would easily cost more than 1 GB additional memory compared to RR.
In this paper, we address the limitations of previous order-selection approaches, culminating in proposing a new algorithm, GraB, which converges faster than RR both in theory and in practice without any blow-up in memory or computational time.
First, we address the theoretical gap by connecting the convergence of permuted-order SGD to the classic herding problem [7]: we show how any algorithm for solving the herding problem can be run on stale stochastic gradients to select an example ordering.
While the greedy selection method of previous work may fail to select a good ordering, using the best herding algorithms from the literature yields a better convergence rate: O ( T  2 / 3 ) compared to RR
’s O ( n 1 / 3 T  2 / 3 ) on smooth non-convex problems and ˜ O
( T  2 ) compared to RR’s ˜ O ( nT  2 ) under PL condition, where T is the total number of iterations.
Building on this, we leverage discrepancy minimization theory [8] to propose an online Gradient Balancing algorithm (GraB) that achieves the same O ( T  2 / 3 ) rate but only requires O ( d ) memory and O ( n ) computation.
Perhaps surprisingly, on multiple deep learning applications we observe GraB not only allows fast minimization of the empirical risk, but also lets the model generalize better.
Our contributions in this paper can be summarized as follows
: • We formulate a general framework of sorting stochastic gradients with herding, and prove SGD with herding converges at a faster rate than Random Reshufﬂing (Section 3).
• We propose an online Gradient Balancing (GraB) algorithm that enjoys the same fast convergence rate as herding SGD while signiﬁcantly reducing the memory from O ( nd ) to O ( d ) and computation from O ( n 2 ) to O ( n ) (Section 5).
• We conduct extensive experiments on MNIST, CIFAR10, WikiText and GLUE with different models.
We demonstrate that GraB converges faster in terms of both training and validation compared to Random Reshufﬂing and other baselines (Section 6).
Related Work
While traditional data ordering research mostly focuses on with-replacement samplers [9–11], withoutreplacement sampling is more common in practice [12].
Random Reshufﬂing (RR) [13] and the related shufﬂe-once (SO) method [14, 15] are among the most popular data permutation methods.
Recht and Ré [16] undertakes the ﬁrst theoretical investigation of RR, while subsequent works like [2, 17] give counter examples where RR orders badly.
HaoChen and Sra [18], Gürbüzbalaban et al.
[19], and Mishchenko et al.
[1] discuss extensively on the conditions needed for RR to beneﬁt.
Some recent works [5, 6] suggest constructing better data permutations than RR via a memory-intensive greedy strategy.
Concretely, Mohtashami et al.
[6] proposes evaluating gradients on all the examples ﬁrst to minimize Equation (2) before starting an epoch, applied to Federated Learning;
Lu et al.
[5] provides an alternative of minimizing Equation (2) using stale gradients from previous epoch to estimate the gradient on each example.
Rajput et al.
[4] introduces an interesting variant to RR by reversing the ordering every other epoch, achieving better rates for quadratics.
Other approaches, such as curriculum learning [20], try to order the data to mimic human learning and improve generalization [21–23]: these approaches differ from our goal of ﬁnding good permutations for minimizing loss in a ﬁnite-sum setting.
Ofﬂine Stale-Gradient Herding
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.
[5] and Mohtashami et al.
[6] propose to use greedy ordering in this way—selecting examples one at a time to minimize Equation (2).
Here, we ﬁrst show that this objective is closely related to the classic herding problem from discrepancy theory.
We demonstrate with examples from the herding formulation that in the adversarial setting, greedy selection is bad at herding and can underperform basic random reshufﬂing.
We conclude this section by proving SGD with proper herding converges faster than random reshufﬂing.
Herding.
The herding problem originates from Welling [24] for sampling from a Markov random ﬁeld that agrees with a given data set.
Its discrete version is later formulated in Harvey and Samadi [7].
Given n vectors { z i } ni =1 that are d -dimensional (i.e. in R d ) and have norm k z i k 2  1 , the herding problem we study 3 is the task of ﬁnding a permutation  ⇤ : [ n ] !
[ n ] that minimizes
[Equation]
State-of-the-art algorithms guarantee an ˜ O (1) bound to this objective [7].
It is straightforward to observe the the formulation of Equation (3) generalizes Equation (2) if we replace each z i with the stochastic gradient computed on the example i .
In other words, any algorithm that minimizes
Equation (3) can be used to minimize Equation (2) and ﬁnd a better data permutation.
The simplest way to do this is with the greedy algorithm for herding, which we show in Algorithm 1: this is inﬁnite
[Table]
essentially what the previous work was doing to select an order based on stale gradients.
, despite its elegance, it was ﬁrst pointed out in Chelidze et
[25] that the greedy can underperform random following.
Statement 1.
There exist n vectors in R 2 such that when applying Algorithm 1, the objective in
Equation (2) becomes ⌦ ( n ) ; a random permutation is guaranteed to achieve O ( p n ) on average.
Herding SGD with Stale Gradients.
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) ofﬂine 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.
[6]) 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.
[5] can be described as running Algorithm 2 using Algorithmas the Herding subroutine.
However, as greedy herding with Algorithm 1 does not have a good worst-case guarantee, we propose using some other herding algorithm instead.
There have been many algorithms proposed that can reduce the Equation (3) to ˜ O
(1) in polynomial time.
Here, we show that this herding-objective bound is sufﬁcient to prove convergence of herded SGD and even yield a better convergence rate than Random Reshufﬂing.
We defer the details of the herding algorithm to
Section 5, as they do not affect the convergence rate of SGD.
We start by stating some assumptions for non-convex optimization, as well as our herding bound.
Assumption 1.
( Smoothness .)
For any example x j in the dataset (for j 2 { 1 , . . . , n } ), the loss for x j is L 2 , 1 -smooth and L 1 -smooth meaning that, for any w , v 2 R d , it holds that
[Equation]
and
[Equation]
The averaged loss function is also L -smooth in that, for any w , v 2 R d , it holds that
[Equation]
Assumption 2.
( PL condition .)
We say the loss function f fulﬁlls the Polyak-Lojasiewicz (PL) condition if there exists µ > 0 such that for any w 2 R d ,
[Equation]
where f ⇤ = inf v 2 R d f ( v
[Equation]
Assumption 3.
( Bounded Gradient Error .)
For any j 2 { 1 , · · · , n } and any w 2 R d ,
[Equation]
Assumption 4.
( Herding Bound .)
In Algorithm 2, the subroutine Herding algorithm has the following property [26]: if given input vectors z 1 , . . .
, z n 2 R d that satisfy
k z i k 2  1 and
P ni =
z i = 0 , Herding outputs a permutation  of { 1 , . . . , n } such that for all k 2 { 1 , . . . , n } , k
P k i =
1 z  ( i ) k 1  H for some constant H .
Remarks on the assumptions.
Assumptions 1, 2, and 3 are commonly used assumptions in the study of permutation-based SGD, except that we separate different smoothness constants with respect to different norms to tighten the ﬁnal bound.
The cross-norm Lipschitz constant L 2 , 1 is often used to obtain tight dimension-independent bounds [27, 28].
Note that these cross-norm assumptions can
[Figure]
of i
is randomly sampled from [0 , 1]
For a given order  the norm of the preﬁx sums k
P kt =
[Equation]
for all k
[Equation]
[Equation]
[Equation] n [Equation] be dropped by incurring extra factors of the dimension d , as L 1  L 2 , 1  p dL for any function.
Assumption 4 holds for multiple algorithms for some H =
˜ O (1) : we will discuss this in depth in
Section 5, while treating H as a parameter here as we want to observe the dependence on this variable in the ﬁnal bound.
Under these assumptions, the convergence rate of Algorithm 2 is given in the following theorem.
Theorem 1.
In Algorithm 2, under Assumptions 1, 3 and 4, if we set ↵ to be
[Equation]
then it converges at the rate
[Equation]
where F = f ( w 1 )  f ⇤ .
Furthermore, under the additional PL condition (Assumption 2), with ↵ as ↵ = min ( [Equation] nµ
[Equation]
where W 0 denotes the Lambert W function, Algorithm 2 converges at the rate
[Equation]
Improved Rate Compared to RR.
Theorem 1 gives us a better rate than random reshufﬂing, in that there is no dependence on the number of examples n in the dominant term.
If we let T = nK be the total number of steps, then Theorem 1 shows Algorithm 2 converges at rate O ( 1 T 2 / 3 + nT ) .
In contrast, random reshufﬂing gives a rate of O ( n 1 / 3 T 2 / 3 + nT ) as given in Mishchenko et al.
[1], Theorem 4.
Similarly, herding improves the rate from ˜ O
( nT 2 ) to ˜ O
( 1 T 2 ) in the PL case.
This gives us an afﬁrmative answer to our question from the intro: a faster rate can be obtained even if we herd with stale gradients estimates, without computing additional gradients.
Enabling Memory-Efﬁcient Herding with Balancing
In Section 3 we introduce the herding problem, and prove that with proper herding, SGD converges at a faster rate than Random Reshufﬂing.
However, a downside of ofﬂine herding SGD is that it requires storing all the vectors (i.e., in the case of SGD, all the stochastic gradients) which incurs
O ( nd ) memory overhead.
In this section, we present background on the vector balancing problem , the solutions of which will enable us to solve the herding problem in an online fashion.
Background: Online Vector Balancing.
The problem of online vector balancing was ﬁrst formulated in Spencer [29]: given n vectors { z i } ni =1 2 R d , arriving one at a time, the goal is to assign a sign ✏ i 2 {  1 , +1 } to each vector upon receiving it so as to minimize max k 2 [ n ] k
P ki =1 ✏ i z
i k 1 .
Over the last few decades, many efforts have been spent on bounding this ` 1 -norm objective.
Speciﬁcally,
Aru et al.
[30] and Bansal and Spencer [31] shows it can be made to be on the order of ˜ O
(1) when the coordinates of each vector are independent .
Bansal et al.
[32] gives a poly (log( n ) , d ) bound when there are dependencies among coordinates.
The adversarial setting with interval discrepancy is discussed in [33].
Other works focus on balancing variants such as stochastic arrival [34] or box discrepancy [35].
Solving Herding via Balancing.
While the objectives of balancing and herding are different, an algorithm for balancing can be used as a subroutine to do herding as well.
Consider the simple case of vectors z 1 , . .
, z n in some order, where
[-=
-][Equation]
= 0 , and suppose we have obtained a group of signs { ✏ i } ni =1 via some balancing algorithm: can we use this to determine an order of { z i } ni =1 that can solve the herding problem?
Harvey and Samadi [7], Chobanyan et al.
[36], and Chobanyan [37] approach this problem by a clever reordering that works as in Algorithm 3.
Concretely, a new order is obtained by ﬁrst taking all the vectors with positive signs in order, followed by all the vectors with negative signs in reversed order from the original sequence.
This process is visualized in Figure 1.
Harvey and Samadi [7] shows that this balancing-and-then-reordering has the following effect on the herding objective in Equation (3).
Theorem 2.
([7], Theorem 10) Consider n vectors z 1 , . .
, z n 2 R d that fulﬁll P ni =
z i {+= 
 and k+} z i k  1 .
Suppose that for some permutation  , the herding objective is bounded by max k [Equation]  H for some constant H .
Suppose we have computed some signs ✏ 1 , . . .
, ✏ n such that the vector-balancing objective max k k
[Equation]
 A for some constant A .
Then if we let  0 be the permutation output by Algorithm 3, the herding objective on this new permutation is bounded by max k
[Equation] Theorem 2 shows that calling Algorithm 3 allows us to construct a new ordering from balancing, which reduces the current herding bound by a constant factor if A  H .
This means that if we have a balancing algorithm that is guaranteed to reduce the balancing objective to less than A , by applying this “reordering” many times, we can also reduce the herding objective to be less than A : this approach is used in many ofﬂine herding algorithms that achieve ˜ O (1) bounds.
GraB: SGD with Gradient Balancing for Good Data Permutation
Building on the balance-then-reorder paradigm from the previous section, in this section we propose an algorithm named GraB that improves SGD with herding by choosing the ordering online —without storing gradients—thus using only O ( d ) memory and O ( n ) computation.
However, the conditions of applying the balancing-reordering framework bring up additional challenges in this online setting:
• Challenge I. To apply Theorem 2, we will need to pre-center all the vectors (as in line 2 of Algorithm 1) to ensure they sum to zero.
This cannot be done online, as we need to make decisions on-the-ﬂy before seeing all the vectors.
• Challenge II.
The reordering strategy, as shown in Theorem 2, only reduces the herding bound by a constant factor, which does not give us the accelerated rate directly since applying Theorem
1: Input: number of epochs K , initialized order  1 , initialized weight w 1 , stale mean m 1 = 0 , step size ↵ . 2: for k = 1 , . .
, K do 3: Initialized indices and running average: l 1 ; r n ; s 0 ; m k + 0 . 4: for t = 1 , . . . , n do 5:
Compute gradient r f ( w tk ; x  k ( t ) ) .
Compute centered gradient and update stale mean: g k,t r f ( w tk ; x  k ( t ) )
m k m k
+ m k +1 + r f ( w tk
; x  k ( t ) ) /n 7:
Compute sign for the current gradient: ✏ k,t
Balancing ( s , g k,t ) 8: if ✏ k,t =
{++1 then 9: s s + g k,t ;  k +
+}
( l )  k ( t ) ; l l + 1 .
10: else 11: s s  g k,t ;  k +1 ( r )  k ( t ) ; r r  1 .
12: end if 13:
Optimizer Step: w ( t +1) k w ( t ) k
↵ r f ( w tk ; x  k ( t ) ) .
14: end for 15: w (1) k
+ w ( n +1) k .
16: end for 17: return w (1) K .
would require Equation (3) to be on the order of ˜ O
We can repeatedly call the balancingreordering subroutines, but that would again require storing all the stochastic gradients.
To overcome challenge I, we apply a two-step stale gradient estimate: for any k , we use the running average of stale gradients to “center” stochastic gradient in epoch k + 1 (this centering is itself stale, and does not guarantee the vectors sum to 0 , but this is still enough to prove convergence).
Then, the online balancing-reordering subroutine will determine the proper order to use in epoch k + 2 .
To address challenge II, we leverage a noisy reordering process: let  k denote the order to use in epoch k , then the next order  k +1 will be obtained from  k .
The intuition here is that if we train the model with enough epochs, the reordering will in general push the herding bound towards A in Theoremwithout storing any additional vectors.
We present the full GraB algorithm in Algorithm 4.
We proceed to provide a convergence guarantee for Algorithm 4.
Since now we are working with balancing rather than herding, in the convergence analysis we replace Assumption 4 with the following one related to balancing.
Assumption 5.
( Balancing Bound. )
There exists a constant A > 0 such that in Algorithm 4, if the subroutine Balancing is given input vectors z 1 , . . .
, z n 2 R d that satisfy k z
i k 2  1 , then
Balancing outputs signs ✏ i 2 {  1 , +1 } such that k [Equation]
We defer discussing the complexity of A to the later part of this section, and ﬁrst present the convergence guarantee for GraB (Algorithm 4) as follows.
Theorem 3.
In Algorithm 4, under Assumptions 1, 3 and 5, if we set ↵ to be
[Equation]
then it converges at the rate
[Equation]
where F = f ( w 1 )  f ⇤ .
Furthermore, under the additional PL condition (Assumption 2) with ↵ as
[Equation]
where W 0 denotes the Lambert W function, it converges at the rate
[Equation]
Comparing Theorem 1 and 3, we can observe GraB
(Algorithm 4) achieves essentially the same fast asymptotic convergence rate as SGD with herding (Algorithm 2).
Asymptotic complexity for A and H .
So far we have covered the main design of GraB. A remaining yet unanswered question is how do A in Assumption 4 and H in Assumption 5 relate to other factors exactly, and what are the concrete algorithms to achieve small A and H .
Some recent progress in the theory community shows that there exists an online algorithm (Algorithm 6) that guarantees with high probability that A is on the order of ˜ O
(1) as shown in the following theorem.
Theorem 4.
(Alweiss et al.
In Algorithm 6, if all the vectors have k z t k  1 and the parameter c is set c = 30 log( nd/  ) , then with probability  , the maximum partial signed sum is bounded by max t k P
[Equation]
This theorem gives us a high-probability bound on A that can easily be converted to a high-probability statement that the result of Theorem 3 holds with A replaced by ˜ O (1) .
Theorem 4 also yields a good ofﬂine herding algorithm and a bound on H : recall from Theorem 2 that given any herding bound
H on an ordering, running Algorithm 3 once allows us to obtain a new herding bound ( A + H ) / 2 .
Given Theorem 4 that a ˜ O (1) bound is guaranteed on A , then if we repeatedly run Algorithmas described in Section 4, we can eventually push the herding bound H towards A , which yields
H =
˜ O (1) .
This can be done with probability 1 ofﬂine by simply restarting on failure; however, as stated in Section 3, this still requires storing all the vectors.
On the normalization of vectors.
Note that a typical balancing subroutine (e.g. Algorithm 6) requires normalized vectors
k z i k  1 . 8 i .
This, in practice, implies we will need to estimate a large enough constant to normalize gradients when we run GraB. To alleviate this issue, we can instead adopt a balancing subroutine that is invariant to the normalizer, for instance, Algorithm 5.
As shown in the next section, this naive balancing algorithm sufﬁces for non-trivial improvements.
Experiment
In this section, we evaluate GraB and other baseline algorithms on multiple machine learning applications.
We compare GraB
(Algorithm 4) with the following baselines: Random Reshufﬂing (RR), Shufﬂe Once (SO) [1], FlipFlop [4] and Greedy Ordering 4 [5].
All the experiments run on an instance conﬁgured with a 4-core Intel(R)
Xeon(R)
2.50GHz CPU, 32GB memory and an NVIDIA
GeForce RTX 2080 Ti GPU.
Although we used Algorithm 6 to show a ˜ O (1) bound can be obtained for herding, in our experiments, we adopt the simple balancing algorithm of Algorithm 5, which performs well in practice.
More results with other balancing algorithms can be found in the Appendix.
Model and Dataset.
We adopt the following model training tasks for evaluation: (1) training logistic regression on MNIST, (2) training LeNet on CIFAR10 [39], (3) training LSTM on wikitext-2 [40], and (4) ﬁnetuning BERT-Tiny [41] on GLUE [42].
Detailed information regarding models, datasets and hyperparameters can be found in Appendix A.
Convergence Speedup.
The convergence plots for each algorithms are summarized in Figure 2.
We observe that in general, GraB is able to converge faster in terms of both training and validation loss on various tasks.
Of all the baseline algorithms, Greedy Ordering is able to achieve comparable convergence speed with respect to epochs to GraB. However, Greedy Ordering consumes much more memory and wall-clock time as shown in the ﬁgure.
In the BERT training case, we observe the
Greedy Ordering runs Out Of Memory (OOM) , causing an error on our machine.
Based on simple calculations, GraB only requires less than 1% the memory used by Greedy Ordering for all the tasks.
It is worth pointing out that throughout the experiments, we do not tune hyperparameters for GraB but let it reuse the hyperparameters from RR (details in the Appendix).
This implies GraB can provide in-place improvement without any additional tuning in practice, and a well-tuned GraB could beneﬁt with even larger margin.
Shufﬂe Once (SO) performs worst in all the cases, and FlipFlop is generally the same as RR, which is aligned with the conclusion made in Rajput et al.
Ablation Study: are good permutations ﬁxed?
In GraB, the example ordering  is carried and improved over epochs, which is motivated from minimizing the herding bound.
This raises another question of whether the found permutation changes at later stage of training, and if so, can we ﬁnd a ﬁxed permutation that outperforms the RR in practice?
Here we conduct an ablation study on two ﬁxed order strategies with GraB.
We include the following two variants:
[Figure]
• 1-step GraB : it only runs GraB for 1 epoch to obtain an ordering, and then use it as a ﬁxed order for the rest of the training.
• Retrain from GraB :
First launch a full run of GraB, and then use the ordering in the ﬁnal epoch as a ﬁxed order in a new run.
We test these new variants on two tasks: LeNet on CIFAR non-convex and logistic regression on
MNIST (convex).
The training loss and validation accuracy curves can be found in Figure 3.
We observe that on both tasks, the 1-step GraB does not work well, which is aligned with our theory and motivation in Section 5 Challenge II.
On the other hand, Retrain from GraB achieves the comparable performance than original GraB on the convex task, but not on the non-convex one.
This is mainly because a good ordering is going to depend on the local optimum the algorithm is approaching and for a convex problem, there is only one such optimum.
It again veriﬁes that GraB ﬁnds better orderings compared to baseline Random Reshufﬂing and Shufﬂe Once.
On the granularity of example ordering.
Throughout this paper, we have discussed how to ﬁnd good data permutations via per-example gradients.
In practice, however, per-example gradients are usually not easy to obtain since data is usually loaded in batches, and many ML libraries (e.g.
PyTorch) directly accumulate gradients over the batch.
A direct workaround is to ﬁx the data within batches and reorder the batches (treating them as coarse-grained examples).
This, however, would compromise the beneﬁts of ordering since the total number of examples n is reduced by a factor of batch size, while the statistical improvement of herding (GraB) is in the order of O ( n  1 / 3 ) .
To alleviate this, we provide two alternatives: (1) Use ML frameworks that support quick per-example gradients computation (e.g. JAX 5 ); (2) Leverage gradient accumulation steps as used in large language model training, i.e., use smaller batch sizes in the code but perform optimizer.step once every few steps so that we can obtain the ﬁner-grained gradients on-the-ﬂy while still optimizing the models with the desired batch size.
This is the default method we used throughout the experiments, we will include more details in the appendix.
Conclusion
In this paper, we formulate a herding framework for data-ordering in SGD.
We prove SGD with herding ﬁnds better data orderings, and converges faster than random reshufﬂing on smooth nonconvex and PL objectives.
We propose an online gradient balancing algorithm named GraB that ﬁnds a better ordering with little compute or memory overhead.
We substantiate our theory and the usefulness of GraB on multiple machine learning applications.
Acknowledgement
This work is supported by NSF-2046760.
Yucheng is supported by Meta PhD Fellowship.
The authors would like to thank Si Yi Meng and anonymous reviewers from NeurIPS 2022 for providing valuable feedbacks on this paper.
[-( k + 4) 
-]
[-In the second step, we apply −⟨ a , b ⟩ = − 12 ∥ a ∥ 2 − 
-]
[-j = t 
-]
[-= ∆ 2 k + 
-]
[-= ∆ 2 k −+ 
-]
[-= ∆ 2 k −≤ 6 n 2 ς 2 + 
-]
[-= ∆ 2 k + (24 n 2 L 2 , ∞ + 
-]