Parametric Density Estimation with Uncertainty using Deep Ensembles

In parametric density estimation , the parameters of a known probability density1 are typically recovered from measurements by maximizing the log-likelihood.2 Prior knowledge of measurement uncertainties is not included in this method , po-3 tentially producing degraded or even biased parameter estimates . We propose4 an efficient two-step , general-purpose approach for parametric density estimation5 using deep ensembles . Feature predictions and their uncertainties are returned6 by a deep ensemble and then combined in an importance weighted maximum7 likelihood estimation to recover parameters representing a known density along8 with their respective errors . To compare the bias-variance tradeoff of different9 approaches , we define an appropriate figure of merit . We illustrate a number of10 use cases for our method in the physical sciences and demonstrate state-of-the-art11 results for X-ray polarimetry that outperform current classical and deep learning12 methods.13 1 INTRODUCTION14 The majority of state-of-the-art NN performances are single ( high-dimensional ) input , multiple-15 output tasks , for instance classifying images ( Krizhevsky et al. , 2012 ) , scene understanding ( Red-16 mon et al. , 2015 ) and voice recognition ( Graves et al. , 2006 ) . These tasks typically involve one input17 vector or image and a single output vector of predictions.18 In parametric density estimation , there is a known probability density that the data ( or latent features19 of the data ) are expected to follow . The goal is to find representative distribution parameters for a20 given dataset . In simple cases where the likelihood is calculable , maximum likelihood estimation21 can be used effectively . In cases where latent features of the data follow a known distribution ( e.g.,22 heights of people in a dataset of photographs ) , NNs can potentially be used to directly estimate the23 distribution parameters . For clarity , we define this direct/end-to-end approach as parametric feature24 density estimation ( PFDE ) . Such an approach requires employing entire datasets ( with potentially25 thousands to millions of high-dimensional examples ) as inputs in order to output a vector of den-26 sity parameters . Furthermore , to be useful these NNs would need to generalize to arbitrarily sized27 dataset-inputs.28 One example of NNs making sense of large dataset-inputs is found in natural language processing.29 Here large text corpora , converted to word vectors ( Pennington et al. , 2014 ; Devlin et al. , 2019 ) ,30 can be input and summarized by single output vectors using recurrent neural networks ( RNNs ) , for31 instance in sentiment analysis ( Can et al. , 2018 ) . However , these problems and RNNs themselves32 contain inductive bias – there is inherent structure in text . Not all information need be given at once33 and a concept of memory or attention is sufficient ( Vaswani et al. , 2017 ) . The same can be said34 about time domain problems , such as audio processing or voice recognition . Memory is inherently35 imperfect – for PFDE , one ideally wants to know all elements of the ensemble at once to make36 the best prediction : sequential inductive bias is undesirable . Ultimately , memory and architectural37 constraints make training NNs for direct PFDE computationally intractable.38 On the other hand , density estimation on data directly ( not on its latent features ) , is computationally39 tractable . Density estimation lets us find a complete statistical model of the data generating process.40 Applying deep learning to density estimation has advanced the field significantly ( Papamakarios,41 2019 ) . Most of the work so far focuses on density estimation where the density is unknown a priori.42 This can be achieved with non-parametric methods such as neural density estimation ( Papamakarios43 et al. , 2018 ) , or with parametric methods such as mixture density networks ( Bishop , 1994 ) . In PFDE,44 however , we have a known probability density over some features of the whole dataset . The features45 may be more difficult to predict accurately in some datapoints than others.46 Typical parametric density estimation does not make use of data uncertainties where some elements47 in the dataset may be more noisy than others . Not including uncertainty information can lead to48 biased or even degraded parameter estimates . The simplest example of parametric density estimation49 using uncertainties is a weighted mean . This is the result of a maximum likelihood estimate for a50 multi-dimensional Gaussian . For density estimation on predicted data features , PFDE , we would51 like a way to quantify the predictive uncertainty . A general solution is offered by deep ensembles52 ( Lakshminarayanan et al. , 2017 ) . While these are not strictly equivalent to a Bayesian approach,53 although they can be made such using appropriate regularization ( Pearce et al. , 2018 ) , they offer54 practical predictive uncertainties , and have been shown to generalize readily ( Fort et al. , 2019 ) .55 Additionally Ovadia et al . ( 2019 ) have shown deep ensembles perform the best across a number56 of uncertainty metrics , including dataset shift , compared to competing methods such as stochastic57 variational inference and Monte Carlo methods.58 In this work , we propose a NN approach that circumvents large dataset-input training or recurrent59 architectures to predict known feature density parameters over large input datasets . We use predic-60 tive uncertainties on features of individual dataset elements as importance weights in a maximum61 likelihood estimation . We will show that estimating known density parameters in a 2-step approach62 provides greater interpretability and flexibility . We are able to predict uncertainties on our density63 parameter estimates using bootstrap methods ( Efron , 1979 ) . Our method is widely applicable to a64 number of applied machine learning fields ; §3 showcases a few important examples.65 Contributions : Our contributions in this paper are as follows : ( 1 ) We introduce a general , flexi-66 ble method for PFDE using NNs . The method can be applied to any domain requiring PFDE . We67 illustrate a number of varied domain examples in the physical sciences in §3 . ( 2 ) In an in-depth68 evaluation we show that our method outperforms not only classical methods for density estimation,69 but also standard NN implementations in an application to X-ray polarimetry . ( 3 ) We investigate the70 bias-variance tradeoff associated with our method and introduce a tuneable hyperparameter to con-71 trol it . Note : In the following we focus on regression examples , ( since unbinned density estimation72 is preferable to binned ) . However , a similar method can be applied to prediction examples where73 softmax class probabilities are used as heteroscedastic aleatoric uncertainty.74 2 IMPORTANCE WEIGHTED ESTIMATION WITH DEEP ENSEMBLES75 2.1 PROBLEM SETUP AND HIGH-LEVEL SUMMARY76 We wish to estimate the feature density parameters of N high dimensional data points { x } :77 f ( { xn } Nn=1 ) . Here x ∈ RD can be any high dimensional data ( e.g . images , time series ) . N is78 arbitrary , although usually large since otherwise density estimation is inaccurate . For example , con-79 sider estimating the mean and variance of human heights from a dataset consisting of photographs80 of people . A person ’ s height in each photograph is the image feature and we know this feature81 approximately follows a Gaussian distribution . We develop a method that can estimate the density82 parameters ( mean and variance ) and generalize to any dataset of photographs.83 In general , the function f mapping the high dimensional data points to the desired density parame-84 ters is unknown , since the high dimensional data is abstracted from its features . Learning f directly85 is typically infeasible because an entire ensemble of inputs { xn } Nn=1 must be processed simultane-86 ously to estimate density parameters , and this approach would have to generalize to arbitrary N and87 density parameter values . We discuss some special cases where this is possible in §1 . However , the88 function g mapping data features yn to the density parameters is known.89 We cast this as a supervised learning problem where we have a datasetD consisting of N data points90 D = { xn , yn } Ntrainn=1 with labels y ∈ RK where x ∈ RD . We want to estimate the density parameters91 ψ1 , ψ2 , ... ψk for an unseen test set g ( { yn } Ntestn=1 ) for arbitrary Ntest.92 The basic recipe that comes to mind is training a single NN to predict output labels { yn } Nn=1 then93 evaluate g directly . This ignores the high variance in single NN predictions ( dependent on train-94 ing/random initialization ) , that some individual examples may be more informative than others , and95 that an objective to predict the most accurate output labels may not be the best for predicting good96 density parameters ( high bias may be introduced , for instance ) .97 Our hybrid approach is as follows . ( i ) Train a deep ensemble of M NNs1 to predict { yn , σn } Nn=198 where σn is the total uncertainty on each prediction yn , ( ii ) use the { σn } Nn=1 as weights in an99 importance weighted maximum likelihood estimate . The next section , §2.2 , describes procedure ( i ) .100 2.2 DEEP ENSEMBLES101 Deep ensembles ( Lakshminarayanan et al. , 2017 ) return robust and accurate supervised learning pre-102 dictions and predictive uncertainties , which enable the best density parameter predictions . These use103 an ensemble of individual NNs ( with different random initializations ) trained to predict features and104 their aleatoric uncertainties . Final predictions and their epistemic uncertainties are then recovered105 by combining the estimates from each of the NNs in the ensemble.106 In regression , deep ensembles model heteroscedastic aleatoric σa uncertainty by modifying the typ-107 ical mean-squared errors ( MSE ) objective to a negative log-likelihood ( NLL ) ( Lakshminarayanan108 et al. , 2017 ) ,109 Loss ( y|x ) = 1 2 logσ2a ( x ) + 1 2σ2a ( x ) ‖y − ŷ ( x ) ) ‖22 . ( 1 ) Extensions using more complex distributions like mixture density networks or heavy tailed distribu-110 tions may be more applicable to certain problems with prior knowledge about the error distribution.111 In practice , the log-likelihood of any exponential family could be used ; we find this simple Gaussian112 approach to be sufficient and robust for regression problems . Our results in §3.4 for a compare a113 Gaussian and Von Mises distribution.114 Epistemic uncertainty σe is modelled using a uniformly weighted ensemble of M NNs each115 trained starting from a different random initialization . The regression prediction and uncertainty116 are approximated by the mean and standard deviation over the M NN ensemble predictions re-117 spectively ( each NN in the ensemble contributes equally ) i.e . ŷ ( x ) = M−1 ∑M m=1 ŷm ( x ) and118 σ2e ( x ) = Var ( { ŷm ( x ) } Mm=1 ) . The epistemic uncertainty is then combined with the aleatoric in119 quadrature to arrive at the total uncertainty : σ2 = σ2a + σ 2 e . Typically M ∼ 5− 15.120 In part ( i ) of our hybrid approach for PFDE , we train a deep ensemble to minimize the NLL ( 1 ) on121 desired features y . We follow the deep ensemble training procedure outlined in Lakshminarayanan122 et al . ( 2017 ) ( with recast loss function from Kendall & Gal ( 2017 ) ) without using adversial examples,123 using the full dataset for each NN . Since the individual density parameters over predicted features124 are the final desired values in PFDE , it is possible that an objective maximizing feature accuracy on125 the validation set is not the true objective . This is possible if the training dataset is biased or the126 model ( 1 ) is highly misspecified for the particular problem . The Kitaguchi et al . ( 2019 ) single CNN127 method in table 1 , §3.4 , shows a clear case of training bias . If de-biasing the training dataset or using128 a more appropriate model is not possible , we have identified two potential ways of ameliorating this129 issue for PFDE:130 1 . Include terms in the individual NN objectives to penalize known sources of bias.131 2 . Select the top M performing NNs , as measured by a criterion that includes density param-132 eter prediction bias on a held out test set.133 In practice both can be used simultaneously . However , the former runs into batch size problems134 ( since one needs a large sample size to accurately estimate bias ) , and the source of bias is not always135 well understood . The latter naturally arises from the use of deep ensembles , but could include its136 own unwanted bias and risk underestimating the epistemic uncertainty . We compare selecting the137 top performing NNs for the ensemble by a domain specific criterion against randomly selecting NNs138 for the ensemble in §3.139 1We note that the NN architecture used will of course depend on the dataset domain . 2.3 IMPORTANCE WEIGHTED LOG-LIKELIHOOD140 Provided a mapping between high dimensional inputs and interpretable features xn 7→ yn , we can141 calculate the density parameters ψ1 , ψ2 , ... ψk by minimizing the appropriate negative log-likelihood142 function p ( { yn } |ψ1 , ψ2 , ... ψk ) . Some feature predictions yn will have greater total predictive uncer-143 tainties , σn . We estimate feature density parameters by incorporating the total uncertainty into an144 importance weighted maximum likelihood estimate . This makes up part ( ii ) of our hybrid method.145 An importance weight quantifies the relative importance of one example over another . Importance146 weighting an element should be the same as if that element were included multiple times in the147 dataset , proportional to its importance weight Karampatziakis & Langford ( 2011 ) . The deep ensem-148 ble , once trained , will act as mapping between high dimensional inputs xn and feature-uncertainty149 output pairs yn , σn . For each input xn there will be M output pairs { ŷnm , ( σa ) nm } Mm=1 , one for150 each NN in the deep ensemble . Both the features ŷnm and aleatoric uncertainty variances ( σa ) 2nm151 can be combined by taking the appropriate mean over m ; this mean may depend on the distribution152 used in ( 1 ) , but for the simple Gaussian case the standard mean is sufficient . Taking the mean results153 in a single output pair ( ŷn , ( σa ) n ) for each input . Epistemic uncertainties are included as in §2.2,154 resulting in the final output ( ŷn , σn ) .155 In order to use all possible information when estimating the desired density parameters ψ1 , ψ2 , ... ψk,156 we define an importance weighted negative log-likelihood function157 Lw ( { ŷn } , ψ1 , ψ2 , . . . , ψk ) = − N∑ n=1 wnlogL ( ŷn|ψ1 , ψ2 , . . . , ψk ) , ( 2 ) 158 wn = σ −λ n ( 3 ) Each individual prediction yn has an associated importance weight wn . The σ−λn term weights each159 yn by its predictive uncertainty . The hyperparamter λ ≥ 0 controls the importance weighting distri-160 bution . A high λ means the yn with the lowest ( estimated ) MSE will dominate the final ensemble161 statistic . As always in estimation problems , there is a trade-off between lower variance predictions162 and more bias . This can be tuned for a specific application using λ ; we discuss the procedure in163 detail in our example application , §3 . Final density parameters are found by minimizing ( 2 ) over the164 domain of the density parameters ψ.165 Typically , the weights in weighted likelihood estimation are determined heuristically ( Hu & Zidek,166 2002 ) . In this example , we choose w = σ−λ since it approximates the simple functional form of167 the likelihood used in a weighted mean estimate ( λ = 2 ) . This weighting choice is also inspired168 by the dispersion parameter used in generalized linear models ( GLMs ) ( Nelder & Wedderburn,169 1972 ) . We expect that this weighting will retain similar robustness properties in terms of model170 fitting , and will generalize well to many domains . However , of course , any decreasing function171 f : R+ → R+ may be used to determine weights , with the most suitable choice of function f172 within a given class of functions ( in our case , parameterized by λ ) to be determined by either cross-173 validation or performance on a holdout set . In some applications it is possible to find the exact174 weighting function [ in prep. , reference deleted to maintain integrity of review process ] . Further175 discussion of weight choice in our application is given in section §3.4.176 Confidence intervals on the density parameters can be calculated using the non-parametric bootstrap177 Efron ( 1979 ) : select N yn , σn pairs with replacement and minimize ( 2 ) . In the limit of many trials178 with different random subsamples , this will give the output distribution on the density parameters.179 2.4 DENSITY PARAMETER REGRESSION180 For a special class of parameterized densities it is possible to find the global minimizer or minimize181 ( 2 ) analytically ( e.g . for a multivariate Gaussian ) . In practice , the majority of parametric densities182 of interest for PFDE are likely to be convex ( exponential families , our application example §3 , etc . ) ,183 so will fall into this special class . In the general case , minimization is performed numerically to find184 locally optimal solutions.185 In this work , we employ Ipopt ( Wächter & Biegler , 2006 ) , an open-source interior-point solver for186 large-scale non-convex optimization problems , to minimize ( 2 ) . This method can be used for convex187 or non-convex parametric density estimates , but only convex ones are guaranteed to be global opti-188 mal . Because Ipopt finds locally optimal solutions , which are highly dependent upon an initial guess189 of the parameters provided to the solver , in the non-convex case , we recommend nested sampling190 Feroz et al . ( 2009 ) to test many initial guesses and then select the best local solution . Constraints191 on the density parameters , for instance if they have a finite domain , can be incorporated for both the192 convex and non-convex case . Of course , any optimizer appropriate for ( 2 ) can be used and this will193 depend on the problem.194 The overall training and evaluation procedure is summarized in Algorithm 1.195 Algorithm 1 : Pseudocode for our PFDE method . 1 : Identify output features yn relevant to the desired density parameter ( s ) ( e.g. , subject height in photographs ) . 2 : Train a deep ensemble of NNs using loss function ( 1 ) to maximise accuracy on the desired output features 3 : Evaluate the density parameter ( s ) using importance weights by minimizing ( 2 ) . 4 : Tune λ hyperparameter for the specific application . 196 3 EXPERIMENTS197 3.1 X-RAY POLARIMETRY198 Measuring X-ray polarization has been a major goal in astrophysics for the last 40 years . X-ray po-199 larization can provide essential measurements of magnetic fields very close to high energy sources,200 such as accreting black holes and astrophysical jets ( Weisskopf , 2018 ) . The recent development201 of photoelectron tracking detectors ( Bellazzini et al. , 2003 ) has greatly improved the prospects of202 doing so . X-ray polarization telescopes with photoelectron tracking detectors directly image elec-203 tron tracks formed from photoelectrons scattered by the incoming X-ray photons . We describe an204 application of our hybrid PFDE method to X-ray polarimetry using photoelectron tracking detec-205 tors . We use data from the upcoming NASA Imaging X-ray Polarization explorer ( IXPE ) ( Sgrò & 206 IXPE Team , 2019 ) as a working example . The problem of recovering polarization parameters from207 a dataset of ( IXPE ) electron track images has recently been announced as an open problem in the208 machine learning community ( Moriakov et al. , 2020 ) .209 The linear polarization of light can be fully described by two degrees of freedom : the polarization210 fraction 0 ≤ Π ≤ 1 , ( 0 % – 100 % ) , and the electric vector position angle −π/2 ≤ φ ≤ π/2.211 These can be thought of as the magnitude and direction of a vector perpendicular to the direction212 of propagation of the light . In imaging X-ray polarimetry , when the detector images an X-ray213 source , it measures individual 2D images of electron tracks excited by incoming X-ray photons.214 The initial directions the electrons travel follow a known probability density that depend on the215 source polarization , and the problem is to recover the polarization parameters Π and φ from the216 collected dataset of 2D track images.217 In the case of IXPE , charge tracks are imaged by hexagonal pixels . Fig . 1 shows some example218 photoelectron tracks at different X-ray energies . Each track represents the interaction of a single219 photon with a single gas molecule . The initial track angle y follows the probability density220 p ( y | Π , φ ) = 1 2π ( 1 + Πcos ( 2 ( y + φ ) ) , ( 4 ) where Π and φ are fixed polarization parameters that depend on the source . By estimating y for a221 large number of tracks , we may recover the original polarization parameters Π and φ , using para-222 metric density estimation.223 Track morphologies vary greatly with energy ( and even for the same energy ) ; this affects how dif-224 ficult it is to recover an accurate intial photoelectron angle y . Low energy tracks are typically less225 elliptical and so more difficult to estimate . For this reason it is essential to incorporate some form of226 quality control in the tracks used for polarization estimates.227 Current IXPE methods estimate individual track y using a moment analysis ( Sgro , 2017 ) . This228 calculates the first , second and third charge moments using the 2D coordinates of the hexagonal229 detector pixels , combining them to extract y . For each track , a single −π ≤ y ≤ π is output.230 The polarization parameters are then estimated using a standard ( unweighted ) MLE . The moment231 analysis additionally outputs an estimate of the track ellipticity , which can be used as a proxy for232 y estimation accuracy . The standard moment analysis uses a track cut to improve polarization re-233 covery – 20 % of the tracks are cut based on ellipticity . NNs have also recently been applied to234 this problem Kitaguchi et al . ( 2019 ) . This approach uses single CNNs for classification on y , with235 binned fits to y histograms to extract polarization parameters and track quality cuts . Our hybrid236 method exhibits significantly improved performance over both the standard IXPE method and this237 basic NN approach.238 3.2 PARAMETRIC FEATURE DENSITY ESTIMATION239 Following §2 , we define CNNs that take single track images as input and ( ŷ , σ̂ ) as output . In this240 case the track angles y are the data features that follow the known density ( 4 ) , the density parameters241 Π ≡ ψ1 , φ ≡ ψ2 , and the CNNs will make up the deep ensemble.242 To make the hexagonal track images admissable inputs to standard CNN architectures , we first243 convert the hexagonal images to square image arrays by shifting every other column and rescaling244 the distance between points , as described in Steppa & Holch ( 2019 ) . Since there are two possible245 shifts ( odd and even rows ) , we apply both and stack the two shifted images , similar to color channels246 in rgb images . We do this to more closely approximate spatial equivariance of the CNN convolution247 kernels in the hexagonal space . At test time , we apply the deep ensemble to the same track 3 times,248 each time rotated by 120◦ in hexagonal space . We find this reduces all relevant prediction bias on ŷ249 ( and later Π , φ ) introduced when converting from hexagonal to square coordinates.250 To recover Π , φwe need to predict 2y , so we use the loss function ( 1 ) but parameterize the true angle251 y as a 2D vector v = ( cos2y , sin2y ) to capture the periodicity . The loss function is as follows:252 Loss ( v , v̂ ) = 1 2 logσ̂2 + 1 2σ̂2 ‖v − v̂‖22 . ( 5 ) The final NN ensembles output the 3-vector ( v̂ , σ̂ ) . In this case the mean over ensemble predictions253 is calculated using the circular mean of { v̂m } Mm=1 . Then ŷ = 12arctan v̂2 v̂1 . To calculate the final254 Π , φ with an ensemble ofM NNs for a given test dataset withN tracks we minimize the importance255 weighted NLL ( 2 ) with likelihood256 L ( ŷn|Π , φ ) = 1 2π ( 1 + Πcos ( 2 ( ŷn + φ ) ) ) . ( 6 ) We can recast this as the convex optimization problem257 minimize x − N∑ n=1 σ̂−λn log ( 1 + vTnx ) subject to ‖x‖2 ≤ 1 ( 7 ) where vn = ( cosŷn , sinŷn ) and x = ( Πcosφ , Πsinφ ) . By recasting ( 2 ) as a convex optimization258 problem , we have a guaranteed globally optimal solution for ( Π , φ ) . We can solve ( 7 ) quickly and259 efficiently using second order Newton methods . In practice we use the robust open source software260 IpOpt , §2.4.261 We also consider a more domain specific , non-Gaussian likelihood function for our loss , ( 5 ) . We262 use the log-likelihood of the von Mises distribution for the NN loss:263 Loss ( v , v̂ ) = log ( I0 ( σ̂ −2 ) ) − 1 σ̂2 vTv̂ , ( 8 ) where I0 is the modified Bessel function of the first kind . This is a close approximation of the264 wrapped Gaussian on the circle . It is more appropriate than the Gaussian ( 5 ) for angular estimates265 since it can capture the π periodicity in ŷ . For very small σ̂ this is equivalent to the Gaussian . We266 compare the results from both losses in §3.4 and table 1.267 3.2.1 FIGURE OF MERIT268 In polarization estimation , we want high recovered Π̂100 % ( and accurate φ ) for a known 100 % 269 polarized source ( Π = 1 ) , and low recovered Π̂0 % for an unpolarized source ( Π = 0 ) . Since there270 is irreducible noise in the tracks , it is impossible for any method to achieve Π̂100 % ∼ 1 , so Π̂meas271 estimates are calibrated to get the final Π̂ for an unknown source2 : Π̂ = Π̂meas/Π̂100 % . We define a272 figure of merit for polarization estimation:273 FoM = 100× Π̂0 % /Π̂100 % . ( 9 ) We use the FoM to evaluate model performance : a lower FoM means better polarization estimation.274 This is effectively a measure of the signal to noise ratio , a simplified extension of the minimum275 detectable polarization ( MDP ) typically defined for X-ray polarization ( Weisskopf et al. , 2010 ) that276 does not preclude biased estimators . It is evaluated on unseen polarized and unpolarized datasets.277 In estimating the FoM , we take the number of tracks N ∼ 360 , 000 so we can compare directly to278 Kitaguchi et al . ( 2019 ) . We average the FoM over 200 independent track dataset samples of size N .279 We use the FoM as the criterion to select the hyperparameter λ in ( 2 ) . In this way we can tradeoff280 accuracy and bias in our Π , φ estimates.281 3.3 NN TRAINING AND SELECTION282 Our training dataset consists of 3 million simulated tracks , examples of which are shown in fig . 1.283 The track energies uniformly span 1.0 − 9.0keV , IXPE ’ s most sensitive range and are unpolarized284 ( uniform track angle distribution ) . Since we don ’ t know a priori what energy each track is , we want285 2Π̂100 % is measured before on a source with the same track energy distribution . NNs that can make predictions for tracks of all energies . This also makes for a more generalizable286 system , since some high energy tracks have similar characteristics to lower energy ones . Each track287 is labelled with its 2D angle vector v.288 We use a ResNet-19 ( He et al. , 2015 ) convolutional NN architecture as our base NN . This particular289 architecture is large enough to overfit the training set , and trains in a reasonable amount of time.290 Before training we preprocess the training data ( square track images ) . We apply pixelwise centering291 and rescaling . We use stochastic gradient descent with momentum and a decaying learning rate292 starting at 1e − 2 . We choose batch sizes 512 , 1024 , 2048 ( tracks per batch ) . We trained for 150293 epochs , using early stopping to prevent overfitting . We use L2-norm regularization 5 × 10−5 . We294 train 30 NNs and compare randomly selecting M = 10 NNs to selecting M = 10 NNs with the top295 MSEs on y for an unseen test dataset spanning all energies to make up our final NN ensemble . The296 results for both methods are shown in table 1.297 3.4 RESULTS298 Table 1 shows the results of our deep ensemble PFDE method alongside the current state of the art299 methods . The single CNN method with optimized cuts , developed in ( Kitaguchi et al. , 2019 ) , pro-300 vides significant improvements in Π100 % over the moment analysis , but adds bias to the unpolarized301 measurement Π0 % , increasing its FoM and making it a worse method for all energies . We perform302 an ablation study over our method , testing a single NN without using weighting when estimating303 ( Π , φ ) ( i.e . wn = 1 ∀n , ( 3 ) ) , an ensemble of NNs without weighting , a randomly selected ensemble304 with weighting , a top MSE selected ensemble with weighting and a von Mises loss weighted en-305 semble . We find a single NN without weighting beats the classical moments and moments with cuts306 baselines . This result is visualized in the right panel of fig . 3.3 for the 6.4keV dataset : the single NN307 shows improved ŷ estimates and thus a density that more closely resembles the ground truth . Using308 an ensemble of NNs improves this result slightly , but the real power of our method comes with the309 importance weights . Our final importance weighted ensemble method , with λ tuned accordingly for310 each energy , significantly outperforms the rest , especially in the power law datasets , where there is311 a reduction in FoM of almost a factor of 1.5 . This shows the power of a simple weighted scheme312 over quality cuts in PFDE , it allows our method to take advantage of higher signal ( Π100 % ) at higher313 energies in the power law datasets . The λ tuning procedure is shown in the left panel of fig.3.3.314 Comparing a randomly selected ensemble with a top MSE selected ensemble we find the results315 are almost identical . Random selection should yield more accurate approximations of the epistemic316 uncertainty and thus better weights , while selecting top performing NN on MSE should improve ŷ317 accuracy . Since the results are identical , but selecting NNs has the potential to bias density estima-318 tion , we recommend randomly selecting NNs . We note that , although not included in the table , a319 single NN with importance weighting performs only slightly worse than than the weighted ensem-320 ble . Since a single NN only produces aleatoric uncertainties , this suggests , as expected , that for a321 correctly specified model aleatoric uncertainties dominate epistemic ones . Finally , the von Mises322 loss shows a small improvement over the simple Gaussian . This is expected , since characterizing323 the predictive uncertainties by a periodic distribution is more appropriate for the polarimetry ap-324 plication , but the improvement is small , suggesting that the Gaussian is a robust starting point for325 many applications . We plan to release further results and more domain specific information for this326 particular application [ reference deleted to maintain integrity of review process ] .327 3.5 OTHER APPLICATIONS328 There are numerous application of PFDE with uncertainty in the physical sciences and engineering.329 In high energy particle physics massive , short-lived particles can be detected by fitting a Cauchy330 distribution to the frequencies of measured decay states . Raw sensor data from hadronic particle331 colliders like the LHC are very noisy with variable uncertainty , meaning our PFDE approach to332 estimate the Cauchy distribution parameters could be very fruitful . This especially true with the333 widespread current use of deep learning in particle physics ( Guest et al. , 2018 ) . Our approach is334 heuristically justified due to the asymptotic efficiency of the maximum likelihood estimator in a335 Cauchy location model ( Cohen Freue , 2007 ) . In manufacturing , GLMs fit to binomial distributions336 are commonly used to assess product quality , or the probability of a product being defunct . Today,337 computer vision is used for much of the inspection ( Rossol , 1983 ) , making our hybrid PFDE method338 a potential step forward . These are just a few application examples – our method may be useful for339 any GLM based method with high dimensional data.340 4 DISCUSSION341 We have proposed a supervised learning framework for parametric feature density estimation . Our342 method uses deep ensembles to predict high dimensional data features , their aleatoric and epistemic343 uncertainties . We estimate feature density parameters by incorporating both of these uncertainties344 into an importance weighted maximum likelihood estimate . We include a tuneable weighting hyper-345 parameter λ , allowing one to control the bias-variance tradeoff for density estimation . Intuitively,346 in many real feature density estimation problems , some high dimensional data points may be much347 more informative than others due to complex noise or differing generative distributions . Our method348 models this explicitly , weighting datapoint features by their predictive uncertainty when estimating349 density parameters . This avoids throwing away valuable data with quality cuts , yielding improved350 density estimates . Our method is scaleable to any feature dataset size and is completely flexible for351 specific domain applications ; most NN architectures can be used . We achieve state-of-the-art results352 over standard deep learning methods and classical algorithms in X-ray polarimetry - a recent open353 problem in ML . We expect our method would provide similar improvements to a number of PFDE354 application fields , including high energy particle physics and manufacturing.355 We perform an ablation study comparing a single NN , a deep ensemble , and various importance356 weighted deep ensembles . A single NN approach or standard deep ensemble improves slightly357 on the classical baselines , but importance weighting by predictive uncertainty provides the main358 improvements to our method . Selecting NNs for the deep ensemble based on quality of density359 estimation provides no additional gain in performance compared to random selection – since it is360 possible performance-based NN selection can degrade epistemic uncertainty estimates , we recom-361 mend randomly selecting NNs for the ensemble . Comparing the Gaussian and von Mises distribution362 for feature prediction we find the standard Gaussian likelihood ( 1 ) an effective and robust approx-363 imation , although results can potentially be improved for specific applications by choosing a more364 appropriate distribution over the predictive uncertainties.365 While our method works well for densities with convex log-likelihoods , non-convex ones will not366 necessarily yield globally optimal solutions and may be very time consuming to evaluate . Future367 Work : Future additions to the method include more complex aleatoric uncertainty modelling . We368 assume a Gaussian distribution for our feature prediction ( 1 ) , but for domain applications where369 there is an expected feature uncertainty , one could use an alternative distribution , or even a mixture370 density network ( Bishop , 1994 ) for more flexibility . In that case the functional form of weighting371 would have to be reconsidered . Additionally , finding the optimal weighting function for specific372 problem applications is likely to yield significant improvements.373 REFERENCES374 Ronaldo Bellazzini , F. Angelini , Luca Baldini , Alessandro Brez , Enrico Costa , Giuseppe Di375 Persio , Luca Latronico , M. M. Massai , Nicola Omodei , Luigi Pacciani , Paolo Soffitta,376 and Gloria Spandre . Novel gaseous x-ray polarimeter : data analysis and simulation.377 In Polarimetry in Astronomy , volume 4843 , pp . 383–393 . International Society for Op-378 tics and Photonics , February 2003. doi : 10.1117/12.459381 . URL https : //www.379 spiedigitallibrary.org/conference-proceedings-of-spie/4843/0000/380 Novel-gaseous-x-ray-polarimeter-data-analysis-and-simulation/381 10.1117/12.459381.short.382 Christopher Bishop . Mixture Density Networks . January 1994 . URL383 https : //www.microsoft.com/en-us/research/publication/384 mixture-density-networks/.385 Ethem F. Can , Aysu Ezen-Can , and Fazli Can . Multilingual Sentiment Analysis : An RNN-Based386 Framework for Limited Data . arXiv:1806.04511 [ cs ] , June 2018 . URL http : //arxiv.org/387 abs/1806.04511 . arXiv : 1806.04511.388 Gabriela V. Cohen Freue . The Pitman estimator of the Cauchy location parameter . Jour-389 nal of Statistical Planning and Inference , 137 ( 6 ) :1900–1913 , June 2007 . ISSN 0378-3758.390 doi : 10.1016/j.jspi.2006.05.002 . URL http : //www.sciencedirect.com/science/391 article/pii/S0378375806001285.392 Jacob Devlin , Ming-Wei Chang , Kenton Lee , and Kristina Toutanova . BERT : Pre-training of Deep393 Bidirectional Transformers for Language Understanding . arXiv:1810.04805 [ cs ] , May 2019.394 URL http : //arxiv.org/abs/1810.04805 . arXiv : 1810.04805.395 B. Efron . Bootstrap Methods : Another Look at the Jackknife . Annals of Statistics , 7 ( 1 ) :1–26,396 January 1979 . ISSN 0090-5364 , 2168-8966. doi : 10.1214/aos/1176344552 . URL https : //397 projecteuclid.org/euclid.aos/1176344552 . Publisher : Institute of Mathematical398 Statistics.399 F. Feroz , M. P. Hobson , and M. Bridges . MultiNest : an efficient and robust Bayesian inference tool400 for cosmology and particle physics . Monthly Notices of the Royal Astronomical Society , 398 ( 4 ) :401 1601–1614 , October 2009 . ISSN 00358711 , 13652966. doi : 10.1111/j.1365-2966.2009.14548.x.402 URL http : //arxiv.org/abs/0809.3437 . arXiv : 0809.3437.403 Stanislav Fort , Huiyi Hu , and Balaji Lakshminarayanan . Deep Ensembles : A Loss Landscape Per-404 spective . arXiv:1912.02757 [ cs , stat ] , December 2019 . URL http : //arxiv.org/abs/405 1912.02757. arXiv : 1912.02757.406 Alex Graves , Santiago Fernández , Faustino Gomez , and Jürgen Schmidhuber . Connectionist tem-407 poral classification : labelling unsegmented sequence data with recurrent neural networks . In408 Proceedings of the 23rd international conference on Machine learning , ICML ’ 06 , pp . 369–409 376 , Pittsburgh , Pennsylvania , USA , June 2006 . Association for Computing Machinery . ISBN410 978-1-59593-383-6. doi : 10.1145/1143844.1143891 . URL https : //doi.org/10.1145/411 1143844.1143891.412 Dan Guest , Kyle Cranmer , and Daniel Whiteson . Deep Learning and its Application to LHC Physics.413 Annual Review of Nuclear and Particle Science , 68 ( 1 ) :161–181 , October 2018 . ISSN 0163-8998,414 1545-4134. doi : 10.1146/annurev-nucl-101917-021019 . URL http : //arxiv.org/abs/415 1806.11484. arXiv : 1806.11484.416 Kaiming He , Xiangyu Zhang , Shaoqing Ren , and Jian Sun . Deep Residual Learning for Image417 Recognition . arXiv:1512.03385 [ cs ] , December 2015 . URL http : //arxiv.org/abs/418 1512.03385. arXiv : 1512.03385.419 Feifang Hu and James V. Zidek . The Weighted Likelihood . The Canadian Journal of Statistics / La420 Revue Canadienne de Statistique , 30 ( 3 ) :347–371 , 2002 . ISSN 0319-5724. doi : 10.2307/3316141.421 URL https : //www.jstor.org/stable/3316141.422 Nikos Karampatziakis and John Langford . Online importance weight aware updates . In Proceedings423 of the Twenty-Seventh Conference on Uncertainty in Artificial Intelligence , UAI ’ 11 , pp . 392–399,424 Barcelona , Spain , July 2011 . AUAI Press . ISBN 978-0-9749039-7-2.425 Alex Kendall and Yarin Gal . What Uncertainties Do We Need in Bayesian Deep Learning for426 Computer Vision ? In I. Guyon , U. V. Luxburg , S. Bengio , H. Wallach , R. Fergus , S. Vish-427 wanathan , and R. Garnett ( eds . ) , Advances in Neural Information Processing Systems 30,428 pp . 5574–5584 . Curran Associates , Inc. , 2017 . URL http : //papers.nips.cc/paper/429 7141-what-uncertainties-do-we-need-in-bayesian-deep-learning-for-computer-vision.430 pdf.431 Takao Kitaguchi , Kevin Black , Teruaki Enoto , Asami Hayato , Joanne E. Hill , Wataru B. Iwakiri,432 Philip Kaaret , Tsunefumi Mizuno , and Toru Tamagawa . A convolutional neural network ap-433 proach for reconstructing polarization information of photoelectric X-ray polarimeters . Nuclear434 Instruments and Methods in Physics Research Section A : Accelerators , Spectrometers , Detectors435 and Associated Equipment , 942:162389 , October 2019 . ISSN 01689002. doi : 10.1016/j.nima.436 2019.162389 . URL http : //arxiv.org/abs/1907.06442 . arXiv : 1907.06442.437 Alex Krizhevsky , Ilya Sutskever , and Geoffrey E Hinton . ImageNet Classification with438 Deep Convolutional Neural Networks . In F. Pereira , C. J. C. Burges , L. Bottou , and439 K. Q. Weinberger ( eds . ) , Advances in Neural Information Processing Systems 25 , pp.440 1097–1105 . Curran Associates , Inc. , 2012 . URL http : //papers.nips.cc/paper/441 4824-imagenet-classification-with-deep-convolutional-neural-networks.442 pdf.443 Balaji Lakshminarayanan , Alexander Pritzel , and Charles Blundell . Simple and scalable predictive444 uncertainty estimation using deep ensembles . In Proceedings of the 31st International Conference445 on Neural Information Processing Systems , NIPS ’ 17 , pp . 6405–6416 , Long Beach , California,446 USA , December 2017 . Curran Associates Inc. ISBN 978-1-5108-6096-4.447 Nikita Moriakov , Ashwin Samudre , Michela Negro , Fabian Gieseke , Sydney Otten , and Luc Hen-448 driks . Inferring astrophysical X-ray polarization with deep learning . arXiv:2005.08126 [ astro-449 ph ] , May 2020 . URL http : //arxiv.org/abs/2005.08126 . arXiv : 2005.08126.450 J . A. Nelder and R. W. M. Wedderburn . Generalized Linear Models . Journal of the Royal Statistical451 Society . Series A ( General ) , 135 ( 3 ) :370–384 , 1972 . ISSN 0035-9238. doi : 10.2307/2344614.452 URL https : //www.jstor.org/stable/2344614 . Publisher : [ Royal Statistical Soci-453 ety , Wiley ] .454 Yaniv Ovadia , Emily Fertig , Jie Ren , Zachary Nado , D. Sculley , Sebastian Nowozin , Joshua V.455 Dillon , Balaji Lakshminarayanan , and Jasper Snoek . Can You Trust Your Model ’ s Uncertainty ? 456 Evaluating Predictive Uncertainty Under Dataset Shift . arXiv:1906.02530 [ cs , stat ] , December457 2019 . URL http : //arxiv.org/abs/1906.02530 . arXiv : 1906.02530.458 George Papamakarios . Neural Density Estimation and Likelihood-free Inference . arXiv:1910.13233459 [ cs , stat ] , October 2019 . URL http : //arxiv.org/abs/1910.13233 . arXiv:460 1910.13233.461 George Papamakarios , Theo Pavlakou , and Iain Murray . Masked Autoregressive Flow for Density462 Estimation . arXiv:1705.07057 [ cs , stat ] , June 2018 . URL http : //arxiv.org/abs/1705.463 07057. arXiv : 1705.07057.464 Tim Pearce , Mohamed Zaki , and Andy Neely . Bayesian Neural Network Ensembles.465 arXiv:1811.12188 [ cs , stat ] , November 2018 . URL http : //arxiv.org/abs/1811.466 12188. arXiv : 1811.12188.467 Jeffrey Pennington , Richard Socher , and Christopher Manning . Glove : Global Vectors for Word468 Representation . In Proceedings of the 2014 Conference on Empirical Methods in Natural Lan-469 guage Processing ( EMNLP ) , pp . 1532–1543 , Doha , Qatar , October 2014 . Association for Com-470 putational Linguistics . doi : 10.3115/v1/D14-1162 . URL https : //www.aclweb.org/471 anthology/D14-1162.472 Joseph Redmon , Santosh Divvala , Ross Girshick , and Ali Farhadi . You Only Look Once : Unified,473 Real-Time Object Detection . arXiv:1506.02640 [ cs ] , June 2015 . URL http : //arxiv.org/474 abs/1506.02640 . arXiv : 1506.02640.475 Lothar Rossol . Computer Vision in Industry . In Alan Pugh ( ed . ) , Robot Vision , International476 Trends in Manufacturing Technology , pp . 11–18 . Springer , Berlin , Heidelberg , 1983 . ISBN 978-477 3-662-09771-7. doi : 10.1007/978-3-662-09771-7 2 . URL https : //doi.org/10.1007/478 978-3-662-09771-7_2.479 Carmelo Sgro . The gas pixel detector on board the IXPE mission . In Oswald H.480 Siegmund ( ed . ) , UV , X-Ray , and Gamma-Ray Space Instrumentation for Astron-481 omy XX , pp . 16 , San Diego , United States , August 2017 . SPIE . ISBN 978-1-482 5106-1251-8 978-1-5106-1252-5. doi : 10.1117/12.2273922 . URL https : //www.483 spiedigitallibrary.org/conference-proceedings-of-spie/10397/484 2273922/The-gas-pixel-detector-on-board-the-IXPE-mission/10.485 1117/12.2273922.full.486 C. Sgrò and IXPE Team . The Imaging X-ray Polarimetry Explorer ( IXPE ) . Nuclear Instru-487 ments and Methods in Physics Research A , 936:212–215 , August 2019 . ISSN 0168-9002. doi:488 10.1016/j.nima.2018.10.111 . URL http : //adsabs.harvard.edu/abs/2019NIMPA.489 936 .. 212S.490 Constantin Steppa and Tim Lukas Holch . HexagDLy - Processing hexagonally sampled data with491 CNNs in PyTorch . SoftwareX , 9:193–198 , January 2019 . ISSN 23527110. doi : 10.1016/j.softx.492 2019.02.010 . URL http : //arxiv.org/abs/1903.01814 . arXiv : 1903.01814.493 Ashish Vaswani , Noam Shazeer , Niki Parmar , Jakob Uszkoreit , Llion Jones , Aidan N Gomez,494 Łukasz Kaiser , and Illia Polosukhin . Attention is All you Need . In I. Guyon , U. V. Luxburg,495 S. Bengio , H. Wallach , R. Fergus , S. Vishwanathan , and R. Garnett ( eds . ) , Advances in Neu-496 ral Information Processing Systems 30 , pp . 5998–6008 . Curran Associates , Inc. , 2017 . URL497 http : //papers.nips.cc/paper/7181-attention-is-all-you-need.pdf.498 Martin Weisskopf . An Overview of X-Ray Polarimetry of Astronomical Sources . Galaxies , 6:33,499 March 2018. doi : 10.3390/galaxies6010033 . URL http : //adsabs.harvard.edu/abs/500 2018Galax ... 6 ... 33W.501 Martin C. Weisskopf , Ronald F. Elsner , and Stephen L. O ’ Dell . On understanding the figures of merit502 for detection and measurement of x-ray polarization . arXiv:1006.3711 [ astro-ph ] , pp . 77320E,503 July 2010. doi : 10.1117/12.857357 . URL http : //arxiv.org/abs/1006.3711 . arXiv:504 1006.3711.505 Andreas Wächter and Lorenz T. Biegler . On the implementation of an interior-point filter line-506 search algorithm for large-scale nonlinear programming . Mathematical Programming , 106 ( 1 ) :507 25–57 , March 2006 . ISSN 1436-4646. doi : 10.1007/s10107-004-0559-y . URL https : //508 doi.org/10.1007/s10107-004-0559-y.509

Summary: This paper addresses the problem of estimating the distribution parameters of features extracted from a set of high dimensional observations, a problem that is common in the physical sciences. To solve this problem, the authors present a deep learning approach that utilises a combination of (i) deep ensemble training, (ii) post hoc model selection, and (iii) importance weighted parameter estimation. First, a deep ensemble is trained to solve a regression task (observation -> feature). During testing, this ensemble is frozen and used to generate feature samples from unseen observations. Using these feature samples, it is possible to estimate the distribution parameters using maximum likelihood estimation. The authors evaluate their method on X-ray polarimetry, and compare it with two other approaches, one of which is also a deep learning approach. On all tasks, the presented method outperforms both baseline approaches.