Scalable Transfer Learning with Expert Models

1 INTRODUCTION . Deep learning has been successful on many computer vision tasks . Unfortunately , this success often requires a large amount of per-task data and compute . To scale deep learning to new vision tasks , practitioners often turn to transfer learning . Transfer learning involves re-using models trained on a large source task , and tuning them on the target task . This can improve both convergence rates ( Ben-David et al. , 2007 ; 2010 ; Blitzer et al. , 2008 ; Du et al. , 2017 ; Kuzborskij & Orabona , 2013 ; Mansour et al. , 2009 ) and empirical performance ( Dai et al. , 2007 ; Donahue et al. , 2014 ; Oquab et al. , 2014 ; Tan et al. , 2018 ) . Transfer learning reduces per-task data or compute requirements , given a large one-off pre-training cost . In practice , this one-off down payment may not be made by the practitioner , since pre-trained networks are made available through platforms like PyTorch and TensorFlow Hub1 . For instance , ImageNet pre-training is popular since it is freely available and works well for many tasks ( Donahue et al. , 2014 ; Oquab et al. , 2014 ; Sharif Razavian et al. , 2014 ) . In contrast to generic homogeneous models ( e.g . most pre-trained ImageNet networks ) , Mixture of Experts ( MoE ) include multiple heterogeneous sub-models ( “ experts ” ) that specialize to sub-problems of the full task . MoEs have been studied for decades ( Eigen et al. , 2013 ; Jacobs & Jordan , 1993 ) , and have also been successful in deep learning ( Shazeer et al. , 2017 ) . Yet , the application of experts for deep transfer learning has been less explored . We study visual transfer with experts , and present a simple , scalable , yet effective strategy . Transfer of specialist models has been studied before . However , they either require expensive retraining on the source dataset for every target task ( Ngiam et al. , 2018 ; Yan et al. , 2020 ) , or operate at a small scale where all experts can be applied simultaneously ( Dvornik et al. , 2020 ) . Further , most of them are tested only on a limited suite of natural single-object classification tasks . We lift these ∗Equal contribution . Order decided by a coin toss . †Work done while interning at Google Research . 1https : //pytorch.org/hub/ and https : //tfhub.dev/ , respectively . constraints , and present a practical approach that scales to hundreds of large experts , while requiring relatively little compute per target task . Our strategy consists of four stages ( fig . 1 ) . ( 1 ) Unconditional pre-training . A single baseline model is trained on the entire upstream data . ( 2 ) Experts training . Multiple experts are pre-trained by exploiting the label hierarchy present in many large-scale image datasets , such as ImageNet and JFT . In addition to entire expert networks , we explore residual adapters that allow all of the expertise to be packed into a single model that can be loaded into memory . These two stages may be expensive , but are done only once . ( 3 ) Expert selection . Applying all experts to each task does not scale well ; some sort of sparsification is required . We focus on inexpensive model selection that can be applied to hundreds or thousands of experts . ( 4 ) Downstream fine-tuning . We take the output of the model selection phase and tune it on the target task . Importantly , this phase does not require revisiting the source dataset , which may be unavailable or expensive to train on . We show that this approach yields remarkably strong performance on many diverse tasks . We evaluate not only on classic vision tasks , but also on the diverse VTAB benchmark of 19 tasks ( Zhai et al. , 2019 ) . Our contributions can be summarized as follows . • We propose a transfer learning algorithm with a large number of experts based on per-task routing via nearest neighbors selection . Once we have amortized the pre-training cost , this algorithm requires little compute per target task , achieving an speed-up of 500×–1000× compared to competing strategies . Also , it can be easily replicated with any large upstream multilabel dataset . • We achieve a mean accuracy improvement of 3.6 % over the state-of-the-art performance on 19 VTAB datasets using ResNet50 networks . Our algorithm offers improvements on every group of tasks : natural , specialized , and structured . Figure 2 summarizes these results . • We explore using sub-networks as experts via residual adapters , allowing all experts to be packed into a single model . Surprisingly these perform almost as well as their full-network counterparts . 2 RELATED WORK . Transfer Learning . Tasks with little training data can benefit from other larger datasets , often from a similar domain . Transfer learning concerns the link between the source and target dataset ( Pan & Yang , 2009 ; Weiss et al. , 2016 ; Tan et al. , 2018 ; Wang , 2018 ) . One family of methods creates a single training dataset , where source instances are re-weighted according to their relevance ( Dai et al. , 2007 ; Pardoe & Stone , 2010 ; Wan et al. , 2011 ; Xu et al. , 2017 ) . A popular method consists of fine-tuning a model that was pre-trained on the source data ( Donahue et al. , 2014 ; Oquab et al. , 2014 ; Sharif Razavian et al. , 2014 ) . Some transfer learning algorithms condition the initial source model on the target dataset itself ( Ngiam et al. , 2018 ; Xie et al. , 2019 ; Yalniz et al. , 2019 ) , while others ( like ours ) are agnostic about the downstream task when the initial model is trained on the source data ( Kolesnikov et al. , 2019 ) . We offer an in-depth comparison with ( Ngiam et al. , 2018 ) in section 6.6 . In the context of few-shot learning , where out-of-the-box fine-tuning may not work , generic representations are sometimes frozen , and simple feature selection ( Dvornik et al. , 2020 ) or model training ( Chen et al. , 2019 ) techniques are applied on top . Instead of relying on fixed universal representations , ( Rebuffi et al. , 2017 ; 2018 ) use small additional modules , or adapters , that incorporate knowledge from several visual domains . Our work also explores this idea . Multi-task Learning . MTL tries to leverage the common aspects of several learning tasks ( Caruana , 1997 ) . A prominent approach uses explicit parameter sharing ; for instance , by means of common low-level layers leading to different heads . Among others , this has been successfully applied to vision ( Zhang et al. , 2014 ) , language ( Liu et al. , 2015 ) , and reinforcement learning ( Fedus et al. , 2019 ) tasks . In addition , a variety of ways to combine task-specific representations have arisen , such as cross-stitch networks ( Misra et al. , 2016 ) , or lateral connections ( Rusu et al. , 2016 ) . A different family of methods impose joint constraints on the –possibly different– models corresponding to each task . We can combine the learning problems via regularization and shared sparsity patterns ( Argyriou et al. , 2007 ; Lounici et al. , 2009 ) , or by imposing some prior knowledge regarding the task structure ( Evgeniou et al. , 2005 ; Jacob et al. , 2009 ; Kim et al. , 2012 ) . 3 THE TRANSFER LEARNING FRAMEWORK . In this section , we describe our transfer learning setup of interest . The high-level goal is to train strong models for arbitrary downstream tasks , possibly under severe data and compute limitations . To do so efficiently , one can offload computation to a previous upstream phase which is executed a priori , without knowing the downstream tasks in advance . Accordingly , the upstream model should not depend on any specific target data . We are mostly interested in the low data regime where downstream tasks contain few datapoints . These restrictions have a practical motivation : we would like to build and deploy universal representations that are easily transferred to a wide range of downstream settings . Any transfer algorithm must implement the following three stages . Upstream Training . Given the upstream data DU , the algorithm first outputs a source model M. The goal is to provide useful initial representations for various new tasks . This stage could actually produce a family of models { Me } rather than a single one . These models might not be disjoint , and could share parameters . The upstream learning problems are auxiliary ; accordingly , DU could include a diverse set of classification , regression , or even synthetic learning instances . Model Selection . When a new downstream task is given , a selection algorithm is applied to choose the upstream model ( s ) to transfer , possibly depending on the downstream data . This phase should be computationally cheap ; thus , the upstream data is no longer available . Sometimes , there is no choice to make ( say , with a single ImageNet representation ) . Alternatively , in models with a complex structure , one may choose which parts , routes , or modules to keep in a data-dependent fashion . Downstream Training . The final stage fine-tunes the selected model using the downstream data , either fully or partially . For neural nets , a new head is added as the output classes are task-specific . Our overall algorithm is depicted in fig . 1 . We give details about each step in the following sections . 4 UPSTREAM TRAINING . 4.1 EXPERT ARCHITECTURES . Our experts should provide feature extractions that are a good starting point to learn future tasks related to the expert ’ s upstream training data . We explore two different model architectures to train such experts . As an obvious choice , we first look at ResNets ( He et al. , 2016b ) . These are powerful models ; however , storing and deploying many of them can be challenging . As an alternative , we also develop more compact adapter modules that can all be assembled in a single architecture . Also , their individual size can be easily customized to meet memory and computational constraints , which makes them an ideal candidate for combining multiple experts in a single model , when needed . We informally refer to these as full and adapter modules ( or experts ) , respectively . Full ResNet Modules . As a base architecture for full experts we use ResNets . In particular , all of our experiments focus on the ResNet50-v2 architecture ( R50 ) ( He et al. , 2016a ) , which sequentially stacks a root block and 4 blocks with ( 3 , 4 , 6 , 3 ) residual units . The initial step in every experiment consists of training a baseline model B on the whole upstream data ( see stage 1 in fig . 1 ) . This baseline is subsequently fine-tuned by both full and adapter experts , but in different ways . A full expert trained on a slice of data is simply the baseline B fine-tuned on that data . The head will later be discarded for transfer . This approach requires as many R50s as there are experts . Adapter Modules . Residual adapters were proposed to adapt a neural network to a particular downstream task without needing to fine-tune the entire network ( Rebuffi et al. , 2017 ) . Instead , we use them to adapt the baseline architecture to slices of the upstream data . Originally , they were 1× 1 convolutions placed after each 3× 3 convolution , with a residual connection . Instead we place them before each of the R50 ’ s blocks . Finally , our adapters have a bottleneck and are non-linear , as in ( Houlsby et al. , 2019 ) . We insert several in parallel into the backbone B . When creating an expert , only the adapters are tuned and the backbone weights are frozen . Figure 3a depicts the ResNet architecture with multiple expert adapters ( a ( i ) 1 , . . . , a ( i ) n ) . Let Fi be the function implemented by the i-th block of the backbone network . We adapt its input by computing the output as xi : = Fi ( xi−1 + a ( i ) e ( xi−1 ) ) , where e = R ( x ) is the identity of the selected expert , given by some routing function R , and x is the original input . During upstream training , the function R may also use the labels in addition to the image , as we discuss in section 4.3 . Figure 3b shows the adapter ’ s bottleneck architecture . An adapter sequentially applies componentsA1 and A2 . Each component performs a group normalization ( N ) ( Wu & He , 2018 ) , a ReLU activation ( A ) ( Glorot et al. , 2011 ) , and a convolution ( C ) ( LeCun et al. , 1989 ) , in that order . Due to the skip connection , the output dimension of A2 ◦A1 must match that of its input , c. However , we can change the output channels k of A1 , in order to limit the amount of parameters . Thus , we set k = c2 so that the number of parameters equals that of a linear adapter . Each adapter only increases the parameter count of the R50 backbone by 6 % . We briefly explored placing these adapters in other locations , or using other variations ( Rebuffi et al. , 2018 ) , but we did not observe any significant improvement . It is important to emphasize that the names Full and Adapters refer to the parameters that are specialized to a given subset of the upstream data before any downstream fine-tuning . In both cases , we fine-tune the entire network , after the model selection phase , to the downstream task .

Summary: The authors address transfer learning scenarios. In particular, the authors resort to training to a diverse set of experts and "cheap" performance proxies to select, for a given task, the relevant expert. This "per-task routing" is conducted via a nearest neighbor classifier based on a reduced representation for each expert. Two variants are considered: (1) full ResNet50 models are used (one for each expert) and (2) "compact adapter modules", which depict expert layers between ResNet block (all experts are learnt simultaneously with the same model backbone).