Guarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer Samples

A driving force behind the diverse applicability of modern machine learning is the ability to extract meaningful features across many sources. However, many practical domains involve data that are non-identically distributed across sources, and statistically dependent within its source, violating vital assumptions in existing theoretical studies. Toward addressing these issues, we establish statistical guarantees for learning general textitnonlinear representations from multiple data sources that admit different input distributions and possibly dependent data. Specifically, we study the sample-complexity of learning T+1 functions f_star(t) circ g_star from a function class mathcal F × mathcal G, where f_star(t) are task specific linear functions and g_star is a shared nonlinear representation. A representation hat g is estimated using N samples from each of T source tasks, and a fine-tuning function hat f(0) is fit using N' samples from a target task passed through hat g. We show that when N gtrsim Cdep (dim(mathcal F) + C(mathcal G)/T), the excess risk of hat f(0) circ hat g on the target task decays as nudiv big(fracdim(mathcal F)N' + fracC(mathcal G)N T big), where Cdep denotes the effect of data dependency, nudiv denotes an (estimatable) measure of textittask-diversity between the source and target tasks, and mathrm C(mathcal G) denotes the complexity of the representation class mathcal G. In particular, our analysis reveals: as the number of tasks T increases, both the sample requirement and risk bound converge to that of r-dimensional regression as if g_star had been given, and the effect of dependency only enters the sample requirement, leaving the risk bound matching the iid setting.