U NCERTAINTY E STIMATION AND C ALIBRATION WITH F INITE -S TATE P ROBABILISTIC RNN S
A BSTRACT
Uncertainty quantiﬁcation is crucial for building reliable and trustable machine learning systems.
We propose to estimate uncertainty in recurrent neural networks (RNNs) via stochastic discrete state transitions over recurrent timesteps.
The uncertainty of the model can be quantiﬁed by running a prediction several times, each time sampling from the recurrent state transition distribution, leading to potentially different results if the model is uncertain.
Alongside uncertainty quantiﬁcation, our proposed method offers several advantages in different settings.
The proposed method can (1) learn deterministic and probabilistic automata from data, (2) learn well-calibrated models on real-world classiﬁcation tasks, (3) improve the performance of out-of-distribution detection, and (4) control the explorationexploitation trade-off in reinforcement learning.
An implementation is available.
I NTRODUCTION
Machine learning models are well-calibrated if the probability associated with the predicted class reﬂects its correctness likelihood relative to the ground truth.
The output probabilities of modern neural networks are often poorly calibrated (Guo et al., 2017).
For instance, typical neural networks with a softmax activation tend to assign high probabilities to out-of-distribution samples (Gal &
Ghahramani, 2016b).
Providing uncertainty estimates is important for model interpretability as it allows users to assess the extent to which they can trust a given prediction (Jiang et al., 2018).
Moreover well-calibrated output probabilities are crucial in several use cases.
For instance, when monitoring medical time-series data (see Figure 1(a)), hospital staff should also be alerted when there is a low-conﬁdence prediction concerning a patient’s health status.
Bayesian neural networks (BNNs), which place a prior distribution on the model’s parameters, are a popular approach to modeling uncertainty.
BNNs often require more parameters, approximate inference, and depend crucially on the choice of prior (Gal, 2016; Lakshminarayanan et al., 2017).
Applying dropout both during training and inference can be interpreted as a BNN and provides a more efﬁcient method for uncertainty quantiﬁcation (Gal & Ghahramani, 2016b).
The dropout probability, however, needs to be tuned and, therefore, leads to a trade-off between predictive error and calibration error.
Sidestepping the challenges of Bayesian NNs, we propose an orthogonal approach to quantify the uncertainty in recurrent neural networks (RNNs).
At each time step, based on the current hidden (and cell) state, the model computes a probability distribution over a ﬁnite set of states.
The next state of the RNN is then drawn from this distribution.
We use the Gumbel softmax trick (Gumbel, 1954; Kendall & Gal, 2017; Jang et al., 2017) to perform Monte-Carlo gradient estimation.
Inspired by the effectiveness of temperature scaling (Guo et al., 2017) which is usually applied to trained models, we learn the temperature τ of the Gumbel softmax distribution during training to control the concentration of the state transition distribution.
Learning τ as a parameter can be seen as entropy regularization (Szegedy et al., 2016;
Pereyra et al., 2017; Jang et al., 2017).
The resulting model, which we name ST- τ , deﬁnes for every input sequence a probability distribution over state-
[Figure]
w h a t a
[Figure]
c r i p t
[Figure]
w h a t a [Figure] t o r y
[Figure]
w h a t [Figure] e
[Figure]
transition paths similar to a probabilistic state machine.
To estimate the model’s uncertainty for a prediction, ST- τ is run multiple times to compute mean and variance of the prediction probabilities.
We explore the behavior of ST- τ in a variety of tasks and settings.
First, we show that ST- τ can learn deterministic and probabilistic automata from data.
Second, we demonstrate on real-world classiﬁcation tasks that ST- τ learns well calibrated models.
Third, we show that ST- τ is competitive in out-of-distribution detection tasks.
Fourth, in a reinforcement learning task, we ﬁnd that ST- τ is able to trade off exploration and exploitation behavior better than existing methods.
Especially the outof-distribution detection and reinforcement learning tasks are not amenable to post-hoc calibration approaches
(Guo et al., 2017) and, therefore, require a method such as ours that is able to calibrate the probabilities during training.
U NCERTAINTY IN R ECURRENT N EURAL N ETWORKS 2. B ACKGROUND
An RNN is a function f deﬁned through a neural network with parameters w that is applied over time steps: at time step t , it reuses the hidden state h t − 1 of the previous time step and the current input x t to compute a new state h t , f : ( h t − 1 , x t ) → h t .
Some RNN variants such as LSTMs have memory cells
c t and apply the function f : ( h t − 1 , c t − 1 , x t ) → h t at each step.
A vanilla RNN maps two identical input sequences to the same state and it is therefore not possible to measure uncertainty of a prediction by running inference multiple times.
Furthermore, it is known that passing h t through a softmax transformation leads to overconﬁdent predictions on out-of-distribution samples and poorly calibrated probabilities (Guo et al., 2017).
In a Bayesian RNN the weight matrices w are drawn from a distribution and, therefore, the output is an average of an inﬁnite number of models.
Unlike vanilla RNNs, Bayesian RNNs are stochastic and it is possible to compute average and variance for a prediction.
Using a prior to integrate out the parameters during training also leads to a regularization effect.
However, there are two major and often debated challenges of BNNs: the right choice of prior and the efﬁcient approximation of the posterior.
With this paper, we side-step these challenges and model the uncertainty of an RNN through probabilistic state transitions between a ﬁnite number of k learnable states s 1 , ..., s k .
Given a state h t , we compute a probability distribution over the learnable states.
Hence, for the same state and input, the
RNN might move to different states in different runs.
Instead of integrating over possible weights, as in the case of BNNs, we sum over all possible state sequences and weigh the classiﬁcation probabilities by the probabilities of these sequences.
Figure 2 illustrates the proposed approach and contrasts it with vanilla and Bayesian RNNs.
The proposed method combines two building blocks.
The ﬁrst is state-regularization (Wang & Niepert, 2019) as a way to compute a probability distribution over a ﬁnite set of states in an RNN.
State-regularization, however, is deterministic and therefore we utilize the second building block, the Gumbel softmax trick (Gumbel, 1954;
Maddison et al., 2017; Jang et al., 2017) to sample from a categorical distribution.
Combining the two blocks allows us to create b a b 0.
accept 0.
0.2 b a b b a b 0.
w ~ reject accept reject accept reject accept reject accept reject
Deterministic State-Regularized RNNs.
State regularization (Wang & Niepert, 2019) extends
RNNs by dividing the computation of the hidden state h t into two components.
The ﬁrst component is an intermediate vector u t ∈ R d computed in the same manner as the standard recurrent component, u t = f
( h t − 1 , c t − 1 , x t ) .
The second component models probabilistic state transitions between a ﬁnite set of k learnable states s 1 , ..., s k , where s i ∈ R d , i ∈ [1 , k ] and which can also be written as a matrix
S t ∈ R
d
× k .
S t is randomly initialized and learnt during backpropagation like any other network weight.
At time step t , given an u t , the transition over next possible states is computed by: θ t = ϕ
( S t , u t ) , where θ t = { θ t, 1 , ..., θ t,k } and ϕ is some pre-deﬁned function.
In Wang & Niepert (2019), ϕ was a matrix-vector product followed by a S OFTMAX function that ensures (cid:80) ki
=1 θ t,i = 1 .
The hidden state h t is then computed by
[Equation]
where g ( · ) is another function, e.g. to compute the average.
Because Equation (1) is deterministic, it cannot capture and estimate epistemic uncertainty.
Monte-Carlo Estimator with Gumbel Trick.
The Gumbel softmax trick is an instance of a pathwise Monte-Carlo gradient estimator (Gumbel, 1954; Maddison et al., 2017; Jang et al., 2017).
With the Gumbel trick, it is possible to draw samples z from a categorical distribution given by paramaters θ , that is, z = ONE HOT (cid:0) arg max i [ γ i + log θ i ] (cid:1) , i ∈ [1 . . .
k ] , where k is the number of categories and γ i are i.i.d.
samples from the G UMBEL (0 , 1) , that is, γ = − log( − log( u )) , u ∼
UNIFORM (0 , 1) .
Because the arg max operator breaks end-to-end differentiability, the categorical distribution z can be approximated using the differentiable softmax function (Jang et al., 2017;
Maddison et al., 2017).
This enables us to draw a k -dimensional sample vector α ∈ ∆ k − 1 , where ∆ k − 1 is the ( k − 1)
-dimensional probability simplex.
2. S TOCHASTIC F INITE -S TATE RNN S (ST- τ )
Our goal is to make state transitions stochastic and uitilize them to measure uncertainty: given an input sequence, the uncertainty is modeled via the probability distribution over possible statetransition paths (see right half of Figure 2).
We can achieve this by setting ϕ to be a matrix-vector product and using θ t to sample from a Gumbel softmax distribution with temperature parameter τ .
Applying Monte Carlo estimation, at each time step t , we sample a distribution over state transition probabilities α t from the Gumbel softmax distribution with current parameter τ , where each state transition has the probability
[Equation]
The resulting α t = { α t, 1 , ..., α t,k } can be seen as a probability distribution that judges how important each learnable state s t is.
The new hidden state h t can now be formed either as
{ } x t c th tu t θ t h t c t 1×d d×k k×d =
GUMBEL-SOFTMAX 1×d u t
T h t T S t
S t
S t x t+1 f φ S t f g an average, h t = (cid:80) ki
=1 α t,i s i (the “soft” Gumbel estimator), or as a one-hot vector, h t =
ONE HOT (cid:0) arg max i [log( θ t,i ) + γ i ] (cid:1) .
For the latter, gradients can be estimated using the straightthrough estimator.
Empirically, we found the average to work better.
By sampling from the Gumbel softmax distribution at each time step, the model is stochastic and it is possible to measure variance across predictions and, therefore, to estimate epistemic uncertainty.
For more theoretical details we refer the reader to Appendix 2.3.
The parameter τ of the Gumbel softmax distribution is learned during training (Jang et al., 2017).
This allows us to directly adapt probabilistic RNNs to the inherent uncertainty of the data.
Intuitively, the parameter τ inﬂuences the concentration of the categorical distribution, that is, the larger τ the more uniform the distribution.
Since we inﬂuence the state transition uncertainty with the learned temperature τ , we refer to our model as ST- τ .
We provide an ablation experiment of learning τ versus keeping it ﬁxed in Appendix E.
Figure 3 illustrates the proposed model.
Given the previous hidden state of an RNN, ﬁrst an intermediate hidden state u t is computed using a standard RNN cell.
Next, the intermediate representation u t is multiplied with k learnable states arranged as a matrix S t , resulting in θ t .
Based on θ t , samples are drawn from a Gumbel softmax distribution with learnable temperature parameter τ .
The sampled probability distribution represents the certainty the model has in moving to the other states.
Running the model on the same input several times (drawing Monte-Carlo samples) allows us to estimate the uncertainty of the ST- τ model.
2. A LEATORIC AND E PISTEMIC U NCERTAINTY
Let Y be a set of class labels and D be a set of training samples.
For a classiﬁcation problem and a given ST- τ model with states { s 1 , ..., s k } , we can write for every y ∈ Y
[Equation]
where h i ∈ { s 1 , ..., s k } and the sum is over all possible paths (state sequences)
Ψ of length n .
Moreover, p ( (cid:104) h 1 , ..., h n (cid:105) | x ) is the probability of path ψ
= (cid:104) h 1 , ..., h n (cid:105) ∈
Ψ given input sequence x and q ( y | h n ) is the probability of class y given that we are in state h n .
Instead of integrating over possible weights, as in the case of BNNs, with ST- τ we integrate (sum) over all possible paths and weigh the class probabilities by the path probabilities.
The above model implicitly deﬁnes a probabilistic ensemble of several deterministic models, each represented by a particular path.
As mentioned in a recent paper about aleatoric and epistemic uncertainty in ML (H¨ullermeier & Waegeman, 2019): “the variance of the predictions produced by an ensemble is a good indicator of the (epistemic) uncertainty in a prediction.”
Let us now make this intuition more concrete using recently proposed measures of aleatoric and epistemic uncertainty (Depeweg et al., 2018); further discussed in (H¨ullermeier & Waegeman, 2019).
In Equation (19) of (H¨ullermeier & Waegeman, 2019) the total uncertainty is deﬁned as the entropy of the predictive posterior distribution
[Equation]
The above term includes both aleatoric and epistemic uncertainty (Depeweg et al., 2018;
H¨ullermeier & Waegeman, 2019).
Now, in the context of Bayesian NNs, where we have a distribution over the weights of a neural network, the expectation of the entropies wrt said distribution q(y 1 | h 1 ) = 0.
q(y 2 | h 1 ) = 0. q(y 1 | h 2 ) = 0.
q(y 2 | h 2 ) = 0.
q(y 1 | h 1 ) = 0.
q(y 2 | h 1 ) = 0. q(y 1 | h 2 ) = 0.
q(y 2 | h 2 ) = 0. is the aleatoric uncertainty (Equation 20 in H¨ullermeier & Waegeman (2019)):
[Equation]
Fixing the parameter weights to particular values eliminates the epistemic uncertainty.
Finally, the epistemic uncertainty is obtained as the difference of the total and aleatoric uncertainty (Equation 21 in H¨ullermeier & Waegeman (2019)):
[Equation]
Now, let us return to ﬁnite-state probabilistic RNNs.
Here, the aleatoric uncertainty is the expectation of the entropies with respect to the distribution over the possible paths Ψ :
[Equation]
where p ( y | ψ, x ) is the probability of class y conditioned on x and a particular path ψ .
The epistemic uncertainty for probabilistic ﬁnite-state RNNs can then be computed by
[Equation]
Probabilistic ﬁnite-state RNNs capture epistemic uncertainty when the equation above is non-zero.
As an example let us take a look at the two ST- τ models given in Figure 4.
Here, we have for input x two class labels y 1 and y 2 , three states ( s i , i ∈ { 0 , 1 , 2 } ), and two paths.
In both cases, we have that H [ p ( y | x )]
= − (0 .
3 log 2 0 .
3 + 0 .
7 log 2 0 .
≈ 0 .
Looking at the term for the aleatoric uncertainty, for the ST- τ depicted on the left side we have E p ( ψ | x ) H
[ p ( y | ψ, x )]
= − (0 .
3 log 2 0 .
3 + 0 .
7 log 2 0 .
≈ 0 .
In contrast, for the ST- τ depicted on the right side we have
E p ( ψ | x ) H [ p ( y | ψ, x )] ≈ 0 .
Consequently, the left side ST- τ has an epistemic uncertainty of u e ( x ) =
but the right side ST- τ exhibits an epistemic uncertainty of u e ( x ) = 0 .
This example illustrates three main observations.
First, we can represent epistemic uncertainty through distributions over possible paths.
Second, the more spiky the transition distributions, the more deterministic the behavior of the ST- τ model, and the more conﬁdent it becomes with its prediction by shrinking the reducible source of uncertainty (epistemic uncertainty).
Third, both models are equally calibrated as their predictive probabilities are, in expectation, identical for the same inputs.
Hence, ST- τ is not merely calibrating predictive probabilities but also captures epistemic uncertainty.
Finally, we want to stress the connection between the parameter τ (the temperature) and the degree of (epistemic) uncertainty of the model.
For small τ the ST- τ model behavior is more deterministic and, therefore, has a lower degree of epistemic uncertainty.
For instance, the ST- τ on the left in Figure 4 has 0 epistemic uncertainty because all transition probabilities are deterministic.
Empirically, we ﬁnd that the temperature and, therefore, the epistemic uncertainty often reduces during training, leaving the irreducible uncertainty (aleatoric) to be the main source of uncertainty.
R ELATED W ORK
Uncertainty .
Uncertainty quantiﬁcation for safety-critical applications (Krzywinski & Altman, 2013) has been explored for deep neural nets in the context of Bayesian learning (Blundell et al., 2015; Gal & Ghahramani, 2016b; Kendall & Gal, 2017; Kwon et al., 2018).
Bayes by Backprop (BBB) (Blundell et al., 2015) is a variational inference scheme for learning a distribution over weights w in neural networks and assumes that the weights are distributed normally, that is, w i ∼ N ( µ, σ 2 ) .
The principles of bayesian neural networks (BNNs) have been applied to RNNs and shown to result in superior performance compared to vanilla RNNs in natural language processing (NLP) tasks (Fortunato et al., 2017).
However, BNNs come with a high computational cost because we need to learn µ and σ for each weight in the network, effectively doubling the number of parameters.
Furthermore, the prior might not be optimal and approximate inference could lead inaccurate estimates (Kuleshov et al., 2018).
Dropout (Hinton et al., 2012; Srivastava et al., 2014) can be seen as a variational approximation of a Gaussian Process (Gal & Ghahramani, 2016b;a).
By leaving dropout activated at prediction time, it can be used to measure uncertainty.
However, the dropout probability needs to be tuned which leads to a trade-off between predictive error and calibration error (see Figure 10 in the Appendix for an empirical example).
Deep ensembles (Lakshminarayanan et al., 2017) offer a non-Bayesian approach to measure uncertainty by training multiple separate networks and ensembling them.
Similar to BNNs, however, deep ensembles require more resources as several different RNNs need to be trained.
We show that ST- τ is competitive to deep ensembles without the resource overhead.
Recent work (Hwang et al., 2020) describes a sample-free uncertainty estimation for Gated Recurrent Units (SP-GRU) (Chung et al., 2014), which estimates uncertainty by performing forward propagation in a series of deterministic linear and nonlinear transformations with exponential family distributions.
ST- τ estimates uncertainties through the stochastic transitions between two consecutive recurrent states.
Calibration .
Platt scaling (Platt et al., 1999) is a calibration method for binary classiﬁcation settings and has been extend to multi-class problems (Zadrozny & Elkan, 2002) and the structured prediction settings (Kuleshov & Liang, 2015).
(Guo et al., 2017) extended the method to calibrate modern deep neural networks, particularly networks with a large number of layers.
In their setup, a temperature parameter for the ﬁnal softmax layer is adjusted only after training.
In contrast, our method learns the temperature and, therefore, the two processes are not decoupled.
In some tasks such as timeseries prediction or RL it is crucial to calibrate during training and not post-hoc.
Deterministic & Probabilistic Automata Extraction .
Deterministic Finite Automata (DFA) have been used to make the behavior of RNNs more transparent.
DFAs can be extracted from RNNs after an RNN is trained by applying clustering algorithms like k -means to the extracted hidden states (Wang et al., 2018) or by applying the exact learning algorithm L ∗ (Weiss et al., 2018).
Posthoc extraction, however, might not recover faithful DFAs.
Instead, (Wang & Niepert, 2019) proposed state-regularized RNNs where the ﬁnite set of states is learned alongside the RNN by using probabilistic state transitions.
Building on this, we use the Gumbel softmax trick to model stochastic state transitions, allowing us to learn probabilistic automata (PAs)
(Rabin, 1963) from data.
Hidden Markov Models (HMMs) & State-Space Models (SSMs) & RNNs .
HMMs are transducer-style probabilistic automata, simpler and more transparent models than RNNs.
(Bridle, 1990) explored how a HMM can be interpreted as an RNNs by using full likelihood scoring for each word model.
(Krakovna & Doshi-Velez, 2016) studied various combinations of HMMs and RNNs to increase the interpretability of RNNs.
There have also been ideas on incorporating RNNs to HMMs to capture complex dynamics (Dai et al., 2017; Doerr et al., 2018).
Another relative group of work is SSMs, e.g., rSLDS (Linderman et al., 2017), Kalman VAE (Fraccaro et al., 2017), PlaNet (Hafner et al., 2019) and RKN (Becker et al., 2019).
They can be extended and viewed as another way to inject stochasticity to RNN-based architectures.
In contrast, ST- τ models stochastic ﬁnite-state transition mechanisms end-to-end in conjunction with modern gradient estimators to directly quantify and calibrate uncertainty and the underlying probabilistic system.
This enables ST- τ to approximate and extract the probabilistic dynamics in RNNs.
E XPERIMENTS
The experiments are grouped into ﬁve categories.
First, we show that it is possible to use ST- τ to learn deterministic and probabilistic automata from language data (Sec. 4.1).
This demonstrates that ST- τ can capture and recover the stochastic behavior of both deterministic and stochastic languages.
Second, we demonstrate on classiﬁcation tasks (Sec. 4.2) that ST- τ performs better than or similar to existing baselines both in terms of predictive quality and model calibration.
Third, we compare ST- τ with existing baselines using out-of-distribution detection tasks (Sec. 4.3).
Fourth, we conduct reinforcement learning experiments where we show that the learned parameter τ can
A
B 0 [0.3] 1 [0.7] 0 [1.0] 1 [1.0] (a) True PA0 [0.7,0.07] 1 [0.29,0.06] 2 0 [0.3,0.07] 1 [0.71,0.06]
[1.0,0.0]1 [1.0,0.0] (b) Extracted PAA
B 0 [0.7] 1 [0.3]
C 0 [0.3] 1 [0.7] 0 [1.0]1 [1.0] 0 [1.0]1 [1.0] (c) True PA1 [0.17] 0 [0.39] 2 1 [0.13]
[1.0] 1 [1.0] (d) Extracted PA- calibrate the exploration-exploitation trade-off during learning (Appendix D), leading to a lower sample complexity.
Fifth, we a regression task ( C).
Models.
We compare the proposed ST- τ method to four existing models.
First, a vanilla LSTM ( LSTM ).
Second, a Bayesian RNN ( BBB ) (Blundell et al., 2015; Fortunato et al., 2017), where each network weight w i is sampled from a Gaussian distribution N ( µ, σ 2 ) , w i = µ + σ · (cid:15) , with σ = log(1 + exp( ρ )) and (cid:15) ∼ N (0 , 1) .
To estimate model uncertainty, we keep the same sampling strategy as employed during training (rather than using w i = µ ).
Third, a RNN that employs
Variational Dropout ( VD ) (Gal & Ghahramani, 2016a).
In variational dropout, the same dropout mask is applied at all time steps for one input.
To use this method to measure uncertainty, we keep the same dropout probability at prediction time (Gal & Ghahramani, 2016a).
Fourth, a deep ensemble of a LSTM base model (Lakshminarayanan et al., 2017).
For ST- τ we compute the new hidden state using the soft version of the Gumbel softmax estimator (Jang et al., 2017).
All models contain only a single LSTM layer, are implemented in Tensorﬂow (Abadi et al., 2015), and use the
A DAM (Kingma & Ba, 2015) optimizer with initial learning rate 0 .
4. D ETERMINISTIC & P ROBABILISTIC A UTOMATA E XTRACTION
We aim to investigate the extent to which ST- τ can learn deterministic and probabilistic automata from sequences generated by regular and stochastic languages.
Since the underlying languages are known and the data is generated using these languages, we can exactly assess whether ST- τ can recover these languages.
We refer the reader to the appendix A for a deﬁnition of deterministic ﬁnite automata (DFA) and probabilistic automata (PA).
The set of languages recognized by DFAs and PAs are referred to as, respectively, regular and stochastic languages.
For the extraction experiments we use the GRU cell as it does not have a cell state.
This allows us to read the Markovian transition probabilities for each state-input symbol pair directly from the trained ST- τ .
Regular Languages .
We conduct experiments on the regular language deﬁned by Tomita grammar 3 .
This language consists of any string without an odd number of consecutive 0’s after an odd number of consecutive l’s (Tomita, 1982).
Initializing τ
= 1 , we train ST- τ (with k = 10 states) to learn a ST- τ model to represent this grammar and then extract DFAs (see Appendix A.4 for the extraction algorithm).
In principle, k ≥ # of classes, the model learns to select a ﬁnite set of (meaningful) states to represent a language, as shown in Appendix Figure 8.
ST- τ is able to learn that the underlying language is deterministic, learning the temperature parameter τ accordingly and the extraction produces the correct underlying DFA.
The details are discussed in the appendix.
Stochastic Languages We explore whether it is possible to recover a PA from a ST- τ model trained on the data generated by a given PA.
While probabilistic deterministic ﬁnite automata (PDFAs) have been extracted previously (Weiss et al., 2019), to the best of our knowledge, this is the ﬁrst work to directly learn a PA, which is more expressive than a PDFA (Denis & Esposito, 2004), by extraction from a trained RNN.
We generate training data for two stochastic languages deﬁned by the PAs shown in Figure 5 (a) and (c).
Using this data, we train a ST- τ with a GRU with k = 4 states
and we directly use the Gumbel softmax distribution to approximate the probability transitions of the underlying PA (see Appendix A.5 for more details and the extraction algorithm).
Figures 5 (b, d) depict the extracted PAs.
For both stochastic languages the extracted PAs indicate that ST- τ is able to learn the probabilistic dynamics of the ground-truth PAs.
[Table]
(lower is better for all metrics).
ST offers the best and reliable trade-off predictive and calibration errors.
Furthermore, it does not require more parameters as BBB (double) or Deep
Ensemble (order of magnitude more) nor does a hyperparameter has to be tuned as in VD
Stochastic predictions are averaged across 10 independent runs and their variance is reported Best and second best results are marked in bold and underlined (PE bold models are signiﬁcantly different at level p ≤ 0 . 005 ). An ablation experiment with post-hoc temperature scaling is in Appendix B [Figure]
M ODEL C ALIBRATION
We evaluate ST- τ ’s prediction and calibration quality on two classiﬁcation tasks.
The ﬁrst task is heartbeat classiﬁcation with 5 classes where we use the MIT-BIH arrhythmia dataset (Goldberger et al., 2000; Moody & Mark, 2001).
It consists of 48 half-hour excerpts of electrocardiogram (ECG) recordings.
To preprocess the recordings we follow Kachuee et al.
(Part III.A).
The second task is sentiment analysis where natural language text is given as input and the problem is binary sentiment classiﬁcation.
We use the IMDB dataset (Maas et al., 2011), which contains reviews of movies that are to be classiﬁed as being positive or negative.
For further dataset and hyperparameter details, please see Appendix B.2.
We compute the output of each model 10 times and report mean and variance.
For the deep ensemble, we train 10 distinct LSTM models.
For VD, we tuned the best dropout rate in { 0 . 05 , 0 . 1 , 0 . 15 } and for BBB we tuned µ = { 0 , 0 . 01 } and ρ = {− 1 , − 2 , − 3 , − 4 } , choosing the best setup by lowest predictive error achieved on validation data.
For ST- τ , we evaluate both, setting the number of states to the number of output classes ( k = 5 / 2 , BIH and IMDB, respectively ) and to a ﬁxed value k = 10 .
We initialize with τ = 1 and use a dense layer before the softmax classiﬁcation.
For more details see Appendix, Table 4.
With a perfectly calibrated model, the probability of the output equals the conﬁdence of the model.
However, many neural networks do not behave this way
Guo et al.
To assess the extent to which a model is calibrated, we use reliability diagrams as well as Expected Calibration Error (ECE) and Maximum Calibration Error (MCE).
Please see
Appendix B.1 for details.
For VD, the best dropout probability on the validation set is 0.05.
Lower is better for all metrics.
For PE, all models are marked bold if there is no signiﬁcant difference at level p ≤ 0 .
005 to the best model.
Results .
The results are summarized in Table 1.
For the BIH dataset, the vanilla LSTM achieves the smallest PE with a signiﬁcant difference to all other models at p ≤ 0 .
005 using an approximate randomization test (Noreen, 1989).
It cannot, however, measure uncertainty and suffers from higher
ECE and MCE.
Similarly, VD exhibits a large MCE.
The situation is reversed for BBB were we ﬁnd a vastly higher PE, but lower MCE.
In contrast, ST- τ achieves overall good results: PE is only
[Table]
has a slightly higher PE than the best models, but has the lowest ECE and MCE offering a good trade-off.
The calibration plots of IMDB can be found in the Appendix, Figure 9.
In addition to achieving consistently competitive results across all metrics,
ST- τ has further advantages compared to the other methods.
The deep ensemble doubles the number of parameters by the number of model copies used.
BBB requires the doubling of parameters and a carefully chosen prior, where ST- τ does only require a slight increase in number of parameters compared to a vanilla LSTM.
VD requires the tuning of the hyperparameter for the dropout probability, which leads to a trade-off between predictive and calibration errors (see Appendix B.3.2).
4. O UT -O F -D ISTRIBUTION D ETECTION
We explore the ability of ST- τ to estimate uncertainty by making it detect out-of-distribution (OOD) samples following prior work (Hendrycks & Gimpel, 2017).
The in-distribution dataset is IMDB
and we use two OOD datasets: the Customer Review test dataset (Hu & Liu, 2004) and the Movie review test dataset (Pang et al., 2002), which consist of, respectively, 500 and 1,000 samples.
As in
Hendrycks & Gimpel (2017), we use the evaluation metrics AUROC and AUPR.
Additionally, we report the accuracy on in-distribution samples.
For VD we select the dropout rate from the values { 0 . 05 , 0 .
1 , 0 . 2 } and for BBB we select the best µ = { 0 , 0 . 01 } and ρ = {− 1 , − 2 , − 3 , − 4 } , based on best AUROC and AUPR.
For ST- τ we used c = 10 and c = 100 .
Beside using the maximum probability of the softmax (MP) as baseline (Hendrycks & Gimpel, 2017), we also consider the variance of the maximum probability (Var-MP) across 10 runs.
The number of in-domain samples is set to be the same as the number of out-of-domain samples from IMDB (Maas et al., 2011).
Hence, a random baseline should achieve 50% AUROC and AUPR.
Results .
Table 2 lists the results.
ST- τ and deep ensembles are the best methods in terms of OOD detection and achieve better results for both MP and Var-MP.
The MP results for ST- τ are among the best showing that the proposed method is able to estimate out-of-distribution uncertainty.
We consider these results encouraging especially considering that we only tuned the number of learnable ﬁnite states c in ST- τ .
Interestingly, a larger number of states improves the variance-based out-ofdistribution detection of ST- τ .
In summary, ST- τ is highly competitive in the OOD task.
D ISCUSSION AND C ONCLUSION
We proposed ST- τ , a novel method to model uncertainty in recurrent neural networks.
ST- τ achieves competitive results relative to other strong baselines (VD, BBB, Deep Ensembles), while circumventing some of their disadvantages, e.g., extensive hyperparameters tuning and doubled number of parameters.
ST- τ provides a novel mechanism to capture the uncertainty from (sequential) data over time steps.
The key characteristic which distinguishes ST- τ from baseline methods is its ability to model discrete and stochastic state transitions using modern gradient estimators at each time step.
A D ETERMINISTIC & P ROBABILISTIC A UTOMATA E XTRACTION
For the ST- τ models we use an embedding layer and an LSTM layer (both with 100 hidden units) and a dense layer which accepts the last hidden state output and has two output neurons (accept, reject).
The training objective aims to minimize a cross-entropy loss for a binary classiﬁcation problem (accept or reject).
A. D EFINITION OF D ETERMINISTIC F INITE A UTOMATA
A Deterministic Finite Automata (DFA) is a 5-tuple ( Q , Σ , δ, q 0 , F ) consisting of a ﬁnite set of states
Q ; a ﬁnite set of input tokens Σ (called the input alphabet); a transition function δ : Q × Σ → Q ; a start state q 0 ; and a set of accept states F ⊆ Q .
Given a DFA and an input, it is possible to follow how an accept or reject state is reached.
DFAs can be extracted from RNNs in order to offer insights into the workings of the RNN, making it more interpretable.
SR -RNNs (Wang & Niepert, 2019) extract DFAs from RNNs by counting the number of transitions that have occurred between a state and its subsequent states, given a certain input (Schellhammer et al., 1998).
However, this extraction method is deterministic and cannot give any uncertainty estimates for the extracted DFA.
By adding stochasticity using the Gumbel softmax distribution, we can additionally offer uncertainty measures for the state transitions.
A. E XTRACTING A UTOMATA FOR R EGULAR L ANGUAGES
With this experiment, we want to explore two features of ST- τ .
First, we want to understand how τ changes as training progresses (see Figure 7 (a)).
At the beginning of training, τ ﬁrst increases, allowing the model to explore the state transitions and select the states which will represent the corresponding grammar (in our example, the model selects 5 out of 10 states to represent Tomita 3, see
Figure 7 (b, c)).
Later, τ decreases and the transition uncertainty is calibrated to have tighter bounds, becoming more deterministic).
Second, we want to see if ST- τ can model the uncertainty in transitions and adaptively learn to calibrate the state transition distributions.
For this, we extract DFAs at two different iterations (see Figure 7 (b, c)).
After 50k iterations, a correct DFA can be extracted.
However, the transitions are not well calibrated.
The ideal transition should be deterministic and have transition probability close to 1.
For example, at state 7, for input “ 1 ”, only 53% of the time the model transitions to the correct state 9.
In contrast, after 250k iterations, the transitions are well calibrated and all transition are almost deterministic.
At the same time, τ has lowered, indicating that the model has become more certain.
A. D EFINITION OF P ROBABILISTIC A UTOMATA
The stochasticity in ST- τ also allows us to extract Probabilistic Automata (PA) (Rabin, 1963) which have a stronger representational power than DFAs.
A PA is deﬁned as a tuple ( Q , Σ , δ, q 0 , F ) consisting of a ﬁnite set of states Q ; a ﬁnite set of input tokens Σ ; a transition function δ : Q × Σ → P , where P is the transition probability for a particular state, and P ( Q ) denotes the power set of Q ; a start state q 0 and a set of accept states F ⊆ Q .
A. E XTRACTING DFA S WITH U NCERTAINTY I NFORMATION
Let p ( s j |
s i , x t ) , i, j ∈ { 1 , ..., k } , x t ∈ Σ be the probability of the transition to state s j , given current state s i and current input symbol x t at a given time step t .
We query each training sample to model and record the transition probability for each input symbol.
The DFA extraction algorithms in (Schellhammer et al., 1998; Wang et al., 2018) are based on the count of transitions.
In contrast, our extraction algorithm utilizes the transition probability, as described in Algorithm 1.
A. E XTRACTING P ROBABILISTIC A UTOMATA
We ﬁrst deﬁne two probabilistic automata as shown in Figure 5(a)(c) to generate samples from it 2 .
We generated 10,170 samples for stochastic language 1 (abbreviated SL-1) with PA-1 (Figure 5(a))
[Figure]
and sample length l ∈ [1 , 9] .
For SL-2 with PA-2, we generated 20,460 samples with sample length l ∈ [1 , 10] .
The learning procedure is described in Algorithm 2.
We use the Gumbel softmax distribution in ST- τ to approximate the probability distribution of next state p ( s j | s i , x t ) .
To achieve this goal, we set the number of states k to an even number and force the ﬁrst half of the state to “reject” state and the second half of states to be “accept”.
This allows us to ensure that the model models both “reject” and “accept” with the same number of states.
In the main part of the paper we report results when setting ST- τ to have k = 4 states for SL=2.
Here, we additionally present the results for k = 6 in Figure 8, which yields comparable results, showing that ST- τ is not overly sensitive to the number of states.
Is is however, helpful to have the same number of accept and reject states.
[Table]
Input: model M , dataset D , alphabet Σ , start token x
Output: transition function δ 1: Initialize an empty dictionary Z [ s j |
s i , x t ]
= Ø 2: Compute the probability distribution over k states when input x 0 : p 1: k = M
( x 0 ) , s j = arg max i ∈{ 1 ,...,k } [ p 1: k ] 3: Set s i =
s j , Update Z [ s j | x 0 ]
= { p j } 4: for x = ( x 1 , x 2 , ...,
x T ) ∈
D do 5: for t ∈ [1 , ..., T ] do 6: p 1: k =
M
( s i , x t ) , s j = arg max i ∈{ 1 ,...,k } [ p 1: k ] , 7:
Set s i =
s j , Update Z [ s j |
s i , x t ] : 8:
Z [ s j | s i , x t ]
= Z [ s j |
s i , x t ] ∪ { p j } , 9: end for 10: end for 11: Compute transition function δ , transition probability mean p u and variance p var : 12: for i, j ∈ { 1 , ..., k } and x t ∈ Σ do 13:
p u
= mean ( Z [( s j | s i , x t )])
14: p var = var ( Z [( s j | s i , x t )]) 15: δ ( s j | s i , x t ) =
arg max j ∈{ 1 ,...,k } p u 16: end for
Input: dataset D
Output: network loss L for x , y ∈ D do initialize h 0 with the 1 st state s 1 , s 1 ∈ S =
{ s 1: k } h
=
s 1 , ˆ y =for x t ∈ x , t ∈ [1 , ..., T ] do z t =
sigma ( W z x t
+ U z h t − 1 ) r t = sigma ( W [Equation] x
[Equation]
) g t = tahn ( W
[Equation]
U
[Equation]
) u t = [Equation] t (cid:12) h t [Equation]
[Equation]
z t ) (cid:12) g
[Equation]
p t =
S u (cid:62) t
ˆ p t ∼ G UMBEL ( p t , τ ) h t
= ˆ p t (cid:62) S ˆ y =
ˆ p t end for if y = 0 then y
=
[0 , 0 , .. 0 , (cid:124) (cid:123)(cid:122) (cid:125) k/1 , 1 , ..(cid:124) (cid:123)(cid:122) (cid:125) k/] else y =
[1 , 1 , .. 1 , (cid:124) (cid:123)(cid:122) (cid:125) k/0 , 0 , ..
(cid:124) (cid:123)(cid:122) (cid:125) k/] end if
Compute cross-entropy loss:
L = CROSS - ENTROPY ( y , ˆ y ) end for
[Figure]
[Figure] [Figure]
C 0 [0.3] 1 [0.7] 0 [1.0] 1 [1.0] 0 [1.0] 1 [1.0] [-(a) Truth PA1 [0.08,0.02] 
-]( [Figure] [Figure] [Figure]
PA
[-[1.0,0.0] 1 [1.0,0.0] (b) Extracted PA1 [0.08] 
-][Figure] [Figure] state A input “0” p ( accept | s
≈ 29 = p ( | p ( |
7 ≈ 71 =
B M ODEL C ALIBRATION
We address the problem of supervised multi-class sequence classiﬁcation with recurrent neural networks.
We follow the deﬁnitions and evaluation metrics as in (Guo et al., 2017).
Given input X ∈ X and ground-truth label Y ∈ Y , a probabilistic classiﬁer m ( X ) =
( ˆ Y , ˆ P ) .
The ˆ Y =
{ ˆ y 1 , ..., ˆ y k } , ˆ P = { ˆ p 1 , ..., ˆ p k } present the predicted class label and conﬁdence (the probability of correctness) over k classes and (cid:80) ki
=1 ˆ p
i
= 1 .
Deﬁnition of Calibration A model is perfectly calibrated if the conﬁdence estimation equals the true probability, that is, P ( ˆ Y
=
Y | ˆ P = p ) = p , p ∈ [0 , 1] .
B. E VALUATION OF C ALIBRATION
Reliability Diagrams (DeGroot & Fienberg, 1983; Niculescu-Mizil & Caruana, 2005) visualise whether a model is over- or under-conﬁdent by grouping predictions into bins according to their prediction probability.
The predictions are grouped into N interval bins (each of of size 1 /N ) and the accuracy of samples y i wrt.
to the ground truth label ˆ y i in each bin b n is computed as:
[Equation]
where i indexes all examples that fall into bin b n .
Let ˆ p
i be the probability for sample y i , then average conﬁdence is deﬁned as
[Equation]
A model is perfectly calibrated if acc ( b n ) = conf ( b n ) , ∀ n and in a diagram the bins would follow the identity function.
Any derivation from this represents miscalibration.
Based on the accuracy and conﬁdence measures, two calibration error metrics have been introduced (Naeini et al., 2015).
Expected Calibration Error (ECE) .
Besides the reliability diagrams, ECE is a convenient tool to have scalar summary statistic of calibration.
It computes the difference between model accuracy and conﬁdence as a weighted average across bins,
[Equation]
where m is the total number of samples.
Maximum Calibration Error (MCE) is particularly important in high-risk applications where reliable conﬁdence measures are absolutely necessary.
It measures the worst-case deviation between accuracy and conﬁdence,
[Equation]
For a perfectly calibrated classiﬁer, the ideal ECE and MCE both equal to 0.
[Table]
[Figure]
[Figure]
VD Prob
[Figure]
[Table]
B. D ATASET AND H YPERPARAMETER DETAILS
For the experiments of Section 4.2, we provide an overview of the used datasets in Table 3 and give details on the different hyperparameters used in the experiments in Table 4.
On BIH, we use the training / test split of (Kachuee et al., 2018), however we additionally split off 10% of the training
[Figure]
data to use as a validation set.
On IMDB, the original dataset consist of 25k training and 25k test samples, we split 2k from training set as validation dataset.
The word dictionary size is limited to 5,000.
B. A DDITIONAL E XPERIMENTS
Here we report three groups of additional or ablation experiments: (1) All baseline method and STτ with directly employing post-training temperature scaling.
The trade-off between predictive and calibrate performance with different dropout ratio in variational dropout (VD).
(3) Additional calibration plots for the IMDB dataset.
B.3. E XPERIMENTS WITH T EMPERATURE S CALING
For classiﬁcation calibration experiments, the post-hoc temperature scaling can also be used to calibrate models.
However, please note, post-hoc temperature scaling can not be used when we need to calibrate a model during training stages, for example, the tasks like DFA or PA extraction, reinforcement learning tasks.
Table 5 reports the results with temperature scaling where temperature is tuned on valid set.
Among the models (Ensemble, BBB, VD) that can estimate uncertainty, ST- τ is very competitive, for example, the second best on both the BIH and the IMDB datsaet.
While LSTM can achieve better predictive performance and sometimes better calibration performance, LSTM is not able to provide any uncertainty information.
B.3. E XPERIMENTS WITH D IFFERENT D ROPOUT R ATE
To gain a better understanding of how crucial the hyperparameter of VD is, we investigate the effect of the dropout probability during prediction with VD.
We perform an experiment where we vary the dropout probability in the range of [0 . 0 , 0 . 9] with increments of 0 . 1 .
In Figures 10 (a, b) we plot the result for BIH and IMDB, respectively, reporting PE, ECE and MCE for the various VD settings as well as ST- τ .
On both datasets, for VD, the point of lowest PE does not necessarily coincide with the points of lowest MCE and/or ECE.
For example on BIH, VD achieves the lowest PE when dropout is switched
[Figure]
off (0.0), but then uncertainty cannot be measured.
On the other hand, choosing the next lowest PE results in a high MCE.
In contrast, ST- τ directly achieves good results for all three metrics.
Similarly, on IMDB, at the point where VD has the lowest PE is also has highest MCE.
In conclusion, VD requires careful tuning which comes with choosing a trade-off between the different metrics [Figure] while
ST- τ achieves good results directly, without any tuning.
B.3. C ALIBRATION P LOTS FOR THE IMDB D ATASET
Figure 9 shows the calibration plots for IMDB.
For the binary classiﬁcation task BBB achieves the best calibration performance and VD achieves the best predictive performance.
It should be noted that ST- τ achieves the best trade-off between predictive and calibration performance without doubling parameters (for BBB) and without tuning dropout rate for (VD).
C R EGRESSION
Calibration plays an important role in probabilistic forecasting.
We consider a time-series forecasting regression task using the individual household electric power consumption dataset.
The goal of the task is to predict the global active power at the current time (t) given the measurement and other features of the previous time step.
The dataset was sampled at the time step of an hour (the original data are given in minutes), which leads to 34,588 samples.
We split it 25,000/2,000/7, for training/validation/test.
One LSTM layer with 100 hidden units is used for the baselines (LSTM,
BBB and VD) and ST- τ (the number of states is set to 10).
The evaluation metric is mean squared error (MSE) and we use the model with lowest MSE on the validation dataset at test stage.
is [Figure]
R e w a r d s
LSTMVD
BBBST
[Figure]
5 Episode(×10 3 )
R e w a r d s
LSTMVD
BBBST- (
[Figure]
MDP
[Figure]
greedy     5 Episode(×10 3 )
R e w a r d s
LSTMVD
BBBST-
[Figure]
[Figure]
[Figure]
R e w a r d s
LSTMVD
BBBST-
[Figure]
[Figure]
Figure 11 presents the performance.
LSTM, VD and ST- τ perform very well at this task, achieving
MSE close to zero.
BBB performs worse than the other methods.
For uncertainty estimation, BBB,
VD and ST- τ are able to provide uncertainty information alongside the predictive score.
BBB gives high uncertainty, while VD and ST- τ are more conﬁdent in their prediction, with ST- τ offering the tightest uncertainty bounds.
D R EINFORCEMENT L EARNING
We explore the ability of ST- τ to learn the exploration-exploitation trade-off in reinforcement learning.
To demonstrate this empirically, we evaluate ST- τ in the continuous control environment CARTPOLE (Figure 13) of OpenAI Gym (Brockman et al., 2016).
The objective of the task is to train an agent to balance a pole for 1 second (which equals 50 time steps) given environment observations
O .
To keep the pole in balance, the agent has to move the pole either left or right, that is, the possible actions are a = { LEFT , RIGHT } .
If the chosen action keeps the pole in balance, the agent receives a reward of 1 and continues; otherwise a reward of 0 is given and the run stops.
The environment can be formulated as a reinforcement learning problem ( S, A, P, R ) where S is a set of states.
Each state s
= { x, ˆ x, β, ˆ β } ∈ S consists of x : cart position,
ˆ x : cart velocity, β : angle position, and ˆ β : angle velocity.
A is the set of actions, P the state transition probability, and R the reward.
We consider two different setups.
In the ﬁrst setup, the CARTPOLE environment is fully observable Markov Decision Process (MDP setup) where the agent has full access to observation O , that is, O = { x, ˆ x, β, ˆ β } ∈ S .
In the second and more difﬁcult setup (POMDP), the environment is partially observable, where O = { x, β } ⊂ S .
For the latter, the agent cannot observe the state cart velocity and angular velocity information.
It has to learn to approximate them using its recurrent connections, and thus needs to retain long-term history (Bakker, 2002; Gomez & Schmidhuber, 2005; Sch¨afer, 2008).
The various RNN-based models are trained to output a distribution over actions at each time step t , that is, π ( a | s t , a t ) , where π is the set of policy.
For selecting the next action, we consider two policies: (1) sampling: a t +1 ∼ π ( a | s t , a t ) , and (2) greedy: a t +1 = arg max π ( a | s t , a t ) .
For all baselines, we use one LSTM layer and one dense layer with softmax to return the distribution over actions.
Each layer has 100 hidden units.
For VD we tuned the dropout rate { 0 . 05 , 0 .
1 , 0 . 2 } and for BBB µ = { 0 , 0 . 01 , 0 .
05 , 0 . 1 } and ρ = {− 1 , − 2 , − 3 , − 4 } .
For ST- τ we simply set the initial temperature value to τ = 1 , the number of possible states to the number of actions ( k = 2 ), and the next action is directly selected based on the Gumbel softmax distribution over states.
[Figure]
Learning the temperature τ is akin to entropy regularization (Szegedy et al., 2016;
Pereyra et al., 2017; Jang et al., 2017) and adjusts the conﬁdence (epistemic uncertainty) during training.
Results are presented in Figure 12.
An important criteria for evaluating RL agents is the sample complexity (Malik et al., 2019), that is, the amount of interactions between agent and environment before a sufﬁciently good reward is obtained.
For both environment setups and both policy types,
ST- τ achieves a higher averaged cumulative reward given lower sample complexity.
Moving from sampling to a greedy policy, ST- τ performance slightly drops.
This is easily explained by the inherent sampling process due to the Gumbel softmax.
Interestingly, it seems it is this sampling process which allows ST- τ to exhibit a lower sampling complexity in the sampling setups compared to LSTM and VD.
In contrast, LSTM and VD show poor performance for the greedy policy setups because they can no longer explore by sampling from the softmax.
BBB consistently performs worse than the other methods and we conjecture that this is due to the much larger number of parameters of this model, leading to a worse sampling complexity.
Moving from the MDP to the POMDP, the average accumulative reward naturally drops, as the agents receive less information, but ST- τ again exhibits the best performance for both policy types.
E A BLATION STUDY ON O UT - OF -D ISTRIBUTION D ETECTION
Figure 14 depicts the results of an ablation study focusing on the number of states k and whether or not the temperature τ was learned.
The results are for the two IMDB Out-of-Distribution Detection (OOD) tasks from section 4.3 .
The results indicate that a smaller number of states is sufﬁcient to capture the ST- τ model’s uncertainty on out-of-distribution data.
Especially when the temperature parameter is learned during training (the green, solid line), ST- τ shows the best results.
Increasing the number of states of ST- τ , gives the model more capacity to model uncertainty.
For 50 states, ﬁxing the temperature to a constant values works better but does not reach the accuracy of STτ models with fewer states and learned temperature.