  At the heart of the standard deep learning training loop is a greedy gradient
step minimizing a given loss. We propose to add a second step to maximize
training generalization. To do this, we optimize the loss of the next training
step. While computing the gradient for this generally is very expensive and
many interesting applications consider non-differentiable parameters (e.g. due
to hard samples), we present a cheap-to-compute and memory-saving reward, the
gradient-alignment reward (GAR), that can guide the optimization. We use this
reward to optimize multiple distributions during model training. First, we
present the application of GAR to choosing the data distribution as a mixture
of multiple dataset splits in a small scale setting. Second, we show that it
can successfully guide learning augmentation strategies competitive with
state-of-the-art augmentation strategies on CIFAR-10 and CIFAR-100.
