Provable Multi-Task Representation Learning by Two-Layer ReLU Neural Networks

TL;DR

This paper proves that multi-task pretraining with two-layer ReLU neural networks enables feature learning of underlying data projections, outperforming single-task training in recovering true features with sample complexity independent of input dimension.

Contribution

It provides the first theoretical proof that nonlinear multi-task training induces feature learning, specifically recovering data projections in high-dimensional spaces.

Findings

Multi-task pretraining induces a pseudo-contrastive loss.

Gradient-based training recovers the true data projection.

Single-task training generally fails to learn all ground-truth features.

Abstract

An increasingly popular machine learning paradigm is to pretrain a neural network (NN) on many tasks offline, then adapt it to downstream tasks, often by re-training only the last linear layer of the network. This approach yields strong downstream performance in a variety of contexts, demonstrating that multitask pretraining leads to effective feature learning. Although several recent theoretical studies have shown that shallow NNs learn meaningful features when either (i) they are trained on a {\em single} task or (ii) they are {\em linear}, very little is known about the closer-to-practice case of {\em nonlinear} NNs trained on {\em multiple} tasks. In this work, we present the first results proving that feature learning occurs during training with a nonlinear model on multiple tasks. Our key insight is that multi-task pretraining induces a pseudo-contrastive loss that favors representations that align points that typically have the same label across tasks. Using this observation, we show that when the tasks are binary classification tasks with labels depending on the projection of the data onto an $r$-dimensional subspace within the $d\gg r$-dimensional input space, a simple gradient-based multitask learning algorithm on a two-layer ReLU NN recovers this projection, allowing for generalization to downstream tasks with sample and neuron complexity independent of $d$. In contrast, we show that with high probability over the draw of a single task, training on this single task cannot guarantee to learn all $r$ ground-truth features.