Auto-Transfer: Learning to Route Transferable Representations

1 INTRODUCTION . Deep learning models have become increasingly good at learning from large amounts of labeled data . However , it is often difficult and expensive to collect sufficient a amount of labeled data for training a deep neural network ( DNN ) . In such scenarios , transfer learning ( Pan & Yang , 2009 ) has emerged as one of the promising learning paradigms that have demonstrated impressive gains in several domains such as vision , natural language , speech , etc. , and tasks such as image classification ( Sun et al. , 2017 ; Mahajan et al. , 2018 ) , object detection ( Girshick , 2015 ; Ren et al. , 2015 ) , segmentation ( Long et al. , 2015 ; He et al. , 2017 ) , question answering ( Min et al. , 2017 ; Chung et al. , 2017 ) , and machine translation ( Zoph et al. , 2016 ; Wang et al. , 2018 ) . Transfer learning utilizes the knowledge from information-rich source tasks to learn a specific ( often information-poor ) target task . There are several ways to transfer knowledge from source task to target task ( Pan & Yang , 2009 ) , but the most widely used approach is fine-tuning ( Sharif Razavian et al. , 2014 ) where the target DNN being trained is initialized with the weights/representations of a source ( often large ) DNN ( e.g . ResNet ( He et al. , 2016 ) ) that has been pre-trained on a large dataset ( e.g . ImageNet ( Deng et al. , 2009 ) ) . In spite of its popularity , fine-tuning may not be ideal when the source and target tasks/networks are heterogeneous i.e . differing feature spaces or distributions ( Ryu et al. , 2020 ; Tsai et al. , 2020 ) . Additionally , the pretrained source network can get overwritten/forgotten which prevents its usage for multiple target tasks simultaneously . Among the myriad of other transfer techniques , the most popular approach involves matching the features of the output ( or gradient of the output ) of the target model to that of the source model ( Jang et al. , 2019 ; Li et al. , 2018 ; Zagoruyko & Komodakis , 2016 ) . In addition to the output features , a few methods attempt to match the features of intermediate states between the source and target models . Here , in this paper , we focus on the latter by guiding the target model with the intermediate source knowledge representations . Besides differing on what features to match , these methods also vary on how to match features , typically based on a regularization approach . While common approaches allow knowledge transfer between heterogeneous tasks/networks , it is also important to recognize that constraining the target DNN representations to be close to certain source DNN representations may be sub-optimal . For example , a source model , trained to classify cats vs dogs may be accessed at different levels to provide internal representations of tiger or wolf images to guide the target task in classifying tigers vs wolves . Since the source model is trained with a large number of parameters and labeled examples of cats and dogs , it will have learned several patterns that distinguish cat images from dog images . It is postulated that concepts or representations such as the shape of the tail , eyes , mouth , whiskers , fur , etc . are useful to differentiate them ( Neyshabur et al. , 2020 ) , and it is further possible to reuse these learned patterns to generalize to new ( related ) tasks by accessing representations at the appropriate level . This example raises three important questions related to knowledge transfer between the source-target models : 1 ) What knowledge to transfer ? 2 ) Where to transfer ? 3 ) How to transfer the source knowledge ? While the what and where have been considered in prior literature ( Jang et al. , 2019 ) , our work takes a novel and principled approach to the questions of what , where and how to transfer knowledge in the transfer learning paradigm . Specifically , and perhaps most importantly , we address the question of how to transfer knowledge , going beyond the standard matching techniques , and take the perspective that it might be best to let the target network decide what source knowledge is useful rather than overwriting one ’ s knowledge to match the source representations . Figure 1 illustrates our approach to knowledge transfer where the question of what and where is addressed by an adversarial multiarmed bandit ( routing function ) and the how is addressed by an aggregation operation detailed later . In building towards these goals , we make the following contributions : • We propose a transfer learning method that takes a novel and principled approach to automatically decide which source layers ( if any ) to receive knowledge from . To achieve this , we propose an adversarial multi-armed bandit ( AMAB ) to learn the parameters of our routing function . • We propose to meaningfully combine feature representations received from the source network with the target network-generated feature representations . Among various aggregation operations that are considered , AMAB also plays a role in selecting the best one . This is in contrast with existing methods that force the target representation to be similar to source representation . • Benefits of the proposed method are demonstrated on multiple datasets . Significant improvements are observed over seven existing benchmark transfer learning methods , particularly when the target dataset is small . For example , in our experiment on ImageNet-based transfer learning on the target Stanford 40 Actions dataset , our auto-transfer learning method achieved more than 15 % improvement in accuracy over the best competitor . 2 RELATED WORK . Transfer learning from a pretrained source model is a well-known approach to handle target tasks with a limited label setup . A key aspect of our work is that we seek to transfer knowledge between heterogeneous DNNs and tasks . Recent work focused on feature and network weight matching to address this problem where the target network is constrained to be near the source network weights and/or feature maps . Network matching based onL2−SP regularization penalizes the ` 2 distance of the pretrained source network weights and weights of the target networks to restrict the search space of the target model and thereby hinder the generalization ( Xuhong et al. , 2018 ) . Recent work ( Li et al. , 2018 ) has shown that it is better to regularize feature maps of the outer layers than the network weights and reweighting the important feature via attention . Furthermore , attention-based feature distillation and selection ( AFDS ) matches the features of the output of the convolutional layers between the source-target models and prunes the unimportant features for computational efficiency . Similar matching can also be applied to match the Jacobians ( change in output with respect to input rather than matching the output ) between source and target networks ( Srinivas & Fleuret , 2018 ) . Previous works ( Dhurandhar et al. , 2018 ; 2020 ) also suggested that rather than matching the output of a complex model , it could also be used to weight training examples of a smaller model . Learning without forgetting ( LwF ) ( Li & Hoiem , 2017 ) leverages the concept of distillation ( Hinton et al. , 2015 ) and takes it further by introducing the concept of stacking additional layers to the source network , retraining the new layers on the target task , and thus adapting to different source and target tasks . SpotTune ( Guo et al. , 2019 ) introduced an adaptive fine-tuning mechanism , where a policy network decides which parts of a network to freeze vs fine-tune . FitNet ( Romero et al. , 2014 ) introduced an alternative to fine-tuning , where the internal feature representations of teacher networks were used as a guide to training the student network by using ` 2 matching loss between the two feature maps . Attention Transfer ( AT ) ( Zagoruyko & Komodakis , 2016 ) used a similar approach to FitNet , except the matching loss was based on attention maps . The most relevant comparison to our work is that of Learning to Transfer ( L2T-ww ) ( Jang et al. , 2019 ) , which matches source and target feature maps but uses a meta-learning based approach to learn weights for useful pairs of source-target layers for feature transfer . Unlike L2T-ww , our method uses a very different principled approach to combine the feature maps in a meaningful way ( instead of feature matching ) and let the target network decide what source knowledge is useful rather than overwriting one ’ s knowledge to match the source representations . Finally , Ji et al . ( 2021 ) uses knowledge distillation based approach to transfer knowledge between source and target networks . 3 AUTO-TRANSFER METHOD . In this section , we describe our main algorithm for Auto-Transfer learning and explain in detail the adversarial bandit approach that dynamically chooses the best way to combine source and target representations in an online manner when the training of the target proceeds . What is the best way to train a target network such that it leverages pre-trained source representations speeding up training on the target task in terms of sample and time efficiency ? We propose a routing framework to answer this : At every target layer , we propose to route one of the source representations from different layers and combine it with a trainable operation ( e.g . a weighted addition ) such that the composite function can be trained together . We propose to use a bandit algorithm to make the routing/combination choices in an online manner , i.e . which source layer ’ s representation to route to a given target layer and how to combine , while the training of the target network proceeds . The bandit algorithm intervenes once every epoch of training to make choices using rewards from evaluation of the combined network on a hold out set , while the latest choice made by the bandit is used by the training algorithm to update the target network parameters on the target task . We empirically show the benefit of this approach with other baselines on standard benchmarks . We now describe this framework of source-target representation transfer along with the online algorithm . 3.1 ROUTING REPRESENTATIONS . For a given image x , let { f1S ( x ) , f2S ( x ) , · · · , fNS ( x ) } and { f1T ( x ) , f2T ( x ) , · · · , fMT ( x ) } be the intermediate feature representations for image x from the source and the target networks , respectively . Let us assume the networks have trainable parameters WS ∈ Rds and WT ∈ Rdt where ds and dt are the total number of trainable parameters of the networks . Clearly , the representations are a function of the trainable parameters of the respective networks . We assume that the source network is pre-trained . These representations could be the output of the convolutional or residual blocks of the source and target networks . Our Key Technique : For the i-th target representation f iT , our proposed method a ) maps i to one of the N intermediate source representations , f jS , or NULL ( zero valued ) representation ; b ) uses Tj , a trainable transformation of the representation f j S , to get f̃ j S , i.e . f̃ j S ( x ) = Tj ( f j S ( x ) ) ; and c ) combines transformed source f̃ jS and the target representations f i T using another trainable operation⊕ chosen from a set of operationsM . LetW ⊕ i , j be the set of trainable parameters associated with the operator chosen . We describe the various possible operations below . The target network uses the combined representation in place of the original i-th target representation : f̃ iT ( x ) = Tj ( f j S ( x ) ) ⊕ f iT ( x ) ( 1 ) In the above equation , the trainable parameters of the operator depend on the i and j ( that dependence is hidden for convenience in notation ) . The set of choices are discrete , that is , P = { [ N ] ∪NULL } ×M where [ N ] denotes set ofN source representations . Each choice has a set of trainable parameters Tj , W ⊕ i , j in addition to the trainable parametersWT of the target network .

Summary: The paper proposes an algorithm to transfer a "source" pre-trained deep model into a new one for a given target task. The idea is to wire in a certain way the intermediate representations of the source model into those of the "target" one. The algorithm considers a (small) discrete space of potential such ways to connect both models (location, source processing, and type of merging), and applies a bandit algorithm to sequentially spend training budget under one of the configurations. Experiments are provided that show big gains with respect to standard approaches (like finetuning the source model).