Active Finetuning: Exploiting Annotation Budget
in the Pretraining-Finetuning Paradigm

We formally define the active finetuning task. As is demonstrated in Fig. 1, a deep neural network model $f(\cdot;w_{0}):\mathcal{X}\rightarrow\mathbb{R}^{C}$ with pretrained weight $w_{0}$ is given, where $\mathcal{X}$ is the data space and $\mathbb{R}^{C}$ is the normalized high-dimensional feature space. We also have access to a large unlabeled data pool $\mathcal{P}^{u}=\{\mathbf{x}_{i}\}_{i\in[N]}\sim p_{u}$ inside data space $\mathcal{X}$ with distribution $p_{u}$, where $[N]=\{1,2,\dots,N\}$. The subset $\mathcal{P}^{u}_{\mathcal{S}}$ for supervised finetuning is selected from $\mathcal{P}^{u}$. It is worth noting that $f(\cdot;w_{0})$ can be pretrained either on $\mathcal{P}^{u}$ or other data sources, e.g. pretrained on ImageNet-1k [38] and finetuned on a subset of CIFAR-10 [23].

In the active finetuning task, we should design a sampling strategy $\mathcal{S}=\{s_{j}\in[N]\}_{j\in[B]}$ to select a subset $\mathcal{P}_{\mathcal{S}}^{u}=\{\mathbf{x}_{s_{j}}\}_{j\in[B]}\subset\mathcal{P}^{u}$ from $\mathcal{P}^{u}$, where $B$ is the annotation budget size for supervised finetuning. The model would have access to the labels $\{\mathbf{y}_{s_{j}}\}_{j\in[B]}\subset\mathcal{Y}$ of this subset through the oracle, obtaining a labeled data pool $\mathcal{P}^{l}_{\mathcal{S}}=\{\mathbf{x}_{s_{j}},\mathbf{y}_{s_{j}}\}_{j\in[B]}$, where $\mathcal{Y}$ is the label space. Afterward, the model $f$ is finetuned on $\mathcal{P}^{l}_{\mathcal{S}}$ supervisedly and the model parameter is updated to $w_{\mathcal{S}}$ after the finetuning. The goal of active finetuning is to find the sampling strategy $\mathcal{S}_{opt}$ minimizing the expected model error $error(f(\mathbf{x};w_{\mathcal{S}}),\mathbf{y})$.

Our active finetuning is different from traditional active learning in: 1) We have access to the pretrained model $f(\cdot;w_{0})$, which will be finetuned, before data selection. 2) The selected samples are applied to the finetuning of the pretrained model $f(\cdot;w_{0})$ instead of from-scratch training. 3) The sampled subset size $|\mathcal{P}_{\mathcal{S}}^{l}|$ is relatively small, less than $10\%$ in most cases. 4) We have no access to any labels such as a random initial labeled set before data selection.

We select samples under the guidance two basic intuitions: 1) bringing close the distributions between the selected subset $\mathcal{P}_{\mathcal{S}}^{u}$ and the original pool $\mathcal{P}^{u}\sim p_{u}$. 2) maintaining the diversity of $\mathcal{P}_{\mathcal{S}}^{u}$. The former ensures the model finetuned on the subset performs similarly with that trained on the full set, while the latter allows the subset to cover corner cases in the full set. In comparison to distribution $p_{u}(\mathbf{x})$ in the data space, it is more feasible to work on its corresponding distribution $p_{f_{u}}(\mathbf{f})$ in the feature space. Through the agency of pretrained model $f(\cdot;w_{0})$, we map each data sample $\mathbf{x}_{i}$ to the high dimensional feature space as $\mathbf{f}_{i}=f(\mathbf{x}_{i};w_{0})$, where $\mathbf{f}_{i}$ is the normalized feature of $\mathbf{x}_{i}$. As a result, we can derive the pool $\mathcal{F}^{u}=\{\mathbf{f}_{i}\}_{i\in[N]}$ from $\mathcal{P}^{u}$ and corresponding distribution $p_{f_{u}}$ of $\mathcal{F}^{u}$.

Similarly, the feature pool $\mathcal{F}_{\mathcal{S}}^{u}$ is also associated with the selected data subset $\mathcal{P}_{\mathcal{S}}^{u}$. We define the corresponding distribution over $\mathcal{F}_{\mathcal{S}}^{u}$ in the feature space as $p_{f_{\mathcal{S}}}$. Our goal is to find the optimal selection strategy $\mathcal{S}$ as follows.

where $D(\cdot,\cdot)$ is some distance metrics between distributions, $R(\cdot)$ is to measure the diversity of a set, and $\lambda$ is a scale to balance these two terms. The first term aims to bring close these two distributions $p_{f_{u}},p_{f_{\mathcal{S}}}$ while the second term is to ensure the diversity of subset.

Unfortunately, it is difficult to directly optimize the discrete selection strategy $\mathcal{S}$, so we alternatively model $p_{f_{\mathcal{S}}}$ with $p_{\theta_{\mathcal{S}}}$, where $\theta_{\mathcal{S}}=\{\theta_{\mathcal{S}}^{j}\}_{j\in[B]}$ are the continuous parameters and $B$ is the annotation budget size. Each $\theta_{\mathcal{S}}^{j}$ after optimization corresponds to the feature of a selected sample $\mathbf{f}_{s_{j}}$. We would find $\mathbf{f}_{s_{j}}$ closest to each $\theta_{\mathcal{S}}^{j}$ after optimization to determine the selection strategy $\mathcal{S}$. Therefore, our goal in Eq. 2 is written as follows.

The difference between extracted sample features $\mathcal{F}_{\mathcal{S}}^{u}=\{\mathbf{f}_{s_{i}}\}$ and our define parameters $\theta_{\mathcal{S}}=\{\theta_{\mathcal{S}}^{j}\}$ is that $\mathbf{f}_{s_{i}}$ is a discrete feature corresponding to a sample in the dataset while $\theta_{\mathcal{S}}^{j}$ is continuous in the feature space.

In the parametric model $p_{\theta_{\mathcal{S}}}$, the distribution is represented by $B$ parameters $\{\theta_{\mathcal{S}}^{j}\}_{j\in[B]}$. We consider it as a mixture model with $B$ components in Eq. 4.

where $\phi_{j}$ is the mixture weight or prior probability $p(\theta_{\mathcal{S}}^{j})$ of the $j$-th component, satisfying $\sum_{j=1}^{B}\phi_{j}=1$. Since $\mathbf{f}$ and $\theta_{\mathcal{S}}^{j}$ both lie in the feature space, we formulate the distribution of each component based on their similarity as Eq. 5.

where $Z_{j}$ is a normalizing constant, $sim(\cdot,\cdot)$ is a similarity metric, and $\tau$ is the temperature scale. We follow the protocol in [47, 6] to apply the cosine similarity between normalized features as the metric $sim(\mathbf{f}_{1},\mathbf{f}_{2})=\mathbf{f}_{1}^{\top}\mathbf{f}_{2},||\mathbf{f}_{1}||_{2}=||\mathbf{f}_{2}||_{2}=1$ and set the temperature $\tau=0.07$ [47, 6] throughout the paper. For each $\mathbf{f}_{i}\in\mathcal{F}^{u}$, there exists a $\theta_{\mathcal{S}}^{c_{i}}$ most similar (and closest) to $\mathbf{f}_{i}$, i.e.

where we keep updating $c_{i}$ in the optimization process.

Since there is a very low temperature ($\tau=0.07$), the gap between the exponential similarity $\exp(sim(\mathbf{f}_{i},\theta_{\mathbf{S}}^{j})/\tau)$ with different $\theta_{\mathbf{S}}^{j}$ is significant. Therefore, it is safe to make the following assumption.

Although it is hard to mathematically prove, this assumption is empirically verified by the results in Fig. 3. In another word, $p(\mathbf{f}_{i}|\theta_{\mathcal{S}}^{c_{i}})\gg p(\mathbf{f}_{i}|\theta_{\mathcal{S}}^{j}),j\neq c_{i},j\in[B]$. We can approximate the parametric model for $\mathbf{f}_{i}\in\mathcal{F}^{u}$ in Eq. 4.

where $\tilde{Z}_{c_{i}}=Z_{c_{i}}/\phi_{c_{i}}$ is a new normalizing constant. We can derive $p_{\theta_{\mathcal{S}}}(\mathbf{f}_{i})\propto\exp(sim(\mathbf{f}_{i},\theta_{\mathcal{S}}^{c_{i}})/\tau)$ from Eq. 7.

The two distributions $p_{f_{u}},p_{\theta_{\mathcal{S}}}$ can be brought close by minimizing the KL-divergence.

where the first term $\underset{\mathbf{f}_{i}\in\mathcal{F}^{u}}{E}\left[\log p_{f_{u}}(\mathbf{f}_{i})\right]$ is a constant without the parameter $\theta_{\mathcal{S}}$. Then, minimizing the KL-divergence $KL(p_{f_{u}}|p_{\theta_{\mathcal{S}}})$ equals to maximizing the second term $\underset{\mathbf{f}_{i}\in\mathcal{F}^{u}}{E}\left[\log p_{\theta_{\mathcal{S}}}(\mathbf{f}_{i})\right]$, and according to Eq. 7, it is also equivalent with maximizing $\underset{\mathbf{f}_{i}\in\mathcal{F}^{u}}{E}\left[\log\exp(sim(\mathbf{f}_{i},\theta_{\mathcal{S}}^{c_{i}})/\tau)\right]=\underset{\mathbf{f}_{i}\in\mathcal{F}^{u}}{E}\left[sim(\mathbf{f}_{i},\theta_{\mathcal{S}}^{c_{i}})/\tau\right]$. Therefore, we derive the first term in Eq. 3 as follows.

However, directly carrying out this optimization leads to a severe collapse problem, i.e. most $\theta_{\mathcal{S}}^{j},j\in[B]$ converge to the same position with the highest density of $\mathbf{f}_{i},i\in[N]$, losing the diversity inside the selected data. To this end, as shown in Eq. 2, we introduce an extra regularization term to ensure the diversity of selected subset. Without bells and whistles, this regularization is implemented by minimizing the similarity between selected samples. We also add an exponential operation to make the optimization process more stable, otherwise some $\theta_{\mathcal{S}}^{j}$ may become outliers.

At this point, we are able to solve Eq. 3 by optimizing the following loss function continuously.

where the balance weight $\lambda$ is empirically set as $1$.

Finally, we directly optimize the loss function in Eq. 11 by gradient descent. When the optimization is finished, we find feature $\{\mathbf{f}_{s_{j}}\}_{j\in[B]}$ with the highest similarity to $\theta_{\mathcal{S}}^{j}$.

The corresponding data samples $\{\mathbf{x}_{s_{j}}\}_{j\in[B]}$ are selected as the subset $\mathcal{P}_{\mathcal{S}}^{u}$ with selection strategy $\mathcal{S}=\{s_{j}\}_{j\in[B]}$.

In this part, we show that optimizing the loss in Eq. 11 is actually minimizing the earth mover’s distance between the distributions of selected subset and full set. This justifies that our optimization in the continuous space is equivalent with bringing close the distribution gap in the discrete data sample space.

After the optimization, we get the features $\mathbf{f}_{s_{j}}$ of selected samples. We deliberately assign the discrete probability distribution $p_{f_{S}}$ as Eq. 13.

where $C_{j}$ is the set of features closest to $f_{s_{j}}$ with $c_{i}$ defined in Eq. 6. The distribution $p_{f_{u}}$ is modeled as a uniform distribution over $\mathcal{F}^{u}$, i.e. $p_{f_{u}}(\mathbf{f}_{i})=\frac{1}{N},\mathbf{f}_{i}\in\mathcal{F}^{u}$.

The earth mover’s distance (EMD) between $p_{f_{u}},p_{f_{S}}$ is written as [24]:

where $\Pi(p_{f_{u}},p_{f_{S}})$ is the set of all possible joint distributions whose marginals are $p_{f_{u}}$ and $p_{f_{S}}$. It is intuitive to come up with the infimum, i.e. each $\mathbf{f}_{i}\sim p_{f_{u}}$ transports to their closest $\mathbf{f}_{s_{j}}\sim p_{f_{\mathcal{S}}}$. The detailed derivation is in the appendix.

In this case, the distance in Eq. 14 becomes

In Eq. 11, we minimize $-sim(\mathbf{f}_{i},\theta_{\mathcal{S}}^{c_{i}})$, and $\mathbf{f}_{s_{c_{i}}}$ is set as the closest $\mathbf{f}_{k}\in\mathcal{F}^{u}$ to $\theta_{\mathcal{S}}^{c_{i}}$ in Eq. 12 after optimization. Then, the distance in Eq. 16 is actually also minimized. Therefore, our optimization method in Sec. 3.3 is equivalent with reducing the earth mover’s distance between the distributions of the original unlabeled pool and selected subset.

Alg. 1 shows how to implement this method to deep learning models. Given a pretrained model, for each image sample $\mathbf{x}_{i}\in\mathcal{P}^{u}$, we extract the last layer [CLS] token feature in the transformer model or global pooling feature in the convolutional model, which is normalized as the high-dimensional feature $\mathbf{f}_{i}=f(\mathbf{x}_{i};w_{0})$. Before the optimization process, the parameter $\theta_{\mathcal{S}}$ is initialized by uniformly sampling $\theta_{\mathcal{S}}^{j},j\in[B]$ at random from the feature pool $\mathcal{F}^{u}=\{\mathbf{f}_{i}\}_{i\in[N]}$. If $|\mathcal{F}^{u}|$ is extremely large, we would randomly select $M$ elements from $\mathcal{F}^{u}$ (e.g. $M$=100,000 for ImageNet dataset) for the each training iteration of our parametric model. In each iteration, we calculate the similarity between sample features and parameters, then update $c_{i}$ in Eq. 6 for each $\mathbf{f}_{i}$ and positive feature set $\{\mathbf{f}_{i}|c_{i}=j\}$ for each $\theta_{\mathcal{S}}^{j}$. Afterwards, we can compute the loss function in Eq. 11 and update the parameters $\theta_{\mathcal{S}}$ by gradient descent. When the optimization process is finished, we find the sample feature $\mathbf{f}_{s_{j}}$ most similar to each parameter $\theta_{\mathcal{S}}^{j}$ (Eq. 12). Those corresponding samples $\{\mathbf{x}_{s_{j}}\}_{j\in[B]}$ are selected for annotation for the following supervised finetuning.

Datasets and Benchmarks For classification task, we choose three classical datasets CIFAR10, CIFAR100 [23], and ImageNet-1k [38] for experiments. CIFAR10 and CIFAR100 contain 60,000 images of 32x32 scale with 10 and 100 categories separately, among which 50,000 images are in the training set and 10,000 images are for test. ImageNet-1k is a large-scale dataset spans 1000 classes, containing 1,281,167 training images and 50,000 validation images. Their training sets are considered as the unlabeled pool $\mathcal{P}^{u}$ for selection. We evaluate the performance with the Top-1 Accuracy metric.

Implementation Details Our method is agnostic with pretraining frameworks and networks. We apply DeiT-Small [43] pretrained with DINO [6] framework on ImageNet-1k [38] in the experiments for its verified popularity and effectiveness. We also attempt other architectures in Sec. 4.3. For all three datasets, we resize images to 224x224 consistent with the pretraining for both data selection and supervised finetuning. In the data selection process, the parameters $\theta_{\mathcal{S}}$ are optimized with Adam [22] optimizer (learning rate 1e-3) until convergence. The experiment details of supervised finetuning are available in the supplementary materials.

Baselines We compare our method with eight counterparts including the following three baselines and five traditional active learning methods.

1. 1.

   Random: The samples for annotation are selected uniformly at random.

2. 2.

   FDS: a.k.a K-Center-Greedy algorithm. It selects the next sample feature farthest from current selections. Proved in [39], it minimizes the gap between the expected loss over the entire pool and the selected subset. In accordance with the pretraining process [6], we apply cosine distance as the distance metric.

3. 3.

   K-Means: We conduct K-Means over the feature pool $\mathcal{F}^{u}$ and choose samples closest to the centers. $K$ equals to the budget size $B$.

We transplant five active learning algorithms CoreSet [39], VAAL [40], LearnLoss [49], TA-VAAL[21] and ALFA-Mix[34] to our active finetuning task. The former three are classical and the latter two are newer, all equipped with open-source codes. We refer readers to the appendix for transplantation details.

Results and Comparison We average our results over three independent runs. The results are shown in Tab. 1. Traditional active learning methods typically fail in the pretraining-finetuning paradigm, which is consistent with the results in [4]. In contrast, our ActiveFT outperforms counterparts on all three datasets with different sampling ratios. On each dataset, the performance gain is especially significant when the sampling ratio is low, since our method can select the most representative samples. This phenomenon is of great practical use because the sampling number for supervised finetuning is usually much smaller than the pool size to save the annotation cost.

Datasets and Benchmarks For segmentation task, we apply ADE20k dataset [51]. It contains 20,210 images for training, 2,000 images for validation, and 3,352 images for test. All images have fine-grained labels with 150 semantic classes. The training set is considered as the unlabeled pool $\mathcal{P}^{u}$ for selection. We evaluate the performance with the mIoU metric.

Implementation Details Same with image classification task, we apply DeiT-Small [43] model pretrained with DINO framework [6] for data selection. The images are resized to 224x224 as well. Since the semantic segmentation task relies more on the local information inside images, we concatenate the [CLS] token features with the average features of other tokens as $\mathbf{f}_{i}$ for data selection. For the segmentation task, Segmenter [41] is adopted for finetuning, which is a pure transformer model. We use the same pretrained DeiT-Small [6] as its backbone. The finetuning details are also available in appendix.

Results and Comparison In Tab. 2, we report the performance of our method when choosing $5\%,10\%$ of training samples for finetuning. Our results are averaged over three independent trials. We compare to three baselines same with image classification. The traditional active learning methods are not included due to their failure on image classification. The performance gain of our data selection method is not as significant as the image classification task. This is understandable because semantic segmentation is a fine-grained vision task, focusing on subtle local visual pattern. In this case, it is hard for a global feature to represent all the details in a scene. However, despite the difficulty, our method still shows notable superiority in comparison with other baselines, reflecting the generality of our method to different tasks.

Data Selection Efficiency It is desirable that the data selection method operates in a time-efficient manner, as close as possible to random selection. In Tab. 3, we compare the required time to select different percentages of training samples from CIFAR100. We do not take FDS into account due to its very poor performance. For those traditional active learning algorithms, both the repetitive model training and data sampling should be counted into the running time, and the former takes the majority. In contrast, our method chooses all the samples in a single-pass, so we do not have to train the model again in the selection process. As a result, our method’s speed is faster than traditional active learning methods by a notable margin. Besides, unlike some active learning methods [39, 49], our method does not need access to ground-truths before the end of all selection, which enables more flexibility of annotator assignment.

Visualization of Selected Samples In Fig. 4, we visualize the feature $\mathbf{f}_{i}$ of each sample in CIFAR10 training set. The dimension of features is reduced by tSNE. The black dots represent the $1\%$ samples selected by different methods. Results demonstrate that our selected samples distribute more similarly with the entire pool in the feature space than other counterparts. It reflects that optimizing our proposed parametric model helps to reduce the distribution gap between the selected subset and original unlabeled pool.

Generality of our Method ActiveFT can fit different pretraining frameworks and models well. In addition to DINO [6] framework and DeiT-Small [43] model, we also apply ActiveFT to a DeiT-Small [43] trained with generative unsupervised pretraining framework iBOT [52] and CNN model ResNet-50 [16] trained with DINO [6]. The models are pretrained on ImageNet-1k and would be finetuned in CIFAR10. Other implementation details are exactly same with Sec. 4.1. Tab 4 shows the significant superiority of our results in comparison with random sampling baseline with different sampling ratios. The results reflect the compatibility of ActiveFT with different unsupervised pretraining frameworks and model architectures.

We discuss the importance of different modules in our method including the update manner of $c_{i}$, the design of diversity regularization, and the effect of temperature.

Update Manner of $c_{i}$ In this part, we discuss the ways to update $c_{i}$ (Eq. 6) which denotes the parameter closest to each sample $\mathbf{f}_{i}$ in the feature space. In Alg. 1, it is updated in each iteration. Alternatively, we remain $c_{i}$ unchanged as the initial state in the optimization process. Results in Tab. 5(a) shows that this stationary strategy does not work well. In this case, it would rely heavily on the initial state. The frequent update of $c_{i}$ could help to relieve some harmful biases inside the initial state.

Regularization Design We try two alternative strategies to design the regularization term $R(\cdot)$ in Eq. 11. S1) No Regularization: We only optimize the first term $D(\cdot,\cdot)$ in Eq. 11. S2) InfoNCE [44]: We get inspiration from [44] to design a contrastive loss to approximate the distribution $p_{f_{u}}$ with $p_{\theta_{\mathcal{S}}}$: $L=-\underset{\mathbf{f}_{i}\in\mathcal{F}^{u}}{E}\left[\log\frac{\exp(\mathbf{f}_{i}^{T}\theta_{\mathcal{S}}^{c_{i}}/\tau)}{\sum_{k\in[N]}\exp(sim(\mathbf{f}_{k}^{T}\theta_{\mathcal{S}}^{c_{i}}/\tau)}\right]$. In Tab. 5(b), we evaluate these three strategies. We find that both S1 and S2 fails, and only our applied strategy S3 succeeds. It justifies our design of the regularization strategy.

Temperature $\tau$ We analyze the effect of different temperatures in Eq. 11. Pointed out in Assumption 1, a small $\tau$ is a pre-requisite for our derivation. Tab. 6 shows the results on CIFAR10 with different temperatures. When the temperature is relatively low (e.g. $\tau<$0.1), the performance of ActiveFT is great. However, as it becomes higher (e.g. $\tau=0.5$), the performance drops. The results are in line with our theoretical derivation.