Meta-Learning to Improve Pre-Training

TL;DR

This paper introduces a gradient-based meta-learning algorithm to optimize pre-training hyperparameters, enhancing model performance across diverse real-world tasks by efficiently tuning task weights and augmentation strategies.

Contribution

It presents a novel method combining implicit differentiation and backpropagation to efficiently meta-learn hyperparameters in the two-stage pre-training and fine-tuning process.

Findings

Improved AUROC by up to 3.9% on protein-protein interaction data.

Enhanced AUROC by up to 1.9% on electrocardiography data.

Demonstrated scalability to high-dimensional hyperparameter spaces.

Abstract

Pre-training (PT) followed by fine-tuning (FT) is an effective method for training neural networks, and has led to significant performance improvements in many domains. PT can incorporate various design choices such as task and data reweighting strategies, augmentation policies, and noise models, all of which can significantly impact the quality of representations learned. The hyperparameters introduced by these strategies therefore must be tuned appropriately. However, setting the values of these hyperparameters is challenging. Most existing methods either struggle to scale to high dimensions, are too slow and memory-intensive, or cannot be directly applied to the two-stage PT and FT learning process. In this work, we propose an efficient, gradient-based algorithm to meta-learn PT hyperparameters. We formalize the PT hyperparameter optimization problem and propose a novel method to obtain PT hyperparameter gradients by combining implicit differentiation and backpropagation through unrolled optimization. We demonstrate that our method improves predictive performance on two real-world domains. First, we optimize high-dimensional task weighting hyperparameters for multitask pre-training on protein-protein interaction graphs and improve AUROC by up to 3.9%. Second, we optimize a data augmentation neural network for self-supervised PT with SimCLR on electrocardiography data and improve AUROC by up to 1.9%.