# The Conditional Prediction Function: A Novel Technique to Control False
Discovery Rate for Complex Models

Yushu Shi, Michael Martens

###### Abstract

In modern scientific research, the objective is often to identify which
variables are associated with an outcome among a large class of potential
predictors. This goal can be achieved by selecting variables in a manner that
controls the the false discovery rate (FDR), the proportion of irrelevant
predictors among the selections. Knockoff filtering is a cutting-edge approach
to variable selection that provides FDR control. Existing knockoff statistics
frequently employ linear models to assess relationships between features and
the response, but the linearity assumption is often violated in real world
applications. This may result in poor power to detect truly prognostic
variables. We introduce a knockoff statistic based on the conditional
prediction function (CPF), which can pair with state-of-art machine learning
predictive models, such as deep neural networks. The CPF statistics can
capture the nonlinear relationships between predictors and outcomes while also
accounting for correlation between features. We illustrate the capability of
the CPF statistics to provide superior power over common knockoff statistics
with continuous, categorical, and survival outcomes using repeated
simulations. Knockoff filtering with the CPF statistics is demonstrated using
(1) a residential building dataset to select predictors for the actual sales
prices and (2) the TCGA dataset to select genes that are correlated with
disease staging in lung cancer patients.

##  BACKGROUND

The frequency and sizes of high dimensional datasets have increased greatly in
recent years, accompanied both by new opportunities to find determinants of
outcomes of interest and by formidable challenges to identify truly prognostic
variables among this myriad of candidates. Usually, many of them are not
expected to affect the outcome and so interest lies in choosing only the
variables that are most likely to be associated with the response. An
appropriate way to formalize this objective is to focus on controlling the FDR
(Benjamini and Hochberg, 1995). Knockoff filtering (Barber and Candès, 2015)
is a popular device for FDR controlled variable selection, particularly since
this control is guaranteed for finite sample sizes rather than appealing to
asymptotic properties for its validity, as other frequently used approaches
do. The general problem involved is to determine, within a vector $\mathbf{X}$
of $p$ features, which ones are associated with the outcome $Y$ of interest.
The commonly used Model-X knockoff approach (Candès et al., 2018) conducts
this evaluation as follows:

  * •

First, construct knockoff copies $\widetilde{\mathbf{X}}$ satisfying the
following two conditions:

    1. 1.

Pairwise exchangeability: for any subset $S\subset\\{1,\dots,p\\},$
$(\mathbf{X},\widetilde{\mathbf{X}})_{swap(S)}\overset{d}{=}(\mathbf{X},\widetilde{\mathbf{X}})$,
i.e. swapping any subset of variables with their knockoffs leaves the joint
distribution invariant;

    2. 2.

Conditional independence: $\widetilde{\mathbf{X}}\perp\\!\\!\\!\perp
Y|\mathbf{X}.$

  * •

Second, a model is fitted using all original and knockoff features, and the
influences of the original and knockoff features are evaluated using
importance statistics, such as the lasso (Tibshirani, 1996) penalty parameter
at which a feature enters the model. The knockoff statistic built on it is the
lasso signed max (LSM) statistic. Because the knockoff features are
independent of the outcome, their importance statistics represent noise
variables that can be contrasted against the original features’ to assess
which original features correlate with the response. From the importance
statistics of the original feature and its knockoff, $Z_{j}$ and
$\widetilde{Z}_{j}$, $j=1,\dots,p$, we can compute knockoff statistics
$W_{j}=f_{j}(Z_{j},\widetilde{Z}_{j}),$ where $f_{j}$ is an anti-symmetric
function. A large positive $W_{j}$ provides evidence that the distribution of
$Y$ depends on $X_{j},$ whereas under the null $W_{j}$ is equally likely to
take on positive and negative values (the ”coin-flip” property).

  * •

Third, features are selected based on the knockoff statistics in a manner that
controls the FDR at the specified level, $q$. Specifically, if we find a
threshold $\tau>0$ by setting

$\tau=\min\left\\{t>0:\frac{\\#\\{j:W_{j}\leq-t\\}}{\\#\\{j:W_{j}\geq
t\\}}\leq q\right\\}.$

Then the procedure selecting the set

$\hat{S}=\\{j:W_{j}\geq\tau\\}$

controls the modified FDR as

$mFDR=\mathbb{E}\left[\frac{|\\{j\in\hat{S}\cap\mathcal{H}_{0}\\}|}{(|\hat{S}|+1/q)}\right]\leq
q.$

A more conservative procedure, knockoffs$+,$ with threshold

$\tau^{+}=\min\left\\{t>0:\frac{1+\\#\\{j:W_{j}\leq-t\\}}{\\{\\#j:W_{j}\geq
t\\}}\leq q\right\\}$

controls the ordinary FDR at level $q$ by selecting

$\hat{S}_{+}=\\{j:W_{j}\geq\tau^{+}\\}.$

Previous work has explored various ways to improve knockoff feature generation
with a guarantee of pairwise exchangeability. Notably, Bates et al. (2021)
introduced a Metropolized knockoff sampler, which handles situations where the
covariates are continuous, heavy-tailed, and follow an Ising model.
Furthermore, other methods took advantage of generative machine learning
models, including Deep Knockoffs (Romano et al., 2019), KnockoffGAN (Jordon et
al., 2019), and variational autoencoder related methods (Liu and Zheng, 2019;
Duarte and Feng, 2020). Compared with the abundance of choices for generating
knockoff counterparts, few studies have attempted to improve knockoff
statistics. Common choices of knockoff statistics are based on a linear
regression based model, such as lasso (Barber and Candès, 2015; Candès et al.,
2018); a common choice is the lasso coefficient difference (LCD) statistic,
which is defined as

$W_{j}^{LCD}=|\beta_{j}|-|\tilde{\beta}_{j}|,$

where $\beta_{j}$ and $\tilde{\beta}_{j}$ are the regression coefficients for
$X_{j}$ and its knockoff $\tilde{X}_{j}$. The lasso-based statistics can yield
good power if the true relationship is linear (Weinstein et al., 2017). But,
if the linearity assumption is violated, these approaches may have poor power
to detect prognostic variables. This paper investigates methods for
constructing knockoff statistics that can yield higher power to resolve
features with linear or nonlinear relationships with the outcomes while
maintaining the FDR control. We propose the conditional prediction function
(CPF) as a basis for building the knockoff importance statistic. The proposed
CPF knockoff statistic is able to capture nonlinear relationships, and will
empower researchers to explore more flexible prediction models for knockoff
filtering, such as deep neural networks (DNNs).

In Section 2, we explain the rationale and summarize the algorithm for
calculating CPF knockoff statistics. The CPF statistic’s ability to capture
non-linear relationship for different types of outcomes is demonstrated
through repeated simulations in Section 3. Section 4 applies the proposed
method to a residential building dataset to select predictors for the actual
sales prices and to the TCGA LUAD cancer dataset to identify genes that are
associated with or more advanced stage of lung cancer. Section 5 concludes the
paper with a short discussion.

##  THE CPF STATISTIC

In this section, we introduce the conditional prediction function (CPF), which
can measure a feature’s importance by changing its value by a small magnitude
while conditioning on all other covariates and computing the prediction
difference before and after the change. The concept of the CPF knockoff
statistic originates from multivariable linear models, where the regression
coefficient signifies how much the mean of the dependent variable changes
given a one-unit shift in the independent variable while holding other
variables in the model constant. The CPF can be paired with many complicated
predictive models, however. In the following text, we assume that we have
already generated high-quality knockoff features $\tilde{\mathbf{X}}$ from $p$
original predictors $\mathbf{X}.$ Calculation of the CPF knockoff statistics
proceeds per the following:

  1. 1.

Combine both the original and the knockoff predictors
$\mathbf{X}^{*}_{[N\times 2p]}=[\mathbf{X},\tilde{\mathbf{X}}]$ in one matrix.
Center and scale continuous predictors.

  2. 2.

Train the prediction model $\mathcal{M}$ for the outcome $Y$ with both the
original and knockoff features.

  3. 3.

Set $U_{i}=0,\quad i=1,\dots,2p.$ Choose the number of percentiles at which we
want to evaluate the CPF, $J$ , and the $\delta$ value for numerical
calculation. Then,

for  _each predictor $X^{*}_{i}$, $i\leftarrow 1$ to $2p$_ do

Sample a subset with $n$ observations.

for  _$k\leftarrow 1$ to $n$_ do

if  _$X_{i}^{*}$ is a continuous covariate_ then

for  _$j\leftarrow 1$ to $J$_ do

Denote the value of $X_{i}^{*}$ at the $100j/(J+1)$ percentile as
$X_{ij}^{*}$. Conditioning on the values of all other predictors
$\mathbf{X}_{k,-i}^{*}$, evaluate the predictive difference for
$X_{ij}^{*}-\delta/2$ and $X_{ij}^{*}+\delta/2$.

$(\star)\quad
U_{i}=U_{i}+\lVert\mathcal{M}(X_{ij}^{*}+\delta/2|\mathbf{X}_{k,-i}^{*})-\mathcal{M}(X_{ij}^{*}-\delta/2|\mathbf{X}_{k,-i}^{*}))\rVert_{2}^{2}/(\delta^{2}J)$

end for

end if

else

$(\star)\quad
U_{i}=U_{i}+||\mathcal{M}(X_{ij}^{*}=1|\mathbf{X}_{k,-i}^{*})-\mathcal{M}(X_{ij}^{*}=0|\mathbf{X}_{k,-i}^{*})||_{2}^{2}$

end if

end for

end for

for  _each pair of the original predictor $X_{m}$ and the corresponding
knockoff $\tilde{X}_{m}$, $m\leftarrow 1$ to $p$_ do

$CPF_{m}=U_{m}-U_{m+p}$

end for

After computing knockoff statistics for each pair of the original and knockoff
covariates, we implement the last step of knockoff filtering. Given the pre-
specified FDR level $q$, we find the threshold value $T$ such that

$T=\min\left\\{t>0:\frac{\\#\\{r:CPF_{r}\leq-t\\}}{\\#\\{r:CPF_{r}\geq
t\\}}\leq q\right\\},$

and return the selected feature set $\hat{S}=\\{r:CPF_{r}\geq T\\}.$

Figure 1: CPF can quantify nonlinear covariate effect. Here the red curve is
the true effect of the covariate. If we fit a linear model, the slope of the
line will be 0, as indicated by the blue line. The black lines indicate the
prediction difference that can be captured by CPF statistics.

To illustrate the benefit of CPF, consider the univariate case where the
predictor and the outcome have a relationship $y=x^{2}$. As shown in Figure 1,
the predictor has a negative slope when less than $0,$ and a positive slope
when greater than $0.$ If one uses a linear model or a linear model based
method (such as lasso, elastic net), this variable will likely be declared
null, as the positive effect when $x>0$ and negative effect when $x<0$ cancel
out. But if one employs CPF as knockoff importance statistics, this nonlinear
relationship will be well captured, since CPF measures the cumulative squared
effect over the space of $x$.

When features are independent and the relationship to the outcome is genuinely
linear, the CPF reduces to $n$ times the squared regression coefficient of a
linear model. However, when the relationship is more complicated, by
evaluating the conditional prediction differences at $J$ percentiles and
averaging over $n$ possible combinations of other predictors, the CPF can
overcome the rigidity of linear models.

When the outcome of interest is continuous, the predicted outcome
$\mathcal{M}(\cdot)$ is a real number. When the outcome of interest is
categorical, $\mathcal{M}(\cdot)$ is a vector of probabilities belonging to
each category, and the increment in each iteration is the squared Euclidean
distance of two probability vectors. When the outcome of interest is a
survival time in the single failure cause setting, the CPF can be calculated
over all the distinct time points in the dataset and $\mathcal{M}(\cdot|t)$ is
the predicted survival function indexed by time $t$. When multiple failure
types, i.e. competing risks, are present, $\mathcal{M}(\cdot|t)$ is the
predicted cumulative incidence function (CIF) for the event type of interest.
When the number of distinct event times $N_{t}$ is large, we can select a
subset of $n_{t}$ time points (say, quintiles of the observed survival times),
$\\{t_{l};l=1,\dots,n_{t}\\}$. For survival outcomes, the $\star$ step in the
algorithm becomes

for  _$l\leftarrow 1$ to $n_{t}$_ do

$U_{i}=U_{i}+||\mathcal{M}(X_{ij}^{*}+\delta/2|t_{l},\mathbf{X}_{k,-i}^{*})-\mathcal{M}(X_{ij}^{*}-\delta/2|t_{l},\mathbf{X}_{k,-i}^{*})||_{2}^{2}/\delta^{2}J$

end for

if $X*_{i}$ is continuous, and

for  _$l\leftarrow 1$ to $n_{t}$_ do

$U_{i}=U_{i}+||\mathcal{M}(X_{ij}^{*}=1|t_{l},\mathbf{X}_{k,-i}^{*})-\mathcal{M}(X_{ij}^{*}=0|t_{l},\mathbf{X}_{k,-i}^{*})||_{2}^{2}$

end for

otherwise.

The CPF resembles the squared gradient of a multivariate function. To
approximate the gradient well, $\delta$ needs to be small relative to the
range of $\mathbf{X}$. Large $J$, $n$ and $n_{t}$ can yield a more accurate
estimation of the predictive ability of the features; yet, the choice of these
parameters may be constrained by the computing resources available.

Linear Relationship Nonlinear Relationship

((a))

Continuous

((b))

((c))

Binary

((d))

((e))

Log-linear Hazard Relationship Non Log-linear Hazard Relationship

((f))

Survival

((g))

((h))

Competing risks

((i))

((j))

Figure 2: Comparison of the FDR and power for CPF and LCD statistics. LCD and
CPF denote standard knockoff filtering using the LCD and CPF statistics; LCD+
and CPF+ denote knockoff+ filtering for these techniques. The left columns
show the results when the relationship between covariates and the outcome is
linear/log-linear, while the right two columns show the results when the
relationship is not.

##  SIMULATION STUDY

To illustrate the ability of the CPF statistics to detect both linear and
nonlinear effects of predictors, we evaluated their performance in conjunction
with DNNs for various outcome types, including continuous, categorical,
survival, and competing risks data. In all simulated datasets, the number of
original features was set at $250,$ while the sample size was $20000.$ Among
the $250$ predictors, ten were genuinely predictive. During the training of
DNNs, early stopping callback was implemented to mitigate overfitting, with
the patience parameter set at 50 for continuous, categorical, and survival
outcomes and at 200 for competing risks data. For continuous and categorical
outcomes, a batch size of 20 was employed, and the validation proportion in
each epoch was designated at 10%. For survival and competing risks outcomes, a
batch size of 50 was used, and the validation proportion in each epoch was set
at 20%. The parameters $J=5$, $n=100$, and $\delta=0.1$ were specified for the
calculation of the CPF statistics from fitted DNNs.

The LCD statistics were selected as the benchmark for comparison, as it has
demonstrated higher power than the LSM statistics in most instances (Candès et
al., 2018). For both continuous and categorical outcomes, the LCD statistics
were derived from lasso results; for survival outcomes, they were based on Cox
lasso results, and for competing risks outcomes, they were obtained from the
Fine-Gray regression model (Fine and Gray, 1999).

The target false discovery rate was set to be $0.2.$ Original features were
generated using either a standard normal distribution or an autoregressive
model with a mean of $\mathbf{0}$ and a covariance matrix with $0.2^{|i-j|}$
entries, where $i$ and $j$ denote column and row indicators, respectively. The
covariance matrix of the autoregressive model contains $1$ along the diagonal,
with entry values diminishing to zero as they move further from this diagonal.
The second-order approximate Gaussian knockoffs, provided by the R package
“knockoff” (Patterson and Sesia, 2020), were employed as the knockoff
generator. Each simulation configuration was repeated $200$ times, and the
results were summarized using box plots. Subsequent subsections will present
the simulation setup for each type of outcome.

###  Continuous Outcome

For continuous outcomes, the true distribution of the response variable was
assumed to follow a linear model with predictors that are either linear or
quadratic functions of X. In the linear setting, the predictors and the
outcome are related as follows:
$y_{i}=\boldsymbol{\beta}^{T}\mathbf{X}^{*}_{i}+100+\epsilon_{i},$
$i=1,2,\dots,n,$ where $\mathbf{X}^{*}_{i}$ represents the vector that
includes only true predictors, and $\boldsymbol{\beta}$ is a vector with all
entries equal to 3\. In the nonlinear setting, the relationship between each
predictor and the outcome is quadratic, given by:
$y_{i}=\sum_{j=1}^{10}\mathbf{X}^{*2}_{ij},$ $i=1,2,\dots,n.$

In the linear setting, a three-layer DNN was employed, where the internal
layer consisted of 8 nodes. For the nonlinear setting, a five-layer DNN was
used, with the numbers of nodes in internal layers specified as (64, 32, 16).
The activation function for both models was chosen to be the Rectified Linear
Unit (ReLU), and for each layer, an L1 regularization term was imposed with a
parameter value of 0.001. The optimizer was configured to be Adam, with a
learning rate set at 0.01, and the loss function was defined as the mean
squared error in the validation data.

###  Categorical Outcome

The data simulation for categorical outcome is similar to the simulation for
continuous outcomes. In the linear relationship setting, the predictors and
the expected value of the outcome follow a linear relationship on the logit-
transformed scale according to
$y_{i}\sim\operatorname{binom}(\operatorname{logistic}(\boldsymbol{\beta}^{T}\mathbf{X}^{*}_{i})),$
$i=1,2,\dots,n.$ $\mathbf{X}^{*}_{i}$ is the vector including only true
predictors and $\boldsymbol{\beta}$ is a vector of $1^{\prime}s.$ In the
nonlinear setting, we first computed the sum of squares of the true predictors
$Z_{i}=\sum_{j=1}^{10}\mathbf{X}^{*2}_{ij},$ then standardized $Z_{i}$ and
generated the outcome $y_{i}\sim\operatorname{binom}(Z_{i}).$

To decrease the training time, we started with the model trained for
continuous outcomes and simply changed the last layer to a sigmoid activation
function. The loss function was chosen to be binary-crossentropy in the
validation data. The learning optimizer was chosen to be Adam, and the
learning rate was 0.01.

###  Survival Outcome

We also applied our method on data with survival outcomes. We used DeepSurv
model (Katzman et al., 2018), which adopts the framework of the Cox
proportional hazards model, for computing the predicted survival based CPF. In
the DeepSurv model, the hazard ratio of an individual is parameterized by the
weights of the network, and the optimization goal is to maximize the log
partial likelihood with some regularization. Again, we use $X_{j}^{*},\quad
j=1,\dots,10,$ to denote the $j$th true predictor and specify that the hazard
for each individual follows
$h_{i}(t|\mathbf{X})=\exp(\sum_{j=1}^{10}0.5X_{ij}^{*}),\quad i=1,\dots,n.$
This satisfies the underlying assumption for the Cox model as the covariate
imposes a time-invariant multiplicative effect on the hazard. We refer to the
scenario in which the covariate affects the hazard through a log-linear
relationship as the log-linear setting. The censoring times were generated
independently from an $\exp(0.05)$ distribution and this resulted in a $5\%$
censoring rate. In simulation scenarios where the covariate effects are not
log-linear, the hazard for the $i$th individual was
$h_{i}(t|\mathbf{X})=\exp(\sum_{j=1}^{10}0.5X_{ij}^{*2}),$ and the censoring
times were generated independently from an $\exp(0.1)$ distribution.

For the scenarios under the log-linear setting, the numbers of nodes were set
to be $(64,32,16,8,1),$ for the non log-linear scenarios, the numbers of nodes
were $(128,64,32,16,1).$ Batch normalization and a dropout rate $0.5$ were
employed. We used Adam optimizer with learning rate 0.0001 and the weight
decay was set to be 0.01.

###  Competing Risks Outcome

To further assess the capabilities of the CPF statistics, we explored the
application with competing risks data. In competing risks data, there are
several types of potential events and the occurrence of one type of outcome
prevents us from observing other types of events. Suppose that type 1 events
are the main interest. Fine and Gray (1999) mimicked the Cox proportional
hazards model by transforming the CIF $F_{1}(t)$ for type 1 into a
subdistribution hazard $\eta_{1}(t),$

$\eta_{1}(t)=-\frac{\partial}{\partial t}\log(1-F_{1}(t)),$ (1)

and assumed the covariates have multiplicative, time-constant effects on the
subdistribution hazard through the form

$\eta_{1}(t|\mathbf{x})=\eta_{01}(t)\exp\left(f(\mathbf{X})\right).$ (2)

Recently, Lee et al. (2018) leveraged DNN for modeling competing risks data
and proposed the DeepHit model, which was employed in computing the predicted
CIF used in CPF statistics in our simulation.

Our simulations assumed that 2 event types exist and that type 1 is the
primary interest. In both log-linear and non log-linear hazards scenarios, we
assumed the baseline probability of experiencing event type 1 was $q=0.5.$ The
influence of covariates on outcomes was encapsulated by the variable
$Z=f(\mathbf{X})$. Specifically, in the log-linear scenario,
$Z_{i}=\sum_{j=1}^{10}X_{ij}^{*},\quad i=1,\dots,n.$ For the non log-linear
scenario, $Z_{i}=\sum_{j=1}^{10}X_{ij}^{*2}-10.$ For both scenarios, $Z_{i}$
is connected to the subdistribution hazard function through equation 2. By
virtue of (1)-(2), the relationship between the baseline and individual CIFs
is

$F_{1}(t|\mathbf{X}_{i})=1-(1-F_{01}(t))^{\exp(Z_{i})}.$ (3)

Generating event times involved a two-step process: first, we generated
probabilities of experiencing a type 1 event for each subject by randomly
sampling from a binomial distribution with a probability of
$F_{1}(\infty|\mathbf{X}_{i})=1-(1-q)^{\exp(Z_{i})};$ subsequently, for those
subjects experiencing type 1 events, we modeled the baseline CIF for event
type 1 as $F_{01}(t)=q(1-\exp(-t))$ and generated their event time from the
cumulative distribution function
$F_{1}(t|\mathbf{X}_{i})/F_{1}(\infty|\mathbf{X}_{i})$, where the numerator
follows (3). For type 2 events, we generated event times from an exponential
distribution with a rate parameter of 1. Censoring times were independently
generated from an $\exp(1)$ distribution.

Within the architecture of the DeepHit, for both the log-linear and non lig-
linear scenarios, the number of internal nodes was configured as
$(128,64,32,16).$ To enhance training performance, we applied batch
normalization and utilized a dropout rate of 0.5. In fitting the DeepHit
model, the number of cutpoints for evaluating the likelihood was set to be 20.
To mitigate overfitting, an L2 penalty with a parameter of 0.1 was
incorporated.

###  Summary of Simulation Results

As shown in Figure 2, the simulation results are consistent across different
outcome types. When the linear assumption or the log-linear assumption is
satisfied, the performance of the CPF is similar to the performance of the
linear model-based LCD. However, when this assumption does not hold, for
example, the relationship between the predictors and outcome is quadratic, the
LCD fails to capture the true relationship and exhibits significant power
loss, while the CPF controls the FDR while offering high power to detect
signals.

##  REAL DATA APPLICATION

Table 1: Genes identified by the CPF method. Genes that corroborate with the selection from LCD statistics results are shown in bold. GENE | RELEVANT LITERATURE
---|---
GAPDH | Puzone et al. (2013); Liu et al. (2017)
RPLP0P6 | Xie et al. (2023); Wang et al. (2018)
RNASE1 | Travaglini et al. (2020); Li et al. (2019b); Mironova and Vlassov (2019)
SFTPC | Li et al. (2007); Wang et al. (2009); Li et al. (2019a)
IGKV4-1 | Wang et al. (2020); Qin et al. (2020); Xia et al. (2022); Pedersen et al. (2022)
IGHJ3 | Isaeva et al. (2019); Feng et al. (2009); Jiang et al. (2022)

###  Identifying predictors of real estate sale price

We analyzed a dataset (Rafiei and Adeli, 2018) consisting of 108 covariates,
aiming to identify predictors of the sale prices of single-family residential
apartments in Tehran, Iran. The dataset encompasses various factors, including
temporal economic covariates as well as diverse project-related physical and
financial attributes. Furthermore, the dataset includes economic indicators
and indices across five periods prior to the start of construction, such as
stock market index and city population.

We divided the 372 observations into training ($70\%$) and testing ($30\%$)
sets. We first fit the lasso model with cross-validation, yielding a mean
squared error of 27767.39 on the testing dataset. Then, we employed the LCD
knockoff statistics with a target false discovery rate of 0.2. The knockoff+
method led to no feature selection, while the standard knockoff procedure
resulted in the selection of a single feature: the unit price at project
initiation.

To compute CPF knockoff statistics, we started by constructing a DNN with four
layers. The internal layers contained 256 and 128 nodes respectively,
utilizing the leakyReLU activation function with an alpha parameter of 0.3.
The final layer adopted the ReLU activation function. Across all layers, batch
normalization was implemented. We opted for the Adam optimizer with a learning
rate of 0.01 and chose batch size to be 50. To prevent overfitting, early
stopping with a patience parameter of 200 was employed. The DNN can achieve a
mean squared error of 20787.06 on the test dataset, which outperforms lasso in
model fitting. When adopting the more conservative knockoff+ approach, similar
to the result from LCD statistics, no features were selected. Meanwhile, with
the standard knockoff method, three features were chosen: equivalent
preliminary estimated construction cost based on prices at the project’s
outset, unit price at project initiation, and actual construction expenses.

###  Determinants of lung adenocarcinoma stage from TCGA dataset

We applied CPF statistics on lung adenocarcinoma gene expression data from The
Cancer Genome Atlas (TCGA). Our primary focus was on determining genes that
express differently for stage I and stages II/III/IV cancer. In this dataset,
there were 330 patients with stage I cancer and 271 patients with stages
II/III/IV cancer. From the original 31,294 features, considering the
constrained sample size, we selected 100 features with the highest variance.
The target FDR here is 0.2. Before fitting the model, the predictors were
standardized. We allocated $70\%$ of the data for training and utilized the
remainder for testing the model. Given the limited sample size, we employed a
three layer neural network, with an internal layer with 100 nodes. The
activation functions were chosen to be ReLU and sigmoid. Training DNN involved
using dropout (dropout rate = 0.2) and L1 regularization with penalty
parameter of 0.1 for kernel and 0.01 for bias in the last layer. The batch
size was chosen to be $100.$ The Adam optimizer with a learning rate of
$0.001$ was used, and early stopping call with patience parameter 200 was
employed to prevent overfitting.

Using cross-validation, the lasso model yielded a c-index of 0.60088 for the
test data, while the DNN achieved a c-index of 0.6774, demonstrating a
superior model fitting than the lasso model. The knockoff+ method did not
select any features for either LCD or CPF knockoff statistics. However, with
the regular knockoff method, 6 features were chosen using the CPF approach,
and LCD selection corroborated two of these (LCD with lasso chose 3 features
in total). Table 1 lists all the genes selected by CPF and recent biomedical
publications that have identified the correlations between the same gene and
lung cancer.

In both examples, the CPF statistics selected more features than LCD
statistics. This behavior can be explained by its capacity in capturing
nonlinear relationship between covariates and outcomes. The real data
application is consistent with simulation study outcomes, demonstrating that
CPF statistics have better power than a widely used existing method.

##  DISCUSSION

Applying the CPF statistics successfully relies on pairing them with a well-
selected and trained machine learning model. This paper is not aimed at
providing a one-size-fits-all solution for variable selection and the
flexibility of the CPF does not eliminate the need for tuning complicated
machine learning models, in particular DNNs. We encourage users to first check
the model fitting metrics of the trained model, including mean squared error,
prediction accuracy, and concordance index and comparing them to the
performance of traditional methods, such as linear model, lasso, and the Cox
model. After verifying that the machine learning model attains better model
fitting, then CPF can be used with it to identify key features that are
predictive of the outcomes.

Complex machine learning models such as DNNs have long been perceived as
“black boxes,” such that the relationships between features and response
variables are hard to discern from the fitted model and which variables are
truly associated with the outcome is unclear. The CPF knockoff statistic
offers improved model interpretability by capturing the nonlinear
relationships between the covariates and outcomes and enabling these
correlations to be identified through knockoff filtering. When applied in
biomedical studies, CPF statistics offer an improved resolution of prognostic
features for health outcomes. This improvement in feature identification will
promote the understanding of the determinants of patient health and,
subsequently, can lead to advances in the prevention, treatment, and
management of disease, thereby improving public health. Also, The proposed
technique will be helpful more broadly in other scientific contexts such as
information technology (IT), engineering, sociology, economics, psychology,
and environmental science, where a large set of candidate variables are being
evaluated for association with an outcome of interest.

## ACKNOWLEDGMENT

This project is partly supported by NSF DMS 2310955. We would like to thank
Dr. Hao Feng of Case Western Reserve University for his help in TCGA data
application.

## References

  * Barber and Candès (2015) Barber, R. F. and Candès, E. J. (2015).  Controlling the false discovery rate via knockoffs.  The Annals of Statistics, 43(5):2055–2085.
  * Bates et al. (2021) Bates, S., Candès, E., Janson, L., and Wang, W. (2021).  Metropolized knockoff sampling.  Journal of the American Statistical Association, 116(535):1413–1427.
  * Benjamini and Hochberg (1995) Benjamini, Y. and Hochberg, Y. (1995).  Controlling the false discovery rate: A practical and powerful approach to multiple testing.  Journal of the Royal Statistical Society. Series B, 57(1):289–300.
  * Candès et al. (2018) Candès, E. J., Fan, Y., Janson, L., and Lv, J. (2018).  Panning for gold: ‘Model-X’ knockoffs for high dimensional controlled variable selection.  Journal of the Royal Statistical Society: Series B, 80(3):551–577.
  * Duarte and Feng (2020) Duarte, M. F. and Feng, S. (2020).  Knockoff-inspired feature selection via generative models.  preprint.
  * Feng et al. (2009) Feng, Y., Xiao, X., Zhu, Z., Streaker, E., Ho, M., Pastan, I., and Dimitrov, D. S. (2009).  A novel human monoclonal antibody that binds with high affinity to mesothelin-expressing cells and kills them by antibody-dependent cell-mediated cytotoxicity.  Molecular Cancer Therapeutics, 8(5):1113–1118.
  * Fine and Gray (1999) Fine, J. P. and Gray, R. J. (1999).  A proportional hazards model for the subdistribution of a competing risk.  Journal of the American Statistical Association, 94(446):496–509.
  * Isaeva et al. (2019) Isaeva, O. I., Sharonov, G. V., Serebrovskaya, E. O., Turchaninova, M. A., Zaretsky, A. R., Shugay, M., and Chudakov, D. M. (2019).  Intratumoral immunoglobulin isotypes predict survival in lung adenocarcinoma subtypes.  Journal for ImmunoTherapy of Cancer, 7(1).
  * Jiang et al. (2022) Jiang, W., Zhou, S., Li, J., Zhu, G., Gao, M., Zhao, K., Zhang, L., Xie, X., Zhao, N., Tian, C., Zhang, Z., Zhao, Y., Guo, Y., Cui, Y., and Liu, P. (2022).  B cell receptor (BCR) diversity and differential cdr3s usage as potential immune indictors of diffuse large B cell lymphoma (DLBCL).  preprint.
  * Jordon et al. (2019) Jordon, J., Yoon, J., and van der Schaar, M. (2019).  KnockoffGAN: Generating knockoffs for feature selection using generative adversarial networks.  In International Conference on Learning Representations.
  * Katzman et al. (2018) Katzman, J. L., Shaham, U., Cloninger, A., Bates, J., Jiang, T., and Kluger, Y. (2018).  DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network.  BMC Medical Research Methodology, 18(1):24.
  * Lee et al. (2018) Lee, C., Zame, W., Yoon, J., and van der Schaar, M. (2018).  DeepHit: A deep learning approach to survival analysis with competing risks.  Proceedings of the AAAI Conference on Artificial Intelligence, 32(1).
  * Li et al. (2019a) Li, B., Meng, Y.-Q., Li, Z., Yin, C., Lin, J.-P., Zhu, D.-J., and Zhang, S.-B. (2019a).  MiR-629-3p-induced downregulation of SFTPC promotes cell proliferation and predicts poor survival in lung adenocarcinoma.  Artificial Cells, Nanomedicine, and Biotechnology, 47(1):3286–3296.  PMID: 31379200.
  * Li et al. (2007) Li, R., Todd, N. W., Qiu, Q., Fan, T., Zhao, R. Y., Rodgers, W. H., Fang, H.-B., Katz, R. L., Stass, S. A., and Jiang, F. (2007).  Genetic deletions in sputum as diagnostic markers for early detection of stage I non–small cell lung cancer.  Clinical Cancer Research, 13(2):482–487.
  * Li et al. (2019b) Li, Y., Yang, Y., Gan, T., Zhou, J., Hu, F., Hao, N., Yuan, B., Chen, Y., and Zhang, M. (2019b).  Extracellular rnas from lung cancer cells activate epithelial cells and induce neutrophil extracellular traps.  International Journal of Oncology, 55(1):69–80.
  * Liu et al. (2017) Liu, K., Tang, Z., Huang, A., Chen, P., Liu, P., Yang, J., Lu, W., Liao, J., Sun, Y., Wen, S., Hu, Y., and Huang, P. (2017).  Glyceraldehyde-3-phosphate dehydrogenase promotes cancer growth and metastasis through upregulation of snail expression.  International Journal of Oncology, 50(1):252–262.
  * Liu and Zheng (2019) Liu, Y. and Zheng, C. (2019).  Deep latent variable models for generating knockoffs.  Stat, 8(1):e260.  e260 sta4.260.
  * Mironova and Vlassov (2019) Mironova, N. and Vlassov, V. (2019).  Surveillance of tumour development: The relationship between tumour-associated RNAs and ribonucleases.  Frontiers in Pharmacology, 10.
  * Patterson and Sesia (2020) Patterson, E. and Sesia, M. (2020).  knockoff: The Knockoff Filter for Controlled Variable Selection.  R package version 0.3.3.
  * Pedersen et al. (2022) Pedersen, S., Jensen, K. P., Honoré, B., Kristensen, S. R., Pedersen, C. H., Szejniuk, W. M., Maltesen, R. G., and Falkmer, U. (2022).  Circulating microvesicles and exosomes in small cell lung cancer by quantitative proteomics.  Clinical Proteomics, 19(1):2.
  * Puzone et al. (2013) Puzone, R., Savarino, G., Salvi, S., Dal Bello, M. G., Barletta, G., Genova, C., Rijavec, E., Sini, C., Esposito, A. I., Ratto, G. B., Truini, M., Grossi, F., and Pfeffer, U. (2013).  Glyceraldehyde-3-phosphate dehydrogenase gene over expression correlates with poor prognosis in non small cell lung cancer patients.  Molecular Cancer, 12(1):97.
  * Qin et al. (2020) Qin, F.-l., Xu, Z.-y., Yuan, L.-q., Chen, W.-j., Wei, J.-b., Sun, Y., and Li, S.-k. (2020).  Novel immune subtypes of lung adenocarcinoma identified through bioinformatic analysis.  FEBS Open Bio, 10(9):1921–1933.
  * Rafiei and Adeli (2018) Rafiei, M. H. and Adeli, H. (2018).  Novel machine-learning model for estimating construction costs considering economic variables and indexes.  Journal of Construction Engineering and Management, 144(12):04018106.
  * Romano et al. (2019) Romano, Y., Sesia, M., and Candès, E. (2019).  Deep knockoffs.  Journal of the American Statistical Association, pages 1–12.
  * Tibshirani (1996) Tibshirani, R. (1996).  Regression shrinkage and selection via the lasso.  Journal of the Royal Statistical Society: Series B, 58(1):267–288.
  * Travaglini et al. (2020) Travaglini, K. J., Nabhan, A. N., Penland, L., Sinha, R., Gillich, A., Sit, R. V., Chang, S., Conley, S. D., Mori, Y., Seita, J., Berry, G. J., Shrager, J. B., Metzger, R. J., Kuo, C. S., Neff, N., Weissman, I. L., Quake, S. R., and Krasnow, M. A. (2020).  A molecular cell atlas of the human lung from single-cell RNA sequencing.  Nature, 587(7835):619–625.
  * Wang et al. (2020) Wang, H., Wang, M.-S., Wang, Y., Huang, Y.-Q., Shi, J.-P., Ding, Z.-L., and Wang, W.-J. (2020).  Prognostic value of immune related genes in lung adenocarcinoma.  Oncology Letters, 20(5):259.
  * Wang et al. (2009) Wang, Y., Kuan, P. J., Xing, C., Cronkhite, J. T., Torres, F., Rosenblatt, R. L., DiMaio, J. M., Kinch, L. N., Grishin, N. V., and Garcia, C. K. (2009).  Genetic defects in surfactant protein a2 are associated with pulmonary fibrosis and lung cancer.  The American Journal of Human Genetics, 84(1):52–59.
  * Wang et al. (2018) Wang, Y., Qian, W., and Yuan, B. (2018).  A graphical model of smoking-induced global instability in lung cancer.  IEEE/ACM Transactions on Computational Biology and Bioinformatics, 15(1):1–14.
  * Weinstein et al. (2017) Weinstein, A., Barber, R., and Candès, E. J. (2017).  A power and prediction analysis for knockoffs with lasso statistics.  preprint arXiv:1712.06465.
  * Xia et al. (2022) Xia, M., Wu, L., Sun, X., Han, X., Yan, H., Huang, J., Zhang, Y., Hu, Z., Zu, Y., Yin, C. C., and Qiu, X. (2022).  Next-generation sequencing revealed a distinct immunoglobulin repertoire with specific mutation hotspots in acute myeloid leukemia.  Biology, 11(2).
  * Xie et al. (2023) Xie, J., Lin, Y., Li, Y., Fang, A., Li, X., Wang, S., and Li, W. (2023).  lncrna trhde-as1 correlated with genomic landscape and clinical outcome in glioma.  Genes, 14(5).


# Logit-Based Ensemble Distribution Distillation for
Robust Autoregressive Sequence Uncertainties

Yassir Fathullah Engineering Department
University of Cambridge
UK  Guoxuan Xia Department of Electrical & Electronic Engineering
Imperial College London
UK  Mark J. F. Gales  Engineering Department
University of Cambridge
UK

###### Abstract

Efficiently and reliably estimating uncertainty is an important objective in
deep learning. It is especially pertinent to autoregressive sequence tasks,
where training and inference costs are typically very high. However, existing
research has predominantly focused on tasks with static data such as image
classification. In this work, we investigate Ensemble Distribution
Distillation (EDD) applied to large-scale natural language sequence-to-
sequence data. EDD aims to compress the superior uncertainty performance of an
expensive (teacher) ensemble into a cheaper (student) single model.
Importantly, the ability to separate knowledge (epistemic) and data
(aleatoric) uncertainty is retained. Existing probability-space approaches to
EDD, however, are difficult to scale to large vocabularies. We show, for
modern transformer architectures on large-scale translation tasks, that
modelling the ensemble logits, instead of softmax probabilities, leads to
significantly better students. Moreover, the students surprisingly even
outperform Deep Ensembles by up to $\sim$10% AUROC on out-of-distribution
detection, whilst matching them at in-distribution translation.

##  1 Introduction

The ability to produce reliable estimates of uncertainty is important to many
tasks in deep learning. When it is costly to make mistakes, a model should
know when to discard a prediction or defer it to a human expert. Although
there is a significant body of research in this area [Ovadia et al., 2019,
Hullermeier and Waegeman, 2021, Yang et al., 2021], it tends to focus on
static data, where outputs are of a fixed dimension, with approaches most
commonly evaluated on image classification [Geifman and El-Yaniv, 2019, Liu et
al., 2020, Yang et al., 2022, Moon et al., 2020, Xia and Bouganis, 2022b]. In
contrast, in this work, we aim to investigate uncertainty estimation for
sequence prediction tasks, such as machine translation, which is a relatively
under-explored domain [Malinin and Gales, 2021].

Large attention-based autoregressive neural networks have recently emerged as
the most competitive approach to many structured sequence-prediction tasks,
especially in translation [Bahdanau et al., 2015, Vaswani et al., 2017, Ott et
al., 2018], and are increasingly being used in practice. However, as the
computational and memory costs of these modern approaches are typically very
large, it is particularly important that approaches for improving the quality
of uncertainties should be efficient. We will focus on one such efficient
approach to better uncertainties, Ensemble Distribution Distillation (EDD)
[Malinin et al., 2020].

Ensembling multiple neural networks trained using different random seeds is a
well-established approach for boosting uncertainty performance
[Lakshminarayanan et al., 2017]. Deep Ensembles have been shown to be
effective over a wide range of data, tasks, and evaluation metrics [Ovadia et
al., 2019, Malinin and Gales, 2021, Kim et al., 2021, Gustafsson et al.,
2020]. Moreover, they are naturally able to decompose total uncertainty into
knowledge (epistemic) and data (aleatoric) uncertainty [Hullermeier and
Waegeman, 2021], which can be useful for different tasks such as active
learning [Gal et al., 2017, Radmard et al., 2021], reinforcement learning
[Depeweg et al., 2018], and out-of-distribution detection [Malinin and Gales,
2021]. However, Deep Ensembles suffer from costs that scale linearly with the
number of members. EDD aims to tackle this by using Knowledge Distillation
(KD) [Hinton et al., 2014] to compress the (teacher) ensemble into a more
efficient (student) single model. Crucially, EDD not only has the student
learn the predictions of the ensemble but also the distribution over
individual ensemble member outputs. By explicitly modelling the diversity over
the ensemble, the student is able to express knowledge and data uncertainty
independently just like the teacher ensemble [Malinin et al., 2020].

However, EDD is not without its challenges. Prior work has shown that EDD
suffers from optimisation issues, meaning it can be difficult to scale to
confident ensembles with large label spaces. Thus, EDD requires a number of
practical modifications in order to be applied to large-scale tasks such as
machine translation [Fathullah et al., 2021, Ryabinin et al., 2021]. Despite
these challenges, the concept behind EDD remains a promising approach for
training single autoregressive models with smaller footprints and the ability
to estimate high-quality, robust uncertainties.

Summary of contributions: In this paper, we focus on an underexplored area of
uncertainty estimation: robust and efficient autoregressive sequence
uncertainties. Specifically, we address the drawbacks of sequence EDD by using
logit-based ensemble distribution distillation (L-EDD). Instead of training a
student to distribution distil the information from an ensemble in the
probability/softmax space, we teach it to perform the same task in the pre-
softmax logit space. Experiments on the En-De WMT’16 and En-Ru WMT’20 machine
translation tasks show that L-EDD, in particular when using a Laplace
distribution, produces strong estimates of sequence uncertainty. L-EDD is able
to outperform EDD and surprisingly even Deep Ensembles on out-of-distribution
(OOD) detection and match them for translation quality. Furthermore, by using
Snapshot Ensembles [Huang et al., 2017], we are able to greatly reduce the
overall training costs of EDD compared to using a Deep Ensemble teacher.

##  2 Background

In this section, we review ensemble-based uncertainty estimation. We follow
with a discussion of how the limitations of ensembles can be addressed using
recently developed distillation techniques for autoregressive sequence tasks
such as machine translation.

###  2.1 Uncertainty Estimation

We adopt a Bayesian perspective on ensembles as this offers a flexible
framework within which uncertainties have an information-theoretic
justification. The posterior over model parameters ${\tt
p}(\bm{\theta}|\mathcal{D})$ is derived given some observed (training) data,
${\cal D}$. Unfortunately, the posterior is often intractable and cannot be
derived for large non-linear networks. Alternatively an approximation ${\tt
q}(\bm{\theta})\approx{\tt p}(\bm{\theta}|\mathcal{D})$ can be used. Samples
from this approximate distribution can then be drawn to generate an ensemble
of models.

Take an ensemble $\\{{\tt P}(\bm{y}|\bm{x},\bm{\theta}^{(m)})\\}_{m=1}^{M}$
sampled from an approximate posterior ${\tt q}(\bm{\theta})$ where each model
maps a variable-length input $\bm{x}\in\mathcal{X}$ into a variable-length
output $\bm{y}\in\mathcal{Y}$ of discrete units. The predictive distribution
is obtained by:

$\displaystyle{\tt P}(\bm{y}|\bm{x},\mathcal{D})=\mathbb{E}_{{\tt
q}(\bm{\theta})}\left[{\tt P}(\bm{y}|\bm{x},\bm{\theta})\right].$ (1)

From this predictive distribution, a measure of total uncertainty can be
estimated using the entropy:

$\displaystyle\mathcal{H}\left[{\tt
P}(\bm{y}|\bm{x},\mathcal{D})\right]=\mathbb{E}_{{\tt
P}(\bm{y}|\bm{x},\mathcal{D})}\left[-\ln{\tt
P}(\bm{y}|\bm{x},\mathcal{D})\right].$ (2)

Furthermore, a measure of disagreement between models, also referred to as
knowledge or epistemic uncertainty, can be estimated by using mutual
information between $\bm{y}$ and $\bm{\theta}$:

$\displaystyle\begin{split}\mathcal{I}\left[\bm{y},\bm{\theta}|\bm{x},\mathcal{D}\right]=\mathbb{E}_{{\tt
q}(\bm{\theta})}\left[{\tt KL}\big{(}{\tt
P}(\bm{y}|\bm{x},\bm{\theta})\;\big{\|}\;{\tt
P}(\bm{y}|\bm{x},\mathcal{D})\big{)}\right].\end{split}$ (3)

This estimate can also be decomposed into a measure of total and data
(aleatoric) uncertainty, as mentioned in Malinin and Gales [2021]. There are
also many other potential measures of knowledge uncertainty such as expected
pairwise KL-divergence or reverse mutual information [Malinin and Gales,
2021], however, for the sake of simplicity we restrict our focus to the
already mentioned eq. (2) and (3) since these represent uncertainties of
differing natures.

Limitations: The discussion has so far assumed one can enumerate all possible
variable-length outputs $\bm{y}\in\mathcal{Y}$ which is not tractable in
autoregressive sequence tasks. Instead, one can approximate the uncertainties
by monte-carlo methods [Notin et al., 2021] and utilising the autoregressive
structure of predictions [Malinin and Gales, 2021]:

$\displaystyle\begin{split}{\tt
P}(\bm{y}|\bm{x},\bm{\theta})=\prod_{l=1}^{L}{\tt
P}(y_{l}|\bm{y}_{<l},\bm{x},\bm{\theta}).\end{split}$ (4)

We refer to Malinin and Gales [2021] for an in-depth discussion and analysis
of approximations for predictive entropy and mutual information for
autoregressive prediction.

###  2.2 Knowledge Distillation

Ensembles $\\{{\tt P}(\bm{y}|\bm{x},\bm{\theta}^{(m)})\\}_{m=1}^{M}$ sampled
from some posterior can be computationally demanding. One approach to
efficiently exploit the information of the ensemble is to use Knowledge
Distillation (KD) to yield a single student model [Hinton et al., 2014, Kim
and Rush, 2016].

Given a reference data pair $(\bm{x},\bm{y})\sim\tilde{{\tt
p}}(\bm{x},\bm{y})$, a standard model might be trained using negative log-
likelihood (NLL):

$\displaystyle\mathcal{L}_{\tt
NLL}(\bm{\theta})=-\frac{1}{L}\sum_{l=1}^{L}\ln{\tt
P}(y_{l}|\bm{y}_{<l},\bm{x},\bm{\theta}).$ (5)

This is referred to as teacher-forcing since during training the model makes
predictions at step $l$ conditioned on the true output $\bm{y}_{<l}$ (rather
than its own previous predictions). Similarly, a student model with parameters
$\bm{\phi}$ can be trained to emulate a teacher ensemble by additionally using
average ensemble categorical/softmax outputs $\bm{\pi}_{l},\pi_{l,k}={\tt
P}(y_{l}=k|\bm{y}_{<l},\bm{x},\mathcal{D})$ as soft labels:

$\displaystyle\begin{split}\mathcal{L}_{\tt
KL}(\bm{\phi})=\frac{1}{L}\sum_{l=1}^{L}{\tt
KL}\big{(}\bm{\pi}_{l}\;\big{\|}\;{\tt
P}(y_{l}|\bm{y}_{<l},\bm{x},\bm{\phi})\big{)}.\end{split}$ (6)

However in practice, one optimises a convex combination of the likelihood and
KL-divergence losses $\mathcal{L}_{\tt KD}(\bm{\phi})=\lambda\mathcal{L}_{\tt
NLL}(\bm{\phi})+(1-\lambda)\mathcal{L}_{\tt
KL}(\bm{\phi}),\medspace\lambda\in[0,1]$ for added supervision and stability.
The probability mass functions in the KL-divergence can also be temperature
scaled to improve optimisation [Hinton et al., 2014]. Note that this criterion
is only considered for the teacher-forcing case, more sophisticated
distillation approaches exist, by sampling $(\bm{x},\bm{y})$ from alternative
distributions, but is outside the scope for this work, see Kim and Rush
[2016], Wong et al. [2016], Lee et al. [2018], Malinin et al. [2017] for
details.

###  2.3 Ensemble Distribution Distillation

Whilst KD has been successful in many sequence tasks, the resulting student is
not able to estimate knowledge uncertainty, since it only models the average
ensemble output. To avoid this issue, Malinin et al. [2020], Ryabinin et al.
[2021] consider the task of distilling the distribution of sequence ensemble
predictions onto a single student. This allows the student to retain both
predictive performance and information about ensemble diversity.

To explain the mechanics behind Ensemble Distribution Distillation (EDD),
consider modelling a distribution over autoregressive ensemble predictions, in
which all $M$ ensemble members share the same back-history $\bm{y}_{<l}$:

$\displaystyle\\{\bm{\pi}_{l}^{(m)}\\}_{m=1}^{M},\medspace\pi_{l,k}^{(m)}={\tt
P}(y_{l}=\omega_{k}|\bm{y}_{<l},\bm{x},\bm{\theta}^{(m)}).$ (7)

Now let an autoregressive student predict the parameters $\bm{\alpha}_{l}$ of
a Dirichlet distribution ${\tt Dir}(\bm{\pi}_{l}|\bm{\alpha}_{l})={\tt
p}(\bm{\pi}_{l}|\bm{y}_{<l},\bm{x},\bm{\phi})$. Since the Dirichlet models a
distribution over categorical distributions it is an ideal candidate for this
task. The distribution distillation loss of such a model is then simply the
result of (negative) log-likelihood:

$\displaystyle\begin{split}\mathcal{L}_{\tt NLL}^{\tt
DD}(\bm{\phi})&=-\frac{1}{MLK}\sum_{m,l}\ln{\tt
Dir}(\bm{\pi}_{l}^{(m)}|\bm{\alpha}_{l})\\\
&\equiv\frac{1}{LK}\sum_{l}\Big{(}\ln
B(\bm{\alpha}_{l})-\sum_{k}\alpha_{l,k}\ln\tilde{\pi}_{l,k}\Big{)},\
\end{split}$ (8)

where $K$ is the number of classes, $B(\bm{\alpha})$ is the beta function and
$\tilde{\pi}_{l,k}$ is the geometric average of the individual ensemble
softmax probabilities in Equation (7).

Whilst this approach was shown to be promising on a small-scale image
classification task in Malinin et al. [2020], following work [Fathullah et
al., 2021, Ryabinin et al., 2021] found that direct application of Equation
(8) encounters optimisation issues when scaled to larger label spaces. This
arises from the way classwise loss gradients are related to teacher class
probabilities. It turns out that, unlike standard distillation, the loss in
Equation (8) induces small gradients for (important) high-probability classes
and large gradients for (unimportant) low-probability classes. This negatively
affects convergence as the number of low-probability classes increases.
Ryabinin et al. [2021] proposed an approach for scaling Dirichlet EDD, where
the student aims to minimise a normalized reverse KL-divergence to a proxy
Dirichlet, which will be used as a baseline in this work.

##  3 Sequence Logit-based EDD

In sequence tasks with a large number of classes, which commonly occurs in
speech recognition and machine translation, the output categorical
distributions are often very sparse and concentrated. Therefore, it often
becomes highly challenging to apply EDD to tasks of this nature. On the other
hand, KD has been shown to work well for larger tasks [Kim and Rush, 2016,
Gaido et al., 2020, Tan et al., 2019, Jiao et al., 2019], but since it only
models the average teacher predictions, it cannot estimate data and knowledge
uncertainties that are important for many downstream tasks such as out-of-
distribution detection.

In this section, we describe a Logit-based Ensemble Distribution Distillation
(L-EDD) approach for autoregressive models which addresses the drawbacks of
both KD and EDD in a single consistent framework and is scalable to sequence
problems with a large number of classes. Consider a set of logits produced by
an ensemble:

$\displaystyle\\{\bm{z}_{l}^{(m)}\\}_{m=1}^{M},\medspace\bm{\pi}_{l}^{(m)}={\tt
Softmax}(\bm{z}_{l}^{(m)}).$ (9)

Traditional distillation approaches thereafter use the logits to produce
categorical probability distributions by applying the softmax function.
However, instead of operating in the probability space, we propose training a
student, with model parameters $\bm{\phi}$, to directly model the logit space
by predicting the mean $\bm{\mu}_{l}$ and scale $\bm{\sigma}_{l}$ parameters
of a diagonal Laplace distribution:

$\displaystyle\begin{split}{\tt p}(\bm{z}|\bm{y}_{<l},\bm{x},\bm{\phi})&={\tt
Lap}(\bm{z}|\bm{\mu}_{l},\bm{\sigma}_{l})\\\
&=\prod_{k}\frac{1}{2\sigma_{l,k}}\exp{-\frac{\absolutevalue{z_{k}-\mu_{l,k}}}{\sigma_{l,k}}}.\end{split}$
(10)

Because we opt for a diagonal distribution, sampling is parallelisable, highly
efficient, and straightforward and allows for the estimation of uncertainties
in exactly the same manner as in standard ensembles. Additionally,
significantly fewer parameters are required compared to using a fully-
specified covariance matrix. Another benefit of the chosen distribution is the
long tails which make the Laplace robust to outliers, unlike the Gaussian
distribution. This robustness also makes it a natural choice for handling the
early stages of training when the student model is randomly initialised and
its output distribution substantially differs from the ensemble logits.

Furthermore, given the set of logits produced by an ensemble, the student
model ${\tt p}(\bm{z}|\bm{y}_{<l},\bm{x},\bm{\phi})$ can be trained by
straightforward application of log-likelihood training:

$\displaystyle\begin{split}\mathcal{L}_{\tt NLL}^{\tt
L-EDD}(\bm{\phi})&=-\frac{1}{MLK}\sum_{m,l}\ln{\tt
Lap}(\bm{z}_{l}^{(m)}|\bm{\mu}_{l},\bm{\sigma}_{l})\\\
&\equiv\frac{1}{MLK}\sum_{m,l,k}\frac{\absolutevalue{z_{l,k}^{(m)}-\mu_{l,k}}}{\sigma_{l,k}}+\ln\sigma_{l,k}.\end{split}$
(11)

We also perform experiments with a student (diagonal) Gaussian distribution on
the logits, variations of which have been explored in static image
classification [Fathullah and Gales, 2022, Lindqvist et al., 2020] but
remained unexplored for autoregressive sequence tasks:

$\displaystyle\begin{split}{\tt
p}(\bm{z}|\bm{y}_{<l},&\bm{x},\bm{\phi})=\mathcal{N}(\bm{z}|\bm{\mu}_{l},\bm{\sigma}_{l}^{2})\\\
&=\prod_{k}\frac{1}{(2\pi\sigma_{l,k}^{2})^{\frac{1}{2}}}\exp{-\frac{(z_{k}-\mu_{l,k})^{2}}{2\sigma_{l,k}^{2}}}.\end{split}$
(12)

Similar to all of the mentioned approaches, this system is also trained using
the log-likelihood objective:

$\displaystyle\begin{split}\mathcal{L}_{\tt NLL}^{\tt
L-EDD}(\bm{\phi})&=-\frac{1}{MLK}\sum_{m,l}\ln\mathcal{N}(\bm{z}_{l}^{(m)}|\bm{\mu}_{l},\bm{\sigma}_{l}^{2}).\end{split}$
(13)

The Gaussian distribution, which induces an L2-norm loss function is much more
sensitive to outliers in the ensemble outputs. This student could potentially
be more challenging to train, but should still be more stable than Dirichlet
EDD.

###  3.1 Practical Considerations

Since the softmax activation function is shift invariant,

$\displaystyle{\tt Softmax}(\bm{z}-\bm{1}b)={\tt
Softmax}(\bm{z})\thickspace\thickspace\forall b\in\mathbb{R},$

one has to consider this property when performing distribution distillation.
Ensemble members are unconstrained along $\bm{1}$, and so can potentially vary
wildly in the logit space, even if they give consistent softmax predictions.
Therefore, logits are normalised by $\tilde{\bm{z}}=\bm{z}-\bm{1}{\tt
LogSumExp}(\bm{z})$. This particular normalisation scheme is not special and
any choice of the normalisation constant such as ${\tt Max}(\bm{z})$ or ${\tt
Mean}(\bm{z})$ would be valid. Next, to ensure that the student can be trained
reliably, we interpolate the knowledge and distribution distillation losses
$\mathcal{L}_{\tt KD}(\bm{\phi})+\beta\mathcal{L}_{\tt NLL}^{\tt
L-EDD}(\bm{\phi})$ (see Eq. 11 and Sec. 2.2).

Furthermore, distributions in logit space often lead to analytically
intractable expectations in probability space. The standard approach to
circumvent this issue is by sampling from the distribution using monte-carlo
approximations. However, in this paper, we opt for an approximative
deterministic approach when computing the predictive distribution (e.g. when
decoding):

$\displaystyle{\tt P}(y_{l}|\bm{y}_{<l},\bm{x},\bm{\phi})$
$\displaystyle=\mathbb{E}_{{\tt
p}(\bm{z}_{l}|\bm{y}_{<l},\bm{x},\bm{\phi})}\left[{\tt
Softmax}(\bm{z}_{l})\right]$ $\displaystyle\approx{\tt
Softmax}(\bm{\mu}_{l}),$

in which we approximate the expectation by just using the mean of the logit
distribution. When performing downstream tasks that require uncertainties we
revert to a stochastic sampling scheme to generate multiple predictions from
the distribution.

##  4 Experiments on Artificial Data

1.0

(a) Dataset.

1.0

(b) Distillation loss

1.0

(c) Dirichlet loss

1.0

(d) Laplace loss

1.0

(e) Ensemble conf.

1.0

(f) Distillation conf.

1.0

(g) Dirichlet conf.

1.0

(h) Laplace conf.

Figure 1: An artificial three-class classification problem with 1000 examples
per class. The top row shows the loss surface contours for various
distillation approaches; darker colours imply lower losses. The bottom row
shows the corresponding confidence contours (the confidence scores are
reported on the contour). This shows that Dirichlet-based EDD is unable to
learn properly whilst our proposed Laplace L-EDD can imitate an ensemble.

This section investigates the proposed Laplace logit-based ensemble
distribution distillation (L-EDD) technique on a static artificial dataset,
see Figure 1(a). The dataset was generated by sampling 3000 data points from
three isotropic Gaussian distributions equally. The location and standard
deviation of the Gaussians were chosen such that there would be regions with
significant overlap and regions where models can be highly confident.

In these exploratory experiments, an ensemble of 10 small neural networks is
first trained by randomly initialising each member. Thereafter, the ensemble
is distilled using KD, EDD and Laplace L-EDD. We perform a qualitative
comparison of these methods by displaying both the loss surface (see Figures
1(c)-1(d)) of each approach and the resulting confidence (maximum softmax
probability) contours ((see Figures 1(e)-1(h))). The loss surface shows how
the student distillation loss varies over the input space, thus providing
useful information about which regions of the data are successfully optimised.
The confidence contours are useful to understand if the system can separate
between each of the three classes. EDD training had to be terminated early as
it diverged due to large gradients originating from the high-confidence
regions (as discussed in Ryabinin et al. [2021]).

Figure 1(e) shows the ensemble confidence contours which clearly trace out
class boundaries and partially separate the three classes. The confidence also
increases as one moves further away from regions of overlap since there is
less uncertainty. This is the behaviour we expect from a properly trained
system and further, expect distilled students to behave similarly. Next, we
knowledge distil the ensemble onto a single student model, see Figures 1(b)
and 1(f). The loss surface shows that the student can optimise the
distillation objective over regions with high overlap well and generate
confidence score contours that are consistent with the teacher ensemble.

However, Ensemble Distribution Distillation completely fails on this very
simple task. Observing the loss surface in Figure 1(c), one can infer that the
Dirichlet student is unable to optimise regions where there is significant
data overlap, instead displaying extremely small losses in regions for which
the teacher ensemble is already confident. This links back to a result in
Ryabinin et al. [2021] in which they find that highly confident teachers can
induce extremely large gradients. This also translates into inaccurate
confidence contours which are unable to separate between classes, especially
in regions of overlap, see Figure 1(g).

The Laplace L-EDD approach circumvents these issues by operating in the logit
space. The resulting loss surface is much more consistent with the knowledge-
distilled student, optimising regions of high overlap, see Figure 1(d).
Similarly, the confidence contours trace out boundaries consistent with the
ensemble but with lower overall confidence. Since the log-likelihood (similar
to KL-divergence) is a mode-covering objective [Minka et al., 2005], the
Laplace student ends up predicting distributions that overestimate the range
of ensemble logits. Couple this with the long tails of the distributions and
the Laplace will often overestimate the variance in logits and produce lower
overall confidence.

##  5 Machine Translation

This section reports on the performance of base transformers [Vaswani et al.,
2017] trained on the En-De WMT’16 dataset, consisting of 4.5 million sentence
pairs covering topics such as news & policy. We use newstest-13 for validation
and newstest-14 for predictive evaluation. For the main task of investigating
out-of-distribution (OOD) detection, we compare the in-distribution (ID)
newstest-14 with one of the publically available Khresmoi-Summary (Khresmoi)
[Dušek et al., 2017], MTNT [Michel and Neubig, 2018] and Kyoto Free
Translation Task (KFTT) [Neubig, 2011] datasets. These datasets relate to
medical articles, Reddit-based noisy conversational text and specialised
Wikipedia articles, respectively. Furthermore, we apply insights from training
these systems to big transformers [Vaswani et al., 2017] trained on the larger
En-Ru WMT’20 consisting of 58 million pairs after processing. In this case, we
use newstest-19 for validation and newstest-20 for evaluation. The OOD
detection task uses newstest-20 as ID and the same OOD datasets as above for
the base transformer.

Data is tokenized using Moses, following Ott et al. [2018]. For WMT’16, a
shared dictionary is trained using Byte Pair Encoding (BPE) with 32,000 merge
operations [Sennrich et al., 2016]. For WMT’20 we learn disjoint dictionaries
using BPE with 40,000 merge operations. The predictive performance
(translation quality) will be evaluated using corpus-level (Sacre)BLEU [Post,
2018], with no post-processing of outputs before being scored. For the main
task of detection, we use the ubiquitous threshold-independent AUROC metric
[Manning and Schütze, 1999], with baseline random detection corresponding to a
score of 50%.

All standard transformers are trained using an inverse square root with a
linear warmup stage. A stronger Deep Ensemble baseline is formed by taking
$M=5$ such models. To avoid the high training cost of Deep Ensembles, we also
train Snapshot Ensembles [Xie et al., 2013, Huang et al., 2017] with a cyclic
learning rate [Smith, 2017] to showcase that distribution distillation can be
achievable with smaller training budgets. Since building Deep Ensembles is
expensive, and the performance difference to Snapshot Ensembles was shown to
be small, we opted to perform most distillation experiments on Snapshot
Ensembles. We compare the proposed L-EDD approaches with KD and EDD and repeat
each experiment 5 times, each with a different Snapshot Ensemble. Finally,
similar to a range of prior work on uncertainty estimation tasks [Malinin et
al., 2020, Malinin and Gales, 2021, Corbière et al., 2019], we do not aim to
achieve state-of-the-art predictive performance but opt for a simpler setup
with a focus on achieving better uncertainty estimation. All setup details are
provided in Appendix A. Hyperparameters were determined on ID validation sets.

Table 1: Model parameter size, relative training time and translation
performance on newstest-14 $\pm$2std (BLEU) for base transformer. We include
two different KD baselines, one for each teacher ensemble.

Model Size Train BLEU $\uparrow$ Time Standard 60.9M 1.0 25.85 $\pm$ 0.17 Deep
Ensemble 304.5M 5.0 26.72 KD (Categorical) 60.9M 5.9 26.70 $\pm$ 0.26 Snapshot
Ensemble 304.5M 1.5 26.54 $\pm$ 0.16 KD (Categorical) 60.9M 1.9 27.02 $\pm$
0.19 EDD (Dirichlet) 60.9M 2.0 26.96 $\pm$ 0.06 L-EDD (Gaussian) 77.8M 1.9
26.90 $\pm$ 0.28 L-EDD (Laplace) 77.8M 1.9 27.08 $\pm$ 0.20

Table 2: OOD detection performance (%AUROC $\uparrow$ $\pm$ 2 std) for base transformer with ID dataset newtest-14 and OOD datasets Khresmoi, MTNT and KFTT. Bold indicates best in a column, underline second best. Laplace L-EDD with knowledge uncertainty (KU), shows superior performance for all OOD datasets even compared to the Deep Ensemble.  Model | Khresmoi | MTNT | KFTT
---|---|---|---
TU | KU | TU | KU | TU | KU
Standard |  47.5 $\pm$ 0.8 | ✗ |  63.5 $\pm$ 1.3 | ✗ |  30.6 $\pm$ 1.2 | ✗
Deep Ensemble | 48.0 | 61.9 | 64.5 | 63.7 | 30.1 | 44.0
KD (Categorical) |  47.9 $\pm$ 1.1 | ✗ |  64.5 $\pm$ 1.3 | ✗ |  29.8 $\pm$ 0.7 | ✗
Snapshot Ensemble |  49.0 $\pm$ 0.6 |  62.6 $\pm$ 1.1 |  63.8 $\pm$ 1.2 |  63.1 $\pm$ 0.7 |  31.7 $\pm$ 0.9 |  47.4 $\pm$ 2.5
KD (Categorical) |  48.0 $\pm$ 1.4 | ✗ |  64.6 $\pm$ 0.9 | ✗ |  31.3 $\pm$ 0.5 | ✗
EDD (Dirichlet) |  49.6 $\pm$ 1.3 |  57.1 $\pm$ 1.4 |  65.1 $\pm$ 1.7 |  65.6 $\pm$ 2.0 |  31.0 $\pm$ 0.9 |  36.2 $\pm$ 1.4
L-EDD (Gaussian) |  59.5 $\pm$ 1.1 |  71.7 $\pm$ 1.9 |  66.3 $\pm$ 1.6 |  64.0 $\pm$ 2.1 |  35.8 $\pm$ 1.2 |  44.0 $\pm$ 0.2
L-EDD (Laplace) |  65.1 $\pm$ 1.8 |  73.1 $\pm$ 1.7 |  65.1 $\pm$ 1.5 |  66.8 $\pm$ 1.8 |  37.8 $\pm$ 0.2 |  48.8 $\pm$ 1.4

###  5.1 Base Transformer Results

Table 1 shows both the efficiency and performance of a wide range of systems
on newstest-14. As expected the performance of the Deep Ensemble trumps both
the Snapshot Ensemble and a standard trained system. Surprisingly, Snapshot
Ensemble distilled students achieve better performance, a pattern also
observed in self-distilled systems and is explored in more detail in Allen-Zhu
and Li [2021].

Next, we compare the threshold-independent out-of-distribution detection
performance of baseline systems with L-EDD models. From Table 2, we observe
that Snapshot Ensembles are able to compete with the Deep equivalent whilst
being more than 3 times cheaper to train. Furthermore, the knowledge-distilled
students are able to match the detection performance of their Deep and
Snapshot ensemble teachers using total uncertainty (TU). This is a natural
result since they were specifically designed to capture the predictive
distribution of their teacher ensemble. However, since KD students are unable
to estimate knowledge uncertainty (KU), they fail to reach ensemble-level
detection performance in all but the MTNT dataset. Similarly, the modified
Dirichlet baseline as described by Ryabinin et al. [2021] is able to achieve
similar detection performance using TU but with the added ability to estimate
KU. And whilst the Dirichlet KU are often better than its TU estimates, they
often fall short when compared to ensembles.

On the other hand, the Laplace & Gaussian L-EDD models are (surprisingly) able
to outperform both ensembles in all three detection splits, producing either
similar or significantly better TU and KU estimates. This may partially be due
to the fact that diagonal Laplace and Gaussian distributions have more
parameters and are more flexible than the Dirichlet, and also because they do
not suffer from the same optimisation issues. Nonetheless, neither reason
explains why L-EDD models can outperform ensembles in detection. We explore
this pattern in Section 6.

Additionally, many models are worse than a random detector, especially for the
KFTT and partially for the Khresmoi dataset. A partial explanation could be
that these datasets contain longer sequences. When decoding, the transformer
models produce more and more confident predictions further along in the output
sequence, causing lower uncertainty scores. Section 6 investigates this effect
and isolates a possible reason behind Laplace’s success in outperforming
ensembles.

As an aside, we observe that estimates of knowledge uncertainty are clearly
important for OOD detection for autoregressive sequence tasks (and this is
corroborated in prior work Malinin and Gales [2021], Ryabinin et al. [2021]).
This is in contrast to recent empirical results on image classification data,
which show the opposite, that measures of knowledge uncertainty are not useful
for indicating distributional shifts [Xia and Bouganis, 2022a, Abe et al.,
2022].

Table 3: Model parameter size, relative training time and translation
performance on newstest-20 $\pm$2std (BLEU) for the big transformer.

Model Size Train BLEU $\uparrow$ Time Standard 271M 1.0 26.28 $\pm$ 0.34 Deep
Ensemble 1.35B 5.0 26.81 Snapshot Ensemble 1.35B 1.5 26.42 $\pm$ 0.23 KD
(Categorical) 271M 2.1 26.73 $\pm$ 0.16 EDD (Dirichlet) 271M 2.2 26.66 $\pm$
0.19 L-EDD (Laplace) 320M 2.2 26.71 $\pm$ 0.18

Table 4: OOD detection performance (%AUROC $\uparrow$ $\pm$ 2 std) for big transformer with ID dataset newtest-20 and OOD datasets Khresmoi, MTNT and KFTT. Bold indicates best in a column, underline second best. Similar to Table 4, L-EDD (Laplace) with KU shows superior performance over all OOD datasets. Model | Khresmoi | MTNT | KFTT
---|---|---|---
TU | KU | TU | KU | TU | KU
Deep Ensemble | 39.3 | 53.2 | 70.8 | 69.0 | 51.0 | 60.3
Snapshot Ensemble |  40.8 $\pm$ 0.5 |  55.0 $\pm$ 0.8 |  70.1 $\pm$ 0.5 |  69.3 $\pm$ 0.9 |  51.1 $\pm$ 0.6 |  60.9 $\pm$ 1.4
KD (Categorical) |  40.4 $\pm$ 0.8 | ✗ |  70.9 $\pm$ 1.0 | ✗ |  50.9 $\pm$ 0.6 | ✗
EDD (Dirichlet) |  41.0 $\pm$ 0.8 |  52.9 $\pm$ 1.3 |  71.3 $\pm$ 0.7 |  69.4 $\pm$ 0.7 |  51.0 $\pm$ 0.8 |  60.0 $\pm$ 1.3
L-EDD (Laplace) |  51.0 $\pm$ 0.9 |  63.4 $\pm$ 1.2 |  72.6 $\pm$ 0.8 |  70.2 $\pm$ 0.6 |  63.2 $\pm$ 1.0 |  70.2 $\pm$ 1.1

###  5.2 Big Transformer Results

In this section, we take the best-performing systems from the previous section
and apply them to the ‘big transformer’ on the larger En-Ru WMT’20 dataset.
Table 3 shows the efficiency and predictive performance on newstest-20. Again,
we observe that the Deep Ensemble outperforms its Snapshot equivalent.
Furthermore, the KD and L-EDD students, distilled from the Snapshot Ensemble
were able to outperform their teacher. However, unlike the smaller-scale
experiments, these students were only able to reach Deep Ensemble performance
within a standard deviation, but were able to do so with a single forward
pass.

From Table 4 we observe a similar pattern in which Deep and Snapshot ensembles
perform equivalently whilst L-EDD (Laplace) is able to significantly
outperform both ensembles in all but the MTNT dataset. Interestingly, unlike
in Table 2 where no model was able to beat a random detector on the KFTT
detection, the larger En-Ru WMT’20 based models are able to differentiate
between newstest-20 and KFTT; switching the ID dataset to newstest-14 does not
affect the results notably.

##  6 Analysis: Ensemble vs Laplace

###  6.1 Augmented Ensemble Uncertainties

Both Sections 5.1 and 5.2 found that L-EDD models overall significantly
outperformed their teacher Snapshot Ensemble and a Deep Ensemble. Therefore,
we propose an alternative experiment to understand the source of L-EDD’s
superior performance. We fit an auxiliary Laplace distribution to a Deep
Ensemble during inference and use the samples from this proxy to perform the
detection task.

Consider a Deep Ensemble which produces a set of normalised logits
$\\{\tilde{\bm{z}}_{l}^{(m)}\\}_{m=1}^{M}$ as in Equation (9). In traditional
uncertainty estimation, these logits would be transformed into categorical
distributions. However, in this experiment, we estimate an auxiliary Laplace
distribution using maximum likelihood (which is the loss-minimising
distribution for a Laplace L-EDD student):

$\displaystyle\tilde{\bm{\mu}}_{l},\tilde{\bm{\sigma}}_{l}=\operatorname*{arg\,max}_{\bm{\mu},\bm{\sigma}}\sum_{m}\ln{\tt
Lap}(\tilde{\bm{z}}_{l}^{(m)}|\bm{\mu},\bm{\sigma}).$ (14)

By sampling new points from this auxiliary distribution, we can estimate total
and knowledge uncertainty:

$\displaystyle\tilde{\bm{\pi}}={\tt Softmax}(\bm{z}),\thickspace\bm{z}\sim{\tt
Lap}(\tilde{\bm{\mu}}_{l},\tilde{\bm{\sigma}}_{l}).$ (15)

The aim of this modified approach to ensemble-based uncertainty estimation is
to investigate whether or not approximating the logits with a Laplace
distribution is the reason behind L-EDD performing better.

Table 5: OOD detection performance (%AUROC $\uparrow$ $\pm$ 2 std) following the same setup as in Table 2. The Laplace augmented ensemble demonstrates much better performance in most cases compared to its standard counterpart. Model | Khresmoi | MTNT | KFTT
---|---|---|---
TU | KU | TU | KU | TU | KU
Deep Ensemble | 48.0 | 61.9 | 64.5 | 63.7 | 30.1 | 44.0
Deep Ensemble (Laplace) | 62.5 | 72.1 | 63.8 | 56.1 | 34.8 | 55.4
L-EDD (Laplace) |  65.1 $\pm$ 1.8 |  73.1 $\pm$ 1.7 |  65.1 $\pm$ 1.5 |  66.8 $\pm$ 1.8 |  37.8 $\pm$ 0.2 |  48.8 $\pm$ 1.4

Table 5 shows the detection performance of the Laplace-modified Deep Ensemble,
following the detection setup in Section 5.1. Clearly, the augmented ensemble
bridges the OOD detection performance gap between standard Deep Ensemble and
L-EDD. This suggests that measures of uncertainty based on directly fit,
logit-space models of an ensemble are better at indicating distributional
shift than directly using the ensemble logits, for autoregressive sequence
prediction. We remark that fully understanding why this is the case would be
an interesting direction of future research, as it may enable further
advancement in autoregressive out-of-distribution detection.

###  6.2 Model Confidence for Longer Sequences

Another reason for Laplace L-EDD’s superior performance could be found in
analysing behaviour with increasing sequence lengths.

(a) PCC = -20%

(b) PCC = -29%

(c) PCC = -3%

(d) PCC = 1%

Figure 2: Length vs uncertainty density plots for ID (newstest-14) and OOD
(Khresmoi) datasets. The left half corresponds to the Deep Ensemble and the
right half to Laplace L-EDD. Each figure caption also shows the Pearson
Correlation Coefficient (PCC) between sequence length and total uncertainty.
Total uncertainty and length are negatively correlated for the ensemble, i.e.
it is more confident on longer sequences, but they are uncorrelated for L-EDD.

Figure 2 shows how Deep Ensemble and L-EDD total uncertainties scale with
output sequence length for both ID (newstest-14) and OOD (Khresmoi) datasets.
Under each figure is also the associated Pearson correlation.

Three observations can be made from this data. The first is that the L-EDD
system consistently outputs much higher total uncertainties. The second is
that the Deep Ensemble displays a negative correlation between total
uncertainty and sequence length. This implies that the ensemble is more
confident in translating longer sequences, but this is also why it fails in
detecting OOD datasets which contain longer sequences. The third and possibly
most significant observation is that the Laplace system shows almost no
correlation between total uncertainty and sequence length for both ID and OOD
datasets, effectively eliminating the length bias, and allowing it to better
differentiate between ID and OOD datasets even when the OOD inputs differ in
length from what the detection system was trained on.

##  7 Conclusion

In this work, we investigate the efficient estimation of uncertainties for
large-scale autoregressive sequence prediction. To this end, we examine
Ensemble Distribution Distillation (EDD) in the logit-space, in order to
bypass optimisation issues found in softmax-space EDD. We perform experiments
using modern transformer models trained to perform large-scale machine
translation. They show that a student model trained to parameterise a Laplace
distribution over logits is able to significantly outperform Deep Ensembles
for OOD detection at a fraction of the inference cost, whilst matching the
ensemble for translation quality. Moreover, we show that the use of Snapshot
Ensembling can greatly reduce the training costs of EDD, without sacrificing
translation performance. We hope that our work can encourage further
investigation into the comparatively less well-explored domain of uncertainty
estimation for structured sequence prediction, on tasks such as machine
translation, image captioning, and automatic speech recognition.

###### Acknowledgements.

Guoxuan Xia is funded jointly by Arm ltd. and EPSRC.

## References

  * Abe et al. [2022] Taiga Abe, E. Kelly Buchanan, Geoff Pleiss, Richard Zemel, and John Patrick Cunningham.  Deep ensembles work, but are they necessary?  In _Advances in Neural Information Processing Systems_ , 2022.
  * Allen-Zhu and Li [2021] Zeyuan Allen-Zhu and Yuanzhi Li.  Towards understanding ensemble, knowledge distillation and self-distillation in deep learning.  In _arXiv_ , 2021.
  * Bahdanau et al. [2015] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural machine translation by jointly learning to align and translate.  In _International Conference on Learning Representations_ , 2015.
  * Corbière et al. [2019] Charles Corbière, Nicolas Thome, Avner Bar-Hen, Matthieu Cord, and Patrick Pérez.  Addressing failure prediction by learning model confidence.  In _Advances in Neural Information Processing Systems 32_. 2019.
  * Depeweg et al. [2018] Stefan Depeweg, Jose Miguel Hernández-Lobato, Finale Doshi-Velez, and Steffen Udluft.  Decomposition of uncertainty in bayesian deep learning for efficient and risk-sensitive learning.  In _International Conference on Learning Representations_ , 2018.
  * Dušek et al. [2017] Ondřej Dušek, Jan Hajič, Jaroslava Hlaváčová, Jindřich Libovický, Pavel Pecina, Aleš Tamchyna, and Zdeňka Urešová.  Khresmoi summary translation test data 2.0.  http://hdl.handle.net/11234/1-2122, 2017.  LINDAT/CLARIAH-CZ digital library at the Institute of Formal and Applied Linguistics (ÚFAL).
  * Fathullah and Gales [2022] Yassir Fathullah and Mark J. F. Gales.  Self-distribution distillation: efficient uncertainty estimation.  In _Proceedings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence_ , volume 180 of _Proceedings of Machine Learning Research_ , 01–05 Aug 2022.
  * Fathullah et al. [2021] Yassir Fathullah, Mark J. F. Gales, and Andrey Malinin.  Ensemble distillation approaches for grammatical error correction.  In _International Conference on Acoustics, Speech, and Signal Processing_ , 2021.
  * Gaido et al. [2020] Marco Gaido, Mattia A. Di Gangi, Matteo Negri, and Marco Turchi.  End-to-end speech-translation with knowledge distillation: FBK@IWSLT2020.  In _International Conference on Spoken Language Translation_ , 2020\.
  * Gal et al. [2017] Yarin Gal, Riashat Islam, and Zoubin Ghahramani.  Deep Bayesian active learning with image data.  In _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , 06–11 Aug 2017.
  * Geifman and El-Yaniv [2019] Yonatan Geifman and Ran El-Yaniv.  SelectiveNet: A deep neural network with an integrated reject option.  In _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , 09–15 Jun 2019.
  * Gustafsson et al. [2020] Fredrik K Gustafsson, Martin Danelljan, and Thomas B Schön.  Evaluating scalable bayesian deep learning methods for robust computer vision.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops_ , 2020.
  * Hinton et al. [2014] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  In _Conference on Neural Information Processing Systems Deep Learning Workshop_ , 2014.
  * Huang et al. [2017] Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E. Hopcroft, and Kilian Q. Weinberger.  Snapshot ensembles: Train 1, get m for free.  In _International Conference on Learning Representations_ , 2017.
  * Hullermeier and Waegeman [2021] Eyke Hullermeier and Willem Waegeman.  Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods.  In _Machine Learning_ , 2021.
  * Jiao et al. [2019] Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu.  TinyBERT: Distilling BERT for natural language understanding.  In _International Conference on Learning Representations_ , 2019.
  * Kim et al. [2021] Jihyo Kim, Jiin Koo, and Sangheum Hwang.  A unified benchmark for the unknown detection capability of deep neural networks.  _ArXiv_ , abs/2112.00337, 2021.
  * Kim and Rush [2016] Yoon Kim and Alexander M. Rush.  Sequence-level knowledge distillation.  In _Conference on Empirical Methods in Natural Language Processing_ , 2016.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _International Conference for Learning Representations_ , 2015\.
  * Lakshminarayanan et al. [2017] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _Conference on Neural Information Processing Systems_ , 2017.
  * Lee et al. [2018] Kimin Lee, Honglak Lee, Kibok Lee, and Jinwoo Shin.  Training confidence-calibrated classifiers for detecting out-of-distribution samples.  In _International Conference on Learning Representations_ , 2018.
  * Lindqvist et al. [2020] Jakob Lindqvist, Amanda Olmin, Fredrik Lindsten, and Lennart Svensson.  A general framework for ensemble distribution distillation.  In _International Workshop on Machine Learning for Signal Processing (MLSP)_ , 2020.
  * Liu et al. [2020] Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan.  Simple and principled uncertainty estimation with deterministic deep learning via distance awareness.  In _Advances in Neural Information Processing Systems_ , volume 33, 2020.
  * Malinin and Gales [2021] Andrey Malinin and Mark J. F. Gales.  Uncertainty estimation in autoregressive structured prediction.  In _International Conference on Learning Representations_ , 2021.
  * Malinin et al. [2017] Andrey Malinin, Anton Ragni, Kate Knill, and Mark J. F. Gales.  Incorporating uncertainty into deep learning for spoken language assessment.  In _Association for Computational Linguistics_ , 2017.
  * Malinin et al. [2020] Andrey Malinin, Bruno Mlodozeniec, and Mark J. F. Gales.  Ensemble distribution distillation.  In _International Conference on Learning Representations_ , 2020.
  * Manning and Schütze [1999] Chris Manning and Hinrich Schütze.  _Foundations of Statistical Natural Language Processing_.  MIT Press, 1999.
  * Michel and Neubig [2018] Paul Michel and Graham Neubig.  Mtnt: A testbed for machine translation of noisy text.  In _Conference on Empirical Methods in Natural Language Processing_ , 2018.
  * Minka et al. [2005] Tom Minka et al.  Divergence measures and message passing.  Technical report, Technical report, Microsoft Research, 2005.
  * Moon et al. [2020] Jooyoung Moon, Jihyo Kim, Younghak Shin, and Sangheum Hwang.  Confidence-aware learning for deep neural networks.  In _International Conference on Machine Learning_ , 2020.
  * Neubig [2011] Graham Neubig.  The Kyoto free translation task.  http://www.phontron.com/kftt, 2011.
  * Notin et al. [2021] Pascal Notin, José Miguel Hernández-Lobato, and Yarin Gal.  Improving black-box optimization in VAE latent space using decoder uncertainty.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Ott et al. [2018] Myle Ott, Sergey Edunov, David Grangier, and Michael Auli.  Scaling neural machine translation.  In _Proceedings of the Third Conference on Machine Translation: Research Papers_ , 2018.
  * Ott et al. [2019] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli.  fairseq: A fast, extensible toolkit for sequence modeling.  In _Proceedings of NAACL-HLT 2019: Demonstrations_ , 2019.
  * Ovadia et al. [2019] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek.  Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift.  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * Post [2018] Matt Post.  A call for clarity in reporting BLEU scores.  In _Association for Computational Linguistics_ , 2018.
  * Radmard et al. [2021] Puria Radmard, Yassir Fathullah, and Aldo Lipani.  Subsequence based deep active learning for named entity recognition.  In _Association for Computational Linguistics_ , 2021.
  * Ryabinin et al. [2021] Max Ryabinin, Andrey Malinin, and Mark J. F. Gales.  Scaling ensemble distribution distillation to many classes with proxy targets.  In _Conference on Neural Information Processing Systems_ , 2021.
  * Sennrich et al. [2016] Rico Sennrich, Barry Haddow, and Alexandra Birch.  Neural machine translation of rare words with subword units.  In _Association for Computational Linguistics_ , 2016.
  * Smith [2017] Leslie N. Smith.  Cyclical learning rates for training neural networks.  In _Winter Conference on Applications of Computer Vision_ , 2017.
  * Tan et al. [2019] Xu Tan, Yi Ren, Di He, Tao Qin, and Tie-Yan Liu.  Multilingual neural machine translation with knowledge distillation.  In _International Conference on Learning Representations_ , 2019.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Conference on Neural Information Processing Systems_ , 2017.
  * Wong et al. [2016] Jeremy H. M. Wong, Mark J. F. Gales, and Yu Wang.  Sequence-level knowledge distillation.  In _IEEE/ACM Transactions on Audio Speech and Language Processing_ , 2016.
  * Xia and Bouganis [2022a] Guoxuan Xia and Christos-Savvas Bouganis.  On the usefulness of deep ensemble diversity for out-of-distribution detection.  _ArXiv_ , abs/2207.07517, 2022a.
  * Xia and Bouganis [2022b] Guoxuan Xia and Christos-Savvas Bouganis.  Augmenting softmax information for selective classification with out-of-distribution data.  In _Proceedings of the Asian Conference on Computer Vision (ACCV)_ , December 2022b.
  * Xie et al. [2013] Jingjing Xie, Bing Xu, and Zhang Chuang.  Horizontal and vertical ensemble with deep representation for classification.  In _International Conference on Machine Learning_ , 2013.
  * Yang et al. [2021] Jingkang Yang, Kaiyang Zhou, Yixuan Li, and Ziwei Liu.  Generalized out-of-distribution detection: A survey.  _arXiv preprint arXiv:2110.11334_ , 2021.
  * Yang et al. [2022] Jingkang Yang, Pengyun Wang, Dejian Zou, Zitang Zhou, Kunyuan Ding, Wenxuan Peng, Haoqi Wang, Guangyao Chen, Bo Li, Yiyou Sun, Xuefeng Du, Kaiyang Zhou, Wayne Zhang, Dan Hendrycks, Yixuan Li, and Ziwei Liu.  OpenOOD: Benchmarking generalized out-of-distribution detection.  In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_ , 2022.

##  Appendix A Experimental Configuration

This section will provide detailed information about the datasets used for
training, development, evaluation and detection. It will also give the exact
training and various hyperparameters used for all models.

###  A.1 Datasets

We utilise two training sets WMT16/20, each with a pair of development and
evaluation datasets based on newstest13/14 and newstest19/20. Additionally, we
utilise three out-of-domain datasets for evaluating detection performance of a
wide range of transformer models, see Table 6. As stated previously, all data
is cleaned and tokenized using Moses111github.com/moses-smt/mosesdecoder. For
WMT16, a shared dictionary is learned using BPE with 32,000 merge operations.
On WMT20 we learn disjoint dictionaries using BPE with 40,000 merge
operations. A consequence of the larger disjoint dictionary on WMT20 is the
significantly lower number of unknown tokens in the OOD datasets.

Table 6: Dataset information together with average source and target sentence
sizes post tokenization and processing. The OOD testsets Khresmoi, MTNT and
KFTT have two quoted numbers for each field as they were processed using
either the En-De WMT16 or En-Ru WMT20 BPE based dictionaries. Additionally,
only source side information is provided for OOD sets as these are only used
for unsupervised uncertainty estimation.

Dataset Type Number of Tokens per Sentence Fraction of Unknown Sentences
Source Target Tokens in Source En-De WMT16 policy, news, web 4.5M 29.5 30.6
0.01% En-De newstest13 news 3.0K 26.0 28.0 0.00% En-De newstest14 3.0K 27.6
29.1 0.00% En-Ru WMT20 policy, news, web 58.4M 27.8 27.5 0.00% En-Ru
newstest19 news 2.0K 29.9 33.4 0.00% En-Ru newstest20 2.0K 30.9 32.5 0.00%
Khresmoi medical 1.0K 30.9/30.3 — 0.78%/0.00% MTNT noisy reddit 1.4K 21.1/21.3
— 0.45%/0.06% KFTT encyclopedia 1.2K 35.4/35.2 — 1.46%/0.01%

###  A.2 En-De WMT16 Training

We use the base transformer from [Vaswani et al., 2017] implemented in fairseq
[Ott et al., 2019] and train it using 4 NVIDIA© A100 with an update frequency
of 32. This is virtually equivalent to training on $4\times 32=128$ GPUs. A
per-gpu batch has a maximum of 3584 tokens. Models are optimized with Adam
[Kingma and Ba, 2015] using $\beta_{1}$ = 0.9, $\beta_{2}$ = 0.98, and
$\epsilon$= 1e-8. We use a similar learning rate schedule to Vaswani et al.
[2017], i.e., the learning rate increases linearly for 4000 warmup steps to a
learning rate dependent on $d_{\tt model}$ after which it is decayed
proportionally to the inverse square root of the number of steps:

$\eta=({\tt step}\cdot d_{\tt model})^{-0.5}\min\left(1,\frac{{\tt step}}{{\tt
warmup}}\right)^{1.5}$

We use label smoothing with 0.1 weight for the uniform prior distribution over
the vocabulary. The last 10 weight checkpoints were averaged. Training was
stopped after 31 epochs corresponding to approximately a total of 18 GPU-
hours. At inference, a beam of 4 with a length-penalty of 0.6 is used for all
models. The Deep Ensemble consists of 5 of such models.

KD of Deep Ensemble: Knowledge distilled models are first initialised by one
of the teacher members and then trained using the knowledge distillation loss
$\mathcal{L}_{\tt KD}$ provided in Section 2.2 with $\lambda=0.50$. The
student was trained with a warmup of 1026 steps (3 epochs), from
$\eta=4.0\times 10^{-4}$ to $\eta=7.0\times 10^{-4}$ after which it decays for
a total of 24 epochs. A temperature of $T=0.8$ was used in the KL-divergence
loss as this was found to be mildly beneficial. All other hyperparameters
match the standard case above.

Snapshot Ensemble: The Snapshot Ensemble was generated by first starting from
the last checkpoint of a standard trained transformer. At this point, a cyclic
triangular learning rate schedule [Smith, 2017] was employed oscillating
between the values of $\eta_{\tt min}=1.0\times 10^{-4}$ and $\eta_{\tt
max}=1.0\times 10^{-3}$ with a period of 3 epochs. Note that the maximum
learning rate in this cyclic phase is notably larger than the peak learning
rate ($7.0\times 10^{-4}$) during standard training This setting was run for
15 epochs generating an ensemble with 5 members.

KD of Snapshot Ensemble: This system was trained using the same parameters as
the Deep Ensemble distilled students but was however, trained for only 12
epochs since it converged faster.

EDD & L-EDD: All of the EDD and L-EDD systems were distribution distilled from
the Snapshot Ensemble using the same setup as "KD of Snapshot Ensemble". We
chose $\beta=0.10$ by evaluating the translation performance of a range of
values $\beta\in\\{0.05,0.10,0.20,0.50\\}$ on the development newstest-13 set,
see Section 3.1.

###  A.3 En-Ru WMT20 Training

We use the big transformer from Vaswani et al. [2017] again implemented in
fairseq and trained using 4 NVIDIA© A100 with an update frequency of 32. A
per-gpu batch has a maximum of 5120 tokens. Dropout was set to a value of 0.10
and weight decay to 0.0001. In this case we train the model for 20 epochs,
corresponding to 53960 update steps and approximately 230 GPU-hours. The last
5 checkpoints were averaged leading to improved performance. At inference, a
beam of 5 with a length-penalty of 1.0 is used for all models.

Snapshot Ensemble: Based on the last checkpoint of a standard trained big
transformer, a triangular cyclic learning rate is utilised, oscillating
between $\eta=5.0\times 10^{-5}$ and $\eta=5.0\times 10^{-4}$ every 2 epochs
for 10 epochs. This results in an ensemble with 5 members.

KD of Snapshot Ensemble: Similar to the previous section, the distillation
student is initialised from its teacher but is trained using a learning rate
warmup of 2698 steps (one epoch) from $\eta=2.0\times 10^{-4}$ to
$\eta=4.0\times 10^{-4}$ after which it decays for a total of 12 epochs. The
last 3 or 5 epochs are averaged, based on development newstest19 performance.

L-EDD: Following distillation, L-EDD (Laplace) models are trained using the
same parameters. The best-found parameter $\beta=0.10$ in the WMT’16
experiments is to be used here. No hyperparameter search is performed at this
stage.