Deep Active Learning with Manifold-Preserving Trajectory Sampling

Active learning is a pivot research area in machine learning, focused on optimizing data annotations to enhance model performance with fewer labelled samples. Most AL methods mainly consider uncertainty as a crucial criterion to intelligently sample data that improves model’s generalization. Such methods prioritize data points with high prediction variance or near the decision boundary, employing techniques like MC-Dropout [2], Query-by-Committee (QBC) [6], and adversarial training [7] to address overconfident deep neural networks [8, 9]. Influence-based AL approaches select data points based on their estimated impact on model performance, using schemes like Learning Loss [10], and the Influence Function [11] that leverages gradient to estimate changes in prediction accuracy [12, 13]. Besides, BADGE [14] also aims to select uncertain data by evaluating gradient. Many deep AL methods resort to auxiliary models to estimate data uncertainty. Typical works include VAAL [9] that uses an auxiliary auto-encoder, and GCNAL [15] that employs a graph network as the auxiliary model. Unlike these methods, the Coreset [16] is free of any auxiliary models, but suffering from a slow optimization process (e.g., solving a classical K-center or 0-1 Knapsack problem) during data selection. Several other works, such as [17], rely on complicated training fashion (e.g., adversarial), and it will be challenging if using such methods on a different data format (e.g., 3D medical images of voxels) other than 2D natural images.

Bayesian Neural Networks (BNNs) are designed to provide robust uncertainty estimates by treating the network’s parameters as probabilistic distributions rather than fixed values. This approach is essential for capturing uncertainty in tasks like active learning. Several works [18, 19, 20] propose to estimate posterior distributions by averaging the training checkpoints. To this end, they use Stochastic Weight Averaging (SWA) to perform the averaging operation, improving the uncertainty estimation. These methods offer practical solutions for reliable uncertainty estimation in deep networks. Notably, our method has a very low level of similarity with these methods, as we propose a brand new solution to estimate posterior considering both labeled and unlabeled data simultaneously.

We first introduce the multi-cycle deep active learning problem setting. Starting with a set of unlabeled samples, $\mathcal{U}$, and a labeled set, $\mathcal{L}$, the objective is to select a subset from $\mathcal{U}$ according to a predefined annotation budget. This chosen subset, ${X_{N}}$, is annotated by experts or equivalent sources, resulting in labeled data: $\mathcal{L}\leftarrow\mathcal{L}\cup\{X_{N},Y_{N}\}$, where ${Y_{N}}$ are the labels. The remaining unlabeled data is updated: $\mathcal{U}\leftarrow\mathcal{U}\setminus{X_{N}}$. In the process of subset selection, we typically need to train a model parameterized with ${\theta}$ as $f(\cdot;\theta)$. Considering a classification task with $C$ classes, Entropy can be used to estimate the prediction uncertainty for a given example $x$ as

$$H(x)=-\underset{c\in C}{\Sigma}p(y=c|x,\theta)\log p(y=c|x,\theta).$$

(1)

where $p(y=c|x,\theta)$ corresponds to the softmax probability of the $c$-th class from the prediction $f(x;\theta)$.

For better estimation of the probability, BNN approaches incorporate the posterior distribution of $\theta$ by

$$p(y=c|x,D)=\int p(y=c|x,\theta)p(\theta|D)d\theta.$$

(2)

where $D$ represents the data distribution. While $p(y=c|x,\theta)$ can be easily calculated from the network, the key problems turn into (1) estimating the posterior distribution $p(\theta|D)$ and (2) sampling from the parameter distribution to approximate the integration. Next, we introduce our solutions for them.

Traditional MC-dropout [2] or other BNN based approaches suggest to estimate the posterior distribution $p(\theta|D)$ by training a model with the available labeled set $\mathcal{L}$. Here we analyze its potential risk and propose our method by the following analysis.

Since a direct calculation of $p(\theta|D)$ is infeasible, we apply the Bayes rule as

$$p(\theta|D)=\frac{p(\theta)p(D|\theta)}{p(D)}\propto p(\theta)p(D|\theta).$$

(3)

Regarding $p(D)$ as data sampling, we further decompose the posterior distribution into the multiplication of the prior term $p(\theta)$ and the likelihood $p(D|\theta)$. Since $p(\theta)$ can be realized by common regularization techniques such as weight decay, the flexible task is to estimate $p(D|\theta)$. Given the data distribution $D$, we convert the problem of sampling the most possible parameters into finding those $\theta$ that maximize the likelihood term.

Denoting the input data by $X$, labels by $Y$ and the deep features by $Z$ with the dependency chain as $X\rightarrow Z\rightarrow Y$ in a discriminative model, we further decompose the likelihood as $p(D|\theta)=p(X,Z,Y|\theta)$. Therefore we have

$$p(D|\theta)=p(X|\theta)p(Z|X,\theta)P(Y|X,Z,\theta).$$

(4)

Given the potentially biased labeled set $\mathcal{L}$ as the training set, existing approaches only have the effect of maximizing the conditional distribution term $p(Y|X,Z,\theta)$, which is sub-optimal in maximizing $p(D|\theta)$.

While label information is only available for $\mathcal{L}$, we propose to correct the data bias by regularizing the second term $p(Z|X,\theta)$ with real distribution calculated by $\mathcal{L}\cup\mathcal{U}$. In this way, we enforce the feature space $Z$ learned from biased input $X$ to represent a more accurate manifold. Specifically, denoting $f_{e}$ as the feature extractor part of the neural network $f$, we aim to align the distribution $Z_{\mathcal{L}}=f_{e}(X_{\mathcal{L}};\theta)$ with $Z_{*}\approx f_{e}(X_{\mathcal{L}\cup\mathcal{U}};\theta)$.

To achieve this, we adopt Maximum Mean Discrepancies $\mathrm{MMD}$ [21], which is originally designed to check whether two samples are from the same distribution. Denoting $\mathcal{H}$ as a class of functions $h:\mathcal{X}\to R$, $\mathrm{MMD}$ is defined as

$$\mathrm{MMD}(Z_{\mathcal{L}},Z_{*})=\sup_{h\in\mathcal{H}}\{\mathop{\mathbb{E}}_{z\sim Z_{\mathcal{L}}}[h(z)]-\mathop{\mathbb{E}}_{z\sim Z_{*}}[h(z)]\}.$$

(5)

By specifying $\mathcal{H}$ as a reproducing kernel Hilbert space, Eq. (5) can be estimated by comparing the square distance between the empirical kernel mean embeddings as

$\displaystyle\mathrm{MMD}^{2}(Z_{\mathcal{L}},Z_{*})\approx\|\frac{1}{|Z_{\mathcal{L}}|}\sum_{z\in Z_{\mathcal{L}}}\kappa(z)-\frac{1}{|Z_{\mathcal{L}\cup\mathcal{U}}|}\sum_{z\in Z_{\mathcal{L}\cup\mathcal{U}}}\kappa(z)\|^{2}.$

(6)

Here we follow existing transfer learning literature by expanding the quadratic term and calculate the kernel functions corresponding to $\kappa$, for which a Gaussian radial basis function (RBF) is adopted. In practice, two batch of examples randomly sampled from $\mathcal{L}$ and $\mathcal{L}\cup\mathcal{U}$ are used to minimize the MMD distance in each training iteration.

To obtain optimal parameters, we simultaneously minimize the standard supervised error $L_{ce}$ (e.g., cross-entropy loss) and the manifold preserving loss as

$$L=L_{ce}(X_{\mathcal{L}})+\lambda\mathrm{MMD}^{2}(Z_{\mathcal{L}},Z_{*}).$$

(7)

We use the optimization trajectory to sample a diverse set of checkpoints by leveraging the path of minimizing Eq. (7) with stochastic gradient descent (SGD). This approach intends to explore different regions of the parameter space that are likely to have a high posterior probability as in Eq. (3). Specifically, we adopt the technique proposed by Stochastic Weight Averaging (SWA) for parameter sampling. By using a cyclic learning rate after the half convergence, a set of checkpoints $\{\theta_{t_{i}}\}_{i=1}^{n}$ are sampled and saved.

Finally, we calculate the prediction probability by averaging the prediction of each sampled parameter as

$$p(y=c|x,D)=\frac{1}{n}\Sigma_{i=1}^{n}p(y=c|x,\theta_{t_{i}}),$$

(8)

where each $\theta_{t_{i}}$ is sampled in the trajectory of minimizing Eq. (7). We substitute Eq. (8) into Eq. (1) to calculate the uncertainty for each unlabeled example from $\mathcal{U}$.

Note that the manifold preserving and parameter averaging approaches are designed only to select uncertain examples in our active learning framework, despite that SWA is more popular as a strategy to improve model accuracy.

Datasets. As most deep AL methods have been evaluated in computer vision challenges, we also adopt four widely used benchmark vision datasets to evaluate our method, including MNIST [22], CIFAR10 [23], SVHN [24], and Mini-ImageNet [25]. In addition, we also incorporate two typical non-vision datasets for the evaluation, namely OpenML-6 [26] and OpenML-155 [26], which are tabular datasets from the OpenML repository, including structured data with mixed types of features.

For each dataset, following the common practice in AL literature, we randomly select a small portion of data as initial samples and annotate them. The number of such such samples is 100 for all the datasets, except the Mini-ImageNet in which we use 1000 initial samples. Then in each AL round (covering both model training and data selection phases), we select 100 unlabeled samples (labeling budget) for all the datasets, except the Mini-ImageNet where we select 1000 unlabeled samples for initial annotation. When MLP or CNN is used as the backbone, within each AL round, we train the model for 100 epochs. When pre-trained ViT is used, we fine-tune it for 1000 epochs within each AL round. We adopt a learning rate of ${1e-3}$ for vision datasets and ${1e-4}$ for non-vision datasets. The batchsize is set to 64 for all the experiments. Notably, we train the MLP and CNN from scratch, whereas we fine-tune the pre-trained ViT following the practice in [31] for a fair comparison. To reduce randomness, we repeat each experiment for 5 times and average the results as the final one.

As shown in Figure 1, across the multiple datasets widely used in computer vision scenarios, the proposed method outperforms the others in terms of the classification accuracy. In addition to this overall evaluation, we have the following observations. First, our method yields solid results when different backbone models are employed, i.e., MLP for MNIST, CNN for CIFAR10 and SVHN, and ViT for Mini-ImageNet. This observation demonstrates the strong adaptability to the mainstream deep learning models. Second, when labeling budget is limited, such that only 100 samples can be annotated at a time (i.e., MNIST, CIFAR10, and SVHN), our method gives consistent better performance than the others, demonstrating its potential in real-world scenarios where labeling cost could be extremely high (e.g., in medical domain). Third, compared to the improvement in MNIST, our method outperforms the others by a more significant margin in CIFAR10 and SVHN, indicating the capability of our method on handling data of complex scene (i.e., CIFAR10 and SVHN have more complex scene than MNIST). Fourth, when labeling budget is relatively high (i.e., 1000 in Mini-ImageNet whereas 100 in the others), our method still shows consistent superiority over the others. This further demonstrates the strong capability of our method since most AL peers fail to show their superiority when labeling budget is high. Last, compared to the AL method that is based on auxiliary models, i.e., GCNAL [15] that uses an auxiliary graph network for data selection in addition to the task model, our method yields more appealing results, suggesting that developing deep AL methods without leveraging auxiliary models could be a promising effort.

In addition, as illustrated in Figure 2 we surprisingly find that our method significantly outperforms the others in non-vision datasets, such as OpenML-6 and OpenML-155. As MLP is used as the backbone model rather than CNN for such structured tabular data, the inferior performance yielded by the other methods indicate that MLP is more impacted by data bias, while our method can effectively handle the bias issue. Notably, in MNIST where MLP is also used (see Fig. 1), the performance gaps between our method and the others are not that big since MNIST is a fairly simple dataset that cannot make the methods show their upper bounds.