Parametric Density Estimation with Uncertainty using Deep Ensembles
Authors: Abel Peirson
Taylor Howell
Marius Aurel Tirlea
Venue:  (2021)
Keywords: Deep ensembles
deep learning
computer vision
density estimation
uncertainty

Abstract: In parametric density estimation, the parameters of a known probability density are typically recovered from measurements by maximizing the log-likelihood. Prior knowledge of measurement uncertainties is not included in this method -- potentially producing degraded or even biased parameter estimates.
We propose an efficient two-step, general-purpose approach for parametric density estimation using deep ensembles. 
Feature predictions and their uncertainties are returned by a deep ensemble and then combined in an importance weighted maximum likelihood estimation to recover parameters representing a known density along with their respective errors. To compare the bias-variance tradeoff of different approaches, we define an appropriate figure of merit.
We illustrate a number of use cases for our method in the physical sciences and demonstrate state-of-the-art results for X-ray polarimetry that outperform current classical and deep learning methods.

34 and a concept of memory or attention is sufficient (Vaswani et al., (017). The same can be said about time domain problems, such audio processing or voice recognition. Memory is inherently Under review as a conference paper at ICLR 2021
PARAMETRIC DENSITY ESTIMATION WITH UNCERTAINTY USING DEEP ENSEMBLES
Anonymous authors Paper under double-blind review
ABSTRACT
In parametric density estimation, the parameters of known probability density are typically recovered from measurements by maximizing the log-likelihodd. Prior knowledge of measurement uncertainties is not included in this method, po- tentially producing degraded or even biased parameter estimates. We propose an efficient two-step. general- purpose approach for parametric density estimation using deep ensembles. Feature predictions and their uncertainties are returned by deep ensemble and then combined in an importance weighted maximum likelihood estimation to recover parameters representing known density along with their respective errors. To compare the bias-variance tradeoff of different approaches, we define appropriate figure of merit. We illustrate number of use cases for our method in the physical sciences and demonstrate state-of-the-art results for X- X-ray polarimetry that outperform current classical and deep learning methods.
2
11
13
14
INTRODUCTION
15
The majority of state-of-the-art NN performances are single (high-dimensional) input, multiple- putput tasks, for instance classifying images (Krizhevsky et al. 2012). scene understanding (Red-) mon et al. 2015) and voice recognition (Graves et al. (2006). These tasks typically involve one input vector or image and single output vector of predictions.
16
17
18
19 In parametric density estimation, there a known probability density that the data (or latent features of the data) are expected to follow. The goal is to find representative distribution parameters for given dataset. In simple cases where the likelihood is calculable, maximum likelihood estimation can be used effectively. In cases where latent features of the data follow a known distribution (e.g., heights of people in a dataset photographs). NNs can potentially be used to directly estimate the approach parametric feature
20
21
22
23
distribution parameters. For clarity, we define this direct/end-to density estimation (PFDE). Such an approach requires employing entire datasets (with potentially thousands millions of digeedimensional examples) inputs in order to output vector of den- sity parameters. Furthermore, to be useful these NNs would need to generalize to arbitrarily sized dataset-inputs.
25
27
29 One example of NNs making sense large dataset-inputs is found in natural language processing. Here large text corpora, converted word vectors (Pennington et al. 2014 Devlin can be input and summarized by single output vectors using recurrent neural networks (RNNs), ),, for
30
2019)
31
32 instance in sentiment analysis (Can al. 2018). However, these problems and RNNs themselves
33 contain inductive bias - there is inherent structure in text. Not all information need be given at once
36 imperfect for PFDF, one ideally wants to know all elements of the ensemble at once to make the best prediction: sequential inductive bias is undesirable. Ultimately, memory and architectural 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 computationally 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)
42
2019). Most the work so far focuses on density estimation where the density is unknown priori.
This can be achieved with non-paramttric methods such neural density estimation (Papamakarios

---

Under review as a conference paper at ICLR 2021
et al. 2018), with parametric methods such mixture density networks (Bishop 1994). In PFDE, however, we have a known probability density over some features of the whole dataset. The features may be more difficult to predict accurately in some datapoints than others.
45
47 Typical parametric density estimation does not make use of data uncertainties where some elements in the dataset may be more noisy than others. Not including uncertainty information can lead to biased or even degraded parameter estimates. The simplest example of parametric density estimation using uncertainties is weighted mean. This is the result of a maximum likelihood estimate for multi-dimensional Gaussian. For density estimation on predicted data features, PFDE, we would like a way to quantify the predictive uncertainty. general solution offered by deep ensembles 52
49
50
51
53
LLkshminarayaaan et al., 2017). While these are not strictly equivalent to a Bayesian approach, although they can made such using appropriate regularization (Pearce et al.] (018), they offer practical predictive uncertainties, and have been shown to generalize readily (Fort et al., 2019). Additionally [Ovddia et al. (2099) have shown deep ensembles perform the best across a number of uncertainty metrics, including dataset shift, compared to competing methods such stochastic variational inference and Monte Carlo methods.
54
55
56
58
59
In this work, we propose a NN approach that circumvents large dataset-input training recurrent architectures to predict known feature density parameters over large input datasets. We use predic- tive uncertainties on features of individual dataset elements importance weights in maximum likelihood estimation. We will show that estimating known density parameters in a 2-step approach provides greater interpretabilityy and flexibility. We are able to predict uncertainties on our density parameter estimates using bootstrap methods (Efron 1979). Our method is widely applicable to a number of applied machine learning fields; 33 showcases few important examples.
60
81
62
63
64
65
Our contributions in this paper are as follows: (1) We introduce a general. flexi- ble method for PFDE using NNs. The method can be applied to any domain requiring PFDE. We illustrate number of varied domain examples in the physical sciences in $3. (2) In an in-depth 69 evaluation we show that our method outperforms not only classical methods for density estimation, but also standard NN implementations an application to X-ray polarimetry. (3) We investigate the bias- tradeoff associated with our method and introduce a tuneable hyperparameter to con- 72 trol it. Note: In the following we focus on regression examples, (since unbinned density estimation
66 Contributions:
87
68
70
73 is preferable to binned). However, similar method can be applied to prediction examples where
74 softmax class probabilities are used as heteroscedastic aleatoric uncertainty.
75
2
IMPORTANCE WEIGHTED ESTIMATION WITH DEEP ENSEMBLES
76
2.1
PROBLEM SETUP AND IIGII-LEVEL SUMMARY
77
We wish to estimate the feature density parameters of N high dimensional data points {x]: f(xx,) JN Here x E can be any high dimensional data (e.g. images, time series). N is 78
79 arbitrary. although usually large since otherwise density estimation is inaccurate. For example, con-
80 sider estimating the mean and variance of human heights from dataset consisting of photographs
81 of people. A person's 's height in each photograph is the image feature and we know this feature
82 approximately follows a Gaussian distribution. We develop a method that can estimate the density parameters (mean and variance) and generalize to any dataset of photographs.
84 In general, the function / mapping the high dimensional data points to the desired density parame- ters is unknown, since the high dimensional data is abstracted from its features. Learning f directly 86 is typically infeasible because an entire ensemble of inputs {x,, \N must be processed simultane-
ously to estimate density parameters, and this approach would have to generalize to arbitrary N and 88 density parameter values. We discuss some special cases where this is possible in $1. However, the function g mapping data features yn to the density parameters is known.
87
89
90
We cast this as a supervised learning problem where we have a dataset D consisting of N data points Musin with labels E & where x E PR We want to estimate the density parameters
31 D (Xn, Un J=1
Un U/2, ./ for an unseen test set g( (.. JNep)) ) for arbitrary Neest-
82
93 The basic recipe that comes to mind is training a single NN to predict output labels {ym h then
94 evaluate y directly. This ignores the high variance in single NN predictions (dependent train-
2

---

Under review as a conference paper at ICLR 2021
ing/random initialization), that some individual examples may be more informative than others, and that an objective to predict the most accurate output labels may not be the best for predicting good density parameters (high bias may be introduced, for instance).
96
97
Our hybrid approach is as follows. (i) Train a deep ensemble of M NN to predict {yn Tn } where on is the total uncertainty on each prediction Vn. (ii) use the {a, l as weights in an importance weighted maximum likelihood estimate. The next section, $2.2, describes procedure (i).
98
99
100
101
2.2
DEEP ENSEMBLES
102
Deep ensembles akshminarayanan et al. 2017 return robust and accurate supervised learning pre- dictions and predictive uncertainties, which enable the best density parameter predictions. These use an ensemble of individual NNs (with different random initializations) trained to predict features and their aleatoric uncertainties. Final predictions and their epistemic uncertainties then recovered by combining the estimates from each of the NNs in the ensemble.
104
105
106
In regression, deep ensembles model heteroscedastic aleatoric d, uncertainty by modifying the typ- ical mean-squared errors (MSE) objective to a negative log-likelihood (NLL) (Lakshminarayanan
108
109
al. 2017
1 2 (xogo((x) +
1 Ily (xx) (x)
Loss(y|x)
2
110
Extensions using more complex distributions like mixture density networks or heavy tailed distribu- tions may be more applicable to certain problems with prior knowledge about the eπror distribution. In practice, the log-likelihood of any exponential family could be used; we find this simple Gaussian approach to be sufficient and robust for regression problems. Our results in 33.4 for a compare a Gaussian and Von Mises distribution.
111
112
113
114
115
Epistemic uncertainty is modelled using a uniformly weighted ensemble of M NNs each trained starting from different random initialization. The regression prediction and uncertainty are approximated by the mean and standard deviation over the M NN ensemble predictions
116
spectively (each NN in the ensemble contributes equally) i.e. (xx = -M- 9m (x( and
118
(x) Var( (ūm (x) ) M=1) The epistemic uncertainty is then combined with the aleatoric in 2= = + o. Typpcally M ~ 5 15.
119
quadrature to arrive at the total uncertainty: x²
120
In part (i) of our hybrid approach for PFDE, train deep ensemble minimize the NLL (1 on desired features y. We follow the deep ensemble training procedure outlined in Lakshminarayane al. (2017) (with recast loss function from Kondll & Gal (007)) without using adversial examples, using the full dataset for each NN. Since the individual density parameters over predicted features are the final desired values in PFDE, it possible that an objective maximizing feature accuracy on the validation set is not the true objective. This is possible if the training dataset is biased or the model (1) is highly misspecified for the particular problem. The Kitaguuh et al. (2019) single CNN method in table I $3.4, shows clear case of training bias. If de-deebbad de-biasing the training dataset using
122
123
124
125
126
127
129
a more appropriate model is not possible, we have identified two potential ways of ameliorating this issue for PFDE:
131
1. Include terms in the individual NN objectives to penalize known sources of bias.
132
2. Select the top M performing NNs, as measured by a criterion that includes density param- eter prediction bias on held out test set.
133
134
In practice both can be used simultaneously. However, the former runs into batch size problems (since one needs a large sample size accurately estimate bias), and the source of bias is not always well understood. The latter naturally arises from the use of deep ensembles, but could include its own unwanted bias and risk underestimating the epistemic uncertainty. We compare selecting the
135
136
137
138
top performing NNs for the ensemble by domain specific criterion against randomly selecting NNs for the ensemble in $3.
139
We note that the NN architecture used will of course depend on the dataset domain.
3

---

Under review as a conference paper at ICLR 2021
2.3
IMPORTANCE WEIGHTED LOG-LIKELIHOOD
141
Provided mapping between high dimensional inputs and interpretable features Xn + In. we can calculate the density parameters v1, 1/2, 4/.%% by minimizing the appropriate negative log- -igliikl function pl (yn) 18,, v/.)... Some feature predictions yn will have greater total predictive uncer- tainties, On- We estimate feature density parameters by incorporating the total uncertainty into an importance weighted maximum likelihood estimate. This makes up part (ii) of our hybrid method.
142
143
144
146
An importance weight quantifies the relative importance of one example over another. Importance weighting an element should be the same as if that element were included multiple times in the dataset, proportional to its importance weight Karampatziakis
148
Langford (011). The deep ensem- ble, once trained, will act mapping between high dimensional inputs x., and feature- feature-uncertainty output pairs 7n- For each input Xn there will be M output pairs (Dnm:(aa))an (aa)um 1= one for each NN in the deep ensemble. Both the features Ino and aleatoric uncertainty variances (aa)).. can be combined by taking the appropriate mean over m,; this mean may depend on the distribution used in (I. but for the simple Gaussian case the standard mean sufficient. Taking the mean results in single output pair (In. (x)).) for each input. Epistemic uncertainties are included as in $2.2, resulting in the final output (fn.
150
152
153
154
155
158
In order to use all possible information when estimating the desired density parameters U1 2/.... we define an importance weighted negative log-likelihood function
157
N
Lw b, 1/1, 1/2, ,2,...,, vh) w. lo.log (2,,,,,,,,,,,,,...... ,,,...,
v/k),
158
Wp n
(3)
Each individual prediction yh has an associated importance weight w,. The o, term weights each Yn by its predictive uncertainty. The hyperparamter 2 0 controls the importance weighting distri- bution. A high À means the yn with the lowest (estimated) MSE will dominate the final ensemble statistic. As always in estimation problems, there is a trade-off between lower variance predictions and more bias. This can be tuned for specific application using \; discuss the procedure in detail in our example application, $3. Final density parameters are found by minimizing (2) over the domain of the density parameters vi.
160
162
164
165
166
Typically, the weights in weighted likelihood estimation are determined heuristically (Hu & Zidek 2.02). In this example, we choose w aー since it approximates the simple functional form of the likelihood used in a weighted mean estimate (A = 2). This weighting choice is also inspired by the dispersion parameter used in generalized linear models (GLMs) (Nelder Wedderbum 1972]. We expect that this weighting will retain similar robustness properties in terms of model fitting, and will generalize well many domains. However, of course, any decreasing functioon R+ -> R+ may be used to determine weights, with the most suitable choice of function s within a given class of functions (in our case, parameterized by X) to be determined by either cross- validation or performance on holdout set. In some applications it is possible to find the exact weighting function [in prep., reference deleted to maintain integrity of review process]. Further discussion of weight choice in application is given in section $3.4.
167
188
169
170
171
172
173
175
176
177
Confidence intervals on the density parameters can be calculated using the non con-aammetriic bootstrap Efron (977); select N yn. oh pairs with replacement and minimize (2). In the limit of many trials with different random subsamples, this will give the output distribution on the density parameters.
179
180
2.4
DENSITY PARA MATAR REGRESSION
For special class of parameterized densities it is possible to find the global minimizer or minimize 2 analytically (c.g. for a multivariate Gaussian). In practice, the majority of parametric densities of interest for PFDE are likely to be convex (exponential families, our application example $3, etc.), so will fall into this special class. In the general case, minimization is performed numerically to find
182
188
194
185
locally optimal solutions.
186
In this work, we employ Ipopt Wächherr & Biegler, 2006), an open-source interior-poie point solver for lcage non-convex optimization problems, minimize 2). This method can be used for convex
187
4

---

Under review as a conference paper at ICLR 2021
guaranteed to be global opti- mal. Because Ipopt finds locally optimal solutions, which are highly dependent upon an initial guess of the parameters provided to the solver, in the non-convex case, recommend nested sampling Feroz et al. 2 to test many initial guesses and then select the best local solution. Constraints on the density parameters, for instance if they have finite domain, be incorporated for both the
non-convex parametric density estimates, but only convex
189
181
192
193
convex and non-convex case. Of course, any optimizer appropriate for ? can be used and this will
194
depend on the problem.
185
The overall training and evaluation procedure is summarized in Algorithm 1.
Algorithm 1: Pseudocode for our PFDE method.
1: Identify output features yn relevant the desired density parameter(s) (e.g.. subject height
in photographs).
196
2: Train deep ensemble of NNs using loss function to maximise accuracy on the desired output features
3: Evaluate the density parameter(s) using importance weights by minimizing
4: Tune À hyperparameter for the specific application.
197
3
EXPERIMENTS
198
3.1
X-RAY POLARIMETRY
199
Measuring X-ray polarization has been major goal in astrophysics for the last 40 years. X-ray po-
200
larization can provide essential measurements of magnetic fields very close to high energy sources, such as accreting black holes and astrophysical jets [Weisskopl] 2018). The recent development of photoelectron tracking detectors (Bellazzini et al. 2003) has greatly improved the prospects of doing so. X-ray polarization telescopes with photoelectron tracking detectors directly image elec- tron tracks formed from photoelectrons scattered by the incoming X-ray photons. We describe an application of our hybrid PFDE method to X-ray polarimetry using photoelectron tracking detec-
202
204
tors. We use data from the upcoming NASA Imaging X-ray Polarization explorer (IXPE) (Sgrò & IXPE Team 2019) as working example. The problem of recovering polarization parameters from a dataset of (IXPE) electron track images has recently been announced as an open problem in the machine learning community (Moriakov al. 2020).
206
208
209
210
The linear polarization of light can be fully described by two degrees of freedom: the polarization fraction
211
< II 1, (0% 100%), and the electric vector position angle -//2 S á I T/2. These can be thought of as the magnitude and direction of a vector perpendicular to the direction of propagation of the light. In imaging X-ray polarimetry, when the detector images an X-ray source, it measures individual 2D images of electron tracks excited by incoming X-ray photons. The initial directions the electrons travel follow a known probability density that depend on the source polarization, and the problem is to recover the polarization parameters II and 6 from the collected dataset of 2D track images.
212
213
215
217
In the case of IXPE, charge tracks are imaged by hexagonal pixels. Fig. I shows some example photoelectron tracks at different X-ray energies. Each track represents the interaction of a single photon with single gas molecule. The initial track angle y follows the probability density
219
1 (1 Ilcos (2(y + +)))
p 11, Φ)
(4)
2T
where II and á are fixed polarization parameters that depend on the source. By estimating for large number of tracks, we may recover the original polarization parameters II and d. using para- metric density estimation.
222
224
Track morphologies vary greatly with energy (and even for the same energy); this affects how dif- ficult it is to recover accurate intial photoelectron angle y. Low energy tracks typically less elliptical and so more difficult to estimate. For this reason it is essential to incorporate some form of quality control in the tracks used for polarization estimates.
226
228
Current IXPE methods estimate individual track y using a moment analysis (Srro 2017). This calculates the first, second and third charge moments using the 2D coordinates of the hexagonal
5

---

Under review as a conference paper at ICLR 2021
2. 70keeV
4. 4.50keV
6,40knV
0,1
-0.2
0.0
0.2
0.2
20.1
0000
090
0.0
0.2
-0.35
0.00
0.25
-0.5
0.0
0.6
0.8
0.0
-0.2
0.4
Figure 1: Example IXPE track images at 2.7, 4.5 and keV energies (columns). The blue lines show initial photoelectron direction; the angle of these lines is y. Color represents the amount of charge deposited in a hexagonal pixel. Track morphology (and thus angle reconstruction) depends strongly on energy.
230
detector pixels, combining them to extract y. For each track, single -π S S π is output. The polarization parameters are then estimated using standard (unweighted) MLE. The moment analysis additionally outputs an estimate of the track ellipticity, which can be used as proxy for y estimation accuracy. The standard moment analysis uses track cut improve polarization re- covery 20% of the tracks are cut based on ellipticity. NNs have also recently been applied to this problem| Kitaguchi et al. (2019). This approach uses single CNNs for classification on y, with binned fits to y histograms to extract polarization parameters and track quality cuts. Our hybrid method exhibits significantly improved performance over both the standard IXPE method and this
232
233
235
237
basic NN approach.
3.2
PARAMETRIC FEATURE DENSITY ESTIMATION
240
Following 32, we define CNNs that take single track images as input and (ú, à) as output. In this
241
case the track angles y are the data features that follow the known density (4), the density parameters
242
11 = 11 Φ = 4/2, and the CNNs will make up the deep ensemble.
243
To make the hexagonal track images admissable inputs to standard CNN architectures, first convert the hexagonal images to square image arrays by shifting every other column and rescaling the distance between points, described in Steppa & Holch (2019). Since there are two possible shifts (odd and even rows), we apply both and stack the two shifted images, similar to color channels rgb images. We do this more closely approximate spatial equivariance of the CNN convolution kernels in the hexagonal space. At test time, we apply the deep ensemble to the same track 3 times, each time rotated by 120" in hexagonal space. We find this reduces all relevant prediction bias on ý (and later 11, Φ) introduced when converting from hexagonal to square coordinates.
244
245
247
249
251
To recover II, we need to predict 2y, so we use the loss function T but parameterize the true angle 3 a 2D vector (cos2y, sin2y) capture the periodicity. The loss function is follows:
1 Loss(v, v) logō2 +
1 26
11
The final NN ensembles output the 3-vector (F, In this case the mean ensemble predictions is calculated using the circular mean of { Then Ò arctan 82, To calculate the final II, e with an ensemble of M NNs for given test dataset with N tracks minimize the importance weighted NLL. (2) with likelihood
256
C(In 11, Φ)
-1 Hcoos(((s, + 6))).
(6)
2x
257
We can recast this as the convex optimization problem
minimize subject to
on Mog(1 + vXx)
(7)
||x
| <

---

Under review as a conference paper at ICLR 2021
Energy
Moments
Mom. w/ cuts
Kitaguchi et al.
Single
Ensemble
IW Ensemble
IW Ensemble (Top MSE)
IW Ensemble (von Mises)
(Random)
RoM (68% CI)
FoM
FoM
FoM
FoM
FoM
FoM
A
FoM
keV
0.78 (1.19)
0.76 (1.16)
2.6 (3.3)
0.75 (1.13) 0.74 (1.13)
0.66 (1.0))
0,66 (1,01)
1.76
0,65 (1.0)
1.24
4.5 keV
0.69 (1,05)
0.67 (1.05)
(1.9)
(0.94)
0.61 (0.94)
0.56 (0.85)
0.56 (0.85)
144
(0.84)
6.4 keV
0.58 (0.88)
0.56 (0.86)
(1.9)))
0.50 (0.75) 0.49 (0.74)
0.45 (0.69)
0.45 (0.69)
1.1
0,44 (0.68)
1,02
8.0 V
0.53 (0.8)
0.51 (0.79)
0.8 (1.3)
0.48 (0.71) 0.46 (0.71)
0.43 (0.66)
0.43 (0.66)
1.88 0.42 (0.65)
1.07
1.12 (1.72)
1.07 (1.64)
(1.64)
1.07
0.89
185 0.85 (1.29)
PL1
1.02 (1.56)
0.97 (1.48)
0.97 (1.46)
0.95 (1.45)
0.79 (1.2)
0.79 (1.2) 1.69
0.78 (1.18)
1.25
Table I: Results on energy selected track image datasets, comparing our method with the current state of the art and including an ablation study. Lower FoM is better. PL1 and PL2 are power law datasets with range spanning 2.0 8.0keV (PL1 dN/dE E-1_ and PL2 dN/dE x ²) All test datasets have 360 thousand tracks each to enable comparison with Kitaguchi et al. 2219), All methods have RMSE, 0.5". Confidence intervals (CI 68%) are calculated using the non norpparaeetric bootstrap note these are not the standard errors on FoM values, standard errors are CT/ /200, except in the case of Kiitauuhhi et al. (2019). For FoMs only the upper CI bound is necessary, since this represents the weighted. All of our method results use the Gaussian loss, (5) exeept for the final column which uses the von Mises loss. All ensembles have M = 10 NN members.
case sigmal noise ratio. IW stands for importance
258 where v, (cosin: sinj.) and = (TIcosé, ITsin)). By recasting (2 as convex optimization problem, we have guaranteed globally optimal solution for (II, é). We can solve 7 quickly and
260
efficiently using second order Newton methods. practice we use the robust open source software IpOpt, $2.4.
282
We also consider a more domain specific, non-Gaussian likelihood function for our loss, 5
We
use the log-lik lggliiklihood of the von Mises distribution for the NN loss:
Loss(v, v, v) = log (Io 2)
62 V'
(8)
284
where In is the modified Bessel function of the first kind. This is a close approximation of the wrapped Gaussian on the circle. It is more appropriate than the Gaussian (5) for angular estimates
286 since it can capture the π periodicity in ÿ. For very small ē this is equivalent to the Gaussian. We
267
compare the results from both losses in $3.4 and table [
3.2.1
FIGURE OF MERIT
In polarization estimation, want high recovered π (and accurate Φ) for known 100% polarized source (II 1), and low recovered 70 for an unpolarized source (IT 0). Since there irreducible noise in the tracks, it is impossible for any method to achieve 111g0% 1, so 11meas estimates are calibrated get the final II for an unknown source" ń' mimens/ 1000 We define figure of merit for polarization estimation:
270
272
FoM 100 f1%
We the FoM evaluate model performance: lower FoM means better polarization estimation. This is effectively a measure of the signal to noise ratio, a simplified extension of the minimum detectable polarization (MDP) typically defined for X-ray polarization (Weisskopf al., 2010) that
275
276
277
does not preclude biased estimators. It is evaluated on unseen polarized and unpolarized datasets. In estimating the FoM, we take the number of tracks N ~ 360, 000 so we can compare directly to Kitaguchi et al. (201)). We average the FoM over 200 independent track dataset samples of size N.
278
279
280
We use the FoM as the criterion select the hyperparameter in (2). In this way we can tradeoff accuracy and bias in our II, & estimates.
3.3
NN TRAINING AND SELECTION
283
Our training dataset consists of 3 million simulated tracks, examples of which are shown in fig. The track energies uniformly span 1.0 9.0ke IXPE's most sensitive range and unpolarized (uniform track angle distribution). Since we don't 't know a priori what energy each track is, we want
285
2 100% is measured before source with the same track energy distribution.

---

Under review as a conference paper at ICLR 2021
1.8
0.8
NN
Normalledd leed Counts [rad-t)
1.6
Mom:
0.6
1.4
6.4 V
PoM
1.2
0.4
1.0
0.2
0.8
0.0
0.5
1.0
1.5
2.0
2.5
3.0
0.0
1.0
-0.5
0.0
0.5
1.0
1.5
Figure 2: Left: FoM function of hyperparameter À for the Mises ensemble on the PL2 dataset. This method is used to select all A. Right: Histogram of 9 predictions for the 6.4Ke polarized dataset, II = L, o = x/ (ground truth). Black shows the ground truth density, (4). to be estimated. Red and blue show the single NN and standard moments estimates respectivcly. A single NN can better predict and thus extract more polarization signal, 1on% resulting in better FoM.
286
NNs that can make predictions for tracks of all energies. This also makes for more generalizable system, since some high energy tracks have similar characteristics to lower energy ones. Each track labelled with its 2D angle vector
287
288
289
We use a ResNet-19 (He et al. 2015) convolutional NN architecture as our base NN. This particular architecttre is large enough to overfit the training set, and trains in reasonable amount of time.
290
291
Before training we preprocess the training data (square track images). We apply pixelwise centering and rescaling. We use stochastic gradient descent with momentum and a decaying learning rate starting at le 2. We choose batch sizes 512, 1024, 2048 (tracks per batch). We trained for 150 epochs, using early stopping to prevent overfitting. We use I---norm regularization
292
10-5, We train 30 NNs and compare randomly selecting M 10 NNs to selecting M 10 NNs with the top MSEs on y for an unseen test dataset spanning all energies to make up our final NN ensemble. The results for both methods are shown in table 1
294
296
3.4
RESULTS
299
Table D shows the results of our deep ensemble PFDE method alongside the current state of the art methods. The single CNN method with optimized cuts, developed in (Kitaguchi et al., 2019), pro- vides significant improvements in 100% over the moment analysis, but adds bias to the unpolarized measurement 1o%, increasing its FoM and making a worse method for all energies. We perform an ablation study over our method, testing a single NN without using weighting when estimating [II Φ) (i.e. Wn. = Vn, B)), an ensemble of NNs without weighting, randomly selected ensemble with weighting, top MSE selected ensemble with weighting and
300
302
304
Mises loss weighted en- semble. We find a single NN without weighting beats the classical moments and moments with cuts baselines. This result visualized in the right panel of fig. 3.3 for the 6.4ke dataset: the single NN shows improved ý estimates and thus a density that more closely resembles the ground truth. Using an ensemble of NNs improves this result slightly, but the real power of our method comes with the importance weights, Our final importance weighted ensemble method, with À tuned accordingly for each energy, significantly outperforms the rest, especially in the power law datasets, where there is a reduction in FoM of almost factor of 1.5. This shows the power of simple weighted scheme over quality cuts PFDE, it allows our method take advantage of higher signal (II 100%) at higher energies in the power law datasets. The À tuning procedure is shown in the left panel of fig 3.3
306
308
309
310
312
313
315
Comparing a randomly selected ensemble with top MSE selected ensemble we find the results are almost identical. Random selection should yield more accurate approximations of the epistemic uncertainty and thus better weights, while selecting top performing NN on MSE should improve ğ accuracy. Since the results are identical, but selecting NNs has the potential to bias density estima-
316
317
319
tion, we recommend randomly selecting NNs. We note that, although not included in the table, a single NN with importance weighting performs only slightly worse than than the weighted ensem-
321
ble. Since a single NN only produces aleatoric uncertainties, this suggests, as expected, that for a correctly specified model aleatoric uncertainties dominate epistemic ones. Finally, the von Mises
8

---

Under review as a conference paper at ICLR 2021
loss shows small improvement over the simple Gaussian. This is expected, since characterizing the predictive uncertainties by a periodic distribution is more appropriate for the polarimetry ap- plication, but the improvement is small, suggesting that the Gaussian is robust starting point for many applications. We plan to release further results and more domain specific information for this particular application [reference deleted to maintain integrity of review process].
324
326
327
328
3.5
OTHER APPLICATIONS
329
There are numerous application of PFDE with uncertainty in the physical sciences and engineering. In high energy particle physics massive, short-lived particles can be detected by fitting a Cauchy distribution to the frequencies of measured decay states. Raw sensor data from hadronic particle colliders like the LHC are very noisy with variable uncertainty, meaning our PFDE approach to estimate the Cauchy distribution parameters could be very fruitful. This especially true with the widespread current use of deep learning in particle physics (Guest et al., 2018). Our approach is heuristically justified due to the asymptotic efficiency of the maximum likelihood estimator in
330
332
334
Cauchy location model (Cohen Freuel 2007). In manufacturing. GLMs fit to binomial distributions commonly used to assess product quality, or the probability of product being defunct. Today, computer vision is used for much of the inspection (Rossol 1983), making our hybrid PFDE method a potential step forward. These are just few application examples our method may be useful for any GLM based method with high dimensional data.
336
337
338
339
340
4
DISCUSSION
342
We have proposed a supervised learning framework for parametric feature density estimation. Our method uses deep ensembles to predict high dimensional data features, their aleatoric and epistemic uncertainties. We estimate feature density parameters by incorporating both of these uncertainties into an importance weighted maximum likelihood estimate. We include tuneable weighting hyper- parameter X, allowing one to control the bias-variance tradeoff for density estimation. Intuitively, in many real feature density estimation problems, some high dimensional data points may be much more informative than others due to complex noise differing generative distributions. Our method models this explicitly, weighting datapoint features by their predictive uncertainty when estimating density parameters. This avoids throwing away valuable data with quality cuts, yielding improved density estimates. Our method is scaleable to any feature dataset size and is completely flexible for specific domain applications: most NN architectures can be used. We achieve sfate-of-the-ard results over standard deep learning methods and classical algorithms in X-ray polarimetry a recent open problem in ML. We expect our method would provide similar improvements to a number of PFDE application fields, including high energy particle physics and manufacturing.
343
345
347
349
351
352
353
354
355
358
We perform an ablation study comparing single NN, deep ensemble, and various importance weighted deep ensembles. A single NN approach or standard deep ensemble improves slightly on the classical baselines, but importance weighting by predictive uncertainty provides the main improvements to our method. Selecting NNs for the deep ensemble based on quality of density estimation provides no additional gain in performance compared to random selection - since it is recom-
357
358
360
possible gerformance-basee NN selection can degrade epistemic uncertainty estimates, mend randomly selecting NNs for the ensemble. Comparing the Gaussian and von Mises distribution for feature prediction we find the standard Gaussian likelihood (1) effective and robust approx- imation, although results can potentially be improved for specific applications by choosing a more appropriate distribution over the predictive uncertainties.
362
363
364
365
366
While our method works well for densities with convex log- lig-liielshoods, non-convex ones will not necessarily yield globally optimal solutions and may be very time consuming evaluate. Future Work: Future additions to the method include more complex aleatoric uncertainty modelling. We assume a Gaussian distribution for our feature prediction (T). but for domain applications where there is an expected feature uncertainty, one could use an alternative distribution, or even a mixture density network (Bishopl 1994) for more flexibility. In that case the functional form of weighting
387
358
389
370
371
372
would have to be reconsidered. Additionally, finding the optimal weighting function for specific problem applications is likely to yield significant improvements.
373

---

Under review as a conference paper at ICLR 2021
REFERENCES
375
Ronaldo Bellazzini, F. Angelini, Luca Baldini, Alessandro Brez, Enrico Costa, Giuseppe Di Persio, Luca Latronico, M. M. Massai, Nicola Omodei, Luigi Pacciani, Paolo Soffitta, and Gloria Spandre. Novel gaseous X-ray polarimeter: data analysis and simulation.
376
378
In Polarimetry in Astronomy, volume 4843, pp. 383-393. International Society for Op- tics and Photonics, February 2003. doi: 10.1117/12.559881. URL https://www
380
Nove g org.:///-eeaaoo...aaaaaaeeeaaaar
382
0 12.1117/22.5991 short.
383
Christopher Bishop.
Mixture
Density
Networks,
January 1994.
URL
385
milttureedeensityymeeetooxx
Ethem F. Can, Aysu Ezen-Can, and Fazli Can. Multilingual Sentiment Analysis: An RNN-Based
387
Framework for Limited Data. 100y:88060111 [cs], June 2018. Uttaat:::/aara org/
abs/1800 04511 arXiv: 1806.04511.
389
Gabriela V. Cohen Freue. The Pitman estimator of the Cauchy location parameter. Jour- nal of Statistical Planning and Inference, 13766):900-1113,, June 2007. ISSN 0378-3758.
390
391
doi: 10. 10.1016/j.sspi.2000.. 05.002. URL http::/ www. www.w.sccl... wttp::/w.ww.....ccc..........o wcciencedireettt com/science/
392
atticie/ps1/$33337550633
Jacob Devlin, Ming-Wei. Chang. Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep
394
Bidirectional Transformers for Language Understanding. 1810.0008 [cs], May 2019.
URL http:://txiiiit org/ abs 1810. 04805
arXiv: 1810.04805.
396
B. Efron.
Bootstrap Methods: Another Look at the Jackknife. Annals of Statistics, 7(1):1-26,
January 1979. ISSN 0090-5364, 2168-8966. doi: 10. 176344555. URL (RR.https:/
398
projecteuclid. sirrteeccii... wttp:// Statistics.
Publisher: Institute of Mathematical
399
400
F. Feroz, M. P. Hobson, and M. Bridges. MultiNest: an efficient and robust Bayesian inference tool for cosmology and particle physics. Monthly Notices of the Royal Astronomical Society, 398(4): 1601-1614, October 2009, ISSN 00358711, 13652966. doi: 10. 13.51/ 14548.xx
401
403
URL
org/ httg://ass/000.00 3437
arXiv: 0809.3437.
Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep Ensembles: A Loss Landscape Per-
405
spective. arXiv: 1912:0911200275 [cs, stat], December 2019. URL
1912. 02757
arXiv: 1912.02757.
407 Alex Graves, Santiago Fernández, Faustino Gomez, and Jürgen Schmidhuber. Connectionist tem- poral classification: labelling unsegmented sequence data with recurrent neural networks. In Proceedings of the 23rd international conference on Machine learning. ICMI. '06, pp. 369- 376, Pittsburgh, Pennsylvania, USA, June 2006. Association for Computing Machinery. ISBN 778-1-59593-383-6 doi: (0.1145/1143444114333 1143844.114388
408
410
URL http::/://cc...r//11...
412
Dan Guest, Kyle Cranmer, and Daniel Whiteson. Deep Learning and its Application to LHC Physics. Annual Review of Nuclear and Particle Science, 68(1):161-181, October 2018. ISSN 0163-8998, URL pett:///axxivn www.:/org.a......
414
415
1545-4134. doi:
416
1806. 11484
arXiv: 1806. 1484.
417 Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. urXiv: [cs], December 2015. URL http::/arxiv. http://wwwwivoogass///// abs/ arXiv: 1512.03385.
419
1512. 03385
Feifang Hu and James V. Zidek. The Weighted Likelihood. The Canadian Journal of Statistics / La Revue Canadienne de Statistique,
421
2002. ISSN 0319-5724, doi: 10.2307/3316141.
URL WtI.http:: www. jstor. wttp://...gg/aaaae/ee1111.
10

---

Under review as a conference paper at ICLR 2021
Nikos Karampatziakis and John Langford. Online importance weight aware updates. In Proceedings of the Twenty- Seventhh Conference on Uncertainty in Artificial Intelligence, UAI' pp. 392-399, Barcelona, Spain, July 2011. AUAI Press. ISBN 978- 278-0-9749039-7-2
424
426
Alex Kendall and Yarin Gal. What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision? In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vish- wanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems 30,
428
429
pp. 5574-5584. Curran Associates, Inc., 2017. URL http: http://papers nips ccppper/ppppr
430
7141 what wetettunneetaannntn les do we need in aon-bayesiarn deepesiiandeee--eeee--eee learning for
431
pdf
432
Takao Kitaguchi, Kevin Black, Teruaki Enoto, Asami Hayato, Joanne E. Hill, Wataru B. Iwakiri, Philip Kaaret, Tsunefumi Mizumo, and Toru Tamagawa. A convolutional neural network ap- proach for reconstructing polarization information of photoelectric X-ray polarimeters. Nuclear Instruments and Methods Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment, 942:162388 October 2019. ISSN 01689002. doi: 10. 10.10a6/j.nima. abs
433
434
435
437
2019.162389. URL
arXiv: 1907.06442.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet Classification with Deep Convolutional Neural Networks. In F. Pereira, C. C. Burges, L. Bottou, and
439
K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 25, pp. ni nips cttpsccc/pppee cc/ppper with deep convolutional. -nerrll networks.
441
1097-1105. Curran Associates, Inc., 2012. URL
442
4824 imagenet wwwsssciissiaaacccccccc
443
pdf
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 6405 6416, Long Beach, California, USA, December 2017. Curran Associates Inc. ISBN 978- 1788--1118-609666
445
446
447
448
Nikita Moriakov, Ashwin Samudre, Michela Negro, Fabian Gieseke, Sydney Otten, and Luc Hen- driks. Inferring astrophysical X-ray polarization with deep learning. arXiv:2005. 08 126 [astro- ph], May 2020. URL http://///aaxxii.orrraaaa///////
450
08126
arXiv: 2005.08 0005.08126..
J. A. Nelder and R. W. M. Wedderbum. Generalized Linear Models. Journal of the Royal Statistical Society. Series (General), 135(3):370-384, 1972. ISSN 0035-9238. doi: 10.2307/2344614
452
URL hatpps:/wwwt jstorr http://ottrraaabbeeeeeeeee cty, Wiley].
Publisher: [Royal Statistical Soci-
454
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian Nowozin, Joshua V. Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift. /rXiv:/906.00 02530 fes, stat], December
456
458
2019, URL BtLpttp::/ /ttt:://axxii...t../m/ org/ab.b.r........... 906. C2530 arXiv: 1906.02530.
459
George Papamakarios. Neural Density Estimation and Likeiooodrrreeeeeeerrrrree Inf 1010::19.1111111111122 [cs, stat]. October 2019. URL Bttp://axxi..orrr//b////1......0
480
arXiv:
1910.13233.
462
George Papamakarios, Theo Pavlakou, and lain Murray. Masked Autoregresssive Flow for Density Estimation. arXiv: /70.000057777 [cs, stat), June 2018. UK.hattp.//aaxiv.orr///b/////
463
464
0 0755
arXiv: 1705.07057.
465
Tim Pearce, Mohamed Zaki, and Andy Neely. Bayesian Neural Network Ensembles.
(cs, stat], November 2018. URL http:://axx.v.orrraa/////
467
12188 arXiv: 1811.12188.
Jeffrey Pennington, Richard Socher, and Christopher Manning, Glove: Global Vectors for Word Representation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Lan- guage Processing (EMNLP), pp. 1532-1543, Doha, Qatar. October 2014. Association for Com- putational Linguistics. doi: 10.3115/v1/D14-11 bnthology/Di4-11
469
471
URL https: /www.w aclweb. org/
11

---

Under review as a conference paper at ICLR 2021
Joseph Redmon, Santosh Divvala, Ross Girshick, and Ali Farhadi. You Only Look Once: Unified, Real- Real-Time Object Detection. arXiv:1506.0260 [cs], June 2015. URRLattp://axxiv.oom
474
abs/ /106. 02640
arXiv: 1506.02640.
Lothar Rossol. Computer Vision in Industry. In Alan Pugh (ed.), Robot Vision, International Trends in Manufacturing Technology, pp. 11-18. Springer, Berlin, Heidelberg, 1983. ISBN 978- 3-662-09771-7. doi: 0000077978-3366-09777777
477
478
URL
479
Carmelo Sgro. The gas pixel detector on board the IXPE mission. In Oswald H. 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-
483
5106-1251-4 978-1-5106-125255
doi: 10.1117/11.227322.. URL https://ww:
485
486
111712..27392 full
487
C. Sgrö and IXPE Team. The Imaging X-ray Polarimetry Explorer (IXPE). Nuclear Instru-
ments and Methods in Physics Research A, 936:212-215, August 2019. ISSN 0168-9002. doi: 2.001//..ima.2115...
489
URL http::/adsabs. http://ddaaa...aaaaa.dddddd/aaaa wdw.aadd..daaa////221 /ttpt//////aaaaaaaaaaaaaaa////aa 9NIMPA.
936.. 212S
Constantin Steppa and Tim Lukas Holch. HcxagDLy Processing hexagonally sampled data with CNNs in PyTorch. SoftwareX, 9:193-198, January 2019. ISSN 30. doi: 10.1016/j.soft. 2019.02.010. http://://aaxxt .rg. 10903 01814 arXiv: 1903.01814.
482
494
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In I. Guyon, U. V. Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neu- ral Information Processing Systems 30, pp. 5998 6008. Curran Associates, Inc., 2017. URL http://papers httpa//cc..ppp.eee.ee...e.eaaaaa
495
496
498
pdf
499
Martin Weisskopf. An Overview of X-Ray Polarimetry of Astronomical Sources. Galaxies, 6:33, March 2018. doi: 10.3390/galxxie660110oo
httt:::/adsabs.aaraaacc
501
201 8Ga laaalax...
........... 1aw.l.x......................... .3w.
502
Martin C. Weisskopf, Ronald F. Elsner, and Stephen L. O'Dell. On understanding the figures of merit for detection and measurement x-ray polarization. arXiv:1006.3711 [astro-ph], pp. 77320E,
503
504
July 2010. doi: 10.1117/12.857357. URL http:://axxiv..rg///bbb..o
3711 arXiv:
505
1006.3711.
506
Andreas Wächter and Lorenz T. Biegler. On the implementation of an interior-point filter line- search algorithm for large-scale nonlinear programming. Mathematical Programming, 106(1);
508
25-57, March 2006. ISSN 1436-4646. doi: 10.10077110107-000005559444 URL https::/
12


Reviews:
[{"review_id": "kC9--ElTg8e", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Paper Summary\nThis 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.\n\n# Assessment Summary\nThis paper presents a flexible, data agnostic, and easy to implement approach to *parametric density estimation*. The paper is clearly written, giving a good understanding of the different components of the approach. I would feel confident to implement this approach myself.\n\nHowever, I believe that the presented methodology, in particular steps (ii) and (iii) achieve the opposite of the authors intentions (see negatives below). Further, the use of an ensemble model is not properly motivated. The experimental section does not give justification to each of the model components. \n\nI therefor cannot support the acceptance of this paper.\n\nFor a future manuscript, I recommend the authors to add an ablation study (see below).\n\n# Positives\n- Flexible approach: The approach is separated into two stages: (1) Training a deep ensemble on a regression task, and (2) maximum likelihood estimation of distribution parameters under the ensemble and an unseen test set. This two stage process makes the approach very flexible. The ensemble is trained once on a data set, and then exploited on multiple test sets - even with different likelihoods - for ML estimation.\n- Data agnostic: The approach does not make any further assumptions about the data, apart from the fact that the ensemble can be trained in a regression task.\n- Easy to implement: The presented approach chains a number of simple components together (deep ensembles, sample reweighting, ML estimation). Each of these components can be found or easily implemented in most common learning frameworks. Advancements in each of these components can have a trickle down effect on this approach.\n- The paper is clearly written.\n\n# Negatives\n- Step (ii) of the approach is used to select a sub-set of models from the ensemble that was trained in step (i). Step 3 & 4 in algorithm 1 gives a hint of how this selection is implemented: On a new data set the features $y$ are estimated from each neural network of the ensemble. The density parameters are then fit to approximate the distribution of features from each network under a given likelihood (I believe it should say maximize instead of minimized in algorithm 1: 3). The models for which the parameters can be best fit will be selected for step (iii). The authors claim that this removes the models with highest bias from the ensemble. However, at least in the way I understood it, this approach will in fact select the most biased models. Imagine a network that always output the same feature values. It will be easy to fit these values using the density parameters.\n- Step (iii) will give highest importance to the most confident model prediction, disregarding whether that prediction is correct or not (we cannot know during test time). A confident but wrong predictor can therefore dominate the loss for a given sample, essentially eliminating the benefit of the ensemble.\n- The experimental evaluation does not give any insights into which components of the approach actually help in performance. A proper ablation study should be carried out (see recommendations below)\n\n# Minor Comments\n- Figure 2 does not follow the formatting guidelines. Figure and table should be separated.\n- Wrong citation command (citet <-> citep)\n  - Section 2.2, 2nd paragraph\n  - Section 2.2, 4th paragraph\n  - Section 2.3, 2nd paragraph\n  - Section 2.3, 4th paragraph\n  - Section 2.3, 5th paragraph\n  - Section 2.4, 2nd paragraph\n  - Section 3.1, 1st paragraph\n  - Section 3.1, 5th paragraph\n  - Section 3.4, 1st paragraph\n  - Section 4, 1st paragraph \n\n# Recommmendations\n- Ablation study: The authors present 3 components to their approach: (1) A deep ensemble, (2) model selection, and (3) importance weighting. Whether each of these components is necessary to achieve the performance as presented in the results section is not clear. To test that, the authors should carry out an ablation study by varying the size of the ensemble, varying the number of M top performing models, and by adding or removing the sample reweighting from step (iii).\n\n--------------------------------------------------------------------------------\n\n# Review Update\nI thank the authors for their thorough response and the additional experiments. Based on these factors I will raise my score from 'clear rejection' (3) to 'okay, but not good enough' (4). I would have liked to score the paper higher, but at this stage I believe the paper is still not ready to be published. The authors acknowledged in their update that the review process helped them to understand their own work better. As a result, some aspects of their approach have been changed (e.g. removing step(ii), changes to step (iii)). I believe changes to the method go beyond the scope of the discussion phase and instead justify resubmission. This would give the authors some more time to get an in depth understanding of their approach as well.\n\n## Author comments on step (ii)\nI thank the authors for clarifying. In their response, the authors claim that they will use a held-out data set with known density parameters. It is then possible to evaluate which models in the ensemble best estimate these density parameters. I have some issues with this claim:\n\n1. This is not made clear in the paper.\n2. The approach assumes that the density parameters are unknown. Adding this assumption will weaken the paper.\n3. The selected models will be biased toward the held-out set.\n\n## Author comments on step (iii), now step (ii)\nIn their update, the authors change the reweighting scheme. Instead of having a model-based weight, the reweighting is now done solely on a per-sample basis. I believe this looks like the right direction to take.\n\n## Ablation study\nThe ablation study is important. One possible addition would be to make a comparison for different ensemble sizes.", "extra_scores": {}, "extra_text": {}}, {"review_id": "jgFL6_0pkCt", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update: After reading the other reviews/responses, I'm keeping my score of 5, due to pervasive concerns about the narrow focus of the experiments and incremental novelty of the method.\n\nThis paper proposes a method that uses ensembles of deep neural networks for parametric feature density estimation and enables control of the bias-variance tradeoff, and demonstrates the method’s performance on an X-ray polarimetry task. The paper is well-presented, the method is straightforward to implement, and the choice of experiments in the natural and physical sciences is interesting and relevant. The case study on X-ray polarimetry shows convincing results (pending some concerns on baselines mentioned below), and the authors make a good argument for the relevance of the method to parametric density estimation in the natural/physical sciences more generally.\n\nThe method appears to be sound and principled, and involves combining epistemic uncertainty (captured by the NN ensemble) with aleatoric uncertainty (estimated directly in the loss function) using quadrature. A drawback of the method is the additional computational cost incurred through training multiple ensembles. The authors mention a connection to Bayesian methods, which I believe could use further elaboration.\n\nBaselines are the standard IXPE method (based on matching moments) and a simpler NN-based approach. Bayesian methods (MCMC or variational inference) for parameter estimation would seem to provide a natural basis for comparison as well -- why were these not considered? A full Bayesian neural network likely isn’t necessary, but perhaps incorporating Bayesian uncertainty over the last layer of weights, or a similar baseline the authors deem appropriate.\n\nWhy was IPOPT chosen, instead of more standard gradient descent methods? Constraints on parameters could also be achieved via transformations of an unconstrained space (e.g. a softplus function to transform the real line to R+).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "jLH1lhK0VWz", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "* **Summary**: \n\nThis paper propose to improve parametric density estimation for unseen environment using uncertainty-aware neural networks, and proposed a practical, two-step algorithm based on DeepEnsemble. Given training data $(y, x)$ and a known density family $p(y|\\theta)$, the goal is to learn $\\theta$ in an unseen environment $\\{x', y'\\}$. \n\nThe basic method proceeds as follows: (step 1) at training time $(y, x)$, train a deep ensemble model $\\{ f_1, .., f_M \\}$ to learn a mapping $f: x \\rightarrow y$, (step 2) at testing time, estimate parameters of interest $\\theta$ by MLE w.r.t. $p(y'|\\theta)$, where $y'$ comes from ensemble predictions $f_1(x')), .., f_M(x'))$.  Based on this basic recipe, authors proposed two augmentations: (a) select only the top performing ensemble members at test time. (b) re-weight the training objective using a function of the predictive variance $w_{nm} \\propto \\sigma_{nm}^{-\\lambda}$, where $\\lambda$ is an application specific parameter.\n\n* **Strength and Weakness**\n   * (Strength) A novel method for density parameter estimation in physics problems that account for uncertainty. \n   * (Strength) An interesting application to X-ray polarization, showing clear advantage over existing approaches. \n   * (Weakness) There exists some theoretical concerns on the soundness of the approach, which should be addressed by adding additional discussion, please see Major Comments. \n   * (Weakness) Insufficient ablation study for the proposed modifications (ensemble member selection and re-weighting), as a result it is unclear the relative contribution of each components. \n\n* **Recommendation**: I recommend reject the manuscript in its current form. While acknowledging the novelty and significant of the application, I find the method a relatively straightforward combination of existing techniques (deep uncertainty and sample re-weighting), without sufficient in-depth analysis (either theoretical or empirical) on the merit the combination for the intended application. There are also some potential theoretical concerns that needs to be addressed. I'm open to adjust my recommendation, assuming these concerns are sufficiently discussed in the paper and additional ablation is conducted.\n\n* **Major Comments**:\n  * Why two-stage approach / use Gaussian likelihood: If I understand correctly, at test time, authors trained a deep model to generate uncertainty-aware predictions $y_{test}$ using Gaussian likelihood, and then conduct parameter estimation by performing MLE over a weighted likelihood constructed using the deep ensemble prediction. I have two concerns over this approach: (1) In the case that distribution of y is not Gaussian (e.g., Equation 4), how to justify learning y using a Gaussian likelihood? Would that lead to issues in uncertainty quantification, since model likelihood is mis-specified? (2) Even if it is admissible to learn y using Gaussian likelihood, are we risking under-estimating uncertainty by using MLE to estimate parameters in the second stage? In comparison, why can't we estimate the model parameters $\\theta=(\\phi, \\Pi)$ jointly with $y$ (e.g., jointly learn $y=f_y(x)$ and $\\theta=f_\\theta(x)$ using deep ensemble by minimizing the correct likelihood $p(y|\\theta)$)? Because by doing so you are learning with respect to the correct likelihood, and the uncertainties can be quantified end-to-end via deep ensemble. It would be good if author can provide discussion clarifying (1), and discuss / compare the method outlined in (2) as an intuitive baseline.\n\n  * Uncertainty under model selection: At test time, author used only the best-performing ensemble members to construct the model likelihood. There might be an concern regarding uncertainty quantification under model selection: since the model selection is not uniform, the predictive uncertainty from the select model is no longer a representative sample of the original deep ensemble. Would this cause issue in terms of uncertainty quantification? It might be good for author to justify this at least empirically by comparing against a baseline with no model selection.\n \n\n* **Minor Comments**:\n  * Notation: This is very minor: author used $k$ for total number of parameters, and $K$ for testing data points. This can be a bit confusing. It might be good to use consistently use lower case for index, and upper case for total number of parameter / samples. So it might be good to use $N_{train}$, $N_{test}$ to indicate sample size, and $K$ for the total number of parameters to estimate.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "60d_gUUWuqK", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose an approach to constructing an unconditional model on the targets (which they refer to as features) on a supervised learning dataset. As I understand it,  such models are useful in physics applications, for example. However, directly learning such a robustly model may be challenging. Thus, the authors proposed a two stage approach where, they first train an ensemble of probabilistic discriminative regression models, and use the predicted means and variances as important weights for max-likelihood estimation of the unconditional (marginal) distribution over the targets. The proposed approach works better on a range of physics-based applications of ML (X-RAY POLARIMETRY).\n\nThis work, however, raises some questions. Firstly, this work is difficult to read due to non-standard machine learning terminology, likely due to the authors background in physics? This makes it challenging to fully understand what the author's aim is (I only managed managed to put to together by the middle of paper on a second reading). I appreciate that writing cross-disciplinary work is challenging, however, I strongly encourage the authors to improve upon the clarity and structure, and to bring the terminology closer to the standard in the ML community. \n\nSecond, why was no comparison made to directly estimating the parameters of the generative model on the targets (predicted features) in the polarimetry dataset? It seems that a natural baseline for a Deep-Ensemble-based importance weighting likelihood loss would be a standard likelihood loss. If this comparison was made - can you please clarify where?\n\nMore generally, I would be keen to see an ablation study, first where a model is trained on targets (predicted features) in the dataset, then on predictions from a neural net, then on an ensemble of predictions (but no importance weighting), and then finally with importance weighting via predicted variance. This will allow the authors the piece apart which part of their contribution is useful and which isn't. In addition, comparison to use the full ensemble vs. the selected ensemble would also be informative.   \n\nFinally, given the the authors claim the method is useful for other application - it would be good to have a demonstration. As it currently stands, this method is only validated on the X-ray polarimetry task.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers are in consensus that the manuscript is not ready for publication in its current form: more comprehensive evaluation, and careful analysis (either theoretical or empirical) of the simple-but-effective methodology would improve the quality further. The discussion was constructive and helped the authors to reason about their work better.

The AC recommends Reject and encourages the authors to take the constructive feedback into consideration . 
Decision: Reject