Representation Consolidation from Multiple Expert Teachers

1 INTRODUCTION . A promising approach to scale transfer learning to diverse downstream vision tasks is to maintain a library of diverse experts pre-trained on different tasks . ( Deshpande et al. , 2021 ) When presented with a novel downstream task , one can select an appropriate expert and fine-tune the representation with small amounts of task-specific data . This strategy has many practical benefits : fine-tuning a task-relevant pre-trained representation is fast , there is no need to store or revisit expert pre-training data . How is such a library of expert representations populated and maintained ? Previous work ( e.g. , Puigcerver et al . ( 2020 ) ; Achille et al . ( 2019 ) ; Deshpande et al . ( 2021 ) ) has assumed a static collection of experts created by training on domain-specific datasets or domain-specific subsets of general dataset . Instead , we would like to automatically enrich the diversity of the expert library by accumulating knowledge from transferred downstream tasks , so that the overall system performance continually increases over time ( life-long meta-learning ) . One approach to growing the library would be to add fine-tuned downstream task-specific models back into the library as a candidate expert for transfer on future tasks . Such a naı̈ve approach clearly does not scale and , at a minimum , requires developing techniques for selecting experts that is sub-linear in the size of the library ( Puigcerver et al. , 2020 ; Achille et al. , 2019 ; Deshpande et al. , 2021 ) . More importantly , as our experiments show in Sec . 4 , task-specific models fine-tuned on small amounts of data do not provide transferable representations . Task-specific models tend to overspecialize and degrade under-utilized features in their representations , and thus under-perform on new tasks compared to generic pre-trained models . To address this , we introduce representation consolidation where the goal is to consolidate knowledge from multiple task-specific teacher models into a single expert student representation which transfers to downstream tasks better than any of the individual teacher representations . There are three factors that make the problem of representation consolidation unique relative to the existing literature : ( 1 ) we assume a representation is to be consolidated from a collection of multiple taskspecific models , ( 2 ) we avoid the need to revisit task-specific training data of these models , and ( 3 ) we focus on improving downstream transfer instead of simply replicating performance on the upstream teacher ’ s task ( s ) . To the best of our knowledge , no prior work has analyzed the downstream transfer aspect of distilled domain-expert representations . To carry out representation consolidation , we utilize multi-teacher multi-task model distillation ( see Fig . 1 ) . We use a simple yet effective method to jointly distill one or several task-specific teacher with a generalist one . Each teacher may operate on a different set of classes , and a multi-head student is trained to emulate all teachers . Previous work on knowledge distillation has focused on the student model ’ s performance on the teacher ’ s task . Instead , we evaluate how well the student representation generalizes to new downstream tasks ( whether related or unrelated to the teachers ’ tasks ) . In this new setting we demonstrate several surprising results : 2 REPRESENTATION CONSOLIDATION . Problem statement . We start with a collection of one or more task-specific image classification models { Mit } Ni=1 as teachers , trained on corresponding datasets { Dit } belonging to some domain ( satellite images , images of flowers , etc. ) . We assume models consist of a feature extractor or backbone φit ( · ) , composed with a classifier head hit ( · ) so thatMit = hit ( φit ( · ) ) . We consolidate the knowledge of these task-specific teachers into a single student representation φs ( · ) using a proxy dataset Dproxy ( e.g. , ImageNet ) , and then evaluate φs ( · ) by training a SVM on top of it ( or fine-tuning it ) on a given downstream Djd chosen from some set { D j d } . Our goal is that the resulting downstream model hjd ( φ j d ( · ) ) achieves good performance , where φ j d denotes the student φs as-is or potentially after tuning on Djd . Fig . 1 highlights how this differs from standard distillation in which the student model hs ( φs ( · ) ) is simply evaluated on the same task its teachers were once trained to perform . Forgetting and representation collapse during distillation . Though our problem statement is novel , can existing approaches solve our problem statement ? We observe ( Sec . 4 ) that neither the task-specific teachers nor students distilled with standard or state-of-the-art methods provide sufficiently transferable representations . They under-perform general pre-training representations ( e.g . with ImageNet ) when evaluated on downstream tasks from the teachers ’ domain , and drastically under-perform on tasks outside the teachers domain . Surprisingly , this holds true even though we use ImageNet-pretrained weights to initialize the teacher and student networks . Thus , we argue that training task-specific teachers or distilling from them suffers from catastrophic forgetting of general knowledge that is crucial for transfer learning . Intuitively , transfer performance depends on how distinguishable different classes are represented in the penultimate layer feature space . In standard distillation , the student only learns from taskspecific teachers trained on smaller datasets and is only required to discriminate the classes in those datasets . This is well suited for the traditional distillation problem that evaluates the student on these same tasks . But for representation consolidation , distinguishing unknown downstream classes requires preserving general features that may not be relevant to the teachers ’ specific task . Our strategy is thus to ensure that the student maintains general features as it learns task-specific ones . Method formulation . A key challenge to our approach is avoiding replay of the images used to train the task-specific teachers . To overcome this , we propose using a proxy dataset and a generalist model that transfers well ( e.g . one pre-trained on ImageNet ) . This model is referred to as a generalist as it is not task specific and contains a wide variety of classes , and use it to help avoid forgetting . More specifically , we use multi-head multi-task distillation . We put N + 1 heads , h0s , . . . , h N s , on top of the student backbone φs ( · ) . In addition to the task-specific teachers hit ( φit ( · ) ) , i ∈ { 1 , . . . , N } used in traditional distillation , we also include the generalist teacher trained on ImageNet ( denoted as h0t ( φ 0 t ( · ) ) ) . The rationale behind this construction is that we learn task knowledge from the task-specific model on the proxy dataset , but also force the student to retain knowledge from the generalist teacher to prevent over-specialization . We learn the student by optimizing the loss : L = ∑ x∈Dproxy ( λ0 Ldistill ( h0t ( φ 0 t ( x ) ) , h 0 s ( φs ( x ) ) ) ︸ ︷︷ ︸ generalist teacher term + N∑ i=1 λi Ldistill ( hit ( φ i t ( x ) ) , h i s ( φs ( x ) ) ) ︸ ︷︷ ︸ traditional distillation ) ( 1 ) with a distillation loss Ldistill . Our key difference from traditional distillation is the generalist teacher term . Setting λ0 = 0 while keeping all other conditions ( proxy dataset , network initialization , etc . ) yields a standard multi-teacher distillation baseline . Note that this is also distinct from ensemble knowledge distillation , where one averages the homogeneous output of every model and distills the mean onto a single student head , since our teacher outputs are heterogeneous ( different tasks ) . Our method is agnostic to the underlying distillation method . We can use the standard knowledge distillation ( KD ) loss ( Hinton et al. , 2015 ) which is cross-entropy with temperature T = 2 , i.e . Ldistill ( pt , ps ) = − ∑C c=1 p ( c ) t log ( p ( c ) s ) where c indexes the C classes , and pt = softmax ( hit ( φ i t ( x ) ) /T ) , ps = softmax ( his ( φs ( x ) ) /T ) . Alternatively , we can use a state-of-theart distillation method , e.g. , CRD ( Tian et al. , 2020 ) . CRD uses the KD loss and adds a contrastive embedding loss between hiCRD , t ( φ i t ( x ) ) and h i CRD , s ( φs ( x ) ) with a pair of trainable linear embedding layers hiCRD , t and h i CRD , s that are used only in training and discarded afterwards . The teacher φ i t ’ s are trained on completely different tasks , so for CRD we use different embedding layers for each teacher-student-head pair . We initialize the student backbone φs and its 0-th head h0s using the generalist model ’ s weights φ0t , h 0 t , whereas other heads are randomly initialized . Since it is important to maintain pre-trained model ’ s representational power , we simply set the loss weights λ0 = 1 and λi = 1N for 1 ≤ i ≤ N . This forces the learned representation φs to balance between learning general and taskspecific features , which benefits future downstream transfer . After training the student we evaluate the resulting representation φs ( · ) on multiple downstream tasks { Did } . For each task j we can either fine-tune the whole model hjd ( φ j d ( · ) ) or keep the student representation fixed and only learn the classifier head hjd ( φs ( · ) ) which is often referred to as a linear probe . 3 EXPERIMENTAL SETUP . Datasets and downstream tasks . We utilize datasets from a variety of domains to generate teachers and downstream tasks : Cars196 ( Krause et al. , 2013 ) , Resisc45 ( Cheng et al. , 2017 ) ( remote sensing images ) , iFood ( Kaur et al. , 2019 ) and Food101 ( Bossard et al. , 2014 ) , iFashion ( Guo et al. , 2019 ) , DTD ( Cimpoi et al. , 2014 ) ( describable textures ) , iNaturalist ( Horn et al. , 2018 ) ( species classification , 2019 challenge version ) , CUB Birds ( Wah et al. , 2011 ) , Flowers ( Nilsback & Zisserman , 2008 ) , Caltech256 ( Griffin et al. , 2007 ) , and Aircrafts ( Maji et al. , 2013 ) . Among these , iFood and Food101 are the same domain , and Birds and Flowers are subdomains of iNaturalist . We checked for near-duplicates between these datasets using perceptual hash ( Buchner , 2021 ) , and found negligible duplication : 1 out of 134k iNaturalist ( 50 % classes we used ) is a duplicate of CUB , and 8 out of 130k/100k images are duplicates between iFood and Food101 . To evaluate if a consolidated student representation has learned features relevant to a specific domain , we require downstream tasks that are related ( in the same domain ) to that of each task-specific teacher . Except for Food101 and iFashion , we split each dataset at random into two disjoint sets which each contain only 50 % of the classes . We take the first half of the dataset , named “ t-split ” , and use as Dit , to train a task-specific teacher . We use the second half of each dataset , named “ dsplit ” , as one of the downstream tasks Djd . We primarily evaluate on few-shot downstream transfer , where we randomly sample 5 training images from each available class in Djd , but always use the entire test set for evaluation . We use all samples for those classes with < 5 images . For iFashion , which is a multi-task multi-label dataset , we use all images but with a subset of labels ( e.g. , those related to clothing category ) as Dit for teacher training . We then use all images with a disjoint set of labels ( e.g. , those related to sleeve style ) as a downstream task Djd for evaluating transfer . For iFashion ’ s few-shot scenario , we randomly subsample 1000 images for downstream training , and evaluate on all test data . We do not train any teacher on Food101 so the complete set of classes are used as a downstream task . During distillation and consolidation , we use ImageNet ( Russakovsky et al. , 2015 ) ( ILSVRC12 ) as the proxy dataset Dproxy for most experiments . However , we also show results with Places365standard ( Zhou et al. , 2018 ) or iNaturalist ( Horn et al. , 2018 ) as the proxy . Compared methods and Criteria . As a reminder , our goal is not to duplicate upstream teacher predictions , but rather to improve transfer performance of distilled representations . Therefore , we evaluate each student ’ s φs using its performance when transferred to various downstream datasets { Djd } and compare multiple methods for initializing the downstream representation φ j d : ( 1 ) ImageNet-pretrained φ0t , a strong baseline for transfer learning . ( 2 ) ImageNet-pretrained φ 0 t fine-tuned on the soft-labels produced by the ImageNet-pretrained model with batchnorms in test mode . This baseline aims to isolate the effect of soft-labels / self-distillation on model performance . ( 3 ) The task-specific teacher φ1t ( or one of the teachers when N > 1 ) without further distillation . ( 4 ) Distillation with N teachers using either standard knowledge distillation ( KD ) or contrastive representation distillation ( CRD ) . This distills only task-specific teacher ( s ) on Dproxy , without the ImageNet teacher h0t ( φ 0 t ( · ) ) . ( 5 ) Our consolidated representation φs which includes h0t ( φ0t ( · ) ) as a teacher , using either KD or CRD as the underlying distillation loss . ( 6 ) A multi-task learning oracle ( MTL ) that jointly trains on ImageNet and teacher data Dt rather than on unlabeled proxy data . For comparison fairness , ( 2-6 ) are all ImageNet pre-trained ( i.e . initialized with ( 1 ) ) before further training . Both ( 4 ) and ( 5 ) are using the same proxy dataset ( i.e . ImageNet for most experiments ) . We use the transfer accuracy on downstream tasks Djd to measure each representation ’ s power . We primarily use the few-shot linear probe ( train a 5-shot linear SVM as hjd over fixed φ ) on D j d ’ s training set ( single training run for full dataset , few-shot performance averaged over 50 random subsampling trials ) . We also verify our results hold when fine-tuning the student representation hjd ( φd ( · ) ) on few-shot or full dataset splits . Implementation details . We will release our code to reproduce this paper upon acceptance . For experimental fairness , we use the same network ( ResNet50 ( He et al. , 2016 ) ) / hyperparameters / ImageNet initialization for the baselines and ours . For more details , please see appendix A .

Summary: - The authors propose a method of learning a consolidated image feature representation from a collection of related task-specific teachers that transfer well on novel recognition tasks.  - To achieve it, the authors utilize multi-teacher multi-task model distillation framework that jointly distills one or several task-specific teachers with a generalist one (trained with Imagenet dataset).   - Each teacher is set to operate on a different set of classes, and a multi-head student is trained to emulate all teachers.  - Experimental results show that the proposed method doesn’t need to revisit original training data of each teacher, but gains better performance.