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

arXiv: https://arxiv.org/abs/2307.06887 · HTML: https://mlsanity.com/paper/2307.06887/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 62 of its 75 references that are in the map.
- First posted on arXiv 13 Jul 2023.

## Builds on (most cited first, 30 of 62)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Multitask Prompted Training Enables Zero-Shot Task Generalization](https://mlsanity.com/paper/2110.08207/) (2021)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [On Exact Computation with an Infinitely Wide Neural Net](https://mlsanity.com/paper/1904.11955/) (2019)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Multi-Task Learning with Deep Neural Networks: A Survey](https://mlsanity.com/paper/2009.09796/) (2020)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data](https://mlsanity.com/paper/1808.01204/) (2018)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Few-Shot Learning via Learning the Representation, Provably](https://mlsanity.com/paper/2002.09434/) (2020)
- [The Benefit of Multitask Representation Learning](https://mlsanity.com/paper/1505.06279/) (2015)
- [Muppet: Massive Multi-task Representations with Pre-Finetuning](https://mlsanity.com/paper/2101.11038/) (2021)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [On the Theory of Transfer Learning: The Importance of Task Diversity](https://mlsanity.com/paper/2006.11650/) (2020)
- [Towards moderate overparameterization: global convergence guarantees for training shallow neural networks](https://mlsanity.com/paper/1902.04674/) (2019)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)

## Built on by (most cited first, 16 of 16)

- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data](https://mlsanity.com/paper/2308.04428/) (2023)
- [Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations](https://mlsanity.com/paper/2406.11828/) (2024)
- [Inductive biases of multi-task learning and finetuning: multiple regimes of feature reuse](https://mlsanity.com/paper/2310.02396/) (2023)
- [Robust Feature Learning for Multi-Index Models in High Dimensions](https://mlsanity.com/paper/2410.16449/) (2024)
- [Guarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer Samples](https://mlsanity.com/paper/2410.11227/) (2024)
- [Operator-Based Generalization Bound for Deep Learning: Insights on Multi-Task Learning](https://mlsanity.com/paper/2512.19184/) (2025)
- [On the Koopman-Based Generalization Bounds for Multi-Task Deep Learning](https://mlsanity.com/paper/2512.19199/) (2025)
- [The Effects of Multi-Task Learning on ReLU Neural Network Functions](https://mlsanity.com/paper/2410.21696/) (2024)
- [Meta-learning of shared linear representations beyond well-specified linear regression](https://mlsanity.com/paper/2501.18975/) (2025)
- [Evidential Uncertainty Probes for Graph Neural Networks](https://mlsanity.com/paper/2503.08097/) (2025)
- [Syntax-Guided Diffusion Language Models with User-Integrated Personalization](https://mlsanity.com/paper/2510.01028/) (2025)
- [Feed m Birds with One Scone: Accelerating Multi-task Gradient Balancing via Bi-level Optimization](https://mlsanity.com/paper/2603.07389/) (2026)
- [Near-optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation](https://mlsanity.com/paper/2605.00473/) (2026)
- [Early Prediction of Future Behavioral Strategy from Process Traces](https://mlsanity.com/paper/2605.30550/) (2026)
- [High-Dimensional Asymptotics and Dataset Selection for Private Transfer Learning](https://mlsanity.com/paper/2610.02578/) (2026)

## Abstract (arXiv)

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 dgg 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.