Representation Learning Beyond Linear Prediction Functions

1 Introduction . It has become a common practice ( Tan et al. , 2018 ) to use a pre-trained network as the representation for a new task with small sample size in various areas including computer vision ( Marmanis et al. , 2015 ) , speech recognition ( Dahl et al. , 2011 ; Jaitly et al. , 2012 ; Howard and Ruder , 2018 ) and machine translation ( Weng et al. , 2020 ) . Most representation learning is based on the assumption that the source tasks and the target task share the same low-dimensional representation . In this paper , following the work of Tripuraneni et al . ( 2020 ) , we assume that each task , indexed by t , generates observations noisily from the mean function f∗t ◦ h∗ , where h∗ is the true representation shared by all the tasks and f∗t is called the prediction function . We assume h ∗ ∈ H , the representation function space and ft ∈ Ft , the prediction function space . Tripuraneni et al . ( 2020 ) proposed a diversity condition which guarantees that the learned representation will generalize to target tasks with any prediction function . Assume we have T source tasks labeled by 1 , . . . , T and F1 = · · · = FT = Fso . We denote the target task by ta and let Et ( f , h ) ∈ R be the excess error of f , h ∈ Ft×H for task t. The diversity condition can be stated as follows : there exists some ν > 0 , such that for any f∗ta ∈ Fta and any h ∈ H , inf fta∈Fta Eta ( fta , h ) Excess error given h for target task ≤ ν inf ft∈Fso , t=1 , ... , T 1 T T∑ t=1 Et ( ft , h ) Excess error given h for source tasks . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . The diversity condition relates the excess error for source tasks and the target task with respect to any fixed representation h. A smaller ν indicates a better transfer from source tasks to the target task . Generally , we say T source tasks achieve diversity over Fta when ν is finite and relatively small . The number of source tasks plays an important role here . To see this , assume Fso = Fta = F is a discrete function space and each function can be arbitrarily different . In this case , we will need the target task to be the same as at least one source task , which in turn requires T ≥ |F| and ν ≥ |F| . So far , it has only been understood that when both source tasks and target task use linear prediction functions , it takes at least d source tasks to be diverse , where d is the number of dimension of the linear mappings . This paper answers the following two open questions with a focus on the deep neural network ( DNN ) models : How does representation learning work when Fso 6= Fta ? How many source tasks do we need to achieve diversity with nonlinear prediction functions ? There are strong practical motivations to answer the above two questions . In practice , researchers use representation learning despite the difference in difficulty levels of source and target tasks . We use a more complex function class for substantially harder target task , which means Fso 6= Fta . This can be reflected as extra layers when a deep neural network model is used as prediction functions . On the other hand , the source task and target task may have different objectives . Representation pretrained on a classification problem , say ImageNet , may be applied to object detection or instance segmentation problems . For instance , Oquab et al . ( 2014 ) trained a DNN on ImageNet and kept all the layers as the representation except for the last linear mapping , while two fully connected layers are used for the target task on object detection . Another motivation for our work is the mismatch between recently developed theories in representation learning and the common practice in empirical studies . Recent papers on the theory of representation learning all require multiple sources tasks to achieve diversity so as to generalize to any target task in F ( Maurer et al. , 2016 ; Du et al. , 2020 ; Tripuraneni et al. , 2020 ) . However , most pretrained networks are only trained on a single task , for example , the ImageNet pretrained network . To this end , we will show that a single multi-class classification problem can be diverse . Lastly , while it is common to simply use linear mapping as the source prediction function , there is no clear theoretical analysis showing whether or not diversity can be achieved with nonlinear prediction function spaces . Main contributions . We summarize the main contributions made by this paper . 1 . We show that diversity over Fso implies diversity over Fta , when both Fso and Fta are DNNs and Fta has more layers . More generally , the same statement holds when Fta is more complicated than Fso , in the way that Fta = F ′ta ◦ ( F⊗mso ) for some positive integer m1 and function class F ′ta . 2 . Turning our attention to the analysis of diversity for non-linear prediction function spaces , we show that for a depth-1 NN , it requires Ω ( 2d ) many source tasks to establish diversity with d being the representation dimension . For general Fso , we provide a lower bound on the number of source tasks required to achieve diversity using the eluder dimension ( Russo and Van Roy , 2013 ) and provide a upper bound using the generalized rank ( Li et al. , 2021 ) . 3 . We show that , from the perspective of achieving diversity , a single source task with multiple outputs can be equivalent to multiple source tasks . While our theories are built on empirical risk minimization , our simulations on DNNs for a multi-variate regression problem show that the qualitative predictions our theory makes still hold when stochastic gradient descent is used for optimization . 2 Preliminaries . We first introduce the mathematical setup of the problem studied in this paper along with the two-phase learning method that we will focus on . 1F⊗T is the T times Cartesian product of F . Problem setup . Let X denote the input space . We assume the same input distribution PX for all tasks , as covariate shift is not the focus of this work . In our representation learning setting , there exists a generic feature representation function h∗ ∈ H : X 7→ Z that is shared across different tasks , where Z is the feature space and H is the representation function space . Since we only consider the different prediction functions , each task , indexed by t , is defined by its prediction function f∗t ∈ Ft : Z 7→ Yt ⊂ [ 0 , 1 ] , where Ft is the prediction function space of task t and Yt is the corresponding output space . The observations Yt = f∗t ◦ h∗ ( X ) + are generated noisily with mean function f∗t ◦ h∗ , where X ∼ PX and is zero-mean noise that is independent of X . Our representation is learned on source tasks f∗so : = ( f ∗ 1 , . . . , f ∗ T ) ∈ F1×· · ·×FT for some positive integer T . We assume that all the prediction function spaces Ft = Fso are the same over the source tasks . We denote the target task by f∗ta ∈ Fta : Z 7→ Yt ⊂ [ 0 , 1 ] , where Fta is the target prediction function space . Unlike the previous papers ( Du et al. , 2020 ; Tripuraneni et al. , 2020 ; Maurer et al. , 2016 ) , which all assume that the same prediction function space is used for all tasks , we generally allow for the possibility that Fso 6= Fta . Learning algorithm . We consider the same two-phase learning method as in Tripuraneni et al . ( 2020 ) . In the first phase ( the training phase ) , n = ( n1 , . . . , nT ) samples from each task are available to learn a good representation . In the second phase ( the test phase ) , we are presented nta samples from the target task to learn its prediction function using the pretrained representation learned in the training phase . We denote a dataset of size n from task ft by Snt = { ( xti , yti ) } ni=1 . We use empirical risk minimization ( ERM ) for both phase . In the training phase , we minimize average risks over { Sntt } Tt=1 : R̂ ( f , h | f∗so ) : = 1∑ t nt T∑ t=1 nt∑ i=1 lso ( ft ◦ h ( xti ) , yti ) , where lso : Y × Y 7→ R is the loss function for the source tasks and f = ( f1 , . . . , fT ) ∈ F⊗Tso . The estimates are given by ( f̂so , ĥ ) ∈ arg minf , h R̂ ( f , h | f∗so ) . In the second phase , we obtain the dataset { xta , i , yta , i } ntai from the target task and our predictor f̂ta is given by arg min f∈Fta R̂ ( f , ĥ | f∗ta ) : = 1 nta nta∑ i=1 lta ( f ◦ ĥ ( xta , i ) , yta , i ) , for some loss function lta on the target task . We also use R ( · , · | · ) for the expectation of the above empirical risks . We denote the generalization error of certain estimates f , h by E ( f , h | · ) : = R ( f , h | · ) −minf ′∈F , h′∈HR ( f ′ , h′ | · ) , where F can be either F⊗Tso or Fta depending on the tasks . Our goal is to bound the generalization error of the target task . For simplicity , our results are presented under square loss functions . However , we show that our results can generalize to different loss functions in Appendix G. We also assume n1 = · · · = nT = nso for some positive integer nso . 2.1 Model complexity . As this paper considers general function classes , our results will be presented in terms of the complexity measures of classes of functions . We follow the previous literature ( Maurer et al. , 2016 ; Tripuraneni et al. , 2020 ) which uses Gaussian complexity . Note that we do not use the more common Rademacher complexity as the proofs require a decomposition theorem that only holds for Gaussian complexity . For a generic vector-valued function class Q containing functions q : Rd 7→ Rr and N data points , XN = ( x1 , . . . , xN ) T , the empirical Gaussian complexity is defined2 as ĜN ( Q ) = Eg [ sup q∈Q 1√ N N∑ i=1 gTi q ( xi ) ] , gi ∼ N ( 0 , Ir ) i.i.d . 2Note that the standard definition has a 1/N factor instead of 1/ √ N . We use the variant so that ĜN does not scale with N in most of the cases we consider and only reflects the complexity of the class . The corresponding population Gaussian complexity is defined as GN ( Q ) = EXN [ ĜN ( Q ) ] , where the expectation is taken over the distribution of XN .

Summary: The paper studies representation learning in a setting where the source tasks and the target task use different prediction function spaces beyond linear functions, and asks the following question: how many source tasks are required to achieve *diversity* (an important quantity related to the generalization ability from source tasks to the target task)?  The paper first shows that diversity over the source tasks' prediction function space implies diversity over the target task's prediction function space. As a result, even when the target task uses multi-layer neural networks as prediction functions, generalization is still guaranteed as long as the source tasks use linear prediction functions. However, when the source tasks use nonlinear prediction functions, there exist some fundamental limits for generalization: the paper proves some hardness results, including a novel lower bound on the number of required source tasks in terms of the eluder dimension. The paper further provides an upper bound in terms of the generalized rank. Furthermore, the paper gives some theoretical results indicating that simpler tasks generalize better. Such insights are validated by numerical experiments.