Detecting Extrapolation with Local Ensembles

1 INTRODUCTION . As machine learning is deployed in increasingly vital areas , there is increasing demand for metrics that draw attention to potentially unreliable predictions . One important source of unreliability is extrapolation . Extrapolation can be formalized in a number of ways : it can refer to making predictions on inputs outside the support of the training data , making predictions with high Bayesian or Frequentist uncertainty , or making predictions that depend strongly on arbitrary choices outside of the learning problem specification ( e.g. , a random seed ) . In this paper , we develop a method for detecting this last form of extrapolation . Specifically , we say that a trained model is extrapolating on a test input if the prediction at this input is underdetermined — meaning that many different predictions are all equally consistent with the constraints posed by the training data and the learning problem specification ( i.e. , the model architecture and the loss function ) . Underdetermination is just one form of extrapolation , but it is particularly relevant in the context of overparameterized model classes ( e.g . deep neural networks ) . Recently , simple ( but computationally expensive ) ensembling methods ( Lakshminarayanan et al. , 2017 ) , which train many models on the same data from different random seeds , have proven highly effective at uncertainty quantification tasks ( Ovadia et al. , 2019 ) . This suggests that underdetermination is a key threat to reliability in deep learning , and motivates flexible methods that can detect underdetermined predictions cheaply . With this motivation , we present local ensembles , a post-hoc method for measuring the extent to which a pre-trained model ’ s prediction is underdetermined for a particular test input . Given a trained model , our method returns an extrapolation score that measures the variability of test predictions across a local ensemble , i.e . a set of local perturbations of the trained model parameters that fit the training data equally well . Local ensembles are a computationally cheap , post-hoc alternative to fully trained ensembles , and do not require special training procedures of approximate ensembling methods that measure related , but distinct , notions of uncertainty ( Gal & Ghahramani , 2015 ; Blundell et al. , 2015 ) . Local ensembles also address a gap in approximate methods for estimating prediction uncertainty . Specifically , whereas exact Bayesian or Frequentist uncertainty includes underdetermination as one component , approximate methods such as Laplace approximations ( MacKay , 1992 ) or influence function-based methods ( Schulam & Saria , 2019 ) break down when underdetermination is present . ∗Some of this work was done while this author was an intern at Google Brain . In contrast , our method leverages the pathology that makes these methods struggle ( an ill-conditioned Hessian ) . Our contributions in this paper are as follows : • We present local ensembles , a test-time method for detecting underdetermination-based extrapolation in overparameterized models . • We demonstrate theoretically that our method approximates the variance of a trained ensemble with local second-order information . • We give a practical method for tractably approximating this quantity , which is simpler and cheaper than alternative second-order reliability methods . • Through experiments aimed at testing underdetermination , we show our method approximates the behavior of trained ensembles , and can detect extrapolation in a range of scenarios . 2 EXTRAPOLATION SCORE AND LOCAL ENSEMBLES . 2.1 SETUP . Let z = ( x , y ) be an example input-output pair , where x is a vector of features and y is a label . We define a model in terms of a loss function L with parameters θ as a sum over training examples ( zi ) n i=1 , i.e. , L ( θ ) = ∑n i ` ( zi , θ ) , where ` is an example-wise loss ( e.g. , mean-squared error or cross entropy ) . Let θ ? be the parameters of the trained model , obtained by , e.g. , minimizing the loss over this dataset , i.e. , θ ? = arg minθ L ( θ ) . We write the prediction function given by parameters θ at an input x as ŷ ( x , θ ) . We consider the problem of auditing a trained model , where unlabeled test points x′ arrive one at a time in a stream , and we wish to assess extrapolation on a point-by-point basis . In this section , we introduce our local ensemble extrapolation score Em ( x′ ) for an unlabeled test point x′ ( significance of m explained below ) . The score is designed to measure the variability that would be induced by randomly choosing predictions from an ensemble of models with similar training loss . Our score has a simple form : it is the norm of the prediction gradient gθ ? ( x′ ) : = ∇θŷ ( x′ , θ ? ) multiplied by a matrix of Hessian eigenvectors spanning a subspace of low curvature Um ( defined in more detail below ) . Em ( x′ ) = ‖U > mgθ ? ( x′ ) ‖2 ( 1 ) Here , we show that this score is proportional to the standard deviation of predictions across a local ensemble of models with near-identical training loss , and demonstrate that this approximation holds in practice . 2.2 DERIVATION . Our derivation proceeds in two steps . First , we define a local ensemble of models with similar training loss , then we state the relationship between our extrapolation score and the variability of predictions within this ensemble . The spectral decomposition of the Hessian Hθ ? plays a key role in our derivation . Let Hθ ? = UΛU > , ( 2 ) where U is a square matrix whose columns are the orthonormal eigenvectors of Hθ ? , written ( ξ ( 1 ) , · · · , ξ ( p ) ) , and Λ is a square , diagonal matrix with the eigenvalues of Hθ ? , written ( λ ( 1 ) , · · · , λ ( p ) ) , along its diagonal . As a convention , we index the eigenvectors and eigenvalues in decreasing order of the eigenvalue magnitude . To construct a local ensemble of loss-preserving models , we exploit the fact that eigenvectors with large corresponding eigenvalues represent directions of high curvature , whereas eigenvectors with small corresponding eigenvalues represent directions of low curvature . In particular , under the assumption that the model has been trained to a local minimum or saddle point , parameter perturbations in flat directions ( those corresponding to small eigenvalues λ ( j ) ) do not change the training loss substantially ( see Fig . 1 ) . 1 We characterize this subspace by the span of eigenvectors with corresponding small eigenvalues . Formally , let m be the eigenvalue index such that the eigenvalues { λ ( j ) : j > m } , are sufficiently small to be considered “ flat ” . 2 We call the subspace spanned by { ξ ( j ) : j ∈ σ } the ensemble subspace . Parameter perturbations in the ensemble subspace generate an ensemble of models with near-identical training loss . 3 Our score Em ( x′ ) characterizes the variance of test predictions with respect to random parameter perturbations ∆θ that occur in the ensemble subspace . We now show that our extrapolation score Em ( x′ ) is , to the first order , proportional to the standard deviation of predictions at a point x′ . Proposition 1 . Let ∆θ be a random parameter perturbation with mean zero and covariance proportional to the identity · I , projected into the ensemble subspace spanned by { ξ ( j ) : j > m } . Let P∆ ( x′ ) be the linearized change in a test prediction at x′ induced by perturbing the learned parameters θ ? by ∆θ : P∆ ( x ′ ) : = gθ ? ( x ′ ) > ∆θ ≈ ŷ ( x′ , θ ? + ∆θ ) − ŷ ( x′ , θ ? ) . Then Em ( x′ ) = −1/2 · SD ( P∆ ( x′ ) ) . Proof . First , we characterize the distribution of the loss-preserving parameter perturbation ∆θ . Let Um be the matrix whose columns { ξ ( j ) : j > m } span the ensemble subspace . Then UmU > m is a projection matrix that projects vectors into the ensemble subspace , so the projected perturbation ∆θ has covariance · UmU > m . From this , we derive the variance of the linearized change in prediction V ar ( P∆ ) = · gθ ? ( x′ ) > UmU > mgθ ? ( x′ ) = · ‖U > mgθ ? ( x′ ) ‖22 = · Em ( x′ ) 2 . 2.5 2.0 1.5 1.0 0.5 0.0 Ensemble Standard Deviation 1.5 1.0 0.5 0.0 0.5 1.0 1.5 Ex tra po la tio n Sc or e R=0.87 Figure 2 : Mean extrapolation score vs. ensemble prediction standard deviation on WineQuality dataset . Shows line of best fit and Pearson coefficient R. Both axes log-scaled . We test this hypothesized relationship on several tabular datasets . We train an ensemble of twenty neural networks using the same architecture and training set , varying only the random seed . Then , for each model in the ensemble , we calculate our local ensemble extrapolation score Em ( x ) for each input x′ in the test set ( see Sec . 3 for details ) . For each x′ , we compare , across the ensemble , the mean value of Em ( x′ ) to the standard deviation ŷ ( x′ ) . In Fig . 2 , we plot these two quantities against each other for one of the datasets , finding a nearly linear relationship . On each dataset , we found a similar , significantly linear relationship ( see Table 1 and Appendix A ) . We note the relationship is weaker with the Diabetes dataset ; the standard deviations of these ensembles are an order of magnitude higher than the other datasets , indicating much noisier data . Finally , we note that we can obtain similar results for the standard deviation of the loss at test points if we redefine g as the loss gradient rather than the prediction gradient . 3 COMPUTING EXTRAPOLATION SCORES . 1To see this , consider a second-order Taylor approximation to the training loss when the parameters are perturbed by a vector ∆θ : L ( θ ? + ∆θ ) − L ( θ ? ) ≈ 12∆ > θ Hθ ? ∆θ = ∑p j=1 λ ( j ) ( ξ > ( j ) ∆θ ) 2 . Note that the first-order term drops out because we assume θ ? lies at a local minimum or saddle point . Perturbations ∆θ that lie in the subspace spanned by eigenvectors ξ ( j ) whose corresponding eigenvalues λ ( j ) are close to zero contribute negligibly to this sum . 2For now , we take m to be given , and discuss tradeoffs for choosing m in practice in Section 3 . 3We note that not all loss-preserving ensembles will necessarily be local ; however , the existence of some such local ensemble is sufficient to enable underdetermination . Dataset Pearson Boston 0.76 Diabetes 0.50 Abalone 0.76 Wine 0.87 Table 1 : Correlation of extrapolation scores and ensemble std . deviations on 4 datasets . We now discuss a practical method for computing our extrapolation scores . The key operation is constructing the set of eigenvectors that span a suitably loss-preserving ensemble subspace ( i.e . have sufficiently small corresponding eigenvalues ) . Because eigenvectors corresponding to large eigenvalues are easier to estimate , we construct this subspace by finding the top m eigenvectors and defining the ensemble subspace as their orthogonal complement . Our method can be implemented with any algorithm that returns the top m eigenvectors . See below for discussion on some tradeoffs in the choice of m , as well as the algorithm we choose ( the Lanczos iteration ( Lanczos , 1950 ) ) . Our method proceeds as follows . For a given choice of m , we calculate the m eigenvectors of H with the largest eigenvalues . These eigenvectors define an m-dimensional subspace which is the orthogonal complement to the ensemble subspace . We use these eigenvectors to construct a matrix that projects gradients into the ensemble subspace . Specifically , let Um⊥ be the matrix whose columns are these large-eigenvalue eigenvectors { ξ ( j ) : j ≤ m } . Then Um⊥U > m⊥ is the projection matrix that projects vectors into the “ large eigenvalue ” subspace , and I − Um⊥U > m⊥ projects into its complement : the ensemble subspace . Now , for any test input x′ , we take the norm of this projected gradient to compute our score Em ( x′ ) = ∥∥ ( I − Um⊥U > m⊥ ) gθ ? ( x′ ) ∥∥2 . The success of this approach depends on the choice of m. Specifically , the extrapolation score Em ( x′ ) is the most sensitive to underdetermination in the region of the trained parameters if we set m to be the smallest index for which the training loss is relatively flat in the implied ensemble subspace . If m is set too low , the ensemble subspace will include well-constrained directions , and Em ( x′ ) will over-estimate the prediction ’ s sensitivity to loss-preserving perturbations . If m is set too high , the ensemble subspace will omit some under-constrained directions , and Em ( x′ ) will be less sensitive . For models where all parameters are well-constrained by the training data , a suitable m may not exist . This will usually not be the case for deep neural network models , which are known to have very ill-conditioned Hessians ( see , e.g. , Sagun et al. , 2017 ) . Whether a direction is “ well-constrained ” is ultimately a judgment call for the user . One potential heuristic is to consider a small parameter perturbation of a fixed norm in the ensemble subspace , and to set a tolerance for how much that perturbation can change the training loss ; for small perturbations , the change in loss is a linear function of the curvature in the direction of the perturbation , which is upper bounded by λ ( m ) . We use the Lanczos iteration to estimate the top m eigenvectors , which presents a number of practical advantages for usage in our scenario . Firstly , it performs well under early stopping , returning good estimates of the top m eigenvectors after m iterations . Secondly , we can cache intermediate steps , meaning that computing the m+ 1-th eigenvector is fast once we have computed the first m. Thirdly , it requires only implicit access to the Hessian through a function which applies matrix multiplication , meaning we can take advantage of efficient Hessian-vector product methods ( Pearlmutter , 1994 ) . Finally , the Lanczos iteration is simple – it can be implemented in less than 20 lines of Python code ( see Appendix B.1 ) . It contains only one hyperparameter , the stopping value m. Fortunately , tuning this parameter is efficient — given a maximum value M , we can try many values m < M at once , by estimating M eigenvectors and then calculating Em by using the first m eigenvectors . The main constraint of our method is space rather than time — while estimating the first m eigenvectors enables easy caching for later use , it may be difficult to work with these eigenvectors in memory as m and model size p increase . This tradeoff informed our choice of m in this paper ; we note in some cases that increasing m further could have improved performance ( see Appendix E ) . This suggests that further work on techniques for mitigating this tradeoff , e.g . online learning of sparse representations ( Wang & Lu , 2016 ; Wang et al. , 2012 ) , could improve the performance of our method . See Appendix B for more details on the Lanczos iteration .

Summary: This paper presents local ensembles, a method for detecting underdetermination when extrapolating to test points. The authors define an extrapolation score which is used to estimate the standard deviation of predictions at test points. The extrapolation score is chosen to represent the variability in predictions that would be generated by models with similar training loss. By considering the eigenvectors of the Hessian that are associated with minimum eigenvalues the directions of the loss surface with minimal curvature are found, and perturbations of the parameters in the subspace of minimal curvature correspond to models with similar training loss.