Research map / Learning theory and algorithms / Sample complexity
Guarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer Samples
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.
Gives sample-complexity and excess-risk guarantees for two-stage ERM representation learning across tasks with non-identical covariate distributions and phi-mixing (dependent) data, with a refined task-diversity measure and a plug-in estimator of it.
Problem
Existing multi-task nonlinear representation learning theory assumes i.i.d. covariates across tasks and independent samples within tasks, and requires per-task sample burn-in proportional to the input dimension; the paper targets these limitations.
Approach
Replaces the uniform task-diversity assumption with a task-diversity coefficient that jointly accounts for covariate coverage (mu_X) and head coverage (mu_F), proves excess-risk bounds under non-identical and phi-mixing covariates via blocking that moves dependence effects into the burn-in, and reduces per-task burn-in from order dim(X) to order r.
Key results
- Multi-task pretraining of the representation improves over single-task imitation learning in the low-target-data regime on all three cartpole metrics.
- Multi-task gains saturate once target trajectories exceed the number of per-task training trajectories (10), consistent with the theory's limiting rate for large T.
- Theory: per-task burn-in for the target head reduces from order dim(X) in prior work to order r, and the risk bound matches the iid rate under phi-mixing.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| pretrain on source tasks and fit only the final layer on target demonstrations instead of training the full network | Multi-task imitation learning policy · Cartpole imitation learning, low-target-data regime | controller-input MSE | Helped |
| pretrain on source tasks and fit only the final layer on target demonstrations instead of training the full network | Multi-task imitation learning policy · Cartpole imitation learning, low-target-data regime | trajectory MSE | Helped |
| pretrain on source tasks and fit only the final layer on target demonstrations instead of training the full network | Multi-task imitation learning policy · Cartpole imitation learning, low-target-data regime | stabilizing trials (%) | Helped |
Limitations
- Task-diversity estimator is shown only consistent asymptotically; finite-sample guarantees left open.
- Analysis assumes realizable setting with subgaussian martingale noise and bounded function classes; the excess-risk bound is for ERM rather than a practical algorithm.
- Pathologies such as class imbalance and low task diversity are not addressed.
Experimental setup
Datasets: PyBullet simulated cartpole demonstrations from linear expert controllers
Benchmarks: PyBullet cartpole imitation learning (controller MSE, trajectory deviation, % balanced trials)
Metrics: MSE of controller inputs trajectory deviation percent of trials stabilized
Baselines: single-task imitation learning
Builds on: multi-task linear representation learning (Du et al. / Tripuraneni et al.) offset Rademacher / martingale offset complexity task diversity (Tripuraneni et al.) blocking for mixing processes
Closest papers by meaning
- Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data (ICLR 2024 · similarity 0.88)
- Robustness of Nonlinear Representation Learning (ICML 2024 · similarity 0.88)
- Generalization Guarantees for Representation Learning via Data-Dependent Gaussian Mixture Priors (ICLR 2025 · similarity 0.86)
- A Unified Framework for Learning with Nonlinear Model Classes from Arbitrary Linear Samples (ICML 2024 · similarity 0.86)
- Sharp Rates in Dependent Learning Theory: Avoiding Sample Size Deflation for the Square Loss (ICML 2024 · similarity 0.86)
- How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs (NeurIPS 2025 · similarity 0.85)
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.