SlowGuess/NeurIPS_2023_OCR huggingface dataset:
MODEL SPIDER: Learning to Rank Pre-Trained Models Efficiently
Yi-Kai Zhang $^{1}$ , Ting-Ji Huang $^{1}$ , Yao-Xiang Ding $^{2}$ , De-Chuan Zhan $^{1}$ , Han-Jia Ye $^{1,\text{忍}}$
$^{1}$ National Key Laboratory for Novel Software Technology, Nanjing University, China
$^{2}$ State Key Lab of CAD & CG, Zhejiang University
{zhangyk,huangtj,zhandc,yehj}@lamda.nju.edu.cn yxding@zju.edu.cn
Figuring out which Pre-Trained Model (PTM) from a model zoo fits the target task is essential to take advantage of plentiful model resources. With the availability of numerous heterogeneous PTMs from diverse fields, efficiently selecting the most suitable one is challenging due to the time-consuming costs of carrying out forward or backward passes over all PTMs. In this paper, we propose MODEL SPIDER, which tokenizes both PTMs and tasks by summarizing their characteristics into vectors to enable efficient PTM selection. By leveraging the approximated performance of PTMs on a separate set of training tasks, MODEL SPIDER learns to construct representation and measure the fitness score between a model-task pair via their representation. The ability to rank relevant PTMs higher than others generalizes to new tasks. With the top-ranked PTM candidates, we further learn to enrich task repr. with their PTM-specific semantics to re-rank the PTMs for better selection. MODEL SPIDER balances efficiency and selection ability, making PTM selection like a spider preying on a web. MODEL SPIDER exhibits promising performance across diverse model zoos, including visual models and Large Language Models (LLMs). Code is available at https://github.com/zhangyikaii/Model-Spider.
Fine-tuning Pre-Trained Models (PTMs) on downstream tasks has shown remarkable improvements in various fields [35, 26, 75, 42, 16], making "pre-training  $\rightarrow$  fine-tuning" the de-facto paradigm in many real-world applications. A model zoo contains diverse PTMs in their architectures and functionalities [1, 12], but a randomly selected PTM makes their helpfulness for a particular downstream task vary unpredictably [80, 70, 102]. One important step to take advantage of PTM resources is to identify the most helpful PTM in a model zoo — estimating and ranking the transferabilities of PTMs — with the downstream task's data accurately and efficiently.
Which PTM is the most helpful? A direct answer is to enumerate all PTMs and evaluate the performance of their corresponding fine-tuned models. However, the high computational cost of the backward steps in fine-tuning makes this solution impractical. Some existing methods estimate proxies of transferability with only forward passes based on the target task's features extracted by PTMs [9, 97, 66, 55, 113, 27, 71, 25, 93]. Nowadays, a public model zoo often contains hundreds and thousands of PTMs [104]. Then, the computational burden of forward passes will be amplified, let alone for the time-consuming forward passes of some complicated PTMs. Therefore, the efficiency of searching helpful PTMs and estimating the transferability should be further emphasized.
In this paper, we propose MODEL SPIDER, the SSpecification InDuced Expression and Ranking of PTMs, for accurate and efficient PTM selection. In detail, we tokenize all PTMs and tasks into vectors that capture their general properties and their relationship with each other. For example,
(a) Forward-Based and Representation/Specification-Based Model Selection
Figure 1: (a) Two strategies for PTM selection. Related works utilize forward-based features and corresponding proxies on the target dataset to evaluate transferability. The representation/specification-based approach with learned model-task pair reduces the requirement for forwarding pass on each PTM. (b) The average efficiency (wall-clock time) vs performance (correlation  $\tau_w$ , the higher, the better) comparison of PTM selection. The circle sizes indicate the memory footprint. Red circles are MODEL SPIDER with different values of the number of PTM-specific features  $k$ , while others are comparison methods. MODEL SPIDER balances efficiency and accuracy well.
(b) Selection ability & Efficiency Comparison
two models pre-trained on NABirds [37] and Caltech-UCSD Birds [100] datasets may have similar abilities in bird recognition. The comprehension abilities of two models pre-trained on XSum [64] dataset, Ax-b, and Ax-g datasets of SuperGLUE benchmark [101] may also be mutually transferable. We can then associate them with similar representation. Then the transferability from a PTM to a task could be approximated by the distance of their repr. without requiring per-PTM forward pass over the downstream task. The success of MODEL SPIDER depends on two key factors. First, how do we obtain representation for tasks and PTMs? The representation of the most helpful PTM should be close to the task one w.r.t. some similarity measures. Then, will a general task repr. weaken the selection ability since it may ignore specific characteristics of a PTM?
In MODEL SPIDER, we learn to construct representation with a general encoder and measure the similarity between them with a Transformer module [98] in a supervised learning manner. We estimate the rankings of PTMs in the model zoo for some historical tasks using rank aggregation. By leveraging the approximated supervision, we pull task representation close to the top-ranked PTM repr. and push unhelpful PTM repr. away based on the transformer-measured similarity. We expect that the ability tookenize and measure similarity could be generalized to unseen tasks. The difference between MODEL SPIDER's representation-based PTM selection with forward-based strategy is illustrated in Figure 1.
The representation generated by general encoders significantly reduces the PTM search time and improves the search performance. If the budget allows, we can extract features of the downstream task by carrying out forward passes over a part of (the top-  $k$  ranked) PTMs, revealing the specific relationship between PTMs and the task. We equip our MODEL SPIDER with the ability to incorporate PTM-specific representation, which re-ranks the PTMs and further improves the selection results. In summary, MODEL SPIDER is suitable for different budget requirements, where the general and task-specific repr. makes a flexible trade-off between efficiency and accuracy, given various forward passes. Figure 1 illustrates a comparison of PTM selection methods w.r.t. both efficiency and accuracy. Our contributions are
- We propose a novel approach MODEL SPIDER to tokenize tasks and PTMs, which is able to rank PTMs in a model zoo given a downstream task efficiently and accurately.
- MODEL SPIDER learns to tokenize and rank PTMs on a separate training set of tasks, and it can incorporate task-specific forward results of some PTMs when resource budgets allow.
- The experiments demonstrate that MODEL SPIDER effectively ranks PTMs and achieves significant improvements on the visual models and the Large Language Models (LLMs).
Efficient PTM Search with Transferability Assessment. Whether a selected PTM is helpful could be formulated as the problem measuring transferability from source data pre-training the PTM to the target downstream task [111, 13, 41, 4, 78]. The current evaluation of transferability relies on a
forward pass of the PTM on the target task, which generates the PTM-specific features on the target task. For example, NCE [97], LEEP [66], LogME [113, 114], PACTran [27], and TransRate [39] estimate negative conditional entropy, log expectation, marginalized likelihood, PAC-Bayesian bound, mutual information to obtain proxy metric of transferability, respectively. Several extensions including  $\mathcal{N}$ -LEEP [55] with Gaussian mixture model on top of PTM features, H-Score [9] utilizing divergence transition matrix to approximate the transferred log-likelihood, and [25, 71, 84] exploring correlations between categories of target task. Auxiliary information such as source clues [6, 93] and gradients of PTMs when back propagating with few steps [85, 74] are also investigated. Although the transferability assessment methods avoid the time-consuming fine-tuning, the forward costs over PTMs also become heavier given diverse and complicated pre-trained model zoos.
Relatedness of Task. Whether a PTM gains improvements after fine-tuning on the downstream task has been verified to depend on the relatedness between tasks both theoretically [10, 11, 60] and empirically [102, 58, 112]. The relatedness could be measured through various ways, such as fully fine-tuning [115], task vectors [2], example-based graphs [48, 29, 86], representation-level similarities [30, 3], and human prior knowledge [44, 76]. Instead of utilizing a pre-defined strategy to measure the relatedness, MODEL SPIDER constructs the representation of PTMs/tasks in vector forms and learns a similarity between them on historical tasks.
Learning to rank predicts the orders of objects usually with a score function [43], and the experience on a training set could be generalized to unseen data [5, 63]. Additional learned metrics or embeddings further improve the ranking ability [62, 110, 15]. The task relatedness can also be modeled as a learning-to-rank problem, where the preference over one PTM over another could be learned from historical rankings of PTMs. However, obtaining the supervision on the training set requires complete fine-tuning over a large number of historical tasks, which either come from a time-consuming transfer learning experience [103] or the output from some specially selected transferability assessment methods [28]. We propose a strong and efficient approximation of the PTM ranking supervision on the training set tasks, and a novel representation-based similarity is applied.
We describe the PTM selection problem by assuming all PTMs are classifiers, and the description could be easily extended to PTMs for other tasks, e.g., regression. Then we discuss several solutions.
3.1 Selecting PTMs from a Model Zoo
Consider we have a target classification task  $\mathcal{T} = \{(\pmb{x}_i, y_i)\}_{i=1}^N$  with  $N$  labeled examples, where the label  $y_i$  of each instance  $\pmb{x}_i$  comes from one of the  $C_{\mathcal{T}}$  classes. Instead of learning on  $\mathcal{T}$  directly, we assume there is a model zoo  $\mathcal{M} = \{f_m = W_m \circ \phi_m\}_{m=1}^M$  containing  $M$  PTMs. A PTM  $f_m$  could be decomposed into two components.  $\phi_m$  is the feature extraction network producing  $d_m$ -dimensional features.  $W_m \in \mathbb{R}^{d_m \times C_m}$  is the top-layer classifier which maps a  $d_m$ -dimensional feature to the confidence score over  $C_m$  classes. PTMs in  $\mathcal{M}$  are trained on source data across various domains. Their feature extractors  $\phi_m$  have diverse architectures, and the corresponding classifiers are pre-trained for different sets of objects. In other words,  $d_m$  and  $C_{m'}$  may differ for a certain pair of  $m$  and  $m'$ . A widely-used way to take advantage of a PTM  $f_m = W_m \circ \phi_m$  in the target task is to fine-tune the feature extractor together with a randomly initialized classifier over  $\mathcal{T}$ . In detail, we minimize the following objective
$$
\hat {f} = \hat {\boldsymbol {W}} \circ \hat {\phi} = \underset {f = \boldsymbol {W} \circ \phi} {\arg \min } \sum_ {i = 1} ^ {N} \ell \left(\boldsymbol {W} ^ {\top} \phi (\boldsymbol {x} _ {i}), y _ {i} \mid \phi_ {m}\right), \tag {1}
$$
where  $\phi$  is initialized with  $\phi_{m}$ . The fine-tuned  $f$  makes prediction with  $\arg \max_{c\in [C]}\hat{\pmb{w}}_c^\top \hat{\pmb{\phi}} (\pmb {x})$ .  $[C] = \{1,\dots ,C\}$  and  $\hat{\pmb{w}}_c$  is the cth column of  $\hat{\pmb{W}}$ . Then, we can rank the helpfulness of PTMs based on the performance of their fine-tuned models. In other words, we obtain  $\hat{f}_m$  following Equation 1 based on the mth PTM  $f_{m}$ , then we calculate the averaged accuracy when predicting over an unseen test set of  $\mathcal{T}$  (the higher, the better), i.e.,
$$
\left. \mathrm {t} _ {\phi_ {m} \rightarrow \mathcal {T}} = \mathbb {E} \left[ \mathbb {I} \left(y = \arg \max  _ {c \in [ C ]} \hat {f} _ {m} (\boldsymbol {x})\right)\right]\right.. \tag {2}
$$
<sup>1</sup>We omit the bias term for simplicity.
$\mathrm{t}_{\phi_m \to \mathcal{T}}$  is also named as the transferability, measuring if the feature extractor  $\phi_m$  in a PTM could be transferred well to the target task with fine-tuning [97, 39].  $\mathbb{I}(\cdot)$  is the indicator function, which outputs 1 if the condition is satisfied. Given  $\boldsymbol{t}_{\mathcal{T}} = \{\mathrm{t}_{\phi_m \to \mathcal{T}}\}_{m=1}^{M}$ , i.e., the transferability for all PTMs, then we can obtain the ground-truth ranking of all PTMs in the model zoo for task  $\mathcal{T}$  and select the top-ranked one. In the PTM selection problem, the goal is to estimate the ranking of all PTMs for a task  $\mathcal{T}$  using  $\hat{\boldsymbol{t}}_{\mathcal{T}} = \{\hat{\mathrm{t}}_{\phi_m \to \mathcal{T}}\}_{m=1}^{M}$ . The evaluation criterion is the similarity between the predicted  $\hat{\boldsymbol{t}}_{\mathcal{T}}$  and the ground-truth  $\boldsymbol{t}_{\mathcal{T}}$ , typically measured by weighted Kendall's  $\tau_w$  [45]. We omit the subscript  $\mathcal{T}$  when it is clear from the context.
3.2 Efficiency Matters in PTM Selection
One direct solution to PTM selection is approximating the ground truth  $t_{\mathcal{T}}$  by fine-tuning all the PTMs over  $\mathcal{T}$ , where a validation set should be split from  $\mathcal{T}$  to estimate Equation 2. Since fine-tuning PTM contains multiple forward and backward passes, the computation burden is astronomical.
A forward pass of a certain PTM's extractor  $\phi_{m}$  over  $\mathcal{T}$  generates the features  $\Phi_T^m = \{\phi_m(\pmb{x}_i)\in$ $\mathbb{R}^{d_m}\}_{(\pmb {x}_i,y_i)\in \mathcal{T}}$ , which is lightweight compared with the backward step. The feature reveals how examples in  $\mathcal{T}$  are distributed from the selected PTM's view, and a more discriminative feature may have a higher transfer potential. As mentioned in section 2, the existing transferability assessment methods estimate  $\mathrm{t}_{\phi_m\to T}$  based on the PTM-specific feature  $\Phi_T^m$  and target labels  $\{y_i\}_{i = 1}^N$  [66, 113, 55, 114]. Precise estimation requires a large  $N$ , which means we need to collect enough examples to identify the most helpful PTMs from a model zoo.
While the previous forward-based transferability assessment methods reduce the time cost, selecting among  $M$  PTMs in the model zoo multiplies the forward cost  $M$  times, making the estimation of  $\hat{t}$  computationally expensive. Moreover, since forward passes for complicated PTMs take longer, selecting a PTM efficiently, especially given a large model zoo, is crucial.
4 MODEL SPIDER
In MODEL SPIDER, we propose to tokenize PTMs and tasks regardless of their complexity, allowing us to efficiently calculate their relatedness based on a certain similarity measure over their representation, which capture general properties and serve as a specification of a model or task, demonstrating which kinds of tasks a model performs well on or what kind of models a task requires. In this section, we first introduce the process of obtaining repr. by learning from a training set of tasks, and the ability to rank PTMs could be generalized to downstream tasks. We then describe the encoder, the similarity measure, and an efficient way to generate supervision during representation learning. Finally, we discuss how MODEL SPIDER can be flexible in incorporating forward pass results of top-ranked PTMs to further improve the representation's semantics and the ranking's quality.
4.1 Learning to Rank PTMs with Representation
In MODEL SPIDER, we learn the model repr.  $\{\pmb{\theta}_m\}_{m=1}^M$ , task repr.  $\pmb{\mu}(\mathcal{T})$ , and the similarity measure  $\mathrm{sim}(\cdot, \cdot)$  in a supervised learning manner based on a separate training set  $\mathcal{D}$ . The training set  $\mathcal{D}$  does not contain overlapping classes with the downstream task  $\mathcal{T}$ .
Specifically, we randomly sample training tasks  $\{\mathcal{T}_i\}$  from  $\mathcal{D}$ . For a given training task  $\mathcal{T}_i$ , we assume that we can obtain the ground-truth ranking  $t_{\mathcal{T}_i} = \{\mathrm{t}_{\phi_m \to \mathcal{T}_i}\}_{m=1}^M$  over the  $M$  PTMs, indicating the helpfulness of each PTM. We will discuss the details of obtaining the supervision  $t_{\mathcal{T}_i}$  later. We then select PTMs for  $\mathcal{T}_i$  based on the similarity between the task repr.  $\pmb{\mu}(\mathcal{T}_i)$  and those  $M$  PTM repr.  $\{\pmb{\theta}_m\}_{m=1}^M$ . We expect the higher the similarity, the more helpful a PTM is for the given task. We use  $\Theta$  to denote all learnable parameters and optimize  $\Theta$  with a ranking loss, which minimizes the discrepancy between the rank  $\hat{t}_{\mathcal{T}_i}$  predicted by the similarity function and the ground-truth  $t_{\mathcal{T}_i}$ :
$$
\min  _ {\Theta} \sum_ {\mathcal {T} _ {i} \sim \mathcal {D}} \ell_ {\text {r a n k}} \left(\hat {\boldsymbol {t}} _ {\mathcal {T} _ {i}} = \left\{\operatorname {s i m} \left(\boldsymbol {\theta} _ {m}, \boldsymbol {\mu} \left(\mathcal {T} _ {i}\right)\right) \right\} _ {m = 1} ^ {M}, \boldsymbol {t} _ {\mathcal {T} _ {i}}\right). \tag {3}
$$
Given  $\pmb{t} \in \mathbb{R}^{M}$ , we use an operator  $\mathrm{dsc}(\cdot)$  to index the elements of  $\pmb{t}$  in a descending order, i.e.,  $\forall m < l$ , we have  $\pmb{t}_{\mathrm{dsc}(m)} \geqslant \pmb{t}_{\mathrm{dsc}(l)}$ .  $\mathrm{dsc}(m)$  is exactly the index of the PTM with  $m$ th largest
Figure 2: An illustration of MODEL SPIDER. The middle part (b) shows the workflow of MODEL SPIDER, which involves tokenizing both PTMs and tasks into a shared space. Plot (c) demonstrates how the model-task similarity calculated based on the representation helps rank PTMs for a given task. In plot (a), when the budget allows, MODEL SPIDER can take advantage of PTM-specific features obtained by performing forward passes of the top- $k$  ranked PTMs on some selected tasks. This improves the quality of task repr. as well as the PTM ranking.
ground-truth score. Based on this, we use the following ranking loss:
$$
\ell_ {\text {r a n k}} (\hat {\boldsymbol {t}}, \boldsymbol {t}) = \sum_ {m = 1} ^ {M} - \log \left(\frac {\exp \left(\hat {\boldsymbol {t}} _ {\mathrm {d s c} (m)}\right)}{\sum_ {l = m} ^ {M} \exp \left(\hat {\boldsymbol {t}} _ {\mathrm {d s c} (l)}\right)}\right). \tag {4}
$$
Equation 4 aims to make the whole order of the predicted  $\hat{t}_{\mathcal{T}_i}$  similar to the ground-truth  $t_{\mathcal{T}_i}$ . So the similarity between the task repr. and that of a higher-ranked PTM indicated by  $t_{\mathcal{T}_i}$  should be larger than the similarity with lower-ranked PTM representation. The underlying intuition is that if a PTM performs well on certain tasks, it is likely to generalize its ability to related tasks. For example, if a PTM excels at bird recognition, it may effectively recognize other flying animals.
For a downstream task  $\mathcal{T}$ , we generate its task repr. with  $\mu(\mathcal{T})$ , and identify the close PTM ones with the learned  $\mathrm{sim}(\cdot, \cdot)$ . Objective Equation 3 also works when the number of examples in a task is small. By learning to rank PTMs for sampled few-shot tasks, MODEL SPIDER can rank helpful models even with limited training data. We will show this ability of MODEL SPIDER in section 5.
4.2 Model and Task Representation for PTM Selection
We encode the general characteristics of tasks and PTMs via two types of representation.
Model Representation. Given a model zoo with  $M$  PTMs, we associate a PTM  $f_{m}$  with a form  $\theta_{m} \in \mathbb{R}^{d}$  encoding rich semantics about the aspects in which  $f_{m}$  excels. Models pre-trained from related datasets or those with similar functionalities are expected to have similar representation.
Task Representation. A  $C_{\mathcal{T}}$ -class task  $\mathcal{T} = \{(x_i, y_i)\}_{i=1}^N$  contains a set of instances and labels. We would like tockenize a task with a mapping  $\mu(\cdot)$ , which outputs a set of vectors  $\mu(\mathcal{T}) \in \mathbb{R}^{d \times C_{\mathcal{T}}}$  one for each class. We implement  $\mu$  with one additional frozen encoder  $\psi$  with an equivalent parameter magnitude as the PTMs in the model zoo.  $\psi$  is pre-trained by self-supervised learning methods [17, 33, 53] and captures the semantics of a broad range of classes. In detail, we extract the features of all instances in the task  $\mathcal{T}$  and take the class centers as the task repr.:
$$
\boldsymbol {\mu} (\mathcal {T}) = \left\{\frac {1}{| \mathbb {I} (y _ {i} = c) |} \sum_ {\left(\boldsymbol {x} _ {i}, y _ {i}\right) \in \mathcal {T}} [ \psi \left(\boldsymbol {x} _ {i}\right) \cdot \mathbb {I} (y _ {i} = c) ] \right\} _ {c \in [ C ]}. \tag {5}
$$
The task repr. expresses the characteristics of a task, e.g., those tasks with semantically similar classes may have similar sets of representation.
Model-Task Similarity. The helpfulness of a PTM w.r.t. a task, i.e., the transferability score, could be estimated based on the similarity of the model-task pairs  $\hat{\mathfrak{t}}_{\phi_m\rightarrow \mathcal{T}} = \mathrm{sim}(\pmb {\theta}_m,\pmb {\mu}(\mathcal{T}))$  , and the PTM selection is complemented by embedding the model and tasks into a space and then identifying close PTM repr. for a task. In MODEL SPIDER, the  $\mathrm{sim}(\cdot ,\cdot)$  is implemented with a one-layer Transformer [98], a self-attention module that enables various inputs. The Transformer consists of alternating layers of multi-head self-attention, multi-layer perceptron, and layer norm blocks. We set the input of the Transformer as the union set of model and task repr.  $z = [\pmb {\theta}_m,\pmb {\mu}(\mathcal{T})]\in \mathbb{R}^{d\times (1 + C)}$
then the similarity  $\hat{\mathfrak{t}}_{\phi_m\to T}$  between model and task ones is:
$$
\sin \left(\theta_ {m}, \boldsymbol {\mu} (\mathcal {T})\right) = \operatorname {F C} \left(\text {t r a n s f o r m e r} (\boldsymbol {z}) [ 0 ]\right), \tag {6}
$$
where [0] is the first output of the Transformer, i.e., the corresponding output of the model representation. We add a Fully Connected (FC) layer to project the intermediate result to a scalar. Learnable parameters  $\Theta$ , including  $\{\theta_m\}_{m=1}^M$ , FC, and weights of the Transformer, are trained via objective in Equation 3.
4.3 Accelerating Training for MODEL SPIDER
The training of MODEL SPIDER in Equation 3 requires a large number of (task  $\mathcal{T}_i$ , PTM ranking  $t_{\mathcal{T}_i}$ ) pairs. Although we could collect enough data for each task, obtaining the ground-truth PTMs rankings, i.e., the helpfulness order of PTMs for each task, is computationally expensive. In addition, using some proxies of  $t_{\mathcal{T}_i}$  may weaken the ability of the MODEL SPIDER. We propose a closer approximation of the ground-truth  $t_{\mathcal{T}_i}$ , which efficiently supervises sampled tasks from  $\mathcal{D}$ .
Approximated Training Supervision. We take advantage of the fact that existing PTM selection methods rely on the PTM-specific features  $\Phi_{\mathcal{T}_i}^m$  to estimate the transferability score w.r.t.  $\mathcal{T}_i$  and produce diverse scores. In other words, a PTM will be placed in different positions based on the scores provided by various methods such as NCE [97], LEEP [66], and LogME [113, 114]. Based on their "relatively good but diverse" ranking results, an intuitive approach to estimate the ground-truth  $t_{\mathcal{T}_i}$  is to ensemble their multiple ranking results into a stronger single order.