In-Loop Meta-Learning with Gradient-Alignment Reward

TL;DR

This paper introduces the gradient-alignment reward (GAR), a computationally efficient method to optimize training strategies by aligning gradients, improving generalization and guiding data distribution and augmentation choices.

Contribution

It proposes GAR as a novel, cheap-to-compute reward for in-loop meta-learning, enabling better data and augmentation strategies during training.

Findings

GAR effectively guides data distribution selection.

GAR-guided augmentation strategies perform competitively.

The method improves training generalization.

Abstract

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.