A linearized framework and a new benchmark for model selection for fine-tuning

Given an initialization $w_{0}$, the weights of the pre-trained model, we can define the linearized model:

which approximates the output of the real model for $w$ close to $w_{0}$. Mu et al. [35] observe that, while in general not accurate, a linear approximation can correctly describe the model throughout fine-tuning since the weights $w$ tend to remain close to the initial value $w_{0}$. Under this linear approximation [29] shows the following proposition,

From Prop. 1, the behavior of the network during fine-tuning is fully characterized by the kernel matrix $\Theta$, which depends on the pre-trained model $f_{w_{0}}$, the data $\mathcal{X}$ and the task labels $\mathcal{Y}$. We then expect to be able to select the best model by looking at these quantities. To show how we can do this, we now derive several results connecting $\Theta$ and $\mathcal{Y}$ to the quantities of relevance for model selection below, i.e. Training time and Generalization on the target task.

Training time. In  [56], it is shown that the loss $L_{t}$ of the linearized model evolves with training over time $t$ as

where we have defined $\delta\mathcal{Y}=\mathcal{Y}-f_{w_{0}}(\mathcal{X})$ to be the initial projection residual. Eq. (3) suggests using the quadratic term $\delta\mathcal{Y}^{T}\Theta\delta\mathcal{Y}$ as a simple estimate of the training speed.

Generalization. The most important criterion for model selection is generalization performance. Unfortunately, we cannot have any close form characterization of generalization error, which depends on test data we do not have. However, in [3] the following bound on the test error is suggested:

We see that if $\mathcal{Y}$ correlates more with the first principal components of variability of the per-sample gradients (so that $\mathcal{Y}\cdot v_{k}$ is larger), then we expect better generalization.

Arora et al. [3] prove that this bound holds with high-probability for a wide-enough randomly initialized 3-layer network. In practice, however, this generalization bound may be vacuous as hypotheses are not satisfied (the network is deeper, and the initialization is not Gaussian). For this reason, rather than using the above quantity as a real bound, we refer to it as an empirical “generalization score”.

Note eq. 3 and eq. 4 contain the similar terms $\delta\mathcal{Y}^{T}\Theta\delta\mathcal{Y}$ and $\delta\mathcal{Y}^{T}\Theta^{-1}\delta\mathcal{Y}$. By diagonalizing $\Theta$ and applying Jensen’s inequality we have the following relation between the two:

Hence, good “generalization score” $\delta\mathcal{Y}^{T}\Theta^{-1}\delta\mathcal{Y}$ implies faster initial fine-tuning, that is, larger $\delta\mathcal{Y}^{T}\Theta\delta\mathcal{Y}$. In general we expect the two quantities to be correlated. Hence, selecting the fastest model to train or the one that generalizes better are correlated objectives. $\mathcal{Y}^{T}\Theta\mathcal{Y}$ is an approximation to $\delta\mathcal{Y}^{T}\Theta\delta\mathcal{Y}$ that uses task labels $\mathcal{Y}$ and kernel $\Theta$, and we use it to derive our model selection scores in \secrefsec:label-correlation. Large value of $\mathcal{Y}^{T}\Theta\mathcal{Y}$ implies better generalization and faster training and it is desirable for a model when fine-tuning.

Should model selection use gradients or features? Our analysis is in terms of the matrix $\Theta$ which depends on the network’s gradients (2), not on its features. In \secrefsec:proof, we show that it suffices to use features (i.e. network activations) in (2) as an approximation to the NTK matrix. Let $[f(x_{i})]_{l}$ denote the feature vector (or activation) extracted from layer $l$ of pre-trained network $f$ after forward pass on image, i.e. after $f(x_{i})$. In analogy with the gradient similarity matrix $\Theta$ of (2), we define the feature similarity matrix $\Theta_{F}$ (which approximates $\Theta$) as follows

We now introduce our two scores for model selection, Label-Gradient correlation and Label-Feature correlation.

Label-Gradient Correlation. From \secrefsec:linearization we know that the following score,

which we call Label-Gradient Correlation (LGC), can be used to estimate both the convergence time (eq. 3) and the generalization ability of a model. Here, “$\cdot$” denotes the dot-product of the matrices (i.e. the sum of Hadamard product of two matrices). $\mathcal{Y}\mathcal{Y}^{T}$ is an $N\times N$ matrix such that $(\mathcal{Y}\mathcal{Y}^{T})_{i,j}=1$ if $x_{i}$ and $x_{j}$ have the same label and $-1$ otherwise. For this reason, we call $\mathcal{Y}\mathcal{Y}^{T}$ the label similarity matrix. On the other hand, $\Theta_{ij}=\nabla_{w}f_{w_{0}}(x_{i})\cdot\nabla_{w}f{w_{0}}(x_{j})$ is the pair-wise similarity matrix of the gradients. Hence, eq. 7 can be interpreted as giving high LG score (i.e., the model is good for the task) if the gradients are similar whenever the labels are also similar, and are different otherwise.

Label-Feature Correlation. Instead of $\Theta$, we can use the approximation $\Theta_{F}$ from (6) and define our Label-Feature Correlation (LFC) score as:

Similarly to the LGC score, this score is higher if samples with the same labels have similar features extracted from the pre-trained network.

Notice that the scores $S_{\text{LG}}$ and $S_{\text{LF}}$ are not normalized. Different pre-training could lead to very different scores if the gradients or the features have a different norm. Also, $\mathcal{Y}\mathcal{Y}^{T}$ used in our scores is specific to binary classification. In practice, we address this as follows: For a multi-class classification problem, let $K_{\mathcal{Y}}$ be the $N\times N$-matrix with $(K_{\mathcal{Y}})_{i,j}=1$ if $x_{i}$ and $x_{j}$ have the same label, and $-1$ otherwise. Let $\mu_{K}$ denote the mean of the entries of $K_{\mathcal{Y}}$, and $\mu_{\Theta}$ the mean of $\Theta$. We define the normalized LGC score as:

We normalize LFC similar to LGC in (8). This can also be interpreted as the Pearson’s Correlation coefficient between the entries of $\Theta$ (or $\Theta_{F}$) and the entries of $K_{\mathcal{Y}}$, justifying the name label-gradient (or label-feature) correlation.

Which features and gradients to use? For LFC, we extract features from the layer before the fully-connected classification layer (for both Resnet-101 [19] and Densenet-169 [22] models in our model zoo of \secrefsec:model_zoo). We use these features to construct our $\Theta_{F}$ and compute the normalized LFC. For LGC, following [35], we use gradients corresponding to the last convolutional layer in the pre-trained network. For a large gradient vector, to perform fast computation of LGC, we take a random projection to $10K$ dimensions and compute the normalized LGC score. This results in a trade-off between accuracy and computation for LGC.

Sampling of target task. Model selection is supposed to be an inexpensive pre-processing step before actual fine-tuning. To reduce its computation, following previous work of RSA [15], we sample the training set of target dataset ${\cal D}$ and pick at most $25$ images per class to compute our model selection scores. Note, test set is hidden from model selection. Our results show, this still allows us to select models that obtain accuracy gain over Imagenet expert (\figreffig:model_zoo_selection_top3), and we need few selections ($<7$ for model zoo size $30$) to select the optimal models (\figreffig:trials) to fine-tune. We include additional implementation details of our model selection methods and other baselines: RSA [15], Domain Similarity [11], LEEP [37], Feature Metrics [49] in \secrefsec:model_selection_impl.

We evaluate model selection and fine-tuning with both, a model zoo of single-domain experts (i.e. models trained on single dataset) and a model zoo of multi-domain experts described below.

Source Datasets. \tabreftab:source_dataset and \tabreftab:suppl_datasets lists the source datasets, i.e. the datasets used for training our model zoo. We include publicly available large source datasets (from $25K$ to $1.8M$ training images) from different domains, e.g. Nwpu-resisc45 [7] consists of aerial imagery, Food-101 [6] and iNaturalist 2019 [21] consist of food, plant images, Places-365 [57] and Google Landmark v2 [39] contain scene images. This allows us to maximize the coverage of our model zoo to different domains and enables more effective transfer when fine-tuning on different target tasks.

Model zoo of single-domain experts. We build a model zoo of a total of $30$ models (Resnet-101 [19] and Densenet-169 [22]) trained on $8$ large image classification datasets (i.e. source datasets). Since each model is trained on a single classification dataset (i.e. domain), we refer to these models as single-domain experts. This results in a model zoo, $\mathcal{F}=\{f^{k}\}_{k=1}^{30}$, to evaluate our model selection.

On each source dataset of \tabreftab:source_dataset, we train Resnet-101 and Densenet-169 models for 90 epochs, with the following hyper-parameters: initial learning rate of 0.1, with decay by $0.1\times$ every 30 epochs, SGD with momentum of $.9$, weight decay of $10^{-4}$ and a batch size 512. We use the training script²²2https://bit.ly/38NMvyu from PyTorch [42] library and ensure that our models are well-trained.

In \tabreftab:source_dataset, we show slightly higher top-1 test accuracy for our models trained on Imagenet [13] when compared to the PyTorch [42] model zoo³³3https://bit.ly/35vZpPE. Our Resnet-101 model trained on Imagenet has $+.17\%$ top-1 test accuracy and our Densenet-169 model has $+.4\%$ top-1 test accuracy vs. PyTorch. On source datasets other than Imagenet, we train our models with ($\checkmark$) and without ($\times$) Imagenet pre-training. This allows us to study the effect of pre-training on a larger dataset when we fine-tune and perform model selection. Note that our Resnet-101 models with ($\checkmark$) Imagenet pre-training have higher accuracy compared to that reported in the literature for all source datasets, except iNaturalist [21] by $-1.03\%$.

Model zoo of multi-domain expert. We also train a Resnet-101 based multi-dataset (or multi-domain) [45] model on the combination of all the $8$ source datasets. Our multi-domain Resnet-101 expert, $f_{w_{s},\{w_{d}\}_{d=1}^{D}}$, uses shared weights (or layers), i.e. $w_{s}$, across different domains (or datasets), and in addition it has some domain-specific parameters, i.e. $\{w_{d}\}_{d=1}^{D}$, for each domain. We have 8 source datasets or domains, so $D=8$ in our benchmark. Note, for fine-tuning we can choose any one of the $D$ domain-specific parameters to fine-tune. For a given multi-domain expert, this results in a model zoo of $D$ models (one per domain) that we can fine-tune, $\mathcal{F}=\{f_{w_{s},w_{1}},f_{w_{s},w_{2}},\cdots,f_{w_{s},w_{D}}\}$.

We experiment with a few different variants of domain-specific parameters – $i)$ Shared: The domain-specific parameters are also shared, therefore we simply train a Resnet-101 on all datasets, $ii)$ Multi-BN: We replace each batch norm in Resnet-101 architecture with a domain-specific batch norm. Note, for a batch norm layer we replace running means, scale and bias parameters, $iii)$ Adapter: We use the domain-specific parallel residual adapters [45] within the Resnet-101 architecture. Our training hyper-parameters for the multi-domain expert are the same as our single-domain expert. The only change is that for every epoch we sample at most $100K$ training images (with replacement if $100K$ exceeds dataset size) from each dataset to balance training between different datasets and to keep the training time tractable. As we show in \tabreftab:multi_domain, Multi-BN model outperforms other multi-domain models and we use it in our subsequent fine-tuning (\secrefsec:finetune) and model selection (\secrefsec:model_selection) experiments.

Target Tasks. We use various target tasks (\tabreftab:suppl_datasets) to study transfer learning from our model zoo of \secrefsec:model_zoo: Cucumber [12], Describable Textures [9], Magnetic Tile Defects [23], iCassava [36], Oxford Flowers 102 [38], Oxford-IIIT Pets [41], European Flood Depth [5], UC Merced Land Use [53]. For few-shot, due to lesser compute needed, we use additional target tasks: CUB-200 [52], Stanford Cars [27] and Belga Logos [25]. Note, while some target tasks have domain overlap with our source datasets, e.g. aerial images of UC Merced Land Use [53], other tasks do not have this overlap, e.g. defect images in Magnetic Tile Defects [23], texture images in Describable Textures [9].

Fine-tuning with single-domain experts in model zoo. For fine-tuning, Imagenet pre-training is a standard technique. Note, most deep learning frameworks, e.g. PyTorch3, MxNet/Gluon⁴⁴4https://gluon-cv.mxnet.io/api/model_zoo.html etc., just have the Imagenet pre-trained models for different architectures in their model zoo. \figreffig:finetune_full shows the top-1 test error obtained by fine-tuning single-domain experts in our model zoo vs. Imagenet expert.

Our fine-tuning hyper-parameters are: $30$ epochs, weight decay of $10^{-4}$, SGD with Nesterov momentum $0.9$, batch size of $32$ and learning rate decay by $0.1\times$ at $15$ and $25$ epochs. We observe that the most important hyper-parameter for test accuracy is the initial learning rate $\eta$, so for each fine-tuning we try $\eta=0.01$, $0.005$, $0.001$ and report the best top-1 test accuracy.

Does fine-tuning with model zoo perform better than fine-tuning a Imagenet expert? While fine-tuning an Imagenet pre-trained model is standard and works well on most target tasks, we show that by fine-tuning models of a large model-zoo we can indeed obtain a lower test error on some target tasks (see models highlighted by black $\leftarrow$ in \figreffig:finetune_full). The reduction in error is more pronounced in the low-data regime. Therefore, we establish that maintaining a model zoo of models trained on different datasets is helpful to transfer to a diverse set of target tasks with different amounts of training data.

We demonstrate gains in the low-data regime by training on a smaller subset of the target task, with only 20, 5 samples per class in \figreffig:finetune_full (i.e., we train in a 20-shot and 5-shot setting). In few-shot cases we still test on the full test set.

Fine-tuning with multi-domain expert. In \secrefsec:model_zoo, we show that fine-tuning can be done by choosing different domain-specific parameters within the multi-domain expert for fine-tuning. In \figreffig:finetune_universal, we fine-tune the multi-domain expert, i.e. Multi-BN of \tabreftab:multi_domain, on our target tasks by choosing different domain-specific parameters to fine-tune. Similar to \figreffig:finetune_full, we show the accuracy gain obtained by fine-tuning multi-domain expert with respect to fine-tuning the standard Resnet-101 pre-trained on Imagenet. We observe that selecting the correct domain to fine-tune, i.e. the correct $w_{d}$, where $d\in\{1,2,\cdots,D\}$ from multi-domain model zoo $\mathcal{F}=\{f_{w_{s},w_{d}}\}_{d=1}^{D}$, is important to obtain high fine-tuning test accuracy on the target task. In \secrefsec:model_selection, we show that model selection algorithms help in selecting the optimal domain-specific parameters for fine-tuning our multi-domain model zoo.

We also observe that fine-tuning with our multi-domain expert improves over the fine-tuning of single-domain model zoo for some tasks, e.g. iCassava: $+1.4\%$ accuracy gain with multi-domain expert compared to $+.72\%$ accuracy gain with single domain model expert over Imagenet expert. However, the comparison between single domain and multi-domain experts and their transfer properties is not the focus of our research and we refer the reader to [32, 44, 45].

In \secrefsec:finetune, using our benchmark we find that fine-tuning with a model zoo, both single-domain and multi-domain domain, improves the test accuracy on the target tasks. Now, we demonstrate that using a model selection algorithm we can select the best model or domain-specific parameters from our model zoos with only a few selections or trials.

Model Selection Algorithms. We use the following scores, $S$, for our model selection methods: LFC (see $S_{\text{LF}}$ defined in \secrefsec:implementation), LGC (see $S_{\text{LG}}$ defined in  (8)), which we introduce in \secrefsec:label-correlation. We compare against alternative measures of model selection and/or task similarity proposed in the literature: Domain similarity [11], Feature metrics [49], LEEP [37] and RSA [15]. Finally, we compare with a simple baseline: Random which selects models randomly for fine-tuning.

Model selection with single-domain model zoo. In \figreffig:model_zoo_selection_top3, we select the top-$3$ experts (i.e. $3$ highest model selection scores) for each model selection method for fine-tuning. We do this for all the target tasks (row) using each model selection method (column). We use the maximum of fine-tuning test accuracy obtained by $3$ selected models to compute accuracy gain with respect to fine-tuning with Resnet-101 Imagenet expert. Ideally, we want the accuracy gain with the model selection method to be high and equal to the “Best Gain” possible for the target task. As seen in \figreffig:model_zoo_selection_top3: LFC, LGC and LEEP obtain high accuracy gain with just 3 selections in both full dataset and 20-shot per class setting. They outperform random selection.

Model selection with multi-domain expert. For our multi-domain expert ( \secrefsec:model_zoo), we use model selection to select the domain-specific parameters to fine-tune for every model selection method. We compute the accuracy gain for fine-tuning using selected domains vs. fine-tuning Imagenet parameters in the multi-domain expert. It is desirable to have high or close to best gain with model selection. Our results in \figreffig:universal_selection_top1, show that LFC and LEEP [37] obtain higher accuracy gain compared to Feature Metrics [49] and Random selection.

Is fine-tuning with model selection faster than brute-force fine-tuning? In \tabreftab:runtime, we show that brute-force fine-tuning is expensive. We can save computation by performing model selection using LFC and LEEP and fine-tuning only the selected top-$3$ models.

How many trials to select the model with best fine-tuning accuracy? In \figreffig:trials, we measure the average of selections or trials, across all target tasks, required to select the best model for fine-tuning from the model zoo. The best model corresponds to the highest fine-tuning test accuracy on target task. Our label correlation and LEEP [37] methods can select the best model in $<7$ trials for our single domain model zoo of $30$ experts and in $<3$ trials for the multi-domain model zoo with $8$ domain experts.

Are model selection scores a good proxy for fine-tuning accuracy? In \figreffig:trials, we show our LFC scores have the highest Spearman’s ranking correlation to the actual fine-tuning accuracy for different experts. Note, we average the correlation for all our target tasks. Our LFC score is a good proxy for ranking by fine-tuning accuracy and it can allow us to select (or reject) models for fine-tuning.