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: This paper presents a novel method for obtaining better representations for transfer learning. Specifically, instead of using a generic representation for various down-stream tasks, this paper proposed to create a family of expert models in pre-training, and selectively choose one expert to generate representation depending on the target transfer learning task. A simple yet effective k-nearest neighbor strategy is used for picking the best-fitting expert. This paper has extensive experiments including pre-training models on two large-scale image datasets, and evaluated on two transfer learning benchmarks (VTAB and datasets from DAT).