Learning to Shuffle: Block Reshuffling and Reversal Schemes for Stochastic Optimization

TL;DR

This paper introduces a novel data shuffling scheme for stochastic gradient descent, combining block reshuffling and paired reversal, guided by a large language model, to improve convergence and stability.

Contribution

It designs a new shuffling algorithm using LLM-guided program evolution, analyzing its structural components for provable improvements over existing schemes.

Findings

Block reshuffling reduces prefix-gradient variance constants.

Paired reversal cancels leading second-order terms, reducing order sensitivity.

Numerical experiments show consistent gains over standard shuffling schemes.

Abstract

Shuffling strategies for stochastic gradient descent (SGD), including incremental gradient, shuffle-once, and random reshuffling, are supported by rigorous convergence analyses for arbitrary within-epoch permutations. In particular, random reshuffling is known to improve optimization constants relative to cyclic and shuffle-once schemes. However, existing theory offers limited guidance on how to design new data-ordering schemes that further improve optimization constants or stability beyond random reshuffling. In this paper, we design a pipeline using a large language model (LLM)-guided program evolution framework to discover an effective shuffling rule for without-replacement SGD. Abstracting from this instance, we identify two fundamental structural components: block reshuffling and paired reversal. We analyze these components separately and show that block reshuffling strictly reduces prefix-gradient variance constants within the unified shuffling framework, yielding provable improvements over random reshuffling under mild conditions. Separately, we show that paired reversal symmetrizes the epoch map and cancels the leading order-dependent second-order term, reducing order sensitivity from quadratic to cubic in the step size. Numerical experiments with the discovered algorithm validate the theory and demonstrate consistent gains over standard shuffling schemes across convex and nonconvex benchmarks.