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

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.

References in corpus (39)

• A Simple Framework for Contrastive Learning of Visual Representations
• An Overview of Multi-Task Learning in Deep Neural Networks
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• A Convergence Theory for Deep Learning via Over-Parameterization
• A Mean Field View of the Landscape of Two-Layers Neural Networks
• Multi-Task Learning with Deep Neural Networks: A Survey
• Breaking the Curse of Dimensionality with Convex Neural Networks
• Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent
• Stochastic Gradient Descent Optimizes Over-parameterized Deep ReLU Networks
• Gradient Descent Finds Global Minima of Deep Neural Networks
• On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport
• Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks
• On the Theory of Transfer Learning: The Importance of Task Diversity
• FedAvg with Fine Tuning: Local Updates Lead to Representation Learning
• Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks
• Multitask Prompt Tuning Enables Parameter-Efficient Transfer Learning
• Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit
• High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation
• Neural Networks can Learn Representations with Gradient Descent
• Neural tangent kernels, transportation mappings, and universal approximation
• The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks
• Poly-time universality and limitations of deep learning
• Approximate is Good Enough: Probabilistic Variants of Dimensional and Margin Complexity
• How Two-Layer Neural Networks Learn, One (Giant) Step at a Time
• A Local Convergence Theory for Mildly Over-Parameterized Two-Layer Neural Network
• Neural Networks Efficiently Learn Low-Dimensional Representations with SGD
• MAML is a Noisy Contrastive Learner in Classification
• A Theoretical Analysis on Feature Learning in Neural Networks: Emergence from Inputs and Advantage over Fixed Features
• Towards Sample-efficient Overparameterized Meta-learning
• MAML and ANIL Provably Learn Representations
• Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks
• SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics
• Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models
• When Hardness of Approximation Meets Hardness of Learning
• Active Multi-Task Representation Learning
• Feature selection with gradient descent on two-layer networks in low-rotation regimes
• First-order ANIL provably learns representations despite overparametrization
• Learning Hierarchical Polynomials with Three-Layer Neural Networks
• Excess Risk of Two-Layer ReLU Neural Networks in Teacher-Student Settings and its Superiority to Kernel Methods