Exploiting a Zoo of Checkpoints for Unseen Tasks
Year: 2022

Abstract:
There are so many models in the literature that it is difficult for practitioners to decide which combinations are likely to be effective for a new task. This paper attempts to address this question by capturing relationships among checkpoints published on the web. We model the space of tasks as a Gaussian process. The covariance can be estimated from checkpoints and unlabeled probing data. With the Gaussian process, we can identify representative checkpoints by a maximum mutual information criterion. This objective is submodular. A greedy method identifies representatives that are likely to “cover” the task space. These representatives generalize to new tasks with superior performance. Empirical evidence is provided for applications from both computational linguistics as well as computer vision1.

heading: 1 Introduction
text: 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.

heading: 2 Motivating the Problem
text: 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 cannot 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.

heading: 3 Modeling the Task Space
text: 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.

heading: 3.1 A Taxonomy of Task Similarities
text: 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 cannot 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.

heading: 3.2 Proposed Approach
text: Following the setup in NCE [21] and LEEP [8], we also assume m input samples X , [x1, . . . ,xm], with target outputs Yi , [y1i , . . . ,y m i ] in the i-th task, and Yj , [y 1 j , . . . ,y m j ] in the j-th task. We wish to measure the dependency between Yi and Yj . One approach is Kernel Alignment
(KA) [22, 23]. In the case of linear kernels, KA can be calculated as
KA(Yi,Yj) = 〈Y>i Yi,Y>j Yj〉
〈Y>i Yi,Y>i Yi〉 · 〈Y>j Yj ,Y>j Yj〉 .
To give some intuitions of the above, let columns of Yi and Yj be one-hot. Then Y>i Yi (same for Y>j Yj) is a binary matrix who assigns 1 for two inputs with the same label, and 0 otherwise. KA measures the dependency between the two tasks by cosine similarity of their assignment patterns.
However, we do not have access to the data labels. Moreover, sometimes the task-specific heads are unavailable, so even a prediction of the labels (denoted as Ŷi and Ŷj) cannot be obtained. We tackle this difficulty by resorting to unsupervised “probing" data [x1, . . . ,xm]. Note that this [x1, . . . ,xm] need not be the inputs that train the checkpoint(s). Regardless of the availability of hi(·) and hj(·), we can extract feature representations by
Fi , [fi(x 1), . . . , fi(x m)], and Fj , [fj(x1), . . . , fj(xm)]. For notation easiness, in the following, we may write fai = fi(x a) and faj = fj(x a) for any input xa.
Then we use
KA(Fi,Fj) = 〈F>i Fi,F>j Fj〉
〈F>i Fi,F>i Fi〉 · 〈F>j Fj ,F>j Fj〉 , (1)
as an approximation to KA(Yi,Yj).
A natural question to ask is how good this approximation is. We assure ourselves by the following reasoning. Let
Ŷi = [ŷ 1 i , . . . , ŷ m i ] , [hi(f 1 i ), . . . ,hi(f m i )], and Ŷj = [ŷ 1 j , . . . , ŷ m j ] , [hj(f 1 j ), . . . ,hj(f m j )],
be the predictions made for each task (Although Ŷi may not be obtained as hi is not available. Same for Ŷj). If each checkpoint works sufficiently well for its own task, then Ŷi is very “close" to Yi, so is Ŷj to Yj . Thus KA(Ŷi, Ŷj) well approximates KA(Yi,Yj).
It remains to validate if KA(Fi,Fj) well approximates KA(Ŷi, Ŷj). The only gap between Fi and Ŷi is a task-specific head hi(·), likewise, hj(·) for Fj and Ŷj . We argue that many task-specific architectures are “well conditioned" [24, 25], which guarantees that pairwise inner products are not distorted too much. This property further results in KA(Fi,Fj) ≈ KA(Ŷi, Ŷj). To be more concrete, consider a simple special case, where hi(f) = Mif and hj(f) = Mjf are linear mappings. By the “well-conditioned" assumption, we have M>i Mi = cI and M > j Mj = dI for some c, d > 0.
Then it is easy to show that KA(Fi,Fj) = KA(Ŷi, Ŷj).
Last but not least, we have the following guarantee on positive definiteness. Proposition 1. Let κ(i, j) = KA(Fi,Fj), then κ 0.
We are not the first to use kernel alignment to understand deep nets. [26] proposes to use linear and RBF Kernel alignment to measure similarity between representations. Although the methodologies appear similar, the goals are quite different. [26] attempts to understand representation similarity, an open-ended problem. In contrast, we are explicitly interested in measuring the dependency between two tasks’ output spaces. And it turns out that KA between their respective features is a sensible surrogate. Moreover, Proposition 1 enables us to construct a Gaussian process on the task space.

heading: 3.3 Example: Characterize Linguistic Task Space using Huggingface Checkpoints
text: We showcase a concrete example using 34 checkpoints from huggingface. They are trained for various tasks (or their combinations). For example, the bert-∗ [27] checkpoints are trained for masked language modeling and next sentence prediction. t5-∗’s are for text-to-text generation. Complete details of these checkpoints and their respective tasks can be found in supplementary material. Note also that multiple checkpoints may target at the same task. These checkpoints are usually named with the same prefix. For example, the t5-∗’s only differ by their architecture configurations (e.g., number of layers). In analogous to sensor placement, this amounts to a “pilot deployment" that places multiple checkpoints at the same location in the task space. It is somewhat redundant to do so. However, they grant us a chance to do sanity check on the estimation of κ. That is, the checkpoints corresponding to the same task should have large κ(i, j) values.
We input training set of wikitext2 as probing data, and extract the contextualized word embeddings after penultimate layer. So task specific layers like softmax classifiers are ignored. Note that these checkpoints employ different tokenizers, often resulting in the same word being split into different multiple sub-words. We adopt the strategy in [28], taking the representation of a word as that of its last sub-word’s. Now columns in Fi are contextualized word representations from the i-th checkpoint. Then we can compute KA(Fi,Fj) for each pair of checkpoints to get the κ matrix. To visualize the relationship between these checkpoints (thus the tasks they target), we apply hierarchical clustering3 to the κ matrix and visualize a dendrogram in fig. 1.
In fig. 1, we observe that checkpoints with the same prefix are usually close, e.g., the t5-∗’s. This is less surprising, since obviously they are trained for the same task. A more surprising discovery is that longformer-base [29] is very similar to roberta-base [28]. A closer check of their documentation informs us that they are both trained for masked language modeling, although longformer-base innovated a novel attention mechanism. On the other hand, adapting a checkpoint to another task can result in significant difference, e.g., roberta-large-mnli against roberta-large. Therefore, the κ indeed reflects our intuitions on tasks.
Finally, there are also a few counter-intuitive cases. For example, gpt is very different from gpt2, but they are both trained for causal language modeling. We conjecture it is their different tokenizers that drive KA(Fi,Fj) to be small. A full understanding of this observation is deferred to future study.
3We use scipy functions: linkage with method=“ward", and dendrogram with color_threshold=0.9.

heading: 3.4 Robustness of the Estimated Covariance
text: Estimating κ relies on probing data. Therefore, it is natural to ask, how robust κ is against size of the probing data. In addition, we have implicitly simplified and assumed that the probing data, as well as the inaccessible training inputs for all tasks, have the same distribution on X .
(a) KL divergence (b) correlation
experiment is also repeated on another corpora, 10K sentences taken from billion-word benchmark, dubbed 1B-10K. This benchmark is created from news crawl, a different genre from wiki.
To compare κ against κ∗, we measure the KL-divergence between two Gaussian distributions with κ∗ and κ as their respective covariances. Another simple measure is the cosine similarity between vectorized vec(κ) and vec(κ∗). We plot KL(κ∗||κ) and cos(vec(κ∗), vec(κ)) against the number of words in the probing text, shown in fig. 2a and 2b respectively. For both plots, the left most part of the x-axis stands for zero probing data. In this case, it is natural to assume κ an identity matrix, i.e., all checkpoints (thus the tasks they were trained for) are independent.
In fig. 2, we observe that as the size of probing data increases, KL divergences quickly converge to 0, while the cosines quickly converge to 1. In addition, let κ∗wiki be the covariance estimated from full wikitext2 training set, and κ∗1B the covariance from full 1B-10K. It turns out that KL(κ ∗ wiki||κ ∗ 1B) = 0.76, not a large value considering the range of KL-divergences in fig. 2a. On the other hand, their cosine similarity is as high as 0.98. To conclude, the estimate of covariance function is reasonably robust against size and genre of probing data.

heading: 4 Exploiting the Task Space
text: We have characterized the task space using a zoo of checkpoints. Now we are ready to selecting checkpoints. The idea is to pick seen tasks that are the most informative about the the task space. Again, remember that this is conducted before any new task is revealed to us.
Suppose we have a budget of using K checkpoints. Let Z (with cardinally n) be the zoo of all checkpoints, and S the set of selected checkpoints. We want to maximize the mutual information between the tasks solved by S, and the remaining of the task space. With some abuse of notation, we denote this mutual information as I(S;Z − S). Here S should be understood as the tasks corresponding to the checkpoints in S. We optimize the following constrained objective,
max S⊂Z:|S|=K
I(S;Z − S). (2)
The above a combinatorial optimization problem. Enumerating all cases is infeasible. Instead, a heuristic is to use greedy search. Starting from an empty S, each time we include a checkpoint i that brings the largest gain to the objective. In specific, denote S̄ = Z − S, Si = S ∪ {i} and S̄i = Z − (S ∪ {i}). The gain δi can be calculated as,
δi = I(Si; S̄i)− I(S; S̄) = [H(Si)−H(Si|S̄i)]− [H(S)−H(S|S̄)] = [H(Si)−H(Z) +H(S̄i)]− [H(S)−H(Z) +H(S̄)] = H({i}|S)−H({i}|S̄i)
(3)
The first term implies that i should surprise the current S the most, thus incorporating more information. Meanwhile, the second term suggests that i should be representative of the remaining of the task space, thus avoiding outliers.
Since we assume a Gaussian process on the task space, any subset of Z is Gaussian distributed. And the two terms in Eq. (3) can be easily calculated. Specifically, let κ(S,S) denote the covariance for subset S, κ(i,S) the row vector of cross-covariance between {i} and S, and κ(S, i) its transpose. Then,
δi = H({i}|S)−H({i}|S̄i) = 1
2 ln 1− κ(i,S)κ(S,S)−1κ(S, i) 1− κ(j, S̄i)κ(S̄i, S̄i)−1κ(S̄i, i)
(4)
The greedy selection process is summarized in Algorithm 1. Denote S(K) as the set of selected checkpoints for a K value. Algorithm 1 guarantees that for any K1 < K2, S(K1) $ S(K2).
Algorithm 1 Maximum Mutual Information (MMI) based Selection of Checkpoints
Input: Checkpoint zoo Z with cardinality n, covariance matrix κ ∈ Rn×n, number of checkpoints to pick, K Output: a set S of selected checkpoints, where |S| = K 1: S ← ∅ 2: for k = 1, . . . ,K do 3: for i ∈ Z − S do 4: Compute information gain δi using Eq. (4) 5: end for 6: i∗ ← arg max δi 7: S ← S ∪ {i∗} 8: end for

heading: 4.1 Discussion on Quality of the Greedy Approximation
text: It is natural to ask how good the greedy method is. At this point, we need to review the definition of submodular functions. A set function f(S) is submodular if for any sets S ⊆ S ′ ⊆ Z and i ∈ Z−S ′, f(S ∪ {i})− f(S) ≥ f(S ′ ∪ {i})− f(S ′). It is not hard to show that the following holds. Proposition 2. I(S;Z − S) is a submodular function of S.
Another useful concept is monotonic function. A function f(S) is monotonic if for any i ∈ Z , f(S ∪ {i}) ≥ f(S). For monotonic submodular functions, it is known that a greedy method is guaranteed to achieve an objective no smaller than 1 − 1/e of the optimal one [30]. In our case, I(S;Z − S) is submodular but not monotonic. As I(Z;Z − S) = 0, indicating that when set S is sufficiently large, further enlarging S will reduce mutual information. It thus becomes less clear if the (1− 1/e) guarantee still holds. In optimal senor placement [15], the authors show that if the covariance is estimated for sufficiently many fine-grained grids in the room, then the mutual information is monotonic up to a reasonably big K. This may suggest that our κ has to be estimated for sufficiently many tasks. In other words, the checkpoint zoo must be rich enough. Nevertheless, it is hard to come up with a rigorous guarantee like [15], as assumptions made in 2D physical space does not necessarily transfer to an abstract task space. However, in experiments, we have found that Algorithm 1 always works in a monotonic regime.

heading: 5 Experiments
text: In the previous sections, we have presented two key components, estimation of κ and MMI based selection of checkpoints. In this section, we experiment with the two components combined. First, we apply algorithm 1 to the κ estimated in section 3.3, and show its effectiveness on multiple linguistic tasks. The baseline we compare against is random selection of checkpoints, and single commonly adopted checkpoint, e.g., bert-base-uncased. Then we extend to image classification tasks. Again we observe constant improvements over random picks, and other straightforward alternative.

heading: 5.1 Linguistic Tasks Using Huggingface Checkpoints
text: We run algorithm 1 on the κ matrix estimated in section 3.3. As K increases from 1 to 5, the S set incrementally includes roberta-base, distilbert-base-uncased, t5-base, bert-base-cased and bart-large. These checkpoints solve multiple tasks: masked language modeling, text-to-text generation, next sentence prediction, and text denoising. Referring to the visualization of κ matrix (left pannel of fig. 1), we find that the selected checkpoints all belong to the big cluster in the lower right block. This matches our intuition, as the checkpoints outside this cluster are very isolated.
Then, we apply the selected checkpoints to a suite of probing tasks collected in [28]. To be selfcontained, we briefly describe the tasks and datasets. (a) Chunking examines spans and boundaries for each word. We use CoNLL2000 shared task dataset [31]. (b) NER predicts the entity type for each word. We use CoNLL2003 shared task dataset [32]. (c) POS Tagging (POS) checks if word representations can capture basic syntax. We use PTB dataset [33]. (d) Semantic Tagging (ST) is interested in identifying the semantic role in context. The original source of dataset can be found in [34]. (e) Event Factuality(EF) aims at determining if an event is factual. We use ItHappened-v2 dataset [35]. Events are given scores between -3 (non-factual) and 3 (factual). For instance, in “I decided not to go.", the word “decided" has a score of 3 whereas “go" has -3. We use Pearson correlation to compare predicted scores against groundtruths. (f-h) Syntactic constituency ancestor tagging: A set of tasks that aim at predicting the constituent label for (f) parent; (g) Grand parent (Gparent); and (h) Great Grand parent (GGparent) of a word in the phrase structure tree. The structure is derived from PTB.
Following [28], we train a softmax on top of the combined word representations (⊗i∈Sfi) for each task. The gradients are not back-propagated through the checkpoints. There are two reasons why we choose to do so. First, for an easier comparison with [28]. Second, there are cases where checkpoints are published as blackboxes, and back-propagation is not allowed. One prominent example is GPT3 [36]. Although we do not have access to GPT-3 in this paper, we think that our experiment design should respect this fact. Another design choice is that fi is taken to be the feature at top layer of the checkpoint. We do not use feature at intermediate layers, again to respect the fact of blackboxes.
The baseline to compare against is random selection. In specific, let R ⊆ Z be a set of randomly picked checkpoints, where |R| = K. The task specific head is trained on top of ⊗i∈Rfi. For K = 1,
R ranges over each single checkpoint in the total 34. For K ≥ 2, we run 20 random picks using different random seeds. Then, the metrics from these random runs can be visualized by violin plots (fig. 3). In each violin plot, the bars on the top and middle indicate best and median performances respectively. We observe an increasing trend in the metrics as K increases from 1 to 5. Importantly, MMI achieves near optimal performance compared with random search, especially when K > 1. Note also that random selection results in large variance in performance. Another interesting baseline is a single checkpoint that are commonly adopted. We plot bert-base-uncased and bert-large-cased (at K = 1). Surprisingly, for most tasks, bert-large-cased is inferior to half of the 34 checkpoints. This echos some discoveries made in [28]. That is, higher layers of bert model is less transferable to new tasks. Instead, stronger transferability is found in middle-layer features.
Following the above discussion, we collect metrics reported in [28] for the best layer of bert-largecased. Since layers capture different levels of information [37], the best layer is often task dependent. It has to be found by brute-force training and evaluation for each task. Tab. 2 lists these metrics. Also reported are metrics achieved by ⊗i∈Sfi at K = 5. We observe that MMI outperforms the best layer of bert-large-cased for most of the tasks, and often wins by a notable margin. This is remarkable, as the checkpoints are selected before seeing any new task, and fixed throughout all tasks.

heading: 5.2 Image Classification Tasks on Cifar100
text: We are motivated by applications in computational linguistics. However, we find that the proposed framework extends to computer vision tasks as well. In this section, we simulate an example using cifar100 dataset.
We create n = 100 “seen" tasks. Each task is a 30-way classification problem. The 30 classes are randomly sampled from the total 100 classes. Each task has 480 (=500-20) training samples per class. There are 20 training samples held out for each class. The reason will be clear soon. A resnet-50 is trained for each of these “seen" tasks, stored as a checkpoint. Then, we create 20 “unseen" tasks in the same way, but each of them has only 10 training samples (part of the aforementioned holdouts) per class. This small training data is to manifest the transferability of the selected checkpoints. Finally, the remaining 10 holdouts per-class are used as probing data to estimate κ. We augment this small set of probing images by applying rotations of ±5◦, ±10◦ and ±15◦ degrees.
Figure 4 shows the estimated κ for the 100 seen tasks. Compared with the κ in fig. 1, it is less structured. That is, tasks tend to be equally distant from each other. This probably suggests that the
space of vision tasks fundamentally differs from that of linguistic tasks. We then apply algorithm 1 on the κ to select up to K = 20 checkpoints. For each task t, a classification head is learned using its 10 × 30 training data. The performance of this task is measured by accuracy on the standard validation data (excluding classes not handled in this task), denoted as acct. As before, a baseline is random selection. In addition, suppose we could peek at the first unseen task. We may identify the top-K best checkpoints for this task, by running training and evaluation. Then we can keep using these top-K checkpoints for all new tasks. This gives us another method to compare against, dubbed “peek".
In fig. 5, we plot, Etacct, the average validation accuracy over tasks against K. The shaded area are accuracies of runs using different random seed. Overall, as expected, the accuracies increase and saturate at larger K. To have a better view, we plot the gains of MMI and “peek" against random. That is, the MMI (and “peek") curve subtracting the average curve of the shaded area in fig. 5. The gains are shown in fig. 6. MMI steadily outperforms the random baseline. Whereas “peek" fluctuates and can be inferior than random at multiple K’s. This may be caused by “overfitting" to the first task that it peeked, which prevents generalization to other tasks. In analogous to senor placement, this amounts to optimizing the placement for one location, which could harm the accuracies for other locations. In contrast, MMI does not exhibit this issue.

heading: 6 Conclusion
text: In this paper, we motivate a problem that addresses two recent facts: 1) huge amount of published checkpoints; 2) emerging new tasks. By drawing analogy from optimal sensor placement, we propose a checkpoint selection framework without seeing any new task. The key steps are: 1) Estimating task relevance with proprietary constraints (e.g., data, task-specifc heads); 2) selection by maximizing mutual information, assuming a Gaussian Process on the tasks. Effectiveness is validated on tasks from computational linguistics as well as computer vision. However, it should be reminded that there are a few assumptions (section 3.2) made for step 1). We partially validated them on text data (section 3.3 and 3.4). A more comprehensive validation on other data (e.g., images) is left for future work. Aside from that, it is also interesting to explore more fine-grained models of the task space (e.g., hierarchical) and experiment with more tasks.

heading: Acknowledgement
text: The authors thank Nelson Liu for helping on data of linguistic experiments. The first author thanks Yifan Wei for her company and support during the pandemic lockdown.

Review summaries:
When no task specific heads are available sample labels cannot be predicted and the dependency between two tasks and thus task transferability cannot me measured. The paper uses a Gaussian process to model the task space and comes up with a closed solution for the mutual information gain. Proposes a greedy selection algorithm for selecting checkpoints to maximize the information gain. Experiments with text and image datasets are performed to show that selecting checkpoint based on this criterion outperforms random selection.

 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.