# Representation Learning Beyond Linear Prediction Functions

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 26 references that are in the map.
- First posted on arXiv 31 May 2021.

## Builds on (most cited first, 12 of 12)

- [A Survey on Deep Transfer Learning](https://mlsanity.com/paper/1808.01974/) (2018)
- [Lipschitz regularity of deep neural networks: analysis and efficient estimation](https://mlsanity.com/paper/1805.10965/) (2018)
- [Efficient and Accurate Estimation of Lipschitz Constants for Deep Neural Networks](https://mlsanity.com/paper/1906.04893/) (2019)
- [Size-Independent Sample Complexity of Neural Networks](https://mlsanity.com/paper/1712.06541/) (2017)
- [Automated Curriculum Learning for Neural Networks](https://mlsanity.com/paper/1704.03003/) (2017)
- [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)
- [On the Theory of Transfer Learning: The Importance of Task Diversity](https://mlsanity.com/paper/2006.11650/) (2020)
- [Model-based Reinforcement Learning and the Eluder Dimension](https://mlsanity.com/paper/1406.1853/) (2014)
- [Provable Model-based Nonlinear Bandit and Reinforcement Learning: Shelve Optimism, Embrace Virtual Curvature](https://mlsanity.com/paper/2102.04168/) (2021)
- [Acquiring Knowledge from Pre-trained Model to Neural Machine Translation](https://mlsanity.com/paper/1912.01774/) (2019)
- [Hypothesis Transfer Learning via Transformation Functions](https://mlsanity.com/paper/1612.01020/) (2016)

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

- [FedAvg with Fine Tuning: Local Updates Lead to Representation Learning](https://mlsanity.com/paper/2205.13692/) (2022)
- [Calibrated Multiple-Output Quantile Regression with Representation Learning](https://mlsanity.com/paper/2110.00816/) (2021)
- [Provable Multi-Task Representation Learning by Two-Layer ReLU Neural Networks](https://mlsanity.com/paper/2307.06887/) (2023)
- [Active Multi-Task Representation Learning](https://mlsanity.com/paper/2202.00911/) (2022)
- [Provable and Efficient Continual Representation Learning](https://mlsanity.com/paper/2203.02026/) (2022)
- [On the Statistical Benefits of Curriculum Learning](https://mlsanity.com/paper/2111.07126/) (2021)
- [Learning an Explicit Hyperparameter Prediction Function Conditioned on Tasks](https://mlsanity.com/paper/2107.02378/) (2021)
- [Improved Active Multi-Task Representation Learning via Lasso](https://mlsanity.com/paper/2306.02556/) (2023)
- [Smoothness Adaptive Hypothesis Transfer Learning](https://mlsanity.com/paper/2402.14966/) (2024)
- [On the Generalization Ability of Unsupervised Pretraining](https://mlsanity.com/paper/2403.06871/) (2024)
- [NeuCo-Bench: A Novel Benchmark Framework for Neural Embeddings in Earth Observation](https://mlsanity.com/paper/2510.17914/) (2025)
- [DAC-MR: Data Augmentation Consistency Based Meta-Regularization for Meta-Learning](https://mlsanity.com/paper/2305.07892/) (2023)
- [Provable Pathways: Learning Multiple Tasks over Multiple Paths](https://mlsanity.com/paper/2303.04338/) (2023)
- [On Hypothesis Transfer Learning of Functional Linear Models](https://mlsanity.com/paper/2206.04277/) (2022)
- [Active Representation Learning for General Task Space with Applications in Robotics](https://mlsanity.com/paper/2306.08942/) (2023)
- [Sample Efficient Myopic Exploration Through Multitask Reinforcement Learning with Diverse Tasks](https://mlsanity.com/paper/2403.01636/) (2024)
- [Guarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer Samples](https://mlsanity.com/paper/2410.11227/) (2024)
- [Bridging Lifelong and Multi-Task Representation Learning via Algorithm and Complexity Measure](https://mlsanity.com/paper/2511.01847/) (2025)
- [Beyond Centralization: Provable Communication Efficient Decentralized Multi-Task Learning](https://mlsanity.com/paper/2512.22675/) (2025)
- [Personalized Alignment Revisited: The Necessity and Sufficiency of User Diversity](https://mlsanity.com/paper/2605.09119/) (2026)
- [Transfer Learning in Nonparametric Regression with Deep ReLU Networks](https://mlsanity.com/paper/2608.20255/) (2026)

## Abstract (arXiv)

Recent papers on the theory of representation learning has shown the importance of a quantity called diversity when generalizing from a set of source tasks to a target task. Most of these papers assume that the function mapping shared representations to predictions is linear, for both source and target tasks. In practice, researchers in deep learning use different numbers of extra layers following the pretrained model based on the difficulty of the new task. This motivates us to ask whether diversity can be achieved when source tasks and the target task use different prediction function spaces beyond linear functions. We show that diversity holds even if the target task uses a neural network with multiple layers, as long as source tasks use linear functions. If source tasks use nonlinear prediction functions, we provide a negative result by showing that depth-1 neural networks with ReLu activation function need exponentially many source tasks to achieve diversity. For a general function class, we find that eluder dimension gives a lower bound on the number of tasks required for diversity. Our theoretical results imply that simpler tasks generalize better. Though our theoretical results are shown for the global minimizer of empirical risks, their qualitative predictions still hold true for gradient-based optimization algorithms as verified by our simulations on deep neural networks.