GALAXY: Graph-based Active Learning at the Extreme

Training deep learning systems can require enormous amounts of labeled data. Active learning aims to reduce this burden by sequentially and adaptively selecting examples for labeling, with the goal of obtaining a relatively small dataset of especially informative examples. The most common approach to active learning is uncertainty sampling. The idea is to train a model based on an initial set of labeled data and then to select unlabeled examples that the model cannot classify with certainty. These examples are then labeled, the model is re-trained using them, and the process is repeated. Uncertainty sampling and its variants can work well when the classes are balanced. However, in many applications datasets may be very unbalanced, containing very rare classes or one very large majority class. As an example, suppose an insurance company would like to train an image-based machine learning system to classify various types of damage to the roofs of buildings (Conathan et al., 2018). It has a large corpus of unlabeled roof images, but the vast majority contain no damage of any sort.

Unfortunately, under extreme class imbalance, uncertainty sampling tends to select examples mostly from the dominant class(es), often leading to very slow learning. In this paper, we take a novel approach specifically targeting the class imbalance problem. Our method is guaranteed to select examples that are both uncertain and class-diverse; i.e., the selected examples are relatively balanced across the classes even if the overall dataset is extremely unbalanced. In a nutshell, our algorithm sorts the examples by their softmax uncertainty scores and applies a bisection procedure to find consecutive pairs of points with differing labels. This procedure encourages finding uncertain points from a diverse set of classes. In contrast, uncertainty sampling focuses on sampling around the model’s decision boundary and therefore will collect a biased sample if this model decision boundary is strongly skewed towards one class. Figure 1 displays the results of one of our experiments, showing that our proposed GALAXY algorithm learns much more rapidly and collects a significantly more diverse dataset than uncertainty sampling.

We make the following contributions in this paper:

• •

  we develop a novel, scalable algorithm GALAXY, tailored to the extreme class imbalance setting, which is frequently encountered in practice,

• •

  GALAXY is easy to implement, requiring relatively simple modifications to commonplace uncertainty sampling approaches,

• •

  we conduct extensive experiments showing that GALAXY outperforms a wide collection of deep active learning algorithms in the imbalanced settings, and

• •

  we give a theoretical analysis showing that GALAXY selects much more class-balanced batches of uncertain examples than traditional uncertainty sampling strategies.

Deep Active Learning: There are two main algorithmic approaches in deep active learning: uncertainty and diversity sampling. Uncertainty sampling queries the unlabeled examples that are most uncertain. Often, uncertainty is quantified by distance to the decision boundary of the current model (e.g., (Tong & Koller, 2001; Kremer et al., 2014; Balcan et al., 2009)). Several variants of uncertainty sampling have been proposed for deep learning (e.g., (Gal et al., 2017; Ducoffe & Precioso, 2018; Beluch et al., 2018).

In a batch setting, uncertainty sampling often leads to querying a set of very similar examples, giving much redundant information. To deal with this issue, diversity sampling queries a batch of diverse examples that are representative of the unlabeled pool. Sener & Savarese (2017) propose a Coreset approach for diversity sampling for deep learning. Others include (Gissin & Shalev-Shwartz, 2019; Geifman & El-Yaniv, 2017). However, under the class imbalance scenarios where collecting minority class examples is crucial, previous work by Coleman et al. (2020) has shown such methods to be less effective as they are expected to collect a subset with equal imbalance to the original dataset.

Recently, significant attention has been given to designing hybrid methods that query a batch of informative and diverse examples. Ash et al. (2019) balances uncertainty and diversity by representing each example by its last layer gradient and aiming to select a batch of examples with large Gram determinant. Citovsky et al. (2021) uses hierarchical agglomerative clustering to cluster the examples in the feature space and then cycles over the clusters querying the examples with smallest margin. Finally, Ash et al. (2021) uses experimental design to find a batch of diverse and uncertain examples.

Class Imbalance Deep Active Learning: A number of recent works have studied active learning in the presence of class imbalance. Coleman et al. (2020) proposes SEALS, a method for the setting of class imbalance and an enormous pool of unlabeled examples. Kothawade et al. (2021) proposes SIMILAR, which picks examples that are most similar with the collected in-distribution examples and most different from the known out-of-distribution examples. Their method achieves this by maximizing the conditional mutual information. Our setting is closest to their out-of-distribution imbalance scenario. Finally, Emam et al. (2021) tackles the class imablance issue by proposing BASE which queries the same number of examples per each predicted class based on margin from decision boundary, where the margin is defined by the distance to the model boundary in the feature space of the neural network.

By contrast with the above methods, our method searches adaptively in the output space within each batch for the best threshold separating two classes and provably produces a class-balanced set of labeled examples. Adaptively searching for the best threshold is especially helpful in the extreme class imbalance setting where the decision boundary of the model is often skewed. If all labels were known, theoretically it may be possible to modify the training algorithm to obtain a model without any skew towards on class. However, this is not possible in active learning where we do not know the labels a priori. In addition, it is expensive and undesirable in practical cases to modify the training algorithm (Roh et al., 2020), making it attractive to work for any off-the-shelf training algorithm (like ours).

Graph-based active learning: There have been a number of proposed graph-based adaptive learning algorithms (e.g., (Zhu et al., 2003a, b; Cesa-Bianchi et al., 2013; Gu & Han, 2012; Dasarathy et al., 2015; Kushnir & Venturi, 2020). Our work is most closely related to (Dasarathy et al., 2015), which proposed $S^{2}$ a graph-based active learning with strong theoretical guarantees. While $S^{2}$ assumes a graph as an input and will perform badly on difficult graphs, our work builds a framework that combines the ideas of $S^{2}$ with deep learning to perform active learning while continually improving the graph. We review this work in more detail in Section 4.

We investigate the pool-based batched active learning setting, where the learner has access to a large pool of unlabeled data examples $X=\{x_{1},x_{2},...,x_{N}\}$ and there is an unknown ground truth label function $f^{\star}:X\rightarrow\{1,2,...,K\}$ giving the label of each example. At each iteration $t$, the learner selects a small batch of $B$ unlabeled examples $\{x_{i}^{(t)}\}_{i=1}^{B}\subseteq X$ from the pool, observes its labels $\{f^{\star}(x_{i}^{(t)})\}_{i=1}^{B}$, and adds the examples to $L$, the set of all the examples that have been labeled so far. After each batch of examples is queried, the learner updates the deep learning model training on all of the examples in $L$ and uses this model to inform which batch of examples are selected in the next round.

We are particularly interested in the extreme class imbalance problem where one class is significantly larger than the other classes. Mathematically,

where $N_{k}=|\{i:f^{\star}(x_{i})=k\}|$ denote the number of examples that belong to the $k$-th class. Here, $\epsilon$ is some small class imbalance factor and each of the in-distribution classes $1$, …, $K-1$ contains far fewer examples than the out-of-distribution $K$-th class. This models scenarios, for example in roof damage classification, self-driving and medical applications, in which a small fraction of the unlabeled examples are from classes of interest and the rest of the examples may be lumped into an “other” or “irrelevant” category.

Henceforth, we denote $f=\mathcal{A}(L)$ to be a model trained on the labeled set $L$, where $\mathcal{A}:\mathcal{L}\rightarrow\mathcal{F}$ is a training algorithm that trains on a labeled set and outputs a classifier.

To begin, we introduce some notation. With slight abuse of notation, we define a undirected graph over the pool as $G=(X,E)$ with vertex set $X$ and edge set $E$, where each node $x$ in the graph is also an example in the pool $X$. Let $P_{ij}(X,E)\subset E$ denote the shortest path connecting $x_{i}$ and $x_{j}$ in the graph $G=(X,E)$, and let $|P_{ij}(X,E)|$ denote its length.¹¹1In the special case when $x_{i}$ and $x_{j}$ are not connected, we define $|P_{ij}(X,E)|=\infty$.

Dasarathy et al. (2015) proposed a graph-based active learning algorithm $S^{2}$ (see Algorithm 1) that aims to identify all the cuts $C=\{(x,y)\in E:f^{\star}(x)\neq f^{\star}(y)\}$, namely every edge that connects an oppositely labeled pair of examples. In particular, if one labeled all of the examples in the cut boundary $\partial C=\{x\in X:\exists e\in C,x\in e\}$, one would be able to classify every example in the pool correctly. As an example, take the linear graph in Figure 2(a) where each node represents a numbered example and is associated with a binary label (black/white). It is thus necessary to query at least seven examples to identify the cut boundary and therefore the labeling.

$S^{2}$ performs an alternating two phased procedure. First, if every pair of connected examples have the same label, $S^{2}$ queries an unlabeled example uniform at random. Second, whenever there exist paths connecting examples with different labels, the algorithm bisects along the shortest among these paths until it identifies a cut. The algorithm then removes the identified edge from the graph.

Dasarathy et al. (2015) has shown that the PAC sample complexity to identify all cuts highly depends on the input graph’s structural properties. As an example consider Figure 2, which depicts several graphs on the same set of examples. In graph (a), one needs to query at least seven examples, while in graph (b) one need only query two examples. Indeed, such a difference can be made arbitrarily large. The work of Dasarathy et al. (2015), however, did not address the major problem of how to obtain an “easier” graph that requires fewer examples queries for active learning.

Our algorithm GALAXY shown in Algorithm 4 blends graph-based active learning and deep active learning through the following two alternating steps

• •

  Given a trained neural network, we construct a graph based on the neural network’s predictions and apply a modified version of $S^{2}$ to it, efficiently collecting an informative batch of labels (Algorithm 4).

• •

  Given a new batch of labeled examples, we train a better neural network model that will be used to construct a better graph for active learning (Algorithm 2).

To construct graphs from a learned neural network in multi-class settings, we take a one-vs-all approach on the output (softmax) space as shown in Algorithm 2. For each class $k$, we build a linear graph $G^{(k)}$ by ranking the model’s confidence margin $\delta^{(k)}_{i}$ on each example $x_{i}\in X$. For a neural network $f_{\theta}$, the confidence margin is simply defined as $\delta^{(k)}_{i}=[f_{\theta}(x_{i})]_{k}-\max_{k^{\prime}}[f_{\theta}(x_{i})]_{k^{\prime}}$, where $[\cdot]_{k}$ denotes the $k$-th element of the softmax vector. We break ties by the confidence scores $[f_{\theta}(x_{i})]_{k}$ themselves (equivalently by $\max_{k^{\prime}}[f_{\theta}(x_{i})]_{k^{\prime}}$). Intuitively, for each graph $G^{(k)}$, we sort examples according to their likelihood to belong to class $k$. Indeed, when $f_{\theta}$ is a perfect classifier on the pool, each linear graph constructed behaves like Figure 2(b), i.e., every example in class $k$ is perfectly separated from all other classes with only one cut in between.

Our algorithm GALAXY shown in Algorithm 4 proceeds in a batched style. For each batch, GALAXY first trains a neural network to obtain graphs constructed by the procedure described above. It then performs $S^{2}$ style bisection-like queries on all of the graphs but with two major differences.

• •

  To accommodate multiple graphs, we treat each linear graph $G^{(k)}$ as a binary one-vs-all graph, where we gather all shortest paths $P_{ij}(X,E^{(k)})$ that connects a queried example in class $k$ and a queried example in any other classes. If such shortest paths exist, we then find the shortest of these shortest paths across all $k\in[K]$ and the bisect the resulting shortest shortest path like in $S^{2}$.

• •

  When no such shortest path exists, instead of querying an example uniform at random as in $S^{2}$, we increase the order of the graphs by Algorithm 3 and perform bisection procedures on the updated graphs. Here, we refer to an $m$-th order linear graph where each example is connected to all of its neighboring $m$ examples from each side. For example, Figure 2(c) shows a graph of order $2$ as opposed to an order $1$ graph shown in Figure 2(b). Intuitively, bisecting after the Connect procedure is equivalent with querying around the discovered cuts. For example in the case of Figure 2(b), after querying examples $5$ and $6$, our algorithm will connect second order edges and query exactly examples $1$ and $3$ as the next two queries.

We conduct experiments under $8$ different class imbalance settings. These settings are generated from three image classification datasets with various class imbalance factors. If the classes are balanced in the dataset, then most active learning strategies (including GALAXY) perform similarly, with relative small differences in performance, so we focus our presentation on unbalanced situations. We will first describe the setups (Section 7.1) before turning to the results in Section 7.2. Finally, we present a comparison with vanilla $S^{2}$ algorithm and demonstrate the importance of reconstructing the graphs in Section 7.3.²²2Code can be found in https://github.com/jifanz/GALAXY.

In this paper, we propose a novel graph-based approach to deep active learning that particularly targets the extreme class imbalance cases. We show that our algorithm GALAXY outperforms all existing methods by collecting a mixture of balanced yet uncertain examples. GALAXY runs on similar time complexity as other uncertainty based methods by retraining the neural network model only after each batch. However, it still requires sequential and synchronous labelling within each batch. This means the human labelling effort cannot be parallelized by multiple annotators. For future work, we would like to incorporate asynchronous labelling and investigate its effect on our algorithm.

We thank Andrew Wagenmaker for insightful discussions. This work has been supported in part by NSF Award 2112471.