Research map / Label-efficient and robust learning / Recognition and long-tail
Provable Multi-Task Representation Learning by Two-Layer ReLU Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves that gradient-based multi-task pretraining (head update then one first-layer gradient step) of a two-layer ReLU network on many binary tasks recovers the r-dimensional label-relevant subspace, enabling downstream generalization with sample and neuron complexity independent of the input dimension d.
Problem
Why and how nonlinear neural networks learn effective shared feature representations when pretrained with gradient descent on many tasks, and how the number of pretraining tasks affects downstream performance.
Approach
Updating task-specific heads before the shared first layer induces a pseudo-contrastive loss that aligns inputs likely to share a label across random tasks, so one gradient step recovers the ground-truth subspace; a random second ReLU layer on top enables generalization to any downstream link function.
Key results
- One gradient step of head-then-representation multi-task training projects first-layer weights onto the r-dimensional ground-truth subspace, with error shrinking with T and samples per task.
- Downstream generalization guarantee: sample and neuron complexity exp(Otilde(r^5)) independent of input dimension d, with the random second layer enabling linear separation of any binary function on the r features.
- Negative results: random features need neuron or sample complexity polynomial in d^r, and with high probability a single pretraining task cannot recover all r features.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use gradient-based multi-task pretraining across diverse tasks | Two-layer ReLU network · Shared-feature binary classification | representation error ratio: O\left( \tfrac{ {r^{3.5}} + \log^{3.5}(m/\delta) }{ d^{1.5} } + 2^r \epsilon \right) | Helped |
| Use the pretrained representation with a random ReLU layer and train a task-specific linear head | Multi-task-pretrained representation · Downstream binary classification on shared features | population hinge loss: \frac{ \exp(\tilde{O}(r^5)) }{ \sqrt{N}} | Helped |
| Pretrain on random features without feature learning | Downstream sparse parity classification · Sparse parity on d inputs | neuron dimension and inverse margin requirement: polynomially large in d^r | Hurt |
| Pretrain on one randomly selected task instead of multiple tasks | Downstream task on shared ground-truth features · Single-task sparse parity pretraining | neuron dimension and inverse margin requirement: Ω(d) | Hurt |
| Add task-specific head updates before updating the representation (ablation) | Multi-task pretraining with a two-layer ReLU network · Ground-truth feature recovery | representation learning | Helped |
Surprising findings
- A single pretraining task, even with infinite samples, can fail to recover all r ground-truth features with high probability over the task draw.
Limitations
- Restricted to binary classification tasks with Gaussian inputs and a two-layer ReLU network with a single representation step.
- Requires the task distribution to satisfy a balance assumption (label correlation 1/2) and d^3-scale sample complexity; exponential complexity in r.
- Biases are not updated during pretraining and the analysis uses a single gradient step.
Experimental setup
Datasets: synthetic Gaussian data with sparse parity and uniform link-function tasks
Metrics: representation learning error (subspace projection) downstream population hinge loss bound
Baselines: random features single-task training
Builds on: neural tangent kernel sparse coding feature learning multi-task representation learning with linear models
Closest papers by meaning
- Inductive biases of multi-task learning and finetuning: multiple regimes of feature reuse (NeurIPS 2024 · similarity 0.88)
- Task structure and nonlinearity jointly determine learned representational geometry (ICLR 2024 · similarity 0.88)
- Make Haste Slowly: A Theory of Emergent Structured Mixed Selectivity in Feature Learning ReLU Networks (ICLR 2025 · similarity 0.86)
- A New Neural Kernel Regime: The Inductive Bias of Multi-Task Learning (NeurIPS 2024 · similarity 0.86)
- Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.86)
- On the Emergence of Cross-Task Linearity in Pretraining-Finetuning Paradigm (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.