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

Area of Science:

• Machine Learning
• Artificial Intelligence
• Deep Learning Theory

Background:

• Multi-task pretraining is a popular machine learning paradigm for adapting neural networks (NNs) to downstream tasks.
• Existing theory confirms feature learning in shallow NNs for single tasks or linear models, but not for nonlinear models trained on multiple tasks.
• Understanding feature learning in nonlinear NNs during multi-task pretraining is crucial for practical applications.

Purpose of the Study:

• To provide the first theoretical proof of feature learning during multi-task pretraining with nonlinear neural networks.
• To analyze the conditions under which multi-task learning algorithms can effectively learn underlying data representations.

Main Methods:

• Investigated a two-layer ReLU neural network trained with a gradient-based multi-task learning algorithm.
• Analyzed binary classification tasks where labels depend on projections onto low-dimensional subspaces.
• Introduced the concept of a pseudo-contrastive loss induced by multi-task pretraining.

Main Results:

• Demonstrated that multi-task pretraining induces a pseudo-contrastive loss, aligning data points with similar labels across tasks.
• Showed that the proposed algorithm recovers the underlying projection, enabling generalization to downstream tasks.
• Proved that sample and neuron complexity for generalization are independent of the subspace dimension.

Conclusions:

• Feature learning is proven to occur in nonlinear neural networks trained on multiple tasks.
• Multi-task pretraining offers advantages over single-task training, which may fail to learn all relevant features.
• This work provides theoretical foundations for the effectiveness of multi-task learning in deep neural networks.