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 paper proposes a task-agnostic Knowledge Distillation method that includes a generalist model as an additional teacher, to limit student forgetting and representation collapse. The authors call this method representation consolidation and claim it can benefit both related and unrelated downstream tasks. The authors experiment with ImageNet and iFood image classification tasks, the accuracy of the proposed method is generally higher than using Knowledge Distillation only.