Exploiting a Zoo of Checkpoints for Unseen Tasks

1 Introduction . There are many model checkpoints published . For example , upon acceptance of this paper , the Huggingface repository2 includes over 18,000 checkpoints trained for dozens of tasks ( The tasks are tagged very coarsely . There will be many more if we tag in a fine-grained way ) . On the other hand , novel tasks can always arise , exciting tremendous interest in meta learning or learning to learn [ 1 , 2 , 3 , 4 ] . This paper does not propose a new learning-to-learn method . Rather , we are interested in the question , “ Can we use the checkpoint zoo to build something that can adapt to unseen tasks ? '' To address the question , we have to understand how tasks are relevant to each other . A paraphrase of task relevance is task transferability [ 5 ] . Intuitively , transferring to a new task is easier if we warm start from a very relevant one . Efforts in this space originates from computer vision applications . Earlier work like Taskonomy [ 5 ] explicitly uses performance metrics ( e.g. , test accuracy ) of transfer learning . Since then , multiple works have developed task descriptors to avoid running training and evaluation as Taskonomy . A few prominent examples are Task2vec [ 6 ] , Attribution Map ( A-Map ) [ 7 ] , and LEEP [ 8 ] . Very recently , Task2vec has also been extended to computational linguistics [ 9 ] . Each of these methods has its own pros and cons . However , these studies mainly aim at interpreting the performance of transfer learning [ 9 , 10 ] . They mainly attack the question , “ For a target task , which source task is the most relevant ? '' . It is different from the question we are interested in . Understanding the relation among tasks is crucial for meta learning and transfer learning . It is known that dissimilar tasks can impose significant challenges for these learning frameworks [ 11 , 12 ] . Aware 1All results can be reproduced using code at https : //github.com/baidu-research/task_space 2https : //huggingface.co/models 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of this fact , several recent works [ 13 , 14 ] develop generalization error bounds for learning new tasks . In particular , they address the question , “ how many tasks has to be seen , so that a meta-trained model can generalize to new tasks ? '' . These works often assume that there is a shared function , mapping data to representations for all tasks . Then they propose their measure of task diversity , which is further employed to bound the generalization error for a new task . These bounds suggest that the seen tasks must be sufficiently diversified to “ cover '' all components that may be used for a future task . However , it is not obvious how to compute these diversity measures in practice . Back to our problem , a zoo of checkpoints solve a collection of seen tasks . We are interested in combining/fusing a subset of them . Motivated by the theoretical works [ 13 , 14 ] , we want the selected checkpoints ( essentially the tasks they solve ) to be diversified . In a perhaps tangential problem [ 15 ] , the authors consider where to place sensors in a room such that they are the most informative about the undeployed locations . They show that the placement should be diversified and well represent the landscape . To this end , we could draw an analogy . Consider the task space as a room , and checkpoints as sensors . The unseen tasks are locations where no sensor is placed but we wish to infer their information . In later sections , we shall see how this analogy helps us design a strategy for selecting checkpoints . A relevant concept to our work is model ensemble , which has a long history [ 16 ] , and still actively being studied recently [ 17 ] . However , model ensemble usually assumes a shared output space for all component models , therefore no notion of multiple ( seen or unseen ) tasks . Another relevant topic is feature selection [ 18 ] . Indeed , if we consider each checkpoint a feature extractor , picking checkpoints amount to identifying useful features . However , feature selection methods usually aim at finding the most relevant features for a target variable . In other words , they only target at a single seen task . 2 Motivating the Problem . Consider a zoo of n checkpoints , where the i-th checkpoint is trained for the i-th task . In this paper , we define the i-th task as a joint distribution pi ( x , y ) to be fit to , where x ∈ X is input and y ∈ Yi is target output . The i-th checkpoint is trained on a dataset Di drawn from pi . Di is often proprietary . Or due to non-trivial preprocessing , it is very difficult to obtain the data exactly as it was . The tasks are assumed to share the same input space , but differ in their target spaces . This is a natural assumption . Consider a zoo of convolutional networks that handle various computer vision tasks [ 5 ] , or a collection of huggingface transformers for multiple linguistic tasks . The X is a space of images for the former case , and text for the later . Suppose there is a universe of ( possibly infinitely many ) tasks , drawn from a “ hyper-distribution '' P . That is ( p1 , . . . , pn , pn+1 , . . . ) ∼ P . Modeling a family of tasks as “ distribution over distributions '' [ 1 , 19 ] has a long history in learning-to-learn literature . The n checkpoints in the zoo have distilled information from p1 , . . . , pn , and we are interested in building a “ stronger '' model from them that can easily generalize to unseen tasks ( pn+1 , . . . ) . It should be emphasized that these new tasks are not revealed to us before we build the “ stronger '' model . To be more specific , let the i-th checkpoint be in the ( commonly seen ) form , hi ◦ fi : X 7→ Yi , where fi is a feature extractor . hi is a task-specific “ head '' for the i-th task , and ◦ denotes function composition . hi is often a “ simple '' function , e.g. , softmax classifier . Note that hi may not necessarily be available . This is true for many checkpoints in huggingface repository . We are interested in selecting a subset S of checkpoints from the zoo before seeing any new task pn+1 , . . . , and building a “ stronger '' feature extractor by concatenating the selected fi ’ s , f , ⊗i∈Sfi . The f should be quickly adaptable to new tasks pn+1 , . . . . In other words , given some small training data for a new task , we want a task-specific “ head '' on top of the f to be “ easily '' learned and work well on test data . Of course , this approach should have some budget constraints . Indeed , when n is big , concatenating all the fi ’ s may be costly in computation , and result in overfitting . At this point , one may propose to use feature selection methods like min-Redundancy-Max-Relevance ( mRMR ) [ 18 ] . That is , to pick the fi ’ s that are the most informative for a new task , meanwhile least redundant mutually . However , in our case , no new task is revealed to us before selecting . Thus we can not simply apply mRMR . Having ruled out feature selection methods , we start to re-think the problem by modeling the space of tasks . To this end , we may make an analogy to optimal sensor placement [ 15 ] . In that application , one is supposed to pick the best locations to place sensors , so that their readouts are the most informative about the unsensed locations . In information theoretic language , this would be max I ( sensed location ; unsensed location ) , where I ( · , · ) is mutual information . In our problem , each checkpoint handles its specific task . It can be considered as a “ sensor '' placed at a task . Then , the checkpoint zoo is a “ pilot deployment '' that places many “ sensors '' on grids of the task space , which can reveal how the tasks are related . Selecting checkpoints is essentially picking a few seen tasks that may be the most informative about the unseen ones . In the following sections , we concretize the above idea by modeling and exploiting the task space . 3 Modeling the Task Space . Theoretical works [ 1 , 19 , 13 ] often do not specify a form for the hyper-distribution P . In this work , we assume P a Gaussian process . So a next step is to define the covariance function , κ ( pi , pj ) , for P . In other words , we want to come up with a similarity measure between the two distributions pi and pj . Moreover , by definition , the κ has to be Positive Semi-Definite ( PSD ) . In the following , we first rule out several options that are less applicable . Then we propose our approach . 3.1 A Taxonomy of Task Similarities . If we had access to data samples Di , and Dj , a natural idea would be comparing pi and pj directly . The difficulty is that the output spaces Yi and Yj are not directly comparable . For example , Yi may be the input images ’ class labels , whereas Yj may be the foreground to be segmented . Therefore , straightforward measures like KL divergence can not be applied . To address this issue , Geometric Dataset Distance ( GDD ) [ 20 ] proposes to estimate Wasserstein distance between pi and pj , by carefully designing the “ transport cost '' between Yi and Yj using Di and Dj . However , Di and Dj are often proprietary . For example , the data to train huggingface transformers are often preprocessed with various tools . It is very hard ( if not impossible ) to recover the exact same Di . Other approaches treat κ ( pi , pj ) as the transferability from the i-th task to the j-th . Taskonomy [ 5 ] empirically computes that by running training and testing for this i-to-j transfer learning problem . Doing that for all ( i , j ) pairs is very costly . Besides , non-trivial normalizations have to be applied as performance metrics can differ for different tasks . To circumvent this brute-force process , people have come up with vector representations of tasks , reducing tasks similarities to distances in this vector space . For example , Task2vec [ 6 ] argues that Fisher information matrix ( essentially gradient w.r.t . parameters in the checkpoint ) is an indicative representation for the corresponding task . However , it assumes that all checkpoints have the same model configuration , which is a very strong assumption and need not be satisfied in practice . Another work , A-Map [ 7 ] is based on the intuition that similar tasks should have similar saliency map , which is the gradient of ( pooled ) feature w.r.t . input . Therefore , it only requires some unlabelled probing data , and model architectures can be different . However , the back-propagation can be more costly . Further more , to compare two saliency maps , their inputs have to be the same . This constraint can be violated for NLP applications , as the same input text is often mapped to two different sequences of embedding vectors using two checkpoints . Recently , NCE [ 21 ] and LEEP [ 8 ] measure task transferability by estimating an approximation of conditional entropy , H ( yj |yi ) = H ( yj ) − I ( yi ; yj ) . Ignoring H ( yj ) and focusing on the mutual information term , we can understand these two methods as measuring the dependency between two target spaces Yi and Yj . However , to estimate H ( yj |yi ) , data labels must be available . For clarity , we list the discussed methods in Tab . 1 . Also , the last row of Tab . 1 highlights some properties a desired method should have . We can see various reasons why these existing methods are not applicable . Nevertheless , LEEP [ 8 ] provides us with a principled way of understanding the relationship between two tasks , i.e. , measuring dependency between their target spaces . We next derive another approximation to this dependency , which is a desired method .

Summary:  The paper considers the question of how to select an ensemble of existing pretrained models to be used as feature extractors, when the task for which they are to be used is not known a priori. It seems to aim less to generalize to a given fixed target task distribution, but rather at generalizing over an expectation over target task distributions.  To achieve this, the authors propose essentially two steps:  Firstly, considering the feature representation of each model from the zoo, and computing a within-model similarity matrix between the features extracted for one fixed model over a set of samples. Then to use the cosine angle between similarity matrices of different models as a similarity between models.  Secondly, to select a fixed number of models (corresponding to a fixed budget) based on mutual information estimates between the set of models and the remaining models. This is intuitively appealing.