Bayesian Approaches to Distribution Regression

Ho Chung Leon Law∗ Dougal J. Sutherland∗ Dino Sejdinovic Seth Flaxman
University of Oxford<EMAIL_ADDRESS>University College London
<EMAIL_ADDRESS>University of Oxford<EMAIL_ADDRESS>Imperial
College London<EMAIL_ADDRESS>
###### Abstract

Distribution regression has recently attracted much interest as a generic
solution to the problem of supervised learning where labels are available at
the group level, rather than at the individual level. Current approaches,
however, do not propagate the uncertainty in observations due to sampling
variability in the groups. This effectively assumes that small and large
groups are estimated equally well, and should have equal weight in the final
regression. We account for this uncertainty with a Bayesian distribution
regression formalism, improving the robustness and performance of the model
when group sizes vary. We frame our models in a neural network style, allowing
for simple MAP inference using backpropagation to learn the parameters, as
well as MCMC-based inference which can fully propagate uncertainty. We
demonstrate our approach on illustrative toy datasets, as well as on a
challenging problem of predicting age from images.

##  1 INTRODUCTION

Distribution regression is the problem of learning a regression function from
samples of a distribution to a single set-level label. For example, we might
attempt to infer the sentiment of texts based on word-level features, to
predict the label of an image based on small patches, or even perform
traditional parametric statistical inference by learning a function from sets
of samples to the parameter values.

Recent years have seen wide-ranging applications of this framework, including
inferring summary statistics in Approximate Bayesian Computation (Mitrovic et
al., 2016), estimating Expectation Propagation messages (Jitkrittum et al.,
2015), predicting the voting behaviour of demographic groups (Flaxman et al.,
2015, 2016), and learning the total mass of dark matter halos from observable
galaxy velocities (Ntampaka et al., 2015, 2016). Closely related distribution
classification problems also include identifying the direction of causal
relationships from data (Lopez-Paz et al., 2015) and classifying text based on
bags of word vectors (Yoshikawa et al., 2014; Kusner et al., 2015).

One particularly appealing approach to the distribution regression problem is
to represent the input set of samples by their kernel mean embedding
(described in Section 2.1), where distributions are represented as single
points in a reproducing kernel Hilbert space. Standard kernel methods can then
be applied for distribution regression, classification, anomaly detection, and
so on. This approach was perhaps first popularized by Muandet et al. (2012);
Szábo et al. (2016) provided a recent learning-theoretic analysis.

11footnotetext: These authors contributed equally.

In this framework, however, each distribution is simply represented by the
empirical mean embedding, ignoring the fact that large sample sets are much
more precisely understood than small ones. Most studies also use point
estimates for their regressions, such as kernel ridge regression or support
vector machines, thus ignoring uncertainty both in the distribution embeddings
and in the regression model.

#### Our Contributions

We propose a set of Bayesian approaches to distribution regression. The
simplest method, similar to that of Flaxman et al. (2015), is to use point
estimates of the input embeddings but account for uncertainty in the
regression model with simple Bayesian linear regression. Alternatively, we can
treat uncertainty in the input embeddings but ignore model uncertainty with
the proposed Bayesian mean shrinkage model, which builds on a recently
proposed Bayesian nonparametric model of uncertainty in kernel mean embeddings
(Flaxman et al., 2016), and then use a sparse representation of the desired
function in the RKHS for prediction in the regression model. This model allows
for a full account of uncertainty in the mean embedding, but requires a point
estimate of the regression function for conjugacy; we thus use backpropagation
to obtain a MAP estimate for it as well as various hyperparameters. We then
combine the treatment of the two sources of uncertainty into a fully Bayesian
model and use Hamiltonian Monte Carlo for efficient inference. Depending on
the inferential goals, each model can be useful. We demonstrate our approaches
on an illustrative toy problem as well as a challenging real-world age
estimation task.

##  2 BACKGROUND

###  2.1 Problem Overview

Distribution regression is the task of learning a classifier or a regression
function that maps probability distributions to labels. The challenge of
distribution regression goes beyond the standard supervised learning setting:
we do not have access to exact input-output pairs since the true inputs,
probability distributions, are observed only through samples from that
distribution:

$\left(\\{x_{j}^{1}\\}_{j=1}^{N_{1}},y_{1}\right),\ldots,\left(\\{x_{j}^{n}\\}_{j=1}^{N_{n}},y_{n}\right),$
(1)

so that each bag $\\{x_{j}^{i}\\}_{j=1}^{N_{i}}$ has a label $y_{i}$ along
with $N_{i}$ individual observations $x_{j}^{i}\in{\mathcal{X}}$. We assume
that the observations $\\{x_{j}^{i}\\}_{j=1}^{N_{i}}$ are i.i.d. samples from
some unobserved distribution ${\mathsf{P}}_{i}$, and that the true label
$y_{i}$ depends only on ${\mathsf{P}}_{i}$. We wish to avoid making any strong
parametric assumptions on the ${\mathsf{P}}_{i}$. For the present work, we
will assume the labels $y_{i}$ are real-valued; Appendix B shows an extension
to binary classification. We typically take the observation space
${\mathcal{X}}$ to be a subset of ${\mathbb{R}}^{p}$, but it could easily be a
structured domain such as text or images, since we access it only through a
kernel (for examples, see e.g. Gärtner, 2008).

We consider the standard approach to distribution regression, which relies on
kernel mean embeddings and kernel ridge regression. For any positive definite
kernel function $k:{\mathcal{X}}\times{\mathcal{X}}\to{\mathbb{R}}$, there
exists a unique reproducing kernel Hilbert space (RKHS) ${\mathcal{H}_{k}}$, a
possibly infinite-dimensional space of functions
$f:{\mathcal{X}}\to{\mathbb{R}}$ where evaluation can be written as an inner
product, and in particular $f(x)=\langle
f,k(\cdot,x)\rangle_{\mathcal{H}_{k}}$ for all
$f\in{\mathcal{H}_{k}},x\in{\mathcal{X}}$. Here
$k(\cdot,x)\in{\mathcal{H}_{k}}$ is a function of one argument, $y\mapsto
k(y,x)$.

Given a probability measure ${\mathsf{P}}$ on ${\mathcal{X}}$, let us define
the kernel mean embedding into ${\mathcal{H}_{k}}$ as

${\mu_{{\mathsf{P}}}}=\int
k\left(\cdot,x\right){\mathsf{P}}(dx)\in{\mathcal{H}_{k}}.$ (2)

Notice that ${\mu_{{\mathsf{P}}}}$ serves as a high- or infinite-dimensional
vector representation of ${\mathsf{P}}$. For the kernel mean embedding of
${\mathsf{P}}$ into ${\mathcal{H}_{k}}$ to be well-defined, it suffices that
$\int\sqrt{k(x,x)}{\mathsf{P}}(dx)<\infty$, which is trivially satisfied for
all ${\mathsf{P}}$ if $k$ is bounded. Analogously to the reproducing property
of RKHS, ${\mu_{{\mathsf{P}}}}$ represents the expectation function on
${\mathcal{H}_{k}}$: $\int h(x){\mathsf{P}}(dx)=\langle
h,{\mu_{{\mathsf{P}}}}\rangle_{\mathcal{H}_{k}}$. For so-called
_characteristic_ kernels (Sriperumbudur et al., 2010), every probability
measure has a unique embedding, and thus ${\mu_{{\mathsf{P}}}}$ completely
determines the corresponding probability measure.

###  2.2 Estimating Mean Embeddings

For a set of samples $\\{x_{j}\\}^{n}_{j=1}$ drawn iid from ${\mathsf{P}}$,
the empirical estimator of ${\mu_{{\mathsf{P}}}}$,
$\widehat{{\mu_{{\mathsf{P}}}}}\in{\mathcal{H}_{k}}$, is given by

$\widehat{{\mu_{{\mathsf{P}}}}}=\mu_{\widehat{\mathsf{P}}}=\int
k\left(\cdot,x\right)\hat{\mathsf{P}}(dx)=\frac{1}{n}\sum_{j=1}^{n}k(\cdot,x_{j}).$
(3)

This is the standard estimator used by previous distribution regression
approaches, which the reproducing property of ${\mathcal{H}_{k}}$ shows us
corresponds to the kernel

$\langle\widehat{{\mu_{{\mathsf{P}}}}_{i}},\widehat{{\mu_{{\mathsf{P}}}}_{j}}\rangle_{\mathcal{H}_{k}}=\frac{1}{N_{i}N_{j}}\sum_{\ell=1}^{N_{i}}\sum_{r=1}^{N_{j}}k(x_{\ell}^{i},x_{r}^{j}).$
(4)

But (3) is an empirical mean estimator in a high- or infinite-dimensional
space, and is thus subject to the well-known _Stein phenomenon_ , so that its
performance is dominated by the James-Stein shrinkage estimators. Indeed,
Muandet et al. (2014) studied shrinkage estimators for mean embeddings, which
can result in substantially improved performance for some tasks (Ramdas and
Wehbe, 2015). Flaxman et al. (2016) proposed a Bayesian analogue of shrinkage
estimators, which we now review.

This approach consists of (1) a Gaussian Process prior
$\mu_{\mathsf{P}}\sim\operatorname{\mathcal{GP}}(m_{0},r(\cdot,\cdot))$ on
${\mathcal{H}_{k}}$, where $r$ is selected to ensure that
${\mu_{{\mathsf{P}}}}\in{\mathcal{H}_{k}}$ almost surely and (2) a normal
likelihood $\widehat{{\mu_{{\mathsf{P}}}}}({\bf
x})\mid{\mu_{{\mathsf{P}}}}({\bf x})\sim\mathcal{N}({\mu_{{\mathsf{P}}}}({\bf
x}),\Sigma)$. Here, conjugacy of the prior and the likelihood leads to a
Gaussian process posterior on the true embedding ${\mu_{{\mathsf{P}}}}$, given
that we have observed $\widehat{{\mu_{{\mathsf{P}}}}}$ at some set of
locations ${\bf x}$. The posterior mean is then essentially identical to a
particular shrinkage estimator of Muandet et al. (2014), but the method
described here has the extra advantage of a closed form uncertainty estimate,
which we utilise in our distributional approach. For the choice of $r$, we use
a Gaussian RBF kernel $k$, and choose either $r=k$ or, following Flaxman et
al. (2016), $r(x,x^{\prime})=\int k(x,z)\,k(z,x^{\prime})\,\nu(dz)$ where
$\nu$ is proportional to a Gaussian measure. For details of our choices, and
why they are sufficient for our purposes, see Appendix A.

This model accounts for the uncertainty based on the number of samples
$N_{i}$, shrinking the embeddings for small sample sizes more. As we will see,
this is essential in the context of distribution regression, particularly when
bag sizes are imbalanced.

###  2.3 Standard Approaches to Distribution Regression

Following Szábo et al. (2016), assume that the probability distributions
${\mathsf{P}}_{i}$ are each drawn randomly from some unknown meta-distribution
over probability distributions, and take a two-stage approach, illustrated as
in Figure 1.

Figure 1: Each bag is summarised by a kernel mean embedding
$\mu_{i}\in\mathcal{H}_{k}$; a regression function
$f:\mathcal{H}_{k}\rightarrow\mathbb{R}$ predicts labels $y_{i}\in\mathbb{R}$.
We propose a Bayesian approach to propagate uncertainty due to the number of
samples in each bag, obtaining posterior credible intervals illustrated in
grey.

Denoting the feature map $k(\cdot,x)\in{\mathcal{H}_{k}}$ by $\phi(x)$, one
uses the empirical kernel mean estimator (3) to separately estimate the mean
of each group:

$\widehat{\mu_{1}}=\frac{1}{N_{1}}\sum_{j=1}^{N_{1}}\phi(x_{j}^{1}),~{}~{}\ldots,~{}~{}\widehat{\mu_{n}}=\frac{1}{N_{n}}\sum_{i=1}^{N_{n}}\phi(x^{n}_{j}).$
(5)

Next, one uses kernel ridge regression (Saunders et al., 1998) to learn a
function $f:{\mathcal{H}_{k}}\to{\mathbb{R}}$, by minimizing the squared loss
with an RKHS complexity penalty:

$\hat{f}=\operatorname*{argmin}_{f\in\mathcal{H}_{K}}\sum_{i}(y_{i}-f(\widehat{\mu_{i}}))^{2}+\lambda\|f\|^{2}_{\mathcal{H}_{K}}.$

Here $K:{\mathcal{H}_{k}}\times{\mathcal{H}_{k}}\to{\mathbb{R}}$ is a “second-
level” kernel on mean embeddings. If $K$ is a linear kernel on the RKHS
${\mathcal{H}_{k}}$, then the resulting method can be interpreted as a linear
(ridge) regression on mean embeddings, which are themselves nonlinear
transformations of the inputs. A nonlinear second-level kernel on
${\mathcal{H}_{k}}$ sometimes improves performance (Muandet et al., 2012;
Szábo et al., 2016).

Distribution regression as described is not scalable for even modestly-sized
datasets, as computing each of the $\mathcal{O}(n^{2})$ entries of the
relevant kernel matrix requires time $\mathcal{O}(N_{i}N_{j})$. Many
applications have thus used variants of random Fourier features (Rahimi and
Recht, 2007). In this paper we instead expand in terms of landmark points
drawn randomly from the observations, yielding radial basis networks
(Broomhead and Lowe, 1988) with mean pooling.

##  3 MODELS

We consider here three different Bayesian models, with each model encoding
different types of uncertainty. We begin with a non-Bayesian RBF network
formulation of the standard approach to distribution regression as a baseline,
before refining this approach to better propagate uncertainty in bag size, as
well as model parameters.

Figure 2: Our baseline model, a RBF network for distribution regression.
$X_{i}$ represents the matrix of samples for bag $i$, while
$k(X_{i},u_{\ell})$ represents the element wise operation on each row of
$X_{i}$, with $b$ representing the batch size for stochastic gradient descent.

###  3.1 Baseline Model

The baseline RBF network formulation we employ here is a variation of the
approaches of Broomhead and Lowe (1988), Que and Belkin (2016), Law et al.
(2017), and Zaheer et al. (2017). As shown in Figure 2, the initial input is a
minibatch consisting of several bags $X_{i}$, each containing $N_{i}$ points.
Each point is then converted to an explicit featurisation, taking the role of
$\phi$ in (5), by a radial basis layer: $x^{i}_{j}\in{\mathbb{R}}^{p}$ is
mapped to

$\phi(x^{i}_{j})=[k(x^{i}_{j},u_{1}),\ldots,k(x^{i}_{j},u_{d})]^{\top}\in\mathbb{R}^{d}$

where $\mathbf{u}=\\{u_{\ell}\\}_{\ell=1}^{d}$ are landmark points. A mean
pooling layer yields the estimated mean embedding $\hat{\mu}_{i}$
corresponding to each of the bags $j$ represented in the minibatch, where
$\hat{\mu}_{i}=\frac{1}{N_{i}}\sum_{j=1}^{N_{i}}\phi(x^{i}_{j})$.111In the
implementation, we stack all of the bags $X_{i}$ into a single matrix of size
$\sum_{j}N_{j}\times d$ for the first layer, then perform pooling via sparse
matrix multiplication. Finally, a fully connected output layer gives real-
valued labels $\hat{y}_{i}=\beta^{\mathsf{T}}\hat{\mu}_{i}+b$. As a loss
function we use the mean square error
$\frac{1}{n}\sum_{i}(\hat{y}_{i}-y_{i})^{2}$. For learning, we use
backpropagation with the Adam optimizer (Kingma and Ba, 2015). To regularise
the network, we use early stopping on a validation set, as well as an $L_{2}$
penalty corresponding to a normal prior on $\beta$.

Figure 3: Our Bayesian mean shrinkage pooling model. This diagram takes
$m_{0}=\mathbf{0}$, $\eta=1$ and $\mathbf{u}=\mathbf{z}$, so that
$R=R_{\mathbf{z}}=R_{\mathbf{zz}}$, and $K_{\mathbf{z}}=K$.

###  3.2 Bayesian Linear Regression Model

The most obvious approach to adding uncertainty to the model of Section 3.1 is
to encode uncertainty over regression parameters $\beta$ only, as follows:

$\beta\sim\mathcal{N}(0,\rho^{2})\qquad
y_{i}\mid\mathbf{x}_{i},\beta\sim\operatorname{\mathcal{N}}(\beta^{\mathsf{T}}\hat{\mu}_{i},\sigma^{2}).$

This is essentially Bayesian linear regression on the empirical mean
embeddings, and is closely related to the model of Flaxman et al. (2015).
Here, we are working directly with the finite-dimensional $\hat{\mu}_{i}$,
unlike the infinite-dimensional $\mu_{i}$ before. Due to the conjugacy of the
model, we can easily obtain the predictive distribution $y_{i}\mid\bf{x_{i}}$,
integrating out the uncertainty over $\beta$. This provides us with
uncertainty intervals for the predictions $y_{i}$.

For model tuning, we can maximise the model evidence, i.e. the marginal log-
likelihood (see Bishop (2006) for details), and use backpropagation through
the network to learn $\sigma$ and $\rho$ and any kernel parameters of
interest.222Note that unlike the other models considered in this paper, we
cannot easily do minibatch stochastic gradient descent, as the marginal log-
likelihood does not decompose for each individual data point.

###  3.3 Bayesian Mean Shrinkage Model

A shortcoming of the prior models, and of the standard approach in Szábo et
al. (2016), is that they ignore uncertainty in the first level of estimation
due to varying number of samples in each bag. Ideally we would estimate not
just the mean embedding per bag, but also a measure of the sample variance, in
order to propagate this information regarding uncertainty from the bag size
through the model. Bayesian tools provide a natural framework for this
problem.

We can use the Bayesian nonparametric prior over kernel mean embeddings
(Flaxman et al., 2016) described in Section 2.2, and observe the empirical
embeddings at the landmark points $\mathbf{u_{i}}$. For $\mathbf{u_{i}}$, we
take a fixed set of landmarks, which we can choose via $k$-means clustering or
sample without replacement (Que and Belkin, 2016). Using the conjugacy of the
model to the Gaussian process prior
$\mu_{i}\sim\operatorname{\mathcal{GP}}(m_{0},\eta r(.,.))$, we obtain a
closed-form posterior Gaussian process whose evaluation at points
$\mathbf{h}=\\{h_{s}\\}_{s=1}^{n_{h}}$ is:

$\mu_{i}(\mathbf{h})\mid\mathbf{x_{i}}\sim\mathcal{N}\left(R_{\mathbf{h}}\left(R+\Sigma_{i}/N_{i}\right)^{-1}(\hat{\mu}_{i}-m_{0})+m_{0},\right.\\\
\left.R_{\mathbf{h}\mathbf{h}}-R_{\mathbf{h}}\left(R+\Sigma_{i}/N_{i}\right)^{-1}R_{\mathbf{h}}^{\top}\right)$

where $R_{st}=\eta r(u_{s},u_{t}),(R_{\mathbf{h}\mathbf{h}})_{st}=\eta
r(h_{s},h_{t}),(R_{\mathbf{h}})_{st}=\eta r(h_{s},u_{t})$, and
$\mathbf{x_{i}}$ denotes the set $\\{x_{j}^{i}\\}_{j=1}^{N_{i}}$. We take the
prior mean $m_{0}$ to be the average of the $\hat{\mu}_{i}$; under a linear
kernel $K$, this means we shrink predictions towards the mean prediction. Note
$\eta$ essentially controls the strength of the shrinkage: a smaller $\eta$
means we shrink more strongly towards $m_{0}$. We take $\Sigma_{i}$ to be the
average of the empirical covariance of
$\\{\varphi(x^{i}_{j})\\}_{j=1}^{N_{i}}$ across all bags, to avoid poor
estimation of $\Sigma_{i}$ for smaller bags. More intuition about the
behaviour of this estimator can be found in Appendix C.

Now, supposing we have normal observation error $\sigma^{2}$, and use a linear
kernel as our second level kernel $K$, we have:

$y_{i}\;|\>\mu_{i},f\sim\mathcal{N}\left(\langle
f,\mu_{i}\rangle_{\mathcal{H}_{k}},\sigma^{2}\right)$ (6)

where $f\in\mathcal{H}_{k}$. Clearly, this is difficult to work with; hence we
parameterise $f$ as $f=\sum_{\ell=1}^{s}\alpha_{\ell}k(\cdot,z_{\ell})$, where
$\mathbf{z}=\\{z_{\ell}\\}^{s}_{\ell=1}$ is a set of landmark points for $f$,
which we can learn or fix. (Appendix D gives a motivation for this
approximation using the representer theorem.) Using the reproducing property,
our likelihood model becomes:

$y_{i}\;|\>\mu_{i},\alpha\sim\operatorname{\mathcal{N}}\left(\alpha^{\mathsf{T}}\mu_{i}(\mathbf{z}),\sigma^{2}\right)$
(7)

where $\mu_{i}(\mathbf{z})=[\mu_{i}(z_{1}),\dots,\mu_{i}(z_{s})]^{\top}$. For
fixed $\alpha$ and $\mathbf{z}$ we can analytically integrate out the
dependence on $\mu_{i}$, and the predictive distribution of a bag label
becomes

$\displaystyle y_{i}\mid\mathbf{x}_{i},\alpha$
$\displaystyle\sim\operatorname{\mathcal{N}}(\xi_{i}^{\alpha},\nu_{i}^{\alpha})$
$\displaystyle\xi_{i}^{\alpha}$
$\displaystyle=\alpha^{\top}R_{\mathbf{z}}\left(R+\frac{\Sigma_{i}}{N_{i}}\right){}^{-1}(\hat{\mu}_{i}-m_{0})+\alpha^{\mathsf{T}}m_{0}$
$\displaystyle\nu_{i}^{\alpha}$
$\displaystyle=\alpha^{\mathsf{T}}\left(R_{\mathbf{zz}}-R_{\mathbf{z}}\left(R+\frac{\Sigma_{i}}{N_{i}}\right)^{-1}R_{\mathbf{z}}^{\mathsf{T}}\right)\alpha+\sigma^{2}.$

The prior $\alpha\sim\mathcal{N}(0,\rho^{2}K_{\mathbf{z}}^{-1})$, where
$K_{\mathbf{z}}$ is the kernel matrix on $\mathbf{z}$, gives the standard
regularisation on $f$ of $\lVert f\rVert_{{\mathcal{H}_{k}}}^{2}$. The log-
likelihood objective becomes

$\frac{1}{2}\sum_{i=1}^{n}\left\\{\log\nu_{i}^{\alpha}+\frac{\left(y_{i}-\xi_{i}^{\alpha}\right)^{2}}{\xi_{i}^{\alpha}}\right\\}+\frac{\alpha^{\mathsf{T}}K_{\mathbf{z}}\alpha}{2\rho^{2}}.$

We can use backpropagation to learn the parameters $\alpha$, $\sigma$, and if
we wish $\eta$, $\mathbf{z}$, and any kernel parameters. The full model is
illustrated in Figure 3. This approach allows us to directly encode
uncertainty based on bag size in the objective function, and gives
probabilistic predictions.

###  3.4 Bayesian Distribution Regression

It is natural to combine the two Bayesian models above, fully propagating
uncertainty in estimation of the mean embedding and of the regression
coefficients $\alpha$. Unfortunately, conjugate Bayesian inference is no
longer available. Thus, we consider a Markov Chain Monte Carlo (MCMC) sampling
based approach, and here use Hamiltonian Monte Carlo (HMC) for efficient
inference, though any MCMC-type scheme would work. Whereas inference above
used gradient descent to maximise the marginal likelihood, with the gradient
calculated using automatic differentiation, here we use automatic
differentiation to calculate the gradient of the joint log-likelihood and
follow this gradient as we perform sampling over the parameters we wish to
infer.

We can still exploit the conjugacy of the mean shrinkage layer, obtaining an
analytic posterior over the mean embeddings. Conditional on the mean
embeddings, we have a Bayesian linear regression model with parameters
$\alpha$. We sample this model with the NUTS HMC sampler (Hoffman and Gelman,
2014; Stan Development Team, 2014).

##  4 RELATED WORK

As previously mentioned, Szábo et al. (2016) provides a thorough learning-
theoretic analysis of the regression model discussed in Section 2.3. This
formalism considering a kernel method on distributions using their embedding
representations, or various scalable approximations to it, has been widely
applied (e.g. Muandet et al., 2012; Yoshikawa et al., 2014; Flaxman et al.,
2015; Jitkrittum et al., 2015; Lopez-Paz et al., 2015; Mitrovic et al., 2016).
There are also several other notions of similarities on distributions in use
(not necessarily falling within the framework of kernel methods and RKHSs), as
well as local smoothing approaches, mostly based on estimates of various
probability metrics (Moreno et al., 2003; Jebara et al., 2004; Póczos et al.,
2011; Oliva et al., 2013; Poczos et al., 2013; Kusner et al., 2015). For a
partial overview, see Sutherland (2016).

Other related problems of learning on instances with group-level labels
include learning with label proportions (Quadrianto et al., 2009; Patrini et
al., 2014), ecological inference (King, 1997; Gelman et al., 2001),
pointillistic pattern search (Ma et al., 2015), multiple instance learning
(Dietterich et al., 1997; Kück and de Freitas, 2005; Zhou et al., 2009;
Krummenacher et al., 2013) and learning with sets (Zaheer et al., 2017).333For
more, also see giorgiopatrini.org/nips15workshop.

There have also been some Bayesian approaches in related contexts, though most
do not follow our setting where the label is a function of the underlying
distribution rather than the observed sample set. Kück and de Freitas (2005)
consider an MCMC method with group-level labels but focus on individual-level
classifiers, while Jackson et al. (2006) use hierarchical Bayesian models on
both individual-level and aggregate data for ecological inference.

Jitkrittum et al. (2015) and Flaxman et al. (2015) quantify the uncertainty of
distribution regression models by interpreting the kernel ridge regression on
embeddings as Gaussian process regression. However, the former’s setting has
no uncertainty in the mean embeddings, while the latter’s treats empirical
embeddings as fixed inputs to the learning problem (as in Section 3.2).

There has also been generic work on input uncertainty in Gaussian process
regression (Girard, 2004; Damianou et al., 2016). These methods could provide
a framework towards allowing for second-level kernels in our models. One could
also, though, consider regression with uncertain inputs as a special case of
distribution regression, where the label is a function of the distribution’s
mean and $N_{i}=1$.

##  5 EXPERIMENTS

We will now demonstrate our various Bayesian approaches: the mean-shrinkage
pooling method with $r=k$ (shrinkage) and with $r(x,x^{\prime})=\int
k(x,z)k(z,x^{\prime})\nu(\mathrm{d}z)$ for $\nu$ proportional to a Gaussian
measure (shrinkageC), Bayesian linear regression (BLR), and the full Bayesian
distribution regression model with $r=k$ (BDR). We also compare the non-
Bayesian baselines RBF network (Section 3.1) and freq-shrinkage, which uses
the shrinkage estimator of Muandet et al. (2014) to estimate mean embeddings.
Code for our methods and to reproduce the experiments is available at
https://github.com/hcllaw/bdr.

We first demonstrate the characteristics of our models on a synthetic dataset,
and then evaluate them on a real life age prediction problem. Throughout, for
simplicity, we take $\mathbf{u}=\mathbf{z}$, i.e.
$R=R_{\mathbf{z}}=R_{\mathbf{zz}}$, and $K_{\mathbf{z}}=K$ – although
$\mathbf{u}$ and $\mathbf{z}$ could be different, with $\mathbf{z}$ learnt.
Here $k$ is the standard RBF kernel. We tune the learning rate, number of
landmarks, bandwidth of the kernel and regularisation parameters on a
validation set. For BDR, we use weakly informative normal priors (possibly
truncated at zero); for other models, we learn the remaining parameters.

###  5.1 Gamma Synthetic Data

We create a synthetic dataset by repeatedly sampling from the following
hierarchical model, where $y_{i}$ is the label for the $i$th bag, each
$x^{i}_{j}\in{\mathbb{R}}^{5}$ has entries i.i.d. according to the given
distribution, and $\varepsilon$ is an added noise term which differs for the
two experiments below:

$\displaystyle y_{i}$ $\displaystyle\sim\mbox{Uniform}(4,8)$
$\displaystyle\left[x^{i}_{j}\right]_{\ell}\mid y_{i}$
$\displaystyle\overset{\mathrm{iid}}{\sim}\frac{1}{y_{i}}\left[\Gamma\left(\frac{y_{i}}{2},\frac{1}{2}\right)\right]+\varepsilon\text{
for }j\in[N_{i}],\ell\in[5].$

In these experiments, we generate $1\,000$ bags for training, $500$ bags for a
validation set for parameter tuning, $500$ bags to use for early-stopping of
the models, and $1\,000$ bags for testing. Tuning is performed to maximize
log-likelihoods for Bayesian models, MSE for non-Bayesian models. Landmark
points $\mathbf{u}$ are chosen via $k$-means (fixed across all models). We
also show results of the Bayes-optimal model, which gives true posteriors
according to the data-generating process; this is the best performance any
model could hope to achieve. Our learning models, which treat the inputs as
five-dimensional, fully nonparametric distributions, are at a substantial
disadvantage even in how they view the data compared to this true model.

Figure 4: Top: negative log-likelihood. Bottom: mean-squared error. For
context, performance of the Bayes-optimal predictor is also shown, and for NLL
‘uniform’ shows the performance of a uniform prediction on the possible
labels. For MSE, the constant overall mean label predictor achieves about 1.3.

Figure 5: Predictions for the varying bag size experiment of Section 5.1. Each
column corresponds to a single prediction method. Each point in an image
represents a single bag, with its horizontal position the true label $y_{i}$,
and its vertical position the predicted label. The black lines show
theoretical perfect predictions. The rows represent different subsets of the
data: the first row shows all bags, the second only bags with $N_{i}=5$, and
so on. Colours represent the predictive standard deviation of each point.

#### Varying bag size: Uncertainty in the inputs.

In order to study the behaviour of our models with varying bag size, we fix
four sizes $N_{i}\in\\{5,20,100,1\,000\\}$. For each generated dataset, $25\%$
of the bags have $N_{i}=20$, and $25\%$ have $N_{i}=100$. Among the other half
of the data, we vary the ratio of $N_{i}=5$ and $N_{i}=1\,000$ bags to
demonstrate the methods’ efficacy at dealing with varied bag sizes: we let
$s_{5}$ be the overall percentage of bags with $N_{i}=5$, ranging from
$s_{5}=0$ (in which case no bags have size $N_{i}=5$) to $s_{5}=50$ (in which
case $50\%$ of the overall bags have size $N_{i}=5$). Here we do not add
additional noise: $\varepsilon=0$.

Results are shown in Figure 4. BDR and shrinkage methods, which take into
account bag size uncertainty, perform well here compared to the other methods.
The full BDR model very slightly outperforms the Bayesian shrinkage models in
both likelihood and in mean-squared error; frequentist shrinkage slightly
outperforms the Bayesian shrinkage models in MSE, likely because it is tuned
for that metric. We also see that the choice of $r$ affects the results; $r=k$
does somewhat better.

Figure 5 demonstrates in more detail the difference between these models. It
shows test set predictions of each model on the bags of different sizes. Here,
we can see explicitly that the shrinkage and BDR models are able to take into
account the bag size, with decreasing variance for larger bag sizes, while the
BLR model gives the same variance for all outputs. Furthermore, the shrinkage
and BDR models can shrink their predictions towards the mean more for smaller
bags than larger ones: this improves performance on the small bags while still
allowing for good predictions on large bags, contrary to the BLR model.

#### Fixed bag size: Uncertainty in the regression model.

The previous experiment showed the efficacy of the shrinkage estimator in our
models, but demonstrated little gain from posterior inference for regression
weights $\beta$ over their MAP estimates, i.e. there is no discernible
improvement of BLR over RBF network. To isolate the effect of quantifying
uncertainty in the regression model, we now consider the case where there is
no variation in bag size at all and normal noise is added onto the
observations. In particular we take $N_{i}=1000$ and
$\varepsilon\sim\operatorname{\mathcal{N}}(0,1)$, and sample landmarks
randomly from the training set.

Results are shown in Table 1. Here, BLR or BDR outperform all other methods on
all runs, highlighting that uncertainty in the regression model is also
important for predictive performance. Importantly, the BDR method performs
well in this regime as well as in the previous one.

Table 1: Results on the fixed bag size dataset, over $10$ dataset draws (standard deviations in parentheses). BLR/BDR perform best on all runs in both metrics. METHOD | MSE | NLL
---|---|---
Optimal | 0.170 (0.009) | 0.401 (0.018)
RBF network | 0.235 (0.014) | –
freq-shrinkage | 0.232 (0.012) | –
shrinkage | 0.237 (0.014) | 0.703 (0.027)
shrinkageC | 0.236 (0.013) | 0.700 (0.029)
BLR | 0.228 (0.012) | 0.681 (0.025)
BDR | 0.227 (0.012) | 0.683 (0.025)

###  5.2 IMDb-WIKI: Age Estimation

Table 2: Results on the grouped IMDb-WIKI dataset over ten runs (standard deviations in parentheses). Here shrinkage methods perform the best across all $10$ runs. METHOD | RMSE | NLL
---|---|---
CNN | 10.25 (0.22) | 3.80 (0.034)
RBF network | 9.51 (0.20) | –
freq-shrinkage | 9.22 (0.19) | –
shrinkage | 9.28 (0.20) | 3.54 (0.021)
BLR | 9.55 (0.19) | 3.68 (0.021)

We now demonstrate our methods on a celebrity age estimation problem, using
the IMDb-WIKI database (Rothe et al., 2016) which consists of $397\,949$
images of $19\,545$ celebrities444We used only the IMDb images, and removed
some implausible images, including one of a cat and several of people with
supposedly negative age, or ages of several hundred years., with corresponding
age labels. This database was constructed by crawling IMDb for images of its
most popular actors and directors, with potentially many images for each
celebrity over time. Rothe et al. (2016) use a convolutional neural network
(CNN) with a VGG-16 architecture to perform 101-way classification, with one
class corresponding to each age in $\\{0,\dots,100\\}$.

We take a different approach, and assume that we are given several images of a
single individual (i.e. samples from the distribution of celebrity images),
and are asked to predict their mean age based on several pictures. For
example, we have 757 images of Brad Pitt from age 27 up to 51, while we have
only 13 images of Chelsea Peretti at ages 35 and 37. Note that 22.5% of bags
have only a single image. We obtain $19\,545$ bags, with each bag containing
between $1$ and $796$ images of a particular celebrity, and the corresponding
bag label calculated from the average of the age labels of the images inside
each bag.

In particular, we use the representation $\varphi(x)$ learnt by the CNN in
Rothe et al. (2016), where $\varphi(x):\ \mathbb{R}^{256\times
256}\to\mathbb{R}^{4096}$ maps from the pixel space of images to the CNN’s
last hidden layer. With these new representations, we can now treat them as
inputs to our radial basis network, shrinkage (taking $r=k$ here) and BLR
models. Although we could also use the full BDR model here, due to the
computational time and memory required to perform proper parameter tuning, we
relegate this to a later study.

We use $9\,820$ bags for training, $2\,948$ bags for early stopping, $2\,946$
for validation and $3\,928$ for testing. Landmarks are sampled without
replacement from the training set.

We repeat the experiment on $10$ different splits of the data, and report the
results in Table 2. The baseline CNN results give performance by averaging the
predictive distribution from the model of Rothe et al. (2016) for each image
of a bag; note that this model was trained on all of the images used here.
From Table 2, we can see that the shrinkage methods have the best performance;
they outperforms all other methods in all $10$ splits of the dataset, in both
metrics. Non-Bayesian shrinkage again yields slightly better RMSEs, likely
because it is tuned for that metric. This demonstrates that modelling bag size
uncertainty is vital.

##  6 CONCLUSION

Supervised learning on groups of observations using kernel mean embeddings
typically disregards sampling variability within groups. To handle this
problem, we construct Bayesian approaches to modelling kernel mean embeddings
within a regression model, and investigate advantages of uncertainty
propagation within different components of the resulting distribution
regression. The ability to take into account the uncertainty in mean embedding
estimates is demonstrated to be key for constructing models with good
predictive performance when group sizes are highly imbalanced. We also
demonstrate that the results of a complex neural network model for age
estimation can be improved by shrinkage.

Our models employ a neural network formulation to provide more expressive
feature representations and learn discriminative embeddings. Doing so makes
our model easy to extend to more complicated featurisations than the simple
RBF network used here. By training with backpropagation, or via approximate
Bayesian methods such as variational inference, we can easily ‘learn the
kernel’ within our framework, for example fine-tuning the deep network of
Section 5.2 rather than using a pre-trained model. We can also apply our
networks to structured settings, learning regression functions on sets of
images, audio, or text. Such models naturally fit into the empirical Bayes
framework.

On the other hand, we might extend our model to more Bayesian feature learning
by placing priors over the kernel hyperparameters, building on classic work on
variational approaches (Barber and Schottky, 1998) and fully Bayesian
inference (Andrieu et al., 2001) in RBF networks. Such approaches are also
possible using other featurisations, e.g. random Fourier features (as in Oliva
et al., 2015).

Future distribution regression approaches will need to account for uncertainty
in observation of the distribution. Our methods provide a strong, generic
building block to do so.

## References

  * Andrieu et al. (2001) Christophe Andrieu, Nando De Freitas, and Arnaud Doucet.  Robust full bayesian learning for radial basis networks.  _Neural Computation_ , 13(10):2359–2407, 2001\.
  * Barber and Schottky (1998) David Barber and Bernhard Schottky.  Radial basis functions: a bayesian treatment.  _NIPS_ , pages 402–408, 1998.
  * Bishop (2006) C.M. Bishop.  _Pattern recognition and machine learning_.  Springer New York, 2006.
  * Broomhead and Lowe (1988) David S Broomhead and David Lowe.  Radial basis functions, multi-variable functional interpolation and adaptive networks.  Technical report, DTIC Document, 1988.
  * Damianou et al. (2016) Andreas C. Damianou, Michalis K. Titsias, and Neil D. Lawrence.  Variational inference for latent variables and uncertain inputs in Gaussian processes.  _JMLR_ , 17(42):1–62, 2016.
  * Dietterich et al. (1997) Thomas G Dietterich, Richard H Lathrop, and Tomás Lozano-Pérez.  Solving the multiple instance problem with axis-parallel rectangles.  _Artificial intelligence_ , 89(1):31–71, 1997\.
  * Flaxman et al. (2015) Seth Flaxman, Yu-Xiang Wang, and Alexander J Smola.  Who supported Obama in 2012?: Ecological inference through distribution regression.  In _KDD_ , pages 289–298. ACM, 2015.
  * Flaxman et al. (2016) Seth Flaxman, Dino Sejdinovic, John P. Cunningham, and Sarah Filippi.  Bayesian learning of kernel embeddings.  In _UAI_ , 2016.
  * Flaxman et al. (2016) Seth Flaxman, Dougal J. Sutherland, Yu-Xiang Wang, and Yee-Whye Teh.  Understanding the 2016 US presidential election using ecological inference and distribution regression with census microdata.  2016\.  arXiv:1611.03787.
  * Gärtner (2008) Thomas Gärtner.  _Kernels for Structured Data_ , volume 72.  World Scientific, Series in Machine Perception and Artificial Intelligence, 2008.
  * Gelman et al. (2001) Andrew Gelman, David K Park, Stephen Ansolabehere, Phillip N Price, and Lorraine C Minnite.  Models, assumptions and model checking in ecological regressions.  _Journal of the Royal Statistical Society: Series A (Statistics in Society)_ , 164(1):101–118, 2001.
  * Girard (2004) Agathe Girard.  _Approximate methods for propagation of uncertainty with Gaussian process models_.  PhD thesis, University of Glasgow, 2004.
  * Hoffman and Gelman (2014) Matthew D. Hoffman and Andrew Gelman.  The no-U-turn sampler: Adaptively setting path lengths in Hamiltonian Monte Carlo.  _JMLR_ , pages 1593–1623, 2014.
  * Jackson et al. (2006) Christopher Jackson, Nicky Best, and Sylvia Richardson.  Improving ecological inference using individual-level data.  _Statistics in medicine_ , 25(12):2136–2159, 2006\.
  * Jebara et al. (2004) Tony Jebara, Risi Imre Kondor, and Andrew Howard.  Probability product kernels.  _JMLR_ , 5:819–844, 2004.
  * Jitkrittum et al. (2015) Wittawat Jitkrittum, Arthur Gretton, Nicolas Heess, S. M. Ali Eslami, Balaji Lakshminarayanan, Dino Sejdinovic, and Zoltán Szabó.  Kernel-Based Just-In-Time Learning for Passing Expectation Propagation Messages.  In _UAI_ , 2015.
  * King (1997) Gary King.  _A Solution to the Ecological Inference Problem_.  Princeton University Press, 1997.  ISBN 0691012407.
  * Kingma and Ba (2015) Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _ICLR_ , 2015.  arXiv:1412.6980.
  * Krummenacher et al. (2013) Gabriel Krummenacher, Cheng Soon Ong, and Joachim M Buhmann.  Ellipsoidal multiple instance learning.  In _ICML (2)_ , pages 73–81, 2013.
  * Kück and de Freitas (2005) Hendrik Kück and Nando de Freitas.  Learning about individuals from group statistics.  In _UAI_ , pages 332–339, 2005.
  * Kusner et al. (2015) Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger.  From word embeddings to document distances.  In _ICML_ , pages 957–966, 2015.
  * Law et al. (2017) H. C. L. Law, C. Yau, and D. Sejdinovic.  Testing and learning on distributions with symmetric noise invariance.  In _NIPS_ , 2017.  arXiv:1703.07596.
  * Lopez-Paz et al. (2015) David Lopez-Paz, Krikamol Muandet, Bernhard Schölkopf, and Ilya Tolstikhin.  Towards a learning theory of cause-effect inference.  In _ICML_ , 2015.
  * Lukić and Beder (2001) Milan Lukić and Jay Beder.  Stochastic processes with sample paths in reproducing kernel hilbert spaces.  _Transactions of the American Mathematical Society_ , 353(10):3945–3969, 2001.
  * Ma et al. (2015) Yifei Ma, Dougal J. Sutherland, Roman Garnett, and Jeff Schneider.  Active pointillistic pattern search.  In _AISTATS_ , 2015.
  * Mitrovic et al. (2016) J. Mitrovic, D. Sejdinovic, and Y.W. Teh.  DR-ABC: Approximate Bayesian Computation with Kernel-Based Distribution Regression.  In _ICML_ , pages 1482–1491, 2016.
  * Moreno et al. (2003) Pedro J Moreno, Purdy P Ho, and Nuno Vasconcelos.  A Kullback-Leibler divergence based kernel for SVM classification in multimedia applications.  In _NIPS_ , 2003.
  * Muandet et al. (2012) Krikamol Muandet, Kenji Fukumizu, Francesco Dinuzzo, and Bernhard Schölkopf.  Learning from distributions via support measure machines.  In _NIPS_ , 2012.  arXiv:1202.6504.
  * Muandet et al. (2014) Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, Arthur Gretton, and Bernhard Schoelkopf.  Kernel mean estimation and stein effect.  In _ICML_ , 2014.
  * Ntampaka et al. (2015) Michelle Ntampaka, Hy Trac, Dougal J. Sutherland, Nicholas Battaglia, Barnabás Póczos, and Jeff Schneider.  A machine learning approach for dynamical mass measurements of galaxy clusters.  _The Astrophysical Journal_ , 803(2):50, 2015\.  ISSN 1538-4357.  arXiv:1410.0686.
  * Ntampaka et al. (2016) Michelle Ntampaka, Hy Trac, Dougal J. Sutherland, S. Fromenteau, B. Poczos, and Jeff Schneider.  Dynamical mass measurements of contaminated galaxy clusters using machine learning.  _The Astrophysical Journal_ , 831(2):135, 2016\.  arXiv:1509.05409.
  * Oliva et al. (2013) Junier B Oliva, Barnabás Póczos, and Jeff Schneider.  Distribution to distribution regression.  In _ICML_ , 2013.
  * Oliva et al. (2015) Junier B Oliva, Avinava Dubey, Barnabás Póczos, Jeff Schneider, and Eric P Xing.  Bayesian nonparametric kernel-learning.  In _AISTATS_ , 2015.  arXiv:1506.08776.
  * Patrini et al. (2014) Giorgio Patrini, Richard Nock, Tiberio Caetano, and Paul Rivera.  (Almost) no label no cry.  In _NIPS_. 2014.
  * Pillai et al. (2007) Natesh S Pillai, Qiang Wu, Feng Liang, Sayan Mukherjee, and Robert L Wolpert.  Characterizing the function space for bayesian kernel models.  _JMLR_ , 8(Aug):1769–1797, 2007.
  * Póczos et al. (2011) Barnabás Póczos, Liang Xiong, and Jeff Schneider.  Nonparametric divergence estimation with applications to machine learning on distributions.  In _UAI_ , 2011.
  * Poczos et al. (2013) Barnabas Poczos, Aarti Singh, Alessandro Rinaldo, and Larry Wasserman.  Distribution-free distribution regression.  In _AISTATS_ , pages 507–515, 2013.  arXiv:1302.0082.
  * Quadrianto et al. (2009) Novi Quadrianto, Alex J Smola, Tiberio S Caetano, and Quoc V Le.  Estimating labels from label proportions.  _JMLR_ , 10:2349–2374, 2009.
  * Que and Belkin (2016) Qichao Que and Mikhail Belkin.  Back to the future: Radial basis function networks revisited.  In _AISTATS_ , 2016.
  * Rahimi and Recht (2007) Ali Rahimi and Benjamin Recht.  Random features for large-scale kernel machines.  In _NIPS_ , pages 1177–1184, 2007.
  * Ramdas and Wehbe (2015) Aaditya Ramdas and Leila Wehbe.  Nonparametric independence testing for small sample sizes.  In _IJCAI_ , 2015.  arXiv:1406.1922.
  * Rothe et al. (2016) Rasmus Rothe, Radu Timofte, and Luc Van Gool.  Deep expectation of real and apparent age from a single image without facial landmarks.  _International Journal of Computer Vision (IJCV)_ , July 2016.
  * Saunders et al. (1998) Craig Saunders, Alexander Gammerman, and Volodya Vovk.  Ridge regression learning algorithm in dual variables.  In _ICML_ , 1998.
  * Schölkopf et al. (2001) Bernhard Schölkopf, Ralf Herbrich, and Alex J. Smola.  A generalized representer theorem.  In _COLT_ , 2001.
  * Sriperumbudur et al. (2010) Bharath K Sriperumbudur, Arthur Gretton, Kenji Fukumizu, Bernhard Schölkopf, and Gert RG Lanckriet.  Hilbert space embeddings and metrics on probability measures.  _JMLR_ , 99:1517–1561, 2010.
  * Stan Development Team (2014) Stan Development Team.  Stan: A c++ library for probability and sampling, version 2.5.0, 2014\.  URL http://mc-stan.org/.
  * Steinwart (2017) Ingo Steinwart.  Convergence types and rates in generic Karhunen-Loéve expansions with applications to sample path properties.  _arXiv preprint arXiv:1403.1040v3_ , March 2017.
  * Sutherland (2016) Dougal J. Sutherland.  _Scalable, Flexible, and Active Learning on Distributions_.  PhD thesis, Carnegie Mellon University, 2016.
  * Szábo et al. (2016) Zoltán Szábo, Bharath K. Sriperumbudur, Barnabás Póczos, and Arthur Gretton.  Leraning theory for distribution regression.  _JMLR_ , 17(152):1–40, 2016.  arXiv:1411.2066.
  * Wahba (1990) Grace Wahba.  _Spline models for observational data_ , volume 59.  Siam, 1990.
  * Yoshikawa et al. (2014) Yuya Yoshikawa, Tomoharu Iwata, and Hiroshi Sawada.  Latent support measure machines for bag-of-words data classification.  In _NIPS_ , pages 1961–1969, 2014.
  * Zaheer et al. (2017) Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola.  Deep sets.  In _NIPS_ , 2017.
  * Zhou et al. (2009) Zhi-Hua Zhou, Yu-Yin Sun, and Yu-Feng Li.  Multi-instance learning by treating instances as non-iid samples.  In _ICML_ , 2009.

##  Appendix A Choice of $r(\cdot,\cdot)$ to ensure
${\mu_{{\mathsf{P}}}}\in{\mathcal{H}_{k}}$

We need to choose an appropriate covariance function $r$, such that
${\mu_{{\mathsf{P}}}}\in{\mathcal{H}_{k}}$, where
${\mu_{{\mathsf{P}}}}\sim\operatorname{\mathcal{GP}}(0,r(\cdot,\cdot))$. In
particular, it is for infinite-dimensional RKHSs not sufficient to define
$r(\cdot,\cdot)=k(\cdot,\cdot)$, as draws from this particular prior are no
longer in ${\mathcal{H}_{k}}$ (Wahba, 1990) (but see below). However, we can
construct

$r(x,y)=\int k(x,z)k(z,y)\nu(dz)$ (8)

where $\nu$ is any finite measure on $\mathcal{X}$. This then ensures
${\mu_{{\mathsf{P}}}}\in{\mathcal{H}_{k}}$ with probability $1$ by the nuclear
dominance (Lukić and Beder, 2001; Pillai et al., 2007) for any stationary
kernel $k$. In particular, Flaxman et al. (2016) provides details when $k$ is
a squared exponential kernel defined by

$k(x,y)=\exp(-\frac{1}{2}(x-y)^{\top}\Sigma_{k}^{-1}(x-y))\quad\quad
x,y\in\mathbb{R}^{p}$

and $\nu(dz)=\exp\left(-\frac{||z||^{2}_{2}}{2\ell^{2}}\right)dz$, i.e. it is
proportional to a Gaussian measure on $\mathbb{R}^{d}$, which provides
$r(\cdot,\cdot)$ with a non-stationary component. In this paper, we take
$\Sigma_{k}=\sigma^{2}I_{p}$, where $\sigma^{2}$ and $\ell$ are tuning
parameters, or parameters that we learn.

Here, the above holds for a general set of stationary kernels, but note that
by taking a convolution of a kernel with itself, it might make the space of
functions that we consider overly smooth (i.e. concentrated on a small part of
$\mathcal{H}_{k}$). In this work, however, we consider only the Gaussian RBF
kernel $k$. In fact, recent work (Steinwart, 2017, Theorem 4.2) actually shows
that in this case, the sample paths almost surely belong to (interpolation)
spaces which are infinitesimally larger than the RKHS of the Gaussian RBF
kernel. This suggests that we can choose $r$ to be an RBF kernel with a length
scale that is infinitesimally bigger than that of $k$; thus, in practice,
taking $r=k$ would suffice and we do observe that it actually performs better
(Fig. 4).

##  Appendix B Framework for Binary Classification

Suppose that our labels $y_{i}\in\\{0,1\\}$, i.e. we are in a binary
classification framework. Then a simple approach to accounting for uncertainty
in the regression parameters is to use bayesian logistic regression, putting
priors on $\beta$, i.e.

$\displaystyle\beta$ $\displaystyle\sim$
$\displaystyle\mathcal{N}(0,\rho^{2})$ $\displaystyle y_{i}$
$\displaystyle\sim$ $\displaystyle Ber(\pi_{i}),\text{ where
}\log\left(\dfrac{\pi_{i}}{1-\pi_{i}}\right)=\beta^{\top}\hat{\mu}_{i}$

however for the mean shrinkage pooling model, if we use the above
$y_{i}\;|\>\mu_{i},\alpha$, we would not be able to obtain an analytical
solution for $p(y_{i}|\mathbf{x_{i}},\alpha)$. Instead we use the probit link
function, as given by:

$Pr(y_{i}=1|\mu_{i},\alpha)=\Phi\left(\alpha^{\top}\mu_{i}(\mathbf{z})\right)$

where $\Phi$ denotes the Cumulative Distribution Function (CDF) of a standard
normal distribution, with
$\mu_{i}(\mathbf{z})=[\mu_{i}(z_{1}),\dots,\mu_{i}(z_{s})]^{\top}$. Then as
before we have

$\mu_{i}(\mathbf{z})\mid\mathbf{x_{i}}\sim\mathcal{N}\left(M_{i},C_{i}\right)$

with $M_{i}$ and $C_{i}$ as defined in section 3.3. Hence, as before

$\displaystyle Pr(y_{i}=1|{\bf x}_{i},\alpha)$ $\displaystyle=$
$\displaystyle\int Pr(y_{i}=1|\mu_{i},\alpha)p(\mu_{i}(\mathbf{z})|{\bf
x}_{i})d\mu_{i}(\mathbf{z})$ $\displaystyle=$ $\displaystyle
c\int\Phi(\alpha^{\top}\mu_{i}(\mathbf{z}))\exp\\{-\frac{1}{2}(\mu_{i}(\mathbf{z})-M_{i})^{\top}C_{i}^{-1}(\mu_{i}(\mathbf{z})-M_{i})\\}d\mu_{i}(\mathbf{z})$
$\displaystyle(\text{with }l_{i}=\mu_{i}(\mathbf{z})-M_{i})$ $\displaystyle=$
$\displaystyle
c\int\Phi(\alpha^{\top}(l_{i}+M_{i}))\exp\\{-\frac{1}{2}(l_{i})^{\top}C_{i}^{-1}(l_{i})\\}dl_{i}$
$\displaystyle=$ $\displaystyle Pr(Y\leq\alpha^{\top}(l_{i}+M_{i}))$

Note here $Y\sim\mathcal{N}(0,1)$ and $l_{i}\sim\mathcal{N}(0,\Sigma_{i})$
Then expanding and rearranging

$\displaystyle Pr(y_{i}=1|{\bf x}_{i},\alpha)$ $\displaystyle=$ $\displaystyle
Pr(Y-\alpha^{\top}l_{i}\leq\alpha^{\top}M_{i})$

Note that since $Y$ and $l_{i}$ independent normal r.v.,
$Y-\alpha^{\top}l_{i}\sim\mathcal{N}(0,1+\alpha^{\top}C_{i}\alpha^{\top})$.
Let $T$ be standard normal, then we have:

$\displaystyle Pr(y_{i}=1|{\bf x}_{i},\alpha)$ $\displaystyle=$ $\displaystyle
Pr(\sqrt{1+\alpha^{\top}C_{i}\alpha}\ T\leq\alpha^{\top}M_{i})$
$\displaystyle=$ $\displaystyle
Pr(T\leq\frac{\alpha^{\top}M_{i}}{\sqrt{1+\alpha^{\top}C_{i}\alpha}})$
$\displaystyle=$
$\displaystyle\Phi\left(\frac{\alpha^{\top}M_{i}}{\sqrt{1+\alpha^{\top}C_{i}\alpha}}\right)$

Hence, we also have:

$\displaystyle Pr(y_{i}=0|{\bf x}_{i},\alpha)$ $\displaystyle=$ $\displaystyle
1-\Phi\left(\frac{\alpha^{\top}M_{i}}{\sqrt{1+\alpha^{\top}C_{i}\alpha}}\right)$

Now placing the prior $\alpha\sim\mathcal{N}(0,\rho^{2}K_{\mathbf{z}}^{-1})$,
we have the following MAP objective:

$\displaystyle J(\alpha)$ $\displaystyle=$
$\displaystyle\log\left[p(\alpha)\prod_{i=1}^{n}p(y_{i}|{\bf
x}_{i},\alpha)\right]$ $\displaystyle=$
$\displaystyle\sum_{i=1}^{n}(1-y_{i})\log(1-\Phi\left(\frac{\alpha^{\top}M_{i}}{\sqrt{1+\alpha^{\top}C_{i}\alpha}}\right))$
$\displaystyle+y_{i}\log(\Phi\left(\frac{\alpha^{\top}M_{i}}{\sqrt{1+\alpha^{\top}C_{i}\alpha}}\right))+\frac{1}{\rho^{2}}\alpha^{\top}K_{\mathbf{z}}\alpha$

Since we have an analytical solution for $Pr(y_{i}=0|{\bf x}_{i},\alpha)$, we
can also use this in HMC for BDR.

##  Appendix C Some more intuition on the shrinkage estimator

In this section, we provide some intuition behind the shrinkage estimator in
section 3.3. Here, for simplicity, we choose $\Sigma_{i}=\tau^{2}I$ for all
bag $i$, and $m_{0}=0$, and consider the case where $\mathbf{z}=\mathbf{u}$,
i.e. $R=R_{\mathbf{z}}=R_{\mathbf{zz}}$. We can then see that if $R$ has
eigendecomposition $U\Lambda U^{T}$, with
$\Lambda=\operatorname{\mathrm{diag}}(\lambda_{k})$, the posterior mean is

$U\operatorname{\mathrm{diag}}\left(\frac{\lambda_{k}}{\lambda_{k}+\tau^{2}/N_{i}}\right)U^{T}(\hat{\mu}_{i}),$

so that large eigenvalues, $\lambda_{k}\gg\tau^{2}/N_{i}$, are essentially
unchanged, while small eigenvalues, $\lambda_{k}\ll\tau^{2}/N_{i}$, are shrunk
towards 0. Likewise, the posterior variance is

$U\operatorname{\mathrm{diag}}\\!\left(\lambda_{k}-\frac{\lambda_{k}^{2}}{\lambda_{k}+\frac{\tau^{2}}{N_{i}}}\right)U^{T}=U\operatorname{\mathrm{diag}}\\!\left(\frac{1}{\frac{N_{i}}{\tau^{2}}+\frac{1}{\lambda_{k}}}\right)U^{T};$

its eigenvalues also decrease as $N_{i}/\tau^{2}$ increases.

##  Appendix D Alternative Motivation for choice of $f$

Here we provide an alternative motivation for the choice of
$f=\sum^{k}_{s=1}\alpha_{s}k(\cdot,z_{s})$. First, consider the following
Bayesian model with a linear kernel $K$ on $\mu_{i}$, where
$f:\mathcal{H}_{k}\rightarrow\mathbb{R}$:

$\displaystyle y_{i}\;|\>\mu_{i},f$ $\displaystyle\sim$
$\displaystyle\mathcal{N}\left(f(\mu_{i}),\sigma^{2}\right).$

Now considering the log-likelihood of
$\\{\mu,Y\\}=\\{\mu_{i},y_{i}\\}^{n}_{i=1}$ (supposing we have these exact
embeddings), we obtain:

$\log p(Y|\mu,f)=\sum^{n}_{i=1}-\frac{1}{2\sigma^{2}}(y_{i}-f(\mu_{i}))^{2}$

To avoid over-fitting, we place a Gaussian prior on $f$, i.e. $-\log
p(f)=\lambda||f||_{\mathcal{H}_{k}}+c$. Minimizing the negative log-likelihood
over $f\in\mathcal{H}_{k}$, we have:

$f^{*}=\text{argmin}_{f\in\mathcal{H}_{k}}\sum^{n}_{i=1}\frac{1}{2\sigma^{2}}(y_{i}-f(\mu_{i}))^{2}+\lambda||f||_{\mathcal{H}_{k}}$

Now this is in the form of an empirical risk minimisation problem. Hence using
the representer theorem (Schölkopf et al., 2001), we have that:

$\displaystyle f=\sum_{j=1}^{n}\gamma_{j}K(.,\mu_{j})$

i.e. we have a finite-dimensional problem to solve. Thus since $K$ is a linear
kernel:

$\displaystyle y_{i}\;|\>\mu_{i},\\{\mu_{j}\\}_{j=1}^{n},\gamma$
$\displaystyle\sim$
$\displaystyle\mathcal{N}\left(\sum_{j=1}^{n}\gamma_{j}\langle\mu_{i},\mu_{j}\rangle_{\mathcal{H}_{k}},\sigma^{2}\right).$

where $\langle\mu_{i},\mu_{j}\rangle_{\mathcal{H}_{k}}$ can be thought of as
the similarity between distributions.

Now we have the same $\operatorname{\mathcal{GP}}$ posterior as in Section
3.3, and we would like to compute $p(y_{i}|\mathbf{x_{i}},\gamma)$. This
suggests we need to integrate out $\mu_{1}$, …$\mu_{n}$. But it is unclear how
to perform this integration, since the $\mu_{i}$ follow Gaussian process
distributions. Hence we can take an approximation to $f$, i.e.
$f=\sum^{k}_{s=1}\alpha_{s}k(\cdot,z_{s})$, which would essentially give us a
dual method with a sparse approximation to $f$.


Selective classification is the task of rejecting inputs a model would predict incorrectly on through a trade-off between input space coverage and model accuracy. Current methods for selective classification impose constraints on either the model architecture or the loss function; this inhibits their usage in practice. In contrast to prior work, we show that state-of-the-art selective classification performance can be attained solely from studying the (discretized) training dynamics of a model. We propose a general framework that, for a given test input, monitors metrics capturing the disagreement with the final predicted label over intermediate models obtained during training; we then reject data points exhibiting too much disagreement at late stages in training. In particular, we instantiate a method that tracks when the label predicted during training stops disagreeing with the final predicted label. Our experimental evaluation shows that our method achieves state-of-the-art accuracy/coverage trade-offs on typical selective classification benchmarks. For example, we improve coverage on CIFAR-10/SVHN by 10.1%/1.5% respectively at a fixed target error of 0.5%.

§ INTRODUCTION

* Checkpointing a model has seen increasing utility in different contexts

* Selective Classification is desirable, but current SOTA (which improved dramatically over previous work) requires modifications to training; not clear it's always applicable

* We expand the literature using checkpoints for alternative use cases by showing how to achieve SOTA selective classification just from checkpoints. The key idea is that SC can be broken into a two step process: first removing points predictions didn't converge on, and removing points labels converged but to the wrong label. We propose a method that deals with the first by using checkpoints to capture convergence in labels.

* Our results on SC can also be seen as an extension to the line of work using variances to capture generalization (in particular variance in gradients); here we are interested in generalization to a specific data point and propose more datapoint specific variances to look at.

Machine learning (ML) is increasingly deployed for in high-stakes decision-making environments, where it is critical to detect inputs that the model could misclassify.
This is particularly true when deploying deep neural networks (DNNs) for applications with low tolerances for (classifying with a wrong label), such as healthcare <cit.>, self-driving <cit.>, and law <cit.>.
This problem setup is captured by the selective classification (SC) framework, which introduces a gating mechanism to abstain from predicting on individual test points <cit.>.
Specifically, SC aims to (i) only accept inputs on which the ML model would achieve high accuracy, while (ii) maintaining high coverage, accepting as many inputs as possible.

Current state-of-the-art selective classification techniques take one of two directions: the first augments the architecture of the underlying ML model <cit.>; the second trains the model using a purposefully adapted loss function <cit.>.
The unifying principle behind these methods is to modify the training stage in order to accommodate selective classification. This lack of practicality hinders the adoption of selective classification. In this work, we instead show that these modifications are unnecessary to attain state-of-the-art selective classification performance.

We present the first framework for selective classification based on neural network training dynamics (see Figure <ref>).
Typical DNNs are trained sequentially using stochastic gradient descent. Hence, as training goes on, the optimization process yields a sequence of intermediate models.
These models, in particular their predictions on a given test point, can then be used to construct a threshold-able score for selective classification.
Intuitively, if we understood defining traits of the intermediate predictions for correctly classified points, we can define a gating mechanisms to reject incorrectly classified points.

Through an analysis of neural network training dynamics,
we first observe that an important property of the intermediate predictions is to assess if they agree with the final prediction.
To that end, we introduce a set of assumptions on how the expected disagreement evolves during training for correctly classified points, alongside the variance of this disagreement. Using these assumptions, we can upper-bound the probability that the test point is correctly classified given its intermediate predictions.
Verifying these assumptions empirically, we observe that the expected disagreement is almost always $0$ (except initially) for correctly classified points, and the variance decreases in a convex manner.

These observations lead to our first proposed method, which imposes a threshold on the last intermediate model whose prediction disagrees with the final label prediction. Noting however the potential for noise inherent in training (and hence the intermediate models' predictions) to hamper the effectiveness of this method, we propose an improved method which thresholds a biased average based on multiple intermediate models disagreeing with the final label prediction. The latter method improves on
state-of-the-art (SOTA) accuracy/coverage trade-offs for selective classification.
Specifically, we improve on the previous SOTA by 10.1% on CIFAR10 and by 1.5% on SVHN at a fixed target error of 0.5%.

Our main contributions are as follows:

* We propose a novel method for SC based on neural network training dynamics (dubbed ) in which we devise an effecting scoring mechanism capturing label disagreement of intermediate models with the final prediction for a particular test point.

* We theoretically motivate the benefits of by deriving an upper-bound on the probability of a test point being correctly classified.

* We perform a comprehensive set of empirical experiments on established SC benchmarks. Our results obtained with demonstrate a highly favorable accuracy/coverage of -based selective classification, especially at low rates.

Our proposed method. During the model training stage, we store checkpoints of intermediate models. At inference time, given a test input, we compute various metrics capturing the agreement of intermediate predictions with the final model prediction. Data points with high agreement are accept while data points with low agreement are rejected.

§ BACKGROUND

Our work considers the standard supervised multi-class classification setup. We assume to have access to a dataset $D = \{(\bm{x}_i,y_i)\}_{i=1}^{M}$ consisting of data points $(\bm{x},y)$ with $\bm{x} \in \mathcal{X}$ and  $y \in \mathcal{Y}$. We refer to $\mathcal{X} = \mathbb{R}^d$ as the covariate space of dimensionality $d$ and $\mathcal{Y} = \{1, 2, \ldots, C\}$ as the label space consisting of $C$ classes. All data points $(\bm{x},y)$ are sampled independently from the underlying distribution $p$ defined over $\mathcal{X} \times \mathcal{Y}$. Our goal is to learn a prediction function $f : \mathcal{X} \rightarrow \mathcal{Y}$ which minimizes the classification risk with respect to the underlying data distribution $p$ and an appropriately chosen loss function $\ell : \mathcal{Y} \times \mathcal{Y} \rightarrow \mathbb{R}$.

§.§ Selective Classification

Setup Selective classification alters the standard setup for classification by introducing a rejection class $\bot$ through a gating mechanism <cit.>. In particular, a selective prediction model introduces a selection function $g:\mathcal{X} \rightarrow \mathbb{R}$ which determines whether a model should predict on a data point $\bm{x}$. Given an acceptance threshold $\tau$, the resulting predictive model can be summarized as follows:

\begin{equation}
(f,g)(\bm{x}) = \begin{cases}
f(\bm{x})  & g(\bm{x}) \geq \tau \\
\bot & \text{otherwise.}
\end{cases}
\end{equation}

Evaluation Metrics Prior work evaluates the performance of a selective classifier $(f,g)$ based on two metrics: the coverage of $(f,g)$ (what fraction of points we predict on) and the selective accuracy of $(f,g)$ on the points it accepts. Successful SC methods aim at obtaining both high accuracy and high coverage. Note that these two metrics are at odds with each other: naïvely increasing accuracy leads to lower coverage and vice-versa. The complete performance profile of a model can be specified using the risk–coverage curve, which defines the risk as a function of coverage <cit.>. These metrics can be formally defined as follows:

\begin{equation}
\text{cov}(f,g) = \frac{|\{\bm{x} : g(\bm{x}) \geq \tau \}|}{|D|} \qquad \qquad
\text{acc}(f,g) = \frac{|\{\bm{x} : f(\bm{x}) = y, g(\bm{x}) \geq \tau \}|}{|\{\bm{x} : g(\bm{x}) \geq \tau \}|}
\end{equation}

§.§ Related Work

Selective Classification
The first work on SC in the context of deep neural networks (DNNs) is the softmax response () mechanism <cit.>. A threshold $\tau$ is applied to the maximum response of the softmax layer $\max_{y \in \mathcal{Y}}f(y|\bm{x})$. Given a confidence parameter $\delta$ and desired risk $r^{*}$, this creates a selective classifier $(f, g)$ whose test error is no larger than $r^{*}$ with probability of at least $1-\delta$. It was further shown by <cit.> that model ensembles can improve upon the basic method. In our work, we leverage intermediate models yielded during training to avoid the costly training of separate models to obtain an ensemble. Whereas is a test-time method, SelectiveNet () <cit.> trains a model to jointly optimize for classification and rejection. A penalty is added to the loss to enforce a coverage constraint using a variant of the interior point method. In a different approach, Deep Gamblers () <cit.> transform the original $C$-class classification problem into a $(C + 1)$-class problem where the $(C + 1)$ -th class represents model abstention. This introduces less architectural changes than SelectiveNet but still involves a different training procedure. Recently, One-Sided Prediction () <cit.> has proposed to learn disjoint classification regions corresponding to each class; rejection is then implicitly defined as the gap between these regions via a constrained learning problem. Other similar approaches include: abstention based on the variance statistics from several dropout-enabled forward passes <cit.>, performing statistical inference for the marginal prediction-set coverage rates using model ensembles <cit.>, confidence prediction using an earlier snapshot of the model <cit.>, and complete precision for SC via classifying only when models consistent with the training data predict the same output <cit.>.

Example Difficulty
A related line of work is identifying difficult examples, or how well a model can generalize to a given unseen example. <cit.> introduce a per-instance empirical consistency score which estimates the probability of predicting the ground truth label with models trained on data subsamples of different sizes. Unlike our approach, however, this requires training a large number of models. <cit.> quantifies example difficulty through the lens of a forgetting event, in which the example is misclassified after being correctly classified. However, the metrics that we introduce in  <ref>, are based on the disagreement of the label at each checkpoint with the final predicted label. Other approaches estimate the example difficulty by: prediction depth of the first layer at which a $k$-NN classifier correctly classifies an example <cit.>, the impact of pruning on model predictions of the example <cit.>, and estimating the leave-one-out influence of each training example on the accuracy of an algorithm by using influence functions <cit.>. Closest to our method, the work of <cit.> utilizes gradients of intermediate models during training to rank examples by  difficulty. In particular, they average pixel-wise variance of gradients for each given input image. Notably, this approach is more costly and less practical than our approach and also does not study the accuracy/coverage trade-off which is of paramount importance to SC.

§ OUR METHOD

We now introduce our selective classification algorithms based on neural network training dynamics. We refer to this class of methods as . These methods are motivated by a theoretical framework that links parameters characterizing training dynamics (e.g., how predictions evolve during training) to the goal of predicting the likelihood of a point being correctly classified.

§.§ A Framework for Rejection

We first build some intuition as to how having access to intermediate states of the model during training can inform our gating mechanism. That is, how training dynamics can help us decide when to abstain from predicting on an input. To do so, we introduce a
bound on the probability of a datapoint being correctly classified based on the predictions intermediate models gave on that datapoint.

Let us denote the test dataset as $\mathcal{D}$. Given a model $f$, let $D' \subset \mathcal{D}$ be the subset which $f$ correctly classifies. We are interested in the setting where $f$ is trained sequentially (e.g., with stochastic gradient descent). We thus also have access to a sequence of $T$ intermediate states for $f$, which we denote $\{f_1,\cdots, f_T\}$. In this sequence, note that $f_T$ is exactly the final model $f$.

Now, let $p_{t}$ represent the random variable for outputs on $D'$ given by an intermediate model $f_t$ where the outputs have been binarized: we have $0$ if the output agrees with the final (correct) prediction and $1$ if not.[We note that $f_t$ does not necessarily output the correct prediction on $D'$, only the final model $f$  does.] In other words, $p_{t}$ is the distribution of labels given by first drawing $\bm{x} \sim D'$ and then outputting $1- \delta_{f_t(\bm{x}),f_T(\bm{x})}$ where $\delta$ denotes the Dirac delta function. Here, $\sim$ corresponds to sampling w.r.t some distribution on $D'$. For example, if all test points we evaluate our model on are taken from a larger finite set $\mathcal{D}$, then $D'$ is also finite and so $\sim$ can correspond to the uniform distribution on $D'$ (as it is well-defined).

Note that as $p_{t}$ is bounded, we always have a well-defined mean and a well-defined variance for $p_{t}$. Furthermore, we always have the variances and expectations of $\{p_{t}\}$ converge to $0$ with increasing $t$ (trivially as $p_{T} = 0$ always and the sequence is finite). To state this formally, let $v_t = \mathbb{V}[p_{t}]$ and $e_t = \mathbb{E}[p_{t}]$ denote the variances and expectations over correctly classified points. Then we have:

For all $ \epsilon > 0$ there exists an $N \in \{1,\ldots, T\}$ such that $v_t < \epsilon$ for all $t > N$. Similarly, for all $\epsilon > 0$ there exists a (possibly different) $N \in \{1,\ldots, T\}$ such that $e_t < \epsilon$ for all $t > N$.

However, the core problem is that we do not know how this convergence in the variance and expectation occurs. More specifically, if we knew the exact values of $e_t$ and $v_t$, we could use the following bound on the probability of a particular datapoint being correctly classified.  Note, we introduce the notation $\{a_1,\ldots,a_T\}$ where $a_t = 1- \delta_{f_t(\bm{x}),f_T(\bm{x})}$ which we call the "label disagreement (at $t$)", which will also be used later in the paper.

Given a datapoint $\bm{x}$,  let $\{a_1,\ldots,a_T\}$ where $a_t = 1- \delta_{f_t(\bm{x}),f_T(\bm{x})}$. Assuming not all $a_t = e_t$ then the probability $\bm{x} \in D'$ is  $\leq \min_{v_t~s.t~a_t \neq e_t}  \frac{v_t}{|a_t - e_t|^2}$.

By Markov's inequality we have the probability of the sequence occuring is trivially $\leq \frac{v_t}{|a_t - \mathbb{E}(p_{t})|^2}$ for any $t$ (a bound on any of the individual events occuring), and so taking the minimum over all these upper-bounds we obtain our upper-bound.

We do not guarantee Lemma <ref> is tight. Though we do take a minimum to make it tighter, this is a minimum over inequalities all derived from Markov's inequality.[One could potentially use information about the distribution of incorrectly classified datapoints to further refine this bound. Also recall that $p_{t}$ is only defined over the correctly classified points.] Despite this potential looseness, using the bound from Lemma <ref>, we can design a naïve selective classification protocol; for a given input $
\bm{x}$, if the upper-bound on the probability of being correct is lower than some threshold $\tau$ reject, else accept. However, as mentioned earlier, we have the following two problems preventing us from readily using this method:

* What is $\mathbb{E}[p_{t}]$ during training on correctly classified points?

* How does $\mathbb{V}[p_{t}]$ change during training?

These problems represent questions about how the predictions on correctly classified points evolve during training, the training dynamics of correctly classified points. More generally, they exemplify how a characterization of the training dynamics of correctly classified datapoints could be useful for selective classification. Informally, the first problem represents knowing how often we predict the true label at step $t$, and the second problem represents knowing how we become more consistent as we continue training.

The rest of this section will focus on the two methods we propose for selective classification. Although we do not give any formal guarantees of correctness for our methods, their design and assumptions are informed by Lemma <ref> and by empirical estimates of $e_t$ and $v_t$ (later described in  <ref>).

§.§ The Minimum Score

In what follows we will, after two simplifying assumptions, derive a candidate selective classification algorithm based off of Lemma <ref>. It is worth remarking that, despite our method being given by a simplification of Lemma <ref>, the performance of our approach is not upper bounded by Lemma <ref>.

In  <ref> we observe that the expected values for the $p_t$ distribution from  <ref>, which we denote $e_t$,  are near $0$ for most $t$;  we will therefore further simplify our setup by assuming $e_t = 0$ for all $t$. Note that if $e_t=0$ and the label disagreement $a_t=1$, then $\frac{v_t}{|a_t - \mathbb{E}[p_t]|^2}=v_t$. Assuming $e_t = 0$ for all $t$, based on Lemma <ref> the measure for acceptance we are interested in is $s_\text{min} = \min_{t~s.t~a_t = 1} v_t$. Note now $s_\text{min}$ depends on what $v_t$ are. From  <ref> we observe $v_t$ monotonically decrease in a convex manner. Thus, we choose to approximate the empirical trend observed in  <ref> with $v_t = 1 - t^k$ (for $k \in (0,1]$), which also decrease in a convex manner. Later, we empirically determine which $k$ is best for our methods.

Our first algorithm for selective classification is hence given by:

* Denote $L = f_T(\bm{x})$, i.e. the label our final model predicts.

* If $\exists t~s.t~a_t =1$ then compute $s_\text{min} = \min_{t~s.t~a_t = 1} v_t$ as per the notation in  <ref> (i.e $a_t = 1$ iff $f_t(x) \neq L$), else accept $\bm{x}$ with prediction $L$.

* If $s > \tau$ accept $\bm{x}$ with prediction $L$, else reject ($\bot$).

Intuitively, as all our candidate $v_t$ decrease, the algorithm imposes a last intermediate model which can output a prediction that disagrees with the final prediction: hereafter, the algorithm must output models that consistently agree with the final prediction. The effectiveness of this method depends on how quickly $v_t$ decrease (assuming $e_t = 0$ always) as opposed to the variance for incorrectly predicted data points. In particular, if there was an index $N$ such that for $t > N$, we have that $f_t(x) = f_T(x)$ for correctly classified points, but never (or rarely) for incorrectly classified points, then this method would successfully separate the correct from the incorrect.

§.§ The Average Score

Note that the previous method, defined by the score could be sensitive to stochasticity in training and hence perform sub-optimally. In particular, this algorithm could be noisy for different training runs (of otherwise identical models) as $\min_{t~s.t~a_t= 1} v_t$ could vary drastically based on the exact sequence of mini-batches used during training.

In light of this potential limitation we propose the following "averaging" algorithm:

* Denote $L = f_T(\bm{x})$, i.e. the label our final model predicts.

* If $\exists t~s.t~a_t =1$, compute $s_\text{avg} = \frac{\sum v_t a_t}{\sum a_t}$, else accept $\bm{x}$ with prediction $L$.

* If $s > \tau$ accept $\bm{x}$ with prediction $L$, else reject ($\bot$).

Intuitively, recalling our previous candidates for $v_t$, computes a weighted average, adding little weight to later disagreements with the final prediction (but still increasing the denominator by the same amount). As we demonstrate in  <ref>, this averaging procedure leads to considerably more robust results.

§ EMPIRICAL EVALUATION

2*Dataset     Target     2c     2c       2c     2c      2c

(lr)3-4 (lr)5-6 (lr)7-8 (lr)9-10 (lr)11-12
Error     Cov.     Error     Cov.     Error      Cov.     Error     Cov.     Error     Cov.     Error

3*CIFAR-10           2%      83.3     1.96     80.6     1.91      75.1     2.09     73.0     2.31     74.2     1.98

1%      79.7     1.05     74.0     1.02      67.2     1.09     64.5     1.02      66.4     1.01

0.5%      74.2     0.49     64.1     0.51       59.3     0.53     57.6     0.48     57.8     0.51

3*SVHN          2%      95.7     1.96     95.8     1.99     95.7     2.06     93.5      2.03     94.8     1.99

1%      91.2     0.99     90.1     1.03      88.4     0.99     86.5     1.04     89.5     1.01

0.5%      83.9     0.50     82.4     0.51      77.3     0.51     79.2     0.51     81.6     0.49

3*Cats & Dogs      2%      90.5     2.03     90.5     1.98      88.2     2.03     84.3     1.94     87.4     1.94

1%      85.6     1.01     85.4     0.98      80.2     0.97     78.0     0.98      81.7     0.98

0.5%     77.5     0.50     78.7     0.49      73.2     0.49     70.5     0.46     74.5      0.48

Performance at low target errors. Our $(s_\text{avg}, 0.05)$ method either matches or provides higher coverage levels of the test set at a fixed target error than other competing methods. Bold entries are within one standard deviation of each other over 5 random runs.

2*Dataset     Target     2c     2c       2c     2c      2c

(lr)3-4 (lr)5-6 (lr)7-8 (lr)9-10 (lr)11-12
Coverage     Cov.     Error     Cov.     Error     Cov.     Error      Cov.     Error     Cov.     Error

100%      100     9.57     100      9.74      99.99     9.58     100     11.07     100     10.81

95%       95.1     6.13     95.1     6.98     95.2     8.74     94.7     8.34     95.1     8.21

90%       90.1     4.16     90.0     4.67     90.5     6.52     89.6     6.45     90.1     6.14

100%     100     4.11     100     4.27     99.97     3.86     100     4.27     100     4.03

95%     94.8     1.80     95.1     1.83     95.1     1.86     95.1     2.53     95.0     2.05

90%     90.1     0.77     90.1     1.01     90.0     1.04     90.1     1.31     90.0     1.06

3*Cats & Dogs
100%     100     5.18     100     5.93     100     5.72     100     7.36     100     6.16

95%     94.9     2.99     95.1      2.97      95.0     3.46     95.2      5.1     95.1     4.28

90%     90.0     1.87     90.0      1.74      90.0     2.28     90.2     3.3     90.0     2.50

Performance at high target coverage. Similar as demonstrated in Table <ref>, $(s_\text{avg}, 0.05)$ matches or outperforms state-of-the-art error rates at fixed coverage levels.

Individual SVHN examples. In the first row, we show four examples from the SVHN test set. For each example we plot both the prediction evolution and the final label agreement metric $a_t$ over all checkpoints in the second row. We note that the first two examples are easy to classify as indicated by a mostly stationary prediction curve, while the last two examples are harder to classify since the prediction is highly unstable across intermediate models.

Performance and distribution of our metrics for different datasets. The first row shows the coverage/error curves for and . We show the distribution of the average score metric for both correctly and incorrectly classified samples in the second row.

Convex weighting

Convex weighting used in $v_t = 1 - t^k$. Overall, we find that (,$k$) is generally robust to the exact choice of $k$ and $k=0.05$ performs best across our experiments.

Coverage/error trade-off of $(s_\text{avg}, 0.05)$ for various checkpointing resolutions. Overall, we find that using 25-50 checkpoints already offers very competitive results. Using multiple hundreds of checkpoints offers diminishing returns but  improves coverage at low error rates.

Monitoring $\mathbb{E}[\cdot]$ and $\mathbb{V}[\cdot]$ for both correctly ($p_t$) and incorrectly ($q_t$) classified points. Correctly classified points have both their expectations and variances quickly decreasing to 0 as training progresses. Incorrectly classified points both exhibit large expectations and variances and stay constant over large periods while training.

§.§ Setup

Datasets & Training We evaluate our proposed approach on image dataset benchmarks that are common in the selective classification literature: CIFAR-10/CIFAR-100 <cit.>, SVHN <cit.>, Cats & Dogs[<https://www.microsoft.com/en-us/download/details.aspx?id=54765>], and GTSRB <cit.>. For each dataset, we train a deep neural network following the ResNet-32 architecture <cit.> and checkpoint each model after processing 50 mini-batches of size 128. All models are trained over 200 epochs. SVHN, Cats & Dogs, and GTSRB are trained using the Adam optimizer with learning rate $10^{-3}$. The CIFAR-10 and CIFAR-100 models are trained using momentum-based stochastic gradient descent with an initial learning rate of $10^{-1}$ on a multi-step learning rate reduction schedule (reduction by $10^{-1}$ after epochs 100 and 150), momentum $0.9$, and weight decay $10^{-4}$. To facilitate comparison with prior work, we match the accuracy of our models at the levels provided in <cit.> (see Table <ref> at 100% coverage  for exact levels).

Baselines We compare our method based on neural network training dynamics () to four common selective classification techniques previously introduced in  <ref>: Softmax Response () <cit.>, SelectiveNet () <cit.>, Deep Gamblers () <cit.>, and One-Sided Prediction () <cit.>. We describe hyper-parameter tuning for these baselines in Appendix <ref>.

§.§ Results

Our empirical results show that computing and thresholding the average score using weighting parameter $k=0.05$ provides a strong score for selective classification. We denote this method $(s_\text{avg}, 0.05)$. In the following, we first study the accuracy/coverage trade-off with comparison to past work. Then, we present exemplary label evolution curves for individual SVHN examples. Finally, we analyze the performance of the metrics proposed in  <ref> in depth and examine our method's sensitivity to the checkpointing resolution and the weighting parameter $k$.

Accuracy/Coverage Trade-off Consistent with standard evaluation schemes for selective classification, our main experimental results examine the accuracy/coverage trade-off of $(s_\text{avg}, 0.05)$. We present our main performance results with comparison to past work in Tables <ref> and <ref> where we demonstrate 's effectiveness on CIFAR-10, SVHN, and Cats & Dogs. In Table <ref>, we document the results obtained by , , , , and for a fixed low target error $e \in \{2\%,1\%,0.5\%\}$. In Table <ref>, we follow a similar setup but instead of fixing the target error we evaluate all methods for a fixed high target coverage $c \in \{100\%,95\%,90\%\}$. Across both experiments, we see that either matches or significantly outperforms existing state-of-the-art selective classification methods. In particular, we improve coverage of CIFAR-10 by 10.1% at a fixed target error of 0.5%.

Individual Evolution Plots To analyze the behavior of the metrics proposed in Section <ref>, we examine label prediction evolution curves for individual SVHN examples in Figure <ref> (more examples provided in Appendix <ref>). In particular, we select two easy-to-classify examples and two ambiguous examples from the test set and plot both the prediction evolution and the final label agreement metric $a_t$ over all checkpoints. We observe that easy examples only show a small degree of oscillation while hard examples show more erratic patterns. This result matches our intuition: our model should produce correct decisions on data points whose prediction is mostly constant throughout training and should reject data points for which intermediate models predict inconsistently.

Choice of Metric As per our theoretical framework and intuition provided in  <ref>, the average score should offer the most competitive selective classification performance. We confirm this finding in Figure <ref> where we plot the coverage/error curves for CIFAR-10, SVHN, GTSRB, and CIFAR-100 for both and . Overall, we find that the average score consistently outperforms the more noisy minimum score . We also show that yields distinct distributional patterns for both correctly and incorrectly classified points. This separation enables strong coverage/accuracy trade-offs via our thresholding procedure.

Checkpoint weighting sensitivity One important hyper-parameter of our method is the weighting of intermediate predictions. Recall from  <ref> that makes use of a convex weighting function $v_t = 1-t^k$. In Figure <ref>, we observe  that $(s_\text{avg}, k)$ is generally robust to the exact choice of $k$ and that $k=0.05$ performs best. At the same time, we find that decreasing $k$ too much leads to a decrease in coverage at fixed accuracy levels. This result confirms that (i) all stages of the training process contain valuable signals for selective classification; and that (ii) early label disagreements should be de-emphasized by our method.

Checkpoint resolution sensitivity The second  important hyper-parameter of our method is the granularity of available checkpoints. We study the sensitivity of $(s_\text{avg}, 0.05)$ with respect to the checkpointing resolution in Figure <ref>. Our experiments demonstrate favorable coverage/error trade-offs starting at 25-50 checkpoints. Further increasing the resolution to multiple hundreds of intermediate models does offer diminishing returns but also leads to better coverage especially at very low error rates.

§.§ Examining the Convergence Behavior of Correct & Incorrect  Points

As alluded to in the assumptions presented in  <ref>, we empirically study how the label disagreement distributions over correctly classified data points ($p_t$) and incorrectly classified data points ($q_t$) evolve during training. From Figure <ref> we observe that the expectations for correctly classified points are $0$ for large time periods, with a sudden initial drop from near $1$ to $0$; this drop becomes elongated for more challenging data sets such as CIFAR-10 and CIFAR-100 (for which the learning rate scheduling is visible in the plot). We also see that the variances follow an analogous decreasing trend. In general, both $e_t = \mathbb{E}[p_t]$ and $v_t = \mathbb{V}[p_t]$ converge monotonically (and convexly) to $0$. In contrast, incorrectly classified points show significantly larger mean and variance levels.

§ DISCUSSION

We briefly discuss additional potential metric choices that we investigated but which lead to SC performance worse than . Note that a more elaborate discussion  of these results is provided in Appendix <ref>.

Incorporating Estimates of $e_t$ into and  Since the results in   <ref> show that $e_t$ is only nearly $0$ almost everywhere, we investigate whether incorporating an estimate of $e_t$ into and leads to additional SC improvements. Our results demonstrate that neither an empirical estimate nor a smooth decay function similar to $v_t$ robustly improves over and (which assumed $e_t = 0$).

Jump Score
We also consider a score which captures the level of disagreement between the predicted label of two successive intermediate models (how much jumping occurred over the course of training). For $j_t = 0$ iff $f_{t}(\bm{x}) = f_{t-1}(\bm{x})$ and $j_t = 1$ otherwise we can compute the jump score as $s_\text{jmp} = 1 - \sum v_t j_t$ and threshold it as in  <ref> and  <ref>.

Variance Score for Continuous Metrics
Finally, we consider monitoring the evolution of continuous metrics that have been shown to be correlated with example difficulty. More specifically, <cit.> show that the metric they propose for example difficulty is correlated with confidence and negative entropy. Moreover, they find that difficult examples (or challenging and underrepresented features) are learned later in the training process. This observation motivates designing a score based on these continuous metrics that penalizes changes later in the training process more heavily.
We consider the maximum softmax class probability known as confidence, the negative entropy and the gap between the most confident classes for each example instead of the model predictions.
Assume that any of these metrics is given by a sequence $z = \{z_1,\ldots,z_T\}$ obtained from $T$ intermediate models. Then we can capture the uniformity of $z$ via a (weighted) variance score $s_\text{var} = \sum_{t} w_t (z_t - \mu)^2$ for mean $\mu = \frac{1}{T}\sum_{t} z_t$ and an increasing weighting sequence $w = \{w_1,\ldots,w_T\}$. We provide theoretical motivations for examining the variance of these continuous metrics in Appendix <ref> along with precise metric definitions.

§ CONCLUSION

We propose a new method to assess model outputs based on the training dynamics of the underlying ML model. Our approach is practical and can be directly implemented using existing checkpointing mechanisms.
The method is non-intrusive, requires no changes to the standard model training, and either matches or significantly outperforms results obtained using previous techniques.
Building on these foundations, future work could (i) investigate speeding up the inference procedure of intermediate models; (ii) study the existence of optimal weightings $v_t$; or (iii) explore whether the framework can be applied to the related but distinct problem of out-of-distribution sample detection.

§ ACKNOWLEDGEMENTS

We thank Roger Grosse, Chris Maddison, Franziska Boenisch, Natalie Dullerud, Jonas Guan, Michael Zhang, and Tom Ginsberg for fruitful discussions.

This work was supported by CIFAR (through a Canada CIFAR AI Chair), by NSERC (under the Discovery Program), and by a gift from Ericsson. We are also grateful to the Vector Institute’s sponsors for their financial support.