In-Loop Meta-Learning with Gradient-Alignment Reward
This work addresses the problem of improving generalization in deep learning training for researchers and practitioners, offering an incremental improvement to existing training loops.
This paper proposes a method to maximize training generalization by optimizing the loss of the next training step, using a novel and computationally efficient gradient-alignment reward (GAR). They apply GAR to optimize data distributions and learning augmentation strategies, achieving competitive results on CIFAR-10 and CIFAR-100.
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.