The benefits of full data shuffle, now with optimal I/O cost: $k$-wise independence and matrix transposition to the rescue

It is known that RandomShuffle, the without replacement version of Stochastic Gradient Descent (SGD), converges faster than with replacement SGD. However, RandomShuffle requires uniformly performing a random permutation of the input sequence, which is known to have high I/O complexity due to data movement across the memory hierarchy. In this paper, we propose a shuffling algorithm with a linear I/O complexity that generates almost-uniformly random permutations with rigorous mathematical guarantees. Specifically, we show that the shuffling algorithm can generate $2$-wise independent permutations. Furthermore, we can extend to $k$-wise independence with a small error in the probability distribution, if the fast memory has at least $k$ memory blocks. These results allow us to reach the same expected theoretical convergence as RandomShuffle while achieving optimal linear I/O cost.

Modern machine learning systems are often trained on enormous datasets, and the order in which training examples are shown to the model can strongly affect learning. Fully shuffling the data is known to improve training compared with sampling examples at random with replacement, but a true full shuffle can be expensive because it requires many slow data movements between memory levels. This becomes a serious bottleneck when datasets are larger than fast memory. We study whether it is possible to keep the training benefits of random shuffling while avoiding its high data movement cost. Our main idea is to use carefully structured shuffles that are not fully random, but still have enough randomness to preserve the theoretical advantages needed for Stochastic Gradient Descent. We design algorithms based on limited independence and matrix transposition that use only a linear number of input/output operations. We prove that these shuffles can match the expected convergence guarantees of fully random reshuffling for strongly convex optimization problems, while achieving optimal data movement cost. This suggests that large-scale training can be made more efficient without giving up the benefits of shuffled data.