Gone Fishing: Neural Active Learning with Fisher Embeddings

1 Introduction . The active learning paradigm considers a sequential , supervised learning scenario in which unlabeled samples are abundant but label acquisition is costly . At each round of active learning , the agent fits its parameters using available labeled data before selecting a batch of unlabeled samples to be labeled and integrated into its training set . A well-chosen batch of samples is one that is maximally informative to the learner , such that it can obtain the best hypothesis possible given a fixed labeling budget . Active learning is well established as an area of machine learning research due to the ubiquity of important real world problems that fit the sample-abundant , label-expensive setting ; commonly cited applications range from medical diagnostics [ 1 , 2 ] to image labeling [ 3 ] . Mitigating large sample complexity requirements is particularly relevant for deep neural networks , which have in recent years achieved impressive success on a wide array of tasks but often require considerable amounts of labeled data . Shifting the focus of active learning to deep neural networks highlights several important problems . For one , most foundational active learning work assumes a convex setting , which is clearly violated by massive nonlinear neural networks . Many of these approaches are computationally expensive , and it is not clear how to adapt them for real-world use [ 4 ] . Further , because neural network training is generally expensive , practical active learning algorithms must be able to work in the batch regime , querying B samples at each round of active learning instead of a single point at a time [ 5 ] . Despite a long history of active learning research , these constraints draw attention to a need for practical , principled batch active learning algorithms for neural networks . Current state-of-the-art methods , like Batch Active Learning by Diverse Gradient Embeddings ( BADGE ) , perform robustly in experiments , but explanations for its behavior are fairly limited [ 6 ] . This drawback makes it unclear how to scale some active learning algorithms into regimes that deviate somewhat from the setting for which they were designed—BADGE , for example , can not be run on regression problems , and as we show in this paper , performs poorly when used in conjunction with a convex model . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . This article adopts a probabilistic perspective of neural active learning . We view neural networks as specifying a conditional probability distribution p ( y | x , ✓ ) over label space Y given example X , where ✓ are the network parameters . This perspective provides theoretical inspiration from the convex regime with which to examine and design neural active learning algorithms . From this viewpoint we motivate and revisit a classic , Fisher-based objective for idealized active selection . We argue that approximately minimizing this objective can be done tractably in the neural regime , despite their overparametrized structure and shifting internal representation . Accordingly , this work helps bridge the divide between algorithms that are performant but not well understood by theory , and those that are theoretically transparent but not computationally tractable . Experimentally , BAIT offers improved performance over baselines in deep classification problems , a trend that is robust to experimental conditions like batch size , model architecture , and dataset . Crucially , BAIT is general purpose , and can be easily extended to regression settings , where many other algorithms can not . It further performs well on both regression and classification with convex models , a paradigm in which other algorithms often struggle . In summary , this paper • puts neural active learning on firm probabilistic grounding , giving a new , rigorous perspective on the functionality of previously proposed algorithms . • provides in-depth empirics that elucidate differences between neural and convex regimes , and discusses simplifying assumptions that are sometimes reasonable in the neural case . • proposes a practical , unifying , high-performing active learning algorithm that leverages these insights in a computationally tractable manner . 2 Related work . Active learning is a very well-studied problem [ 7–9 ] . There are two main sample selection approaches , diversity and uncertainty sampling , which are successful respectively for large and small batch sizes . Diversity sampling strategies aim to select batches of data that best represent the space . In a deep learning context , these algorithms typically embed unlabeled samples using the neural network ’ s penultimate layer and select a subset of samples that might act as a proxy for the entire dataset [ 10 , 11 ] . [ 12 ] proposed inducing batch diversity using a generative adversarial network formulation , selecting samples that are maximally indistinguishable from the pool of unlabeled examples . There is also a rich body of work on batch active learning [ 13–17 ] . These methods typically formulate batch selection as an optimization that minimizes an upper-bound on some notion of model loss . Efficiently adapting parameters to an incrementally larger training set is not an issue in convex settings . Accordingly , with linear models , it is more common to use uncertainty sampling and a batch size of one . A frequently used approach is to query samples that lie closest to the current model ’ s decision boundary , a quantity that ’ s considered inversely proportional to uncertainty [ 18–20 ] . Some similar methods offer theoretical guarantees on statistical consistency [ 9 , 21 ] . Other algorithms quantify uncertainty using the entropy of the predicted distribution over classes , or as the size of the expected gradient induced by observing the label corresponding with a candidate sample [ 22 ] . The latter is known to be related to the T -optimality criterion in experimental design , but is unable to account for batch diversity . Similar approaches have been modified for use with neural networks as well . For example , [ 23 ] exercise Dropout to sample weights to approximate the posterior distribution over labels , and use it to identify samples that reduce model uncertainty . Adversarial example generation has been used to approximate the distance between a sample and the decision boundary [ 24 ] . Model ensembling has also been used to approximate sample uncertainty , where the predictive variance across constituent models can be used to inform a sample selection strategy [ 25 ] . There are a variety of algorithms that are meant to combine uncertainty and diversity sampling [ 26 ] . This trade-off is sometimes framed as its own optimization problem , for example using a metalearning approach that hybridizes both strategies [ 27 , 28 ] . Among these is active learning by learning , which uses a bandit approach to select which query rule to employ at any given round of active learning [ 28 ] . BADGE , described in detail in Section 5.1 , also combines uncertainty and diversity sampling , and is considered state-of-the-art for deep neural networks . 3 Notation and Setup . We consider a standard setup for batch active learning with neural network models , where there is an instance space X , label space Y , and a distribution D over X ⇥ Y . We use DX to denote the marginal distribution over the instance space and DY|X ( x ) to denote the conditional distribution over labels given example x . For learning , we are given access to a pool U = { xi } ni=1 ⇠ DX of unlabeled examples and we have the ability to request the label for any point x 2 U . In the tth round of batch active learning , we select a collection { x ( t ) j } Bj=1 ⇢ U of B examples ( B is the batch size ) and request the labels y ( t ) j ⇠ DY|X ( x ( t ) j ) for all examples in the batch . We use these labeled examples to update our neural network model and then we proceed to the next round . In this setup , the ultimate objective is to achieve low loss on the data distribution D , that is we hope our learned parameters b✓ nearly minimize E ( x , y ) ⇠D ` ( x , y ; b✓ ) , where ` is some loss function like the cross entropy loss for classification . We always consider this objective in our experiments , but for algorithm development it is helpful to instead consider the fixed-design or transductive setting , where the goal is to instead minimize LU ( ✓ ) = Ex⇠UEy⇠DY|X ( x ) ` ( x , y ; b✓ ) , essentially treating the unlabeled samples U as the entire distribution . Note that these two objectives can typically be related by generalization arguments . 4 Probabilistic Perspective . We consider neural networks as specifying a probability distribution p ( y | x , ✓ ) over the label space Y given an example x , where ✓ are the network parameters . Adopting this view , it is most natural to use the loss function ` ( x , y ; ✓ ) = log p ( y | x , ✓ ) , choosing parameters that maximize the likelihood of observed labeled data . In classification problems , for example , we apply the softmax operation to the output of the network and then evaluate the cross-entropy loss with the ground truth label . For regression problems , we use the square loss , which treats the neural network as specifying a Gaussian distribution for each x. Bayesian linear regression . As a warm-up , it is illustrative to consider an experimental design setting with Bayesian linear regression . We consider a d-dimensional linear regression problem where we assume the parameter vector ✓ ? has prior distribution N ( 0 , 1I ) and the conditional distribution DY|X ( x ) = p ( · | x , ✓ ? ) = N ( h✓ ? , xi , 2 ) is Gaussian . For any set of labeled data { xj , yj } mj=1 , the resulting maximum a posteriori ( MAP ) estimate is given by ridge regression with regularizer 2 : b✓ = argmin ✓ mX j=1 ( hxj , ✓i yj ) 2 + 2k✓k22 ( 1 ) In the experimental design setting , we have unlabeled data U = { xi } ni=1 and our goal is to select a set S ⇢ U of B points so that the resulting MAP estimate has the lowest Bayes risk . Letting ⌃ = 1 n Pn i=1 xix > i denote the second moment matrix of the unlabeled data , the Bayes risk is BayesRisk ( S ) = E h ( b✓S ✓ ? ) > ⌃ ( b✓S ✓ ? ) i , ( 2 ) where b✓S is the MAP estimate after querying for labels on subset S , and the expectation is with respect to the noise in the labels and the prior over ✓ ? . Lemma 1 in the Appendix shows that for a subset S ⇢ U , letting ⇤S = P x2S xx > + 2I , the Bayes risk in this setting is exactly : BayesRisk ( S ) = 2 tr ( ⇤ 1S ⌃ ) . ( 3 ) Observe that the RHS does not depend on labels , implying that minimizing the RHS over subsets S is feasible and optimal selection strategy under this criteria . This also verifies that multiple batches of active learning are not required for Bayesian linear regression , although this observation does not carry forward to the neural setting . BAIT is designed to approximately minimize this objective . Classical regime . An objective similar to ( 3 ) also emerges naturally in the analysis of maximum likelihood estimators ( MLE ) in the convex regime . Here , classical statistical theory posits that the model is well-specified , so that there is some parameter ✓ ? such that DY|X ( x ) = p ( · | x , ✓ ? ) for each x 2 X . It is also common to impose regularity conditions including strong convexity of the loss function LU ( ✓ ) [ 29 , 30 ] . While these conditions certainly do not hold in the neural setting , BAIT builds on much of this classical technology . The key quantity is the Fisher information matrix I ( x ; ✓ ) : = Ey⇠p ( ·|x , ✓ ) r2 ` ( x , y ; ✓ ) which is known to determine the asymptotic distribution of the maximum likelihood estimator [ 29 ] . In many probabilistic models , including linear and logistic regression , the hessian of the loss function does not depend on the label y , which we assume going forward . For active learning in the classical setup , [ 4 ] give a two-phase sampling scheme with provably near-optimal performance . In the first phase , the algorithm samples a batch of B points uniformly at random , requests their labels , and optimizes the log-likelihood to obtain an initial estimate ✓1 . In the idealized version of the second phase , a batch of B points is chosen to optimize argmin S⇢U , |S|B tr 0 @ X x2S I ( x ; ✓1 ) ! 1 IU ( ✓1 ) 1 A ( 4 ) where IU ( ✓1 ) is the Fisher over all samples , P x2U I ( x ; ✓1 ) . This combinatorial problem is intractable in general , so [ 4 ] instead solve a semidefinite relaxation ( SDP ) . They request labels on the obtained batch B of points and re-fit the model to obtain the final estimate b✓ . For their setup , they prove the statistical properties of this two-phase estimator are near optimal . Despite the theoretical properties , solving an SDP is not feasible in high dimensions ; instead we provide a new , greedy algorithm for minimizing the objective that is usable in the neural regime . Several other works have also looked at this objective , either from a purely theoretical perspective [ 4 , 31 ] or via relaxations [ 32–34 ] . However , some of these do not ensure batch diversity , and none have been extended to the neural regime . This formulation essentially generalizes ( 3 ) , since in linear regression the Fisher information I ( x ; ✓ ) is the covariance matrix xx > / 2 : the objectives in ( 3 ) and ( 4 ) differ only in their use of the regularizer controlled by . That is , essentially the same objective can be derived from two different perspectives , making it a compelling object for active selection . As such , our starting point for the neural setting is the ideal-but-intractable optimization problem in ( 4 ) .

Summary: The paper introduces an algorithm (BAIT) for batch active learning, designed for and tested on neural networks. It centres around an objective in eq 4, argmin_select trace(I_select^-1, I_all), with I as the Fisher information matrix of the final layer parameters of a neural network, an objective that has been previously derived/analysed in simpler settings. The contribution here is applying it to deep neural networks. It draws links with a current SOTA method, BADGE, and shows slight performance benefits over it.